def test_deepspeed_multigpu_stage_3(tmp_path):
    """Test to ensure ZeRO Stage 3 works with a parallel model."""
    model = ModelParallelBoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        strategy=DeepSpeedStrategy(stage=3),
        accelerator="gpu",
        devices=2,
        fast_dev_run=True,
        precision="16-mixed",
        enable_progress_bar=False,
        enable_model_summary=False,
    )
    trainer.test(model)
    trainer.fit(model)

    _assert_save_model_is_equal(model, tmp_path, trainer)
Test to ensure ZeRO Stage 3 works with a parallel model.
tests/tests_pytorch/strategies/test_deepspeed.py