code: def _allreduce_word_embedding_grads(
      model: List[torch.nn.Module], config: TransformerConfig
  ):
      """
      All-reduce word embedding grads.
  
      Reduce grads across first and last stages to ensure that word_embeddings parameters stay in
      sync. This should only run for models that support pipelined model parallelism (BERT and GPT).
      """
  
      if (
          parallel_state.is_rank_in_embedding_group(ignore_virtual=True)
          and parallel_state.get_pipeline_model_parallel_world_size()
          > 1
      ):
          if parallel_state.is_pipeline_first_stage(
              ignore_virtual=True
          ):
              model_module = model[0]
          elif parallel_state.is_pipeline_last_stage(
              ignore_virtual=True
          ):
              model_module = model[-1]
          else:  # We do not support the interleaved schedule for T5 yet.
              model_module = model[0]
  
          # Look for module with 'pre_process' attribute to get around the fact that DDP and
          # other wrapper classes inherit from non-core MegatronModule that has
          # 'share_embeddings_and_output_weights' and 'shared_embedding_or_output_weight'
          # attributes already, causing get_attr_wrapped_model() to not unwrap anything here.
          # TODO: Clean this up once the wrapper classes inherit from core MegatronModule.
          model_module = get_attr_wrapped_model(
              model_module, "pre_process", return_model_obj=True
          )
          if model_module.share_embeddings_and_output_weights:
              weight = model_module.shared_embedding_or_output_weight()
              grad = weight.main_grad
              torch.distributed.all_reduce(
                  grad, group=parallel_state.get_embedding_group()
              )
  
file_path: megatron-lm/megatron/core/distributed/finalize_model_grads.py
imports: from typing import List
  import torch
  from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
  from ..None import parallel_state
  from ..transformer.transformer_config import TransformerConfig
  from ..utils import get_attr_wrapped_model, get_model_config
label_desc: Use image node_id 1 for calling a global function with example usage: _allreduce_word_embedding_grads(model, config) without return types