def test_module_init_context(distributed, compile, precision, expected_dtype, tmp_path):
    """Test that the module under the init-context gets moved to the right device and dtype."""

    class Model(FSDP2Model):
        def on_train_start(self):
            assert self.model.w1.weight.device == torch.device("cuda", self.local_rank)
            assert self.model.w1.weight.dtype == expected_dtype
            optimizer = self.optimizers(use_pl_optimizer=False)
            assert optimizer.param_groups[0]["params"][0].device.type == "cuda"

    def _run_setup_assertions(empty_init, expected_device):
        trainer = Trainer(
            default_root_dir=tmp_path,
            accelerator="cuda",
            devices=2,
            strategy=ModelParallelStrategy(),
            precision=precision,
            max_steps=1,
            barebones=True,
            enable_checkpointing=False,
            logger=False,
        )
        with trainer.init_module(empty_init=empty_init):
            model = Model(compile=compile)

        # The model is on the CPU/meta-device until after `ModelParallelStrategy.setup()`
        assert model.model.w1.weight.device == expected_device
        assert model.model.w1.weight.dtype == expected_dtype
        trainer.fit(model)

    # Case 1: No empty init
    _run_setup_assertions(empty_init=False, expected_device=torch.device("cpu"))

    # Case 2: Empty-init with meta device
    _run_setup_assertions(empty_init=True, expected_device=torch.device("meta"))
Test that the module under the init-context gets moved to the right device and dtype.
tests/tests_pytorch/strategies/test_model_parallel_integration.py