Spaces:
Runtime error
Runtime error
| import logging | |
| from typing import List, Optional | |
| from torch.testing._internal import common_utils | |
| logging.getLogger("torch").setLevel(logging.WARNING) | |
| from apex.transformer import parallel_state | |
| from apex.transformer.pipeline_parallel.utils import ( | |
| _reconfigure_microbatch_calculator, | |
| get_micro_batch_size, | |
| get_num_microbatches, | |
| get_current_global_batch_size, | |
| update_num_microbatches, | |
| ) | |
| from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase | |
| from apex.transformer.testing.distributed_test_base import UccDistributedTestBase | |
| logging.getLogger("apex").setLevel(logging.WARNING) | |
| class MicrobatchCalculatorTestBase: | |
| GLOBAL_BATCH_SIZE: int = 1024 | |
| MICRO_BATCH_SIZE: int = 1 | |
| def _test(self, rampup_batch_size: Optional[List[int]]) -> None: | |
| for data_parallel_size in range(1, self.world_size + 1): | |
| expected_global_batch_size = self.GLOBAL_BATCH_SIZE | |
| expected_micro_batch_size = self.MICRO_BATCH_SIZE | |
| if rampup_batch_size: | |
| expected_global_batch_size = rampup_batch_size[0] | |
| num_consumed_samples = 0 | |
| step_of_global_batch_size = rampup_batch_size[1] | |
| threshold = rampup_batch_size[2] | |
| if data_parallel_size > 1 and data_parallel_size % 2 != 0: | |
| continue | |
| if self.world_size % data_parallel_size != 0: | |
| continue | |
| msg = f"data_parallel_size: {data_parallel_size}" | |
| parallel_state.initialize_model_parallel( | |
| tensor_model_parallel_size_=self.world_size // data_parallel_size, | |
| pipeline_model_parallel_size_=1, | |
| ) | |
| self.assertEqual(data_parallel_size, parallel_state.get_data_parallel_world_size(), msg=msg) | |
| _reconfigure_microbatch_calculator( | |
| self.rank, | |
| rampup_batch_size, | |
| self.GLOBAL_BATCH_SIZE, | |
| self.MICRO_BATCH_SIZE, | |
| data_parallel_size, | |
| ) | |
| self.assertEqual(get_micro_batch_size(), expected_micro_batch_size, msg=msg) | |
| self.assertEqual(get_num_microbatches(), expected_global_batch_size / expected_micro_batch_size / data_parallel_size, msg=msg) | |
| current_global_batch_size = get_current_global_batch_size() | |
| self.assertEqual(current_global_batch_size, expected_global_batch_size, msg=msg) | |
| # Make sure `global_batch_size` equals to the final global batch size after | |
| # certain number of updates. | |
| if rampup_batch_size: | |
| update_num_microbatches(current_global_batch_size) | |
| for i in range(100): | |
| current_global_batch_size = get_current_global_batch_size() | |
| update_num_microbatches(current_global_batch_size) | |
| current_global_batch_size = get_current_global_batch_size() | |
| self.assertEqual(get_current_global_batch_size(), self.GLOBAL_BATCH_SIZE, msg=msg) | |
| parallel_state.destroy_model_parallel() | |
| def test_constant_microbatch_calculator(self): | |
| self._test(rampup_batch_size=None) | |
| def test_dynamic_microbatch_calculator(self): | |
| self._test(rampup_batch_size=[256, 128, 500]) | |
| class NcclMicrobatchCalculatorTest(MicrobatchCalculatorTestBase, NcclDistributedTestBase): pass | |
| class UccMicrobatchCalculatorTest(MicrobatchCalculatorTestBase, UccDistributedTestBase): pass | |
| if __name__ == "__main__": | |
| common_utils.run_tests() | |