code: def main(
      rank: int,
      world_size: int,
      init_method: str = "tcp://localhost:23456",
  ):
      """
      #### Run the training on the node with rank `rank`.
      """
      # Initialize PyTorch distributed process group
      with monit.section("Distributed"):
          torch.distributed.init_process_group(
              "nccl",
              timeout=datetime.timedelta(seconds=30),
              init_method=init_method,
              rank=rank,
              world_size=world_size,
          )
  
      # Set current device
      device = torch.device(f"cuda:{rank}")
      torch.cuda.set_device(device)
  
      # Create the experiment
      experiment.create(
          name="zero3_neox",
          writers={"screen", "labml"},
          distributed_world_size=world_size,
          distributed_rank=rank,
      )
  
      # Create configurations
      conf = Configs()
  
      # Load configurations
      experiment.configs(
          conf,
          {
              "model": "Zero3",
              "optimizer": "Zero3Adam",
              "device": device,
              "rank": rank,
              "world_size": world_size,
              "learning_rate": 3e-4,
              "max_seq_len": 128,
              "batch_size": 16,
          },
      )
  
      # Start the experiment
      with experiment.start():
          # Initialize the model. Do this before the loop for cleaner logs.
          _ = conf.model
  
          # Train the model
          for epoch in monit.loop(conf.epochs):
              conf.train_epoch()
              tracker.new_line()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/scaling/zero3/finetune_neox.py
imports: import datetime
  import torch
  import torch.distributed
  from labml import experiment, monit, tracker
  from labml.configs import option
  from labml.logger import inspect
  from labml_nn.neox.samples.finetune import PipelineParallelTrainerConf
label_desc: Use image node_id 3 for calling a global function with example usage: main(rank, world_size, init_method) without return types