def test_deepspeed_collate_checkpoint(tmpdir): """Test to ensure that with DeepSpeed Stage 3 we can collate the sharded checkpoints into a single file.""" model = BoringModel() trainer = Trainer(default_root_dir=tmpdir, plugins=[DeepSpeedPlugin(stage=3)], gpus=2, fast_dev_run=True, precision=16) trainer.fit(model) checkpoint_path = os.path.join(tmpdir, "model.pt") checkpoint_path = trainer.accelerator.broadcast(checkpoint_path) trainer.save_checkpoint(checkpoint_path) trainer.accelerator.barrier() if trainer.is_global_zero: # ensure function call works output_path = os.path.join(tmpdir, "single_model.pt") convert_zero_checkpoint_to_fp32_state_dict(checkpoint_path, output_path) _assert_checkpoint_equal(model, output_path)
def test_deepspeed_collate_checkpoint(tmpdir): """Test to ensure that with DeepSpeed Stage 3 we can collate the sharded checkpoints into a single file.""" model = BoringModel() trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), accelerator="gpu", devices=2, fast_dev_run=True, precision=16, enable_progress_bar=False, enable_model_summary=False, ) trainer.fit(model) checkpoint_path = os.path.join(tmpdir, "model.pt") checkpoint_path = trainer.strategy.broadcast(checkpoint_path) trainer.save_checkpoint(checkpoint_path) trainer.strategy.barrier() if trainer.is_global_zero: # ensure function call works output_path = os.path.join(tmpdir, "single_model.pt") convert_zero_checkpoint_to_fp32_state_dict(checkpoint_path, output_path) _assert_checkpoint_equal(model, output_path)