def test_deepspeed_init_module_with_stage_3(empty_init):
    """Tests how `.init_module()` behaves with ZeRO stage 3."""
    strategy = DeepSpeedStrategy(stage=3)
    fabric = Fabric(accelerator="cuda", devices=2, strategy=strategy, precision="bf16-true")
    fabric.launch()

    with mock.patch("deepspeed.zero.Init") as zero_init_mock, fabric.init_module(empty_init=empty_init):
        BoringModel()
    fabric.barrier()
    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_fabric/strategies/test_deepspeed_integration.py