code: def main() -> None:
      """
      An example for initializing a torchrec sharded embedding bag with a
      pretrained embedding weight.
      Environment assumptions:
      - The embedding weight fits in the RAM of a single host, but may OOM if all
        processes on the host load the embedding weight simultaneously.
      - For simplicity, the demo assumes a single-host, multi-process environment.
      """
      dist.init_process_group(backend="nccl")
      local_rank = int(os.environ["LOCAL_RANK"])
      device = torch.device(f"cuda:{local_rank}")
      torch.cuda.set_device(device)
      pg = dist.group.WORLD
      assert pg is not None
      dist.barrier()
  
      if dist.get_rank() == 0:
          # Pretend that we are loading the pretrained embedding weight from a parquet file on rank 0.
          emb = torch.rand(2000000, 64)
          # Share the tensor to local peers via shared memory
          emb = share_tensor_via_shm(tensor=emb)
      else:
          # Received the tensor shared by rank 0 via shared memory
          emb = share_tensor_via_shm(tensor=None)
  
      assert emb.is_shared()
  
      # For demo purpose, the entire model is an embedding bag collection with a
      # single embedding bag.
      ebc = EmbeddingBagCollection(
          device=torch.device("meta"),
          tables=[
              torchrec.EmbeddingBagConfig(
                  name="emb",
                  embedding_dim=64,
                  num_embeddings=2000000,
                  feature_names=["f"],
                  pooling=torchrec.PoolingType.SUM,
              )
          ],
      )
  
      # Create a rowwise sharding plan
      sharders = cast(
          List[ModuleSharder[nn.Module]],
          [
              EmbeddingBagCollectionSharder(
                  fused_params={
                      "optimizer": EmbOptimType.EXACT_ROWWISE_ADAGRAD,
                      "learning_rate": 0.01,
                      "eps": 0.01,
                  }
              )
          ],
      )
      plan = EmbeddingShardingPlanner(
          topology=Topology(
              world_size=dist.get_world_size(),
              compute_device=device.type,
          ),
          constraints={
              "emb": ParameterConstraints(
                  sharding_types=[ShardingType.ROW_WISE.value]
              )
          },
      ).collective_plan(
          ebc,
          sharders,
          pg,
      )
      print(plan)
  
      # Initialize dmp which shards the embedding bag
      dmp = DistributedModelParallel(
          module=ebc,
          device=device,
          plan=plan,
          sharders=sharders,
      )
      print(
          "Finished initializing DistributedModelParallel. "
          f"Current device utilization: {torch.cuda.memory_allocated() / 1_000_000} MB"
      )
  
      # For each shard in sharded tensors, load from the corresponding slice from
      # the pretrained weights in shared memory.
      for rank in range(dist.get_world_size()):
          if dist.get_rank() == rank:
              for _, t in dmp.state_dict().items():
                  for shard in t.local_shards():
                      offsets = shard.metadata.shard_offsets
                      lengths = shard.metadata.shard_sizes
                      src = emb[
                          offsets[0] : offsets[0] + lengths[0],
                          offsets[1] : offsets[1] + lengths[1],
                      ]
                      shard.tensor.copy_(src)
              dist.barrier()
          else:
              dist.barrier()
  
imports: import copyreg
  import io
  import os
  import pickle
  import uuid
  from typing import cast, List, Optional
  import torch
  import torch.distributed
  import torch.distributed.launcher
  import torchrec
  from fbgemm_gpu.split_embedding_configs import EmbOptimType
  from torch import nn
  from torch.multiprocessing.reductions import reduce_storage, reduce_typed_storage, reduce_typed_storage_child
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
  from torchrec.distributed.planner.types import ParameterConstraints
  from torchrec.distributed.types import ModuleSharder, ShardingType
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types