code: def _shard_modules(  # noqa: C901
      module: nn.Module,
      env: Optional[ShardingEnv] = None,
      device: Optional[torch.device] = None,
      plan: Optional[ShardingPlan] = None,
      sharders: Optional[List[ModuleSharder[torch.nn.Module]]] = None,
      init_params: Optional[bool] = False,
  ) -> nn.Module:
      """
      See shard_modules
      """
  
      torch._C._log_api_usage_once("torchrec.distributed.shard_modules")
  
      if sharders is None:
          sharders = get_default_sharders()
  
      if env is None:
          pg = dist.GroupMember.WORLD
          assert pg is not None, "Process group is not initialized"
          env = ShardingEnv.from_process_group(pg)
  
      if device is None:
          if torch.cuda.is_available():
              device = torch.device(torch.cuda.current_device())
          else:
              device = torch.device("cpu")
  
      sharder_map: Dict[Type[nn.Module], ModuleSharder[nn.Module]] = {
          sharder.module_type: sharder for sharder in sharders
      }
  
      if plan is None:
          planner = EmbeddingShardingPlanner(
              topology=Topology(
                  local_world_size=get_local_size(env.world_size),
                  world_size=env.world_size,
                  compute_device=device.type,
              )
          )
          pg = env.process_group
          if pg is not None:
              plan = planner.collective_plan(module, sharders, pg)
          else:
              plan = planner.plan(module, sharders)
  
      if type(module) in sharder_map:
          # If the top level module is itself a shardable module, return the sharded variant.
          # Note, we cannot do an inplace replacement in this case.
          return sharder_map[type(module)].shard(
              module, plan.get_plan_for_module(""), env, device
          )
  
      def _replace(_model: nn.Module, path: str = "") -> None:
          for child_name, child in _model.named_children():
              child_path = _join_module_path(path, child_name)
              if type(child) in sharder_map:
                  assert plan is not None
                  sharded_params = plan.get_plan_for_module(child_path)
                  if sharded_params is not None:
                      sharded_module = sharder_map[type(child)].shard(
                          child, sharded_params, env, device
                      )
                      _model.register_module(
                          child_name,
                          sharded_module,
                      )
              else:
                  _replace(child, child_path)
  
      _replace(module)
      if init_params and device is not None and device.type != "meta":
          init_parameters(module, device)
          module = module.to(device)
  
      return module
  
file_path: torchrec/torchrec/distributed/shard.py
imports: from typing import Callable, Dict, List, Optional, Type, Union
  import torch
  import torch.distributed
  from torch import nn
  from torch.distributed._composable.contract import contract
  from torchrec.distributed.comm import get_local_size
  from torchrec.distributed.model_parallel import get_default_sharders
  from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
  from torchrec.distributed.sharding_plan import get_module_to_default_sharders, ParameterShardingGenerator
  from torchrec.distributed.types import ModuleSharder, ModuleShardingPlan, ShardingEnv, ShardingPlan
  from torchrec.distributed.utils import init_parameters
label_desc: Use image node_id 5 for calling a global function with example usage: _shard_modules(module, env, device, plan, sharders, init_params) and returns: module, sharder_map