diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 020f180920a4..b123d91198be 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -4637,7 +4637,7 @@ def _load_checkpoint(self, autoep_partitioned_experts = False allowed_missing_keys = None if self.zero_optimization_partition_weights() and not load_optimizer_states and not self.has_moe_layers: - checkpoint['module'] = get_fp32_state_dict_from_zero_checkpoint(load_dir) + checkpoint['module'] = get_fp32_state_dict_from_zero_checkpoint(load_dir, tag=tag) fetch_z3_params = True if is_pipe_parallel: diff --git a/tests/unit/checkpoint/test_zero_optimizer.py b/tests/unit/checkpoint/test_zero_optimizer.py index 3611c21f67b6..90fd82e900c4 100644 --- a/tests/unit/checkpoint/test_zero_optimizer.py +++ b/tests/unit/checkpoint/test_zero_optimizer.py @@ -20,6 +20,69 @@ import pytest +class TestZeROCheckpointTag(DistributedTest): + world_size = [1, 2] + + @pytest.mark.parametrize('tag,save_latest,load_optimizer_states,load_module_only', [ + ('earlier', True, False, False), + ('earlier', False, False, False), + ('earlier', True, False, True), + ('earlier', False, False, True), + ('earlier', True, True, False), + (None, True, False, False), + ]) + def test_load_requested_tag(self, tmpdir, tag, save_latest, load_optimizer_states, load_module_only): + config = { + 'train_micro_batch_size_per_gpu': 1, + 'zero_allow_untested_optimizer': True, + 'zero_optimization': { + 'stage': 3, + 'reduce_bucket_size': 1000, + 'stage3_prefetch_bucket_size': 1000, + }, + } + + def make_engine(): + model = torch.nn.Linear(4, 2) + optimizer = torch.optim.Adam(model.parameters(), lr=0.1) + return deepspeed.initialize(model=model, optimizer=optimizer, config=config)[0] + + def snapshot(engine): + with deepspeed.zero.GatheredParameters(list(engine.module.parameters())): + return {name: value.detach().clone() for name, value in engine.module.state_dict().items()} + + source = make_engine() + target = None + try: + snapshots = {} + for checkpoint_tag in ('earlier', 'later'): + loss = source(torch.ones(1, 4, device=source.device)).square().mean() + source.backward(loss) + source.step() + snapshots[checkpoint_tag] = snapshot(source) + source.save_checkpoint(tmpdir, + tag=checkpoint_tag, + client_state={'label': checkpoint_tag}, + save_latest=save_latest) + + assert any(not torch.equal(snapshots['earlier'][name], value) + for name, value in snapshots['later'].items()) + target = make_engine() + load_path, client_state = target.load_checkpoint(tmpdir, + tag=tag, + load_optimizer_states=load_optimizer_states, + load_module_only=load_module_only) + expected_tag = tag if tag is not None else 'later' + assert load_path is not None + assert client_state['label'] == expected_tag + for name, value in snapshot(target).items(): + torch.testing.assert_close(value, snapshots[expected_tag][name], rtol=0, atol=0) + finally: + if target is not None: + target.destroy() + source.destroy() + + class TestZeROCheckpoint(DistributedTest): world_size = 2