code: def finalize_model_grads(model: List[torch.nn.Module]):
      """
      All-reduce all model grads across DP replicas, layernorm grads for sequence parallelism,
      embedding grads across first and last pipeline stages (if not tied), and expert grads
      for expert parallelism.
      """
  
      config = get_model_config(model[0])
  
      # All-reduce / reduce-scatter across DP replicas.
      if config.timers is not None:
          config.timers("all-grads-sync", log_level=1).start(
              barrier=config.barrier_with_L1_time
          )
      for model_chunk in model:
          model_chunk.finish_grad_sync()
      if config.timers is not None:
          config.timers("all-grads-sync").stop()
  
      # All-reduce layer-norm grads (for sequence parallelism).
      if config.timers is not None:
          config.timers(
              "layernorm-grads-all-reduce", log_level=1
          ).start(barrier=config.barrier_with_L1_time)
      _allreduce_layernorm_grads(model, config)
      if config.timers is not None:
          config.timers("layernorm-grads-all-reduce").stop()
  
      # All-reduce embedding grads (for pipeline parallelism).
      if config.timers is not None:
          config.timers(
              "embedding-grads-all-reduce", log_level=1
          ).start(barrier=config.barrier_with_L1_time)
      _allreduce_embedding_grads(model, config)
      if config.timers is not None:
          config.timers("embedding-grads-all-reduce").stop()
  
      # All-reduce expert grads (for expert parallelism).
      if config.timers is not None:
          config.timers("expert-grads-all-reduce", log_level=1).start(
              barrier=config.barrier_with_L1_time
          )
      _allreduce_expert_grads(model, config)
      if config.timers is not None:
          config.timers("expert-grads-all-reduce").stop()
  
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 6 for calling a global function with example usage: finalize_model_grads(model) without return types