code: def start_sparse_adam_worker(
      rank,
      device,
      world_size,
      weight,
      tensor_dev="cpu",
      has_zero_grad=False,
      backend="gloo",
      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"
      )
  
      if device.type == "cuda":
          th.cuda.set_device(device)
  
      th.distributed.init_process_group(
          backend=backend,
          init_method=dist_init_method,
          world_size=world_size,
          rank=rank,
      )
  
      init_weight = th.empty((num_embs, emb_dim))
      th.manual_seed(0)
      th.nn.init.uniform_(init_weight, -1.0, 1.0)
      dgl_emb = NodeEmbedding(
          num_embs,
          emb_dim,
          "test",
          init_func=initializer,
          device=tensor_dev,
      )
      dgl_emb.all_set_embedding(init_weight)
  
      if has_zero_grad:
          dgl_emb_zero = NodeEmbedding(
              num_embs,
              emb_dim,
              "zero",
              init_func=initializer,
              device=tensor_dev,
          )
          dgl_adam = SparseAdam(params=[dgl_emb, dgl_emb_zero], lr=0.01)
      else:
          dgl_adam = SparseAdam(params=[dgl_emb], 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,)).to(tensor_dev)
      else:
          idx = th.randint(0, num_embs, size=(4,)).to(tensor_dev)
      dgl_value = dgl_emb(idx, device)
      labels = th.ones((4,)).long().to(device)
      dgl_loss = th.nn.functional.cross_entropy(dgl_value, labels)
      dgl_adam.zero_grad()
      dgl_loss.backward()
      dgl_adam.step()
      th.distributed.barrier()
      dgl_weight = dgl_emb.all_get_embedding().detach()
      after_step = dgl_emb(idx, device).cpu()
  
      if rank == 0:
          dgl_value = dgl_value.detach().cpu()
          assert F.allclose(dgl_value, after_step) is False
          weight[:] = dgl_weight[:]
      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 6 for calling a global function with example usage: start_sparse_adam_worker(rank, device, world_size, weight, tensor_dev, has_zero_grad, backend, num_embs, emb_dim, zero_comm) without return types