def test_deepspeed_multigpu(tmp_path):
    """Test to ensure that DeepSpeed with multiple GPUs works and deepspeed distributed is initialized correctly."""
    import deepspeed

    model = BoringModel()
    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,
    )

    with mock.patch.object(
        model, "configure_optimizers", wraps=model.configure_optimizers
    ) as mock_configure_optimizers:
        trainer.test(model)
    assert mock_configure_optimizers.call_count == 0

    with mock.patch("deepspeed.init_distributed", wraps=deepspeed.init_distributed) as mock_deepspeed_distributed:
        trainer.fit(model)
    mock_deepspeed_distributed.assert_called_once()

    _assert_save_model_is_equal(model, tmp_path, trainer)
Test to ensure that DeepSpeed with multiple GPUs works and deepspeed distributed is initialized correctly.
tests/tests_pytorch/strategies/test_deepspeed.py