diff --git a/examples/deepseek-v4/gfx942/README.md b/examples/deepseek-v4/gfx942/README.md new file mode 100644 index 000000000..f9a76b991 --- /dev/null +++ b/examples/deepseek-v4/gfx942/README.md @@ -0,0 +1,260 @@ +# DeepSeek-V4 SFT on gfx942 (MI308X / CDNA3) + +Runs a 4-layer DeepSeek-V4 SFT at **128k context** exercising **all three** V4 attention +branches — dense+SWA, CSA, HCA — on a single 8×MI308X node. + +Upstream's V4 support targets gfx950 / CDNA4 and, as shipped, only ever ran at +`seq_length=4096`. Several things that are invisible at 4k become hard blockers at 128k. +This directory carries the fixes and a one-click launcher. + +```bash +# inside a rocm/primus:v26.5-pytorch2.12-te2.15 container, from the repo root +bash examples/deepseek-v4/gfx942/run_128k_dense_hca_csa.sh +``` + +The script is self-contained: it copies Megatron-LM out of the image (or fetches the +submodule), builds the SFT dataset on first run, and resolves every path relative to +itself. The only external requirement is a **local DeepSeek-V4 tokenizer directory** — +point `V4_TOKENIZER` at it, or leave it at `/apps/DeepSeek-V4-Flash`. Only +`tokenizer.json` / `tokenizer_config.json` are read. + +## Scope — read this before quoting any number + +* **Weights are randomly initialised.** No DeepSeek-V4 Megatron checkpoint exists, and + Megatron-Bridge registers importers only for V2-Lite/V2/V3. This is *SFT-shaped training + from scratch*, not fine-tuning. The tokenizer and the prompt/response loss mask are real. +* **4 layers, not 43.** `compress_ratios [0, 0, 4, 128]` was chosen to cover one layer of + each branch type. The released model is 3 dense / 21 CSA / 20 HCA. +* **1M does not fit on one node.** See "Long context" below for how far it gets and why. + +## Results + +8×MI308X, 4 layers, `[0, 0, 4, 128]`, `triton_v2`, 10 steps: + +| | TP=8 / CP=1 (first working recipe) | **CP=8 / TP=1 / EP=8 (current)** | +|---|---|---| +| peak memory | 188.63 GB (98.25%) | **42.30 GB (22%)** | +| step time | 9.4 s | **3.83 s** | +| loss | 11.935 → 11.777, monotone | 11.886 → 11.686, monotone | +| nan iterations | 0 | 0 | + +### Why CP, not TP + +This is the single most important finding here, and it is not the obvious choice. + +TP shards along the head axis. But the tensors that actually dominate at long context +**have no head axis**, so TP cannot touch them: + +* the indexer's `scores` is `[B, S, P]` — heads are already summed out +* V4's KV is a **single MQA latent** `[B, S, 1, head_dim]`, broadcast to H heads by a + stride-0 view +* MoE / MLP / residual activations scale with `S` alone +* the LM head's logits are `[S, B, vocab]` + +CP shards the sequence and therefore shards every one of them. Measured: 4.5× less memory +and 2.5× faster than the TP=8 recipe, same model, same math. + +Both can be 8 on 8 GPUs because the expert side decomposes as `ETP × EP × PP`, which does +not include CP (`parallel_state.py:794`), while the attention side needs +`TP × PP × CP ≤ world_size`. + +## What had to change + +### 1. P14 — shard attention heads across TP + +`deepseek_v4_layer_specs.py`, `deepseek_v4_attention.py`, `indexer.py`, +`deepseek_v4_transformer_config.py` + +V4 built `linear_q_up_proj` with `gather_output=True` and set +`num_attention_heads_per_partition = num_heads` (no `divide()`). TP therefore sharded +**weights only** — the `[B, S, H, head_dim]` query was replicated on every rank. The source +called head-sharded attention "tracked in P14"; it was never implemented. + +Enabled by `v4_shard_attention_heads: true` (default **false**). The grouped-O projection +makes this clean: it is block-diagonal over `o_groups`, and `linear_o_b` is already +row-parallel, so at TP=8 with `o_groups=8` each rank owns one group and the row-parallel +all-reduce sums them — verified in float64, residual 2.5e-15. + +Superseded in practice by the CP recipe above, but kept: it is correct, and it is what +makes TP>1 meaningful at all. + +### 2. Fused indexer scoring at the real head count + +`v4_attention_kernels/_triton_common/indexer_score.py`, `indexer_score_post.py` + +Both Triton indexer kernels gated on `_SUPPORTED_H = (1, 2, 4, 8, 16)`, while the released +config and Primus's own `deepseek_v4_base.yaml` set `index_n_heads=64`. **At the real width +these kernels were unreachable** and the indexer silently fell back to an eager einsum that +materialises `[B, S, H, P]`: 0.5 GiB at 4k, **512 GiB at 128k**. + +* `_SUPPORTED_H` extended to 32 and 64 (fwd matches eager to 3e-7, bwd to bf16 noise) +* `k_tile` hoisted out of the unrolled head loop — it never depended on `h`, so H=64 was + doing 63 redundant tile loads per block + +### 3. int64 offsets in the indexer kernels + +Every offset was int32. `s_offs * P` wraps negative once `S*P` exceeds 2³¹, which for CSA +(`P = S/4`) happens at `S ≈ 92682` — the store then lands out of bounds *silently*. +Bisected: 64k clean, 96k NaN, 128k NaN. 15 offsets promoted. + +### 4. A dead 16 GiB allocation in the CSA path + +The CSA branch built a dense `[S, S]` sliding-window mask and passed it to `_csa_forward`, +whose own docstring says it is *"retained in the signature for back-compat but unused"* — +the function `del`s it on entry. 16 GiB at 128k, allocated to be thrown away. + +### 5. gfx942 LDS budget + +`v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py` + +The sparse-MLA backward is tuned for gfx950's 160 KB LDS; gfx942 has 64 KB and the stock +staging asks for 73728 B, so the kernel fails to **compile**. `PRIMUS_DSA_BWD_NUM_STAGES=1` +disables Triton's LDS multi-buffering. Measured: this is the *only* kernel switch needed. + +Also added `PRIMUS_DSA_BWD_R_CHUNK`. The backward's rank-chunk width is hard-coded to 256, +tuned for short sequences where the per-chunk buffers are small and dq reload traffic +dominates. At long context that inverts: the buffers scale with `total_tokens * R_CHUNK`, +so at `S_local = 131072` the `interm` buffer alone is 36 GiB. + +### 6. Context parallelism for all three branches + +`deepseek_v4_cp.py` (new), `deepseek_v4_attention.py`, `indexer.py`, +`v4_sparse_mla_adapter.py`, `sft/forward_step.py` + +A CP rank needs the `d_window` post-RoPE KV rows left of its shard plus its global row +offset; the sparse-MLA adapter then validates the window against **global** positions while +indexing the local `[boundary ++ local]` buffer. + +Three bugs were found getting CSA and HCA correct, all of which produce a *correct forward* +and a wrong gradient: + +1. **The local sliding window was not exchanged.** CP was wired only into the dense branch, + but CSA and HCA each run a sliding window over raw tokens too, and that window straddles + the shard edge. Factored into `_cp_prepend_boundary`, used by all three. +2. **`_AllGatherPool.backward` sliced instead of reduce-scattering.** Rank r's pool rows are + read by the queries of every rank at or after r, so each holds a partial gradient; + taking only this rank's block dropped all the downstream contributions. Invisible in the + forward, compounding over steps (loss drift 2e-5 → 4.5e-4 over three steps). +3. **The indexer built its key pool from local hidden only.** `indexer_compressor(hidden)` + produced `P_local` while the attention pool was already `P_global`, so the top-K indices + named the wrong columns and a query could never select history from an earlier rank. + +Verified at 8k with a uniform dataset, TP=1, EP=1 (so weight init and data are identical +across runs). Spread across CP=1/2/4, against a measured CP noise floor of 5e-5: + +| config | CP=1 | CP=2 | CP=4 | spread | +|---|---|---|---|---| +| dense only (control) | 11.88296 | — | 11.88293 | 5e-5 | +| dense + HCA | 11.88916 | 11.88918 | 11.88914 | 5e-5 | +| dense + CSA + HCA | 11.86170 | 11.86179 | 11.86169 | 1e-4 | + +CSA sits slightly above the floor and does not grow with steps — top-K is a discrete +selection, so a bf16-level perturbation occasionally swaps the 512th and 513th column. + +### 7. Memory: measure, don't guess + +Every attempt to reason about where the memory went was wrong until it was profiled. A +`sitecustomize.py` probe (injected via `PYTHONPATH`, so no code change) recording allocation +stacks gave the actual composition of a 179.67 GiB peak, and the top entry was not in +attention at all: + +| GiB | site | | +|---:|---|---| +| 31.56 | `deepseek_v4_model.py` LM head logits | `S × vocab`, fixed below | +| 24.24 | `v4_sparse_mla_adapter.py` topk index matrix | int64, fixed below | +| 18.00 | `_rope_pad_q` | real | +| 16.00 | sparse-MLA fwd kernel | real | +| 16.00 | RMSNorm | real | +| 12.00 | `hc_expand` | real | + +Fixes that followed: + +* **Chunked linear + cross-entropy for V4.** Primus already ships this + (`patches/fused_linear_ce_patches.py`) but it hooks `GPTModel._postprocess`, and + `DeepseekV4Model` derives from `LanguageModule` — so it never applied. Wired in behind + `FUSED_LINEAR_CE=1`, restricted to TP=1 (the chunked path matmuls against the full output + weight, which is only equivalent when the vocab is not sharded). +* **int32 topk index matrix.** `torch.full((B,S,P), -1)` materialised an 8.6 GiB constant + just to be a `torch.where` else-branch, and every intermediate was int64 although + `_pad_topk_64` casts to int32 immediately. ~24 GiB at 1M. +* **HCA concatenated on the head-broadcast views.** `torch.cat` on a stride-0 expanded view + materialises the H-fold copy the broadcast exists to avoid — 8.51 GiB for K and another + 8.51 GiB for V, per HCA layer, of which the consumer reads 136 MiB (it takes `k_bh[:, 0]` + and never reads `v_bh` at all). +* **Output BSHD→BHSD→BSHD round trip.** The adapter returned a `.contiguous()` BHSD copy + and every caller immediately made a BSHD copy of it. Both are now views. +* **`dk = zeros(B, H, Skv, D)` in both backwards**, 63/64 of it zeros, allocated only to + match the broadcast view's shape. The adapter now takes the un-broadcast latent. + +### 8. SFT plumbing + +* `mock_data` is fatal under `stage: sft` — the mock-data patch force-installs + `NullTokenizer`, whose `text_to_ids` is `int(x)` per whitespace token. +* `train_data_path` must be non-null or the pretrain data-prep hook demands `HF_TOKEN`. +* `rope_type: rope` — the base preset says `yarn`, but common-attention asserts `rope`. +* `moe_router_enable_expert_bias: false` — expert bias requires `sigmoid` scoring, which + conflicts with V4's `sqrtsoftplus`. +* `create_attention_mask_in_dataloader: false` — otherwise a `[S, S]` bool tensor. +* Contiguous CP sharding of the batch and a CP-aware loss reduction. + +## A pattern worth naming + +`DeepseekV4Model` derives from `LanguageModule` rather than `GPTModel`; +`DeepseekV4TransformerBlock` deliberately bypasses `TransformerBlock.__init__`; +`DeepseekV4Attention` fully overrides `MultiLatentAttention.forward`. Each is defensible on +its own. The cumulative effect is that **whole classes of upstream optimisation silently do +nothing on V4**, with no error and no warning: + +| optimisation | upstream hook point | why it missed V4 | +|---|---|---| +| chunked linear+CE | `GPTModel._postprocess` | V4 derives from `LanguageModule` | +| TE CPU activation offload | `TransformerBlock.__init__` / `.forward` | V4 bypasses both | +| fine-grained activation offload | `GPTModel.forward`, `attention.py` | both bypassed | +| head-sharded TP | `Attention.__init__`'s `divide()` | V4 set the full head count | + +The first two are fixed here. When adding anything to V4, check whether the upstream +version of it is reachable before concluding it does not help. + +## Known limits + +* **`turbo` backend does not run on gfx942.** Its FlyDSL kernels emit `permlane16/32_swap` + and `mfma_f32_16x16x32_bf16`, both CDNA4-only. The permlane butterfly has a `ds_bpermute` + equivalent, but the MFMA tile shape does not — CDNA3's `mfma_f32_16x16x16bf16_1k` has half + the K depth, so 25 call sites plus an inline-asm block would need re-tiling. `gluon*` and + `flydsl_v1` hard-assert gfx950. **`triton_v2` is the fastest usable backend.** +* **turbo and the fused indexer want opposite TP.** turbo asserts `num_heads % 32 == 0` + (TP ≤ 2 at H=64) while the fused indexer needs H_local ≤ 16 (TP ≥ 4). +* **`FUSED_LINEAR_CE=1` requires `overlap_param_gather: false`.** The chunked backward + issues one `autograd.grad` per chunk, changing each parameter's backward-hook firing + count, which desyncs the distributed optimizer's overlapped all-gather. +* **Streaming indexer top-K (`PRIMUS_INDEXER_TOPK_CHUNK`) is not a general memory switch.** + It is bit-exact and costs nothing at chunk=8192, but it only helps when + `S_local × P_global` makes `scores` a dominant term. At 128k with CP=8 that tensor is + 1 GiB and chunking it saves exactly zero. + +## Long context + +| seq | S_local (CP=8) | result | +|---|---|---| +| 128k | 16384 | 42.30 GB, 3.83 s/step | +| 512k | 65536 | one step at 175.7 GB and **432 s/step**, then OOM | +| 1M | 131072 | OOM | + +**1M does not fit on one node, and narrowing the model does not fix it.** Measured: + +| narrowing | outcome | +|---|---| +| `num_attention_heads` 64 → 8 | OOM at 182.81 GiB — peak unchanged | +| `hc_mult` 4 → 2 | OOM at 180.69 GiB | +| `kv_channels` 512 → 256 | OOM at 182.84 GiB (steady state only 56 GB) | + +Steady-state memory drops a lot in each case, but the failure watermark does not move: the +binding terms are the ones that do not scale with model width. After the section-7 fixes the +watermark is 167.38 GiB with 5.63 GiB free, and what remains is genuine computation state, +not waste. Two nodes at CP=16 halves `S_local` and is the clean path. + +One open item: at `PRIMUS_DSA_BWD_R_CHUNK=64` — the only setting whose workspace fits — the +8-GPU run dies with `HSA_STATUS_ERROR_MEMORY_FAULT`. The kernel is *not* at fault: a +standalone harness at the identical shapes (both CSA `topk=640` and HCA `topk=8320`, with +30% `-1` sentinels, under 130 GiB of ballast) completes correctly, and an exhaustive audit +found no int32 overflow. The remaining difference is the distributed context. Unresolved. diff --git a/examples/deepseek-v4/gfx942/README_full_4k_multinode.md b/examples/deepseek-v4/gfx942/README_full_4k_multinode.md new file mode 100644 index 000000000..dbcc08d45 --- /dev/null +++ b/examples/deepseek-v4/gfx942/README_full_4k_multinode.md @@ -0,0 +1,146 @@ +# Full DeepSeek-V4-Flash 4k SFT — multi-node on gfx942 (MI308X / CDNA3) + +One-click launcher: [`run_full_4k_multinode.sh`](./run_full_4k_multinode.sh). + +This trains the **complete** DeepSeek-V4-Flash model — 43 decoder layers + 1 MTP, +256 experts (top-6), `compress_ratios` cycling dense / CSA(4) / HCA(128) — at 4k +sequence length, across 3 nodes (24× MI308X). It is the full model, not the 4-layer +cut used by the single-node 128k smoke. + +## TL;DR + +```bash +# In each node's container, from the Primus repo root. Only two vars differ per node. +# node 0 (master): +MASTER_ADDR= NODE_RANK=0 bash examples/deepseek-v4/gfx942/run_full_4k_multinode.sh +# node 1: +MASTER_ADDR= NODE_RANK=1 bash examples/deepseek-v4/gfx942/run_full_4k_multinode.sh +# node 2: +MASTER_ADDR= NODE_RANK=2 bash examples/deepseek-v4/gfx942/run_full_4k_multinode.sh +``` + +`MASTER_ADDR` is the master node's IP on the socket NIC (see Networking). Node +addresses are **not** hardcoded in the script. + +## Why 3 nodes + +The model does not fit on 1 or 2 nodes on this GPU, and the binding term is the +**expert optimizer state**: + +- Experts are sharded `ETP × EP × PP` ways. Their optimizer state (fp32 master + + Adam m/v) **cannot** be sharded across data parallelism, because at EP=8 the + expert data-parallel size is 1 (EP consumes the parallelism DP would use). +- On 2 nodes (16 GPUs) `EP × PP ≤ 16`, giving ~18 B experts/card and an + unshardable optimizer state that fits neither GPU nor host. +- On 3 nodes (24 GPUs) `EP × PP = 24` → ~12 B experts/card, and with the optimizer + offloaded to the host it fits. This is the minimum viable configuration. + +Parallelism used: + +| domain | layout | +|-----------|--------------------------------| +| attention | TP=1 · PP=3 · CP=1 → DP=8 | +| expert | ETP=1 · EP=8 · PP=3 (24-way) | +| PP layout | `Et*14｜t*14｜t*15mL` | + +## The two fixes that make it work + +Both were diagnosed from measurement (py-spy stacks + dmesg), not guessed. + +### 1. `NCCL_ALGO=Ring` — cross-node collective deadlock + +Symptom: the first step hangs with GPUs at 0%, stuck in a 1-int `all_reduce` on the +cross-node model-parallel group (`logical_and_across_model_parallel_group`). + +Cause: NCCL defaults a **small** `all_reduce` to the **Tree** algorithm, and +Tree-over-TCP-socket across nodes **deadlocks** on this fabric. + +Fix: force `NCCL_ALGO=Ring` (plus `RCCL_USE_AMD_SMI_LIB=1` for fabric-topology +probing). This matches a known-working 6-node run on the same fabric, which passed +all-reduce / all-gather / all-to-all over the management NIC in ~10 s. + +### 2. `optimizer_offload_fraction=0.75` — host OOM + +Symptom: a single rank is SIGKILLed by the OOM-killer (dmesg: `python invoked +oom-killer`, anon-rss ~160 GB); the surviving ranks then hang forever in the +collective above, waiting on the dead peer. (This is why fix #1's effect was +masked at first — the "deadlock" was really a dead peer.) + +Cause: the optimizer state is **151 GB/rank** (of which the expert part is 147 GB, +unshardable). At `fraction=1.0` that is ~1208 GB/node, and pinned-memory allocation +peaks push it to ~1570 GB — which overran the node that started with the least free +RAM. + +Fix (calculated, not a seesaw): `fraction=0.75` puts + +- **113 GB/rank on the host** → ~1178 GB/node peak, comfortably under a 3 TB host + (~460 GB headroom on all three nodes), and +- **38 GB/rank back on the GPU** → ~109 GB of 192 GB used (83 GB headroom). + +Both sides are safe, and it keeps `exp_avg_sq` in fp32 (the config's NaN guard) so +numerics are unchanged. + +## Networking + +RDMA on this fabric is unusable (PFC unconfigured → `IBV_WC_RETRY_EXC_ERR`, and an +ionic GDA hang), so NCCL runs over **TCP on the management NIC**. The script pins +the interface for all three transports — set `NCCL_SOCKET_IFNAME` if your NIC is not +`ens50f0`: + +```bash +NCCL_SOCKET_IFNAME= MASTER_ADDR= NODE_RANK= bash .../run_full_4k_multinode.sh +``` + +`GLOO_SOCKET_IFNAME` and `TP_SOCKET_IFNAME` default to the same NIC. All three are +required — miss one and rendezvous or the TP group binds the wrong interface and +hangs. + +## gfx942 kernel settings + +- `PRIMUS_DSA_BWD_NUM_STAGES=1` — the V4 sparse-MLA backward Triton kernel is tuned + for gfx950/CDNA4 (160 KB LDS); gfx942 has 64 KB, so LDS multi-buffering is turned + off or the kernel fails to compile. +- `PRIMUS_INDEXER_TRITON_FULL=1` — use the fused indexer path instead of an eager + einsum that materialises `[B,S,H,P]`. +- attention backend `triton_v2` for both the dense/HCA and CSA paths (the turbo + sparse-MLA backend is CDNA4-only). + +## Verified result + +10-step and 100-step runs both completed cleanly on 3× MI308X nodes over the socket +network, from random init: + +| metric | 10-step | 100-step | +|---------------------|----------------|----------------| +| lm loss | 11.90 → 11.21 | 11.90 → 9.19 | +| grad norm | 20.0 → 14.9 | 20.0 → 2.8 | +| nan iterations | 0 (all steps) | 0 (all steps) | +| GPU peak mem | ~150 GB / 192 | ~150 GB / 192 | +| exit | 0 | 0 | + +Loss decreases monotonically and grad norm is stable — no NaN, no backward errors. + +## Requirements + +- 3 nodes × 8 MI308X (gfx942), 192 GB/GPU, ~3 TB host RAM each. +- Container image `rocm/primus:v26.5-pytorch2.12-te2.15` (or equivalent), started + with `--network host --privileged --device /dev/kfd --device /dev/dri`, `/apps` + (or wherever the repo lives) mounted. +- A local DeepSeek-V4 tokenizer directory (default `/apps/DeepSeek-V4-Flash`); only + `tokenizer.json` / `tokenizer_config.json` are read. Weights are **not** loaded — + this trains from random init (no V4 Megatron checkpoint exists). Nothing is saved. + +## Overridable knobs + +| var | default | meaning | +|-------------------------------------|------------------|----------------------------------| +| `MASTER_ADDR` | (required) | master node IP on the socket NIC | +| `NODE_RANK` | (required) | this node's rank, 0..NNODES-1 | +| `NNODES` | 3 | node count | +| `MASTER_PORT` | 29710 | rendezvous port | +| `NCCL_SOCKET_IFNAME` | ens50f0 | socket NIC | +| `V4_TOKENIZER` | /apps/DeepSeek-V4-Flash | tokenizer dir | +| `TRAIN_ITERS` | 10 | training steps | +| `GBS` | 24 | global batch size | +| `PRIMUS_LR` | 1.0e-6 | learning rate | +| `PRIMUS_OPTIMIZER_OFFLOAD_FRACTION` | 0.75 | optimizer host-offload fraction | diff --git a/examples/deepseek-v4/gfx942/prepare_sft_data.py b/examples/deepseek-v4/gfx942/prepare_sft_data.py new file mode 100644 index 000000000..1748a6ad2 --- /dev/null +++ b/examples/deepseek-v4/gfx942/prepare_sft_data.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python +"""Build the long-form SFT jsonl used by the gfx942 DeepSeek-V4 recipes. + +Non-packed SFT pads every sample to `seq_length` and truncates above it, so rows are built +deliberately LONGER than the target: otherwise a short Alpaca row at 128k would be 99.9% +padding with loss_mask=0, i.e. a benchmark of padding rather than of attention. + +Rows are alpaca-format {"instruction", "input", "output"}; the formatter masks the prompt +and supervises the response, so putting the bulk of the text in "output" yields a genuine +supervised span. + +Usage: python prepare_sft_data.py --tokenizer --out-dir [--lengths 4096 131072] +""" + +import argparse +import json +import os +import sys + + +def load_corpus(tokenizer_dir): + """Natural-language text. Falls back to local docs if the Hub is unreachable.""" + try: + from datasets import load_dataset + + src = load_dataset("tatsu-lab/alpaca", split="train") + blob = "".join(f"{r['instruction']} {r['input']} {r['output']}\n" for r in src) + print(f"[data] alpaca: {len(src)} rows, {len(blob)} chars", flush=True) + return blob + except Exception as e: # noqa: BLE001 + print(f"[data] alpaca unavailable ({type(e).__name__}); falling back to repo docs", flush=True) + here = os.path.dirname(os.path.abspath(__file__)) + root = os.path.abspath(os.path.join(here, "..", "..", "..")) + parts = [] + for base, _, files in os.walk(os.path.join(root, "docs")): + for fn in files: + if fn.endswith((".md", ".py", ".txt")): + try: + parts.append(open(os.path.join(base, fn), encoding="utf-8", errors="ignore").read()) + except OSError: + pass + return "\n".join(parts) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--tokenizer", required=True, help="local HF tokenizer directory") + ap.add_argument("--out-dir", required=True) + ap.add_argument("--lengths", type=int, nargs="+", default=[4096, 32768, 131072]) + ap.add_argument("--rows", type=int, default=16) + args = ap.parse_args() + + from transformers import AutoTokenizer + + os.makedirs(args.out_dir, exist_ok=True) + tok = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + + blob = load_corpus(args.tokenizer) + if len(blob) < 100_000: + sys.exit(f"[data] corpus too small: {len(blob)} chars") + + # Calibrate chars-per-token once rather than tokenizing gigabytes. + sample = blob[:400_000] + cpt = len(sample) / len(tok(sample, add_special_tokens=False)["input_ids"]) + print(f"[data] chars/token = {cpt:.3f}", flush=True) + + cursor = 0 + + def take(n_chars, cur): + out, need = [], n_chars + while need > 0: + end = min(cur + need, len(blob)) + out.append(blob[cur:end]) + need -= end - cur + cur = 0 if end >= len(blob) else end + return "".join(out), cur + + for target in args.lengths: + name = f"sft_{target}.jsonl" + path = os.path.join(args.out_dir, name) + if os.path.exists(path): + print(f"[data] {name} exists, skipping", flush=True) + continue + n_chars = int(target * 1.35 * cpt) # headroom so truncation, not padding, sets the length + with open(path, "w", encoding="utf-8") as f: + for _ in range(args.rows): + text, cursor = take(n_chars, cursor) + f.write( + json.dumps( + {"instruction": "Continue the following document.", "input": "", "output": text}, + ensure_ascii=False, + ) + + "\n" + ) + first = json.loads(open(path, encoding="utf-8").readline()) + ntok = len(tok(first["output"], add_special_tokens=False)["input_ids"]) + ok = "OK" if ntok >= target else "SHORT!" + print( + f"[data] [{ok}] {name}: {args.rows} rows, " + f"{os.path.getsize(path)/2**20:.1f} MB, row0={ntok} tokens (target {target})", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/deepseek-v4/gfx942/run_128k_dense_hca_csa.sh b/examples/deepseek-v4/gfx942/run_128k_dense_hca_csa.sh new file mode 100755 index 000000000..c79dda50e --- /dev/null +++ b/examples/deepseek-v4/gfx942/run_128k_dense_hca_csa.sh @@ -0,0 +1,142 @@ +#!/bin/bash +# One-click 128k DeepSeek-V4 SFT on a single 8x MI308X (gfx942 / CDNA3) node. +# +# Model: 4 decoder layers with compress_ratios [0, 0, 4, 128] -- so the run exercises all +# three V4 attention branches (dense+SWA, CSA, HCA), not just the cheap dense one. +# +# Parallelism is CP=8 / TP=1 / EP=8. That is NOT the obvious choice and it matters a lot: +# the same recipe at TP=8 / CP=1 peaks at 188.63 GB and 9.4 s/step, this one at 42.30 GB +# and 3.83 s/step. The reason is that the dominant tensors do not have a head axis, so TP +# cannot shard them -- the indexer's `scores` is [B, S, P] (heads are already summed out), +# V4's KV is a single MQA latent, and the MoE / residual activations scale with S alone. +# CP shards the sequence and therefore shards all of them. +# +# Everything is resolved relative to this script, so the repo can live anywhere. Run it +# from inside a rocm/primus:v26.5-pytorch2.12-te2.15 container: +# +# bash examples/deepseek-v4/gfx942/run_128k_dense_hca_csa.sh +# +# Required: a local DeepSeek-V4 tokenizer directory. Point V4_TOKENIZER at it, or drop it +# at /apps/DeepSeek-V4-Flash. Only tokenizer.json / tokenizer_config.json are read -- the +# weights are NOT loaded (no V4 Megatron checkpoint exists; this trains from random init). +# +# Why each gfx942-specific setting is needed is documented in +# examples/deepseek-v4/gfx942/README.md. +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "${HERE}/../../.." && pwd)" +cd "${REPO}" + +# ---- prerequisites ----------------------------------------------------------------- +V4_TOKENIZER="${V4_TOKENIZER:-/apps/DeepSeek-V4-Flash}" +DATA_DIR="${DATA_DIR:-${REPO}/data/sft}" +BACKEND_PATH="${BACKEND_PATH:-${REPO}/third_party/Megatron-LM}" + +if [ ! -f "${V4_TOKENIZER}/tokenizer.json" ]; then + echo "[error] no tokenizer at ${V4_TOKENIZER}. Set V4_TOKENIZER= (needs tokenizer.json)." >&2 + exit 1 +fi + +# Megatron-LM: the repo pins it as a submodule, but the container image already ships the +# exact pinned commit, so reuse that instead of a network fetch when it is available. +if [ ! -f "${BACKEND_PATH}/megatron/training/arguments.py" ]; then + if [ -f /workspace/Primus/third_party/Megatron-LM/megatron/training/arguments.py ]; then + echo "[setup] copying Megatron-LM from the image" + mkdir -p "${REPO}/third_party" + cp -a /workspace/Primus/third_party/Megatron-LM "${REPO}/third_party/" + rm -f "${BACKEND_PATH}/.git" + else + echo "[setup] fetching Megatron-LM submodule" + git -C "${REPO}" submodule update --init third_party/Megatron-LM + fi +fi + +if [ ! -f "${DATA_DIR}/sft_131072.jsonl" ]; then + echo "[setup] building SFT data (one-off, a few minutes)" + python "${HERE}/prepare_sft_data.py" --tokenizer "${V4_TOKENIZER}" \ + --out-dir "${DATA_DIR}" --lengths 131072 --rows 16 +fi + +# ---- launcher ---------------------------------------------------------------------- +export PRIMUS_LAUNCHER=direct +unset SLURM_JOB_ID SLURM_JOBID SLURM_NODELIST 2>/dev/null || true +export NNODES=1 NODE_RANK=0 MASTER_ADDR=localhost +export MASTER_PORT="${MASTER_PORT:-29517}" +export GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +export BACKEND_PATH +export PRIMUS_OUTPUT_ROOT="${PRIMUS_OUTPUT_ROOT:-${REPO}/output}" +export PRIMUS_TEAM="${PRIMUS_TEAM:-amd}" +export PRIMUS_USER="${PRIMUS_USER:-$(whoami)}" +export LOCAL_RANKS="${LOCAL_RANKS:---local-ranks-filter 0,1,2,3,4,5,6,7}" +export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" +# Let an OOM surface as a clean HIP out-of-memory instead of an opaque +# HSA_STATUS_ERROR_EXCEPTION + GPU coredump. +export HSA_NO_SCRATCH_RECLAIM="${HSA_NO_SCRATCH_RECLAIM:-0}" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-lo}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-lo}" + +# ---- gfx942 (CDNA3) kernel settings ------------------------------------------------- +# The V4 sparse-MLA backward Triton kernel is tuned for gfx950/CDNA4 (160 KB LDS). gfx942 +# has 64 KB and the stock pipeline staging asks for 73728 B, so it fails to COMPILE. +# num_stages=1 turns off Triton's LDS multi-buffering. Measured: this is the only kernel +# switch needed -- PRIMUS_HC_TRITON and PRIMUS_DSA_DKV_SAFE can stay at their defaults. +export PRIMUS_DSA_BWD_NUM_STAGES="${PRIMUS_DSA_BWD_NUM_STAGES:-1}" +# Fused indexer scoring: without it the CSA indexer falls back to an eager einsum that +# materialises [B, S, H, P] -- 512 GiB at 128k. +export PRIMUS_INDEXER_TRITON_FULL="${PRIMUS_INDEXER_TRITON_FULL:-1}" + +# ---- memory --------------------------------------------------------------------------- +# Chunked linear + cross-entropy. The LM head's logits are [S, B, vocab]; at 128k with +# vocab 129280 that is 31.6 GiB before the loss even upcasts, and it was measured as the +# single largest allocation in the step -- larger than any attention tensor. Off by +# default upstream; this recipe wants it. +export FUSED_LINEAR_CE="${FUSED_LINEAR_CE:-1}" +export FUSED_CE_CHUNK="${FUSED_CE_CHUNK:-4096}" +# REQUIRED with FUSED_LINEAR_CE: the chunked backward issues one autograd.grad per chunk, +# which changes each parameter's backward-hook firing count and desyncs the distributed +# optimizer's overlapped parameter all-gather. DeepseekV4Model raises if you forget. +export PRIMUS_OVERLAP_PARAM_GATHER="${PRIMUS_OVERLAP_PARAM_GATHER:-false}" + +# ---- model / parallelism ------------------------------------------------------------ +export PRIMUS_SEQ_LENGTH=131072 +export PRIMUS_MAX_POSITION_EMBEDDINGS=131072 +export SFT_JSONL="${DATA_DIR}/sft_131072.jsonl" +export PRIMUS_TOTAL_LAYERS=4 +export PRIMUS_COMPRESS_RATIOS="[0, 0, 4, 128, 0]" # dense, dense, CSA, HCA (+MTP slot) +export PRIMUS_RECOMPUTE_LAYERS=4 +export PRIMUS_NUM_EXPERTS="${PRIMUS_NUM_EXPERTS:-8}" +export PRIMUS_MOE_TOPK="${PRIMUS_MOE_TOPK:-1}" +# CP=8 shards the sequence; EP=8 shards the experts over the same 8 ranks (the expert side +# decomposes as ETP*EP*PP, which does not include CP, so both can be 8 on 8 GPUs). +# P14 head sharding is off because TP=1 -- see the header for why TP is the wrong lever. +export PRIMUS_SHARD_HEADS=false +export PRIMUS_TP=1 +export PRIMUS_ETP=1 +export PRIMUS_EP=8 +export PRIMUS_CP=8 +export MBS=1 +export GBS=1 +# Random init at 128k diverges at the default 1e-5 (NaN on step 2); 1e-6 is stable. +export PRIMUS_LR="${PRIMUS_LR:-1.0e-6}" +export TRAIN_ITERS="${TRAIN_ITERS:-10}" +export V4_TOKENIZER +export PRIMUS_EXP_NAME="${PRIMUS_EXP_NAME:-dsv4_4layer_128k_cp8}" + +EXP="${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash_4layer-BF16-sft.yaml}" +LOGDIR="${PRIMUS_OUTPUT_ROOT}/${PRIMUS_TEAM}/${PRIMUS_USER}/${PRIMUS_EXP_NAME}" +mkdir -p "${LOGDIR}" + +echo "[run] repo=${REPO}" +echo "[run] seq=${PRIMUS_SEQ_LENGTH} layers=${PRIMUS_TOTAL_LAYERS} ratios=${PRIMUS_COMPRESS_RATIOS}" +echo "[run] CP=${PRIMUS_CP} TP=${PRIMUS_TP} EP=${PRIMUS_EP} iters=${TRAIN_ITERS}" +echo "[run] expect ~42 GB peak, ~3.8 s/step" +echo "[run] log=${LOGDIR}/log_node0.txt" + +./primus-cli direct -- train pretrain --config "${EXP}" \ + --manual_gc True \ + --manual_gc_interval 100 \ + --pp_warmup False --sequence_parallel False \ + --log_avg_skip_iterations 3 \ + --backend_path "${BACKEND_PATH}" \ + 2>&1 | tee "${LOGDIR}/log_node0.txt" diff --git a/examples/deepseek-v4/gfx942/run_full_4k_multinode.sh b/examples/deepseek-v4/gfx942/run_full_4k_multinode.sh new file mode 100755 index 000000000..e52d75e86 --- /dev/null +++ b/examples/deepseek-v4/gfx942/run_full_4k_multinode.sh @@ -0,0 +1,178 @@ +#!/bin/bash +# One-click FULL DeepSeek-V4-Flash SFT at 4k, multi-node on gfx942 / CDNA3 (MI308X, 192 GB). +# +# This trains the COMPLETE model -- 43 decoder layers + 1 MTP, 256 experts (top-6), +# compress_ratios cycling dense / CSA(4) / HCA(128) -- not the 4-layer cut. On this GPU +# it does not fit on one node: the expert optimizer state is unshardable across data +# parallelism (EP consumes the parallelism DP would use), so a single node cannot hold +# the whole model + optimizer. Three nodes (24 GPUs) is the minimum that does. +# +# attention domain: TP=1 * PP=3 * CP=1 -> DP = 24/3 = 8 +# expert domain: ETP=1 * EP=8 * PP=3 -> experts sharded 24-way +# optimizer: CPU-offloaded at fraction 0.75 (see the memory note below) +# +# NODE ADDRESSES ARE NOT HARDCODED. Each node exports MASTER_ADDR + NODE_RANK, then runs +# this same script. Example on a 3-node set (run in each node's container): +# +# # node 0 (also the master): +# MASTER_ADDR= NODE_RANK=0 bash examples/deepseek-v4/gfx942/run_full_4k_multinode.sh +# # node 1: +# MASTER_ADDR= NODE_RANK=1 bash examples/deepseek-v4/gfx942/run_full_4k_multinode.sh +# # node 2: +# MASTER_ADDR= NODE_RANK=2 bash examples/deepseek-v4/gfx942/run_full_4k_multinode.sh +# +# Overridable knobs (all have sane defaults): +# MASTER_ADDR master node IP/host (REQUIRED, no default) +# NODE_RANK this node's rank 0..N-1 (REQUIRED, no default) +# NNODES node count (default 3) +# MASTER_PORT rendezvous port (default 29710) +# NCCL_SOCKET_IFNAME socket NIC (default ens50f0 -- the management NIC) +# V4_TOKENIZER local tokenizer dir (default /apps/DeepSeek-V4-Flash) +# TRAIN_ITERS steps (default 10) +# PRIMUS_OPTIMIZER_OFFLOAD_FRACTION (default 0.75) +# +# Why each gfx942 / socket / memory setting is needed is documented in +# examples/deepseek-v4/gfx942/README_full_4k_multinode.md. +set -uo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "${HERE}/../../.." && pwd)" +cd "${REPO}" || exit 1 + +# ---- required inputs --------------------------------------------------------------- +if [ -z "${MASTER_ADDR:-}" ]; then + echo "[error] MASTER_ADDR is required (the master node's IP/host). e.g. MASTER_ADDR=10.0.0.1 NODE_RANK=0 bash $0" >&2 + exit 1 +fi +if [ -z "${NODE_RANK:-}" ]; then + echo "[error] NODE_RANK is required (this node's rank, 0..NNODES-1)." >&2 + exit 1 +fi + +# ---- prerequisites ----------------------------------------------------------------- +V4_TOKENIZER="${V4_TOKENIZER:-/apps/DeepSeek-V4-Flash}" +DATA_DIR="${DATA_DIR:-${REPO}/data/sft}" +BACKEND_PATH="${BACKEND_PATH:-${REPO}/third_party/Megatron-LM}" + +if [ ! -f "${V4_TOKENIZER}/tokenizer.json" ]; then + echo "[error] no tokenizer at ${V4_TOKENIZER}. Set V4_TOKENIZER= (needs tokenizer.json)." >&2 + exit 1 +fi + +# Megatron-LM: the container image already ships the exact pinned commit; reuse it +# instead of a network fetch when the submodule is not checked out. +if [ ! -f "${BACKEND_PATH}/megatron/training/arguments.py" ]; then + if [ -f /workspace/Primus/third_party/Megatron-LM/megatron/training/arguments.py ]; then + echo "[setup] copying Megatron-LM from the image" + mkdir -p "${REPO}/third_party" + cp -a /workspace/Primus/third_party/Megatron-LM "${REPO}/third_party/" + rm -f "${BACKEND_PATH}/.git" + else + echo "[setup] fetching Megatron-LM submodule" + git -C "${REPO}" submodule update --init third_party/Megatron-LM + fi +fi + +SFT_JSONL="${SFT_JSONL:-${DATA_DIR}/sft_4096.jsonl}" +if [ ! -f "${SFT_JSONL}" ]; then + echo "[setup] building 4k SFT data (one-off)" + python "${HERE}/prepare_sft_data.py" --tokenizer "${V4_TOKENIZER}" \ + --out-dir "${DATA_DIR}" --lengths 4096 --rows 64 +fi + +# ---- launcher / rendezvous --------------------------------------------------------- +export PRIMUS_LAUNCHER=direct +unset SLURM_JOB_ID SLURM_JOBID SLURM_NODELIST 2>/dev/null || true +export NNODES="${NNODES:-3}" +export NODE_RANK MASTER_ADDR +export MASTER_PORT="${MASTER_PORT:-29710}" +export GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +export BACKEND_PATH +export PRIMUS_OUTPUT_ROOT="${PRIMUS_OUTPUT_ROOT:-${REPO}/output}" +export PRIMUS_TEAM="${PRIMUS_TEAM:-amd}" +export PRIMUS_USER="${PRIMUS_USER:-$(whoami)}" + +# ---- networking: socket over the management NIC ------------------------------------ +# RDMA on this fabric is unusable (PFC unconfigured -> IBV_WC_RETRY_EXC_ERR, ionic GDA +# hang), so NCCL runs over TCP on the management NIC. Pin the interface for NCCL, Gloo +# AND torch.distributed's TP-group sockets -- all three, or rendezvous / TP groups bind +# the wrong interface and hang. +export USING_AINIC="${USING_AINIC:-0}" +export NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-1}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-ens50f0}" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${NCCL_SOCKET_IFNAME}}" +export TP_SOCKET_IFNAME="${TP_SOCKET_IFNAME:-${NCCL_SOCKET_IFNAME}}" +export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" +export HSA_NO_SCRATCH_RECLAIM="${HSA_NO_SCRATCH_RECLAIM:-0}" + +# ROOT-CAUSE FIX for the first-step hang at logical_and_across_model_parallel_group (a +# 1-int all_reduce on the cross-node model-parallel group). NCCL defaults a small +# all_reduce to the Tree algorithm, and Tree-over-TCP-socket across nodes DEADLOCKS on +# this fabric. Forcing Ring fixes it (proven on a working 6-node run on the same fabric: +# all-reduce/all-gather/all-to-all all passed over ens50f0). RCCL_USE_AMD_SMI_LIB=1 is +# the fabric-topology probe that run needs too. +export NCCL_ALGO="${NCCL_ALGO:-Ring}" +export RCCL_USE_AMD_SMI_LIB="${RCCL_USE_AMD_SMI_LIB:-1}" +export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" +export TORCH_NCCL_ASYNC_ERROR_HANDLING="${TORCH_NCCL_ASYNC_ERROR_HANDLING:-1}" +export TORCH_NCCL_DUMP_ON_TIMEOUT="${TORCH_NCCL_DUMP_ON_TIMEOUT:-1}" + +# ---- gfx942 (CDNA3) kernel settings ------------------------------------------------ +# V4 sparse-MLA backward Triton kernel is gfx950-tuned (needs 160 KB LDS); gfx942 has +# 64 KB, so num_stages=1 disables LDS multi-buffering to let it compile. The full indexer +# Triton path avoids an eager einsum that materialises [B,S,H,P]. +export PRIMUS_DSA_BWD_NUM_STAGES="${PRIMUS_DSA_BWD_NUM_STAGES:-1}" +export PRIMUS_INDEXER_TRITON_FULL="${PRIMUS_INDEXER_TRITON_FULL:-1}" + +# ---- parallelism: full model over 3 nodes ------------------------------------------ +# The flash launcher's NNODES=3 branch sets PP=3/EP=8 and the layout Et*14|t*14|t*15mL. +export PRIMUS_TP="${PRIMUS_TP:-1}" +export PRIMUS_PP="${PRIMUS_PP:-3}" +export PRIMUS_EP="${PRIMUS_EP:-8}" +export PRIMUS_CP="${PRIMUS_CP:-1}" + +# ---- optimizer CPU offload (the lever that makes the full model fit) ---------------- +# The expert optimizer state is 147 GB/rank and cannot be sharded across DP (expert_dp=1), +# so it must leave the GPU. fraction=0.75 is CALCULATED, not guessed: it puts 113 GB/rank +# on the host (~1178 GB/node peak incl. pinned-memory overhead -> safe on a 3 TB host) and +# 38 GB/rank back on the GPU (~109 GB of 192 -> 83 GB headroom). fraction=1.0 overran the +# host and got a rank OOM-killed; lower fractions push the GPU to OOM. 0.75 is the middle. +export PRIMUS_OPTIMIZER_CPU_OFFLOAD="${PRIMUS_OPTIMIZER_CPU_OFFLOAD:-true}" +export PRIMUS_OPTIMIZER_OFFLOAD_FRACTION="${PRIMUS_OPTIMIZER_OFFLOAD_FRACTION:-0.75}" + +# ---- model / data (full 43L + MTP, 256 experts, 4k) -------------------------------- +export EXP="${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash_4layer-BF16-sft.yaml}" +export PRIMUS_SEQ_LENGTH="${PRIMUS_SEQ_LENGTH:-4096}" +export PRIMUS_MAX_POSITION_EMBEDDINGS="${PRIMUS_MAX_POSITION_EMBEDDINGS:-4096}" +export SFT_JSONL +export V4_TOKENIZER +export PRIMUS_LR="${PRIMUS_LR:-1.0e-6}" +export MBS="${MBS:-1}" +export GBS="${GBS:-24}" +export TRAIN_ITERS="${TRAIN_ITERS:-10}" +export PRIMUS_EXP_NAME="${PRIMUS_EXP_NAME:-dsv4_flash_full_4k_${NNODES}node}" + +# ---- SFT plumbing (full-model SFT vs pretrain defaults) ---------------------------- +# pp_warmup trips SFT forward_step; mock_data installs a NullTokenizer fatal to real SFT; +# expert bias needs sigmoid scoring which conflicts with V4's sqrtsoftplus. +export PP_WARMUP="${PP_WARMUP:-False}" +export MOCK_DATA="${MOCK_DATA:-False}" +export PRIMUS_MOE_ENABLE_EXPERT_BIAS="${PRIMUS_MOE_ENABLE_EXPERT_BIAS:-False}" + +# ---- attention backend: triton_v2 (turbo won't compile on gfx942) ------------------ +export USE_V4_ATTENTION_BACKEND="${USE_V4_ATTENTION_BACKEND:-triton_v2}" +export USE_V4_CSA_ATTENTION_BACKEND="${USE_V4_CSA_ATTENTION_BACKEND:-triton_v2}" +export USE_TURBO_ATTENTION="${USE_TURBO_ATTENTION:-False}" + +LOGDIR="${PRIMUS_OUTPUT_ROOT}/${PRIMUS_TEAM}/${PRIMUS_USER}/${PRIMUS_EXP_NAME}" +mkdir -p "${LOGDIR}" +LOG="${LOGDIR}/log_node${NODE_RANK}.txt" + +echo "[run] repo=${REPO}" +echo "[run] FULL model: 43L+MTP, 256 experts, seq=${PRIMUS_SEQ_LENGTH}" +echo "[run] nodes=${NNODES} rank=${NODE_RANK} master=${MASTER_ADDR}:${MASTER_PORT} nic=${NCCL_SOCKET_IFNAME}" +echo "[run] TP=${PRIMUS_TP} PP=${PRIMUS_PP} EP=${PRIMUS_EP} CP=${PRIMUS_CP} offload_frac=${PRIMUS_OPTIMIZER_OFFLOAD_FRACTION}" +echo "[run] iters=${TRAIN_ITERS} gbs=${GBS} lr=${PRIMUS_LR}" +echo "[run] log=${LOG}" + +exec bash examples/deepseek-v4/run_deepseek_v4_flash.sh > "${LOG}" 2>&1 diff --git a/examples/deepseek-v4/run_deepseek_v4.sh b/examples/deepseek-v4/run_deepseek_v4.sh index 0f878d2fc..e822099cd 100755 --- a/examples/deepseek-v4/run_deepseek_v4.sh +++ b/examples/deepseek-v4/run_deepseek_v4.sh @@ -340,7 +340,7 @@ fi --v4_grouped_experts_support_clamped_swiglu "$PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU" \ --compress_ratios "$PRIMUS_COMPRESS_RATIOS" \ --mtp_num_layers "${MTP_NUM_LAYERS:-0}" \ - --mock_data True \ + --mock_data "${MOCK_DATA:-True}" \ --enable_primus_turbo "$ENABLE_PRIMUS_TURBO" \ --use_turbo_attention "$USE_TURBO_ATTENTION" \ --use_v4_attention_backend "$USE_V4_ATTENTION_BACKEND" \ diff --git a/examples/deepseek-v4/run_deepseek_v4_flash.sh b/examples/deepseek-v4/run_deepseek_v4_flash.sh index 8b4be6392..7715f135a 100755 --- a/examples/deepseek-v4/run_deepseek_v4_flash.sh +++ b/examples/deepseek-v4/run_deepseek_v4_flash.sh @@ -85,6 +85,48 @@ elif [ "$NNODES" -eq 4 ]; then else export PRIMUS_PP_LAYOUT='Et*10|t*11|t*11|t*11L' fi +elif [ "$NNODES" -eq 3 ]; then + # 3 nodes = 24 GPUs. PP=3/EP=8: experts sharded EP*PP=24 ways -> 12B experts/card, + # optimizer 171 GB/card (too big for GPU) -> offload to host (CPU side ~1.7 TB/node, + # comfortably under 3 TB, unlike 2-node's 2.6 TB which OOM'd). 43 decoder layers + MTP + # across 3 PP stages. + export PRIMUS_TP=${PRIMUS_TP:-1} + export PRIMUS_PP=${PRIMUS_PP:-3} + export PRIMUS_EP=${PRIMUS_EP:-8} + export PRIMUS_RECOMPUTE_LAYERS=${PRIMUS_RECOMPUTE_LAYERS:-43} + if [ -z "${PRIMUS_PP_LAYOUT:-}" ]; then + if [ "$MTP_NUM_LAYERS" -eq 1 ]; then + export PRIMUS_PP_LAYOUT='Et*14|t*14|t*15mL' + else + export PRIMUS_PP_LAYOUT='Et*14|t*14|t*15L' + fi + fi +elif [ "$NNODES" -eq 2 ]; then + # Single-pair 2-node (16 GPUs). Params (42.8B/card at EP=8/PP=1, measured) do not fit + # on one card, and CP does not shard params -- only PP (layers) and EP (experts) do. + # PP=2 halves per-card params to ~21B; EP=8 shards the 256 experts. Full recompute + # keeps activations small. 43 decoder layers + 1 MTP split across 2 PP stages. + export PRIMUS_TP=${PRIMUS_TP:-1} + export PRIMUS_PP=${PRIMUS_PP:-2} + export PRIMUS_EP=${PRIMUS_EP:-8} + export PRIMUS_RECOMPUTE_LAYERS=${PRIMUS_RECOMPUTE_LAYERS:-43} + # Layout follows PRIMUS_PP: PP=4 on 16 GPUs shards experts 32-way (EP*PP), same as the + # 4-node config, so the full model fits in bf16. Honor a caller-provided layout. + if [ -z "${PRIMUS_PP_LAYOUT:-}" ]; then + if [ "${PRIMUS_PP}" -eq 4 ]; then + if [ "$MTP_NUM_LAYERS" -eq 1 ]; then + export PRIMUS_PP_LAYOUT='Et*10|t*11|t*11|t*11mL' + else + export PRIMUS_PP_LAYOUT='Et*10|t*11|t*11|t*11L' + fi + else + if [ "$MTP_NUM_LAYERS" -eq 1 ]; then + export PRIMUS_PP_LAYOUT='Et*21|t*22mL' + else + export PRIMUS_PP_LAYOUT='Et*21|t*22L' + fi + fi + fi fi export MBS=${MBS:-1} diff --git a/examples/deepseek-v4/run_v4_4layer_sft.sh b/examples/deepseek-v4/run_v4_4layer_sft.sh new file mode 100755 index 000000000..e02d765d8 --- /dev/null +++ b/examples/deepseek-v4/run_v4_4layer_sft.sh @@ -0,0 +1,109 @@ +#!/bin/bash +# Single-node (8x MI308X / gfx942) driver for 4-layer DeepSeek-V4 native SFT. +# +# Why not run_deepseek_v4_flash.sh directly: that script targets multi-node SLURM +# and hardcodes `--mock_data True` plus `--moe_router_force_load_balancing True` +# as literal CLI args. mock_data is fatal under stage:sft (it force-installs +# NullTokenizer), and force-LB freezes the router. Everything else it sets is +# reproduced here or in the experiment yaml. +# +# Usage (run from anywhere; paths resolve relative to the repo): +# SEQ=4096 SFT_JSONL=/sft_4k.jsonl bash run_v4_4layer_sft.sh +# SEQ=131072 SFT_JSONL=/sft_128k.jsonl bash run_v4_4layer_sft.sh +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "${HERE}/../.." && pwd)" +cd "${REPO}" || exit 1 + +# ---- launcher: single node, already inside the container ---- +export PRIMUS_LAUNCHER=direct +unset SLURM_JOB_ID SLURM_JOBID SLURM_NODELIST 2>/dev/null || true +export NNODES=1 +export NODE_RANK=0 +export MASTER_ADDR=localhost +export MASTER_PORT=${MASTER_PORT:-29517} +export GPUS_PER_NODE=${GPUS_PER_NODE:-8} +export BACKEND_PATH=${BACKEND_PATH:-${REPO}/third_party/Megatron-LM} +export PRIMUS_OUTPUT_ROOT=${PRIMUS_OUTPUT_ROOT:-${REPO}/output} +export PRIMUS_TEAM=${PRIMUS_TEAM:-amd} +export PRIMUS_USER=${PRIMUS_USER:-$(whoami)} + +# Show OOM tracebacks from every rank, not just rank 0. +export LOCAL_RANKS=${LOCAL_RANKS:-"--local-ranks-filter 0,1,2,3,4,5,6,7"} +# Big variable-size buffers fragment the caching allocator badly at long context. +export PYTORCH_ALLOC_CONF=${PYTORCH_ALLOC_CONF:-expandable_segments:True} +export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-1} +export NVTE_CK_USES_BWD_V3=${NVTE_CK_USES_BWD_V3:-1} +export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-lo} +export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-lo} + +# ---- gfx942 / CDNA3 LDS safety ------------------------------------------------ +# The V4 sparse-MLA backward Triton kernel is tuned for gfx950 / CDNA4 (160 KB LDS). +# gfx942 / CDNA3 has 64 KB, and the stock pipeline staging asks for 73728 B, so the +# kernel fails to COMPILE (not to run): +# triton.runtime.errors.OutOfResources: shared memory, Required: 73728, limit: 65536 +# num_stages=1 disables Triton's LDS multi-buffering and brings it under the limit. +# +# Measured on MI308X: this is the ONLY switch needed. PRIMUS_HC_TRITON and +# PRIMUS_DSA_DKV_SAFE can both stay at their defaults -- an earlier version of this +# script forced them off after mis-attributing a `Required: 131072` failure (which +# actually came from the triton_v1 ATTENTION backend, not from hyper-connection) to +# the hc kernel. Forcing PRIMUS_HC_TRITON=0 is not needed and costs performance, +# because it drops hyper-connection onto a PyTorch fallback. +_GFX_ARCH=$(rocm_agent_enumerator 2>/dev/null | grep -m1 -oE 'gfx[0-9]+' || true) +if [ "${_GFX_ARCH}" = "gfx942" ]; then + export PRIMUS_DSA_BWD_NUM_STAGES=${PRIMUS_DSA_BWD_NUM_STAGES:-1} + echo "[run] gfx942 detected: PRIMUS_DSA_BWD_NUM_STAGES=$PRIMUS_DSA_BWD_NUM_STAGES (64 KB LDS)" +fi + +# V4 triton kernel knobs (mirrors run_deepseek_v4_flash.sh). +export PRIMUS_ROPE_TRITON=${PRIMUS_ROPE_TRITON:-1} +export PRIMUS_SINKHORN_TRITON=${PRIMUS_SINKHORN_TRITON:-1} +export PRIMUS_HC_TRITON=${PRIMUS_HC_TRITON:-1} +export PRIMUS_INDEXER_TRITON=${PRIMUS_INDEXER_TRITON:-1} +export PRIMUS_V4_ROUTER_TRITON=${PRIMUS_V4_ROUTER_TRITON:-1} +export PRIMUS_STACK_GROUPED_WEIGHT_TRITON=${PRIMUS_STACK_GROUPED_WEIGHT_TRITON:-1} +export PRIMUS_V4_ATTN_BWD_USE_SPLIT=${PRIMUS_V4_ATTN_BWD_USE_SPLIT:-1} +export PRIMUS_V4_CSA_BWD_SEGREDUCE=${PRIMUS_V4_CSA_BWD_SEGREDUCE:-1} + +# ---- experiment knobs ---- +export PRIMUS_SEQ_LENGTH=${SEQ:-4096} +export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_SEQ_LENGTH} +export SFT_JSONL=${SFT_JSONL:-${REPO}/data/sft/sft_4096.jsonl} +export MBS=${MBS:-1} +export GBS=${GBS:-8} +export TRAIN_ITERS=${TRAIN_ITERS:-10} +export PRIMUS_TP=${PRIMUS_TP:-1} +export PRIMUS_PP=${PRIMUS_PP:-1} +export PRIMUS_EP=${PRIMUS_EP:-8} +export PRIMUS_ETP=${PRIMUS_ETP:-1} +export PRIMUS_CP=${PRIMUS_CP:-1} +export PRIMUS_SHARD_HEADS=${PRIMUS_SHARD_HEADS:-false} +export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-4} +export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-256} +export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-"[0, 0, 128, 128]"} +export PRIMUS_RECOMPUTE_LAYERS=${PRIMUS_RECOMPUTE_LAYERS:-4} +export MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-0} +export PRIMUS_HC_MULT=${PRIMUS_HC_MULT:-4} +export PRIMUS_USE_TURBO_DEEPEP=${PRIMUS_USE_TURBO_DEEPEP:-False} +export PRIMUS_USE_V4_ATTENTION_BACKEND=${PRIMUS_USE_V4_ATTENTION_BACKEND:-triton_v2} +export PRIMUS_USE_V4_CSA_ATTENTION_BACKEND=${PRIMUS_USE_V4_CSA_ATTENTION_BACKEND:-triton_v2} + +export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-dsv4_4layer_sft_seq${PRIMUS_SEQ_LENGTH}_tp${PRIMUS_TP}_ep${PRIMUS_EP}} +EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash_4layer-BF16-sft.yaml} + +LOGDIR="$PRIMUS_OUTPUT_ROOT/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" +mkdir -p "$LOGDIR" +echo "[run] EXP=$EXP SEQ=$PRIMUS_SEQ_LENGTH TP=$PRIMUS_TP EP=$PRIMUS_EP GBS=$GBS log=$LOGDIR" + +# EXTRA_ARGS is intentionally left unquoted so it splits into separate CLI flags. +# shellcheck disable=SC2086 +./primus-cli direct -- train pretrain --config "$EXP" \ + --manual_gc True \ + --manual_gc_interval 100 \ + --pp_warmup False --sequence_parallel False \ + --log_avg_skip_iterations 3 \ + --backend_path "$BACKEND_PATH" \ + ${EXTRA_ARGS:-} \ + 2>&1 | tee "$LOGDIR/log_node0.txt" diff --git a/examples/megatron/configs/MI355X/deepseek_v4_flash_4layer-BF16-sft.yaml b/examples/megatron/configs/MI355X/deepseek_v4_flash_4layer-BF16-sft.yaml new file mode 100644 index 000000000..0b2175fa1 --- /dev/null +++ b/examples/megatron/configs/MI355X/deepseek_v4_flash_4layer-BF16-sft.yaml @@ -0,0 +1,199 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:deepseek_v4_flash_4layer-sft} +workspace: ${PRIMUS_WORKSPACE:./output} + +# DeepSeek-V4 Flash, cut to 4 layers, native (Bridge-free) SFT, for long-context +# smoke runs on 8x MI308X (gfx942, 192 GiB/GPU). +# +# Derived from deepseek_v4_flash-BF16-pretrain.yaml. Deltas vs that file: +# * config: sft_trainer.yaml (carries `stage: sft` -> MegatronSFTTrainer) +# * real local tokenizer + local jsonl dataset (mock_data is FATAL under sft: +# the mock_data patch force-sets NullTokenizer, whose text_to_ids does +# int(x) on whitespace-split text and dies on natural language) +# * 4 layers, MTP off, compress_ratios truncated to 4 entries +# * gfx942-safe kernels (turbo sparse-MLA is gfx950-only) +# * random init: no DeepSeek-V4 Megatron checkpoint exists, and +# Megatron-Bridge has no V4 importer. Same precedent as the in-tree +# qwen3_235B_A22B_4layer-BF16-sft.yaml smoke config. + +modules: + pre_trainer: + framework: megatron + config: sft_trainer.yaml + model: ${PRIMUS_MODEL:deepseek_v4_flash}.yaml + overrides: + # ---------- log ---------- + wandb_project: "Primus_DeepSeek_V4_SFT" + stderr_sink_level: DEBUG + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + log_interval: 1 + disable_wandb: true + disable_tensorboard: true + + # ---------- hyper parameters ---------- + train_iters: ${TRAIN_ITERS:10} + micro_batch_size: ${MBS:1} + global_batch_size: ${GBS:8} + seq_length: ${PRIMUS_SEQ_LENGTH:4096} + max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:4096} + lr: ${PRIMUS_LR:1.0e-5} + min_lr: 0.0 + lr_warmup_iters: 0 + lr_decay_iters: ${TRAIN_ITERS:10} + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + init_method_std: 0.008 + norm_epsilon: 1.0e-6 + + # ---------- parallel ---------- + # V4 attention does NOT shard heads across TP + # (deepseek_v4_attention.py:547 sets num_attention_heads_per_partition = + # num_heads, no divide()), so TP only shards weights and the vocab-parallel + # output layer. TP=1 is the only config the V4 kernels are gate-tested on. + tensor_model_parallel_size: ${PRIMUS_TP:1} + pipeline_model_parallel_size: ${PRIMUS_PP:1} + expert_model_parallel_size: ${PRIMUS_EP:8} + expert_tensor_parallel_size: ${PRIMUS_ETP:1} + # Context parallel. Supported for the DENSE branch only (compress_ratio 0): + # deepseek_v4_cp.py exchanges the left boundary KV and the sparse-MLA adapter + # validates the window against global positions. HCA (128) and CSA (4) layers have + # no CP path yet, so compress_ratios must be all-zero when this is > 1 -- the + # attention module raises if it is not. seq_length must divide by cp_size. + context_parallel_size: ${PRIMUS_CP:1} + overlap_grad_reduce: ${PRIMUS_OVERLAP_GRAD_REDUCE:true} + overlap_param_gather: ${PRIMUS_OVERLAP_PARAM_GATHER:true} + gradient_accumulation_fusion: true + + # ---------- data (SFT: real tokenizer + real local jsonl) ---------- + mock_data: false + # Non-null purely to make the pretrain data-prep hook early-return + # (runner/helpers/hooks/train/pretrain/megatron/prepare.py:108); it would + # otherwise try to download+tokenize bookcorpus and demand HF_TOKEN. + # The SFT dataset provider never reads this — it reads sft_dataset_name. + train_data_path: ["unused_placeholder"] + valid_data_path: null + test_data_path: null + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: ${V4_TOKENIZER:/apps/DeepSeek-V4-Flash} + trust_remote_code: true + sft_dataset_name: ${SFT_JSONL:./data/sft/sft_4096.jsonl} + sft_conversation_format: alpaca + enable_packed_sequences: false + use_packed_attention: false # TE thd varlen path hangs/SIGABRTs on ROCm + num_workers: 4 + dataloader_type: cyclic + create_attention_mask_in_dataloader: false # would be an [S,S] bool tensor + # = 16 GiB at 128k, 1 TiB at 1M + + # ---------- 4-layer cut ---------- + num_layers: ${PRIMUS_TOTAL_LAYERS:4} + compress_ratios: ${PRIMUS_COMPRESS_RATIOS:"[0, 0, 128, 128]"} + mtp_num_layers: ${MTP_NUM_LAYERS:0} + num_experts: ${PRIMUS_NUM_EXPERTS:256} + index_topk: ${PRIMUS_INDEX_TOPK:512} + moe_router_topk: ${PRIMUS_MOE_TOPK:6} + moe_ffn_hidden_size: ${PRIMUS_MOE_FFN:2048} + q_lora_rank: ${PRIMUS_Q_LORA:1024} + o_lora_rank: ${PRIMUS_O_LORA:1024} + + # ---------- DeepSeek-V4 specific ---------- + # P14: shard attention + indexer heads across TP (not just weights). + # At TP=8 this puts the indexer at 8 local heads, which is inside the fused + # Triton scoring kernel's _SUPPORTED_H -- so CSA gets the fused path for free. + v4_shard_attention_heads: ${PRIMUS_SHARD_HEADS:false} + + hybrid_attention_enabled: true + attn_sink: true + hc_use_sinkhorn: true + hc_mult: ${PRIMUS_HC_MULT:4} + # NARROWING KNOB (default = stock V4-Flash width, 64). + # V4 attention does not shard heads across TP, so the query tensor + # [B, S, H, 512] is materialized at FULL width on every rank: 64 KiB/token + # at H=64. Three such tensors are live inside one layer, so at S=1M that is + # 192 GiB = the entire card. Reducing H is the only lever that shrinks it. + # Any run with H != 64 is NOT DeepSeek-V4-Flash and must be labelled so. + num_attention_heads: ${PRIMUS_NUM_HEADS:64} + kv_channels: ${PRIMUS_HEAD_DIM:512} # MLA latent head_dim; 512 = stock V4-Flash + mtp_use_separate_hc_head: true + moe_router_score_function: sqrtsoftplus + # deepseek_v4_base.yaml enables expert bias (noaux_tc), but Megatron + # asserts expert bias requires moe_router_score_function=sigmoid, which + # conflicts with V4's sqrtsoftplus. run_deepseek_v4.sh defaults + # PRIMUS_MOE_ENABLE_EXPERT_BIAS=False for exactly this reason. + moe_router_enable_expert_bias: false + swiglu_limit: 10.0 + + # ---------- optimizer ---------- + # The pretrain yaml uses bf16 for grads + both Adam moments. That is fine at + # seq 4k, but bf16 has an 8-bit mantissa and the grad accumulation runs over + # seq_length tokens -- at 128k it produces NaN in the grad norm on iteration 2 + # (measured: iter1 loss 11.92 / grad-norm 16.45, iter2 NaN, reproducible with + # both HCA and all-dense compress_ratios, at 68% memory so not an OOM). + # fp32 grads + fp32 exp_avg_sq fix it. + use_precision_aware_optimizer: true + main_grads_dtype: ${PRIMUS_MAIN_GRADS_DTYPE:fp32} + main_params_dtype: ${PRIMUS_MAIN_PARAMS_DTYPE:fp32} + exp_avg_dtype: ${PRIMUS_EXP_AVG_DTYPE:bf16} + exp_avg_sq_dtype: ${PRIMUS_EXP_AVG_SQ_DTYPE:fp32} + use_distributed_optimizer: true + + enable_experimental: true + apply_rope_fusion: false + # deepseek_v4_base.yaml sets rope_type: yarn, but the common-attention path + # asserts rope_type == "rope" ("Common attention only support + # rope_type=\"rope\", but got yarn"). run_deepseek_v4.sh passes + # `--rope_type rope` for exactly this reason; the V4 compressed branches + # apply their own YaRN internally via compress_rope_theta. + rope_type: rope + + # ---------- recompute (mandatory at long context) ---------- + recompute_granularity: ${PRIMUS_RECOMPUTE_GRANULARITY:full} + recompute_method: ${PRIMUS_RECOMPUTE_METHOD:block} + recompute_num_layers: ${PRIMUS_RECOMPUTE_LAYERS:4} + + # ---------- ckpt: random init, save nothing ---------- + finetune: false + auto_continue_train: false + load: null + save: null + no_save_optim: true + no_save_rng: true + disable_last_saving: true + ckpt_format: torch + eval_iters: 0 + eval_interval: 1000000 + + # ---------- LoRA off: full-parameter SFT ---------- + sft_sanitize_nan_grads: ${PRIMUS_SANITIZE_NAN:false} + + lora: + enabled: false + + # ---------- kernels: gfx942-safe ---------- + # turbo / gluon sparse-MLA backends are gfx950/CDNA4-only. + enable_primus_turbo: true + use_turbo_attention: false # must stay off so the dense (cr=0) + # path uses the V4 sparse-MLA backend + use_turbo_grouped_gemm: true + use_turbo_rms_norm: true + use_turbo_deepep: ${PRIMUS_USE_TURBO_DEEPEP:false} # default CU count is 80 + # == 100% of MI308X's CUs + moe_shared_expert_overlap: false + moe_router_dtype: fp32 + turbo_sync_free_moe_stage: 1 + + use_v4_attention_backend: ${PRIMUS_USE_V4_ATTENTION_BACKEND:triton_v2} + use_v4_csa_attention_backend: ${PRIMUS_USE_V4_CSA_ATTENTION_BACKEND:triton_v2} + use_v4_fp8_indexer: false + use_v4_compiled_sinkhorn: false + + moe_use_fused_router_with_aux_score: true + moe_permute_fusion: true + + # ---------- cross entropy ---------- + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py index b06eac1ea..b71dc3f50 100644 --- a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py @@ -60,6 +60,15 @@ import ast import logging from contextlib import nullcontext + +# CPU activation offloading helper, re-exported by Megatron's transformer_block when +# TransformerEngine is present. None when TE is unavailable. +try: + from megatron.core.transformer.transformer_block import ( + get_cpu_offload_context as _get_cpu_offload_context, + ) +except ImportError: # pragma: no cover + _get_cpu_offload_context = None from dataclasses import dataclass from typing import Callable, List, Optional, Union @@ -838,6 +847,30 @@ def __init__( self.pg_collection = pg_collection # Required by pipeline schedules (same contract as TransformerBlock). self.input_tensor = None + + # CPU activation offloading. The parent __init__ is bypassed (see the class + # docstring), so this has to be set up here or `self.offload_context` simply does + # not exist and enabling --cpu-offloading-num-layers dies with an AttributeError. + # TE's context installs saved_tensors_hooks, which intercept every tensor saved + # for backward inside it -- including V4's custom autograd Functions, since + # offloading is opt-out (`mark_not_offload`) rather than opt-in. + self.offload_context, self.group_prefetch_offload_commit_async = nullcontext(), None + if _get_cpu_offload_context is not None: + self.offload_context, self.group_prefetch_offload_commit_async = _get_cpu_offload_context( + config.cpu_offloading, + config.cpu_offloading_num_layers, + config.num_layers, + config.cpu_offloading_activations, + config.cpu_offloading_weights, + config.cpu_offloading_double_buffering, + ) + config._cpu_offloading_context = self.offload_context if config.cpu_offloading else None + elif getattr(config, "cpu_offloading", False): + raise RuntimeError( + "cpu_offloading requires TransformerEngine's get_cpu_offload_context, " + "which is unavailable in this build." + ) + logger.info( "[DeepSeek-V4] decoder block initialized (pre_process=%s post_process=%s).", pre_process, @@ -1153,12 +1186,26 @@ def forward( if recompute_local is not None and local_idx in recompute_local: x = self._forward_layer_checkpointed(layer, x, position_ids, token_ids, global_idx) else: - with self._layer_fp8_context(global_idx): + # ``self.offload_context`` comes from TransformerBlock.__init__ (TE's + # get_cpu_offload_context). This loop replaces the parent's, so it has to + # re-enter that context itself -- without this the context object exists + # but is never active and --cpu-offloading-num-layers is a silent no-op. + # TE v2 installs saved_tensors_hooks, which intercept EVERY tensor saved + # for backward inside the region (offload is opt-out via + # `mark_not_offload`), so V4's custom autograd Functions are covered too. + # Mirrors TransformerBlock.forward (transformer_block.py:828-850). + with self.offload_context, self._layer_fp8_context(global_idx): x, _ = layer( x, position_ids=position_ids, token_ids=token_ids, ) + if ( + torch.is_grad_enabled() + and self.config.cpu_offloading + and self.group_prefetch_offload_commit_async is not None + ): + x = self.group_prefetch_offload_commit_async(x) # Final HC collapse on post_process stage; non-final stages # forward the multi-stream form through PP P2P. diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py index 07a9b0013..6b93b2074 100644 --- a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py @@ -263,6 +263,39 @@ def _build_linear_projection_spec( ) +def v4_shard_heads(config) -> bool: + """P14: shard attention heads across TP instead of gathering them back. + + Off by default, so every existing recipe keeps the byte-identical + ``gather_output=True`` behaviour. When on, each TP rank owns + ``num_attention_heads / tp`` heads end to end -- q-up stays sharded, the + attention kernels run on the local head slice, and the grouped-O down + projection holds only this rank's ``o_groups / tp`` groups, whose partial + results the row-parallel ``linear_o_b`` all-reduces. + + Requires ``num_attention_heads % tp == 0`` and ``o_groups % tp == 0``: the + grouped-O split is defined over ``o_groups``, so a rank must own whole groups. + """ + if not bool(getattr(config, "v4_shard_attention_heads", False)): + return False + tp = int(getattr(config, "tensor_model_parallel_size", 1) or 1) + if tp <= 1: + return False + heads = int(config.num_attention_heads) + groups = int(getattr(config, "o_groups", 1) or 1) + if heads % tp != 0: + raise ValueError( + f"v4_shard_attention_heads requires num_attention_heads ({heads}) " + f"divisible by tensor_model_parallel_size ({tp})." + ) + if groups % tp != 0: + raise ValueError( + f"v4_shard_attention_heads requires o_groups ({groups}) divisible by " + f"tensor_model_parallel_size ({tp}); a rank must own whole grouped-O groups." + ) + return True + + def _build_column_parallel_spec( *, config: DeepSeekV4TransformerConfig, @@ -431,6 +464,10 @@ def _build_v4_attention_submodules( provider=provider, in_features=q_lora_rank, out_features=q_out, + # P14: keep the head slice local instead of all-gathering it back to full + # width. This is the change that makes TP shard attention ACTIVATIONS, not + # just weights -- see v4_shard_heads(). + gather_output=not v4_shard_heads(config), ), linear_kv=_build_linear_projection_spec( config=config, @@ -444,18 +481,38 @@ def _build_v4_attention_submodules( if o_lora_rank > 0: n_per_group = q_out // o_groups - submods.linear_o_a = _build_linear_projection_spec( - config=config, - provider=provider, - in_features=n_per_group, - out_features=o_groups * o_lora_rank, - ) - submods.linear_o_b = _build_row_parallel_spec( - config=config, - provider=provider, - in_features=o_groups * o_lora_rank, - out_features=hidden_size, - ) + if v4_shard_heads(config): + # Each rank owns o_groups/tp whole groups. Sharding the o_a OUTPUT dim + # (o_groups * o_lora_rank) by TP hands each rank exactly those groups' + # rows, and linear_o_b -- already row-parallel over the same axis -- + # all-reduces the per-group partial sums into the full output. + submods.linear_o_a = _build_column_parallel_spec( + config=config, + provider=provider, + in_features=n_per_group, + out_features=o_groups * o_lora_rank, + gather_output=False, + ) + submods.linear_o_b = _build_row_parallel_spec( + config=config, + provider=provider, + in_features=o_groups * o_lora_rank, + out_features=hidden_size, + input_is_parallel=True, + ) + else: + submods.linear_o_a = _build_linear_projection_spec( + config=config, + provider=provider, + in_features=n_per_group, + out_features=o_groups * o_lora_rank, + ) + submods.linear_o_b = _build_row_parallel_spec( + config=config, + provider=provider, + in_features=o_groups * o_lora_rank, + out_features=hidden_size, + ) else: submods.linear_proj = _build_row_parallel_spec( config=config, diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py index 8964bb5c2..83f2f560e 100644 --- a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py @@ -12,6 +12,7 @@ construction path. """ +import os from typing import Literal, Optional, Union from megatron.core import tensor_parallel @@ -37,6 +38,11 @@ ) +def _v4_fused_ce_enabled() -> bool: + """Chunked linear+CE switch, shared spelling with the GPTModel patch.""" + return os.environ.get("FUSED_LINEAR_CE", "0") == "1" + + class DeepseekV4Model(LanguageModule): """DeepSeek-V4 language model rooted on LanguageModule.""" @@ -308,6 +314,42 @@ def forward( scale_logits_fn=self._scale_logits if getattr(self.config, "use_mup", False) else None, ) + # Chunked linear + cross-entropy. The stock path below materialises the full + # [s, b, vocab] logits, which at long context is the single largest allocation in + # the whole step -- measured at 1M with CP=8: s_local=131072 x vocab=129280 x 2 B + # = 31.56 GiB, bigger than any attention tensor and bigger than the whole model. + # Primus already ships this as patches/fused_linear_ce_patches.py, but that patch + # hooks ``GPTModel._postprocess`` and DeepseekV4Model derives from LanguageModule, + # so it never applied here. Off by default; FUSED_LINEAR_CE=1 turns it on. + # + # Restricted to TP=1: the chunked path does a plain matmul against the full output + # weight, which is only equivalent to ``output_layer`` when the vocab is not + # sharded. Under vocab-parallel TP it would feed already-gathered logits into a + # vocab-parallel cross-entropy and double-count. + if labels is not None and _v4_fused_ce_enabled(): + from megatron.core import parallel_state + + if parallel_state.get_tensor_model_parallel_world_size() == 1: + from megatron.training import get_args + + if getattr(get_args(), "overlap_param_gather", False): + raise RuntimeError( + "FUSED_LINEAR_CE=1 is incompatible with overlap_param_gather. The chunked " + "backward issues one torch.autograd.grad per sequence chunk, which changes " + "how many times each parameter's backward hook fires; the distributed " + "optimizer's overlapped parameter all-gather then desyncs and fails in " + "start_param_sync with an empty error message. Set " + "overlap_param_gather: false (PRIMUS_OVERLAP_PARAM_GATHER=false)." + ) + + from primus.backends.megatron.patches.fused_linear_ce_patches import ( + _chunk_size, + _ChunkedLinearCE, + ) + + w = output_weight if output_weight is not None else self.output_layer.weight + return _ChunkedLinearCE.apply(hidden_states, w, labels, self, _chunk_size()) + logits, _ = self.output_layer( hidden_states, weight=output_weight, diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py index f3d40315b..0dfa7655d 100644 --- a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py @@ -154,6 +154,13 @@ class DeepSeekV4TransformerConfig(MLATransformerConfig): o_lora_rank: int = 0 # ---- DeepSeek-V4 MoE routing / expert extras ---- + # P14: shard attention heads across tensor parallel instead of gathering them. + # Off by default so existing recipes are byte-identical. When on, TP shards the + # [B, S, H, head_dim] query ACTIVATION (not just weights), which is what limits + # long-context training on V4. Requires num_attention_heads % tp == 0 and + # o_groups % tp == 0. + v4_shard_attention_heads: bool = False + num_hash_layers: int = 0 hash_routing_seed: int = 0 diff --git a/primus/backends/megatron/core/transformer/deepseek_v4_attention.py b/primus/backends/megatron/core/transformer/deepseek_v4_attention.py index 543f7fbd9..c33b0b875 100644 --- a/primus/backends/megatron/core/transformer/deepseek_v4_attention.py +++ b/primus/backends/megatron/core/transformer/deepseek_v4_attention.py @@ -110,6 +110,20 @@ logger = logging.getLogger(__name__) +def _v4_get_cp_group(): + from primus.backends.megatron.core.transformer.deepseek_v4_cp import get_cp_group + + return get_cp_group() + + +def _v4_exchange_boundary_kv(kv, d_window, cp_group): + from primus.backends.megatron.core.transformer.deepseek_v4_cp import ( + exchange_boundary_kv, + ) + + return exchange_boundary_kv(kv, d_window, cp_group) + + def _require_gfx950() -> None: """Assert the current device is gfx950 / CDNA4 before using the gluon backend. @@ -542,9 +556,24 @@ def __init__( self.pg_collection = pg_collection # ---- shape fields (read by helpers in this class) ---- + # ---- P14: head sharding across TP ---------------------------------- + # Default (v4_shard_attention_heads off): every rank materialises all + # `num_heads` heads, because linear_q_up_proj gathers its output back to full + # width. TP then shards weights only, and the [B, S, H, head_dim] query is + # replicated -- 64 KiB/token at V4-Flash width, which is what caps the usable + # sequence length. With P14 on, each rank owns num_heads/tp heads end to end. + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import ( + v4_shard_heads as _v4_shard_heads, + ) + + self.shard_heads = _v4_shard_heads(config) + self.tp_size = int(getattr(config, "tensor_model_parallel_size", 1) or 1) if self.shard_heads else 1 + num_heads_local = num_heads // self.tp_size + self.hidden_size = hidden_size - self.num_heads = num_heads - self.num_attention_heads_per_partition = num_heads + self.num_heads = num_heads_local + self.num_heads_global = num_heads + self.num_attention_heads_per_partition = num_heads_local self.num_query_groups_per_partition = 1 # single-latent KV self.head_dim = head_dim self.rotary_dim = rotary_dim @@ -641,7 +670,8 @@ def __init__( # TE-fused sink primitive can land as a new spec field once it # actually replaces the inline path. if attn_sink_enabled: - self.attn_sink = nn.Parameter(torch.zeros(num_heads)) + # One sink per head, so it shards with the heads under P14. + self.attn_sink = nn.Parameter(torch.zeros(num_heads_local)) else: self.register_parameter("attn_sink", None) @@ -923,6 +953,10 @@ def _build_indexer(self, spec: Optional[Union[ModuleSpec, type]]) -> nn.Module: index_topk=index_topk, compress_ratio=self.compress_ratio, use_fp8_qk=bool(getattr(self.config, "use_v4_fp8_indexer", False)), + # P14: hand the indexer the TP group so it can shard its heads and + # all-reduce the partial score sums. None (the default) keeps the + # replicated, unsharded behaviour. + tp_group=self._v4_tp_group() if self.shard_heads else None, ) if spec is None: return Indexer(**kwargs) @@ -932,6 +966,21 @@ def _build_indexer(self, spec: Optional[Union[ModuleSpec, type]]) -> nn.Module: # internals # ------------------------------------------------------------------ + @staticmethod + def _v4_tp_group(): + """Tensor-parallel process group, or None when TP is off / dist is not up.""" + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + return None + from megatron.core import parallel_state + + try: + g = parallel_state.get_tensor_model_parallel_group() + except (AssertionError, RuntimeError): + return None + return g if g is not None and g.size() > 1 else None + @property def rope(self) -> DualRoPE: return self._rope[0] @@ -1177,8 +1226,11 @@ def _grouped_o_projection(self, attn: torch.Tensor) -> torch.Tensor: ``.weight`` after ``build_module``.) """ B, S, H, Dh = attn.shape - G = self.o_groups - attn_g = attn.reshape(B, S, G, (H * Dh) // G) # [B, S, G, H*Dh/G] + # Under P14 this rank owns o_groups/tp groups and H is already the local head + # count, so H*Dh/G_local is the same n_per_group as the unsharded path -- the + # group WIDTH is a property of o_groups, not of TP. + G = self.o_groups // self.tp_size + attn_g = attn.reshape(B, S, G, (H * Dh) // G) # [B, S, G_local, H*Dh/G_local] wo_a = self.linear_o_a weight = wo_a.weight if hasattr(wo_a, "weight") else None @@ -1188,7 +1240,7 @@ def _grouped_o_projection(self, attn: torch.Tensor) -> torch.Tensor: o = _projection_forward(wo_a, attn_g.reshape(B, S, -1)) o = o.view(B, S, G * self.o_lora_rank) else: - wo_a_w = weight.view(G, self.o_lora_rank, (H * Dh) // G) + wo_a_w = weight.view(G, self.o_lora_rank, (H * Dh) // G) # G is local if _v4_o_a_fp8_enabled(self.config): o = _fp8_grouped_o_a(attn_g, wo_a_w) # per-group MXFP8 else: @@ -1211,7 +1263,35 @@ def _build_compressed_pool(self, hidden: torch.Tensor) -> torch.Tensor: Returns ``[B, P, head_dim]`` where ``P = S // compress_ratio``. """ device = hidden.device - pooled = self.compressor(hidden) # [B, P, head_dim] + cp_group = _v4_get_cp_group() + if cp_group is None: + pooled = self.compressor(hidden) # [B, P, head_dim] + else: + # ---- context parallel ------------------------------------------------ + # Each rank compresses only its own rows, then the pools are all-gathered + # so every query can see the whole sequence's compressed history. Rank + # order IS sequence order here (single BSHD sequence, S_local a multiple + # of ratio), so no seq-major/rank-major remap is needed. + from primus.backends.megatron.core.transformer.deepseek_v4_cp import ( + build_global_pool, + compressor_boundary_rows, + exchange_boundary_kv, + ) + + nb = compressor_boundary_rows(self.compress_ratio, bool(self.compressor.overlap)) + if nb > 0: + # Overlap mode stitches window i with window i-1, which at a shard + # boundary lives on the left neighbour. Prepend those rows, compress, + # then drop the extra leading pool row they produced. + bnd = exchange_boundary_kv( + hidden.reshape(hidden.shape[0], hidden.shape[1], 1, hidden.shape[2]), + nb, + cp_group, + ).reshape(hidden.shape[0], nb, hidden.shape[2]) + pooled_local = self.compressor(torch.cat([bnd, hidden], dim=1))[:, 1:] + else: + pooled_local = self.compressor(hidden) + pooled = build_global_pool(pooled_local, cp_group) B, P = pooled.shape[0], pooled.shape[1] # Compress-base partial RoPE on compressed indices [0..P). Positions are @@ -1232,10 +1312,13 @@ def _hca_extra_kv( ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Build the HCA (compress_ratio == 128) compressed branch. - Returns ``(extra_k_bh, extra_v_bh, extra_mask)`` where the + Returns ``(extra_k_bh, extra_v_bh, pool, extra_mask)`` where the compressed pool is broadcast across H heads (single-latent compressor output) and the additive mask is shape ``[S, P]`` - (broadcasts over B, H). + (broadcasts over B, H). ``pool`` is the pre-broadcast ``[B, P, head_dim]`` + latent, which the caller concatenates onto the raw KV latent -- see + :meth:`_cp_prepend_boundary` for why the concat must not happen on the + broadcast views. Per the techblog: pool position ``s`` covers raw tokens ``[s*ratio, (s+1)*ratio)``; query at raw token ``t`` may attend @@ -1252,7 +1335,7 @@ def _hca_extra_kv( pool_bh = pool_h.transpose(1, 2) extra_mask = self._hca_extra_mask_cached(S, P, device, dtype) - return pool_bh, pool_bh, extra_mask # K = V = compressed pool + return pool_bh, pool_bh, pool, extra_mask # K = V = compressed pool def _hca_extra_mask_cached(self, S: int, P: int, device, dtype): """HCA additive causal mask ``[S, P]``, cached (data-independent). @@ -1263,6 +1346,19 @@ def _hca_extra_mask_cached(self, S: int, P: int, device, dtype): compressed-layer forward. Bit-identical. PRIMUS_COMPRESS_MASK_CACHE=0 forces the eager rebuild. """ + cp_group = _v4_get_cp_group() + if cp_group is not None: + # Under CP the pool is global but the queries are this rank's slice, so the + # visibility test must use global positions. Not cached: it depends on the + # rank, and rebuilding an [S, P] byte mask once per layer is cheap next to + # the attention itself. + from primus.backends.megatron.core.transformer.deepseek_v4_cp import ( + compressed_causal_mask, + ) + + return compressed_causal_mask( + S, P, cp_group.rank() * S, self.compress_ratio, device=device, dtype=dtype + ) if os.environ.get("PRIMUS_COMPRESS_MASK_CACHE", "1") == "0": t = torch.arange(S, device=device).unsqueeze(1) s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * self.compress_ratio - 1 @@ -1286,6 +1382,7 @@ def _csa_forward( k_local_bh: torch.Tensor, # [B, H, S, head_dim] v_local_bh: torch.Tensor, # [B, H, S, head_dim] local_mask: torch.Tensor, # [S, S] — built by caller; unused here, see below + kv: Optional[torch.Tensor] = None, # [B, S, 1, head_dim] post-RoPE latent; CP only ) -> torch.Tensor: """CSA (compress_ratio == 4) joint local-SWA + sparse-compressed attention. @@ -1315,6 +1412,26 @@ def _csa_forward( B, H, S, Dh = q_bh.shape dtype = hidden.dtype + # 0) Context parallel: like the dense and HCA branches, the raw-token sliding + # window straddles the shard edge, so this rank needs the left neighbour's + # trailing `d_window` post-RoPE KV rows. The pool half is already handled + # (all-gathered to global + the indexer scores against global positions). + cp_dwindow = cp_global_start = 0 + kv_latent = None + if _v4_get_cp_group() is not None: + if self._csa_backend != "triton_v2": + raise NotImplementedError( + "DeepSeek-V4 CSA context parallelism is only wired through the " + f"triton_v2 CSA backend; got '{self._csa_backend}'. The other CSA " + "backends build the local window themselves and would silently drop " + "the cross-shard part of it." + ) + if kv is None: + raise RuntimeError("_csa_forward needs the single-latent kv under CP") + k_local_bh, v_local_bh, kv_latent, cp_dwindow, cp_global_start = self._cp_prepend_boundary( + kv, B, S + ) + # 1) Compressed pool with compress-base RoPE. pool = self._build_compressed_pool(hidden) # [B, P, head_dim] P = pool.shape[1] @@ -1378,10 +1495,13 @@ def _csa_forward( scale=self._attention_scale(), ) if be == "triton_v2": + # Hand the kernel the un-broadcast latent when CP gave us one: it reads a + # single key row per position anyway, and the broadcast view's gradient would + # be an 8 GiB [B, H, Skv, D] buffer that is zero except at head 0. return v4_csa_attention_v2( q_bh, - k_local_bh, - v_local_bh, + k_local_bh if kv_latent is None else kv_latent, + v_local_bh if kv_latent is None else None, pool, topk_idxs=topk_idxs, sink=self.attn_sink, @@ -1389,6 +1509,9 @@ def _csa_forward( attn_dropout=self.attn_dropout, training=self.training, scale=self._attention_scale(), + cp_dwindow=cp_dwindow, + cp_global_start=cp_global_start, + k_is_latent=kv_latent is not None, ) if be == "flydsl_v1": return self._v4_csa_attention_flydsl( @@ -1458,7 +1581,21 @@ def _csa_forward( # public forward # ------------------------------------------------------------------ - def _attention_backend_forward(self, q_bh, k, v, *, additive_mask, hca_local_seqlen, S, device, dtype): + def _attention_backend_forward( + self, + q_bh, + k, + v, + *, + additive_mask, + hca_local_seqlen, + S, + device, + dtype, + cp_dwindow=0, + cp_global_start=0, + k_latent=None, + ): """Dense (cr=0) / HCA (cr=128) dispatch on ``use_v4_attention_backend``.""" be = self._attn_backend if be == "gluon": @@ -1514,10 +1651,18 @@ def _attention_backend_forward(self, q_bh, k, v, *, additive_mask, hca_local_seq hca_local_seqlen=hca_local_seqlen, ) if be == "triton_v2": + # This kernel reads one key row per position (single-latent MQA), so hand it + # the un-broadcast [B, Skv, 1, D] latent when we have it. The head-broadcast + # view is free forward, but its gradient would be a [B, H, Skv, D] buffer that + # is zero except at head 0 -- 8.5 GiB at 1M with CP=8. Only this backend takes + # the latent form; the others still get the broadcast views. return v4_attention_v2( q_bh, - k, - v, + k if k_latent is None else k_latent, + v if k_latent is None else None, + cp_dwindow=cp_dwindow, + cp_global_start=cp_global_start, + k_is_latent=k_latent is not None, sink=self.attn_sink, swa_window=int(self.attn_sliding_window), additive_mask=additive_mask, @@ -1553,6 +1698,45 @@ def _attention_backend_forward(self, q_bh, k, v, *, additive_mask, hca_local_seq mask = local_mask if additive_mask is None else torch.cat([local_mask, additive_mask], dim=-1) return self._attention_forward(q_bh, k, v, mask) + def _cp_prepend_boundary(self, kv, B, S): + """Prepend the left neighbour's trailing window rows to the local KV. + + Every branch that runs a sliding window over RAW tokens needs this, not just + the dense one: a query near the shard start would otherwise lose the part of + its window that lives on the previous rank. Returns + ``(k_bh, v_bh, kv_latent, cp_dwindow, cp_global_start)``; with CP off it returns + the unmodified head-expanded views and ``(0, 0)``, which reproduces the non-CP + path exactly. + + The concat happens on the SINGLE-LATENT ``[B, S, 1, D]`` tensor and the expand + after. Concatenating the head-expanded ``[B, H, S, D]`` view instead would + materialise a real H-fold tensor for both K and V -- 8.6 GB each at 128k rows + with H=64 -- where the expand is otherwise free. K and V are the same tensor in + V4's single-latent design, so one buffer serves both. + + ``kv_latent`` is that pre-expand ``[B, Skv, 1, D]`` buffer. Callers that need to + concatenate anything else onto the key axis (HCA appends its compressed pool) + MUST concatenate onto this and expand afterwards, for exactly the reason above: + ``torch.cat`` on a stride-0 expanded view materialises the H-fold copy that the + expand was avoiding. + """ + cp_group = _v4_get_cp_group() + if cp_group is None: + kv_bh = kv.expand(B, S, self.num_heads, self.head_dim).transpose(1, 2) + return kv_bh, kv_bh, kv, 0, 0 + if self._attn_backend != "triton_v2": + raise NotImplementedError( + "DeepSeek-V4 context parallelism is only wired through the triton_v2 " + f"backend (the others do not take cp_dwindow/cp_global_start); got " + f"'{self._attn_backend}'. Set USE_V4_ATTENTION_BACKEND=triton_v2." + ) + cp_dwindow = int(self.attn_sliding_window) + cp_global_start = cp_group.rank() * S + boundary_kv = _v4_exchange_boundary_kv(kv, cp_dwindow, cp_group) + kv_full = torch.cat([boundary_kv, kv], dim=1) # [B, d_window + S, 1, D] + kv_full_bh = kv_full.expand(B, cp_dwindow + S, self.num_heads, self.head_dim).transpose(1, 2) + return kv_full_bh, kv_full_bh, kv_full, cp_dwindow, cp_global_start + def forward( self, hidden: torch.Tensor, @@ -1598,6 +1782,13 @@ def forward( v_local_bh = v_h.transpose(1, 2) if self.compress_ratio == 0: + # ---- context parallel (dense / SWA branch) ---------------------- + # This branch is index-driven, so CP needs only the d_window post-RoPE KV rows + # left of this shard plus the shard's global offset; the kernel is unchanged. + # cp_dwindow == cp_global_start == 0 reproduces the non-CP path exactly. + k_local_bh, v_local_bh, kv_latent, cp_dwindow, cp_global_start = self._cp_prepend_boundary( + kv, B, S + ) out_bh = self._attention_backend_forward( q_bh, k_local_bh, @@ -1607,23 +1798,47 @@ def forward( S=S, device=device, dtype=dtype, + cp_dwindow=cp_dwindow, + cp_global_start=cp_global_start, + k_latent=kv_latent, ) elif self.compress_ratio == 128: # HCA: the local SWA branch and the compressed-pool branch share ONE # softmax with ONE sink column; concatenate the pool to the local # keys and pass the pool-only additive mask. - extra_k_bh, extra_v_bh, extra_mask = self._hca_extra_kv(hidden) - k_full = torch.cat([k_local_bh, extra_k_bh], dim=2) # along Sk - v_full = torch.cat([v_local_bh, extra_v_bh], dim=2) + # + # Under CP the LOCAL half needs the same left-boundary rows the dense branch + # takes: the pool being global is not enough, because the local SWA still runs + # over raw tokens that straddle the shard edge. The local segment then grows to + # `cp_dwindow + S`, which is what `hca_local_seqlen` has to report -- the adapter + # uses it as the base offset for the pool columns (`base + hca_local_seqlen + ps`), + # so the [S, P] pool mask stays valid unchanged. + _, _, kv_latent, cp_dwindow, cp_global_start = self._cp_prepend_boundary(kv, B, S) + _, _, pool, extra_mask = self._hca_extra_kv(hidden) + # Concatenate the compressed pool onto the raw KV on the SINGLE-LATENT axis, + # then expand across heads -- the expand is a stride-0 view and costs nothing. + # Doing it the other way round (cat on the already-broadcast [B, H, Sk, D] + # views) materialises the H-fold copy the broadcast exists to avoid: at 1M + # with CP=8 that is 8.51 GiB for K and another 8.51 GiB for V, per HCA layer, + # of which the consumer reads 136 MiB -- the sparse-MLA adapter takes only + # `k_bh[:, 0]`, and never reads `v_bh` at all (its backward returns dv=None, + # because V4 is single-latent and the V-side gradient is structurally zero). + # K and V are the same object here for the same reason. + Sk = kv_latent.shape[1] + pool.shape[1] + kv_cat = torch.cat([kv_latent, pool.unsqueeze(2)], dim=1) # [B, Sk, 1, D] + k_full = v_full = kv_cat.expand(B, Sk, self.num_heads, self.head_dim).transpose(1, 2) out_bh = self._attention_backend_forward( q_bh, k_full, v_full, additive_mask=extra_mask, - hca_local_seqlen=S, + hca_local_seqlen=cp_dwindow + S, S=S, device=device, dtype=dtype, + cp_dwindow=cp_dwindow, + cp_global_start=cp_global_start, + k_latent=kv_cat, ) elif self.compress_ratio == 4: # CSA cannot use ``core_attention``: the per-query top-K @@ -1632,8 +1847,12 @@ def forward( # attention — there is no flash-attn kernel that reads a # different per-query subset of keys from a pool. Stays on # eager-Python under plan-3 (a custom kernel is required). - local_mask = self._local_mask(S, device=device, dtype=dtype) - out_bh = self._csa_forward(hidden, q_bh, k_local_bh, v_local_bh, local_mask) + # `_csa_forward` documents `local_mask` as retained for back-compat and + # `del`s it on entry -- the reference op rebuilds the SWA mask from + # `swa_window` itself. Materialising it here costs a dense [S, S] byte + # tensor for nothing: 16 GiB at S=131072, which is what made CSA OOM at + # 128k. Pass None; the callee never reads it. + out_bh = self._csa_forward(hidden, q_bh, k_local_bh, v_local_bh, None, kv) else: # Guarded by __init__; included for static-analysis completeness. raise ValueError(f"Unsupported compress_ratio {self.compress_ratio}") diff --git a/primus/backends/megatron/core/transformer/deepseek_v4_cp.py b/primus/backends/megatron/core/transformer/deepseek_v4_cp.py new file mode 100644 index 000000000..88beaa6e9 --- /dev/null +++ b/primus/backends/megatron/core/transformer/deepseek_v4_cp.py @@ -0,0 +1,211 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Context-parallel support for the DeepSeek-V4 dense / SWA attention branch. + +V4 attention materialises the query tensor at FULL head width on every tensor-parallel rank +(``deepseek_v4_attention.py`` sets ``num_attention_heads_per_partition = num_heads`` with no +``divide()``), so TP shards weights but not attention activations. Context parallelism is the +only lever that shrinks the per-rank sequence, and therefore the only way past the long-context +memory wall: measured on 8x MI308X, an unmodified 4-layer V4-Flash needs ~396 GB at 128k +against a 192 GiB card, while 32k fits in 145 GB. + +For the dense (``compress_ratio == 0``) branch the CP contract is small, because that branch is +already index-driven -- causality and the sliding window live entirely in the index matrix the +sparse-MLA adapter builds, not in the kernel. A CP rank therefore needs exactly two things: + + 1. the ``d_window`` post-RoPE KV rows immediately left of its shard, so a query near the shard + start can still see its full window (:class:`LeftBoundaryExchange`), and + 2. its global row offset, so the adapter can validate the window against global positions + while indexing into the local ``[boundary ++ local]`` buffer + (``cp_dwindow`` / ``cp_global_start`` in ``v4_sparse_mla_adapter``). + +Exchanging post-RoPE KV rather than pre-projection hidden states is deliberate: the neighbour +has already applied RoPE with the correct global positions, and it moves ``d_window`` rows +instead of a full hidden block. + +Ported from NVIDIA/Megatron-LM PR #5087 (`csa_cp_utils.py`), whose CP path is THD-only; this +is the BSHD-shaped equivalent for Primus's V4 attention. +""" + + +import torch +import torch.distributed as dist + + +class LeftBoundaryExchange(torch.autograd.Function): + """Receive the previous CP rank's trailing ``d_window`` rows; scatter grads back. + + Forward is one batched isend/irecv step around the CP ring. Backward returns the boundary + gradient to the rank that actually owns those rows, so they accumulate where the parameters + that produced them live. + """ + + @staticmethod + def forward(ctx, tensor: torch.Tensor, d_window: int, cp_group): + cp_size = cp_group.size() + cp_rank = cp_group.rank() + ctx.cp_group = cp_group + ctx.d_window = d_window + ctx.input_shape = tensor.shape + if tensor.shape[0] < d_window: + raise RuntimeError( + "DeepSeek-V4 CP boundary exchange needs local rows >= d_window: " + f"local_rows={tensor.shape[0]}, d_window={d_window}. Reduce " + "context_parallel_size or the sliding window." + ) + boundary = tensor.new_zeros((d_window,) + tuple(tensor.shape[1:])) + + ops = [] + if cp_rank > 0: + ops.append( + dist.P2POp(dist.irecv, boundary, dist.get_global_rank(cp_group, cp_rank - 1), cp_group) + ) + if cp_rank + 1 < cp_size: + send_tail = tensor[-d_window:].contiguous() + ops.append( + dist.P2POp(dist.isend, send_tail, dist.get_global_rank(cp_group, cp_rank + 1), cp_group) + ) + if ops: + for req in dist.batch_isend_irecv(ops): + req.wait() + return boundary + + @staticmethod + def backward(ctx, grad_boundary: torch.Tensor): + cp_group = ctx.cp_group + cp_size = cp_group.size() + cp_rank = cp_group.rank() + grad_input = grad_boundary.new_zeros(ctx.input_shape) + + ops = [] + recv_grad = None + if cp_rank > 0: + ops.append( + dist.P2POp( + dist.isend, + grad_boundary.contiguous(), + dist.get_global_rank(cp_group, cp_rank - 1), + cp_group, + ) + ) + if cp_rank + 1 < cp_size: + recv_grad = grad_boundary.new_empty(grad_boundary.shape) + ops.append( + dist.P2POp(dist.irecv, recv_grad, dist.get_global_rank(cp_group, cp_rank + 1), cp_group) + ) + if ops: + for req in dist.batch_isend_irecv(ops): + req.wait() + if recv_grad is not None: + grad_input[-ctx.d_window :] = recv_grad + return grad_input, None, None + + +def get_cp_group(): + """The context-parallel process group, or None when CP is off / torch.distributed is not up.""" + if not dist.is_available() or not dist.is_initialized(): + return None + try: + from megatron.core import parallel_state + except ImportError: + return None + try: + group = parallel_state.get_context_parallel_group() + except (AssertionError, RuntimeError): + return None + if group is None or group.size() <= 1: + return None + return group + + +def exchange_boundary_kv(kv_bshd: torch.Tensor, d_window: int, cp_group) -> torch.Tensor: + """Boundary KV for a ``[B, S, 1, head_dim]`` post-RoPE latent. + + Returns ``[B, d_window, 1, head_dim]``. Rank 0 gets zeros, which the adapter's + global-position validity mask then excludes -- no separate special case is needed. + """ + B, S, G, Dh = kv_bshd.shape + if B != 1: + raise RuntimeError(f"DeepSeek-V4 CP currently assumes micro_batch_size=1, got B={B}.") + flat = kv_bshd.reshape(S, G * Dh) + boundary = LeftBoundaryExchange.apply(flat, int(d_window), cp_group) + return boundary.reshape(1, int(d_window), G, Dh) + + +class _AllGatherPool(torch.autograd.Function): + """All-gather the per-rank compressed pool into the global, sequence-ordered pool. + + Concatenating in rank order IS sequence order here: this path is BSHD with one + sequence, every rank owns a contiguous block of `S_total / cp_size` rows, and that + block length is a multiple of `ratio`, so compressed group boundaries never straddle + a rank boundary. (Upstream's THD path needs a seq-major -> rank-major remap precisely + because ragged packed sequences break that property; here it is free.) + + Backward is a reduce-scatter, NOT a plain slice. Rank r's pool rows are read by the + queries of every rank at or after r, so each of those ranks holds a partial gradient + for them; slicing this rank's block out of its OWN grad_out would keep only the + contribution from its own queries and silently drop the rest. That error is invisible + in the forward and compounds over training steps -- measured as loss drift growing + 2e-5 -> 4.5e-4 across three steps against a 5e-5 CP noise floor. + """ + + @staticmethod + def forward(ctx, pool_local: torch.Tensor, cp_group): + cp_size = cp_group.size() + ctx.cp_group = cp_group + ctx.cp_rank = cp_group.rank() + ctx.p_local = pool_local.shape[1] + gathered = [torch.empty_like(pool_local) for _ in range(cp_size)] + dist.all_gather(gathered, pool_local.contiguous(), group=cp_group) + return torch.cat(gathered, dim=1) + + @staticmethod + def backward(ctx, grad_out): + grad = grad_out.contiguous() + dist.all_reduce(grad, group=ctx.cp_group) + lo = ctx.cp_rank * ctx.p_local + return grad[:, lo : lo + ctx.p_local].contiguous(), None + + +def build_global_pool(pool_local: torch.Tensor, cp_group) -> torch.Tensor: + """`[B, P_local, D]` -> `[B, P_local * cp_size, D]` in sequence order.""" + return _AllGatherPool.apply(pool_local, cp_group) + + +def compressor_boundary_rows(compress_ratio: int, overlap: bool) -> int: + """Hidden rows this rank must receive from its left neighbour before compressing. + + The compressor pools each window independently EXCEPT in overlap mode (V4 uses it for + ratio 4 / CSA), where window i is stitched with the previous window's second channel + half. At a CP boundary that previous window lives on the left neighbour, so `ratio` + hidden rows have to come across. Non-overlap (ratio 128 / HCA) is purely local. + """ + return int(compress_ratio) if overlap else 0 + + +def compressed_causal_mask( + s_local: int, p_global: int, global_start: int, ratio: int, *, device, dtype +) -> torch.Tensor: + """`[S_local, P_global]` additive mask against GLOBAL query positions. + + Pool slot `s` covers raw tokens `[s*ratio, (s+1)*ratio)`, so a query at global token + `t` may attend to it iff `(s+1)*ratio - 1 <= t`. Under CP the query's global position + is `global_start + local_row`, which is the only change from the non-CP form. + """ + t = torch.arange(s_local, device=device).unsqueeze(1) + int(global_start) + s_end = (torch.arange(p_global, device=device).unsqueeze(0) + 1) * int(ratio) - 1 + return torch.where(s_end <= t, 0.0, float("-inf")).to(dtype) + + +__all__ = [ + "LeftBoundaryExchange", + "get_cp_group", + "exchange_boundary_kv", + "build_global_pool", + "compressor_boundary_rows", + "compressed_causal_mask", +] diff --git a/primus/backends/megatron/core/transformer/indexer.py b/primus/backends/megatron/core/transformer/indexer.py index bc06ca642..d0f970f2f 100644 --- a/primus/backends/megatron/core/transformer/indexer.py +++ b/primus/backends/megatron/core/transformer/indexer.py @@ -46,7 +46,7 @@ import logging import os -from typing import Tuple +from typing import Optional, Tuple import torch import torch.nn as nn @@ -111,6 +111,20 @@ def fake_quantize_fp8_e4m3(x: torch.Tensor) -> torch.Tensor: is_triton_path_enabled as _indexer_tail_triton_enabled, ) + +def _indexer_topk_chunk() -> int: + """Pool-column chunk width for the streaming top-K; 0 (default) = one-shot. + + Off by default so existing recipes keep the exact one-shot ``torch.topk`` numerics. + Set ``PRIMUS_INDEXER_TOPK_CHUNK=32768`` (or smaller) for long context, where the + full ``[B, S, P]`` score row is what does not fit. + """ + try: + return max(0, int(os.environ.get("PRIMUS_INDEXER_TOPK_CHUNK", "0"))) + except ValueError: + return 0 + + # MXFP4 block size (E2M1 data + E8M0 per-32 block scales). _MXFP4_BLOCK = 32 @@ -210,10 +224,27 @@ def __init__( compress_ratio: int = 4, dq_rank: int = None, use_fp8_qk: bool = False, + tp_group=None, ) -> None: super().__init__() self.hidden_size = hidden_size self.index_head_dim = index_head_dim + # ---- P14: shard indexer heads across TP -------------------------------- + # The score is a SUM over heads, so heads can be split across TP ranks and the + # partial sums all-reduced. Two things follow: + # * each rank's local head count is index_n_heads / tp, which for V4-Flash at + # tp=8 is 8 -- inside `_SUPPORTED_H` of the fused Triton scoring kernel, so + # the fused path becomes usable without touching that kernel; and + # * the [B, S, H, P] einsum intermediate shrinks by tp on every rank. + self.tp_group = tp_group + self.tp_size = tp_group.size() if tp_group is not None else 1 + if self.tp_size > 1 and index_n_heads % self.tp_size != 0: + raise ValueError( + f"indexer head sharding requires index_n_heads ({index_n_heads}) divisible " + f"by tensor_model_parallel_size ({self.tp_size})." + ) + self.index_n_heads_global = index_n_heads + index_n_heads = index_n_heads // self.tp_size self.index_n_heads = index_n_heads self.index_topk = index_topk self.compress_ratio = compress_ratio @@ -268,12 +299,27 @@ def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): # ------------------------------------------------------------------ + def _cp_query_offset(self, s_local: int) -> int: + """This rank's first GLOBAL query row. + + Under context parallel the pool is all-gathered (global) while the queries are a + contiguous local slice, so every causal test has to be made against the query's + global position. 0 when CP is off. + """ + from primus.backends.megatron.core.transformer.deepseek_v4_cp import ( + get_cp_group, + ) + + g = get_cp_group() + return 0 if g is None else g.rank() * int(s_local) + def _causal_mask( self, n_queries: int, n_pool: int, device: torch.device, dtype: torch.dtype, + q_offset: Optional[int] = None, ) -> torch.Tensor: """Return ``[n_queries, n_pool]`` mask: 0.0 if pool position ``s`` is allowed for query ``t``, ``-inf`` otherwise. @@ -281,7 +327,12 @@ def _causal_mask( A compressed position ``s`` covers raw tokens ``[s*ratio, (s+1)*ratio)``; a query at raw token ``t`` may attend to ``s`` iff its window end ``(s+1)*ratio - 1 <= t``. + + ``q_offset`` defaults to this rank's CP query offset. The streaming top-K path + overrides it to fold in a pool-column offset -- see :meth:`forward`. """ + if q_offset is None: + q_offset = self._cp_query_offset(n_queries) # The mask depends only on (n_queries, n_pool, compress_ratio, dtype) — all # fixed per run — so cache it instead of rebuilding arange + where every # call. PRIMUS_INDEXER_MASK_CACHE=0 forces the eager rebuild. @@ -290,11 +341,13 @@ def _causal_mask( cache = getattr(self, "_causal_mask_cache", None) if cache is None: cache = self._causal_mask_cache = {} - key = (n_queries, n_pool, device, dtype) + key = (n_queries, n_pool, device, dtype, q_offset) cached = cache.get(key) if cached is not None: return cached - t_idx = torch.arange(n_queries, device=device).unsqueeze(1) # [t, 1] + # + CP offset: under context parallel the pool is global but these queries are + # this rank's slice, so causality is judged on the GLOBAL query position. + t_idx = torch.arange(n_queries, device=device).unsqueeze(1) + q_offset # [t, 1] s_end = (torch.arange(n_pool, device=device).unsqueeze(0) + 1) * self.compress_ratio - 1 # [1, s] allowed = s_end <= t_idx # [t, s] bool mask = torch.where(allowed, 0.0, float("-inf")).to(dtype) @@ -326,7 +379,35 @@ def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: Hd = self.index_head_dim # 1) K^{IComp}: pool hidden via the mini-Compressor → [B, P, Hd] - k_icomp = self.indexer_compressor(hidden) # [B, P, Hd] + # + # Under context parallelism this pool MUST be all-gathered to the global P, for + # two reasons: the top-K indices it produces are used to address the attention's + # own (already global) pool, so a local-P index would name the wrong column; and + # a query could otherwise never select compressed history owned by an earlier + # rank. Mirrors DeepseekV4Attention._build_compressed_pool -- same boundary-row + # rule (overlap mode stitches window i with window i-1, which at a shard edge + # lives on the left neighbour), same rank-order == sequence-order argument. + from primus.backends.megatron.core.transformer.deepseek_v4_cp import ( + get_cp_group, + ) + + cp_group = get_cp_group() + if cp_group is None: + k_icomp = self.indexer_compressor(hidden) # [B, P, Hd] + else: + from primus.backends.megatron.core.transformer.deepseek_v4_cp import ( + build_global_pool, + compressor_boundary_rows, + exchange_boundary_kv, + ) + + nb = compressor_boundary_rows(self.compress_ratio, bool(self.indexer_compressor.overlap)) + if nb > 0: + bnd = exchange_boundary_kv(hidden.reshape(B, S, 1, D), nb, cp_group).reshape(B, nb, D) + k_local = self.indexer_compressor(torch.cat([bnd, hidden], dim=1))[:, 1:] + else: + k_local = self.indexer_compressor(hidden) + k_icomp = build_global_pool(k_local, cp_group) P = k_icomp.shape[1] k_icomp = k_icomp.unsqueeze(2) # [B, P, 1, Hd] @@ -336,7 +417,7 @@ def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: # FP8 (paper / NVIDIA backend.linear) when enabled, else the bf16 nn.Linear. proj = _fp8_linear if _indexer_fp8_proj_enabled() else (lambda lin, x: lin(x)) if self._fuse_qw_proj: - dqw = proj(self.w_dq_w, hidden) # [B, S, dq_rank + H] in one GEMM + dqw = proj(self.w_dq_w, hidden) # [B, S, dq_rank + local H] in one GEMM q_q = dqw[..., : self.dq_rank] # [B, S, dq_rank] w_i = dqw[..., self.dq_rank :] # [B, S, H] else: @@ -365,9 +446,76 @@ def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: q_i = fake_quantize_fp8_e4m3(q_i) k_icomp_2d = fake_quantize_fp8_e4m3(k_icomp_2d) + def _scores_for(k2d: torch.Tensor, p_off: int) -> torch.Tensor: + """Masked scores ``[B, S, Pc]`` for pool columns ``[p_off, p_off + Pc)``. + + The causal test is ``(p_global + 1) * ratio - 1 <= t_global``. Substituting + ``p_global = p_local + p_off`` turns it into + ``(p_local + 1) * ratio - 1 <= t_global - p_off * ratio``, so a pool-column + offset is EXACTLY a shift of the query offset. That is why the streaming path + needs no chunk-awareness inside the Triton kernels -- it reuses ``q_offset``. + """ + q_off = self._cp_query_offset(S) - int(p_off) * self.compress_ratio + Pc = k2d.shape[1] + if _indexer_fp4_enabled(): + dot_c = _fp4_qk_gemm(q_i, k2d) + sc = (F.relu(dot_c) * w_i.unsqueeze(-1)).sum(dim=2) + return sc + self._causal_mask(S, Pc, sc.device, sc.dtype, q_offset=q_off).unsqueeze(0) + if _indexer_triton_full_enabled() and _indexer_triton_full_supported(q_i, k2d, w_i): + return indexer_score_triton( + q_i, + k2d, + w_i, + compress_ratio=self.compress_ratio, + out_dtype=hidden.dtype, + q_offset=q_off, + ) + dot_c = torch.einsum("bshd,bpd->bshp", q_i, k2d) + sc = (F.relu(dot_c) * w_i.unsqueeze(-1)).sum(dim=2) + return sc + self._causal_mask(S, Pc, sc.device, sc.dtype, q_offset=q_off).unsqueeze(0) + # Phase 5: FP4 CSA-indexer QK. Real MXFP4 GEMM for the QK product (paper # §2.3.4/§5.2.1: "QK multiplied entirely in FP4"), then the eager # ReLU/weight/sum tail (w_i + tail stay BF16/FP32 — only the QK is FP4). + # Streaming (chunked) top-K. `scores` is [B, S, P] and P grows with the GLOBAL + # sequence even under CP (the pool is all-gathered), so it is the 1M wall: at + # S_local=131072 / P=262144 it is 64 GiB per rank, while torch.topk's own extra + # peak is only 0.75 GiB -- the tensor, not the selection, is the problem. + # Chunking over P keeps a running top-K and never materialises the full row: + # peak drops to [B, S, chunk] + [B, S, K]. + chunk = _indexer_topk_chunk() + if chunk > 0 and P > chunk: + if _indexer_tail_triton_enabled() and not _indexer_triton_full_enabled(): + raise NotImplementedError( + "Streaming indexer top-K does not support the tail-fused path " + "(PRIMUS_INDEXER_TRITON): indexer_score_post applies the causal mask " + "itself with no pool-column offset hook. Use " + "PRIMUS_INDEXER_TRITON_FULL=1 instead." + ) + run_v = run_i = None + for lo in range(0, P, chunk): + hi = min(lo + chunk, P) + sc = _scores_for(k_icomp_2d[:, lo:hi], lo) + # Same P14 reduction as below, just per chunk: every element is still + # reduced exactly once, so the result is unchanged. + if self.tp_size > 1: + import torch.distributed as _dist + + _dist.all_reduce(sc, group=self.tp_group) + v, i = sc.topk(min(K, hi - lo), dim=-1) + del sc + i = i + lo # chunk-local column -> GLOBAL pool column + if run_v is None: + run_v, run_i = v, i + else: + cv = torch.cat([run_v, v], dim=-1) + ci = torch.cat([run_i, i], dim=-1) + run_v, sel = cv.topk(min(K, cv.shape[-1]), dim=-1) + run_i = torch.gather(ci, -1, sel) + topk_scores, topk_idxs = run_v, run_i + topk_eff = topk_scores.shape[-1] + return self._finalize_topk(topk_idxs, topk_scores, topk_eff, K, B, S) + if _indexer_fp4_enabled(): dot = _fp4_qk_gemm(q_i, k_icomp_2d) # [B, S, H, P], real FP4 matmul relu = F.relu(dot) @@ -381,6 +529,7 @@ def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: w_i, compress_ratio=self.compress_ratio, out_dtype=hidden.dtype, + q_offset=self._cp_query_offset(S), ) else: dot = torch.einsum("bshd,bpd->bshp", q_i, k_icomp_2d) @@ -397,9 +546,21 @@ def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: mask = self._causal_mask(S, P, scores.device, scores.dtype) # [S, P] scores = scores + mask.unsqueeze(0) # [B, S, P] + # P14: with heads sharded, `scores` holds only this rank's partial head sum, so + # reduce before the top-k -- the selection must see the full-head score. The + # causal mask is 0 / -inf and survives the sum unchanged (-inf + finite = -inf), + # so it does not need to be re-applied or divided out. + if self.tp_size > 1: + import torch.distributed as _dist + + _dist.all_reduce(scores, group=self.tp_group) + topk_eff = min(K, P) topk_scores, topk_idxs = scores.topk(topk_eff, dim=-1) # [B, S, topk_eff] + return self._finalize_topk(topk_idxs, topk_scores, topk_eff, K, B, S) + def _finalize_topk(self, topk_idxs, topk_scores, topk_eff, K, B, S): + """Sentinel + pad, shared by the one-shot and streaming top-K paths.""" # 5) Replace selections that are still -inf (i.e. fewer than K valid # pool positions for very early queries) with sentinel ``-1`` so # callers can drop them. diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py index c82543d75..aa1fbcc31 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py @@ -66,10 +66,18 @@ import triton import triton.language as tl -# Indexer scoring runs at V4-Flash widths [B=1, S=4096, P=1024, H=8, -# Hd=128]. H is small enough to be a compile-time constant in the -# kernel; supported values are documented here. -_SUPPORTED_H = (1, 2, 4, 8, 16) +# H is a compile-time constant (the h loop is `tl.static_range`), so every value +# here costs a separate compilation. The list originally stopped at 16 because the +# kernel was written and benchmarked against a claimed "V4-Flash production" width of +# H=8 -- but the released DeepSeek-V4-Flash config and Primus's own +# deepseek_v4_base.yaml both set index_n_heads=64, so at the real width this kernel +# was unreachable and the indexer silently fell back to the eager einsum. That +# fallback materialises [B, S, H, P]: 0.5 GiB at S=4096, but 512 GiB at S=131072, +# which is what made CSA impossible at long context. +# +# 32 and 64 are added so the real width is covered. They only unroll the h loop +# further; the per-iteration tile shapes are unchanged. +_SUPPORTED_H = (1, 2, 4, 8, 16, 32, 64) # --------------------------------------------------------------------------- @@ -88,6 +96,7 @@ def _indexer_score_fwd_kernel( P, H: tl.constexpr, HD: tl.constexpr, + Q_OFFSET, COMPRESS_RATIO: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_P: tl.constexpr, @@ -126,30 +135,41 @@ def _indexer_score_fwd_kernel( s_mask = s_offs < S p_mask = p_offs < P + # int64 offsets: `s_offs` comes from tl.arange (int32) and the row stride into + # SCORES/Q is P and H*HD. At V4-Flash CSA widths P = S/4, so S*P crosses 2**31 at + # S ~= 92682 -- measured: 64k clean, 96k produces NaN. In int32 the product wraps + # negative and the store lands out of bounds, silently. Promote the row index once; + # the column term stays int32 because it is bounded by P. + s_offs64 = s_offs.to(tl.int64) + p_offs64 = p_offs.to(tl.int64) + hd_idx = tl.arange(0, HD) acc = tl.zeros((BLOCK_S, BLOCK_P), dtype=tl.float32) # Unroll over heads (H is small and constexpr). + # k does not depend on h, so load it once instead of once per unrolled head. + # At H=64 that is 63 redundant [BLOCK_P, HD] loads removed from the inner loop. + k_tile = tl.load( + K_PTR + pid_b.to(tl.int64) * P * HD + p_offs64[:, None] * HD + hd_idx[None, :], + mask=p_mask[:, None], + other=0.0, + ).to(tl.float32) + k_tile_t = tl.trans(k_tile) + for h in tl.static_range(0, H): # q [BLOCK_S, HD]: q_i[pid_b, s_offs, h, :] q_tile = tl.load( - Q_PTR + pid_b * S * H * HD + s_offs[:, None] * H * HD + h * HD + hd_idx[None, :], + Q_PTR + pid_b.to(tl.int64) * S * H * HD + s_offs64[:, None] * H * HD + h * HD + hd_idx[None, :], mask=s_mask[:, None], other=0.0, ).to(tl.float32) - # k [BLOCK_P, HD]: k_icomp[pid_b, p_offs, :] - k_tile = tl.load( - K_PTR + pid_b * P * HD + p_offs[:, None] * HD + hd_idx[None, :], - mask=p_mask[:, None], - other=0.0, - ).to(tl.float32) # dot [BLOCK_S, BLOCK_P] = q @ k.T - dot = tl.dot(q_tile, tl.trans(k_tile), out_dtype=tl.float32) + dot = tl.dot(q_tile, k_tile_t, out_dtype=tl.float32) dot = tl.maximum(dot, 0.0) # relu # w [BLOCK_S]: w_i[pid_b, s_offs, h] w_h = tl.load( - W_PTR + pid_b * S * H + s_offs * H + h, + W_PTR + pid_b.to(tl.int64) * S * H + s_offs64 * H + h, mask=s_mask, other=0.0, ).to(tl.float32) @@ -157,14 +177,17 @@ def _indexer_score_fwd_kernel( # Apply causal mask inline. Allowed iff `(p + 1) * cr - 1 <= s`, # i.e. the pool position's window end is no later than the query. - s_arr = s_offs[:, None] + # Under context parallel this rank holds a slice of the queries but the FULL pool, + # so visibility must be judged on the query's GLOBAL position. Q_OFFSET is 0 without + # CP, which reproduces the original expression exactly. + s_arr = s_offs[:, None] + Q_OFFSET p_arr = p_offs[None, :] allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr NEG_INF = -float("inf") acc = tl.where(allowed, acc, NEG_INF) tl.store( - SCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :], + SCORES_PTR + pid_b.to(tl.int64) * S * P + s_offs64[:, None] * P + p_offs64[None, :], acc.to(OUT_DTYPE), mask=s_mask[:, None] & p_mask[None, :], ) @@ -184,6 +207,7 @@ def _indexer_score_bwd_kernel( P, H: tl.constexpr, HD: tl.constexpr, + Q_OFFSET, COMPRESS_RATIO: tl.constexpr, BLOCK_S: tl.constexpr, BLOCK_P: tl.constexpr, @@ -222,41 +246,52 @@ def _indexer_score_bwd_kernel( s_mask = s_offs < S p_mask = p_offs < P + # See the forward kernel: S*P crosses 2**31 at S ~= 92682 for CSA (P = S/4), so the + # int32 row-stride product wraps negative and every load/store lands out of bounds. + s_offs64 = s_offs.to(tl.int64) + p_offs64 = p_offs.to(tl.int64) + hd_idx = tl.arange(0, HD) # Load dscores [BLOCK_S, BLOCK_P] dmasked = tl.load( - DSCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :], + DSCORES_PTR + pid_b.to(tl.int64) * S * P + s_offs64[:, None] * P + p_offs64[None, :], mask=s_mask[:, None] & p_mask[None, :], other=0.0, ).to(tl.float32) # Apply causal mask (zero out invalid positions). - s_arr = s_offs[:, None] + # Under context parallel this rank holds a slice of the queries but the FULL pool, + # so visibility must be judged on the query's GLOBAL position. Q_OFFSET is 0 without + # CP, which reproduces the original expression exactly. + s_arr = s_offs[:, None] + Q_OFFSET p_arr = p_offs[None, :] allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr d_acc = tl.where(allowed, dmasked, 0.0) # Unroll over heads. + # h-invariant; hoisted out of the unrolled loop (see the forward kernel). + k_tile = tl.load( + K_PTR + pid_b.to(tl.int64) * P * HD + p_offs64[:, None] * HD + hd_idx[None, :], + mask=p_mask[:, None], + other=0.0, + ).to(tl.float32) + k_tile_t = tl.trans(k_tile) + for h in tl.static_range(0, H): - # Reload q, k for this head (FlashAttention-style recompute). + # Reload q for this head (FlashAttention-style recompute). q_tile = tl.load( - Q_PTR + pid_b * S * H * HD + s_offs[:, None] * H * HD + h * HD + hd_idx[None, :], + Q_PTR + pid_b.to(tl.int64) * S * H * HD + s_offs64[:, None] * H * HD + h * HD + hd_idx[None, :], mask=s_mask[:, None], other=0.0, ).to(tl.float32) - k_tile = tl.load( - K_PTR + pid_b * P * HD + p_offs[:, None] * HD + hd_idx[None, :], - mask=p_mask[:, None], - other=0.0, - ).to(tl.float32) w_h = tl.load( - W_PTR + pid_b * S * H + s_offs * H + h, + W_PTR + pid_b.to(tl.int64) * S * H + s_offs64 * H + h, mask=s_mask, other=0.0, ).to(tl.float32) # Recompute dot, relu, relu_dot. - dot = tl.dot(q_tile, tl.trans(k_tile), out_dtype=tl.float32) + dot = tl.dot(q_tile, k_tile_t, out_dtype=tl.float32) relu_dot = tl.maximum(dot, 0.0) # [BLOCK_S, BLOCK_P] relu_mask = dot > 0.0 @@ -276,17 +311,17 @@ def _indexer_score_bwd_kernel( # Stores with atomic_add since multiple p_tiles/s_tiles touch # the same locations. tl.atomic_add( - DQ_PTR + pid_b * S * H * HD + s_offs[:, None] * H * HD + h * HD + hd_idx[None, :], + DQ_PTR + pid_b.to(tl.int64) * S * H * HD + s_offs64[:, None] * H * HD + h * HD + hd_idx[None, :], d_q, mask=s_mask[:, None], ) tl.atomic_add( - DK_PTR + pid_b * P * HD + p_offs[:, None] * HD + hd_idx[None, :], + DK_PTR + pid_b.to(tl.int64) * P * HD + p_offs64[:, None] * HD + hd_idx[None, :], d_k, mask=p_mask[:, None], ) tl.atomic_add( - DW_PTR + pid_b * S * H + s_offs * H + h, + DW_PTR + pid_b.to(tl.int64) * S * H + s_offs64 * H + h, dw_h, mask=s_mask, ) @@ -346,6 +381,7 @@ def forward( # type: ignore[override] w_i: torch.Tensor, compress_ratio: int, out_dtype: torch.dtype, + q_offset: int = 0, ): if q_i.dim() != 4: raise ValueError(f"q_i must be [B, S, H, Hd], got shape {tuple(q_i.shape)}") @@ -389,6 +425,7 @@ def forward( # type: ignore[override] P, H=H, HD=HD, + Q_OFFSET=int(q_offset), COMPRESS_RATIO=int(compress_ratio), BLOCK_S=block_s, BLOCK_P=block_p, @@ -402,6 +439,7 @@ def forward( # type: ignore[override] ctx.save_for_backward(q_c, k_c, w_c) ctx.compress_ratio = int(compress_ratio) + ctx.q_offset = int(q_offset) ctx.shape = (B, S, P, H, HD) ctx.in_dtypes = (q_i.dtype, k_icomp.dtype, w_i.dtype) return scores @@ -411,6 +449,7 @@ def backward(ctx, d_scores): # type: ignore[override] q_c, k_c, w_c = ctx.saved_tensors B, S, P, H, HD = ctx.shape compress_ratio = ctx.compress_ratio + q_offset = ctx.q_offset q_dtype, k_dtype, w_dtype = ctx.in_dtypes d_scores = d_scores.contiguous() @@ -435,12 +474,13 @@ def backward(ctx, d_scores): # type: ignore[override] P, H=H, HD=HD, + Q_OFFSET=int(q_offset), COMPRESS_RATIO=int(compress_ratio), BLOCK_S=block_s, BLOCK_P=block_p, ) - return d_q_fp32.to(q_dtype), d_k_fp32.to(k_dtype), d_w_fp32.to(w_dtype), None, None + return d_q_fp32.to(q_dtype), d_k_fp32.to(k_dtype), d_w_fp32.to(w_dtype), None, None, None # --------------------------------------------------------------------------- @@ -486,13 +526,14 @@ def indexer_score_triton( *, compress_ratio: int, out_dtype: torch.dtype, + q_offset: int = 0, ) -> torch.Tensor: """Compute Indexer scores via the fused Triton kernel. Returns ``scores [B, S, P]`` of dtype ``out_dtype``. Masked positions hold ``-inf``. """ - return IndexerScoreFn.apply(q_i, k_icomp, w_i, compress_ratio, out_dtype) + return IndexerScoreFn.apply(q_i, k_icomp, w_i, compress_ratio, out_dtype, q_offset) __all__ = [ diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py index f1f4c61e3..543f77a30 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py @@ -64,7 +64,9 @@ # H is small and known per call site -- V4-Flash uses H=8. Same # supported set as P38. -_SUPPORTED_H = (1, 2, 4, 8, 16) +# See indexer_score.py: the real DeepSeek-V4 index_n_heads is 64, not the 8 this +# kernel was written against, so 32 / 64 are needed for the path to be reachable. +_SUPPORTED_H = (1, 2, 4, 8, 16, 32, 64) # --------------------------------------------------------------------------- diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py index 1256d7ade..50f148da0 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py @@ -72,7 +72,29 @@ def sparse_mla_bwd_v4_triton(q, kv, o, do, topk_indices, lse, attn_sink=None, kv R_CHUNK = min(topk, 1536) else: R_CHUNK = min(256, topk) - BH_DQ, TK_DQ = 64, 16 + # The defaults above are tuned for SHORT sequences, where the per-chunk buffers are + # small and the dq read-modify-write traffic dominates. At long context that trade + # inverts hard: the buffers below scale with total_tokens * R_CHUNK, so at + # S_local = 131072 (1M context, CP=8) R_CHUNK=256 makes `interm` alone + # 131072 * 256 * 576 * 2 B = 36 GiB, plus 8 GiB for chunk_dS/chunk_P -- a 44 GiB + # workspace spent to avoid some dq reloads. PRIMUS_DSA_BWD_R_CHUNK trades that back. + # Chunking is a partition of the same computation, so any value is numerically + # equivalent; note that R_CHUNK % 128 != 0 also disables the fused dKV path below. + _r_override = os.environ.get("PRIMUS_DSA_BWD_R_CHUNK", "") + if _r_override: + R_CHUNK = max(1, min(int(_r_override), topk)) + # BH_DQ x D_V is the dominant LDS term of _bwd_chunk_dq_store_ds. With the V4 + # latent head_dim of 512, BH_DQ=64 needs 64*512*2 = 65536 B for the Q tile + # alone -- exactly the whole 64 KB LDS budget of gfx942/CDNA3 -- and the kernel + # asks for 73728 B total, so it fails to compile there. gfx950/CDNA4 has 160 KB + # and is unaffected. Expose the head-block so CDNA3 can halve it (32 -> 32 KB + # Q tile). Default is unchanged. + # Triton multi-buffers LDS operand tiles across pipeline stages; on gfx942 + # (64 KB LDS vs gfx950's 160 KB) the default staging pushes this kernel to + # 73728 B. num_stages=1 disables the double-buffering. 0 = leave to Triton. + _BWD_NUM_STAGES = int(os.environ.get("PRIMUS_DSA_BWD_NUM_STAGES", "0")) or None + BH_DQ = int(os.environ.get("PRIMUS_DSA_BWD_BLOCK_H", "64")) + TK_DQ = int(os.environ.get("PRIMUS_DSA_BWD_TILE_K", "16")) # dKV-intermediate tiling. The default (BH_DKV=32, TK_DKV=64) is best for high # head counts (H>=128) and for chunk widths that are not 128-aligned. For low # head counts (H<=64) with a 128-aligned chunk, a wider TILE_K=128 over a single @@ -176,6 +198,7 @@ def sparse_mla_bwd_v4_triton(q, kv, o, do, topk_indices, lse, attn_sink=None, kv IS_FIRST_CHUNK=is_first, num_warps=4, waves_per_eu=1, + num_stages=_BWD_NUM_STAGES, ) # The wide dKV kernel MUST compile with ping-pong off to fit LDS, even @@ -205,6 +228,7 @@ def sparse_mla_bwd_v4_triton(q, kv, o, do, topk_indices, lse, attn_sink=None, kv D_ROPE=rope_rank, HAS_ROPE=HAS_ROPE, num_warps=4, + num_stages=_BWD_NUM_STAGES, ) inv_ptr, inv_data = all_csr[chunk_idx] @@ -219,6 +243,7 @@ def sparse_mla_bwd_v4_triton(q, kv, o, do, topk_indices, lse, attn_sink=None, kv D_ROPE=rope_rank, HAS_ROPE=HAS_ROPE, num_warps=4, + num_stages=_BWD_NUM_STAGES, ) d_sink = None diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py index f31815887..a5c0d5302 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py @@ -35,6 +35,47 @@ _ROPE_PAD = 64 # dummy separate-rope block (zeros); the kernels need D_ROPE > 0 +def _rope_pad_q(q_bh: torch.Tensor) -> torch.Tensor: + """``[B, H, S, D]`` -> ``[B*S, H, D + _ROPE_PAD]``, pad block zeroed. + + Allocates the padded buffer and strided-copies the query into its first ``D`` + columns, rather than ``cat([q_bh.permute(...).reshape(...), zeros], -1)``. The + permuted view is non-contiguous, so that ``.reshape`` materialises a FULL extra + copy before the cat allocates the result on top of it: the cat form peaks at + ~3.2x the query size where this one peaks at ~2.1x. At 1M with CP=8 the query is + ``[1, 64, 131072, 512]`` = 8 GiB, and that difference (~9 GiB) is exactly what + put the run over a 192 GiB card -- it OOM'd here with 176.11 GiB already live. + """ + B, H, S, D = q_bh.shape + out = q_bh.new_zeros(B, S, H, D + _ROPE_PAD) + out[..., :D] = q_bh.permute(0, 2, 1, 3) # strided copy, no contiguous temp + return out.reshape(B * S, H, D + _ROPE_PAD) # contiguous -> free view + + +def _bshd_slice_to_bhsd(dq_g: torch.Tensor, B: int, S: int, H: int, D: int) -> torch.Tensor: + """Drop the rope pad off ``[B*S, H, D+_ROPE_PAD]`` and return ``[B, H, S, D]``. + + The inverse of :func:`_rope_pad_q`, and it has to dodge the same trap. Written the + obvious way -- ``dq_g[:, :, :D].reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous()`` + -- it costs TWO full copies: the padded slice is non-contiguous so ``.reshape`` + materialises one, and ``.contiguous()`` after the permute materialises another. At 1M + with CP=8 that is 8.00 GiB each, in the backward pass, which is where the run died. + Allocating the destination and doing one strided copy through a permuted view of it + leaves a single 8.00 GiB allocation. + """ + out = dq_g.new_empty(B, H, S, D) + out.permute(0, 2, 1, 3).copy_(dq_g[:, :, :D].unflatten(0, (B, S))) + return out + + +def _rope_pad_kv(kv512: torch.Tensor) -> torch.Tensor: + """``[N, 1, D]`` -> ``[N, 1, D + _ROPE_PAD]``, pad block zeroed. See :func:`_rope_pad_q`.""" + N, G, D = kv512.shape + out = kv512.new_zeros(N, G, D + _ROPE_PAD) + out[..., :D] = kv512 + return out + + def _pad_topk_64(topk: torch.Tensor) -> torch.Tensor: """Pad the topk width to a multiple of 64 with -1 so a backend whose dKV tiling is 64-wide (e.g. gluon) stays valid (HCA 128+32=160 -> 192).""" @@ -47,25 +88,52 @@ def _pad_topk_64(topk: torch.Tensor) -> torch.Tensor: return topk.contiguous() -def _build_csa_topk(topk_idxs: torch.Tensor, S: int, P: int, W: int) -> torch.Tensor: - """Flat topk [B*S, W+K] over the per-batch [local ++ pool] buffer. +def _build_csa_topk( + topk_idxs: torch.Tensor, + S: int, + Skv: int, + P: int, + W: int, + cp_dwindow: int = 0, + cp_global_start: int = 0, +) -> torch.Tensor: + """Flat topk [B*S, W+K] over the per-batch [raw ++ pool] buffer. ``topk_idxs`` [B, S, K] holds pool indices in [0, P) (or -1). Batch ``b`` - occupies rows ``[b*(S+P) : (b+1)*(S+P))`` (local 0..S-1, pool S..S+P-1). + occupies rows ``[b*(Skv+P) : (b+1)*(Skv+P))`` (raw 0..Skv-1, pool Skv..Skv+P-1). + + ``S`` is the QUERY count, ``Skv`` the raw-token KV count. They differ only under + context parallelism, where the raw buffer is ``[boundary ++ local]`` and so + ``Skv == cp_dwindow + S``: the sliding window is then validated against GLOBAL + positions (a token must not attend before the sequence start) but indexed in LOCAL + buffer coordinates. The pool is already global, so ``topk_idxs`` needs no shift. + With ``cp_dwindow == cp_global_start == 0`` and ``Skv == S`` this is byte-identical + to the non-CP form. """ B, _, K = topk_idxs.shape device = topk_idxs.device - base = (torch.arange(B, device=device) * (S + P)).view(B, 1, 1) - - win_pos = torch.arange(S, device=device).view(S, 1) - W + 1 + torch.arange(W, device=device).view(1, W) - win_valid = win_pos >= 0 + # int32 throughout -- see the note in _V4SparseMLAAttnFn.forward. The result is cast + # to int32 regardless, so int64 intermediates only double the transient peak. + idx_dtype = torch.int32 + neg1 = torch.tensor(-1, device=device, dtype=idx_dtype) + base = (torch.arange(B, device=device, dtype=idx_dtype) * (Skv + P)).view(B, 1, 1) + + gpos = ( + torch.arange(S, device=device, dtype=idx_dtype).view(S, 1) + + int(cp_global_start) + - W + + 1 + + torch.arange(W, device=device, dtype=idx_dtype).view(1, W) + ) + win_valid = gpos >= 0 + win_pos = gpos - int(cp_global_start) + int(cp_dwindow) win_idx = base + win_pos.view(1, S, W) - win_idx = torch.where(win_valid.view(1, S, W), win_idx, torch.full_like(win_idx, -1)) + win_idx = torch.where(win_valid.view(1, S, W), win_idx, neg1) pool_valid = topk_idxs >= 0 - pool_idx = torch.where(pool_valid, base + S + topk_idxs, torch.full_like(topk_idxs, -1)) + pool_idx = torch.where(pool_valid, base + Skv + topk_idxs.to(idx_dtype), neg1) - return torch.cat([win_idx, pool_idx], dim=2).reshape(B * S, W + K).to(torch.int32).contiguous() + return torch.cat([win_idx, pool_idx], dim=2).reshape(B * S, W + K).contiguous() class _V4SparseMLACSAFn(torch.autograd.Function): @@ -75,47 +143,61 @@ class _V4SparseMLACSAFn(torch.autograd.Function): def forward( # type: ignore[override] ctx, q_bh: torch.Tensor, # [B, H, S, D] - k_local_bh: torch.Tensor, # [B, H, S, D] (single MQA latent, head-broadcast) - v_local_bh: torch.Tensor, # [B, H, S, D] (== k_local in V4) + k_local_bh: torch.Tensor, # [B, H, S, D] broadcast, or [B, Skv, 1, D] if k_is_latent + v_local_bh: Optional[torch.Tensor], # == k_local in V4; None when k_is_latent pool: torch.Tensor, # [B, P, D] topk_idxs: torch.Tensor, # [B, S, K] pool indices, -1 = invalid sink: Optional[torch.Tensor], # [H] fp32 or None swa_window: int, scale: float, + cp_dwindow: int, + cp_global_start: int, + k_is_latent: bool, fwd_fn: Callable, bwd_fn: Callable, ) -> torch.Tensor: B, H, S, D = q_bh.shape + # Under CP the raw KV buffer is [boundary ++ local], so it is LONGER than the + # query count; every KV-side extent below is Skv, not S. + # `k_is_latent` means the caller handed us the un-broadcast [B, Skv, 1, D] latent + # rather than its head-broadcast view -- see _V4SparseMLAAttnFn.forward for why + # (the broadcast's gradient would be a [B, H, Skv, D] buffer that is zero except + # at head 0: 8.01 GiB at 1M with CP=8, which is what the backward died on). + ctx.k_is_latent = bool(k_is_latent) + Skv = k_local_bh.shape[1] if k_is_latent else k_local_bh.shape[2] P = pool.shape[1] W = int(swa_window) assert q_bh.dtype == torch.bfloat16, "sparse-MLA adapter requires bf16" assert W > 0, "sparse-MLA adapter requires swa_window > 0" - latent = k_local_bh[:, 0, :, :] # [B, S, D] + latent = k_local_bh[:, :, 0, :] if k_is_latent else k_local_bh[:, 0, :, :] # [B, Skv, D] - z_q = torch.zeros(B * S, H, _ROPE_PAD, device=q_bh.device, dtype=q_bh.dtype) - q_g = torch.cat([q_bh.permute(0, 2, 1, 3).reshape(B * S, H, D), z_q], dim=-1).contiguous() + q_g = _rope_pad_q(q_bh) + kv_g = _rope_pad_kv(torch.cat([latent, pool], dim=1).reshape(B * (Skv + P), 1, D)) - kv512 = torch.cat([latent, pool], dim=1).reshape(B * (S + P), 1, D) - z_kv = torch.zeros(B * (S + P), 1, _ROPE_PAD, device=q_bh.device, dtype=q_bh.dtype) - kv_g = torch.cat([kv512, z_kv], dim=-1).contiguous() - - topk_g = _pad_topk_64(_build_csa_topk(topk_idxs, S, P, W)) + topk_g = _pad_topk_64(_build_csa_topk(topk_idxs, S, Skv, P, W, int(cp_dwindow), int(cp_global_start))) sink_arg = sink.float().contiguous() if sink is not None else None o_g, lse = fwd_fn(q_g, kv_g, topk_g, attn_sink=sink_arg, kv_lora_rank=D, scale=float(scale)) ctx.save_for_backward(q_g, kv_g, o_g, lse, topk_g, sink_arg if sink is not None else q_g.new_empty(0)) - ctx.shapes = (B, H, S, D, P, W) + ctx.shapes = (B, H, S, Skv, D, P, W) ctx.scale = float(scale) ctx.sink_was_none = sink is None ctx.bwd_fn = bwd_fn - return o_g.reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous() + # Return a VIEW, not a copy. The kernel's `o_g` is contiguous [B*S, H, D], i.e. + # already BSHD; every caller immediately does `out_bh.transpose(1, 2).contiguous()` + # to get back to BSHD, so a `.contiguous()` here would materialise BHSD only for + # the caller to materialise BSHD again -- two full copies of the output that + # cancel out. Without it, the caller's transpose restores the original strides and + # its `.contiguous()` becomes a no-op. At 1M with CP=8 each copy is 8.00 GiB, and + # this is the allocation the run actually died on. + return o_g.reshape(B, S, H, D).permute(0, 2, 1, 3) @staticmethod def backward(ctx, grad_o_bh: torch.Tensor): # type: ignore[override] q_g, kv_g, o_g, lse, topk_g, sink_saved = ctx.saved_tensors - B, H, S, D, P, W = ctx.shapes + B, H, S, Skv, D, P, W = ctx.shapes sink_arg = None if ctx.sink_was_none else sink_saved grad_o_g = grad_o_bh.permute(0, 2, 1, 3).reshape(B * S, H, D).contiguous() @@ -123,13 +205,18 @@ def backward(ctx, grad_o_bh: torch.Tensor): # type: ignore[override] q_g, kv_g, o_g, grad_o_g, topk_g, lse, attn_sink=sink_arg, kv_lora_rank=D, scale=ctx.scale ) - dq_bh = dq_g[:, :, :D].reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous() - dkv512 = dkv_g[:, 0, :D].reshape(B, S + P, D) - dlatent = dkv512[:, :S, :] - dpool = dkv512[:, S:, :].contiguous() + dq_bh = _bshd_slice_to_bhsd(dq_g, B, S, H, D) + dkv512 = dkv_g[:, 0, :D].reshape(B, Skv + P, D) + dlatent = dkv512[:, :Skv, :] + dpool = dkv512[:, Skv:, :].contiguous() - dk_local = torch.zeros(B, H, S, D, device=dq_bh.device, dtype=dq_bh.dtype) - dk_local[:, 0, :, :] = dlatent.to(dq_bh.dtype) + if ctx.k_is_latent: + # Gradient has the latent's own [B, Skv, 1, D] shape -- 128 MiB at 1M/CP=8 + # instead of an 8.01 GiB buffer that is 63/64 zeros. Nothing to memset. + dk_local = dlatent.to(dq_bh.dtype).unsqueeze(2) + else: + dk_local = torch.zeros(B, H, Skv, D, device=dq_bh.device, dtype=dq_bh.dtype) + dk_local[:, 0, :, :] = dlatent.to(dq_bh.dtype) # V4 is single-latent (K = V = kv): the kernel returns one combined # ``dkv`` which we route entirely through ``dk_local``. The V branch # gradient is structurally zero, so we return ``None`` for it instead @@ -144,8 +231,23 @@ def backward(ctx, grad_o_bh: torch.Tensor): # type: ignore[override] if not ctx.sink_was_none and dsink is not None: dsink_out = dsink.to(sink_saved.dtype) - # forward args: (q, k_local, v_local, pool, topk_idxs, sink, swa_window, scale, fwd_fn, bwd_fn) - return dq_bh, dk_local, dv_local, dpool.to(dq_bh.dtype), None, dsink_out, None, None, None, None + # forward args: (q, k_local, v_local, pool, topk_idxs, sink, swa_window, scale, + # cp_dwindow, cp_global_start, k_is_latent, fwd_fn, bwd_fn) + return ( + dq_bh, + dk_local, + dv_local, + dpool.to(dq_bh.dtype), + None, + dsink_out, + None, + None, + None, + None, + None, + None, + None, + ) class _V4SparseMLAAttnFn(torch.autograd.Function): @@ -155,48 +257,73 @@ class _V4SparseMLAAttnFn(torch.autograd.Function): def forward( # type: ignore[override] ctx, q_bh: torch.Tensor, # [B, H, S, D] - k_bh: torch.Tensor, # [B, H, Skv, D] (Skv = S for cr=0; S+P for HCA) - v_bh: torch.Tensor, # [B, H, Skv, D] (== k_bh in V4) + k_bh: torch.Tensor, # [B, H, Skv, D], or the [B, Skv, 1, D] latent if k_is_latent + v_bh: Optional[torch.Tensor], # == k_bh in V4; None when k_is_latent sink: Optional[torch.Tensor], swa_window: int, additive_mask: Optional[torch.Tensor], # [S, P] pool-only mask (HCA) or None scale: float, hca_local_seqlen: int, + cp_dwindow: int, + cp_global_start: int, + k_is_latent: bool, fwd_fn: Callable, bwd_fn: Callable, ) -> torch.Tensor: B, H, S, D = q_bh.shape - Skv = k_bh.shape[2] + # V4 is single-latent MQA: this kernel reads ONE key row per position, and the H + # query heads all use it. Callers may therefore hand us the raw [B, Skv, 1, D] + # latent (k_is_latent) instead of its [B, H, Skv, D] head-broadcast view. That is + # strictly better: the broadcast view is free going forward, but autograd would + # make us return a [B, H, Skv, D] gradient of which 63/64 is zeros -- 8.5 GiB at + # 1M with CP=8, which is what the backward pass died on. + ctx.k_is_latent = bool(k_is_latent) + Skv = k_bh.shape[1] if k_is_latent else k_bh.shape[2] W = int(swa_window) assert q_bh.dtype == torch.bfloat16, "sparse-MLA adapter requires bf16" assert W > 0, "sparse-MLA adapter requires swa_window > 0" device = q_bh.device - base = (torch.arange(B, device=device) * Skv).view(B, 1, 1) - win_pos = ( - torch.arange(S, device=device).view(S, 1) - W + 1 + torch.arange(W, device=device).view(1, W) + # int32 throughout: this index matrix is cast to int32 by _pad_topk_64 anyway, and + # the largest value is B*(Skv+P) -- 139264 at 1M with CP=8, nowhere near the int32 + # range. Building it in torch's default int64 doubled every intermediate below, + # which at 1M is ~8.6 GiB apiece. + idx_dtype = torch.int32 + neg1 = torch.tensor(-1, device=device, dtype=idx_dtype) + base = (torch.arange(B, device=device, dtype=idx_dtype) * Skv).view(B, 1, 1) + # Context parallel: this rank owns global rows [cp_global_start, +S), and its KV + # buffer is [boundary ++ local] with `cp_dwindow` boundary rows received from the + # left neighbour. So the window is validated against GLOBAL positions (a token must + # not attend before the sequence start) but indexed in LOCAL buffer coordinates. + # With cp_dwindow == cp_global_start == 0 this is byte-identical to the non-CP form. + gpos = ( + torch.arange(S, device=device, dtype=idx_dtype).view(S, 1) + + int(cp_global_start) + - W + + 1 + + torch.arange(W, device=device, dtype=idx_dtype).view(1, W) ) - win_valid = win_pos >= 0 + win_valid = gpos >= 0 + win_pos = gpos - int(cp_global_start) + int(cp_dwindow) win_idx = base + win_pos.view(1, S, W) - win_idx = torch.where(win_valid.view(1, S, W), win_idx, torch.full_like(win_idx, -1)) + win_idx = torch.where(win_valid.view(1, S, W), win_idx, neg1) if hca_local_seqlen > 0 and additive_mask is not None: P = Skv - int(hca_local_seqlen) vis = (additive_mask == 0).view(1, S, P) - ps = torch.arange(P, device=device).view(1, 1, P) - pool_idx = torch.where( - vis, base + hca_local_seqlen + ps, torch.full((B, S, P), -1, device=device) - ) + ps = torch.arange(P, device=device, dtype=idx_dtype).view(1, 1, P) + # `neg1` is a 0-dim tensor, not torch.full((B, S, P), -1): the full form + # materialised an entire [B, S, P] constant just to be the `where` else-branch + # -- 8.6 GiB at 1M, thrown away immediately. + pool_idx = torch.where(vis, base + hca_local_seqlen + ps, neg1) topk = torch.cat([win_idx, pool_idx], dim=2) else: topk = win_idx - topk_g = _pad_topk_64(topk.reshape(B * S, -1).to(torch.int32)) + topk_g = _pad_topk_64(topk.reshape(B * S, -1)) - z_q = torch.zeros(B * S, H, _ROPE_PAD, device=device, dtype=q_bh.dtype) - q_g = torch.cat([q_bh.permute(0, 2, 1, 3).reshape(B * S, H, D), z_q], dim=-1).contiguous() - kv512 = k_bh[:, 0, :, :].reshape(B * Skv, 1, D) - z_kv = torch.zeros(B * Skv, 1, _ROPE_PAD, device=device, dtype=q_bh.dtype) - kv_g = torch.cat([kv512, z_kv], dim=-1).contiguous() + q_g = _rope_pad_q(q_bh) + latent = k_bh[:, :, 0, :] if k_is_latent else k_bh[:, 0, :, :] # [B, Skv, D] + kv_g = _rope_pad_kv(latent.reshape(B * Skv, 1, D)) sink_arg = sink.float().contiguous() if sink is not None else None o_g, lse = fwd_fn(q_g, kv_g, topk_g, attn_sink=sink_arg, kv_lora_rank=D, scale=float(scale)) @@ -206,7 +333,14 @@ def forward( # type: ignore[override] ctx.scale = float(scale) ctx.sink_was_none = sink is None ctx.bwd_fn = bwd_fn - return o_g.reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous() + # Return a VIEW, not a copy. The kernel's `o_g` is contiguous [B*S, H, D], i.e. + # already BSHD; every caller immediately does `out_bh.transpose(1, 2).contiguous()` + # to get back to BSHD, so a `.contiguous()` here would materialise BHSD only for + # the caller to materialise BSHD again -- two full copies of the output that + # cancel out. Without it, the caller's transpose restores the original strides and + # its `.contiguous()` becomes a no-op. At 1M with CP=8 each copy is 8.00 GiB, and + # this is the allocation the run actually died on. + return o_g.reshape(B, S, H, D).permute(0, 2, 1, 3) @staticmethod def backward(ctx, grad_o_bh: torch.Tensor): # type: ignore[override] @@ -219,21 +353,41 @@ def backward(ctx, grad_o_bh: torch.Tensor): # type: ignore[override] q_g, kv_g, o_g, grad_o_g, topk_g, lse, attn_sink=sink_arg, kv_lora_rank=D, scale=ctx.scale ) - dq_bh = dq_g[:, :, :D].reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous() + dq_bh = _bshd_slice_to_bhsd(dq_g, B, S, H, D) dkv = dkv_g[:, 0, :D].reshape(B, Skv, D) - dk_bh = torch.zeros(B, H, Skv, D, device=dq_bh.device, dtype=dq_bh.dtype) - dk_bh[:, 0, :, :] = dkv.to(dq_bh.dtype) - # Single-latent (K = V): route the combined ``dkv`` through ``dk_bh``; - # the V branch gradient is structurally zero, so return ``None`` and - # skip the big [B, H, Skv, D] memset (see the CSA branch note). + # Single-latent (K = V): route the combined ``dkv`` through the K slot; the V + # branch gradient is structurally zero, so return ``None`` for it. + if ctx.k_is_latent: + # The caller passed the [B, Skv, 1, D] latent, so the gradient has that shape + # too -- 136 MiB at 1M/CP=8 instead of the 8.5 GiB [B, H, Skv, D] buffer below, + # of which only head 0 was ever nonzero. No zeros to allocate or memset. + dk_bh = dkv.to(dq_bh.dtype).unsqueeze(2) + else: + dk_bh = torch.zeros(B, H, Skv, D, device=dq_bh.device, dtype=dq_bh.dtype) + dk_bh[:, 0, :, :] = dkv.to(dq_bh.dtype) dv_bh = None dsink_out = None if not ctx.sink_was_none and dsink is not None: dsink_out = dsink.to(sink_saved.dtype) - # forward args: (q, k, v, sink, swa_window, additive_mask, scale, hca_local_seqlen, fwd_fn, bwd_fn) - return dq_bh, dk_bh, dv_bh, dsink_out, None, None, None, None, None, None + # forward args: (q, k, v, sink, swa_window, additive_mask, scale, hca_local_seqlen, + # cp_dwindow, cp_global_start, k_is_latent, fwd_fn, bwd_fn) + return ( + dq_bh, + dk_bh, + dv_bh, + dsink_out, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) def make_csa_from_pool(fwd_fn: Callable, bwd_fn: Callable) -> Callable: @@ -251,6 +405,9 @@ def _csa_from_pool( attn_dropout, training, scale, + cp_dwindow=0, + cp_global_start=0, + k_is_latent=False, ): if attn_dropout > 0.0 and training: raise NotImplementedError( @@ -258,7 +415,19 @@ def _csa_from_pool( f"(V4 trains with attn_dropout=0). Got attn_dropout={attn_dropout}, training={training}." ) return _V4SparseMLACSAFn.apply( - q_bh, k_local_bh, v_local_bh, pool, topk_idxs, sink, int(swa_window), float(scale), fwd_fn, bwd_fn + q_bh, + k_local_bh, + v_local_bh, + pool, + topk_idxs, + sink, + int(swa_window), + float(scale), + int(cp_dwindow), + int(cp_global_start), + bool(k_is_latent), + fwd_fn, + bwd_fn, ) return _csa_from_pool @@ -279,6 +448,9 @@ def _attention( training, scale, hca_local_seqlen=0, + cp_dwindow=0, + cp_global_start=0, + k_is_latent=False, ): if attn_dropout > 0.0 and training: raise NotImplementedError( @@ -286,7 +458,19 @@ def _attention( f"(V4 trains with attn_dropout=0). Got attn_dropout={attn_dropout}, training={training}." ) return _V4SparseMLAAttnFn.apply( - q, k, v, sink, int(swa_window), additive_mask, float(scale), int(hca_local_seqlen), fwd_fn, bwd_fn + q, + k, + v, + sink, + int(swa_window), + additive_mask, + float(scale), + int(hca_local_seqlen), + int(cp_dwindow), + int(cp_global_start), + bool(k_is_latent), + fwd_fn, + bwd_fn, ) return _attention diff --git a/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py b/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py index 9fb0c2de5..084ab63e5 100644 --- a/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py +++ b/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py @@ -226,6 +226,15 @@ def patched_schedule(*args, **kwargs): getattr(get_args(ctx), "model_type", None) == "deepseek_v4" and int(getattr(get_args(ctx), "num_hash_layers", 0) or 0) > 0 and int(getattr(get_args(ctx), "pipeline_model_parallel_size", 1) or 1) > 1 + # SFT does NOT need this. It exists only to compensate for pretrain_gpt.get_batch + # gating tokens to first/last PP stages. The SFT dataloader is built on every PP + # stage (is_distributed=True, no per-stage gating) and yields the full dict batch + # (with input_ids) on each, so every stage's hash router already has its tokens. + # Worse, in SFT this patch calls pretrain_gpt.get_batch -> get_batch_on_this_tp_rank, + # which reads data["tokens"] while the SFT batch has "input_ids" -> KeyError, and it + # would also double-consume the shared data_iterator. So it must be a no-op for SFT. + and getattr(get_args(ctx), "stage", None) != "sft" + and not getattr(get_args(ctx), "sft", False) ), # Ordered after pp_dump_data so its schedule_wrapper does not double-wrap; # see ``pp_dump_data_patches.py`` for the priority=100 anchor. diff --git a/primus/backends/megatron/sft/forward_step.py b/primus/backends/megatron/sft/forward_step.py index bf71a08b9..32e2d3e17 100644 --- a/primus/backends/megatron/sft/forward_step.py +++ b/primus/backends/megatron/sft/forward_step.py @@ -169,6 +169,13 @@ def _move_to_runtime_device(tensor: torch.Tensor) -> torch.Tensor: return tensor +def _sft_get_cp_group(): + """CP process group, or None when context parallelism is off.""" + from primus.backends.megatron.core.transformer.deepseek_v4_cp import get_cp_group + + return get_cp_group() + + def _empty_loss_result(device: torch.device | None = None) -> Tuple[torch.Tensor, torch.Tensor, dict]: """Build a no-op loss tuple with the expected Megatron shape.""" if device is None: @@ -336,6 +343,26 @@ def forward_step(data_iterator: Iterator, model, return_schedule_plan: bool = Fa position_ids = torch.arange(seq_len, dtype=torch.long, device=tokens.device) position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) + # ---- context parallel: contiguous sequence sharding ------------------- + # V4's dense branch consumes a CONTIGUOUS shard plus a boundary window (see + # deepseek_v4_cp.py), not Megatron's load-balanced 2-chunk ring layout, so we slice + # directly instead of calling get_batch_on_this_cp_rank. position_ids stay GLOBAL -- + # slicing the 0..seq_len-1 stream is exactly what RoPE and the window-validity mask + # need on this rank. + cp_group = _sft_get_cp_group() + if cp_group is not None: + cp_size, cp_rank = cp_group.size(), cp_group.rank() + if seq_len % cp_size != 0: + raise RuntimeError( + f"seq_length={seq_len} must be divisible by context_parallel_size={cp_size}." + ) + l_local = seq_len // cp_size + sl = slice(cp_rank * l_local, (cp_rank + 1) * l_local) + tokens = tokens[:, sl].contiguous() + labels = labels[:, sl].contiguous() + loss_mask = loss_mask[:, sl].contiguous() + position_ids = position_ids[:, sl].contiguous() + # attention_mask: None for causal mask (standard GPT autoregressive) attention_mask = None @@ -374,14 +401,30 @@ def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor, model=None) # This is crucial for proper loss averaging across micro-batches num_tokens = loss_mask.sum().clone().detach().to(torch.int) + # Context parallel: each rank saw only its shard. Sum both the loss and the + # token count across the CP group so every rank reports the FULL sequence's + # numbers -- otherwise the logged loss is a per-shard partial and the + # normalisation by num_tokens is wrong. The gradient is unaffected by this + # reduction: it is applied to detached copies for reporting, while `loss` + # itself is scaled by cp_size to cancel the averaging that Megatron's DDP + # performs over the data-parallel-with-context-parallel group. + cp_group_loss = _sft_get_cp_group() + if cp_group_loss is not None: + summed = torch.stack([loss.detach().float(), num_tokens.float()]) + torch.distributed.all_reduce(summed, group=cp_group_loss) + loss_report, tokens_report = summed[0], summed[1].to(torch.int) + loss = loss * cp_group_loss.size() + else: + loss_report, tokens_report = loss.clone().detach(), num_tokens + # Create reporting loss for logging # Format: [loss_value, num_tokens] concatenated # This allows Megatron to compute proper weighted average across DP ranks - reporting_loss = torch.cat([loss.clone().detach().view(1), num_tokens.view(1)]) + reporting_loss = torch.cat([loss_report.view(1), tokens_report.view(1)]) # Return standard Megatron loss function signature # (loss, num_tokens, metrics_dict) - return (loss, num_tokens, {"lm loss": reporting_loss}) + return (loss, tokens_report, {"lm loss": reporting_loss}) _pre_forward_canary(model) diff --git a/primus/configs/modules/megatron/trainer_base.yaml b/primus/configs/modules/megatron/trainer_base.yaml index 6830dd8b0..fc488e920 100755 --- a/primus/configs/modules/megatron/trainer_base.yaml +++ b/primus/configs/modules/megatron/trainer_base.yaml @@ -110,8 +110,8 @@ muon_use_nesterov: false muon_split_qkv: true muon_momentum: 0.95 -optimizer_cpu_offload: false -optimizer_offload_fraction: 1.0 # float +optimizer_cpu_offload: ${PRIMUS_OPTIMIZER_CPU_OFFLOAD:false} +optimizer_offload_fraction: ${PRIMUS_OPTIMIZER_OFFLOAD_FRACTION:1.0} # float use_torch_optimizer_for_cpu_offload: false overlap_cpu_optimizer_d2h_h2d: false pin_cpu_grads: true