code: def _build_key_size_numel_dictionaries(keys, data):
      """Build the size on rank 0 and broadcast."""
      max_dim = _MAX_DATA_DIM
      sizes = [0 for _ in range(max_dim) for _ in keys]
  
      # Pack the sizes on rank zero.
      if (
          not gpc.is_initialized(ParallelMode.TENSOR)
          or gpc.get_local_rank(ParallelMode.TENSOR) == 0
      ):
          offset = 0
          for key in keys:
              assert (
                  data[key].dim() < max_dim
              ), "you should increase MAX_DATA_DIM"
              size = data[key].size()
              for i, s in enumerate(size):
                  sizes[i + offset] = s
              offset += max_dim
  
      # Move to GPU and broadcast.
      sizes_cuda = torch.cuda.LongTensor(sizes)
      torch.distributed.broadcast(
          sizes_cuda,
          gpc.get_ranks_in_group(ParallelMode.TENSOR)[0],
          group=gpc.get_group(ParallelMode.TENSOR),
      )
  
      # Move back to cpu and unpack.
      sizes_cpu = sizes_cuda.cpu()
      key_size = {}
      key_numel = {}
      total_numel = 0
      offset = 0
      for key in keys:
          i = 0
          size = []
          numel = 1
          while sizes_cpu[offset + i] > 0:
              this_size = sizes_cpu[offset + i]
              size.append(this_size)
              numel *= this_size
              i += 1
          key_size[key] = size
          key_numel[key] = numel
          total_numel += numel
          offset += max_dim
  
      return key_size, key_numel, total_numel
  
file_path: colossalai/examples/tutorial/sequence_parallel/data/bert_helper.py
imports: import torch
  from colossalai.legacy.context import ParallelMode
  from colossalai.legacy.core import global_context
label_desc: Use image node_id 1 for calling a global function with example usage: _build_key_size_numel_dictionaries(keys, data) and returns: key_size,  key_numel,  total_numel