willychan21/ParallelKernelBench_Problems
ParallelKernelBench (benchmark) Reference problems for ParallelKernelBench: a benchmark for LLM-generated multi-GPU CUDA kernels. This dataset contains 87 reference implementations in reference/ and the input tensor specification in utils/input_output_tensors.py. Files Path Description data/problems.parquet One row per problem (tabular access) reference/*.py Reference solution() implementations utils/input_output_tensors.py Input/output tensor… See the full description on the dataset page: https://huggingface.co/datasets/willychan21/ParallelKernelBench_Problems.
0263
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(21 0.0,22 device=next((t.device for t in grad_tensors if t is not None), torch.device("cuda", 0)),23 dtype=torch.float32,24 )25 return acc26 27 28def _fsdp2_reduce_group(29 grad_tensors: List[torch.Tensor],30 norm_type: float,31 reduce_groups: List[tuple],32) -> torch.Tensor:33 p = float(norm_type)34 val = _local_pth_sum(grad_tensors, p)35 for _, group in reduce_groups:36 if group is not None:37 dist.all_reduce(val, op=dist.ReduceOp.SUM, group=group)38 return val39 40 41def solution(42 non_ep_grad_tensors: List[torch.Tensor],43 ep_grad_tensors: List[torch.Tensor],44 max_norm: float,45 norm_type: float = 2.0,46 ep_size: int = 1,47 fsdp_group: Optional[dist.ProcessGroup] = None,48 ep_fsdp_group: Optional[dist.ProcessGroup] = None,49 ep_group: Optional[dist.ProcessGroup] = None,50) -> torch.Tensor:51 if ep_size > 1 and ep_grad_tensors:52 scale = 1.0 / float(ep_size)53 for t in ep_grad_tensors:54 if t is not None:55 t.detach().mul_(scale)56 57 non_ep_total = _fsdp2_reduce_group(58 non_ep_grad_tensors,59 norm_type=norm_type,60 reduce_groups=[("fsdp", fsdp_group)],61 )62 63 ep_total = _fsdp2_reduce_group(64 ep_grad_tensors,65 norm_type=norm_type,66 reduce_groups=[("ep_fsdp", ep_fsdp_group), ("ep", ep_group)],67 )68 69 total_norm = (non_ep_total + ep_total) ** (1.0 / float(norm_type))70 71 if total_norm > max_norm:72 coef = (max_norm / total_norm)73 for t in non_ep_grad_tensors:74 if t is not None:75 t.mul_(coef.to(t.device))76 for t in ep_grad_tensors:77 if t is not None:78 t.mul_(coef.to(t.device))79 80 return total_norm81