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.
0266
1from __future__ import annotations2 3import math4 5import torch6import torch.distributed as dist7import torch.nn.functional as F8from torch import Tensor9from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors10 11 12def solution(13 X_local: Tensor,14 y_local: Tensor,15 W1: Tensor,16 b1: Tensor,17 W2: Tensor,18 b2: Tensor,19 exp_avg_part: Tensor,20 exp_avg_sq_part: Tensor,21 lr: float,22 beta1: float,23 beta2: float,24 eps: float,25 step: int,26) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:27 world_size = dist.get_world_size()28 rank = dist.get_rank()29 30 templates = [W1, b1, W2, b2]31 flat_p = _flatten_dense_tensors(templates)32 dist.broadcast(flat_p, src=0)33 34 param_views = _unflatten_dense_tensors(flat_p, templates)35 params = [t.detach().requires_grad_(True) for t in param_views]36 37 part = exp_avg_part.numel()38 assert flat_p.numel() == part * world_size39 40 m_part = exp_avg_part.clone()41 v_part = exp_avg_sq_part.clone()42 43 h = F.relu(F.linear(X_local, params[0], params[1]))44 out = F.linear(h, params[2], params[3])45 loss = F.mse_loss(out, y_local)46 loss.backward()47 48 flat_g = _flatten_dense_tensors([p.grad for p in params]).contiguous()49 g_part = torch.empty(part, dtype=flat_g.dtype, device=flat_g.device)50 dist.reduce_scatter_tensor(g_part, flat_g, op=dist.ReduceOp.SUM)51 g_part.div_(world_size)52 53 start = rank * part54 w_part = flat_p[start : start + part].clone()55 56 assert step >= 157 bc1 = 1.0 - math.pow(beta1, step)58 bc2 = 1.0 - math.pow(beta2, step)59 60 m_part.mul_(beta1).add_(g_part, alpha=1.0 - beta1)61 v_part.mul_(beta2).addcmul_(g_part, g_part, value=1.0 - beta2)62 m_hat = m_part / bc163 v_hat = v_part / bc264 w_part.add_(m_hat.div(v_hat.sqrt().add(eps)).mul(-lr))65 66 gathered = torch.empty_like(flat_p)67 dist.all_gather_into_tensor(gathered, w_part.contiguous())68 flat_p.copy_(gathered)69 70 out_params = _unflatten_dense_tensors(flat_p, templates)71 return (*out_params, m_part, v_part)72