code: def broadcast_data(keys, data, datatype):
      """Broadcast data from rank zero of each model parallel group to the
      members of the same model parallel group.
  
      Arguments:
          keys: list of keys in the data disctionary to be broadcasted
          data: data dictionary of string keys and cpu tensor values.
          datatype: torch data type of all tensors in data associated
                    with keys.
      """
      # Build (key, size) and (key, number of elements) dictionaries along
      # with the total number of elements on all ranks.
      (
          key_size,
          key_numel,
          total_numel,
      ) = _build_key_size_numel_dictionaries(keys, data)
      # Pack on rank zero.
      if get_tensor_model_parallel_rank() == 0:
          # Check that all keys have the same data type.
          _check_data_types(keys, data, datatype)
          # Flatten the data associated with the keys
          flatten_data = torch.cat(
              [data[key].contiguous().view(-1) for key in keys], dim=0
          ).cuda()
      else:
          flatten_data = torch.empty(
              total_numel,
              device=torch.cuda.current_device(),
              dtype=datatype,
          )
  
      # Broadcast
      torch.distributed.broadcast(
          flatten_data,
          get_tensor_model_parallel_src_rank(),
          group=get_tensor_model_parallel_group(),
      )
  
      # Unpack
      output = {}
      offset = 0
      for key in keys:
          size = key_size[key]
          numel = key_numel[key]
          output[key] = flatten_data.narrow(0, offset, numel).view(size)
          offset += numel
  
      return output
  
imports: import torch
  from apex.transformer.parallel_state import get_tensor_model_parallel_group
  from apex.transformer.parallel_state import get_tensor_model_parallel_rank
  from apex.transformer.parallel_state import get_tensor_model_parallel_src_rank
label_desc: Use image node_id 3 for calling a global function with example usage: broadcast_data(keys, data, datatype) and returns: output