def _configure_using_config_file(deepspeed_config):
global num_layers, PARTITION_ACTIVATIONS, CONTIGUOUS_CHECKPOINTING, \
PA_TO_CPU, SYNCHRONIZE, PROFILE_TIME
config = DeepSpeedConfig(deepspeed_config).activation_checkpointing_config
print(config.repr())
PARTITION_ACTIVATIONS = config.partition_activations
CONTIGUOUS_CHECKPOINTING = config.contiguous_memory_optimization
num_layers = config.number_checkpoints
PA_TO_CPU = config.cpu_checkpointing
SYNCHRONIZE = config.synchronize_checkpoint_boundary
PROFILE_TIME = config.profile
def _test_batch_config(num_ranks, batch, micro_batch, gas, success):
assert dist.get_world_size() == num_ranks, \
'The test assumes a world size of f{num_ranks}'
ds_batch_config = 'tests/unit/ds_batch_config.json'
ds_config = DeepSpeedConfig(ds_batch_config)
#test cases when all parameters are provided
status = _run_batch_config(ds_config,
train_batch=batch,
micro_batch=micro_batch,
gas=gas)
_batch_assert(status, ds_config, batch, micro_batch, gas, success)
#test cases when two out of three parameters are provided
status = _run_batch_config(ds_config,
train_batch=batch,
micro_batch=micro_batch)
_batch_assert(status, ds_config, batch, micro_batch, gas, success)
if success:
#when gas is provided with one more parameter
status = _run_batch_config(ds_config, train_batch=batch, gas=gas)
_batch_assert(status, ds_config, batch, micro_batch, gas, success)
status = _run_batch_config(ds_config,
micro_batch=micro_batch,
gas=gas)
_batch_assert(status, ds_config, batch, micro_batch, gas, success)
#test the case when only micro_batch or train_batch is provided
if gas == 1:
status = _run_batch_config(ds_config, micro_batch=micro_batch)
_batch_assert(status, ds_config, batch, micro_batch, gas,
success)
status = _run_batch_config(ds_config, train_batch=batch)
_batch_assert(status, ds_config, batch, micro_batch, gas,
success)
else:
#when only gas is provided
status = _run_batch_config(ds_config, gas=gas)
_batch_assert(status, ds_config, batch, micro_batch, gas, success)
#when gas is provided with something else and gas does not divide batch
if gas != 1:
status = _run_batch_config(ds_config,
train_batch=batch,
gas=gas)
_batch_assert(status, ds_config, batch, micro_batch, gas,
success)
def _configure_with_arguments(self, args, mpu):
self.local_rank = args.local_rank if hasattr(args, 'local_rank') else 0
self._config = DeepSpeedConfig(args.deepspeed_config, mpu)
