code: def start_torch_adam_worker(
      rank,
      world_size,
      weight,
      has_zero_grad=False,
      num_embs=128,
      emb_dim=10,
      zero_comm=True,
  ):
      print("start sparse worker for adam {}".format(rank))
      dist_init_method = "tcp://{master_ip}:{master_port}".format(
          master_ip="127.0.0.1", master_port="12345"
      )
      backend = "gloo"
  
      th.distributed.init_process_group(
          backend=backend,
          init_method=dist_init_method,
          world_size=world_size,
          rank=rank,
      )
  
      torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
      th.manual_seed(0)
      th.nn.init.uniform_(torch_emb.weight, -1.0, 1.0)
      torch_emb = th.nn.parallel.DistributedDataParallel(torch_emb)
      if has_zero_grad:
          torch_emb_zero = th.nn.Embedding(
              num_embs, emb_dim, sparse=True
          )
          torch_emb_zero = torch_emb_zero.to(tensor_dev)
          th.manual_seed(0)
          th.nn.init.uniform_(torch_emb_zero.weight, -1.0, 1.0)
          torch_emb_zero = th.nn.parallel.DistributedDataParallel(
              torch_emb_zero
          )
          torch_adam = th.optim.SparseAdam(
              list(torch_emb.module.parameters())
              + list(torch_emb_zero.module.parameters()),
              lr=0.01,
          )
      else:
          torch_adam = th.optim.SparseAdam(
              list(torch_emb.module.parameters()), lr=0.01
          )
  
      th.manual_seed(rank)
      if zero_comm:
          start = (num_embs // world_size) * rank
          end = (num_embs // world_size) * (rank + 1)
          idx = th.randint(start, end, size=(4,))
      else:
          idx = th.randint(0, num_embs, size=(4,))
      labels = th.ones((4,)).long()
      torch_value = torch_emb(idx)
      torch_loss = th.nn.functional.cross_entropy(torch_value, labels)
      torch_adam.zero_grad()
      torch_loss.backward()
      torch_adam.step()
      th.distributed.barrier()
  
      if rank == 0:
          weight[:] = torch_emb.module.weight.cpu()[:]
      th.distributed.barrier()
  
imports: import os
  import unittest
  import backend
  import pytest
  import torch
  import torch.multiprocessing
  from dgl.nn import NodeEmbedding
  from dgl.optim import SparseAdagrad, SparseAdam
label_desc: Use image node_id 7 for calling a global function with example usage: start_torch_adam_worker(rank, world_size, weight, has_zero_grad, num_embs, emb_dim, zero_comm) without return types