@RunIf(min_gpus=2, deepspeed=True, standalone=True)
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), gpus=2, fast_dev_run=True, precision=16)
    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:
        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)
Test to ensure that with DeepSpeed Stage 3 we can collate the sharded checkpoints into a single file.
@RunIf(min_gpus=2, deepspeed=True, standalone=True)
def test_deepspeed_collate_checkpoint(tmpdir):
    
    model = BoringModel()
    trainer = Trainer(default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), gpus=2, fast_dev_run=True, precision=16)
    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:
        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)