willychan21/ParallelKernelBench_Problems

1from typing import List, Optional2 3import torch4import torch.distributed as dist5 6 7def _local_pth_sum(grad_tensors: List[torch.Tensor], p: float) -> torch.Tensor:8    dev = None9    acc = None10    for g in grad_tensors:11        if g is None:12            continue13        g_local = g14        if dev is None:15            dev = g_local.device16            acc = torch.tensor(0.0, device=dev, dtype=torch.float32)17        gn = torch.norm(g_local.detach().to(torch.float32), p=p)18        acc = acc + (gn ** p)19    if acc is None:20        acc = torch.tensor(0.0, device=next((t.device for t in grad_tensors if t is not None), torch.device("cuda", 0)), dtype=torch.float32)21    return acc22 23 24def _fsdp2_reduce_group(25    grad_tensors: List[torch.Tensor],26    norm_type: float,27    reduce_groups: List[tuple],28) -> torch.Tensor:29    p = float(norm_type)30    val = _local_pth_sum(grad_tensors, p)31    for _, group in reduce_groups:32        if group is not None:33            dist.all_reduce(val, op=dist.ReduceOp.SUM, group=group)34    return val35 36 37def solution(38    grad_tensors: List[torch.Tensor],39    max_norm: float,40    norm_type: float = 2.0,41    fsdp_group: Optional[dist.ProcessGroup] = None,42) -> torch.Tensor:43    reduce_groups = [("fsdp", fsdp_group)]44    total_p = _fsdp2_reduce_group(grad_tensors, norm_type, reduce_groups)45    total_norm = total_p ** (1.0 / float(norm_type))46 47    if total_norm > max_norm:48        coef = (max_norm / total_norm)49        for t in grad_tensors:50            if t is not None:51                t.mul_(coef.to(t.device))52 53    return total_norm54