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 math4from typing import Sequence5 6import torch7import torch.distributed as dist8import torch.nn.functional as F9from torch import Tensor10from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors11 12 13def solution(14 X_local: Tensor,15 y_local: Tensor,16 flat_param_shard: Tensor,17 param_shapes: Sequence[tuple[int, ...]],18 exp_avg_shard: Tensor,19 exp_avg_sq_shard: Tensor,20 lr: float,21 beta1: float,22 beta2: float,23 eps: float,24 weight_decay: float,25 step: int,26) -> tuple[Tensor, Tensor, Tensor]:27 assert step >= 128 29 world_size = dist.get_world_size()30 p = flat_param_shard.numel()31 32 device = flat_param_shard.device33 dtype = flat_param_shard.dtype34 35 templates = [torch.zeros(shape, dtype=dtype, device=device) for shape in param_shapes]36 full_flat = torch.empty(world_size * p, dtype=dtype, device=device)37 dist.all_gather_into_tensor(full_flat, flat_param_shard.contiguous())38 39 params_f = _unflatten_dense_tensors(full_flat, templates)40 params = [t.detach().requires_grad_(True) for t in params_f]41 42 h = F.relu(F.linear(X_local, params[0], params[1]))43 out = F.linear(h, params[2], params[3])44 loss = F.mse_loss(out, y_local)45 loss.backward()46 47 flat_g = _flatten_dense_tensors([x.grad for x in params])48 g_shard = torch.empty(p, dtype=flat_g.dtype, device=flat_g.device)49 dist.reduce_scatter_tensor(g_shard, flat_g.contiguous(), op=dist.ReduceOp.SUM)50 g_shard.div_(world_size)51 52 m = exp_avg_shard.clone()53 v = exp_avg_sq_shard.clone()54 theta = flat_param_shard.clone()55 56 m.mul_(beta1).add_(g_shard, alpha=1.0 - beta1)57 v.mul_(beta2).addcmul_(g_shard, g_shard, value=1.0 - beta2)58 bc1 = 1.0 - math.pow(beta1, step)59 bc2 = 1.0 - math.pow(beta2, step)60 m_hat = m / bc161 v_hat = v / bc262 denom = v_hat.sqrt().add(eps)63 64 theta.add_(m_hat.div(denom), alpha=-lr)65 theta.add_(flat_param_shard, alpha=-lr * weight_decay)66 67 return theta, m, v