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

    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.
    
EVA-CLIP-18B/shinji/training/fp16.py