code: def clip_grad_norm_distributed_optimizer(
      optimizer, max_norm, norm_type=2
  ):
      """Clips gradient norm of parameters in distributed optimizer
  
      This is a wrapper around DistributedFusedAdam.clip_grad_norm with
      added functionality to handle model parallel parameters.
  
      Arguments:
          parameters (DistributedFusedAdam): distributed optimizer
          max_norm (float or int): max norm of the gradients
          norm_type (float or int): type of the used p-norm. Currently
              only 2-norm is supported.
  
      Returns:
          Total norm of the parameters (viewed as a single vector).
  
      """
      assert isinstance(optimizer, DistributedFusedAdam)
  
      # Filter parameters based on:
      #   - parameter should not be shared
      #   - should not be a replica due to tensor model parallelism
      params_for_norm = []
      for param in optimizer.parameters():
          is_not_shared = param_is_not_shared(param)
          is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(
              param
          )
          if is_not_shared and is_not_tp_duplicate:
              params_for_norm.append(param)
  
      # Compute grad norm
      # Note: DistributedFusedAdam caches grad norm to avoid redundant
      # communication.
      optimizer.grad_norm(
          parameters=params_for_norm, norm_type=norm_type
      )
  
      return optimizer.clip_grad_norm(max_norm, norm_type=norm_type)
  
imports: import itertools
  import torch
  from torch import inf
  from nemo.collections.nlp.modules.common.megatron.module import param_is_not_shared
  from nemo.utils import logging
label_desc: Use image node_id 3 for calling a global function with example usage: clip_grad_norm_distributed_optimizer(optimizer, max_norm, norm_type) and returns: optimizer