class_docstr: Tests ``fully_shard`` initialization.
class_docstr_tok: ["Tests","``","fully_shard","``","initialization","."]
docstr: Tests materializing a meta-device module.
docstr_tok: ["Tests","materializing","a","meta-device","module","."]
code: def test_materialize_meta_module(self):
      """Tests materializing a meta-device module."""
  
      def _param_init_fn(module: nn.Module):
          """
          This is an example ``param_init_fn`` for composable FSDP.
  
          TODO: This function is not satisfactory because:
          (1) This requires guarding with ``_is_fsdp_flattened()``. This
          guard is needed to avoid re-initializing parameters for nested
          cases since some initialization methods strictly require non-1D
          shape (e.g. ``kaiming_uniform_()``), while FSDP replaces the
          original parameters with their 1D shards.
          (2) This requires module-by-module traversal and manual ``setattr``
          usage as opposed to first calling ``module.to_empty()`` and then
          initializing each parameter after. The latter will override the
          initialization of already-initialized nested parameters. In other
          words, this parameter initialization function must strictly modify
          only the parameters on meta device.
          """
          torch.manual_seed(0)
          for submodule in module.modules():
              for param_name, param in submodule.named_parameters(
                  recurse=False
              ):
                  if not _is_fsdp_flattened(param) and param.is_meta:
                      materialized_param = nn.Parameter(
                          torch.empty_like(
                              param, device=torch.device("cuda")
                          )
                      )
                      nn.init.uniform_(materialized_param)
                      setattr(submodule, param_name, materialized_param)
  
      composable_module = CompositeParamModel(
          device=torch.device("meta")
      )
      meta_model = CompositeParamModel(device=torch.device("meta"))
      fsdp_wrapped_model = FSDP(
          meta_model,
          auto_wrap_policy=ModuleWrapPolicy({UnitModule}),
          param_init_fn=_param_init_fn,
          use_orig_params=True,
      )
      fully_shard(
          composable_module,
          policy=ModuleWrapPolicy({UnitModule}),
          param_init_fn=_param_init_fn,
      )
      for (
          (composable_param_name, composable_param),
          (fsdp_wrapped_param_name, fsdp_wrapped_param),
      ) in zip(
          composable_module.named_parameters(),
          fsdp_wrapped_model.named_parameters(),
      ):
          self.assertEqual(
              composable_param_name,
              clean_tensor_name(fsdp_wrapped_param_name),
          )
          self.assertEqual(
              composable_param.device,
              torch.device("cuda", torch.cuda.current_device()),
          )
          self.assertEqual(composable_param, fsdp_wrapped_param)
  
imports: import copy
  import sys
  from typing import Optional
  import torch
  import torch.distributed
  import torch.distributed.fsdp._traversal_utils
  import torch.nn
  from torch.distributed._composable import fully_shard
  from torch.distributed.fsdp import BackwardPrefetch, FullyShardedDataParallel
  from torch.distributed.fsdp._common_utils import _is_fsdp_flattened, clean_tensor_name
  from torch.distributed.fsdp.wrap import _Policy, CustomPolicy, ModuleWrapPolicy
  from torch.testing._internal.common_dist_composable import CompositeParamModel, FakeSequential, NestedSequentialModel, UnitModule
  from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
  from torch.testing._internal.common_fsdp import FSDPTest
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 8 for calling the TestInitialization obj's underlying member method code with example usage: obj.test_materialize_meta_module() without return types