def test_deepspeed_init_module_with_stage_3():
    """Tests how `.init_module()` behaves with ZeRO stage 3."""
    trainer = Trainer(
        accelerator="cuda", devices=2, strategy="deepspeed_stage_3", precision="bf16-mixed", fast_dev_run=1
    )
    model = ModelParallelBoringModel()
    with mock.patch("deepspeed.zero.Init") as zero_init_mock:
        trainer.fit(model)

    zero_init_mock.assert_called_once_with(enabled=True, remote_device=None, config_dict_or_path=ANY)
Tests how `.init_module()` behaves with ZeRO stage 3.
tests/tests_pytorch/strategies/test_deepspeed.py