def test_deepspeed_init_module_with_stages_1_2(stage):
    """Tests how `.init_module()` behaves with ZeRO stages 1 and 2."""
    strategy = DeepSpeedStrategy(stage=stage)
    trainer = Trainer(accelerator="cuda", devices=2, strategy=strategy, 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=False, remote_device=None, config_dict_or_path=ANY)
    assert model.layer.weight.dtype == torch.bfloat16
Tests how `.init_module()` behaves with ZeRO stages 1 and 2.
tests/tests_pytorch/strategies/test_deepspeed.py