code: def clip_grad_norm(params, max_norm: float = 0, sync: bool = False):
      """
      Clips grad norms.
  
      During combination with FSDP, will also ensure that grad norms are aggregated
      across all workers, since each worker only stores their shard of the
      gradients.
  
      :param params:
          Parameters whose gradients we wish to clip
      :param max_norm:
          Maximum norm we wish the gradients to have. If non-positive, then
          we will not perform clipping.
      :param sync:
          Boolean indicating whether we should aggregate across the distributed
          group. Used only in combination with FSDP.
  
      :returns:
          The gradient norm across all parameters, before clipping.
      """
      if isinstance(params, torch.Tensor):
          params = [params]
      # make sure any generators are expanded
      params = list(params)
      # if syncing we need to manually perform the clipping so that we aggregrate
      # properly
      if max_norm > 0 and not sync:
          return torch.nn.utils.clip_grad_norm_(params, max_norm)
      else:
          normsq = sum(
              p.grad.data.norm() ** 2
              for p in params
              if p.grad is not None
          )
          if sync:
              # also need to get the norms from all the other sharded works in FSDP
              import torch.distributed as dist
  
              dist.all_reduce(normsq)
          grad_norm = normsq.sqrt()
          if max_norm > 0:
              clip_coef = max_norm / (grad_norm + 1e-6)
              for p in params:
                  p.grad.detach().mul_(clip_coef)
  
          return grad_norm
  
imports: import math
  from itertools import chain
  from typing import Optional
  import parlai.utils.logging
label_desc: Use image node_id 1 for calling a global function with example usage: clip_grad_norm(params, max_norm, sync) and returns: torch, grad_norm