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

    with (
        mock.patch("deepspeed.zero.Init") as zero_init_mock,
        mock.patch("torch.Tensor.uniform_") as init_mock,
        fabric.init_module(empty_init=empty_init),
    ):
        model = BoringModel()

    zero_init_mock.assert_called_with(enabled=False, remote_device=None, config_dict_or_path=ANY)
    assert init_mock.call_count == int(not empty_init)
    assert model.layer.weight.dtype == torch.bfloat16
Tests how `.init_module()` behaves with ZeRO stages 1 and 2.
tests/tests_fabric/strategies/test_deepspeed_integration.py