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
1"""2Utility functions for creating input tensors and saving output tensors.3 4These functions are shared across different worker scripts to ensure5consistent tensor creation and saving behavior.6"""7 8import os9import json10import copy11import importlib.util12import math13 14import torch15import torch.distributed as dist16 17def save_tensor(output, logs_dir: str, rank: int) -> str:18 """19 Save output tensor(s) to file.20 21 Handles:22 - Single tensor: saves as rank_X.pt23 - Tuple/list of tensors: saves as dict with keys 'output_0', 'output_1', etc.24 - Dict: saves as-is25 """26 os.makedirs(logs_dir, exist_ok=True)27 path = os.path.join(logs_dir, f"rank_{rank}.pt")28 29 # Handle different output types30 if isinstance(output, torch.Tensor):31 # Single tensor32 torch.save(output.detach().cpu(), path)33 elif isinstance(output, (tuple, list)):34 # Multiple tensors - save as dict35 output_dict = {f'output_{i}': t.detach().cpu() if isinstance(t, torch.Tensor) else t 36 for i, t in enumerate(output)}37 torch.save(output_dict, path)38 elif isinstance(output, dict):39 # Dict - convert tensors to CPU40 output_dict = {k: v.detach().cpu() if isinstance(v, torch.Tensor) else v 41 for k, v in output.items()}42 torch.save(output_dict, path)43 else:44 # Fallback: try to save as-is45 torch.save(output, path)46 47 return path48 49# ---------------------------------------------------------------------------50# INPUT TENSOR STANDARD (tuple-only)51# ---------------------------------------------------------------------------52# create_input_tensor() returns a tuple unpacked as solution_fn(*x). Entries are usually tensors53# but may include Python scalars / dicts / dataclasses (e.g. problem 4, problems 100–105).54# - solution(tensor) for single-tensor problems: x is (tensor,)55# - solution(t1, t2) for multi-arg problems: x is (t1, t2, ...)56# Problems 100–105: solution(rank, world_size, cfg, input_ids).57# Output from solution_fn may still be a single tensor or a tuple; save_tensor() handles both.58# ---------------------------------------------------------------------------59 60def _seed(problem_id: int, rank: int, trial: int = 0) -> None:61 """trial varies RNG across eval runs; trial=0 matches the historical single-run seed."""62 torch.manual_seed(42 + problem_id * 1000 + rank + trial * 1_000_003)63 64_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))65_REF_MODULES_CACHE: dict[int, object] = {}66 67def _round_up_multiple(n: int, m: int) -> int:68 return ((n + m - 1) // m) * m69 70def _ddp_mlp_shapes_divisible_by_dp(N: int, world_size: int) -> tuple[int, int, int]:71 """Pick (d_in, hidden, d_out) so W1,b1,W2,b2 total numel is divisible by world_size (ZeRO partitions)."""72 d_in = max(16, min(N, 256))73 d_out = max(8, min(N // 4, 256))74 hidden = max(32, min(N // 2, 512))75 for _ in range(1024):76 numel = hidden * d_in + hidden + d_out * hidden + d_out77 if numel % world_size == 0:78 return d_in, hidden, d_out79 hidden += 180 raise RuntimeError(f"Could not align MLP parameter numel with world_size={world_size}")81 82def _factor_tp_fsdp(world_size: int) -> tuple[int, int]:83 """Choose ``N_TP × N_FSDP == world_size``, preferring both factors ≥ 2."""84 for n_tp in range(2, world_size):85 if world_size % n_tp == 0:86 n_fsdp = world_size // n_tp87 if n_fsdp >= 2:88 return n_tp, n_fsdp89 return 1, world_size90 91def _moe_narrow_num_experts(world_size: int) -> int:92 """Largest ``E < world_size`` with ``world_size % E == 0`` (narrow EP / DP-over-EP)."""93 for E in range(world_size // 2, 1, -1):94 if world_size % E == 0:95 return E96 return 197 98def _linear(in_features: int, out_features: int, dtype: torch.dtype, device) -> torch.nn.Linear:99 return torch.nn.Linear(in_features, out_features).to(device=device, dtype=dtype)100 101def _load_reference_module(problem_id: int):102 if problem_id in _REF_MODULES_CACHE:103 return _REF_MODULES_CACHE[problem_id]104 stem = {105 100: "100_deepseek_v3_671b_tp_attn_ep_moe",106 101: "101_gemma3_27b_tp_attn_tp_mlp",107 102: "102_llama32_3b_tp_attn_tp_mlp",108 103: "103_olmo_3_32b_tp_attn_tp_mlp",109 104: "104_qwen3_235b_tp_attn_ep_moe",110 105: "105_qwen3_code_flash_30b_tp_attn_ep_moe",111 106: "106_deepseek_v3_671b_cp_ulysses_attn_ep_moe",112 107: "107_gemma3_27b_cp_ulysses_attn_tp_mlp",113 108: "108_llama32_3b_cp_ulysses_attn_tp_mlp",114 109: "109_olmo_3_32b_cp_ulysses_attn_tp_mlp",115 110: "110_qwen3_235b_cp_ulysses_attn_ep_moe",116 111: "111_qwen3_code_flash_30b_cp_ulysses_attn_ep_moe",117 }[problem_id]118 path = os.path.join(_PROJECT_ROOT, "reference", f"{stem}.py")119 spec = importlib.util.spec_from_file_location(f"ref_{stem}", path)120 mod = importlib.util.module_from_spec(spec)121 spec.loader.exec_module(mod)122 _REF_MODULES_CACHE[problem_id] = mod123 return mod124 125def _align_model_args_100(cfg, world_size: int) -> None:126 """ModelArgs for reference/100: TP/EP divisibility constraints."""127 cfg.n_layers = 2128 for attr in ("dim", "inter_dim", "moe_inter_dim"):129 v = getattr(cfg, attr)130 if v % world_size:131 setattr(cfg, attr, _round_up_multiple(v, world_size))132 if cfg.vocab_size % world_size:133 cfg.vocab_size = _round_up_multiple(cfg.vocab_size, world_size)134 if cfg.n_heads % world_size:135 cfg.n_heads = _round_up_multiple(cfg.n_heads, world_size)136 if cfg.n_routed_experts % world_size:137 cfg.n_routed_experts = _round_up_multiple(cfg.n_routed_experts, world_size)138 shared = cfg.n_shared_experts * cfg.moe_inter_dim139 guard = 0140 while shared % world_size and guard < 4096:141 cfg.moe_inter_dim += 1142 shared = cfg.n_shared_experts * cfg.moe_inter_dim143 guard += 1144 145def _common_attn_dims(base_shape, world_size):146 """Shared (B, T, num_heads, head_dim) from base_shape (M, N)."""147 M, N = base_shape148 B, T = max(1, M // 64), max(1, N // 64)149 num_heads = 8150 head_dim = 64151 assert num_heads % world_size == 0, f"num_heads ({num_heads}) must be divisible by world_size ({world_size})"152 return B, T, num_heads, head_dim153 154def _build_cp_groups():155 """CP-only (problem 54): the CP group is just WORLD."""156 return dist.group.WORLD157 158def _build_tp_cp_groups(tp_size: int):159 """160 Build TP / CP groups for problem 55, following Megatron order='tp-cp'.161 Rank layout: [cp0_tp0, cp0_tp1, ..., cp1_tp0, cp1_tp1, ...]162 """163 world_size = dist.get_world_size()164 rank = dist.get_rank()165 cp_size = world_size // tp_size166 167 tp_group = None168 cp_group = None169 170 # TP groups: contiguous blocks of tp_size within each CP index.171 for cp_idx in range(cp_size):172 ranks = list(range(cp_idx * tp_size, (cp_idx + 1) * tp_size))173 g = dist.new_group(ranks=ranks)174 if rank in ranks:175 tp_group = g176 177 # CP groups: same TP position across CP partitions.178 for tp_idx in range(tp_size):179 ranks = [tp_idx + cp_idx * tp_size for cp_idx in range(cp_size)]180 g = dist.new_group(ranks=ranks)181 if rank in ranks:182 cp_group = g183 184 assert tp_group is not None and cp_group is not None185 return tp_group, cp_group, cp_size186 187def _build_cp_pp_groups(pp_size: int):188 """189 Build CP / PP groups for problem 56, following Megatron order with TP=DP=1.190 Rank layout: [pp0_cp0, pp0_cp1, ..., pp1_cp0, pp1_cp1, ...]191 """192 world_size = dist.get_world_size()193 rank = dist.get_rank()194 cp_size = world_size // pp_size195 196 cp_group = None197 pp_group = None198 199 # CP groups: contiguous stage-local blocks.200 for pp_idx in range(pp_size):201 ranks = list(range(pp_idx * cp_size, (pp_idx + 1) * cp_size))202 g = dist.new_group(ranks=ranks)203 if rank in ranks:204 cp_group = g205 206 # PP groups: same CP rank across pipeline stages.207 for cp_idx in range(cp_size):208 ranks = [cp_idx + pp_idx * cp_size for pp_idx in range(pp_size)]209 g = dist.new_group(ranks=ranks)210 if rank in ranks:211 pp_group = g212 213 assert cp_group is not None and pp_group is not None214 cp_rank = dist.get_rank(cp_group)215 pp_rank = dist.get_rank(pp_group)216 return cp_group, pp_group, cp_rank, pp_rank, cp_size217 218def _build_cp_dp_groups(dp_size: int):219 """220 Build CP / DP / DP-with-CP groups for problem 57 (backward), following Megatron order 'cp-dp'.221 Rank layout: [dp0_cp0, dp0_cp1, ..., dp1_cp0, dp1_cp1, ...]222 """223 world_size = dist.get_world_size()224 rank = dist.get_rank()225 cp_size = world_size // dp_size226 227 dp_cp_group = dist.new_group(ranks=list(range(world_size)))228 229 dp_group = None230 cp_group = None231 232 # DP groups: same CP position across DP replicas.233 for cp_idx in range(cp_size):234 ranks = [cp_idx + dp_idx * cp_size for dp_idx in range(dp_size)]235 g = dist.new_group(ranks=ranks)236 if rank in ranks:237 dp_group = g238 239 # CP groups: contiguous CP shards inside one DP replica.240 for dp_idx in range(dp_size):241 ranks = list(range(dp_idx * cp_size, (dp_idx + 1) * cp_size))242 g = dist.new_group(ranks=ranks)243 if rank in ranks:244 cp_group = g245 246 assert dp_group is not None and cp_group is not None247 cp_rank = dist.get_rank(cp_group)248 dp_rank = dist.get_rank(dp_group)249 return cp_group, dp_group, dp_cp_group, cp_rank, dp_rank, cp_size250 251def _build_polar_azimuth_groups(azimuth_size: int):252 """253 Build a 2D polar/azimuth process grid.254 Rank layout: [polar0_az0, polar0_az1, ..., polar1_az0, polar1_az1, ...]255 """256 world_size = dist.get_world_size()257 rank = dist.get_rank()258 polar_size = world_size // azimuth_size259 260 azimuth_group = None261 polar_group = None262 263 for polar_idx in range(polar_size):264 ranks = list(range(polar_idx * azimuth_size, (polar_idx + 1) * azimuth_size))265 g = dist.new_group(ranks=ranks)266 if rank in ranks:267 azimuth_group = g268 269 for azimuth_idx in range(azimuth_size):270 ranks = [polar_idx * azimuth_size + azimuth_idx for polar_idx in range(polar_size)]271 g = dist.new_group(ranks=ranks)272 if rank in ranks:273 polar_group = g274 275 assert azimuth_group is not None and polar_group is not None276 azimuth_rank = dist.get_rank(azimuth_group)277 polar_rank = dist.get_rank(polar_group)278 return azimuth_group, polar_group, azimuth_rank, polar_rank, azimuth_size, polar_size279 280def create_input_tensor(281 rank: int,282 world_size: int,283 problem_id: int,284 base_shape: tuple,285 dtype: torch.dtype,286 trial: int = 0,287 device=None,288):289 """290 Create appropriate input tensors for this problem. Always returns a tuple of tensors.291 292 base_shape is typically (M, N) from the worker args (e.g. 1024, 1024).293 Derived dimensions are hardcoded where needed for consistency.294 295 Args:296 rank: Process rank (0..world_size-1)297 world_size: Total number of processes298 problem_id: Problem ID (e.g. 1–105) from reference filename299 base_shape: Base tensor shape tuple (e.g., (M, N))300 dtype: Tensor data type301 trial: Non-negative index; changes RNG for problems that use random inputs (trial=0 is legacy behavior).302 device: PyTorch device or device string. If None, uses torch.device("cuda", rank)303 """304 if device is None:305 dev = torch.device("cuda", rank)306 elif isinstance(device, str):307 dev = device308 else:309 dev = device310 311 M, N = base_shape312 val = float(rank + 1)313 314 # 1-8: collectives315 if problem_id in [1, 2, 3, 6]:316 return (torch.full(base_shape, val, dtype=dtype, device=dev),)317 elif problem_id == 4:318 return (torch.full(base_shape, val, dtype=dtype, device=dev), 0)319 elif problem_id == 5:320 src = 0321 if rank == src:322 chunks = [torch.full(base_shape, float(i + 1), dtype=dtype, device=dev) for i in range(world_size)]323 return (torch.stack(chunks, dim=0),)324 return (torch.zeros(base_shape, dtype=dtype, device=dev),)325 elif problem_id == 7:326 return (torch.full((world_size * M,) + base_shape[1:], val, dtype=dtype, device=dev),)327 elif problem_id == 8:328 chunks = [torch.full(base_shape, float(rank * 10 + d), dtype=dtype, device=dev) for d in range(world_size)]329 return (torch.stack(chunks, dim=0),)330 331 # 9: layernorm_backward332 elif problem_id == 9:333 _seed(problem_id, rank, trial)334 B, H = base_shape335 X_hat = torch.randn((B, H), dtype=dtype, device=dev)336 X_hat = X_hat / (X_hat.norm(dim=-1, keepdim=True) + 1e-5)337 dY = torch.randn((B, H), dtype=dtype, device=dev)338 return (X_hat, dY)339 340 # 10: embedding_lookup341 elif problem_id == 10:342 _seed(problem_id, rank, trial)343 shard_size, embed_dim = base_shape344 local_shard = torch.randn((shard_size, embed_dim), dtype=dtype, device=dev)345 indices = torch.randint(0, world_size * shard_size, (shard_size,), dtype=torch.long, device=dev)346 return (indices, local_shard)347 348 # 11: allgather_gemm_AT349 elif problem_id == 11:350 _seed(problem_id, rank, trial)351 K = 512352 K_local = K // world_size353 A_local = torch.randn((M, K_local), dtype=dtype, device=dev)354 B = torch.randn((K, N), dtype=dtype, device=dev)355 return (A_local, B)356 357 # 12: allgather_gemm358 elif problem_id == 12:359 _seed(problem_id, rank, trial)360 K = 512361 K_local = K // world_size362 A_local = torch.randn((M, K_local), dtype=dtype, device=dev)363 B = torch.randn((K, N), dtype=dtype, device=dev)364 return (A_local, B)365 366 # 13: gemm_allreduce367 elif problem_id == 13:368 _seed(problem_id, rank, trial)369 K = 512370 A_local = torch.randn((M, K), dtype=dtype, device=dev)371 B_local = torch.randn((K, N), dtype=dtype, device=dev)372 return (A_local, B_local)373 374 # 14: gemm_allgather375 elif problem_id == 14:376 _seed(problem_id, rank, trial)377 K = 512378 N_local = N // world_size379 A = torch.randn((M, K), dtype=dtype, device=dev)380 B = torch.randn((K, N_local), dtype=dtype, device=dev)381 return (A, B)382 383 # 15: combined_sharded_gemms384 elif problem_id == 15:385 _seed(problem_id, rank, trial)386 M_rows = _round_up_multiple(M, world_size)387 H = _round_up_multiple(256, world_size)388 H_local = H // world_size389 F = 512390 x_local = torch.randn((M_rows, H_local), dtype=dtype, device=dev)391 W1 = torch.randn((H, F), dtype=dtype, device=dev)392 W2 = torch.randn((F, H), dtype=dtype, device=dev)393 return (x_local, W1, W2)394 395 # 16: gemm_reducescatter396 elif problem_id == 16:397 _seed(problem_id, rank, trial)398 K = 512399 K_local = K // world_size400 A_local = torch.randn((M, K_local), dtype=dtype, device=dev)401 B_local = torch.randn((K_local, N), dtype=dtype, device=dev)402 return (A_local, B_local)403 404 # 17: rope_allgather405 elif problem_id == 17:406 _seed(problem_id, rank, trial)407 B, T, num_heads, head_dim = _common_attn_dims(base_shape, world_size)408 S_local = max(1, T // world_size)409 q_local = torch.randn((B, S_local, num_heads, head_dim), dtype=dtype, device=dev)410 k_local = torch.randn((B, S_local, num_heads, head_dim), dtype=dtype, device=dev)411 cos_local = torch.randn((B, S_local, head_dim), dtype=dtype, device=dev)412 sin_local = torch.randn((B, S_local, head_dim), dtype=dtype, device=dev)413 return (q_local, k_local, cos_local, sin_local)414 415 # 18: rms_norm416 elif problem_id == 18:417 _seed(problem_id, rank, trial)418 hidden = torch.randn(base_shape, dtype=dtype, device=dev)419 weight = torch.randn((N,), dtype=dtype, device=dev)420 return (hidden, weight, 1e-5)421 422 # 19: blocked_fp8_quantize423 elif problem_id == 19:424 _seed(problem_id, rank, trial)425 return (torch.randn(base_shape, dtype=dtype, device=dev), 128)426 427 # 20: blocked_fp8_dequantize428 elif problem_id == 20:429 _seed(problem_id, rank, trial)430 chunk_numel = M * N431 block_size = 128432 num_blocks_per_chunk = chunk_numel // block_size433 local_y = torch.randn((world_size, M, N), dtype=dtype, device=dev)434 local_s = torch.randn((world_size, num_blocks_per_chunk), dtype=dtype, device=dev)435 return (local_y, local_s, block_size)436 437 # 21: clip_grad_norm_no_ep438 elif problem_id == 21:439 _seed(problem_id, rank, trial)440 grad_tensors = [torch.randn(base_shape, dtype=dtype, device=dev) for _ in range(3)]441 return (grad_tensors, 1.0, 2.0, None)442 443 # 22: clip_grad_norm_ep444 elif problem_id == 22:445 _seed(problem_id, rank, trial)446 non_ep = [torch.randn(base_shape, dtype=dtype, device=dev)]447 ep_size = max(1, world_size // 2)448 ep = [torch.randn(base_shape, dtype=dtype, device=dev)]449 return (non_ep, ep, 1.0, 2.0, ep_size, None, None, None)450 451 # 23: grad_acc_loss452 elif problem_id == 23:453 _seed(problem_id, rank, trial)454 loss = torch.randn((), dtype=dtype, device=dev)455 local_valid = torch.tensor(M * N, dtype=torch.long, device=dev)456 global_valid = torch.tensor(world_size * M * N, dtype=torch.long, device=dev)457 grad_normalized_loss = torch.ones((), dtype=dtype, device=dev)458 grad_loss_sum = torch.zeros((), dtype=dtype, device=dev)459 return (loss, local_valid, global_valid, grad_normalized_loss, grad_loss_sum)460 461 # 24: load_balancing_loss_fn462 elif problem_id == 24:463 _seed(problem_id, rank, trial)464 num_experts = 8465 gate_logits = torch.randn((M, num_experts), dtype=dtype, device=dev)466 return (gate_logits, num_experts, 2, None)467 468 # 25: importance_sampling_loss469 elif problem_id == 25:470 _seed(problem_id, rank, trial)471 vocab_size = 32000472 hidden_states = torch.randn((M, N), dtype=dtype, device=dev)473 weight = torch.randn((vocab_size, N), dtype=dtype, device=dev)474 labels = torch.randint(0, vocab_size, (M,), dtype=torch.long, device=dev)475 old_logprobs = torch.randn((M,), dtype=dtype, device=dev)476 advantages = torch.randn((M,), dtype=dtype, device=dev)477 return (hidden_states, weight, labels, old_logprobs, advantages, -100)478 479 # 26: moe_token_preprocess480 elif problem_id == 26:481 _seed(problem_id, rank, trial)482 num_experts = 8483 topk = 2484 selected_experts = torch.randint(0, num_experts, (M, topk), device=dev)485 expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts).float().permute(2, 1, 0)486 return (expert_mask, num_experts, None)487 488 # 27: moe_all2all_primitive489 elif problem_id == 27:490 _seed(problem_id, rank, trial)491 local_tokens = M492 hidden_dim = N493 local_tensor = torch.randn((local_tokens, hidden_dim), dtype=dtype, device=dev)494 chunk = local_tokens // world_size495 input_split_sizes = [chunk] * world_size496 if local_tokens % world_size:497 input_split_sizes[-1] += local_tokens % world_size498 output_split_sizes = list(input_split_sizes)499 return (local_tensor, input_split_sizes, output_split_sizes, None)500 501 # 28: moe_pre_all2all502 elif problem_id == 28:503 _seed(problem_id, rank, trial)504 num_experts = 8505 assert num_experts % world_size == 0, (506 f"problem 28 needs num_experts ({num_experts}) divisible by world_size ({world_size})"507 )508 topk = 2509 hidden_states = torch.randn((M, N), dtype=dtype, device=dev)510 expert_mask = torch.zeros((num_experts, topk, M), dtype=torch.long, device=dev)511 for j in range(M):512 experts = torch.randperm(num_experts, device=dev)[:topk]513 for i, e in enumerate(experts):514 expert_mask[e, i, j] = 1515 expert_mask = expert_mask.float()516 routing_map_bool = expert_mask.sum(dim=1) > 0517 total_permuted = int(routing_map_bool.sum().item())518 chunk = total_permuted // world_size519 input_splits = [chunk] * world_size520 if total_permuted % world_size:521 input_splits[-1] += total_permuted % world_size522 output_splits = list(input_splits)523 num_local_experts = num_experts // world_size524 n_slots = world_size * num_local_experts525 base = total_permuted // n_slots526 rem_tp = total_permuted % n_slots527 flat = torch.full((n_slots,), base, dtype=torch.long, device=dev)528 flat[:rem_tp] += 1529 num_global_tokens_per_local_expert = flat.view(world_size, num_local_experts)530 return (hidden_states, expert_mask, num_experts, input_splits, output_splits, num_global_tokens_per_local_expert, None)531 532 # 29: moe_post_all2all533 elif problem_id == 29:534 _seed(problem_id, rank, trial)535 num_experts = 8536 assert num_experts % world_size == 0, (537 f"problem 29 needs num_experts ({num_experts}) divisible by world_size ({world_size})"538 )539 topk = 2540 num_tokens = M541 routing_map = torch.zeros((num_experts, num_tokens), dtype=torch.bool, device=dev)542 for j in range(num_tokens):543 experts = torch.randperm(num_experts, device=dev)[:topk]544 routing_map[experts, j] = True545 num_routed = int(routing_map.sum().item())546 routing_weights = torch.zeros((num_tokens, topk), dtype=dtype, device=dev)547 selected_experts = torch.zeros((num_tokens, topk), dtype=torch.long, device=dev)548 for j in range(num_tokens):549 idx = torch.where(routing_map[:, j])[0][:topk]550 selected_experts[j] = idx551 w = torch.randn((topk,), dtype=dtype, device=dev).softmax(dim=0)552 routing_weights[j, :] = w553 expert_outputs = torch.randn((num_routed, N), dtype=dtype, device=dev)554 chunk = num_routed // world_size555 input_splits = [chunk] * world_size556 if num_routed % world_size:557 input_splits[-1] += num_routed % world_size558 output_splits = list(input_splits)559 num_local_experts = num_experts // world_size560 n_slots = world_size * num_local_experts561 base = num_routed // n_slots562 rem_nr = num_routed % n_slots563 flat = torch.full((n_slots,), base, dtype=torch.long, device=dev)564 flat[:rem_nr] += 1565 num_global_tokens_per_local_expert = flat.view(world_size, num_local_experts)566 perm = torch.zeros(num_routed, dtype=torch.long, device=dev)567 idx = 0568 for e in range(num_experts):569 for t in range(num_tokens):570 if routing_map[e, t]:571 perm[idx] = t572 idx += 1573 org_hidden_states_shape = torch.Size([num_tokens, N])574 return (expert_outputs, routing_weights, selected_experts, num_experts, input_splits, output_splits, num_global_tokens_per_local_expert, routing_map, perm, org_hidden_states_shape, None)575 576 # 30: moe_epgroupgemm_lora_backward577 elif problem_id == 30:578 _seed(problem_id, rank, trial)579 r, in_f, out_f = 8, N, N580 grad_fc1_1 = torch.randn((r, in_f), dtype=dtype, device=dev)581 grad_fc1_2 = torch.randn((r, in_f), dtype=dtype, device=dev)582 grad_fc2 = torch.randn((out_f, r), dtype=dtype, device=dev)583 return (grad_fc1_1, grad_fc1_2, grad_fc2, None)584 585 # 31: fused_moe_fwd586 elif problem_id == 31:587 _seed(problem_id, rank, trial)588 num_experts = 8589 top_k = 2590 hidden_dim = N591 inter_dim = 128592 hidden_states = torch.randn((M, hidden_dim), dtype=dtype, device=dev)593 gate_weight = torch.randn((num_experts, hidden_dim), dtype=dtype, device=dev)594 gate_bias = torch.randn((num_experts,), dtype=dtype, device=dev)595 gate_proj = _linear(hidden_dim, inter_dim, dtype, dev)596 up_proj = _linear(hidden_dim, inter_dim, dtype, dev)597 down_proj = _linear(inter_dim, hidden_dim, dtype, dev)598 return (hidden_states, gate_weight, gate_bias, gate_proj, up_proj, down_proj, num_experts, top_k, None)599 600 # 32: fused_moe_fwd_lora601 elif problem_id == 32:602 _seed(problem_id, rank, trial)603 num_experts = 8604 top_k = 2605 hidden_dim = N606 inter_dim = 128607 lora_r = 8608 hidden_states = torch.randn((M, hidden_dim), dtype=dtype, device=dev)609 gate_weight = torch.randn((num_experts, hidden_dim), dtype=dtype, device=dev)610 gate_bias = torch.randn((num_experts,), dtype=dtype, device=dev)611 gate_proj = _linear(hidden_dim, inter_dim, dtype, dev)612 up_proj = _linear(hidden_dim, inter_dim, dtype, dev)613 down_proj = _linear(inter_dim, hidden_dim, dtype, dev)614 lora_gate_A = torch.randn((lora_r, hidden_dim), dtype=dtype, device=dev)615 lora_gate_B = torch.randn((inter_dim, lora_r), dtype=dtype, device=dev)616 lora_up_A = torch.randn((lora_r, hidden_dim), dtype=dtype, device=dev)617 lora_up_B = torch.randn((inter_dim, lora_r), dtype=dtype, device=dev)618 lora_down_A = torch.randn((lora_r, inter_dim), dtype=dtype, device=dev)619 lora_down_B = torch.randn((hidden_dim, lora_r), dtype=dtype, device=dev)620 return (621 hidden_states,622 gate_weight,623 gate_bias,624 gate_proj,625 up_proj,626 down_proj,627 lora_gate_A,628 lora_gate_B,629 lora_up_A,630 lora_up_B,631 lora_down_A,632 lora_down_B,633 num_experts,634 top_k,635 None,636 )637 638 # 33: ulysses_all_to_all_tensor_primitive639 elif problem_id == 33:640 _seed(problem_id, rank, trial)641 x = torch.randn(base_shape, dtype=dtype, device=dev)642 return (x, 0, 1, None)643 644 # 34: ulysses_all_gather_into_tensor_primitive645 elif problem_id == 34:646 _seed(problem_id, rank, trial)647 x = torch.randn(base_shape, dtype=dtype, device=dev)648 return (x, None)649 650 # 35: ulysses_all_gather_variable_primitive651 elif problem_id == 35:652 _seed(problem_id, rank, trial)653 x = torch.randn(base_shape, dtype=dtype, device=dev)654 return (x, 0, None)655 656 # 36: ulysses_gather_seq_scatter_heads657 elif problem_id == 36:658 _seed(problem_id, rank, trial)659 B, T, num_heads, head_dim = _common_attn_dims(base_shape, world_size)660 x = torch.randn((B, T, num_heads, head_dim), dtype=dtype, device=dev)661 return (x, 1, 2, None, 0)662 663 # 37: ulysses_gather_heads_scatter_seq664 elif problem_id == 37:665 _seed(problem_id, rank, trial)666 B, T, num_heads, head_dim = _common_attn_dims(base_shape, world_size)667 x = torch.randn((B, T, num_heads, head_dim), dtype=dtype, device=dev)668 return (x, 1, 2, None)669 670 # 38: ulysses_gather_seq_scatter_heads_qkv671 elif problem_id == 38:672 _seed(problem_id, rank, trial)673 B, T, num_heads, head_dim = _common_attn_dims(base_shape, world_size)674 qkv = torch.randn((B, T, 3 * num_heads * head_dim), dtype=dtype, device=dev)675 return (qkv, 1, None, None, True)676 677 # 39: ulysses_attention_e2e678 elif problem_id == 39:679 _seed(problem_id, rank, trial)680 B, T, num_heads, head_dim = _common_attn_dims(base_shape, world_size)681 S_local = max(1, T // world_size)682 H = num_heads * head_dim683 hidden_states = torch.randn((B, S_local, H), dtype=dtype, device=dev)684 w_qkv = torch.randn((3 * num_heads * head_dim, H), dtype=dtype, device=dev)685 w_o = torch.randn((H, num_heads * head_dim), dtype=dtype, device=dev)686 return (hidden_states, w_qkv, w_o, None, num_heads, False)687 688 # 40: ddp689 elif problem_id == 40:690 _seed(problem_id, 0, trial)691 n_total = _round_up_multiple(max(M, world_size), world_size)692 chunk = n_total // world_size693 d_in = max(16, min(N, 256))694 hidden = max(32, min(N // 2, 512))695 d_out = max(8, min(N // 4, 256))696 697 full_X = torch.randn((n_total, d_in), dtype=dtype, device=dev)698 full_y = torch.randn((n_total, d_out), dtype=dtype, device=dev)699 sl = slice(rank * chunk, (rank + 1) * chunk)700 X_local = full_X[sl].contiguous()701 y_local = full_y[sl].contiguous()702 703 def _init_param(shape: tuple) -> torch.Tensor:704 if rank == 0:705 return torch.randn(shape, dtype=dtype, device=dev)706 return torch.zeros(shape, dtype=dtype, device=dev)707 708 W1 = _init_param((hidden, d_in))709 b1 = _init_param((hidden,))710 W2 = _init_param((d_out, hidden))711 b2 = _init_param((d_out,))712 713 z = torch.zeros714 exp_avg_W1 = z((hidden, d_in), dtype=dtype, device=dev)715 exp_avg_b1 = z((hidden,), dtype=dtype, device=dev)716 exp_avg_W2 = z((d_out, hidden), dtype=dtype, device=dev)717 exp_avg_b2 = z((d_out,), dtype=dtype, device=dev)718 exp_avg_sq_W1 = z((hidden, d_in), dtype=dtype, device=dev)719 exp_avg_sq_b1 = z((hidden,), dtype=dtype, device=dev)720 exp_avg_sq_W2 = z((d_out, hidden), dtype=dtype, device=dev)721 exp_avg_sq_b2 = z((d_out,), dtype=dtype, device=dev)722 723 lr = 1e-3724 beta1 = 0.9725 beta2 = 0.999726 eps = 1e-8727 adam_step = 1 + (trial % 7)728 return (729 X_local,730 y_local,731 W1,732 b1,733 W2,734 b2,735 exp_avg_W1,736 exp_avg_b1,737 exp_avg_W2,738 exp_avg_b2,739 exp_avg_sq_W1,740 exp_avg_sq_b1,741 exp_avg_sq_W2,742 exp_avg_sq_b2,743 lr,744 beta1,745 beta2,746 eps,747 adam_step,748 )749 750 # 41: zero1_optimizer_shard751 elif problem_id == 41:752 _seed(problem_id, 0, trial)753 n_total = _round_up_multiple(max(M, world_size), world_size)754 chunk = n_total // world_size755 d_in, hidden, d_out = _ddp_mlp_shapes_divisible_by_dp(N, world_size)756 part_numel = (hidden * d_in + hidden + d_out * hidden + d_out) // world_size757 758 full_X = torch.randn((n_total, d_in), dtype=dtype, device=dev)759 full_y = torch.randn((n_total, d_out), dtype=dtype, device=dev)760 sl = slice(rank * chunk, (rank + 1) * chunk)761 X_local = full_X[sl].contiguous()762 y_local = full_y[sl].contiguous()763 764 def _init_param(shape: tuple) -> torch.Tensor:765 if rank == 0:766 return torch.randn(shape, dtype=dtype, device=dev)767 return torch.zeros(shape, dtype=dtype, device=dev)768 769 W1 = _init_param((hidden, d_in))770 b1 = _init_param((hidden,))771 W2 = _init_param((d_out, hidden))772 b2 = _init_param((d_out,))773 774 z = torch.zeros775 exp_avg_part = z((part_numel,), dtype=dtype, device=dev)776 exp_avg_sq_part = z((part_numel,), dtype=dtype, device=dev)777 778 lr = 1e-3779 beta1 = 0.9780 beta2 = 0.999781 eps = 1e-8782 adam_step = 1 + (trial % 7)783 return (784 X_local,785 y_local,786 W1,787 b1,788 W2,789 b2,790 exp_avg_part,791 exp_avg_sq_part,792 lr,793 beta1,794 beta2,795 eps,796 adam_step,797 )798 799 # 42: zero2_optimizer_shard_grad800 elif problem_id == 42:801 _seed(problem_id, 0, trial)802 n_total = _round_up_multiple(max(M, world_size), world_size)803 chunk = n_total // world_size804 d_in, hidden, d_out = _ddp_mlp_shapes_divisible_by_dp(N, world_size)805 part_numel = (hidden * d_in + hidden + d_out * hidden + d_out) // world_size806 807 full_X = torch.randn((n_total, d_in), dtype=dtype, device=dev)808 full_y = torch.randn((n_total, d_out), dtype=dtype, device=dev)809 sl = slice(rank * chunk, (rank + 1) * chunk)810 X_local = full_X[sl].contiguous()811 y_local = full_y[sl].contiguous()812 813 def _init_param(shape: tuple) -> torch.Tensor:814 if rank == 0:815 return torch.randn(shape, dtype=dtype, device=dev)816 return torch.zeros(shape, dtype=dtype, device=dev)817 818 W1 = _init_param((hidden, d_in))819 b1 = _init_param((hidden,))820 W2 = _init_param((d_out, hidden))821 b2 = _init_param((d_out,))822 823 z = torch.zeros824 exp_avg_part = z((part_numel,), dtype=dtype, device=dev)825 exp_avg_sq_part = z((part_numel,), dtype=dtype, device=dev)826 827 lr = 1e-3828 beta1 = 0.9829 beta2 = 0.999830 eps = 1e-8831 adam_step = 1 + (trial % 7)832 return (833 X_local,834 y_local,835 W1,836 b1,837 W2,838 b2,839 exp_avg_part,840 exp_avg_sq_part,841 lr,842 beta1,843 beta2,844 eps,845 adam_step,846 )847 848 # 43: fused_adam_grad_unshard_allgather849 elif problem_id == 43:850 _seed(problem_id, 0, trial)851 P = max(64, min(M * 64, 4096))852 full_grad = torch.randn(P * world_size, dtype=dtype, device=dev)853 grad_shard = full_grad[rank * P : (rank + 1) * P].contiguous()854 full_master = torch.randn(P * world_size, dtype=dtype, device=dev)855 master_shard = full_master[rank * P : (rank + 1) * P].contiguous()856 exp_avg = torch.zeros(P, dtype=dtype, device=dev)857 exp_avg_sq = torch.zeros(P, dtype=dtype, device=dev)858 lr = 1e-3859 beta1 = 0.9860 beta2 = 0.999861 eps = 1e-8862 adam_step = 1 + (trial % 7)863 return (864 grad_shard,865 master_shard,866 exp_avg,867 exp_avg_sq,868 lr,869 beta1,870 beta2,871 eps,872 adam_step,873 )874 875 # 44: quantized_grad_allreduce876 elif problem_id == 44:877 _seed(problem_id, rank, trial)878 n_el = max(M * N, world_size * 64)879 flat_grad = torch.randn((n_el,), dtype=dtype, device=dev)880 block_size = min(128, max(16, max(N // 4, 16)))881 return (flat_grad, block_size)882 883 # 45: reducescatter_fused_rmsnorm884 elif problem_id == 45:885 _seed(problem_id, 0, trial)886 hidden = max(32, min(N, 128))887 rows = max(2, max(M, world_size) // max(world_size, 4))888 chunk = rows * hidden889 gamma = torch.randn((hidden,), dtype=dtype, device=dev)890 _seed(problem_id, rank, trial)891 rs_input = torch.randn((chunk * world_size,), dtype=dtype, device=dev)892 eps = 1e-5893 return (rs_input, gamma, eps)894 895 # 46: fsdp_adamw_sharded896 elif problem_id == 46:897 _seed(problem_id, 0, trial)898 d_in, hidden, d_out = _ddp_mlp_shapes_divisible_by_dp(N, world_size)899 total_numel = hidden * d_in + hidden + d_out * hidden + d_out900 part = total_numel // world_size901 902 full_param = torch.randn(total_numel, dtype=dtype, device=dev)903 flat_param_shard = full_param[rank * part : (rank + 1) * part].contiguous()904 full_grad = torch.randn(total_numel, dtype=dtype, device=dev)905 flat_grad_shard = full_grad[rank * part : (rank + 1) * part].contiguous()906 exp_avg_shard = torch.zeros(part, dtype=dtype, device=dev)907 exp_avg_sq_shard = torch.zeros(part, dtype=dtype, device=dev)908 909 lr = 1e-3910 beta1 = 0.9911 beta2 = 0.999912 eps = 1e-8913 weight_decay = 0.01914 adam_step = 1 + (trial % 7)915 return (916 flat_param_shard,917 flat_grad_shard,918 exp_avg_shard,919 exp_avg_sq_shard,920 lr,921 beta1,922 beta2,923 eps,924 weight_decay,925 adam_step,926 )927 928 # 47: fsdp_step_e2e929 elif problem_id == 47:930 _seed(problem_id, 0, trial)931 n_total = _round_up_multiple(max(M, world_size), world_size)932 chunk = n_total // world_size933 d_in, hidden, d_out = _ddp_mlp_shapes_divisible_by_dp(N, world_size)934 total_numel = hidden * d_in + hidden + d_out * hidden + d_out935 part = total_numel // world_size936 937 full_X = torch.randn((n_total, d_in), dtype=dtype, device=dev)938 full_y = torch.randn((n_total, d_out), dtype=dtype, device=dev)939 sl = slice(rank * chunk, (rank + 1) * chunk)940 X_local = full_X[sl].contiguous()941 y_local = full_y[sl].contiguous()942 943 def _init_param(shape: tuple) -> torch.Tensor:944 if rank == 0:945 return torch.randn(shape, dtype=dtype, device=dev)946 return torch.zeros(shape, dtype=dtype, device=dev)947 948 W1 = _init_param((hidden, d_in))949 b1 = _init_param((hidden,))950 W2 = _init_param((d_out, hidden))951 b2 = _init_param((d_out,))952 953 full_fp = torch.cat([W1.reshape(-1), b1.reshape(-1), W2.reshape(-1), b2.reshape(-1)])954 flat_param_shard = full_fp[rank * part : (rank + 1) * part].contiguous()955 956 exp_avg_shard = torch.zeros(part, dtype=dtype, device=dev)957 exp_avg_sq_shard = torch.zeros(part, dtype=dtype, device=dev)958 959 param_shapes = ((hidden, d_in), (hidden,), (d_out, hidden), (d_out,))960 lr = 1e-3961 beta1 = 0.9962 beta2 = 0.999963 eps = 1e-8964 weight_decay = 0.01965 adam_step = 1 + (trial % 7)966 return (967 X_local,968 y_local,969 flat_param_shard,970 param_shapes,971 exp_avg_shard,972 exp_avg_sq_shard,973 lr,974 beta1,975 beta2,976 eps,977 weight_decay,978 adam_step,979 )980 981 # 48: fsdp_and_tp982 elif problem_id == 48:983 _seed(problem_id, 0, trial)984 n_tp, n_fsdp = _factor_tp_fsdp(world_size)985 base_d = max(32, min(N, 256))986 D = _round_up_multiple(base_d, math.lcm(n_tp, n_fsdp))987 D_ff = _round_up_multiple(max(64, M), n_tp)988 B_total = _round_up_multiple(max(M * 2, world_size * 2), n_fsdp)989 B_fsdp = B_total // n_fsdp990 991 tp_rank = rank % n_tp992 fsdp_rank = rank // n_tp993 994 full_x = torch.randn(B_total, D, dtype=dtype, device=dev)995 x_local = full_x[fsdp_rank * B_fsdp : (fsdp_rank + 1) * B_fsdp].contiguous()996 997 full_W1 = torch.randn(D, D_ff, dtype=dtype, device=dev)998 full_W2 = torch.randn(D, D_ff, dtype=dtype, device=dev)999 full_W3 = torch.randn(D_ff, D, dtype=dtype, device=dev)1000 1001 dr = D // n_fsdp1002 dc = D_ff // n_tp1003 rr = D_ff // n_tp1004 cr = D // n_fsdp1005 1006 W1_shard = full_W1[1007 fsdp_rank * dr : (fsdp_rank + 1) * dr, tp_rank * dc : (tp_rank + 1) * dc1008 ].contiguous()1009 W2_shard = full_W2[1010 fsdp_rank * dr : (fsdp_rank + 1) * dr, tp_rank * dc : (tp_rank + 1) * dc1011 ].contiguous()1012 W3_shard = full_W3[1013 tp_rank * rr : (tp_rank + 1) * rr, fsdp_rank * cr : (fsdp_rank + 1) * cr1014 ].contiguous()1015 1016 return (x_local, W1_shard, W2_shard, W3_shard, n_tp, n_fsdp)1017 1018 # 49: moe_ep_balanced1019 elif problem_id == 49:1020 _seed(problem_id, rank, trial)1021 num_experts = max(1, world_size)1022 top_k = 21023 hidden_dim = N1024 inter_dim = 1281025 hidden_states = torch.randn((M, hidden_dim), dtype=dtype, device=dev)1026 gate_weight = torch.randn((num_experts, hidden_dim), dtype=dtype, device=dev)1027 gate_bias = torch.randn((num_experts,), dtype=dtype, device=dev)1028 gate_proj = _linear(hidden_dim, inter_dim, dtype, dev)1029 up_proj = _linear(hidden_dim, inter_dim, dtype, dev)1030 down_proj = _linear(inter_dim, hidden_dim, dtype, dev)1031 return (1032 hidden_states,1033 gate_weight,1034 gate_bias,1035 gate_proj,1036 up_proj,1037 down_proj,1038 num_experts,1039 top_k,1040 None,1041 )1042 1043 # 50: moe_ep_wide1044 elif problem_id == 50:1045 _seed(problem_id, rank, trial)1046 num_experts = world_size * 21047 top_k = 21048 hidden_dim = N1049 inter_dim = 1281050 hidden_states = torch.randn((M, hidden_dim), dtype=dtype, device=dev)1051 gate_weight = torch.randn((num_experts, hidden_dim), dtype=dtype, device=dev)1052 gate_bias = torch.randn((num_experts,), dtype=dtype, device=dev)1053 gate_proj = _linear(hidden_dim, inter_dim, dtype, dev)1054 up_proj = _linear(hidden_dim, inter_dim, dtype, dev)1055 down_proj = _linear(inter_dim, hidden_dim, dtype, dev)1056 return (1057 hidden_states,1058 gate_weight,1059 gate_bias,1060 gate_proj,1061 up_proj,1062 down_proj,1063 num_experts,1064 top_k,1065 None,1066 )1067 1068 # 51: moe_ep_narrow1069 elif problem_id == 51:1070 _seed(problem_id, rank, trial)1071 num_experts = _moe_narrow_num_experts(world_size)1072 top_k = min(2, num_experts)1073 hidden_dim = N1074 inter_dim = 1281075 hidden_states = torch.randn((M, hidden_dim), dtype=dtype, device=dev)1076 gate_weight = torch.randn((num_experts, hidden_dim), dtype=dtype, device=dev)1077 gate_bias = torch.randn((num_experts,), dtype=dtype, device=dev)1078 gate_proj = _linear(hidden_dim, inter_dim, dtype, dev)1079 up_proj = _linear(hidden_dim, inter_dim, dtype, dev)1080 down_proj = _linear(inter_dim, hidden_dim, dtype, dev)1081 return (1082 hidden_states,1083 gate_weight,1084 gate_bias,1085 gate_proj,1086 up_proj,1087 down_proj,1088 num_experts,1089 top_k,1090 None,1091 )1092 1093 # 52: fp8_reduce_scatter_grads1094 elif problem_id == 52:1095 _seed(problem_id, rank, trial)1096 P = max(64, min(M * 64, 4096))1097 flat_grads = torch.randn(P * world_size, dtype=dtype, device=dev)1098 amax_history = torch.full((16,), 1e-8, dtype=torch.bfloat16, device=dev)1099 return (flat_grads, amax_history)1100 1101 # 53: fp8_allgather_params1102 elif problem_id == 53:1103 _seed(problem_id, rank, trial)1104 P = max(64, min(M * 64, 4096))1105 flat_param_shard = torch.randn(P, dtype=dtype, device=dev)1106 amax_history = torch.full((16,), 1e-8, dtype=torch.bfloat16, device=dev)1107 return (flat_param_shard, amax_history)1108 1109 # 54: ring_attention1110 elif problem_id == 54:1111 _seed(problem_id, rank, trial)1112 B, T, num_heads, head_dim = _common_attn_dims(base_shape, world_size)1113 S_local = max(1, T // world_size)1114 q = torch.randn((B, S_local, num_heads, head_dim), dtype=dtype, device=dev)1115 k = torch.randn((B, S_local, num_heads, head_dim), dtype=dtype, device=dev)1116 v = torch.randn((B, S_local, num_heads, head_dim), dtype=dtype, device=dev)1117 return (q, k, v, None, True, None)1118 1119 # 55: ring_attention_tp1120 elif problem_id == 55:1121 _seed(problem_id, rank, trial)1122 num_heads = 81123 head_dim = 641124 hidden_size = num_heads * head_dim1125 tp_size = min(2, world_size)1126 assert world_size % tp_size == 01127 assert num_heads % tp_size == 01128 tp_group, cp_group, cp_size = _build_tp_cp_groups(tp_size)1129 tp_rank = dist.get_rank(tp_group)1130 cp_rank = dist.get_rank(cp_group)1131 1132 B = max(1, M // 64)1133 T = max(1, N // 64)1134 S_local = max(1, T // cp_size)1135 heads_local = num_heads // tp_size1136 1137 torch.manual_seed(42 + 56 * 1000 + cp_rank + trial * 1_000_003)1138 hidden_states = torch.randn((B, S_local, hidden_size), dtype=dtype, device=dev)1139 torch.manual_seed(42 + 56 * 1000 + 10000 + tp_rank + trial * 1_000_003)1140 w_qkv = torch.randn((3 * heads_local * head_dim, hidden_size), dtype=dtype, device=dev) * 0.021141 w_o = torch.randn((hidden_size, heads_local * head_dim), dtype=dtype, device=dev) * 0.021142 1143 return (hidden_states, w_qkv, w_o, num_heads, None, True, tp_group, cp_group)1144 1145 # 56: ring_attention_pp1146 elif problem_id == 56:1147 _seed(problem_id, rank, trial)1148 num_heads = 81149 head_dim = 641150 hidden_size = num_heads * head_dim1151 pp_size = min(2, world_size)1152 assert world_size % pp_size == 01153 cp_group, pp_group, cp_rank, pp_rank, cp_size = _build_cp_pp_groups(pp_size)1154 1155 B = max(1, M // 64)1156 T = max(1, N // 64)1157 S_local = max(1, T // cp_size)1158 1159 torch.manual_seed(42 + 57 * 1000 + cp_rank + trial * 1_000_003)1160 hidden_states = torch.randn((B, S_local, hidden_size), dtype=dtype, device=dev)1161 torch.manual_seed(42 + 57 * 1000 + 20000 + pp_rank + trial * 1_000_003)1162 w_qkv = torch.randn((3 * num_heads * head_dim, hidden_size), dtype=dtype, device=dev) * 0.021163 w_o = torch.randn((hidden_size, num_heads * head_dim), dtype=dtype, device=dev) * 0.021164 1165 return (hidden_states, w_qkv, w_o, num_heads, None, True, cp_group, pp_group)1166 1167 # 57: ring_attention_backward_dp1168 elif problem_id == 57:1169 _seed(problem_id, rank, trial)1170 B, T, num_heads, head_dim = _common_attn_dims(base_shape, world_size)1171 dp_size = min(2, world_size)1172 assert world_size % dp_size == 01173 cp_group, dp_group, _, cp_rank, dp_rank, cp_size = _build_cp_dp_groups(dp_size)1174 S_local = max(1, T // cp_size)1175 1176 torch.manual_seed(42 + 58 * 1000 + dp_rank * 100 + cp_rank + trial * 1_000_003)1177 q = torch.randn((B, S_local, num_heads, head_dim), dtype=dtype, device=dev)1178 k = torch.randn((B, S_local, num_heads, head_dim), dtype=dtype, device=dev)1179 v = torch.randn((B, S_local, num_heads, head_dim), dtype=dtype, device=dev)1180 dout = torch.randn((B, S_local, num_heads, head_dim), dtype=dtype, device=dev)1181 1182 scale = head_dim ** -0.51183 qh = q.transpose(1, 2).float()1184 kh = k.transpose(1, 2).float()1185 vh = v.transpose(1, 2).float()1186 scores = torch.matmul(qh, kh.transpose(-2, -1)) * scale1187 softmax_lse = torch.logsumexp(scores, dim=-1)1188 out = torch.matmul(torch.softmax(scores, dim=-1), vh).transpose(1, 2).contiguous()1189 out = out.to(dtype)1190 1191 return (dout, q, k, v, out, softmax_lse, None, False, cp_group, dp_group)1192 1193 # 58: openclip_contrastive_loss1194 elif problem_id == 58:1195 _seed(problem_id, rank, trial)1196 B_local = max(1, M // max(world_size, 1))1197 D = max(16, N)1198 image_features = torch.randn((B_local, D), dtype=dtype, device=dev)1199 text_features = torch.randn((B_local, D), dtype=dtype, device=dev)1200 image_features = torch.nn.functional.normalize(image_features, dim=-1).contiguous()