code: def _allreduce_layernorm_grads(
      model: List[torch.nn.Module], config: TransformerConfig
  ):
      """
      All-reduce layernorm grads (for sequence parallelism).
      """
  
      # All-reduce layernorm parameters across model parallel nodes
      # when sequence parallelism is used
      if (
          parallel_state.get_tensor_model_parallel_world_size() > 1
          and config.sequence_parallel
      ):
          grads = []
          for model_chunk in model:
              for param in get_attr_wrapped_model(
                  model_chunk, "parameters"
              )():
                  if getattr(param, "sequence_parallel", False):
                      grad = param.main_grad
                      grads.append(grad.data)
          coalesced = _flatten_dense_tensors(grads)
          torch.distributed.all_reduce(
              coalesced,
              group=parallel_state.get_tensor_model_parallel_group(),
          )
          for buf, synced in zip(
              grads, _unflatten_dense_tensors(coalesced, grads)
          ):
              buf.copy_(synced)
  
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 4 for calling a global function with example usage: _allreduce_layernorm_grads(model, config) without return types