code: def sharding_single_rank_test(
      rank: int,
      world_size: int,
      model_class: TestSparseNNBase,
      embedding_groups: Dict[str, List[str]],
      tables: List[EmbeddingTableConfig],
      sharders: List[ModuleSharder[nn.Module]],
      backend: str,
      optim: EmbOptimType,
      weighted_tables: Optional[List[EmbeddingTableConfig]] = None,
      constraints: Optional[Dict[str, ParameterConstraints]] = None,
      local_size: Optional[int] = None,
      qcomms_config: Optional[QCommsConfig] = None,
      apply_optimizer_in_backward_config: Optional[
          Dict[str, Tuple[Type[torch.optim.Optimizer], Dict[str, Any]]]
      ] = None,
      variable_batch_size: bool = False,
      batch_size: int = 4,
      feature_processor_modules: Optional[
          Dict[str, torch.nn.Module]
      ] = None,
      variable_batch_per_feature: bool = False,
      global_constant_batch: bool = False,
  ) -> None:
      with MultiProcessContext(
          rank, world_size, backend, local_size
      ) as ctx:
          # Generate model & inputs.
          (global_model, inputs) = gen_model_and_input(
              model_class=model_class,
              tables=tables,
              generate=cast(
                  VariableBatchModelInputCallable,
                  ModelInput.generate_variable_batch_input,
              )
              if variable_batch_per_feature
              else ModelInput.generate,
              weighted_tables=weighted_tables,
              embedding_groups=embedding_groups,
              world_size=world_size,
              num_float_features=16,
              variable_batch_size=variable_batch_size,
              batch_size=batch_size,
              feature_processor_modules=feature_processor_modules,
              global_constant_batch=global_constant_batch,
          )
          global_model = global_model.to(ctx.device)
          global_input = inputs[0][0].to(ctx.device)
          local_input = inputs[0][1][rank].to(ctx.device)
  
          # Shard model.
          local_model = model_class(
              tables=cast(List[BaseEmbeddingConfig], tables),
              weighted_tables=cast(
                  List[BaseEmbeddingConfig], weighted_tables
              ),
              embedding_groups=embedding_groups,
              dense_device=ctx.device,
              sparse_device=torch.device("meta"),
              num_float_features=16,
              feature_processor_modules=feature_processor_modules,
          )
  
          global_model_named_params_as_dict = dict(
              global_model.named_parameters()
          )
          local_model_named_params_as_dict = dict(
              local_model.named_parameters()
          )
  
          if apply_optimizer_in_backward_config is not None:
              for apply_optim_name, (
                  optimizer_type,
                  optimizer_kwargs,
              ) in apply_optimizer_in_backward_config.items():
                  for (
                      name,
                      param,
                  ) in global_model_named_params_as_dict.items():
                      if name not in apply_optim_name:
                          continue
                      assert name in local_model_named_params_as_dict
                      local_param = local_model_named_params_as_dict[
                          name
                      ]
                      apply_optimizer_in_backward(
                          optimizer_type, [param], optimizer_kwargs
                      )
                      apply_optimizer_in_backward(
                          optimizer_type,
                          [local_param],
                          optimizer_kwargs,
                      )
  
          planner = EmbeddingShardingPlanner(
              topology=Topology(
                  world_size,
                  ctx.device.type,
                  local_world_size=ctx.local_size,
              ),
              constraints=constraints,
          )
          plan: ShardingPlan = planner.collective_plan(
              local_model, sharders, ctx.pg
          )
          """
          Simulating multiple nodes on a single node. However, metadata information and
          tensor placement must still be consistent. Here we overwrite this to do so.
  
          NOTE:
              inter/intra process groups should still behave as expected.
  
          TODO: may need to add some checks that only does this if we're running on a
          single GPU (which should be most cases).
          """
  
          for group in plan.plan:
              for _, parameter_sharding in cast(
                  EmbeddingModuleShardingPlan, plan.plan[group]
              ).items():
                  if (
                      parameter_sharding.sharding_type
                      in {
                          ShardingType.TABLE_ROW_WISE.value,
                          ShardingType.TABLE_COLUMN_WISE.value,
                      }
                      and ctx.device.type != "cpu"
                  ):
                      sharding_spec = parameter_sharding.sharding_spec
                      if sharding_spec is not None:
                          # pyre-ignore
                          for shard in sharding_spec.shards:
                              placement = shard.placement
                              rank: Optional[int] = placement.rank()
                              assert rank is not None
                              shard.placement = (
                                  torch.distributed._remote_device(
                                      f"rank:{rank}/cuda:{rank}"
                                  )
                              )
  
          local_model = DistributedModelParallel(
              local_model,
              env=ShardingEnv.from_process_group(ctx.pg),
              plan=plan,
              sharders=sharders,
              device=ctx.device,
          )
  
          dense_optim = KeyedOptimizerWrapper(
              dict(
                  in_backward_optimizer_filter(
                      local_model.named_parameters()
                  )
              ),
              lambda params: torch.optim.SGD(params, lr=0.1),
          )
          local_opt = CombinedOptimizer(
              [local_model.fused_optimizer, dense_optim]
          )
  
          # Load model state from the global model.
          copy_state_dict(
              local_model.state_dict(), global_model.state_dict()
          )
  
          # Run a single training step of the sharded model.
          local_pred = gen_full_pred_after_one_step(
              local_model, local_opt, local_input
          )
  
          all_local_pred = []
          for _ in range(world_size):
              all_local_pred.append(torch.empty_like(local_pred))
          dist.all_gather(all_local_pred, local_pred, group=ctx.pg)
  
          # Run second training step of the unsharded model.
          assert optim == EmbOptimType.EXACT_SGD
          global_opt = torch.optim.SGD(
              global_model.parameters(), lr=0.1
          )
  
          global_pred = gen_full_pred_after_one_step(
              global_model, global_opt, global_input
          )
  
          # Compare predictions of sharded vs unsharded models.
          if qcomms_config is None:
              torch.testing.assert_close(
                  global_pred, torch.cat(all_local_pred)
              )
          else:
              # With quantized comms, we can relax constraints a bit
              rtol = 0.003
              if CommType.FP8 in [
                  qcomms_config.forward_precision,
                  qcomms_config.backward_precision,
              ]:
                  rtol = 0.05
              atol = global_pred.max().item() * rtol
              torch.testing.assert_close(
                  global_pred,
                  torch.cat(all_local_pred),
                  rtol=rtol,
                  atol=atol,
              )
  
imports: from enum import Enum
  from typing import Any, cast, Dict, List, Optional, Protocol, Tuple, Type, Union
  import torch
  import torch.distributed
  import torch.nn
  from fbgemm_gpu.split_embedding_configs import EmbOptimType
  from torch.distributed.optim import _apply_optimizer_in_backward
  from torchrec.distributed.embedding_types import EmbeddingTableConfig
  from torchrec.distributed.fbgemm_qcomm_codec import CommType, get_qcomm_codecs_registry, QCommsConfig
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.planner import EmbeddingShardingPlanner, ParameterConstraints, Topology
  from torchrec.distributed.test_utils.multi_process import MultiProcessContext
  from torchrec.distributed.test_utils.test_model import ModelInput, TestEBCSharder, TestEBSharder, TestETCSharder, TestETSharder, TestSparseNNBase
  from torchrec.distributed.types import EmbeddingModuleShardingPlan, ModuleSharder, ShardedTensor, ShardingEnv, ShardingPlan, ShardingType
  from torchrec.modules.embedding_configs import BaseEmbeddingConfig, EmbeddingBagConfig
  from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizerWrapper
  from torchrec.optim.optimizers import in_backward_optimizer_filter
label_desc: Use image node_id 4 for calling a global function with example usage: sharding_single_rank_test(rank, world_size, model_class, embedding_groups, tables, sharders, backend, optim, weighted_tables, constraints, local_size, qcomms_config, apply_optimizer_in_backward_config, variable_batch_size, batch_size, feature_processor_modules, variable_batch_per_feature, global_constant_batch) without return types