code: def main(argv: List[str]) -> None:
      """
      Trains, validates, and tests a Bert4Rec Model
      (https://arxiv.org/abs/1904.06690). The Bert4Rec model contains both data parallel
      components (e.g. transformation blocks) and model parallel
      components (e.g. item embeddings). The Bert4Rec model is pipelined so that dataloading,
      data-parallel to model-parallel comms, and forward/backward are overlapped. Can be
      run with either a random dataloader or the movielens dataset
      (https://grouplens.org/datasets/movielens/).
  
      Args:
          argv (List[str]): command line args.
  
      Returns:
          None.
      """
  
      args = parse_args(argv)
      use_dmp = args.mode == "dmp"
      metric_ks: List[int] = (
          [1, 5, 10, 20, 50, 100]
          if args.dataset_name != "random"
          else [1, 5, 10]
      )
      rank = int(os.environ["LOCAL_RANK"])
      if torch.cuda.is_available():
          device = torch.device(f"cuda:{rank}")
          backend = "nccl"
          torch.cuda.set_device(device)
      else:
          device = torch.device("cpu")
          backend = "gloo"
  
      if not torch.distributed.is_initialized():
          dist.init_process_group(backend=backend)
  
      world_size = dist.get_world_size()
  
      raw_data = get_raw_dataframe(
          args.dataset_name,
          args.random_user_count,
          args.random_item_count,
          args.random_size,
          args.min_rating,
          args.dataset_path,
      )
  
      df = Bert4RecPreprocsser(
          raw_data,
          args.min_rating,
          args.min_user_count,
          args.min_item_count,
          args.dataset_name,
          args.max_len,
          args.mask_prob,
          args.dupe_factor,
      ).get_processed_dataframes()
      # 0 for padding, item_count + 1 for mask
      vocab_size = len(df["smap"]) + 2
      bert4recDataloader = Bert4RecDataloader(
          df,
          args.train_batch_size,
          args.val_batch_size,
          args.test_batch_size,
      )
      (
          train_loader,
          val_loader,
          test_loader,
      ) = bert4recDataloader.get_pytorch_dataloaders(rank, world_size)
      model_bert4rec = BERT4Rec(
          vocab_size,
          args.max_len,
          emb_dim=args.emb_dim,
          nhead=args.nhead,
          num_layers=args.num_layers,
      ).to(device)
      if use_dmp:
          fused_params: Dict[str, Any] = {}
          fused_params["optimizer"] = EmbOptimType.ADAM
          fused_params["learning_rate"] = args.lr
          fused_params["weight_decay"] = args.weight_decay
          model = DMP(
              module=model_bert4rec,
              device=device,
              sharders=[
                  cast(
                      ModuleSharder[nn.Module],
                      EmbeddingCollectionSharder(fused_params),
                  )
              ],
          )
          dense_optimizer = KeyedOptimizerWrapper(
              dict(
                  in_backward_optimizer_filter(model.named_parameters())
              ),
              lambda params: optim.Adam(
                  params, lr=args.lr, weight_decay=args.weight_decay
              ),
          )
  
          optimizer = CombinedOptimizer(
              [model.fused_optimizer, dense_optimizer]
          )
      else:
          device_ids = [rank] if backend == "nccl" else None
          """
          Another way to do DDP is to specify the sharding_type for TorchRec as Data_parallel
          Here we provide an example of how to do it:
          First we constraint the sharding_types to only use data_parallel in sharder,
          then we use DMP to wrap it:
  
              sharding_types = [ShardingType.DATA_PARALLEL.value]
              constraints[
                  "item_embedding"
              ] = torchrec.distributed.planner.ParameterConstraints(sharding_types=sharding_types)
              sharders = [
                  cast(ModuleSharder[nn.Module], EmbeddingCollectionSharder(fused_params))
              ]
              pg = dist.GroupMember.WORLD
              model = DMP(
                  module=model_bert4rec,
                  device=device,
                  plan=torchrec.distributed.planner.EmbeddingShardingPlanner(
                  topology=torchrec.distributed.planner.Topology(
                      world_size=world_size,
                      compute_device=device.type,
                  ),
                  constraints=constraints
              ).collective_plan(model_bert4rec, sharders, pg),
                  sharders=sharders,
              )
  
          """
          model = DDP(model_bert4rec, device_ids=device_ids)
          optimizer = optim.Adam(
              model.parameters(),
              lr=args.lr,
              weight_decay=args.weight_decay,
          )
      lr_scheduler = optim.lr_scheduler.StepLR(
          optimizer, step_size=args.decay_step, gamma=args.gamma
      )
  
      train_val_test(
          model,
          train_loader,
          val_loader,
          test_loader,
          device,
          # pyre-fixme[6]: For 6th param expected `Adam` but got
          #  `Union[CombinedOptimizer, Adam]`.
          optimizer,
          lr_scheduler,
          args.num_epochs,
          metric_ks,
          args.export_root,
      )
  
imports: import argparse
  import os
  import sys
  from typing import Any, cast, Dict, List, Union
  import numpy
  import torch
  import torch.nn
  import torch.optim
  import torch.utils.data
  from fbgemm_gpu.split_embedding_configs import EmbOptimType
  from torch import distributed
  from torch.nn.parallel import DistributedDataParallel
  from torchrec.distributed.embedding import EmbeddingCollectionSharder
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.types import ModuleSharder
  from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizerWrapper
  from torchrec.optim.optimizers import in_backward_optimizer_filter
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
  from tqdm import tqdm
label_desc: Use image node_id 8 for calling a global function with example usage: main(argv) without return types