From 3c9d128751720a2daf68997aa0c8bb07a8db4ac6 Mon Sep 17 00:00:00 2001 From: Willy Chan Date: Thu, 25 Jun 2026 10:53:00 -0700 Subject: [PATCH 1/5] better triton example --- kernelgen/generate_kernel.py | 4 +- kernelgen/prompts.toml | 4 +- kernelgen/solutions_new_add_triton.py | 80 +++++++++++++++++++++++++++ 3 files changed, 84 insertions(+), 4 deletions(-) create mode 100644 kernelgen/solutions_new_add_triton.py diff --git a/kernelgen/generate_kernel.py b/kernelgen/generate_kernel.py index b818160..b46db95 100755 --- a/kernelgen/generate_kernel.py +++ b/kernelgen/generate_kernel.py @@ -55,11 +55,11 @@ # Curated allowlist: --model must be a key. Value is the API backend. ALLOWED_MODELS: dict[str, Provider] = { - "gemini-3-flash-preview": "google", # evaluated + "gemini-3-flash-preview": "google", "gemini-3-pro-preview": "google", "zai-org/GLM-5.1": "together", "deepseek-ai/DeepSeek-V4-Pro": "together", - "Qwen/Qwen3-Coder-Next-FP8": "together", # evaluated + "Qwen/Qwen3-Coder-Next-FP8": "together", "claude-sonnet-4-20250514": "anthropic", "claude-opus-4-20250514": "anthropic", "gpt-4.1": "openai", diff --git a/kernelgen/prompts.toml b/kernelgen/prompts.toml index 01170fa..d2ad657 100755 --- a/kernelgen/prompts.toml +++ b/kernelgen/prompts.toml @@ -72,9 +72,9 @@ Hard requirements: """ [[backends.triton.in_context_examples]] -label = "Example A — elementwise add (reference vs custom extension)" +label = "Example A — elementwise add (reference vs Triton + symmetric memory UVA)" pytorch_solution = "kernelgen/solution_ex_add.py" -custom_solution = "kernelgen/solution_new_add_cuda.py" +custom_solution = "kernelgen/solutions_new_add_triton.py" [backends.parallelkittens] backend_display = "ParallelKittens / ThunderKittens CUDA integrations" diff --git a/kernelgen/solutions_new_add_triton.py b/kernelgen/solutions_new_add_triton.py new file mode 100644 index 0000000..af6c071 --- /dev/null +++ b/kernelgen/solutions_new_add_triton.py @@ -0,0 +1,80 @@ +""" +Distributed element-wise vector add across two ranks using a Triton kernel with +symmetric memory (UVA). Each rank adds its local buffer to the peer's buffer +via a device-visible pointer from symm_mem rendezvous. World size must be 2. +""" + +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +import triton +import triton.language as tl + +BLOCK_SIZE = 1024 + + +@triton.jit +def symmetric_add_kernel( + local_ptr, + remote_ptr, + out_ptr, + n, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs < n + + local = tl.load(local_ptr + offs, mask=mask) + remote_base = tl.cast(remote_ptr, tl.pointer_type(tl.float32)) + remote = tl.load(remote_base + offs, mask=mask) + tl.store(out_ptr + offs, local + remote, mask=mask) + + +_symm_cache = None + + +def _get_symm_state(n: int, dtype: torch.dtype, device: torch.device): + global _symm_cache + if _symm_cache is not None: + c = _symm_cache + if c["n"] == n and c["dtype"] == dtype: + return c["buf"], c["hdl"], c["out"] + + buf = symm_mem.empty(n, device=device, dtype=dtype) + hdl = symm_mem.rendezvous(buf, dist.group.WORLD) + out = torch.empty(n, device=device, dtype=dtype) + _symm_cache = {"n": n, "dtype": dtype, "buf": buf, "hdl": hdl, "out": out} + return buf, hdl, out + + +def _launch_symmetric_add(local: torch.Tensor, remote_ptr: int, out: torch.Tensor, n: int) -> None: + grid = (triton.cdiv(n, BLOCK_SIZE),) + symmetric_add_kernel[grid]( + local, + remote_ptr, + out, + n, + BLOCK_SIZE=BLOCK_SIZE, + ) + + +@torch.no_grad() +def solution(tensor: torch.Tensor) -> torch.Tensor: + assert tensor.is_cuda and tensor.is_contiguous() + assert tensor.dtype == torch.float32 + assert dist.is_initialized() + assert dist.get_world_size() == 2 + + rank = dist.get_rank() + peer = 1 - rank + n = tensor.numel() + + buf, hdl, out = _get_symm_state(n, tensor.dtype, tensor.device) + buf.copy_(tensor.reshape(-1)) + hdl.barrier(channel=0) + + remote_ptr = int(hdl.buffer_ptrs[peer]) + _launch_symmetric_add(buf, remote_ptr, out, n) + + return out.reshape_as(tensor) From 3b4a26a50fc88d8e95e65f6c447edd66a9c501af Mon Sep 17 00:00:00 2001 From: Willy Chan Date: Thu, 25 Jun 2026 12:17:01 -0700 Subject: [PATCH 2/5] make 18 more natural tp split --- utils/input_output_tensors.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/utils/input_output_tensors.py b/utils/input_output_tensors.py index 46bd6dd..ccec4a5 100755 --- a/utils/input_output_tensors.py +++ b/utils/input_output_tensors.py @@ -400,8 +400,10 @@ def create_input_tensor( # 18: rms_norm elif problem_id == 18: _seed(problem_id, rank, trial) - hidden = torch.randn(base_shape, dtype=dtype, device=dev) - weight = torch.randn((N,), dtype=dtype, device=dev) + hidden_dim = _round_up_multiple(N, world_size) + local_hidden_dim = hidden_dim // world_size + hidden = torch.randn((M, local_hidden_dim), dtype=dtype, device=dev) + weight = torch.randn((local_hidden_dim,), dtype=dtype, device=dev) return (hidden, weight, 1e-5) # 19: blocked_fp8_quantize From b816e1dc0801f5fc7c2e7330d75a3a255a9724da Mon Sep 17 00:00:00 2001 From: Willy Chan Date: Thu, 25 Jun 2026 12:17:34 -0700 Subject: [PATCH 3/5] more general prompt and changed default model to GLM 5.2 --- kernelgen/generate_kernel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/kernelgen/generate_kernel.py b/kernelgen/generate_kernel.py index b46db95..1c45699 100755 --- a/kernelgen/generate_kernel.py +++ b/kernelgen/generate_kernel.py @@ -57,7 +57,7 @@ ALLOWED_MODELS: dict[str, Provider] = { "gemini-3-flash-preview": "google", "gemini-3-pro-preview": "google", - "zai-org/GLM-5.1": "together", + "zai-org/GLM-5.2": "together", "deepseek-ai/DeepSeek-V4-Pro": "together", "Qwen/Qwen3-Coder-Next-FP8": "together", "claude-sonnet-4-20250514": "anthropic", @@ -67,7 +67,7 @@ "o3": "openai", } -_DEFAULT_MODEL = "gemini-2.5-flash" +_DEFAULT_MODEL = "zai-org/GLM-5.2" # Chat providers apply a server default max output length when ``max_tokens`` is omitted. # Long CUDA/TK sources then stop mid-file (finish_reason=length). @@ -76,7 +76,7 @@ # System instruction for the model; task text is the assembled user prompt from TOML. # Default system instruction for --backend cuda (and any backend without generate_kernel_system_prompt). -SYSTEM_PROMPT = """You are an expert CUDA and distributed systems engineer. +SYSTEM_PROMPT = """You are an expert NVIDIA kernel and distributed systems engineer. Hard requirements for every answer: - Replace NCCL / torch.distributed collectives with custom CUDA using torch.distributed._symmetric_memory (symm_mem), UVA device pointers, and utils.cuda_helpers.compile_cuda_extension for JIT compilation. From 663e410bbd4ff324a0fcae84f2d4169dcdfa631b Mon Sep 17 00:00:00 2001 From: Willy Chan Date: Thu, 25 Jun 2026 12:35:43 -0700 Subject: [PATCH 4/5] expanded credits in readme --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 1482e71..264dd72 100755 --- a/README.md +++ b/README.md @@ -314,4 +314,4 @@ If you use ParallelKernelBench in your work, please cite: } ``` -Special thanks to Stuart Sul, William Hu, Austin Silveria, Hayden Prairie, Jonah Yi for their invaluable assistance and feedback! And of course, thank you to Together AI for sponsoring the compute for this project. +Special thanks to Stuart Sul, William Hu, Austin Silveria, Hayden Prairie, Jonah Yi,Sokserey Sun, Omar Abul-Hassan, Vincent Yip, Kyle Roh for their invaluable assistance and paper feedback! And of course, thank you to Together AI for sponsoring the compute for this project. From 15066159d9447aa98d63e80ae4b52e83c29c5944 Mon Sep 17 00:00:00 2001 From: Willy Chan Date: Sun, 26 Jul 2026 20:52:15 -0700 Subject: [PATCH 5/5] Fix flagged problems so each is distinct and kernel-relevant Every problem keeps its slot; references and input generators are fixed so each one tests a genuinely different multi-GPU kernel skill: - 19/20 (blocked fp8 quant/dequant): references were already Triton kernels, contradicting the PyTorch+NCCL-reference premise; rewritten in pure PyTorch. 20's generator now produces genuinely quantized inputs (positive per-block absmax scales) instead of randn scales. - 26 (moe_token_preprocess): split metadata stays on-device instead of being returned via .tolist()/CPU copies. - 30 (moe_epgroupgemm_lora_backward): was three bare all_reduces (problem 1 x3, misnamed); now computes shared LoRA adapter grads via ragged grouped GEMMs over local experts + EP all-reduce, without mutating inputs. - 34 (ulysses all-gather primitive): was byte-identical to problem 2; now gathers along dim 1 (Ulysses sequence dim), forcing gather+reorder fusion. - 41/42: header comments making the ZeRO-1 vs ZeRO-2 comm-pattern pair explicit. - 44/52/53 (quantized comms): previously quantize->dequantize locally then bf16 collective, so the named technique was untestable; now int8/fp8 payloads and scales actually cross the wire (two-phase quantized all-reduce, fp8 all-to-all reduce-scatter, fp8 all-gather). - 49 (moe_ep_balanced): was byte-identical code to 31/50; now the static capacity-based balanced dispatch (fixed equal splits, zero host-side size exchange, drop/pad to capacity). - 50 (moe_ep_wide): now hosts genuinely distinct per-local-expert weights (stacked tensors) with a ragged grouped-GEMM expert stage. - 63 (deepmd kalman): fp32 instead of f64, no all_gather_object, no in-place mutation of inputs (which drifted across timed iterations), weight exchange via one fixed-size all_gather_into_tensor. - 64 (gnn_neighbor_sampling): reference was a per-node host loop with .item() syncs and np.unique on CPU over a perfectly regular synthetic graph; now fully vectorized on-device (deterministic first-k selection) over an irregular random graph with random per-rank seed frontiers. - 78 (tile-parallel VAE decode): all_gather_object shape exchange replaced with tensor-based exchange, payload duplication (repeat(world_size) + all_to_all) replaced with a padded all-gather, Python blend loops vectorized with broadcast ramps. - 82: trailing-whitespace cleanup. - 83/84 (chunked log-prob top-k): longer sequences and larger vocab so the chunked pipeline (the point of these variants vs 82) processes many chunks. - 24 (load_balancing_loss_fn): generator now feeds multi-layer gate-logit tuples plus a real attention mask so the masked normalization path is exercised and there is meaningful device compute ahead of the scalar all-reduce. Co-Authored-By: Claude Fable 5 --- reference/19_blocked_fp8_quantize.py | 59 ++- reference/20_blocked_fp8_dequantize.py | 36 +- reference/26_moe_token_preprocess.py | 33 +- reference/30_moe_epgroupgemm_lora_backward.py | 29 +- ...lysses_all_gather_into_tensor_primitive.py | 17 +- reference/41_zero1_optimizer_shard.py | 4 + reference/42_zero2_optimizer_shard_grad.py | 3 + reference/44_quantized_grad_allreduce.py | 61 +++- reference/49_moe_ep_balanced.py | 345 +++--------------- reference/50_moe_ep_wide.py | 259 ++++--------- reference/52_fp8_reduce_scatter_grads.py | 36 +- reference/53_fp8_allgather_params.py | 32 +- .../63_deepmd_kalman_filter_optimizer.py | 61 ++-- reference/64_gnn_neighbor_sampling.py | 182 ++++----- .../78_magi1_tile_parallel_vae_decode.py | 82 +++-- reference/82_vocab_parallel_log_prob_topk.py | 124 +++---- utils/input_output_tensors.py | 121 +++--- 17 files changed, 639 insertions(+), 845 deletions(-) diff --git a/reference/19_blocked_fp8_quantize.py b/reference/19_blocked_fp8_quantize.py index d120966..ef4775f 100755 --- a/reference/19_blocked_fp8_quantize.py +++ b/reference/19_blocked_fp8_quantize.py @@ -1,54 +1,47 @@ +# Blocked FP8 (E4M3) quantization + all-gather. +# +# Each rank quantizes its local tensor in contiguous blocks of `block_size` +# elements (per-block absmax scaling to the E4M3 range), then all ranks +# all-gather both the quantized payload and the per-block scales. The fp8 +# payload crosses the wire as raw bytes (uint8 view). +from typing import Tuple + import torch import torch.distributed as dist -import triton -import triton.language as tl -from typing import Tuple -@triton.jit -def block_fp8_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(axis=0) - offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) +_FP8_E4M3_MAX = 448.0 - x = tl.load(x_ptr + offs).to(tl.float32) - - # FP8 E4M3 max value is 448.0 - s = tl.max(tl.abs(x)) / 448.0 - - # Prevent division by zero if all elements in the block are 0 - s_safe = tl.where(s == 0.0, 1.0, s) - - y = (x / s_safe).to(y_ptr.dtype.element_ty) - - tl.store(y_ptr + offs, y) - tl.store(s_ptr + pid, s) +@torch.no_grad() def solution(local_tensor: torch.Tensor, block_size: int = 128) -> Tuple[torch.Tensor, torch.Tensor]: assert local_tensor.size(-1) % block_size == 0, "Last dimension must be divisible by block_size" - - y_local = torch.empty_like(local_tensor, dtype=torch.float8_e4m3fn) - s_local = local_tensor.new_empty( - *local_tensor.size()[:-1], local_tensor.size(-1) // block_size, dtype=torch.float32 - ) - - grid = (triton.cdiv(local_tensor.numel(), block_size),) - block_fp8_quant_kernel[grid](local_tensor, y_local, s_local, BLOCK_SIZE=block_size) - + + orig_shape = local_tensor.shape + x = local_tensor.float().contiguous().reshape(-1, block_size) + + s = x.abs().amax(dim=1) / _FP8_E4M3_MAX + # Prevent division by zero if all elements in the block are 0 + s_safe = torch.where(s == 0.0, torch.ones_like(s), s) + + y_local = (x / s_safe.unsqueeze(1)).to(torch.float8_e4m3fn).reshape(orig_shape) + s_local = s.reshape(*orig_shape[:-1], orig_shape[-1] // block_size) + if dist.is_initialized(): world_size = dist.get_world_size() - + y_local_u8 = y_local.view(torch.uint8) y_gather_u8 = [torch.empty_like(y_local_u8) for _ in range(world_size)] dist.all_gather(y_gather_u8, y_local_u8) - + y_global_u8 = torch.cat(y_gather_u8, dim=0) y_global = y_global_u8.view(torch.float8_e4m3fn) - + s_gather = [torch.empty_like(s_local) for _ in range(world_size)] dist.all_gather(s_gather, s_local) - + s_global = torch.cat(s_gather, dim=0) else: y_global = y_local s_global = s_local - + return y_global, s_global diff --git a/reference/20_blocked_fp8_dequantize.py b/reference/20_blocked_fp8_dequantize.py index 2218843..f8b06a8 100755 --- a/reference/20_blocked_fp8_dequantize.py +++ b/reference/20_blocked_fp8_dequantize.py @@ -1,21 +1,15 @@ +# Blocked FP8 dequantize + all-to-all. +# +# Each rank holds `world_size` quantized chunks (one destined for each peer) +# plus their per-block scales. The reference dequantizes locally (per-block +# scale multiply) and exchanges the reconstructed chunks with a single +# all_to_all. A fused kernel can instead ship the compact fp8 payload and +# dequantize on the receiving side. import torch import torch.distributed as dist -import triton -import triton.language as tl - -@triton.jit -def block_fp8_dequant_kernel(y_ptr, s_ptr, x_ptr, num_elements, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(axis=0) - offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offs < num_elements - - s = tl.load(s_ptr + pid) - - y = tl.load(y_ptr + offs, mask=mask).to(tl.float32) - - tl.store(x_ptr + offs, y * s, mask=mask) +@torch.no_grad() def solution( local_y: torch.Tensor, local_s: torch.Tensor, @@ -25,22 +19,14 @@ def solution( chunk_shape = local_y.shape[1:] chunk_numel = local_y.numel() // world_size - num_elements = local_y.numel() assert chunk_numel % block_size == 0, ( f"Chunk size {chunk_numel} must be divisible by block_size ({block_size})" ) - y_flat = local_y.view(-1) - s_flat = local_s.view(-1) - x_flat = torch.empty(num_elements, device=local_y.device, dtype=torch.float32) - - if num_elements > 0: - grid = (triton.cdiv(num_elements, block_size),) - block_fp8_dequant_kernel[grid]( - y_flat, s_flat, x_flat, num_elements, BLOCK_SIZE=block_size - ) + y_blocks = local_y.float().contiguous().reshape(-1, block_size) + s_flat = local_s.float().contiguous().reshape(-1, 1) + x = (y_blocks * s_flat).reshape(world_size, *chunk_shape) - x = x_flat.view(world_size, *chunk_shape) out = torch.empty_like(x) dist.all_to_all_single(out, x) diff --git a/reference/26_moe_token_preprocess.py b/reference/26_moe_token_preprocess.py index fd1bc5e..3558355 100755 --- a/reference/26_moe_token_preprocess.py +++ b/reference/26_moe_token_preprocess.py @@ -1,4 +1,12 @@ -from typing import List, Optional, Tuple +# MoE all-to-all split-metadata computation, kept entirely on-device. +# +# From the per-token expert mask, compute the send split sizes (tokens this +# rank routes to each peer), all-gather the per-expert token histogram, and +# derive the receive split sizes plus the per-(source rank, local expert) +# token counts. Everything stays on the GPU: reductions over the expert mask +# plus a single small all-gather — a fusion target for kernel-side metadata +# computation. +from typing import Optional, Tuple import torch import torch.distributed as dist @@ -8,13 +16,13 @@ def _preprocess_impl( expert_mask: torch.Tensor, num_experts: int, ep_group: dist.ProcessGroup, -) -> Tuple[List[int], List[int], torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: ep_size = ep_group.size() num_local_experts = num_experts // ep_size rank = dist.get_rank(ep_group) - num_local_tokens_per_expert = expert_mask.sum(dim=(1, 2)) + num_local_tokens_per_expert = expert_mask.sum(dim=(1, 2)).to(torch.long) - input_splits = num_local_tokens_per_expert.reshape(ep_size, num_local_experts).sum(dim=1).tolist() + input_splits = num_local_tokens_per_expert.reshape(ep_size, num_local_experts).sum(dim=1) num_global_tokens_per_expert = torch.empty( ep_size, @@ -27,14 +35,10 @@ def _preprocess_impl( start_idx, end_idx = rank * num_local_experts, (rank + 1) * num_local_experts num_global_tokens_per_local_expert = num_global_tokens_per_expert[:, start_idx:end_idx].contiguous() - output_splits = num_global_tokens_per_local_expert.sum(dim=1).tolist() + output_splits = num_global_tokens_per_local_expert.sum(dim=1) - num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0).to( - torch.device("cpu"), non_blocking=True - ) - num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(-1, num_local_experts).to( - torch.device("cpu"), non_blocking=True - ) + num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0) + num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(-1, num_local_experts) return input_splits, output_splits, num_global_tokens_per_local_expert, num_global_sum_tokens_per_local_expert @@ -45,9 +49,4 @@ def solution( group: Optional[dist.ProcessGroup] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: group = group or dist.group.WORLD - input_splits, output_splits, per_local_expert, sum_per_local_expert = _preprocess_impl( - expert_mask=expert_mask, num_experts=num_experts, ep_group=group - ) - input_splits_t = torch.tensor(input_splits, dtype=torch.long) - output_splits_t = torch.tensor(output_splits, dtype=torch.long) - return (input_splits_t, output_splits_t, per_local_expert, sum_per_local_expert) + return _preprocess_impl(expert_mask=expert_mask, num_experts=num_experts, ep_group=group) diff --git a/reference/30_moe_epgroupgemm_lora_backward.py b/reference/30_moe_epgroupgemm_lora_backward.py index b9fc4ca..0de17a0 100755 --- a/reference/30_moe_epgroupgemm_lora_backward.py +++ b/reference/30_moe_epgroupgemm_lora_backward.py @@ -1,3 +1,12 @@ +# EP grouped-GEMM LoRA backward: per-expert LoRA adapter gradients + all-reduce. +# +# In expert-parallel MoE fine-tuning with shared LoRA adapters, each rank holds +# the activations / output-gradients of its local experts. The gradient of a +# shared LoRA-A (resp. LoRA-B) matrix is the sum over local experts of +# dY_e^T @ X_e — a grouped GEMM reduced over the expert dimension — followed by +# an all-reduce over the EP group (every rank contributes its local experts' +# partial sums). The GEMMs and the three all-reduces are independent per +# adapter, so a fused kernel can overlap compute with communication. from typing import Optional, Tuple import torch @@ -5,12 +14,26 @@ def solution( - grad_fc1_1_lora_A: torch.Tensor, - grad_fc1_2_lora_A: torch.Tensor, - grad_fc2_lora_B: torch.Tensor, + expert_inputs: torch.Tensor, # (E_local, T, d_in) inputs to fc1 LoRA-A + grad_fc1_1_out: torch.Tensor, # (E_local, T, r) grad wrt fc1 gate LoRA-A output + grad_fc1_2_out: torch.Tensor, # (E_local, T, r) grad wrt fc1 up LoRA-A output + fc2_lora_acts: torch.Tensor, # (E_local, T, r) fc2 LoRA-A activations + grad_expert_outputs: torch.Tensor, # (E_local, T, d_out) grad wrt fc2 LoRA-B output group: Optional[dist.ProcessGroup] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: group = group or dist.group.WORLD + + # Grouped GEMMs: sum over the local-expert dimension. + grad_fc1_1_lora_A = torch.bmm( + grad_fc1_1_out.transpose(1, 2), expert_inputs + ).sum(dim=0) # (r, d_in) + grad_fc1_2_lora_A = torch.bmm( + grad_fc1_2_out.transpose(1, 2), expert_inputs + ).sum(dim=0) # (r, d_in) + grad_fc2_lora_B = torch.bmm( + grad_expert_outputs.transpose(1, 2), fc2_lora_acts + ).sum(dim=0) # (d_out, r) + dist.all_reduce(grad_fc1_1_lora_A, op=dist.ReduceOp.SUM, group=group) dist.all_reduce(grad_fc1_2_lora_A, op=dist.ReduceOp.SUM, group=group) dist.all_reduce(grad_fc2_lora_B, op=dist.ReduceOp.SUM, group=group) diff --git a/reference/34_ulysses_all_gather_into_tensor_primitive.py b/reference/34_ulysses_all_gather_into_tensor_primitive.py index 3c2f6f2..25e4fba 100755 --- a/reference/34_ulysses_all_gather_into_tensor_primitive.py +++ b/reference/34_ulysses_all_gather_into_tensor_primitive.py @@ -1,3 +1,9 @@ +# Ulysses all-gather along an arbitrary dim. +# +# all_gather_into_tensor only concatenates along dim 0; the Ulysses sequence +# dim is usually not dim 0, so a rank-stacked gather must be followed by a +# reorder/copy to concatenate along `gather_dim`. A fused kernel can write +# gathered chunks directly into their strided destination. from typing import Optional import torch @@ -6,6 +12,7 @@ def solution( x: torch.Tensor, + gather_dim: int, group: Optional[dist.ProcessGroup] = None, ) -> torch.Tensor: group = group or dist.group.WORLD @@ -14,8 +21,8 @@ def solution( return x.contiguous() x = x.contiguous() - dim_size = list(x.size()) - dim_size[0] = dim_size[0] * world_size - output = torch.empty(dim_size, dtype=x.dtype, device=x.device) - dist.all_gather_into_tensor(output, x, group=group) - return output + stacked = torch.empty( + (world_size,) + tuple(x.size()), dtype=x.dtype, device=x.device + ) + dist.all_gather_into_tensor(stacked, x, group=group) + return torch.cat(list(stacked.unbind(0)), dim=gather_dim).contiguous() diff --git a/reference/41_zero1_optimizer_shard.py b/reference/41_zero1_optimizer_shard.py index 4e2ed0b..97ad565 100755 --- a/reference/41_zero1_optimizer_shard.py +++ b/reference/41_zero1_optimizer_shard.py @@ -1,3 +1,7 @@ +# ZeRO-1: optimizer-state sharding only. Gradients are ALL-REDUCED in full on +# every rank (each rank then slices out its partition); compare problem 42 +# (ZeRO-2), which replaces the all-reduce with a reduce-scatter so each rank +# only ever materializes its own gradient shard. from __future__ import annotations import math diff --git a/reference/42_zero2_optimizer_shard_grad.py b/reference/42_zero2_optimizer_shard_grad.py index e0e821f..e133c22 100755 --- a/reference/42_zero2_optimizer_shard_grad.py +++ b/reference/42_zero2_optimizer_shard_grad.py @@ -1,3 +1,6 @@ +# ZeRO-2: optimizer-state + gradient sharding. Gradients are REDUCE-SCATTERED +# so each rank only materializes its own shard; compare problem 41 (ZeRO-1), +# which all-reduces the full flat gradient on every rank before slicing. from __future__ import annotations import math diff --git a/reference/44_quantized_grad_allreduce.py b/reference/44_quantized_grad_allreduce.py index 3d3755e..dce25dc 100755 --- a/reference/44_quantized_grad_allreduce.py +++ b/reference/44_quantized_grad_allreduce.py @@ -1,5 +1,14 @@ +# Block-wise int8-quantized gradient all-reduce (quantized payload on the wire). +# +# Two-phase quantized all-reduce: each rank block-quantizes its gradient to +# int8, exchanges int8 chunks + fp32 scales with an all-to-all +# (reduce-scatter phase), dequantizes and averages its shard, re-quantizes the +# result, and all-gathers int8 shards + scales (all-gather phase). Only int8 +# payloads and per-block scales cross the wire. from __future__ import annotations +from typing import Tuple + import torch import torch.distributed as dist import torch.nn.functional as F @@ -7,20 +16,17 @@ @torch.no_grad() -def _block_int8_quant_dequant(x_flat: Tensor, block_size: int) -> Tensor: - n = x_flat.numel() - if n == 0: - return x_flat.clone() - flat = x_flat.contiguous().reshape(-1) - pad = (-n) % block_size - if pad: - flat = F.pad(flat, (0, pad)) - nb = flat.numel() // block_size - xv = flat.view(nb, block_size) +def _block_int8_quant(x_flat: Tensor, block_size: int) -> Tuple[Tensor, Tensor]: + """Quantize a flat fp tensor (numel divisible by block_size) to int8 + fp32 scales.""" + xv = x_flat.view(-1, block_size) scales = xv.abs().amax(dim=1).float().clamp(min=1e-8) / 127.0 q = (xv.float() / scales.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8) - out = (q.float() * scales.unsqueeze(1)).reshape(-1) - return out[:n] + return q.reshape(-1), scales + + +@torch.no_grad() +def _block_int8_dequant(q: Tensor, scales: Tensor, block_size: int) -> Tensor: + return q.view(-1, block_size).float() * scales.float().unsqueeze(1) @torch.no_grad() @@ -33,10 +39,31 @@ def solution( world_size = dist.get_world_size() orig_shape = flat_grad.shape x = flat_grad.reshape(-1) + n = x.numel() + + # Pad so every rank shard is a whole number of blocks. + shard = -(-n // (world_size * block_size)) * block_size + padded = world_size * shard + if padded != n: + x = F.pad(x, (0, padded - n)) + + # Phase 1: quantized reduce-scatter via all_to_all of int8 chunks + scales. + q, scales = _block_int8_quant(x, block_size) + q_recv = torch.empty_like(q) + dist.all_to_all_single(q_recv, q) + s_recv = torch.empty_like(scales) + dist.all_to_all_single(s_recv, scales) + + # Dequantize the world_size received chunks and average this rank's shard. + chunks = _block_int8_dequant(q_recv, s_recv, block_size).reshape(world_size, shard) + shard_sum = chunks.sum(dim=0).div_(world_size) - rec = _block_int8_quant_dequant(x, block_size) - acc = rec.float() - dist.all_reduce(acc, op=dist.ReduceOp.SUM) - acc.div_(world_size) + # Phase 2: quantized all-gather of the reduced shards. + q2, scales2 = _block_int8_quant(shard_sum, block_size) + q_full = torch.empty(padded, dtype=torch.int8, device=x.device) + dist.all_gather_into_tensor(q_full, q2) + s_full = torch.empty(padded // block_size, dtype=torch.float32, device=x.device) + dist.all_gather_into_tensor(s_full, scales2) - return acc.to(dtype=flat_grad.dtype).reshape(orig_shape) + out = _block_int8_dequant(q_full, s_full, block_size).reshape(-1)[:n] + return out.to(dtype=flat_grad.dtype).reshape(orig_shape) diff --git a/reference/49_moe_ep_balanced.py b/reference/49_moe_ep_balanced.py index 5808ff7..d7dede5 100755 --- a/reference/49_moe_ep_balanced.py +++ b/reference/49_moe_ep_balanced.py @@ -1,246 +1,20 @@ -# Expert-parallel (EP) fused MoE forward — BALANCED EP. +# Expert-parallel (EP) fused MoE forward — BALANCED EP with STATIC dispatch. # -# Same kernel as problem 31 (router -> permute -> all_to_all dispatch -> per-expert -# SiLU MLP -> all_to_all combine -> unpermute), but the harness sets -# num_experts == world_size, i.e. exactly one expert per rank. With uniform routing -# this gives the balanced all_to_all dispatch pattern (each rank sends/receives a -# roughly equal token count). -from typing import List, Optional, Tuple, Union +# num_experts == world_size (one expert per rank). Unlike the dynamic-split +# MoE problems (31 / 50 / 51), balanced EP uses a fixed per-expert capacity: +# every rank sends exactly `capacity` token slots to every peer, so the +# all_to_all uses equal static splits and needs NO host-side size exchange +# (no .tolist() / no metadata round-trip). Tokens beyond capacity are dropped, +# unused slots are zero-padded — the standard capacity-factor trade-off. The +# whole pipeline (route -> scatter into slots -> a2a -> expert MLP -> a2a -> +# weighted combine) is device-side and communication is perfectly regular, +# which makes it an ideal target for fused compute/comm kernels. +from typing import Optional import torch import torch.distributed as dist -class _AllToAll(torch.autograd.Function): - @staticmethod - def forward(ctx, group, input, output_split_sizes, input_split_sizes): - ctx.group = group - ctx.output_split_sizes = output_split_sizes - ctx.input_split_sizes = input_split_sizes - if dist.get_world_size(group=group) == 1: - return input.contiguous() - input = input.contiguous() - if output_split_sizes is None: - output = torch.empty_like(input) - else: - output = torch.empty( - size=(sum(output_split_sizes), input.size(1)), - dtype=input.dtype, - device=input.device, - ) - dist.all_to_all_single( - output, - input, - output_split_sizes=output_split_sizes, - input_split_sizes=input_split_sizes, - group=group, - ) - return output - - @staticmethod - def backward(ctx, grad_output): - return ( - None, - _AllToAll.apply( - ctx.group, grad_output, ctx.input_split_sizes, ctx.output_split_sizes - ), - None, - None, - ) - - -def _all_to_all( - group: dist.ProcessGroup, - input: torch.Tensor, - output_split_sizes: Optional[List[int]], - input_split_sizes: Optional[List[int]], -) -> torch.Tensor: - return _AllToAll.apply(group, input, output_split_sizes, input_split_sizes) - - -def _preprocess( - expert_mask: torch.Tensor, - num_experts: int, - ep_group: dist.ProcessGroup, -) -> Tuple[List[int], List[int], torch.Tensor, torch.Tensor]: - ep_size = ep_group.size() - num_local_experts = num_experts // ep_size - rank = dist.get_rank(ep_group) - num_local_tokens_per_expert = expert_mask.sum(dim=(1, 2)) - input_splits = ( - num_local_tokens_per_expert.reshape(ep_size, num_local_experts).sum(dim=1).tolist() - ) - num_local_tokens_per_expert_flat = num_local_tokens_per_expert.contiguous().view(-1) - output_size = ep_size * num_local_tokens_per_expert_flat.numel() - num_global_tokens_per_expert_flat = torch.empty( - output_size, - dtype=num_local_tokens_per_expert.dtype, - device=num_local_tokens_per_expert.device, - ) - dist.all_gather_into_tensor( - num_global_tokens_per_expert_flat, num_local_tokens_per_expert_flat, group=ep_group - ) - num_global_tokens_per_expert = num_global_tokens_per_expert_flat.view( - ep_size, num_local_tokens_per_expert.size(0) - ) - start_idx, end_idx = rank * num_local_experts, (rank + 1) * num_local_experts - num_global_tokens_per_local_expert = num_global_tokens_per_expert[ - :, start_idx:end_idx - ].contiguous() - output_splits = num_global_tokens_per_local_expert.sum(dim=1).tolist() - num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum( - dim=0 - ).to(torch.device("cpu"), non_blocking=True) - num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view( - -1, num_local_experts - ).to(torch.device("cpu"), non_blocking=True) - return ( - input_splits, - output_splits, - num_global_tokens_per_local_expert, - num_global_sum_tokens_per_local_expert, - ) - - -def _permute( - tokens: torch.Tensor, routing_map: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: - num_tokens, _ = tokens.shape - num_experts = routing_map.shape[0] - routing_map = routing_map.bool() - token_indices = ( - torch.arange(num_tokens, device=routing_map.device) - .unsqueeze(0) - .expand(num_experts, -1) - ) - sorted_indices = token_indices.masked_select(routing_map) - permuted_input = tokens.index_select(0, sorted_indices) - return permuted_input, sorted_indices - - -def _sort_chunks_by_idxs( - input: torch.Tensor, - split_sizes: Union[torch.Tensor, List[int]], - sorted_idxs: List[int], -) -> torch.Tensor: - if isinstance(split_sizes, torch.Tensor): - split_sizes = split_sizes.tolist() - chunks = torch.split(input, split_sizes, dim=0) - return torch.cat([chunks[i] for i in sorted_idxs], dim=0) - - -def _generate_weights_idx( - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - num_experts: int, -) -> torch.Tensor: - num_tokens, topk = routing_weights.shape - weights_idx = torch.zeros( - (num_tokens, num_experts), - dtype=routing_weights.dtype, - device=routing_weights.device, - ) - weights_idx.scatter_add_(1, selected_experts, routing_weights) - return weights_idx - - -def _unpermute( - tokens: torch.Tensor, - routing_weights: torch.Tensor, - hidden_states_shape: torch.Size, - permutation_mapping: torch.Tensor, - routing_map: torch.Tensor, -) -> torch.Tensor: - tokens_weight = routing_weights.T.contiguous().masked_select(routing_map.bool()) - tokens = tokens * tokens_weight.unsqueeze(-1) - hidden_dim = hidden_states_shape[-1] - unpermuted_tokens = torch.zeros( - hidden_states_shape, device=tokens.device, dtype=tokens.dtype - ) - expanded_mapping = permutation_mapping.unsqueeze(1).expand(-1, hidden_dim) - unpermuted_tokens.scatter_add_(0, expanded_mapping, tokens) - return unpermuted_tokens - - -def token_pre_all2all( - hidden_states: torch.Tensor, - expert_mask: torch.Tensor, - num_experts: int, - input_splits: List[int], - output_splits: List[int], - num_global_tokens_per_local_expert: torch.Tensor, - group: Optional[dist.ProcessGroup] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Size]: - group = group or dist.group.WORLD - hidden_dim = hidden_states.size(-1) - hidden_states = hidden_states.reshape(-1, hidden_dim) - org_hidden_states_shape = hidden_states.shape - routing_map = expert_mask.sum(dim=1) - - local_permuted_hidden_states, local_input_permutation_mapping = _permute( - hidden_states, routing_map - ) - expected_tokens = sum(input_splits) - actual_tokens = local_permuted_hidden_states.shape[0] - if expected_tokens != actual_tokens: - raise RuntimeError( - f"EP split mismatch: input_splits sum ({expected_tokens}) != " - f"permuted tokens ({actual_tokens})" - ) - - global_permuted_hidden_states = _all_to_all( - group, local_permuted_hidden_states, output_splits, input_splits - ) - num_local_experts = num_experts // dist.get_world_size(group) - permute_order = ( - torch.arange(num_experts).reshape(-1, num_local_experts).T.ravel().tolist() - ) - split_sizes = num_global_tokens_per_local_expert.ravel().tolist() - global_permuted_hidden_states = _sort_chunks_by_idxs( - global_permuted_hidden_states, split_sizes, permute_order - ) - return ( - global_permuted_hidden_states, - routing_map, - local_input_permutation_mapping, - org_hidden_states_shape, - ) - - -def tokens_post_all2all( - expert_outputs: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - num_experts: int, - input_splits: List[int], - output_splits: List[int], - num_global_tokens_per_local_expert: torch.Tensor, - routing_map: torch.Tensor, - local_input_permutation_mapping: torch.Tensor, - org_hidden_states_shape: torch.Size, - group: Optional[dist.ProcessGroup] = None, -) -> torch.Tensor: - group = group or dist.group.WORLD - num_local_experts = num_experts // dist.get_world_size(group) - unpermute_order = ( - torch.arange(num_experts).reshape(num_local_experts, -1).T.ravel().tolist() - ) - split_sizes = num_global_tokens_per_local_expert.T.ravel().tolist() - expert_outputs = _sort_chunks_by_idxs( - expert_outputs, split_sizes, unpermute_order - ) - unpermute_outputs = _all_to_all(group, expert_outputs, input_splits, output_splits) - weights_idx = _generate_weights_idx(routing_weights, selected_experts, num_experts) - unpermute_outputs = _unpermute( - unpermute_outputs, - weights_idx, - org_hidden_states_shape, - local_input_permutation_mapping, - routing_map, - ) - return unpermute_outputs - - def expert_forward( x: torch.Tensor, gate_proj: torch.nn.Linear, @@ -252,6 +26,7 @@ def expert_forward( return down_proj(gate * up) +@torch.no_grad() def solution( hidden_states: torch.Tensor, gate_weight: torch.Tensor, @@ -264,53 +39,51 @@ def solution( group: Optional[dist.ProcessGroup] = None, ) -> torch.Tensor: group = group or dist.group.WORLD - hidden_dim = hidden_states.size(-1) - num_tokens = hidden_states.reshape(-1, hidden_dim).size(0) + world_size = dist.get_world_size(group) + assert num_experts == world_size, "balanced EP hosts exactly one expert per rank" - router_logits = torch.nn.functional.linear( - hidden_states.reshape(-1, hidden_dim), gate_weight, gate_bias - ) - routing_weights, selected_experts = torch.topk( - torch.softmax(router_logits, dim=-1), top_k, dim=-1 - ) - expert_mask = torch.nn.functional.one_hot( - selected_experts, num_classes=num_experts - ).permute(2, 1, 0) - - input_splits, output_splits, num_global_tokens_per_local_expert, _ = _preprocess( - expert_mask, num_experts, group - ) - - ( - global_permuted_hidden_states, - routing_map, - local_input_permutation_mapping, - org_hidden_states_shape, - ) = token_pre_all2all( - hidden_states, - expert_mask, - num_experts, - input_splits, - output_splits, - num_global_tokens_per_local_expert, - group, - ) - - expert_outputs = expert_forward( - global_permuted_hidden_states, gate_proj, up_proj, down_proj - ) - - out = tokens_post_all2all( - expert_outputs, - routing_weights, - selected_experts, - num_experts, - input_splits, - output_splits, - num_global_tokens_per_local_expert, - routing_map, - local_input_permutation_mapping, - org_hidden_states_shape, - group, - ) - return out + hidden_dim = hidden_states.size(-1) + tokens = hidden_states.reshape(-1, hidden_dim) + num_tokens = tokens.size(0) + device = tokens.device + + router_logits = torch.nn.functional.linear(tokens, gate_weight, gate_bias) + probs = torch.softmax(router_logits.float(), dim=-1) + routing_weights, selected_experts = torch.topk(probs, top_k, dim=-1) + routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) + + # Fixed per-expert capacity (capacity factor 1.0 over routed slots). + capacity = (num_tokens * top_k + num_experts - 1) // num_experts + + # Position of each routed slot inside its expert queue (token-major order). + flat_experts = selected_experts.reshape(-1) # (num_tokens * top_k,) + one_hot = torch.nn.functional.one_hot(flat_experts, num_experts) + pos_in_expert = (one_hot.cumsum(dim=0) - 1).mul(one_hot).sum(dim=-1) + keep = pos_in_expert < capacity + + # Scatter kept tokens into their (expert, slot) positions. + send_buf = tokens.new_zeros(num_experts * capacity, hidden_dim) + slot = flat_experts * capacity + pos_in_expert + expanded_tokens = tokens.repeat_interleave(top_k, dim=0) + send_buf[slot[keep]] = expanded_tokens[keep] + + # Static equal-split all_to_all: chunk e goes to rank e. + recv_buf = torch.empty_like(send_buf) + dist.all_to_all_single(recv_buf, send_buf, group=group) + + # This rank's expert processes all received slots (world_size * capacity). + expert_out = expert_forward(recv_buf, gate_proj, up_proj, down_proj) + + # Return processed slots to their source ranks (same static splits). + combine_buf = torch.empty_like(expert_out) + dist.all_to_all_single(combine_buf, expert_out, group=group) + + # Weighted combine; dropped slots contribute zero. + slot_out = tokens.new_zeros(num_tokens * top_k, hidden_dim) + slot_out[keep] = combine_buf[slot[keep]] + combined = ( + slot_out.view(num_tokens, top_k, hidden_dim) + * routing_weights.unsqueeze(-1).to(slot_out.dtype) + ).sum(dim=1) + + return combined.view_as(hidden_states) diff --git a/reference/50_moe_ep_wide.py b/reference/50_moe_ep_wide.py index 737917d..603cb27 100755 --- a/reference/50_moe_ep_wide.py +++ b/reference/50_moe_ep_wide.py @@ -1,68 +1,51 @@ -# Expert-parallel (EP) fused MoE forward — WIDE EP. +# Expert-parallel (EP) fused MoE forward — WIDE EP (multiple experts per rank). # -# Same kernel as problem 31 (router -> permute -> all_to_all dispatch -> per-expert -# SiLU MLP -> all_to_all combine -> unpermute), but the harness sets -# num_experts == world_size * 2, i.e. multiple experts hosted per rank. This stresses -# the per-rank local expert loop and the larger, more fragmented all_to_all token -# dispatch relative to the balanced case. +# num_experts == world_size * 2, so each rank hosts num_local_experts = 2 +# DISTINCT experts with their own weights (stacked as (E_local, ...) tensors). +# Dispatch uses dynamic variable splits like problem 31, but the local compute +# is a grouped GEMM over ragged per-expert token groups — received tokens are +# regrouped from (source rank, expert) order to expert-major order, each local +# expert runs its own SiLU MLP on its segment, and results are regrouped and +# returned. The ragged grouped GEMM + regrouping is the wide-EP-specific +# fusion opportunity that the single-expert problems (31 / 49 / 51) don't have. from typing import List, Optional, Tuple, Union import torch import torch.distributed as dist -class _AllToAll(torch.autograd.Function): - @staticmethod - def forward(ctx, group, input, output_split_sizes, input_split_sizes): - ctx.group = group - ctx.output_split_sizes = output_split_sizes - ctx.input_split_sizes = input_split_sizes - if dist.get_world_size(group=group) == 1: - return input.contiguous() - input = input.contiguous() - if output_split_sizes is None: - output = torch.empty_like(input) - else: - output = torch.empty( - size=(sum(output_split_sizes), input.size(1)), - dtype=input.dtype, - device=input.device, - ) - dist.all_to_all_single( - output, - input, - output_split_sizes=output_split_sizes, - input_split_sizes=input_split_sizes, - group=group, - ) - return output - - @staticmethod - def backward(ctx, grad_output): - return ( - None, - _AllToAll.apply( - ctx.group, grad_output, ctx.input_split_sizes, ctx.output_split_sizes - ), - None, - None, - ) - - def _all_to_all( group: dist.ProcessGroup, input: torch.Tensor, output_split_sizes: Optional[List[int]], input_split_sizes: Optional[List[int]], ) -> torch.Tensor: - return _AllToAll.apply(group, input, output_split_sizes, input_split_sizes) + if dist.get_world_size(group=group) == 1: + return input.contiguous() + input = input.contiguous() + if output_split_sizes is None: + output = torch.empty_like(input) + else: + output = torch.empty( + size=(sum(output_split_sizes), input.size(1)), + dtype=input.dtype, + device=input.device, + ) + dist.all_to_all_single( + output, + input, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + ) + return output def _preprocess( expert_mask: torch.Tensor, num_experts: int, ep_group: dist.ProcessGroup, -) -> Tuple[List[int], List[int], torch.Tensor, torch.Tensor]: +) -> Tuple[List[int], List[int], torch.Tensor]: ep_size = ep_group.size() num_local_experts = num_experts // ep_size rank = dist.get_rank(ep_group) @@ -88,18 +71,10 @@ def _preprocess( :, start_idx:end_idx ].contiguous() output_splits = num_global_tokens_per_local_expert.sum(dim=1).tolist() - num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum( - dim=0 - ).to(torch.device("cpu"), non_blocking=True) num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view( -1, num_local_experts - ).to(torch.device("cpu"), non_blocking=True) - return ( - input_splits, - output_splits, - num_global_tokens_per_local_expert, - num_global_sum_tokens_per_local_expert, - ) + ).to(torch.device("cpu")) + return input_splits, output_splits, num_global_tokens_per_local_expert def _permute( @@ -162,114 +137,43 @@ def _unpermute( return unpermuted_tokens -def token_pre_all2all( - hidden_states: torch.Tensor, - expert_mask: torch.Tensor, - num_experts: int, - input_splits: List[int], - output_splits: List[int], - num_global_tokens_per_local_expert: torch.Tensor, - group: Optional[dist.ProcessGroup] = None, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Size]: - group = group or dist.group.WORLD - hidden_dim = hidden_states.size(-1) - hidden_states = hidden_states.reshape(-1, hidden_dim) - org_hidden_states_shape = hidden_states.shape - routing_map = expert_mask.sum(dim=1) - - local_permuted_hidden_states, local_input_permutation_mapping = _permute( - hidden_states, routing_map - ) - expected_tokens = sum(input_splits) - actual_tokens = local_permuted_hidden_states.shape[0] - if expected_tokens != actual_tokens: - raise RuntimeError( - f"EP split mismatch: input_splits sum ({expected_tokens}) != " - f"permuted tokens ({actual_tokens})" - ) - - global_permuted_hidden_states = _all_to_all( - group, local_permuted_hidden_states, output_splits, input_splits - ) - num_local_experts = num_experts // dist.get_world_size(group) - permute_order = ( - torch.arange(num_experts).reshape(-1, num_local_experts).T.ravel().tolist() - ) - split_sizes = num_global_tokens_per_local_expert.ravel().tolist() - global_permuted_hidden_states = _sort_chunks_by_idxs( - global_permuted_hidden_states, split_sizes, permute_order - ) - return ( - global_permuted_hidden_states, - routing_map, - local_input_permutation_mapping, - org_hidden_states_shape, - ) - - -def tokens_post_all2all( - expert_outputs: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - num_experts: int, - input_splits: List[int], - output_splits: List[int], - num_global_tokens_per_local_expert: torch.Tensor, - routing_map: torch.Tensor, - local_input_permutation_mapping: torch.Tensor, - org_hidden_states_shape: torch.Size, - group: Optional[dist.ProcessGroup] = None, -) -> torch.Tensor: - group = group or dist.group.WORLD - num_local_experts = num_experts // dist.get_world_size(group) - unpermute_order = ( - torch.arange(num_experts).reshape(num_local_experts, -1).T.ravel().tolist() - ) - split_sizes = num_global_tokens_per_local_expert.T.ravel().tolist() - expert_outputs = _sort_chunks_by_idxs( - expert_outputs, split_sizes, unpermute_order - ) - unpermute_outputs = _all_to_all(group, expert_outputs, input_splits, output_splits) - weights_idx = _generate_weights_idx(routing_weights, selected_experts, num_experts) - unpermute_outputs = _unpermute( - unpermute_outputs, - weights_idx, - org_hidden_states_shape, - local_input_permutation_mapping, - routing_map, - ) - return unpermute_outputs - - -def expert_forward( +def _grouped_expert_forward( x: torch.Tensor, - gate_proj: torch.nn.Linear, - up_proj: torch.nn.Linear, - down_proj: torch.nn.Linear, + tokens_per_local_expert: List[int], + w_gate: torch.Tensor, # (E_local, inter_dim, hidden_dim) + w_up: torch.Tensor, # (E_local, inter_dim, hidden_dim) + w_down: torch.Tensor, # (E_local, hidden_dim, inter_dim) ) -> torch.Tensor: - gate = torch.nn.functional.silu(gate_proj(x)) - up = up_proj(x) - return down_proj(gate * up) + """Ragged grouped GEMM: expert e's SiLU MLP on its token segment.""" + outputs = [] + for e, seg in enumerate(torch.split(x, tokens_per_local_expert, dim=0)): + gate = torch.nn.functional.silu(seg @ w_gate[e].T) + up = seg @ w_up[e].T + outputs.append((gate * up) @ w_down[e].T) + return torch.cat(outputs, dim=0) +@torch.no_grad() def solution( hidden_states: torch.Tensor, gate_weight: torch.Tensor, gate_bias: Optional[torch.Tensor], - gate_proj: torch.nn.Linear, - up_proj: torch.nn.Linear, - down_proj: torch.nn.Linear, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, num_experts: int, top_k: int, group: Optional[dist.ProcessGroup] = None, ) -> torch.Tensor: group = group or dist.group.WORLD + world_size = dist.get_world_size(group) + num_local_experts = num_experts // world_size + assert num_local_experts >= 2, "wide EP hosts multiple experts per rank" + hidden_dim = hidden_states.size(-1) - num_tokens = hidden_states.reshape(-1, hidden_dim).size(0) + tokens = hidden_states.reshape(-1, hidden_dim) - router_logits = torch.nn.functional.linear( - hidden_states.reshape(-1, hidden_dim), gate_weight, gate_bias - ) + router_logits = torch.nn.functional.linear(tokens, gate_weight, gate_bias) routing_weights, selected_experts = torch.topk( torch.softmax(router_logits, dim=-1), top_k, dim=-1 ) @@ -277,40 +181,39 @@ def solution( selected_experts, num_classes=num_experts ).permute(2, 1, 0) - input_splits, output_splits, num_global_tokens_per_local_expert, _ = _preprocess( + input_splits, output_splits, num_global_tokens_per_local_expert = _preprocess( expert_mask, num_experts, group ) - ( - global_permuted_hidden_states, - routing_map, - local_input_permutation_mapping, - org_hidden_states_shape, - ) = token_pre_all2all( - hidden_states, - expert_mask, - num_experts, - input_splits, - output_splits, - num_global_tokens_per_local_expert, - group, + routing_map = expert_mask.sum(dim=1) + permuted, permutation_mapping = _permute(tokens, routing_map) + org_shape = tokens.shape + + global_tokens = _all_to_all(group, permuted, output_splits, input_splits) + + # Regroup from (source rank, expert)-order to expert-major order. + permute_order = ( + torch.arange(num_experts).reshape(-1, num_local_experts).T.ravel().tolist() ) + split_sizes = num_global_tokens_per_local_expert.ravel().tolist() + global_tokens = _sort_chunks_by_idxs(global_tokens, split_sizes, permute_order) - expert_outputs = expert_forward( - global_permuted_hidden_states, gate_proj, up_proj, down_proj + tokens_per_local_expert = ( + num_global_tokens_per_local_expert.sum(dim=0).tolist() + ) + expert_outputs = _grouped_expert_forward( + global_tokens, tokens_per_local_expert, w_gate, w_up, w_down ) - out = tokens_post_all2all( - expert_outputs, - routing_weights, - selected_experts, - num_experts, - input_splits, - output_splits, - num_global_tokens_per_local_expert, - routing_map, - local_input_permutation_mapping, - org_hidden_states_shape, - group, + # Regroup back to (source rank, expert)-order and return tokens. + unpermute_order = ( + torch.arange(num_experts).reshape(num_local_experts, -1).T.ravel().tolist() ) - return out + back_split_sizes = num_global_tokens_per_local_expert.T.ravel().tolist() + expert_outputs = _sort_chunks_by_idxs(expert_outputs, back_split_sizes, unpermute_order) + + returned = _all_to_all(group, expert_outputs, input_splits, output_splits) + + weights_idx = _generate_weights_idx(routing_weights, selected_experts, num_experts) + out = _unpermute(returned, weights_idx, org_shape, permutation_mapping, routing_map) + return out.view_as(hidden_states) diff --git a/reference/52_fp8_reduce_scatter_grads.py b/reference/52_fp8_reduce_scatter_grads.py index c87f662..62b208b 100755 --- a/reference/52_fp8_reduce_scatter_grads.py +++ b/reference/52_fp8_reduce_scatter_grads.py @@ -1,3 +1,11 @@ +# FP8 (E4M3) gradient reduce-scatter with the fp8 payload on the wire. +# +# Each rank quantizes its flat gradient with a delayed-scaling factor (amax +# history), exchanges the fp8 chunks (as raw bytes) plus each rank's scalar +# scale with an all-to-all / all-gather, then dequantizes the received chunks +# with the sender's scale and averages them into the local shard. NCCL cannot +# reduce fp8 directly, so dequant + reduce happens on the receiving rank — +# a natural target for a fused kernel. from __future__ import annotations import torch @@ -14,14 +22,6 @@ def _update_amax_history(amax_history: Tensor, cur_abs_max: Tensor) -> Tensor: return out -@torch.no_grad() -def _fp8_round_trip_bf16(x: Tensor, scale: Tensor) -> Tensor: - xf = x.float() - qs = xf / scale - q = qs.to(torch.float8_e4m3fn) - return (q.float() * scale).to(dtype=x.dtype) - - @torch.no_grad() def solution(flat_grads: Tensor, amax_history: Tensor) -> tuple[Tensor, Tensor]: world_size = dist.get_world_size() @@ -32,10 +32,20 @@ def solution(flat_grads: Tensor, amax_history: Tensor) -> tuple[Tensor, Tensor]: updated_hist = _update_amax_history(amax_history, cur_abs_max) scale = updated_hist.max().clamp(min=1e-12).to(torch.float32) / _FP8_E4M3_MAX - recon = _fp8_round_trip_bf16(flat_grads, scale) - out_shard = torch.empty(shard_elems, dtype=flat_grads.dtype, device=flat_grads.device) - dist.reduce_scatter_tensor(out_shard, recon.contiguous(), op=dist.ReduceOp.SUM) - out_shard.div_(world_size) + # Quantize the full local gradient with the rank-local delayed scale. + q = (flat_grads.float() / scale).to(torch.float8_e4m3fn).contiguous() + + # Every peer needs the sender's scale to dequantize. + scales = torch.empty(world_size, dtype=torch.float32, device=flat_grads.device) + dist.all_gather_into_tensor(scales, scale.reshape(1)) + + # Ship fp8 chunks as raw bytes: rank r receives chunk r from every peer. + q_recv_u8 = torch.empty(n, dtype=torch.uint8, device=flat_grads.device) + dist.all_to_all_single(q_recv_u8, q.view(torch.uint8)) + + # Dequantize each received chunk with its sender's scale, then average. + chunks = q_recv_u8.view(torch.float8_e4m3fn).float().reshape(world_size, shard_elems) + out_shard = (chunks * scales.unsqueeze(1)).sum(dim=0).div_(world_size) - return out_shard, updated_hist + return out_shard.to(dtype=flat_grads.dtype), updated_hist diff --git a/reference/53_fp8_allgather_params.py b/reference/53_fp8_allgather_params.py index 377346c..b784874 100755 --- a/reference/53_fp8_allgather_params.py +++ b/reference/53_fp8_allgather_params.py @@ -1,3 +1,10 @@ +# FP8 (E4M3) parameter all-gather with the fp8 payload on the wire. +# +# Each rank quantizes its parameter shard with a delayed-scaling factor (amax +# history), all-gathers the fp8 bytes plus each rank's scalar scale, and +# dequantizes the assembled full parameter with per-sender scales. The wire +# traffic is 1/2 of bf16; dequantization on the receiving side is the fusion +# opportunity. from __future__ import annotations import torch @@ -14,14 +21,6 @@ def _update_amax_history(amax_history: Tensor, cur_abs_max: Tensor) -> Tensor: return out -@torch.no_grad() -def _fp8_round_trip_bf16(x: Tensor, scale: Tensor) -> Tensor: - xf = x.float() - qs = xf / scale - q = qs.to(torch.float8_e4m3fn) - return (q.float() * scale).to(dtype=x.dtype) - - @torch.no_grad() def solution(flat_param_shard: Tensor, amax_history: Tensor) -> tuple[Tensor, Tensor]: world_size = dist.get_world_size() @@ -31,9 +30,18 @@ def solution(flat_param_shard: Tensor, amax_history: Tensor) -> tuple[Tensor, Te updated_hist = _update_amax_history(amax_history, cur_abs_max) scale = updated_hist.max().clamp(min=1e-12).to(torch.float32) / _FP8_E4M3_MAX - recon = _fp8_round_trip_bf16(flat_param_shard, scale) - full = torch.empty(world_size * p, dtype=flat_param_shard.dtype, device=flat_param_shard.device) - dist.all_gather_into_tensor(full, recon.contiguous()) + q = (flat_param_shard.float() / scale).to(torch.float8_e4m3fn).contiguous() + + scales = torch.empty(world_size, dtype=torch.float32, device=flat_param_shard.device) + dist.all_gather_into_tensor(scales, scale.reshape(1)) + + q_full_u8 = torch.empty(world_size * p, dtype=torch.uint8, device=flat_param_shard.device) + dist.all_gather_into_tensor(q_full_u8, q.view(torch.uint8)) + + full = ( + q_full_u8.view(torch.float8_e4m3fn).float().reshape(world_size, p) + * scales.unsqueeze(1) + ).reshape(-1) - return full, updated_hist + return full.to(dtype=flat_param_shard.dtype), updated_hist diff --git a/reference/63_deepmd_kalman_filter_optimizer.py b/reference/63_deepmd_kalman_filter_optimizer.py index 47ef2f6..01cd313 100755 --- a/reference/63_deepmd_kalman_filter_optimizer.py +++ b/reference/63_deepmd_kalman_filter_optimizer.py @@ -1,3 +1,11 @@ +# DeePMD-style layer-wise Kalman filter optimizer step (data-parallel). +# +# Each rank holds the same block structure (num_blocks blocks of identical +# size). The innovation denominator is all-reduced across ranks, each rank +# updates its local weight blocks and covariance blocks, and the updated +# weights are all-gathered so every rank sees every rank's blocks. Inputs are +# never mutated; per-block updates are independent (batched-GEMM / fusion +# opportunity) and the weight all-gather is a single fixed-size collective. from typing import List, Tuple import torch @@ -14,10 +22,12 @@ def solution( kalman_nue: float = 0.9987, ) -> Tuple[List[torch.Tensor], List[torch.Tensor], torch.Tensor]: weights_num = len(weights) - lam = torch.as_tensor(kalman_lambda, dtype=weights[0].dtype, device=weights[0].device) - err = error.to(device=weights[0].device, dtype=weights[0].dtype) + dtype = weights[0].dtype + device = weights[0].device + lam = torch.as_tensor(kalman_lambda, dtype=dtype, device=device) + err = error.to(device=device, dtype=dtype) - tmp = 0 + tmp = torch.zeros((1, 1), dtype=dtype, device=device) for i in range(weights_num): tmp = tmp + (lam + torch.matmul(torch.matmul(H[i].T, P[i]), H[i])) @@ -26,37 +36,30 @@ def solution( A = 1 / tmp + new_weights: List[torch.Tensor] = [] + new_P: List[torch.Tensor] = [] for i in range(weights_num): K = torch.matmul(P[i], H[i]) - weights[i] = weights[i] + A * err * K - P[i] = (1 / lam) * (P[i] - A * torch.matmul(K, K.T)) + new_weights.append(weights[i] + A * err * K) + new_P.append((1 / lam) * (P[i] - A * torch.matmul(K, K.T))) + out_weights = new_weights if dist.is_initialized(): - device = weights[0].device - local_shape = [tensor.shape[0] for tensor in weights] + world_size = dist.get_world_size() + block_sizes = [int(w.shape[0]) for w in new_weights] + local_flat = torch.cat([w.reshape(-1) for w in new_weights], dim=0).contiguous() - shape_list = [ - torch.zeros_like(torch.empty(1), dtype=torch.float64, device=device) - for _ in range(dist.get_world_size()) - ] - dist.all_gather_object(shape_list, local_shape) + gathered = torch.empty( + world_size * local_flat.numel(), dtype=dtype, device=device + ) + dist.all_gather_into_tensor(gathered, local_flat) - weight_tensor = torch.cat([w.reshape(-1) for w in weights], dim=0).to(torch.float64) - world_shape = [sum(inner_list) for inner_list in shape_list] - - weight_list = [ - torch.zeros(world_shape[i], dtype=torch.float64, device=device) - for i in range(len(world_shape)) - ] - dist.all_gather(weight_list, weight_tensor) - - result: List[torch.Tensor] = [] - for i in range(dist.get_world_size()): - result = result + [ - t.reshape(-1, 1).to(dtype=weights[0].dtype) - for t in torch.split(weight_list[i], shape_list[i]) - ] - weights = result + out_weights = [] + for r in range(world_size): + rank_flat = gathered[r * local_flat.numel() : (r + 1) * local_flat.numel()] + out_weights.extend( + t.reshape(-1, 1) for t in torch.split(rank_flat, block_sizes) + ) kalman_lambda_next = ( torch.as_tensor(kalman_nue, dtype=lam.dtype, device=lam.device) * lam @@ -64,4 +67,4 @@ def solution( - torch.as_tensor(kalman_nue, dtype=lam.dtype, device=lam.device) ) - return weights, P, kalman_lambda_next + return out_weights, new_P, kalman_lambda_next diff --git a/reference/64_gnn_neighbor_sampling.py b/reference/64_gnn_neighbor_sampling.py index f1cfaad..66fc521 100755 --- a/reference/64_gnn_neighbor_sampling.py +++ b/reference/64_gnn_neighbor_sampling.py @@ -1,6 +1,14 @@ +# Distributed multi-hop GNN neighbor sampling over a partitioned CSC graph. +# +# Each hop: route frontier nodes to their owner rank (all_to_all), take the +# first `fanout` neighbors of each received node from the local CSC shard +# (deterministic selection so candidates can be checked bit-for-bit), +# return the sampled neighborhoods to the requesting ranks (all_to_all), +# dedup the new frontier, and finally relabel the sampled subgraph. All graph +# work is vectorized on-device; the host only touches the split-size metadata +# that all_to_all requires. from typing import List, Optional, Tuple -import numpy as np import torch import torch.distributed as dist @@ -11,45 +19,53 @@ def _sample_one_hop_csc_dist( colptr: torch.Tensor, row: torch.Tensor, replace: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, List[int]]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Deterministically take the first min(k, deg) neighbors of each input node. + + Returns (input_nodes ++ sampled_neighbors, sampled_edge_ids, per-node counts). + """ + device = input_nodes.device n = input_nodes.numel() - sampled_nodes = [] - sampled_edges = [] - cumsum = [n] - - for i in range(n): - v = int(input_nodes[i].item()) - start = int(colptr[v].item()) - end = int(colptr[v + 1].item()) - deg = end - start - take = min(k, deg) if k >= 0 else deg - - if take > 0: - perm = torch.arange(take, device=input_nodes.device) - sampled_nodes.append(row[start:end].index_select(0, perm)) - sampled_edges.append(torch.arange(start, end, device=input_nodes.device).index_select(0, perm)) - - cumsum.append(cumsum[-1] + take) - - nbr_tensor = ( - torch.cat(sampled_nodes) - if sampled_nodes - else torch.empty(0, dtype=torch.long, device=input_nodes.device) - ) - eid_tensor = ( - torch.cat(sampled_edges) - if sampled_edges - else torch.empty(0, dtype=torch.long, device=input_nodes.device) - ) - return torch.cat([input_nodes, nbr_tensor]), eid_tensor, cumsum + if n == 0: + empty = torch.empty(0, dtype=torch.long, device=device) + return input_nodes, empty, empty + + start = colptr[input_nodes] + end = colptr[input_nodes + 1] + deg = end - start + take = deg.clamp(max=k) if k >= 0 else deg + + total = int(take.sum().item()) + if total == 0: + empty = torch.empty(0, dtype=torch.long, device=device) + return input_nodes, empty, take + + seg = torch.repeat_interleave(torch.arange(n, device=device), take) + out_ptr = torch.zeros(n + 1, dtype=torch.long, device=device) + out_ptr[1:] = take.cumsum(0) + local_off = torch.arange(total, device=device) - out_ptr[seg] + eid = start[seg] + local_off + nbr = row[eid] + return torch.cat([input_nodes, nbr]), eid, take def _remove_duplicates(out_node: torch.Tensor, node: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Keep first occurrences of cat([node, out_node]) in encounter order.""" num_nodes = node.numel() - node_combined = torch.cat([node, out_node]) - _, idx = np.unique(node_combined.cpu().numpy(), return_index=True) - idx = torch.from_numpy(idx).to(node.device).sort().values - node = node_combined[idx] + combined = torch.cat([node, out_node]) + uniq, inverse = torch.unique(combined, return_inverse=True) + first_idx = torch.full( + (uniq.numel(),), combined.numel(), dtype=torch.long, device=combined.device + ) + first_idx.scatter_reduce_( + 0, + inverse, + torch.arange(combined.numel(), device=combined.device), + reduce="amin", + include_self=True, + ) + order = first_idx.sort().values + node = combined[order] src = node[num_nodes:] return src, node @@ -75,16 +91,14 @@ def _relabel_neighborhood( def _exchange_nodes( - send_nodes_list: List[torch.Tensor], + send_nodes: torch.Tensor, + send_counts: torch.Tensor, group: dist.ProcessGroup, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - world_size = dist.get_world_size(group) - device = send_nodes_list[0].device - send_counts = torch.tensor([x.numel() for x in send_nodes_list], dtype=torch.long, device=device) +) -> Tuple[torch.Tensor, torch.Tensor]: + device = send_nodes.device recv_counts = torch.empty_like(send_counts) dist.all_to_all_single(recv_counts, send_counts, group=group) - send_nodes = torch.cat(send_nodes_list) if send_nodes_list else torch.empty(0, dtype=torch.long, device=device) recv_nodes = torch.empty(int(recv_counts.sum().item()), dtype=torch.long, device=device) dist.all_to_all_single( recv_nodes, @@ -93,7 +107,7 @@ def _exchange_nodes( output_split_sizes=recv_counts.cpu().tolist(), group=group, ) - return recv_nodes, send_counts, recv_counts + return recv_nodes, recv_counts def _exchange_replies( @@ -103,14 +117,13 @@ def _exchange_replies( recv_counts: torch.Tensor, group: dist.ProcessGroup, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - world_size = dist.get_world_size(group) device = sampled_nodes.device - recv_splits = recv_counts.cpu().tolist() - send_node_counts = torch.empty(world_size, dtype=torch.long, device=device) - offset = 0 - for r, count in enumerate(recv_splits): - send_node_counts[r] = sampled_counts[offset : offset + count].sum() - offset += count + # sampled_counts is grouped by requesting rank (recv_counts entries per rank). + rank_bounds = torch.zeros(recv_counts.numel() + 1, dtype=torch.long, device=device) + rank_bounds[1:] = recv_counts.cumsum(0) + counts_cumsum = torch.zeros(sampled_counts.numel() + 1, dtype=torch.long, device=device) + counts_cumsum[1:] = sampled_counts.cumsum(0) + send_node_counts = counts_cumsum[rank_bounds[1:]] - counts_cumsum[rank_bounds[:-1]] reply_node_counts = torch.empty_like(send_node_counts) dist.all_to_all_single(reply_node_counts, send_node_counts, group=group) @@ -139,7 +152,7 @@ def _exchange_replies( dist.all_to_all_single( reply_counts, sampled_counts, - input_split_sizes=recv_splits, + input_split_sizes=recv_counts.cpu().tolist(), output_split_sizes=reply_count_counts.cpu().tolist(), group=group, ) @@ -171,55 +184,44 @@ def solution( if src.numel() == 0: break + # Stable-partition the frontier by owner rank. partition_ids = node_to_rank[src].to(torch.long) - partition_orders = torch.empty_like(partition_ids) - send_nodes_list = [] - send_pos_list = [] - for r in range(world_size): - pos = (partition_ids == r).nonzero(as_tuple=False).flatten() - partition_orders[pos] = torch.arange(pos.numel(), dtype=torch.long, device=device) - send_nodes_list.append(src[pos]) - send_pos_list.append(pos) - - recv_nodes, send_counts, recv_counts = _exchange_nodes(send_nodes_list, group) - node_out, edge_out, cumsum = _sample_one_hop_csc_dist( - recv_nodes, int(fanout), local_adj_row_ptr, local_adj_col, replace - ) + order = torch.argsort(partition_ids, stable=True) + send_nodes = src[order] + send_counts = torch.bincount(partition_ids, minlength=world_size) - seed_size = recv_nodes.numel() - sampled_nodes = node_out[seed_size:] - sampled_counts = torch.tensor( - np.subtract(np.array(cumsum[1:]), np.array(cumsum[:-1])), - dtype=torch.long, - device=device, + recv_nodes, recv_counts = _exchange_nodes(send_nodes, send_counts, group) + node_out, edge_out, sampled_counts = _sample_one_hop_csc_dist( + recv_nodes, int(fanout), local_adj_row_ptr, local_adj_col, replace ) + sampled_nodes = node_out[recv_nodes.numel():] reply_nodes, reply_edges, reply_counts = _exchange_replies( sampled_nodes, edge_out, sampled_counts, recv_counts, group ) - rank_offsets = torch.cat( - [send_counts.new_zeros(1), torch.cumsum(send_counts, dim=0)[:-1]] - ) - grouped_index = rank_offsets[partition_ids] + partition_orders - node_chunks = list(torch.split(reply_nodes, reply_counts.cpu().tolist())) - edge_chunks = list(torch.split(reply_edges, reply_counts.cpu().tolist())) - - ordered_nodes = [] - ordered_edges = [] - ordered_dst = [] - for idx in grouped_index.tolist(): - ordered_nodes.append(node_chunks[idx]) - ordered_edges.append(edge_chunks[idx]) - for dst_node, count in zip(src, reply_counts[grouped_index]): - ordered_dst.append(dst_node.repeat(int(count.item()))) - - out_node = torch.cat(ordered_nodes) if ordered_nodes else seed.new_empty(0) - out_edge = torch.cat(ordered_edges) if ordered_edges else seed.new_empty(0) - out_dst = torch.cat(ordered_dst) if ordered_dst else seed.new_empty(0) - if out_node.numel() == 0: + # Reply chunks arrive in send order (frontier sorted by owner rank); + # build gather indices that reorder them back to original frontier order. + chunk_off = torch.zeros(reply_counts.numel() + 1, dtype=torch.long, device=device) + chunk_off[1:] = reply_counts.cumsum(0) + + total = int(reply_counts.sum().item()) + if total == 0: break + inv = torch.empty_like(order) + inv[order] = torch.arange(order.numel(), device=device) + take = reply_counts[inv] # counts per frontier node in original order + seg = torch.repeat_interleave(torch.arange(src.numel(), device=device), take) + out_ptr = torch.zeros(src.numel() + 1, dtype=torch.long, device=device) + out_ptr[1:] = take.cumsum(0) + local_off = torch.arange(total, device=device) - out_ptr[seg] + gidx = chunk_off[inv][seg] + local_off + + out_node = reply_nodes[gidx] + out_edge = reply_edges[gidx] + out_dst = torch.repeat_interleave(src, take) + src, node = _remove_duplicates(out_node, node) node_with_dupl.append(out_node) dst_with_dupl.append(out_dst) @@ -228,4 +230,4 @@ def solution( node_dupl = torch.cat(node_with_dupl) dst_dupl = torch.cat(dst_with_dupl) row, col = _relabel_neighborhood(node, dst_dupl, node_dupl) - return node, row, col, torch.cat(edge) \ No newline at end of file + return node, row, col, torch.cat(edge) diff --git a/reference/78_magi1_tile_parallel_vae_decode.py b/reference/78_magi1_tile_parallel_vae_decode.py index ba4aff1..a54d19d 100755 --- a/reference/78_magi1_tile_parallel_vae_decode.py +++ b/reference/78_magi1_tile_parallel_vae_decode.py @@ -50,16 +50,22 @@ def _gather_tiles( return tiles world_size = dist.get_world_size(group=group) - local_shapes = [tuple(tile.shape) for tile in tiles] - all_shapes: List[List[Tuple[int, ...]]] = [[] for _ in range(world_size)] - dist.all_gather_object(all_shapes, local_shapes, group=group) - - local_flat = ( - torch.cat([tile.reshape(-1).contiguous() for tile in tiles], dim=0) - if tiles - else template.new_empty(0) + device = template.device + max_tiles = -(-len(global_order) // world_size) + + # Exchange tile shapes as a fixed-size tensor (5 dims each, -1 padded). + shape_buf = torch.full((max_tiles, 5), -1, dtype=torch.long, device=device) + for i, tile in enumerate(tiles): + shape_buf[i] = torch.tensor(tile.shape, dtype=torch.long, device=device) + all_shape_buf = torch.empty( + (world_size * max_tiles, 5), dtype=torch.long, device=device ) - local_size = int(local_flat.numel()) + dist.all_gather_into_tensor(all_shape_buf, shape_buf, group=group) + all_shapes_list = all_shape_buf.view(world_size, max_tiles, 5).cpu().tolist() + all_shapes: List[List[Tuple[int, ...]]] = [ + [tuple(s) for s in shapes if s[0] >= 0] for shapes in all_shapes_list + ] + rank_sizes: List[int] = [] for shapes in all_shapes: total = 0 @@ -69,20 +75,22 @@ def _gather_tiles( numel *= size total += numel rank_sizes.append(total) - send = local_flat.repeat(world_size) - recv = template.new_empty(sum(rank_sizes)) - dist.all_to_all_single( - recv, - send, - output_split_sizes=rank_sizes, - input_split_sizes=[local_size] * world_size, - group=group, + + # Pad flattened payloads to the max rank size and all-gather once. + local_flat = ( + torch.cat([tile.reshape(-1).contiguous() for tile in tiles], dim=0) + if tiles + else template.new_empty(0) ) + max_size = max(rank_sizes) + send = template.new_zeros(max_size) + send[: local_flat.numel()] = local_flat + recv = template.new_empty(world_size * max_size) + dist.all_gather_into_tensor(recv, send, group=group) gathered: List[torch.Tensor] = [] - offset = 0 - for shapes, total in zip(all_shapes, rank_sizes): - rank_buf = recv[offset : offset + total] + for r, shapes in enumerate(all_shapes): + rank_buf = recv[r * max_size : r * max_size + rank_sizes[r]] rank_offset = 0 for shape in shapes: numel = 1 @@ -90,33 +98,47 @@ def _gather_tiles( numel *= size gathered.append(rank_buf[rank_offset : rank_offset + numel].view(shape)) rank_offset += numel - offset += total by_index = {tile_idx: tile for tile_idx, tile in zip(global_order, gathered)} return [by_index[idx] for idx in sorted(by_index)] +def _ramp(extent: int, dim: int, like: torch.Tensor) -> torch.Tensor: + """Blend ramp [0, 1) broadcastable along `dim` of a 5-D tensor.""" + shape = [1] * 5 + shape[dim] = extent + return ( + torch.arange(extent, device=like.device, dtype=like.dtype) / extent + ).reshape(shape) + + def _blend_t(prev: torch.Tensor, cur: torch.Tensor, extent: int) -> torch.Tensor: extent = min(prev.shape[2], cur.shape[2], extent) - for idx in range(extent): - ratio = idx / extent - cur[:, :, idx] = prev[:, :, -extent + idx] * (1.0 - ratio) + cur[:, :, idx] * ratio + if extent > 0: + ratio = _ramp(extent, 2, cur) + cur[:, :, :extent] = ( + prev[:, :, -extent:] * (1.0 - ratio) + cur[:, :, :extent] * ratio + ) return cur def _blend_h(prev: torch.Tensor, cur: torch.Tensor, extent: int) -> torch.Tensor: extent = min(prev.shape[3], cur.shape[3], extent) - for idx in range(extent): - ratio = idx / extent - cur[:, :, :, idx] = prev[:, :, :, -extent + idx] * (1.0 - ratio) + cur[:, :, :, idx] * ratio + if extent > 0: + ratio = _ramp(extent, 3, cur) + cur[:, :, :, :extent] = ( + prev[:, :, :, -extent:] * (1.0 - ratio) + cur[:, :, :, :extent] * ratio + ) return cur def _blend_w(prev: torch.Tensor, cur: torch.Tensor, extent: int) -> torch.Tensor: extent = min(prev.shape[4], cur.shape[4], extent) - for idx in range(extent): - ratio = idx / extent - cur[:, :, :, :, idx] = prev[:, :, :, :, -extent + idx] * (1.0 - ratio) + cur[:, :, :, :, idx] * ratio + if extent > 0: + ratio = _ramp(extent, 4, cur) + cur[:, :, :, :, :extent] = ( + prev[:, :, :, :, -extent:] * (1.0 - ratio) + cur[:, :, :, :, :extent] * ratio + ) return cur diff --git a/reference/82_vocab_parallel_log_prob_topk.py b/reference/82_vocab_parallel_log_prob_topk.py index 9e33c43..91deced 100755 --- a/reference/82_vocab_parallel_log_prob_topk.py +++ b/reference/82_vocab_parallel_log_prob_topk.py @@ -10,103 +10,103 @@ def _apply_top_k_top_p( top_k: Optional[int], top_p: float, ) -> torch.Tensor: - need_k = top_k is not None and top_k > 0 - need_p = top_p is not None and top_p < 1.0 - - if not need_k and not need_p: - return logits - - original_shape = logits.shape - vocab_size = logits.shape[-1] + need_k = top_k is not None and top_k > 0 + need_p = top_p is not None and top_p < 1.0 + + if not need_k and not need_p: + return logits + + original_shape = logits.shape + vocab_size = logits.shape[-1] logits_2d = logits.reshape(-1, vocab_size) if need_k: top_k = min(int(top_k), vocab_size) - - if need_k and not need_p: - top_k_values, _ = torch.topk(logits_2d, top_k, dim=-1) - threshold = top_k_values[..., -1:].expand_as(logits_2d) - keep_mask = logits_2d >= threshold - filtered = torch.where( - keep_mask, - logits_2d, - torch.full_like(logits_2d, float("-inf")), - ) + + if need_k and not need_p: + top_k_values, _ = torch.topk(logits_2d, top_k, dim=-1) + threshold = top_k_values[..., -1:].expand_as(logits_2d) + keep_mask = logits_2d >= threshold + filtered = torch.where( + keep_mask, + logits_2d, + torch.full_like(logits_2d, float("-inf")), + ) return filtered.reshape(original_shape) - - logits_sort, logits_idx = logits_2d.sort(dim=-1, descending=False) - - top_k_mask = None - if need_k: - top_k_index = logits_sort.size(-1) - top_k - threshold = logits_sort.gather( - -1, - torch.full( - logits_sort.shape[:-1], - top_k_index, - device=logits_2d.device, - dtype=torch.long, - ).unsqueeze(-1), - ) - top_k_mask = logits_sort >= threshold - logits_sort = logits_sort.masked_fill(~top_k_mask, float("-inf")) - - probs_sort = logits_sort.softmax(dim=-1) - probs_sum = torch.cumsum(probs_sort, dim=-1) - top_p_mask = probs_sum > 1 - top_p - top_p_mask[..., -1] = True # always keep at least one token - logits_sort = logits_sort.masked_fill(~top_p_mask, float("-inf")) - - filtered = logits_sort.scatter(dim=-1, index=logits_idx, src=logits_sort) + + logits_sort, logits_idx = logits_2d.sort(dim=-1, descending=False) + + top_k_mask = None + if need_k: + top_k_index = logits_sort.size(-1) - top_k + threshold = logits_sort.gather( + -1, + torch.full( + logits_sort.shape[:-1], + top_k_index, + device=logits_2d.device, + dtype=torch.long, + ).unsqueeze(-1), + ) + top_k_mask = logits_sort >= threshold + logits_sort = logits_sort.masked_fill(~top_k_mask, float("-inf")) + + probs_sort = logits_sort.softmax(dim=-1) + probs_sum = torch.cumsum(probs_sort, dim=-1) + top_p_mask = probs_sum > 1 - top_p + top_p_mask[..., -1] = True # always keep at least one token + logits_sort = logits_sort.masked_fill(~top_p_mask, float("-inf")) + + filtered = logits_sort.scatter(dim=-1, index=logits_idx, src=logits_sort) return filtered.reshape(original_shape) def _all_to_all_vp_to_seq( vocab_parallel_logits: torch.Tensor, - tp_group: dist.ProcessGroup, + tp_group: dist.ProcessGroup, ) -> torch.Tensor: - world_size = dist.get_world_size(tp_group) + world_size = dist.get_world_size(tp_group) num_tokens, local_vocab = vocab_parallel_logits.shape local_tokens = num_tokens // world_size - + input_flat = vocab_parallel_logits.contiguous().flatten() - output_flat = torch.empty_like(input_flat) - dist.all_to_all_single(output_flat, input_flat, group=tp_group) - + output_flat = torch.empty_like(input_flat) + dist.all_to_all_single(output_flat, input_flat, group=tp_group) + output = output_flat.view(world_size, local_tokens, local_vocab) return output.permute(1, 0, 2).reshape(local_tokens, world_size * local_vocab) - -@torch.no_grad() -def solution( + +@torch.no_grad() +def solution( vocab_parallel_logits: torch.Tensor, target: torch.Tensor, tp_group: Optional[dist.ProcessGroup] = None, - top_k: Optional[int] = None, - top_p: float = 1.0, + top_k: Optional[int] = None, + top_p: float = 1.0, ) -> torch.Tensor: tp_group = tp_group or dist.group.WORLD - world_size = dist.get_world_size(tp_group) - rank = dist.get_rank(tp_group) + world_size = dist.get_world_size(tp_group) + rank = dist.get_rank(tp_group) batch, seq_len, local_vocab = vocab_parallel_logits.shape num_tokens = batch * seq_len - + if num_tokens % world_size != 0: - raise ValueError( + raise ValueError( f"B*S={num_tokens} must be divisible by tensor parallel size {world_size}" - ) + ) local_tokens = num_tokens // world_size - + logits_2d = vocab_parallel_logits.reshape(num_tokens, local_vocab) target_flat = target.reshape(-1) target_local = target_flat[rank * local_tokens : (rank + 1) * local_tokens] seq_parallel_logits = _all_to_all_vp_to_seq(logits_2d, tp_group) logits = _apply_top_k_top_p(seq_parallel_logits, top_k=top_k, top_p=top_p) - log_probs = F.log_softmax(logits.to(dtype=torch.float32), dim=-1) + log_probs = F.log_softmax(logits.to(dtype=torch.float32), dim=-1) token_logprobs = torch.gather(log_probs, -1, target_local.unsqueeze(-1)) token_logprobs = token_logprobs.squeeze(-1) gathered = [torch.empty_like(token_logprobs) for _ in range(world_size)] - dist.all_gather(gathered, token_logprobs, group=tp_group) + dist.all_gather(gathered, token_logprobs, group=tp_group) return torch.cat(gathered, dim=0).reshape(batch, seq_len) \ No newline at end of file diff --git a/utils/input_output_tensors.py b/utils/input_output_tensors.py index ccec4a5..bd6c583 100755 --- a/utils/input_output_tensors.py +++ b/utils/input_output_tensors.py @@ -48,7 +48,7 @@ def save_tensor(output, logs_dir: str, rank: int) -> str: # INPUT TENSOR STANDARD (tuple-only) # --------------------------------------------------------------------------- # create_input_tensor() returns a tuple unpacked as solution_fn(*x). Entries are usually tensors -# but may include Python scalars / dicts / dataclasses (e.g. problem 4, problems 100–105). +# but may include Python scalars / lists / nn.Modules (e.g. problems 4, 21, 31, 70). # - solution(tensor) for single-tensor problems: x is (tensor,) # - solution(t1, t2) for multi-arg problems: x is (t1, t2, ...) # Output from solution_fn may still be a single tensor or a tuple; save_tensor() handles both. @@ -411,14 +411,17 @@ def create_input_tensor( _seed(problem_id, rank, trial) return (torch.randn(base_shape, dtype=dtype, device=dev), 128) - # 20: blocked_fp8_dequantize + # 20: blocked_fp8_dequantize — inputs are genuinely quantized: per-block + # absmax scales (positive) and quantized mantissas derived from them. elif problem_id == 20: _seed(problem_id, rank, trial) chunk_numel = M * N block_size = 128 - num_blocks_per_chunk = chunk_numel // block_size - local_y = torch.randn((world_size, M, N), dtype=dtype, device=dev) - local_s = torch.randn((world_size, num_blocks_per_chunk), dtype=dtype, device=dev) + x = torch.randn((world_size, M, N), dtype=torch.float32, device=dev) + xb = x.reshape(-1, block_size) + s = (xb.abs().amax(dim=1) / 448.0).clamp(min=1e-8) + local_y = (xb / s.unsqueeze(1)).reshape(world_size, M, N).to(dtype) + local_s = s.reshape(world_size, chunk_numel // block_size).to(dtype) return (local_y, local_s, block_size) # 21: clip_grad_norm_no_ep @@ -446,12 +449,22 @@ def create_input_tensor( grad_loss_sum = torch.zeros((), dtype=dtype, device=dev) return (loss, local_valid, global_valid, grad_normalized_loss, grad_loss_sum) - # 24: load_balancing_loss_fn + # 24: load_balancing_loss_fn — multi-layer gate logits (tuple) plus a real + # attention mask so the masked normalization path is exercised; the local + # compute (softmax/one-hot/masked reductions over L*B*S tokens) is the + # fusion target ahead of the scalar all-reduce. elif problem_id == 24: _seed(problem_id, rank, trial) num_experts = 8 - gate_logits = torch.randn((M, num_experts), dtype=dtype, device=dev) - return (gate_logits, num_experts, 2, None) + num_layers = 4 + batch = 8 + seq = max(32, M // 8) + gate_logits = tuple( + torch.randn((batch * seq, num_experts), dtype=dtype, device=dev) + for _ in range(num_layers) + ) + attention_mask = (torch.rand((batch, seq), device=dev) > 0.1).to(dtype) + return (gate_logits, num_experts, 2, attention_mask) # 25: importance_sampling_loss elif problem_id == 25: @@ -561,14 +574,20 @@ def create_input_tensor( org_hidden_states_shape = torch.Size([num_tokens, N]) 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) - # 30: moe_epgroupgemm_lora_backward + # 30: moe_epgroupgemm_lora_backward — per-local-expert activations and + # output-grads; the solution computes the shared LoRA adapter grads via + # grouped GEMMs then all-reduces them across the EP group. elif problem_id == 30: _seed(problem_id, rank, trial) - r, in_f, out_f = 8, N, N - grad_fc1_1 = torch.randn((r, in_f), dtype=dtype, device=dev) - grad_fc1_2 = torch.randn((r, in_f), dtype=dtype, device=dev) - grad_fc2 = torch.randn((out_f, r), dtype=dtype, device=dev) - return (grad_fc1_1, grad_fc1_2, grad_fc2, None) + num_local_experts = 4 + tokens = max(64, M // 4) + r, d_in, d_out = 8, N, N + expert_inputs = torch.randn((num_local_experts, tokens, d_in), dtype=dtype, device=dev) + grad_fc1_1_out = torch.randn((num_local_experts, tokens, r), dtype=dtype, device=dev) + grad_fc1_2_out = torch.randn((num_local_experts, tokens, r), dtype=dtype, device=dev) + fc2_lora_acts = torch.randn((num_local_experts, tokens, r), dtype=dtype, device=dev) + grad_expert_outputs = torch.randn((num_local_experts, tokens, d_out), dtype=dtype, device=dev) + return (expert_inputs, grad_fc1_1_out, grad_fc1_2_out, fc2_lora_acts, grad_expert_outputs, None) # 31: fused_moe_fwd elif problem_id == 31: @@ -629,11 +648,13 @@ def create_input_tensor( x = torch.randn(base_shape, dtype=dtype, device=dev) return (x, 0, 1, None) - # 34: ulysses_all_gather_into_tensor_primitive + # 34: ulysses_all_gather_into_tensor_primitive — gather along dim 1 (the + # Ulysses sequence dim), which forces a gather + strided-write reorder + # rather than the plain dim-0 all_gather of problem 2. elif problem_id == 34: _seed(problem_id, rank, trial) x = torch.randn(base_shape, dtype=dtype, device=dev) - return (x, None) + return (x, 1, None) # 35: ulysses_all_gather_variable_primitive elif problem_id == 35: @@ -1012,26 +1033,28 @@ def _init_param(shape: tuple) -> torch.Tensor: None, ) - # 50: moe_ep_wide + # 50: moe_ep_wide — each rank hosts num_local_experts DISTINCT experts; + # weights are stacked (E_local, ...) tensors driving a ragged grouped GEMM. elif problem_id == 50: _seed(problem_id, rank, trial) num_experts = world_size * 2 + num_local_experts = 2 top_k = 2 hidden_dim = N inter_dim = 128 hidden_states = torch.randn((M, hidden_dim), dtype=dtype, device=dev) gate_weight = torch.randn((num_experts, hidden_dim), dtype=dtype, device=dev) gate_bias = torch.randn((num_experts,), dtype=dtype, device=dev) - gate_proj = _linear(hidden_dim, inter_dim, dtype, dev) - up_proj = _linear(hidden_dim, inter_dim, dtype, dev) - down_proj = _linear(inter_dim, hidden_dim, dtype, dev) + w_gate = torch.randn((num_local_experts, inter_dim, hidden_dim), dtype=dtype, device=dev) * 0.02 + w_up = torch.randn((num_local_experts, inter_dim, hidden_dim), dtype=dtype, device=dev) * 0.02 + w_down = torch.randn((num_local_experts, hidden_dim, inter_dim), dtype=dtype, device=dev) * 0.02 return ( hidden_states, gate_weight, gate_bias, - gate_proj, - up_proj, - down_proj, + w_gate, + w_up, + w_down, num_experts, top_k, None, @@ -1287,7 +1310,8 @@ def _init_param(shape: tuple) -> torch.Tensor: return (x, psi, weight, groups, nlon_out, nlon_in, azimuth_group, polar_group, bias) - # 63: deepmd_kalman_filter_optimizer + # 63: deepmd_kalman_filter_optimizer (fp32; uniform block structure so the + # weight exchange is a single fixed-size all-gather) elif problem_id == 63: _seed(problem_id, rank, trial) num_blocks = 8 @@ -1296,34 +1320,38 @@ def _init_param(shape: tuple) -> torch.Tensor: weights = [] P = [] for _ in range(num_blocks): - h = torch.randn((block, 1), dtype=torch.float64, device=dev) * 0.01 - w = torch.randn((block, 1), dtype=torch.float64, device=dev) - p = torch.eye(block, dtype=torch.float64, device=dev) + h = torch.randn((block, 1), dtype=torch.float32, device=dev) * 0.01 + w = torch.randn((block, 1), dtype=torch.float32, device=dev) + p = torch.eye(block, dtype=torch.float32, device=dev) H.append(h) weights.append(w) P.append(p) - error = torch.randn((), dtype=torch.float64, device=dev) + error = torch.randn((), dtype=torch.float32, device=dev) kalman_lambda = 0.98 kalman_nue = 0.9987 return (H, error, weights, P, kalman_lambda, kalman_nue) - # 64: gnn_neighbor_sampling + # 64: gnn_neighbor_sampling — irregular random graph (variable degrees, + # random neighbors) shared by all ranks; per-rank random seed frontier. elif problem_id == 64: _seed(problem_id, 0, trial) num_nodes = _round_up_multiple(max(10000, min(M, 100000)), world_size) - degree = 16 + avg_degree = 16 fanouts = [10, 5] node_to_rank = (torch.arange(num_nodes, device=dev, dtype=torch.long) % world_size).contiguous() - colptr = (torch.arange(num_nodes + 1, device=dev, dtype=torch.long) * degree).contiguous() - node_ids = torch.arange(num_nodes, device=dev, dtype=torch.long).unsqueeze(1) - offsets = torch.arange(1, degree + 1, device=dev, dtype=torch.long).unsqueeze(0) - row = ((node_ids + offsets) % num_nodes).reshape(-1).contiguous() + # Same RNG stream on every rank -> identical replicated graph. + deg = torch.randint(1, 2 * avg_degree, (num_nodes,), dtype=torch.long, device=dev) + colptr = torch.zeros(num_nodes + 1, dtype=torch.long, device=dev) + colptr[1:] = deg.cumsum(0) + row = torch.randint(0, num_nodes, (int(deg.sum().item()),), dtype=torch.long, device=dev) seeds_per_rank = max(512, min(N // max(world_size, 1), 2048)) - start = rank * seeds_per_rank - seed_nodes = (torch.arange(seeds_per_rank, device=dev, dtype=torch.long) + start) % num_nodes - return (seed_nodes.contiguous(), fanouts, colptr.contiguous(), row, node_to_rank, None, False) + # Distinct stream from the graph RNG (offset), varying per rank; + # randperm keeps the seed frontier duplicate-free. + torch.manual_seed(42 + problem_id * 1000 + 50000 + rank + trial * 1_000_003) + seed_nodes = torch.randperm(num_nodes, device=dev)[:seeds_per_rank].to(torch.long) + return (seed_nodes.contiguous(), fanouts, colptr.contiguous(), row.contiguous(), node_to_rank, None, False) # 65: gnn_feature_exchange_all2all elif problem_id == 65: @@ -1760,13 +1788,15 @@ def _init_param(shape: tuple) -> torch.Tensor: top_p = 0.9 return (logits.contiguous(), target.contiguous(), None, top_k, top_p) - # 83: vocab_parallel_log_prob_topk_chunked + # 83: vocab_parallel_log_prob_topk_chunked — long sequences and a larger + # vocab so the sequence-chunked pipeline (the point of this variant vs 82) + # actually processes many chunks. elif problem_id == 83: _seed(problem_id, rank, trial) - batch = max(1, min(M // 512, 4)) - seq_len = max(world_size, min(M // 128, 32)) + batch = max(1, min(M // 256, 4)) + seq_len = max(world_size, min(M // 4, 512)) seq_len = _round_up_multiple(seq_len, world_size) - vocab_size = 256 + vocab_size = 2048 local_vocab = vocab_size // world_size logits = torch.randn((batch, seq_len, local_vocab), dtype=dtype, device=dev) @@ -1791,13 +1821,14 @@ def _init_param(shape: tuple) -> torch.Tensor: chunk_size, ) - # 84: vocab_parallel_log_prob_topk_chunked_backward + # 84: vocab_parallel_log_prob_topk_chunked_backward — same long-sequence + # chunked regime as 83. elif problem_id == 84: _seed(problem_id, rank, trial) - batch = max(1, min(M // 512, 4)) - seq_len = max(world_size, min(M // 128, 32)) + batch = max(1, min(M // 256, 4)) + seq_len = max(world_size, min(M // 4, 512)) seq_len = _round_up_multiple(seq_len, world_size) - vocab_size = 256 + vocab_size = 2048 local_vocab = vocab_size // world_size logits = torch.randn((batch, seq_len, local_vocab), dtype=dtype, device=dev)