diff --git a/.gitignore b/.gitignore index 673ea19e..911fd414 100644 --- a/.gitignore +++ b/.gitignore @@ -22,6 +22,7 @@ local/* # Environments .env .venv +.venv-sglang/ env/ venv/ ENV/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3f2c9cf9..4fad13c9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -14,6 +14,9 @@ repos: - id: check-toml - id: check-ast - id: check-added-large-files + # The semantic audit ledger intentionally retains evidence for every + # reviewed scope. Keep the size guard for all other additions. + exclude: ^test_audit_decisions\.json$ - id: check-merge-conflict - id: check-shebang-scripts-are-executable - id: detect-private-key @@ -37,3 +40,10 @@ repos: args: ["--config", ".codespellrc"] additional_dependencies: - tomli + - repo: local + hooks: + - id: check-public-tree + name: check public tree for private references + entry: python scripts/check_public_tree.py + language: system + pass_filenames: false diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 62af9fa4..913aaf3b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -47,6 +47,7 @@ The hooks include: - **ruff-format** — code formatting (line length 120) - **codespell** — catches typos - **trailing-whitespace / end-of-file-fixer** — file hygiene +- **check-public-tree** — rejects private paths, cluster references, and internal identifiers CI runs the same hooks, so if pre-commit passes locally, CI will too. @@ -56,12 +57,18 @@ Additional guidelines: ## Testing +Test additions, consolidation, and removal follow the evidence-based process in [TEST_TRIAGE.md](TEST_TRIAGE.md). + ```bash -# Run unit tests -pytest tests/ +# Create the repository-local Python 3.12 test environment from uv.lock +uv venv .venv --python 3.12 +UV_PROJECT_ENVIRONMENT=.venv uv sync --group test --no-install-project + +# Run unit tests against this checkout rather than an ambient xorl install +PYTHONPATH="$PWD/src" .venv/bin/python -m pytest tests/ # Run a specific test file -pytest tests/server/api_server/test_checkpoint_paths.py -v +PYTHONPATH="$PWD/src" .venv/bin/python -m pytest tests/server/api_server/test_checkpoint_paths.py -v ``` Tests must pass locally before requesting review. diff --git a/README.md b/README.md index 63084028..1efe6e63 100644 --- a/README.md +++ b/README.md @@ -76,14 +76,26 @@ The repo includes two git submodules under `submodules/` (needed for server / on - **[xorl-client](https://github.com/togethercomputer/xorl-client)** — Lightweight Python SDK (no PyTorch dependency) for driving the xorl training server. Provides `ServiceClient`, `TrainingClient`, `SamplingClient`, and `RestClient` with async-first `APIFuture` semantics, automatic request ordering, and Tinker API compatibility. - **[xorl-sglang](https://github.com/togethercomputer/xorl-sglang)** — XoRL's fork of [SGLang](https://github.com/sgl-project/sglang) with NCCL-based weight sync endpoints, MoE routing data export (R3), and numerical alignment flags for online RL. -Install individually: +Install the client in the default environment. Keep SGLang in its own +Torch-2.11 environment; its compiled kernel wheel is not ABI-compatible with +the default Torch-2.12 profile: ```bash pip install -e submodules/xorl-client -pip install -e "submodules/xorl-sglang/python[all]" +uv venv .venv-sglang --python 3.12 +uv pip install --python .venv-sglang/bin/python -e submodules/xorl-sglang/python +uv pip install --python .venv-sglang/bin/python \ + torchdata==0.11.0 nvidia-cutlass-dsl==4.5.2 quack-kernels==0.5.0 +uv pip install --python .venv-sglang/bin/python --no-deps -e . +uv pip install --python .venv-sglang/bin/python pytest +PYTHONPATH=src:submodules/xorl-sglang/python XORL_REQUIRE_SGL_KERNEL=1 \ + .venv-sglang/bin/python -m pytest -q tests/ops/test_sgl_kernel_smoke.py ``` -Or use the bundled `pyproject.sglang.toml` which pins PyTorch to 2.9.1 (required by sglang) and installs everything together: +The bundled `pyproject.sglang.toml` provides the same combined profile for uv. +It also overrides SGLang's newer Quack/CUTLASS metadata with the versions used +by XoRL's trainer source; the exact SGLang kernel smoke above validates that +boundary. **uv:** ```bash @@ -92,15 +104,7 @@ uv sync source .venv/bin/activate ``` -**conda:** -```bash -conda create -n xorl-sglang python=3.12 -conda activate xorl-sglang -cp pyproject.sglang.toml pyproject.toml -pip install -e . -``` - -> **Note:** The default `pyproject.toml` uses PyTorch 2.10.0. sglang requires PyTorch 2.9.1, so the two cannot coexist in the same environment unless you use `pyproject.sglang.toml`. +> **Note:** The default `pyproject.toml` uses Torch 2.12.1. Pinned SGLang requires Torch 2.11.0; do not install `sglang-kernel` into the default environment. See the [installation guide](https://togethercomputer.github.io/xorl/getting-started/installation/) for full setup including optional dependencies (DeepEP, Flash Attention). diff --git a/TEST_TRIAGE.md b/TEST_TRIAGE.md new file mode 100644 index 00000000..d5d8c3fb --- /dev/null +++ b/TEST_TRIAGE.md @@ -0,0 +1,5516 @@ +# Test triage + +The goal is a smaller suite with a stronger failure signal, not a small suite by itself. Slowness, GPU requirements, +or a large test file are not sufficient reasons to delete a test. + +## Decision standard + +For every candidate, write down the plausible production regression that should make it fail. Then choose one action: + +- **Keep** when it protects a supported behavior, numerical invariant, public integration, or failure path. +- **Consolidate** when the behavior matters but an equivalent contract is already exercised elsewhere. +- **Rewrite** when the intended contract matters but the test is stale, tautological, or coupled to implementation text. +- **Relocate** when it is a benchmark, diagnostic, campaign check, or external-model certification rather than a repository + pass/fail test. +- **Remove** only when it has no distinct observable contract, is fully subsumed, or cannot run from a clean checkout. + +A removal needs concrete evidence: the covering test, the missing tracked fixture, the tautology, or the reason no production +regression can make it fail. Candidate signals from automation are not evidence on their own. + +## Semantic review heuristics + +- Prefer one test per supported behavior or failure boundary. Put equivalent input spellings and invalid literals in a table + inside that test unless they reach genuinely different production branches. +- A private-helper test must map to a supported producer/consumer path. An escape hatch or layout with no matching runtime + endpoint is not a compatibility contract. +- Do not multiply tests by model layers, expert indices, tensor sizes, or environment variables when the production logic is + shape- or index-independent. Keep representative boundaries and any case that changes ownership or control flow. +- Treat a large synthetic sweep as certification or stress coverage unless its scale can reveal a distinct repository + regression. Move certification out of the default suite and remove repeated scale-only cases. +- When a direct unit fixture bypasses builder or loader state, reproduce the production admission state explicitly; otherwise + the test may be asserting a branch it never entered. + +## Workflow + +1. Generate a fresh static inventory: + + ```bash + python scripts/audit_tests.py --format json > /tmp/xorl-test-audit.json + ``` + +2. Review one subsystem at a time. Add the decision and evidence to `test_audit_decisions.json` before editing. +3. For consolidation, map every removed assertion to its surviving behavioral test. +4. Compare collection before and after, then run the affected surviving tests. +5. Keep removal waves reviewable. Do not combine unrelated product changes with test cleanup. + +The scanner intentionally over-reports tests with no inline assertion, conditional skips, source inspection, or exact duplicate +bodies. Fixtures and helper assertions can make those tests valuable; a human must inspect each candidate. + +## Initial wave + +The first wave removes tests that are fully subsumed, permanently skipped in a clean clone, or unable to report a failure. +It relocates the official GLM-5.2 checkpoint inventory to `certification/glm52/`. Valuable but stale checks remain in the +decision ledger as proposed rewrites rather than being deleted. + +## Second wave + +The second wave removes four fully covered or self-referential helper tests, consolidates three duplicated loss contracts, +and replaces the remaining implementation-source inspection with a behavioral RoPE discriminator. It also repairs the +server-argument test harness so temporary import stubs restore only the keys they replace instead of unloading native +extensions imported during collection. + +After this wave the static inventory contains 2,646 test definitions, 79 review candidates, one accepted duplicate wrapper +pair, and no source-inspection candidates. The decision ledger records the covering evidence and keeps parametrized behavior +for both loss implementations. + +## Third wave: semantic contracts + +The third wave reviews four high-density modules by production contract rather than scanner signal. It consolidates config +override matrices, checkpoint-key spellings, weight-sync precedence rules, and direct-EP selection fragments. It removes +orphan P2P fused layouts that have no receiver locator, a layer-by-expert scale sweep of layer-independent logic, and a +single-receiver FP8 case covered by the retained multi-receiver contract. + +Those four modules collect 167 tests instead of 225. The static repository inventory now contains 2,626 definitions; all +invalid config values and global expert indices remain covered. The review also found and repaired an exact-GLM fixture that +omitted the internal admission marker installed by the production model builder. + +## Fourth wave: geometry versus branch coverage + +The fourth wave targets parametrization that varied sizes without reaching new code. V4 indexer coverage now keeps the basic, +batched, production-head/top-k, and C128 geometries; its separate real-config sweep was already covered by the score and top-k +contracts. Families-v2 RMSNorm keeps tail, aligned, and deep-tile shapes at row-count boundaries, while its shipped-size +dispatch check reflects that rows cannot affect the decision below the split-tile threshold. + +The decode GDN suite no longer invokes internal warp/group launch configurations that no production caller can select, and +its scaled-then-normalized K regime is not treated as a separate input contract. Pipeline placement now checks each mapping +formula rather than multiplying the same arithmetic across PP sizes. + +These four modules collect 44 tests instead of 141. All 44 retained tests pass on GPU. Repository collection is 3,056 items +and the static inventory is 2,624 definitions with 23 curated decisions. + +## Fifth wave: tables are contracts, not test suites + +The fifth wave collapses parser and policy truth tables into one contract each, preserving every input row. Fused selected- +logprob parity replaces independent dtype × bias × temperature products with pairwise branch coverage while retaining all +irregular shapes and production vocabularies. The batch-invariant GEMM table still executes every dtype/shape combination but +reports one doctrine-level bit-neutrality contract. + +Sparse MLA retains all four compiled top-k specializations, the batched case, and the production 64-head geometry. Its sink +coverage keeps the zero boundary and mixed-sign values; a separate large-sink test still proves that the kernel consumes the +sink rather than ignoring it. + +These four modules collect 44 tests instead of 101. The retained set reports 43 passes and one optional DeepGEMM skip on the +current GPU host. Repository collection is now 2,999 items; the static inventory remains 2,624 definitions with 27 curated +decisions. + +## Sixth wave: cross-subsystem semantic cleanup + +The sixth wave spans LoRA gradient ownership, torch.compile, inference API normalization, packing, quantization validation, +and FP8 weight selection. LoRA quantized experts now use pairwise format/backend coverage while retaining every format, +backend label, and producer branch. Explicit monkeypatch contexts preserve per-case distributed-state isolation. + +The compile suite removes a fullgraph probe that swallowed every exception and two non-collected, print-only benchmark +routines. Its three retained contracts perform real AOT/Inductor compilation at block, decoder-layer, and full-model levels. +Inference endpoint detection keeps its integrated HF-config normalization contract rather than duplicating private helpers. + +Packing removes an algebraic restatement of row minimization and a Python-list-indexing tautology; datum order is now checked +against document boundaries in actual packed rows. Quantization aliases and rejection literals are grouped by boundary. FP8 +projection selection is one buffer-level contract spanning supported families and negative cases instead of fourteen narrow +cases. + +The six focused suites report 106 passes. Repository collection is 2,926 items and the static inventory is 2,610 definitions +with 33 curated decisions. + +## Seventh wave: separate backend math from topology + +The seventh wave targets distributed adapter-autograd and expert QLoRA matrices. The unquantized distributed suite still +executes every backend at EP2, but the four-GPU eFSDP composition uses one shipped Quack representative rather than repeating +the same backend dimension. The default all-owner layout likewise uses one representative after backend math is established. +Quantized expert autograd replaces the three-backend by three-format Cartesian product with pairwise coverage that retains +every backend and format; DeepEP and exact projection-subset contracts remain independent. + +Expert QLoRA capability, unsupported-semantic, model-family, and invalid-target tables now report one test per behavioral +boundary. Every table row still executes, including construction of all three model-specific expert families. The two files +collect 29 tests instead of 60 and eliminate twelve distributed subprocess launches. The retained QLoRA suite reports 17 +passes; the distributed gate reports 10 passes and two optional DeepEP skips on eight H100 GPUs. Repository collection is +now 2,895 items; the static inventory remains 2,610 definitions with 35 curated decisions. + +## Eighth wave: make numerical tests exercise nontrivial numerics + +The eighth wave strengthens MoE-LoRA backend parity before removing its weaker smokes. The previous cross-backend gradient +test left every LoRA-B factor at its zero initializer, so all LoRA-A gradients were trivially zero. The retained comparison +uses nonzero adapters and compares eager with both native and Triton outputs plus every A/B gradient. It subsumes separate +GPU forward/backward smokes; explicit zero-delta and nonzero-effect reference contracts remain. Backend-independent adapter +construction uses one Quack representative, while registration and numerical contracts continue to cover all backends. + +RoPE validation still checks the exact fp32 frequency table for every registered initializer. Shared forward consumption now +uses default and YaRN to cover unit and non-unit attention scaling. The native-cache comparison is restricted to default +RoPE, the only type for which `_rope_native` changes `RotaryEmbedding` control flow; five non-default cases previously +compared two executions of the same stock branch. The focused suites report 16 MoE-LoRA and 7 RoPE passes. Repository +collection is now 2,864 items and the static inventory is 2,608 definitions with 37 curated decisions. + +## Ninth wave: FP8 launch branches and transactional corruption boundaries + +The ninth wave reduces FP8 grouped-MoE geometry by kernel control flow. Same-NK coverage keeps small- and large-N launch +branches, aligned and padded K, an empty expert, M spanning several blocks, and output tails in two cases. Wgrad retains the +same split plus a separate all-empty case for its `max_K == 0` early return. Both grouped backends still perform a real +optimizer step, while the dense-plus-MoE integration uses the default Triton-grouped representative instead of repeating +the already-qualified backend dimension. + +FP8-linear block sizes and recipe rows remain intact because they produce different padding, scale grids, SmoothQuant, and +amax behavior; each table now reports one numerical contract. Adapter-optimizer resume similarly retains every identity, +layout, coverage, dtype, shape, step, and staged-state corruption, but groups them into three transactional contracts using +isolated checkpoint directories. The focused gates report 41 FP8 passes with one opt-in DeepGEMM skip and 26 optimizer-resume +passes. +Repository collection is now 2,842 items; the static inventory remains 2,608 definitions with 40 curated decisions. + +## Tenth wave: checkpoint gates and resolver truth tables + +The tenth wave rewrites gradient-checkpointing coverage around the actual gate. Class and enable-time defaults are one +contract across base and MoE layers. Nondefault propagation uses one representative because the model method performs the +same assignment for each already-validated string. The outer-gate test adds the previously missing method predicate and now +proves that selective checkpointing does not invoke the full-layer checkpoint even when training and the feature flag are +both enabled. + +MoE routing-position auto regimes and explicit true/false aliases retain every row, but each resolver boundary is one test. +The two focused files collect 12 tests instead of 27, and all 12 pass. + +## Eleventh wave: exact-construction admission boundaries + +The eleventh wave groups GLM-5.2 exact MoE dependency, EP-topology, and adapter-shape admission failures. Every missing exact +component, sparse-MLA/dispatch requirement, uninitialized or wrong-sized EP state, and individual rank/alpha violation still +fails before adapter mutation. A doubly invalid rank-16/alpha-16 row was removed because the rank-only and alpha-only cases +already exercise both sides of the same combined predicate. + +Exact shared-expert construction keeps all seven shape, TP, rank, alpha, bias, and adaptive-noise rejection boundaries in one +contract. The two focused files collect 14 tests instead of 28 and report 12 passes with two SGLang-dependent skips. +Repository collection is now 2,813 items and the static inventory is 2,607 definitions with 44 curated decisions. + +## Twelfth wave: test complete loader transactions, not their fields separately + +The twelfth wave replaces fragmented QLoRA expert-loader tests with one integrated contract per quantization format. Every +retained geometry still loads, but each load now validates packed bytes, scales, and all three projections together. NVFP4 +also checks per-expert amax, absorbed global scales, and dequantized shape/dtype. This reduces fourteen repeated synthetic +checkpoint loads to six and removes an uncollected, print-only timing/`__main__` harness with no pass/fail threshold. + +FP8 external-config coverage retains all Transformer-Engine-only recipe keys and ModelOpt QARL nesting variants inside their +own fail-closed contracts. The focused QLoRA and config suites collect 12 tests instead of 29, and all 12 pass. Repository +collection is now 2,796 items and the static inventory is 2,601 definitions with 46 curated decisions. + +## Thirteenth wave: remove masked axes and weaker backend smokes + +The thirteenth wave keeps both RMSNorm kernel geometries but reports the family funnel as one bitwise contract. Its module +parity matrix already enables the trunk contract in every row, which fixes family dispatch independently of the separate +batch-invariant flag; that masked axis is removed while native, SGLang, and fused-SGLang modes remain. Dense and MoE Qwen3.5 +call-site truth tables retain every layer/mode row inside one contract per predicate. + +OPD chunking keeps disabled, one-chunk, and over-partitioned boundaries rather than arbitrary intermediate positive counts. +All VERL estimator aliases and both streaming backends remain. Three direct compiled-function tuple/shape smokes are removed +because retained end-to-end reverse/forward KL tests execute the same functions and additionally validate formulas and +diagnostics. The focused gates report 27 OPD and 34 RMSNorm passes. Repository collection is now 2,765 items and the static +inventory is 2,598 definitions with 49 curated decisions. + +## Fourteenth wave: collapse arbitrary geometry and transactional row inflation + +The fourteenth wave reduces NVFP4 forward coverage to its two supported dtypes because the three legal matrix shapes all +entered the same flatten-and-block path. Exact equality with the independent quantization reference subsumes separate +Gaussian-error and lossy smokes, while the E2M1 grid, STE, invalid rank/K layout, registry, expert isolation, fused-half +scale, and metadata contracts remain. EP wrapper signatures now inspect every entry in both live registries inside one API +contract rather than reporting one node per backend. + +Canonical MoE keeps the two-contributor base tree and the production sixteen-contributor tree; intermediate CPU reference +widths selected no new branch. Its distributed test retains two and eight contributors plus the independent EP16 packed +transport gate. A second auto-transport test was removed because its two assertions were verbatim subsets of the adjacent +admission contract. GLM-5.2 QLoRA keeps independent rank-only and alpha-only failures but removes their doubly-invalid +conjunction, and its construction-mode table still builds a fresh model for every rejected mode. Exact-attention checkpoint +tests still run both arrival orders and both weight/scale members for duplicate, missing, dtype, and shape failures, but each +transactional behavior reports one contract. + +The five files collect 53 tests instead of 80. The focused gates report 52 passes and one optional CUDA-path skip. +Repository collection is now 2,738 items and the static inventory is 2,593 definitions with 54 curated decisions. + +## Fifteenth wave: remove scale certification and test-the-helper arithmetic + +The fifteenth wave replaces the Quack EP distribution-by-score Cartesian product with three pairwise cases that retain +balanced, empty-expert, extreme-skew, score-free, and score-scaled behavior. Eager/native MoE parity keeps ordinary, top-k +one/four, and single-token routing boundaries. Two size-only configurations and a separate E64/H512/I1024 test were removed: +they selected no repository branch, and the latter only repeated parity with larger allocation and looser tolerances. + +Full-reduce mean now crosses its actual one-dimension versus multi-dimension branch with both supported dtypes, without a +third arbitrary tensor rank. V2 RMSNorm retains shipped and deep hidden sizes at low and high row counts, all residual modes, +and explicit split/fused execution; the intermediate row count selected no dispatch or kernel boundary. LoRA merge keeps +zero and nonzero production merges for every original dtype in both linear and MoE implementations. A supposed precision +test was removed because it never invoked either production merge method and only compared two formulas local to the test. + +The five files collect 15 tests instead of 42, and all 15 pass on the current GPU host. Repository collection is now 2,711 +items and the static inventory is 2,591 definitions with 59 curated decisions. + +## Sixteenth wave: distinguish repository regressions from dependency canaries + +The sixteenth wave reports all supported EP backend gradient domains as one local contract while retaining the real two-rank +autograd and synchronization gate. Exact-attention construction keeps rank-only and alpha-only rejection, removes their +doubly-invalid conjunction, and still builds a fresh model for every incomplete execution mode. Every active-LoRA component +flag remains independently cleared against all four admission predicates inside one conjunction contract. + +Mamba2 exact-chunk/tail, multi-chunk/tail, aligned/unaligned packed boundaries, forward values, and gradients all remain. +The test asserting that the Transformers 5.5.3 fallback bug still exists was removed: an upstream fix is not an XoRL +regression and should not fail the repository. The sequential SSD recurrence remains the multi-chunk oracle. Qwen3-MoE input +norm family coverage now varies only layer zero versus later layers; the old mode axis was ignored by the capture stub and +could not affect the call site. FlashQLA retains every M/batch and chunk-chaining execution but reports each certification +gate once. + +The six files collect 43 tests instead of 67. The focused gates report 41 passes and two optional mamba-ssm skips. +Repository collection is now 2,687 items and the static inventory is 2,590 definitions with 65 curated decisions. + +## Seventeenth wave: remove dominated loss references and group receiver truth tables + +The seventeenth wave keeps all four full-weight/adapters admission combinations and all three optimizer-load spellings while +reporting each argument boundary once. EP checkpoint selection still rejects every missing/ambiguous named dimension and +restores both legacy and PP-parent meshes. Native GLM FP8 validation still rejects each nonofficial quantization field. + +Streaming forward-KL backward parity already compares forward values and gradients to the independent dense oracle, so the +separate forward-only check was removed. The direct compiled-reference unit was also removed because the retained end-to-end +OPD dispatch compares both streaming aliases with the compiled backend including gradients. Chunk invariance now directly +compares a seven-wide multi-chunk execution with a forty-wide single-chunk execution; the old 40-versus-40000 one-chunk row +added no branch and the 40000-versus-itself row was tautological. TokenPartial microbatch additivity uses one caller-supplied +denominator because denominator values select no reducer branch; scale-one and both token/sequence zero-denominator contracts +remain independent. + +The five files collect 55 tests instead of 74, and all 55 pass. Repository collection is now 2,668 items and the static +inventory is 2,588 definitions with 70 curated decisions. + +## Eighteenth wave: replace stochastic convergence experiments with mechanism checks + +The eighteenth wave removes five QLoRA random-target convergence experiments that collectively ran roughly 750 optimizer +iterations. They included a 400-step rank-accumulation comparison and a loose reset-within-two-times threshold; neither maps +to deterministic repository control flow. Both quantization formats still exercise storage, memory, forward, and backward. +NVFP4 and Block-FP8 merges, prequantized loading, requantization, injection, LoRA-only optimizer reset, and non-LoRA state +preservation remain. The scheduler integration now proves an off-boundary no-op and an on-boundary packed-weight change, +LoRA-B reset, and stale optimizer-state removal directly after one state-populating step. + +Mooncake metadata still rejects every missing/invalid field, LoRA manifests still reject every nonexact scalar, Nemotron-H +checkpoint construction still covers indivisible experts, invalid rank, and invalid EP size, and DeepSeek training still +rejects each unsupported mode. These receiver truth tables now report one contract apiece. + +The five files collect 45 tests instead of 62, and all 45 pass. Repository collection is now 2,651 items and the static +inventory is 2,583 definitions with 75 curated decisions. + +## Nineteenth wave: force real reference paths and delete refactor tombstones + +The nineteenth wave removes GLM-5 construction smokes already subsumed by the default-shape and end-to-end model contracts. +It also fixes a false oracle: a test labeled as TileLang-versus-torch passed a pure causal mask, but that mask is accepted by +the TileLang fast path. The retained GPU contract now opts into blocked scoring to force the independent torch path before +comparing outputs. All shared-factor LoRA autograd executions remain, reported as one behavioral contract. + +The simulator's consolidated validator now owns built-in pack discovery, schema, sanitation, and exact raw/promotable +goldens. DeepEP topology and preflight scenarios are complete truth tables rather than ten separately reported micro-tests. +Packing drops a test-only abstract subclass, a brittle exact-key whitelist, an arbitrary 100-sample repetition, and repeated +roundtrip modes already covered by focused tests; numpy normalization and one composition roundtrip remain. Session cache-path +canonicalization still runs every direction inside one contract. + +Finally, thirteen launcher tests explicitly targeting APIs removed by the launcher refactor are deleted instead of being +collected forever as skips. The live launcher suite retains seven address, readiness, command, parsing, and migration tests. +The seven focused files collect 154 tests and all 154 pass. Repository collection is now 2,619 items and the static inventory +is 2,555 definitions with 82 curated decisions. + +## Twentieth wave: test artifact transactions and behavioral truth tables + +The twentieth wave joins checkpoint metadata writers to their actual compatibility readers. QARL buffer metadata and +pipeline-stage key unions are now written, inspected, and validated as complete transactions instead of testing producers +and consumers against separate fixtures. Token diagnostics similarly combine shape and target-ranking invariants, plus the +two equivalent disabled-input exits, without dropping any output field or hidden-component coverage. + +Muon keeps real optimizer autotuning and drops the helper-only duplicate. Quack tuned/untuned dispatch, SM90 dtype backend +selection, and cautious optimizer-family routing retain every execution as truth tables. API constructor field echoes and +Pydantic assignment smokes are removed because retained endpoint tests cover defaults, aliases, serialization, registration, +and payload forwarding at the application boundary. The heartbeat check no longer sleeps on wall-clock time. + +Trainer clipping, vote scaling, and target-token preference tables retain every original case. Router tie-policy aliases are +grouped, and the invalid policy test now asserts the real construction-time failure. Distillation cache host/device rank-3 +gathers and both bounds failures are grouped by behavior. DeepSeek-V4 successful name mappings form one complete mapping +contract while unknown/MTP rejection stays separate. The nine focused files collect 152 tests and all 152 pass. Repository +collection is now 2,597 items and the static inventory is 2,536 definitions with 91 curated decisions. + +## Twenty-first wave: keep transactions, collapse helper reporting + +The twenty-first wave keeps request-processor transport, Mooncake cleanup, model-ID propagation, registration, packing, +statistics, and error boundaries. Backend timing-field preservation now lives in the main model-pass contract, and five +identical sequential forward smokes are removed from an error test because the dedicated statistics contract already proves +exact operation accounting. Runner dispatch drops a second routing-slice test whose expert IDs, logits, offsets, and metadata +cleanup were all asserted by the retained test. + +Pipeline profiling retains every interval-union case, nonzero and zero bubble formula, rejection boundary, P2P topology, +patch lifecycle, and CUDA-event integration, but reports the pure-math cases as behavioral truth tables. The OPD driver still +checks chunk order/tails and every student weight-version outcome without separate nodes per outcome. Sparse-delta capture and +writing retain both traversal fields and all duplicate, dtype, length, and range failures inside receiver contracts. The +adapter-coordination suite and quantized export transforms were audited and kept because their tests cover distinct rollback, +trust-root, sharded-state, and tensor-conversion paths. + +The five focused files collect 82 tests and all 82 pass. Repository collection is now 2,583 items and the static inventory is +2,523 definitions with 96 curated decisions. + +## Twenty-second wave: prefer end-to-end numerical contracts over arithmetic fragments + +The twenty-second wave removes the densest remaining arithmetic fragments from EP gradient clipping. Retained +classify-then-clip tests already prove skip-FSDP and ordinary grouping, combined L2 clipping, uniform scaling, and the +no-double-division regression. Infinity norm, missing gradients, shared replicas, dispatch, mixed DTensor meshes, explicit +foreach behavior, and live two- and three-rank reductions remain. Separate single-group, 3-4-5, no-clip, and duplicate +mixed-mesh examples selected no additional branch. + +DistSignSGD now proves signing and forced-SUM communication in one AVG-input transaction; its weight-decay step subsumes the +plain update, and its state-dict roundtrip subsumes a separate state-empty smoke. Unsupported HSDP, folded sequence +parallelism, and EP topologies still fail against fresh models inside one truth table. Scheduler coverage retains every +warmup, decay, floor, cosine, and validation branch while dropping phase examples already contained in longer traces. + +NVFP4 ownership is now explicit: the op suite owns the independent quantization reference and exact 2D/3D STE arithmetic, +while QARL wrapper suites own configuration, injection, lossy production forwards, parameter restoration, and gradients. +Direct private-helper/shadow tests, a looser grid property, and a supported-format boolean smoke were dominated by those +contracts. Non-16 group sizes, dense versus expert target selection, FP8-on-MoE rejection, and quantization-disable parity +all remain. Tensor-parallel FP8 model building still executes both included and excluded lm-head outcomes in one contract. + +The eight focused files collect 67 tests instead of 105, and all 67 pass, including both live distributed gradient gates. +Repository collection is now 2,545 items and the static inventory is 2,485 definitions with 101 curated decisions. + +## Twenty-third wave: replace QARL context fragments and convergence heuristics + +The twenty-third wave turns QARL activation override coverage into state transactions. Enabled and disabled modes, distinct +per-module restoration, exception cleanup, and exclusion of ordinary linears now execute together; nested contexts remain a +separate reentrancy contract. W4A4 activation quantization retains its independent forward reference and a nonuniform +upstream-gradient STE check, which subsumes the all-ones gradient example. The exception path proves both `triton_w4a4` +selection and restoration, while activation-off and non-Triton no-op cases form one conjunction-boundary table. + +The dense QARL training smoke no longer trains sixteen steps against a synthetic target and asserts that loss happens to +decrease. One real AdamW step now proves finite loss, gradients through both wrapped projections, parameter mutation, changed +logprobs, persistent summary metadata, and exact checkpoint restoration. The weight-sync handler's mismatch response already +enters the block-size validator, so the weaker direct failure unit is removed. + +Cautious SignSGD and AnyPrecisionAdamW keep mixed-coordinate production steps that contain aligned and misaligned elements +and assert the exact mask/update. Separate all-aligned comparisons selected no different branch and are removed. The five +focused files collect 30 tests instead of 39, and all 30 pass. Repository collection is now 2,536 items and the static +inventory is 2,476 definitions with 105 curated decisions. + +## Twenty-fourth wave: make GLM selector tests describe production behavior + +The twenty-fourth wave removes an IndexShare test that only proved a recursive mapper defined inside the test file preserves +non-tensor Python identity. The retained dense-producer/shared-consumer model forward runs that mapper from simulated FSDP +pre-hooks and proves one context survives all layers, is consumed without recomputation, and is cleaned up after the forward. + +GLM-5 selector shape, dtype, range, final-row validity, and sorted-sentinel ordering now form one output contract. Dense and +one-head-chunk scoring still execute the same diagonal-only additive mask inside one behavioral table. A CUDA- and +TileLang-gated test that never invoked TileLang is now an honest CPU mask-classification contract for valid prefixes versus +interior holes. The standalone sparse-attention shape smoke is removed because retained full-model sparse-versus-dense +parity and Ulysses integration already exercise the forward while checking numerical output, query/KV locality, indices, +offsets, masks, and output shape. + +Six focused retained contracts pass. Repository collection is now 2,532 items and the static inventory is 2,472 definitions +with 108 curated decisions. + +## Twenty-fifth wave: report exact-model admission as family contracts + +The twenty-fifth wave retains every exact Qwen3.5 execution while removing pytest item inflation. Dense and MoE configs both +resolve the full certified attention, router, lm-head, RMSNorm, activation, RoPE, cast, sparse-MLA, and cross-entropy program +inside one family contract. World16 HSDP-plus-EP, world8 EP, and single-GPU dense topologies still validate from fresh config +objects inside one accepted-topology table. + +Model-scope admission likewise keeps dense, MoE, and Hugging Face outer-config snapshots in one accepted-scope contract. +Both dense and MoE hidden-size near misses still fail independently in one rejection contract. The file collects 21 tests +instead of 27, all 21 pass, and no configuration execution was dropped. Repository collection is now 2,526 items and the +static inventory is 2,469 definitions with 110 curated decisions. + +## Twenty-sixth wave: consolidate weight-sync setup, preserve adapter transactions + +The twenty-sixth wave joins requested-adapter materialization and current-adapter fallback in one state contract, and joins +dense-buffer chunking with the cap predicate it consumes. A tied-weight root extraction is removed because the retained +prior-module test performs the same extraction and alias assertions before additionally proving the later duplicate is +skipped. Direct Nemotron-H prefix mapping and sparse-delta request-field echoes are removed because retained unfuse, +remote-backend, and end-to-end request transactions verify those values at their actual consumers. + +Three copied sparse-delta fake endpoint/backend harnesses are replaced by one transport transaction. Baseline post-only +ordering/accounting, explicit baseline configuration, and FP8 KV-cache postprocess metadata all still execute and assert +backend configuration, normalized cache epoch, endpoint results, pause/resume order, posted paths, and weight version. + +The 48-test adapter-manager suite was audited and kept. Its size comes from distinct gradient-ownership, staging, atomic +commit, abort, clipping, collective, poisoning, publication, trust-root, session-spec, rollback, eviction, mixed-rank, and +optimizer-state boundaries rather than literal variation; all 48 pass. The weight-sync file collects 30 tests instead of 37, +and all 30 pass. Repository collection is now 2,519 items and the static inventory is 2,462 definitions with 113 curated +decisions. + +## Twenty-seventh wave: distinguish model smokes from numerical certification matrices + +The twenty-seventh wave reports MiniMax-M3 language-model prefix and w1/w2/w3 expert-key aliases as one classifier contract, +retaining all strings. The DSv4 base-model C0 shape smoke is removed because the retained causal-LM C0 forward/backward runs +the same base model and additionally checks logits, finite loss, required gradients, and intentionally frozen +hyperconnection parameters. The distinct C128 compressed-attention forward remains. The two focused model-support files +collect 21 tests instead of 24, and all 21 pass. + +Cross-engine RMSNorm shapes are valuable certification inputs rather than arbitrary examples, so none are removed. All four +adversarial shapes, both family funnels, residual modes, trunk lane, zero-centered twin, and families-v2 candidate executions +remain, but each invariant reports one test rather than one item per row. That suite would collect 9 tests instead of 34 in +an environment with SGLang batch-invariant ops; the current venv dependency-skips the module. The locally available fused +suite now reports BF16 and FP32 residual/no-residual parity as two dtype-complete contracts, and both pass. + +Repository collection is now 2,514 items and the static inventory is 2,461 definitions with 115 curated decisions. + +## Twenty-eighth wave: keep transport boundaries, remove routing and telemetry examples + +The twenty-eighth wave keeps the weight-sync boundaries that can corrupt state or strand a receiver: primary-to-fallback +health ordering, NCCL initialization, two-phase completion metadata, mixed-dtype byte flattening, chunking, receiver-fenced +work lifetime, multi-rank load-format rejection, sparse-delta byte-change encoding, unchanged-bucket suppression, baseline +priming and rollback, FP8 KV-cache metadata, and per-rank packed paths. It removes a standalone primary-health URL example, +folds configured load-format forwarding into the existing bucket-routing transaction, and makes the receiver-fence contract +also prove cleanup on group destruction instead of repeating the entire hybrid-broadcast harness. + +Sparse-delta drops a factory `isinstance` mirror and a single-path replication example already executed by streaming +transfer. The retained per-rank prepacked path contract now owns unique-file accounting. Post-only import avoidance, +prepacked-only streaming rejection, and valid prepacked post-only initialization all still execute as one policy table. The +fixture now uses a validated loopback literal rather than relying on the fake hostname `infer-0` to resolve before mocked +HTTP calls reach production URL-safety validation. + +Trainer IB-device selection retains global-rank, local-rank, single-device fallback, explicit physical-GPU, numeric +`CUDA_VISIBLE_DEVICES`, and empty-entry autodiscovery outcomes in one environment-isolated precedence contract. Two private +telemetry tests that only copied fields into a dictionary or added three counters are removed; abort-marker lifecycle and +distributed peer-failure gathering remain. The three focused files collect 21 tests instead of 35, and all 21 pass. +Repository collection is now 2,500 items and the static inventory is 2,447 definitions with 119 curated decisions. + +## Twenty-ninth wave: parse composed configurations, not one scalar per test + +The twenty-ninth wave keeps configuration parsing as a production boundary while removing one-node-per-field reporting. +Flat adapter-ownership, nested removed ZORL field, and removed ZORL-section inputs still fail through the real server loader +inside one rejection table. Both shipped MoE LoRA examples and all five shipped Qwen MoE QLoRA examples still parse in a +clean subprocess and compare source YAML with normalized Quack, expert-target, and shared-LoRA values, but no longer create +one pytest item per file. + +The server loader now proves both SignSGD spellings in composed nested configurations that also carry checkpoint optimizer +policy, forward/backward prefetch, HSDP deferral, packing alignment, activation memory limits, and adapter state-load mode. +The runtime attributes and serialized model/train/LoRA dictionaries remain checked. Explicit Mooncake, legacy Mooncake +alias, and filesystem R3 success modes form one transport table while the invalid-directory failure remains separate. +Explicit and automatic MoE routing-weight placement likewise share one defaulting contract. + +The training CLI parser now accepts both sign optimizers while simultaneously preserving multipack fields, model numerical +alignment, FSDP reduction dtype, and parameter-upcast policy from production-shaped YAML. Muon conversion, legacy alias +transforms, FP8/QARL modes, automatic checkpoint resolution, load-optimizer defaults, and every incompatibility rejection +remain independent. The two focused files collect 50 tests instead of 70, and all 50 pass. Repository collection is now +2,480 items and the static inventory is 2,432 definitions with 123 curated decisions. + +## Thirtieth wave: turn numerical-contract fragments into guard tables + +The thirtieth wave retains every GatedDeltaNet exact-convolution rejection branch while reporting decode cache, context +parallelism, missing short convolution, and convolution bias as one unsupported-input table. Independent weight-packing +order, production routing, forward/backward references, kernel and end-to-end determinism, state scoping, checkpoint +recompute, and optional SGLang parity remain separate. + +Batch-invariant trunk wrapping likewise keeps no-match, ordinary-LoRA, custom-Linear, and FP16-weight failures inside one +admission table. Bias-free and biased forward comparisons still execute against the persistent GEMM, and both backward +comparisons still execute against cuBLAS autograd, without parametrized item inflation. A standalone RMSNorm loud-failure +test is removed because the retained multi-op grad-requiring interpose contract already calls RMSNorm and requires the same +failure. + +This audit also exposed an outdated oracle. The wrapper explicitly arms the RMSNorm contract lane, but the selection test +asserted that global dispatch stayed disabled. The assertion now matches the documented implementation, and the autouse +fixture establishes disabled state both before and after every test to prevent order dependence. The two focused files +collect 26 tests instead of 35; 25 pass and the optional SGLang comparison dependency-skips. Repository collection is now +2,471 items and the static inventory is 2,425 definitions with 126 curated decisions. + +## Thirty-first wave: test dispatch and export transactions instead of registry snapshots + +The thirty-first wave removes an attention check that restated `is_flash_attention` as the same registry-membership +expression used by the implementation. The retained FA4-only reload transaction now proves `flash_attention_2`, +`flash_attention_3`, and `flash_attention_4` registration and mask-family detection directly. Registered eager/native +resolution, non-flash eager fallback, and unavailable-flash rejection now share one resolver-boundary contract. Varlen, +paged-KV, SGLang, FA4, eager-head-layout, and cross-attention rejection paths remain independent. + +Quantized export drops a checked-in example's literal path-and-field snapshot; generic config override parsing and the +subprocess CLI export still validate the parser and its real consumer. Existing FP8 scales, MTP config metadata, MTP tensor +namespaces, and unfolded QARL state all continue to construct source directories and fail through one preflight table. +Every tensor split/fusion/remap, sharded index, BF16 island, duplicate-name failure, QARL fold, and block-size failure remains. + +The QARL export parity contract no longer trains eight arbitrary iterations against a synthetic target. One real AdamW step +now establishes finite loss and changed target logprobs before folding, block-FP8 export, dequantization, and exact logprob +comparison. The two focused files collect 32 tests instead of 38; 30 pass and two flash-backend tests dependency-skip. +Repository collection is now 2,465 items and the static inventory is 2,419 definitions with 129 curated decisions. + +## Thirty-second wave: express API schema examples as conversion tables + +The thirty-second wave keeps the API schema as a validation boundary without reporting each literal row separately. +Create-model ZORL, nested adapter-ownership, and create-session ZORL payloads still enter their actual Pydantic request types +and assert the exact field path plus migration message inside one rejection table. Unknown rolling-client and nested LoRA +fields retain their separate forward-compatibility contract. + +TensorData conversion now proves rank-one passthrough for model and loss fields, valid rank-two and rank-three recursive +nesting, mismatched-shape fallback, and empty higher-rank fallback in one contract. Exact nested values remain asserted, so +the sequence-length classification invariant consumed by the packer is unchanged. Datum validation, request aliases, +optimizer/session schemas, required response fields, and serialization roundtrips stay intact. The file collects six tests +instead of twelve, and all six pass. Repository collection is now 2,459 items and the static inventory is 2,415 definitions +with 131 curated decisions. + +## Thirty-third wave: prefer full EP and API ownership transactions + +The thirty-third wave removes the weaker of two identical DeepEP exclusion setups. The retained contract now requires the +all-to-all path and the precise statement that order and rounding differ without asserting an unverifiable mechanism. A +direct private pair-to-slot ordering unit is also removed because the retained full slot-combine transaction routes random +expert pairs, reorders them, performs the weighted reduction, and compares the final tensor with an independent slot-ordered +reference. FP8 exclusion, flag-off stock dispatch, top-k presentation, empty-rank behavior, training autograd, weight-cache +lifecycle, and live gradient parity remain separate. + +At the API boundary, current optimizer fields and legacy Adam aliases still reach the orchestrator payload and response +metrics in one transaction. Sampler paths embedded in `xorl://` URIs and plain paths with explicit request model IDs still +load and enter their distinct cleanup-tracking buckets in one ownership contract. Base-model repository IDs, Hugging Face +cache paths in both directions, distinct models, ordinary paths, and `None` likewise form one canonicalization contract. + +The four focused files collect 54 tests instead of 59; 53 pass and the live GPU gradient comparison dependency-skips. +Repository collection is now 2,454 items and the static inventory is 2,410 definitions with 134 curated decisions. + +## Thirty-fourth wave: collapse FP8 topology shape ladders into mechanism boundaries + +The thirty-fourth wave reduces the full-weight FP8 E2E matrix from fourteen expensive subprocess tests to eight distinct +mechanism contracts. The shorter Ulysses run is dominated by the retained longer packed Ulysses run. Three intermediate +hybrid context datasets and a basic hybrid run add sequence length or sample-shape variation without selecting a different +FP8, Ulysses, or Ring branch; the retained 4096-token long-tail multipack case exercises that composition with heterogeneous +near-full bins. The four-GPU DeepEP checkpoint-resume cross-product is also removed: dense FP8 checkpoint restoration and +DeepEP EP/eFSDP FP8 execution remain independently covered. + +A live run exposed a stale oracle in the retained baseline. Training completed two optimizer steps and all eight eligible +linears used FP8, while `lm_head` correctly remained unused because Qwen3's resolved numerical program intentionally keeps +the output head in FP32. The helper now requires exactly that production contract instead of demanding an impossible 9/9. +The shared E2E configuration generator also now applies the `extra_data` and `extra_model` mappings already passed by the +retained packed-context and DeepEP cases, rather than rejecting those calls during Python argument binding. + +The focused file collects eight tests instead of fourteen. Its one-GPU retained baseline passes live; the remaining costly +multi-GPU matrix was collection-checked rather than executed. Repository collection is now 2,448 items and the static +inventory is 2,404 definitions with 137 curated decisions. + +## Thirty-fifth wave: attach loader and cache assertions to their real transactions + +The thirty-fifth wave folds GLM-5 architecture-alias registration and loader selection into the retained local Hugging Face +config load. That transaction now constructs the actual `Glm5Config`, checks both supported architecture names and their +class relationship, and verifies the selected checkpoint loader. Two standalone membership and description snapshots are +removed. Configured-layer, MTP-boundary, far-out-of-range, and non-layer checkpoint keys likewise pass through both +normalization and the early disk-read skip hook in one handler contract instead of two copies of the same key table. + +Two direct OPD hidden-cache helper examples are removed in favor of their retained consumer transactions. The gathered-SP +writer already filters an interior valid target and asserts its persisted cache index. The multi-rank writer already gathers +local and remote chunks, orders them by logical slice, persists the concatenated tensor, and checks per-sample indices. +Packed-segment splitting, contributor ownership, Mooncake producer/consumer roundtrips, metric collectives, FSDP lm-head +anchoring, and diagnostic artifacts remain independent. + +The two focused files collect 53 tests instead of 58, and all 53 pass. Repository collection is now 2,443 items and the +static inventory is 2,399 definitions with 140 curated decisions. + +## Thirty-sixth wave: remove the private-helper layer beneath P2P transactions + +The thirty-sixth wave keeps all FP8 byte-parity, receiver-layout, cached-prepare, multi-sender, memory-registration, +failure-cleanup, and completion contracts while removing ten lower-value reports from the 78-test P2P protocol suite. +Generic TP slicing is already proved by the retained transfer that writes exact row slices to two receiver pointers; shape +incompatibility is already rejected by a retained full transfer before the engine runs; and unsliced sources pass through +many real transfers. Three direct calls to the private slicer are therefore removed, while every specialized Qwen +linear-attention and FP8 layout transformation stays independent. + +HTTP and remote-declared initialization failures now share one failure contract. Implicit all-rank and explicit sender sets +share one capability contract, and list, deep, and forced-reuse locator alias policies share one scatter-copy contract. +Repeated assertions inside dense-owner and filtered-buffer loops are deleted. Flush-cache preservation and weight-version +propagation now reach the actual completion payload together rather than one test stopping at backend configuration. + +Finally, real multi-sender initialization already adopts and validates a nonzero rank's scattered tensor map, and the +all-filtered direct-EP transfer already uses a nonzero source rank with a real locator and proves no engine transfer occurs. +Their direct state-assignment and empty-bucket no-exception smokes are removed. The file collects 68 tests instead of 78, +and all 68 pass. Repository collection is now 2,433 items and the static inventory is 2,389 definitions with 143 curated +decisions. + +## Thirty-seventh wave: carry endpoint configuration through consumers + +The thirty-seventh wave turns worker-port selection into one registration-to-consumer transaction. The explicit worker +registration still checks both control and worker health, then uses the returned endpoint for LoRA load, unload, and loaded +adapter discovery. Two tests that constructed an endpoint by hand only to repeat the final URL are removed. + +FP8 receiver detection now reports policy rather than one config literal per test. The retained rich skip-list transaction +already proves default format, dynamic activation, block size, language-model prefix normalization, weight-suffix removal, +and vision exclusion, so a minimal default-dictionary snapshot is removed. MTP, invalid activation schemes, UE8M0 scale +storage, and BF16 MTP form one receiver-admission contract. Compressed-tensors rejection and explicit BF16 no-op behavior +likewise share the sync-quantization setter boundary. + +Packing drops an exact strategy-tuple snapshot because every strategy still executes through document, token, position, +capacity, utilization, balance, order, and determinism contracts. A convenience-wrapper smoke is also removed: the retained +full pipeline calls `pack_samples`, asserts packed boundaries, simulates output, and unpacks every sample, while direct packed +and unpacked contracts cover both mode branches. The three focused files collect 62 tests instead of 71, and all 62 pass. +Repository collection is now 2,424 items and the static inventory is 2,380 definitions with 146 curated decisions. + +## Thirty-eighth wave: report admission policies, not success and failure fragments + +The thirty-eighth wave resolves six no-observable-outcome signals by attaching each success path to its failure policy. +Pipeline schedules now map style, single-stage status, and split-backward behavior in one metadata table, while all admitted +and rejected virtual-stage/microbatch configurations enter one schedule-admission contract. This preserves every schedule +and validation branch while removing three separate reports. + +DeepEP's preflight now accepts an identity dispatch/combine roundtrip and rejects corruption after the same setup. Launcher +readiness likewise accepts a set ready event and independently fails fast on worker exit in one lifecycle contract. +Blackwell FP8 policy rejects no override, rejects a missing validation artifact, and accepts an explicit validated override +together. Exact Qwen3.5 MoE admission similarly checks structural defaults before all invalid implementation, dispatch, and +async-combine overrides. + +The exact active-LoRA snapshot guard now rejects both full-weight publication paths before downstream work and then proves +ordinary models remain unrestricted, eliminating a standalone absence-of-error test. The six focused files collect 64 tests +instead of 73, and all 64 pass. Repository collection is now 2,415 items and the static inventory is 2,371 definitions with +150 curated decisions. + +## Thirty-ninth wave: test FP8 quantizers by numerical path + +The thirty-ninth wave removes a weak FP8 output snapshot because the retained independent Slime-reference contract now also +checks CPU placement while already proving emitted names, dtypes, scale shape, exact FP8 bytes, exact scales, and dequantized +parity. A direct contiguous-storage predicate example is also removed: stack-versus-single quantization proves grouping is +numerically transparent, while expert workspace and streaming transactions exercise grouped stacks at their consumers. + +Three CUDA stack tests become one target-and-parity contract. The same nontrivial tensor is quantized to CPU and CUDA +targets, target-specific copy telemetry is checked, and both results are compared bitwise with the CPU quantizer. Expert +projection, skip-list, CPU-workspace, streaming, LoRA folding, partial-block padding, and direct-EP cases stay independent. + +Two-dimensional DTensor save materialization now proves both all-rank and writer-only results inside one four-rank CPU-mesh +transaction. This removes a repeated process-group launch while preserving the distinct one-dimensional mesh contract. The +two focused files collect 48 tests instead of 53, and all 48 pass, including the live CUDA cases. Repository collection is +now 2,410 items and the static inventory is 2,366 definitions with 153 curated decisions. + +## Fortieth wave: replace arbitrary examples with branch-relevant contracts + +The fortieth wave removes twelve reports and several uncollected examples across grouped GEMM, MoE primitives, adapter +checkpointing, adapter coordination, and routing replay. A grouped-GEMM example that allocated roughly a gigabyte for one +aligned operand is replaced by a compact unaligned case that actually reaches the kernel's K/N masking path. Generic +single-group repetitions and a separate dtype/shape smoke disappear; FP16 and BF16 numerics, transpose-B, unequal groups, +zero-K handling, noncontiguous input, and device rejection remain in the two real numerical contracts. + +MoE primitives no longer repeat the same histogram after flattening, test an invalid overlapping slot map, or separately +smoke gather, scatter, and add-gather before their retained numerical references and roundtrip. The multi-block gather is +smaller but now crosses an unaligned hidden-width boundary, and the full routing transaction checks the exact output rather +than merely checking that it is nonzero. Non-gated backend and activation rejections now share one constructor policy. + +Adapter checkpoint save/load containment, missing-tensor rollback, saved dtype, and current learning-rate persistence are +proved at their transaction boundaries rather than by paired fragments. Coordinator load, save, and evicted-model path +traversal share one output-root policy, while direct adapter registration and materialized session registration share one +broadcast transaction. Routing replay combines padded and unpadded sequence-parallel cases, materialized input types, and +nested Qwen top-k discovery with base64 shape inference while retaining packed-document, zigzag, truncation, and NumPy +integration paths. + +The six focused files collect 85 tests instead of 97, and all 85 pass, including the live CUDA kernels. Repository +collection is now 2,398 items and the expected static inventory is 2,354 definitions with 158 curated decisions. + +## Forty-first wave: express propagation as end-to-end policy + +The forty-first wave reduces four previously untouched suites from 51 collected reports to 38. Server optimizer wiring now +has one arguments contract, one initialization contract, one optimizer-step policy, one adapter path, and one dispatcher +boundary. Explicit, default, malformed, partial, omitted, and non-Adam cases still run, but they no longer publish one test +result per dictionary variation. + +Runner save-state and save-LoRA handlers now prove the shared nonresident-adapter materialization rule in one transaction. +Uniform-versus-rank-local forward/backward failure semantics still run on rank-zero and worker paths, but no longer create +two parameterized reports for identical policy. LoRA checkpoint tests fold SGLang layout inspection into its actual +roundtrip and exercise hybrid-shared and all-owner expert layouts through one adapter-manager transaction. + +Native block-FP8 state construction and dtype application now form one byte-exact lifecycle contract. CPU execution, +explicit materialization admission, and import laziness likewise share one fail-closed contract. Kernel dispatch, gradient +admission, invalid prequantized pairs, state dictionaries, DCP metadata, exception restoration, and FSDP composition remain +independent because they guard different mutation or integration boundaries. + +All 38 focused tests pass. Repository collection is now 2,385 items and the static inventory is 2,342 definitions +with 162 curated decisions. + +## Forty-second wave: delete test-only proofs and retain production contracts + +The forty-second wave reduces four untouched suites from 43 reports to 26. Adapter-gradient ownership configuration, +fingerprint invariance, and replica-domain admission now report policy transactions rather than one result per example. A +self-contained "analytical reference" test is removed entirely: it defined scale/clip/AdamW equations inside the test and +checked them against PyTorch, while the retained adapter-manager transaction already validates the production optimizer +step, parameters, moments, clipping, scratch lifecycle, and publication state against those equations. + +DRGRPO now checks seeded forward value, gradient norm, finiteness, and metric schema in one numerical contract. Zero +advantages, fully ignored labels, and empty sequences share one zero-loss boundary; positive-KL admission and effect share +another. Temperature/K3 behavior, advantage direction, and microbatch composition remain independent. + +Outbound endpoint allowlisting, DNS pinning, and malformed-target rejection now form one security policy. Generic path +resolution and environment/explicit artifact roots form another. Diagnostic inputs and compile-worker code/protocol +boundaries remain separate. Router GEMM reference, empty input, and dtype admission now share one kernel contract, as do +top-k renormalization, cast-only behavior, and input dtype; batch invariance, backward, leading dimensions, and MoEBlock +consumer paths remain independent. + +All 26 focused tests pass, including live CUDA. Repository collection is now 2,368 items and the expected static inventory +is 2,325 definitions with 166 curated decisions. + +## Forty-third wave: make instrumentation tests follow event lifecycles + +The forty-third wave reduces phase, component, CUDA, and activation-offload instrumentation from 22 reports to 9. Phase +ordering now covers canonical, custom, and empty inputs together. Phase-time and memory summaries each prove empty behavior, +ordering, float normalization, and single-rank aggregate fields through one complete output map. + +Direct decoder-name and nested-submodule helper tests are removed. The retained live GLM/Qwen hook transaction discovers +three decoder layers, records present forward/backward components, and omits absent indexer/shared-expert components through +the actual timer consumer. Disabled and malformed manual-CUDA modes share one API policy; successful forward/recompute +events and an unrecorded pair share one drain lifecycle. Activation-offload values and empty-context omission likewise +execute through one consume lifecycle. + +All 9 focused tests pass, including live CUDA. Repository collection is now 2,355 items and the expected static inventory +is 2,312 definitions with 170 curated decisions. + +## Forty-fourth wave: report checkpoint and sharding selectors as policies + +The forty-fourth wave reduces four checkpoint/sharding suites from 28 reports to 19. DCP synchronization now expresses +NCCl-to-Gloo creation/caching and Gloo-default behavior in one backend policy. Metadata synchronization expresses disabled +PP, caller-supplied PP group, and global Gloo fallback in one selector contract. + +The meta EP slicing regression now uses the real `already_local=True` skip-loading path while checking local shape, meta +device, dtype, requires-grad, shard metadata, and unrelated replication together. Two malformed exact-GLM dispositions +still execute but no longer generate parameterized duplicate reports. Exact meta allocation, real already-local factor-bank +sharding, reduction metadata, and indivisible geometry remain independent. + +A manager active-slot shape smoke is removed because the retained manager registration/forward tests already assert +rank-specific local factor shapes and the sharded-state suite retains pack/unpack, coordinate initialization, discovery, +and a real two-rank uneven DTensor transaction. Checkpoint zero-meta failures now cover pre-load and post-restore stages in +one lifecycle, while initial checkpoint optimizer-default and weights-only modes share one runner policy. + +All 19 focused tests pass, including the two-rank Gloo transaction. Repository collection is now 2,346 items and the +expected static inventory is 2,304 definitions with 174 curated decisions. + +## Forty-fifth wave: separate MoE TP modes from backend contracts + +The forty-fifth wave reduces the largest untouched suite, MoE tensor-parallel simulation, from 14 reports to 9. Environment, +no-EP/TP1 admission, EP rejection, and layer filtering now form one policy. Direct accumulation, BF16 shard reduction, +cache reduction, and a TP1-to-TP2 simulation override run against their corresponding independent references through one +eager-mode matrix instead of rebuilding identical experts and routing data four times. + +Carried reshaped shard metadata and flat diagnostic captures now come from the same MoEBlock execution and are both checked +against the reconstructed sum. Triton, Triton-plus-SGL reduction, DeepGEMM, SGLang fused-expert layout, and SGLang runner +contracts remain independent because each invokes a different backend interface. The ordinary MoEBlock consumer also +remains independent of the expert-level simulations. + +All 9 focused tests pass. Repository collection is now 2,341 items and the expected static inventory is 2,299 definitions +with 176 curated decisions. + +## Forty-sixth wave: require DeepSeek-V4 tests to prove behavior + +The forty-sixth wave reduces the untouched DeepSeek-V4 MoE suite from 11 reports to 6. A shared-MLP smoke that explicitly +declined to assert any clamp effect is removed; the retained forced-gate case proves bounded output and now also checks the +forward shape. Routed-expert clamp propagation remains separate because it crosses the MoE expert backend. + +Non-hash constructor properties, forward/backward numerics, selection-only bias behavior, and shared-expert contribution +now form one model transaction. Hash-layer table/bias structure, missing-input rejection, table-driven forward, and gate +gradient behavior likewise form one transaction. Hash record-to-replay backward and unknown replay-stage failure remain +independent because they protect replay state transitions rather than ordinary routing. + +All 6 focused tests pass. Repository collection is now 2,336 items and the expected static inventory is 2,294 definitions +with 178 curated decisions. + +## Forty-seventh wave: turn selection fragments into consumer transactions + +The forty-seventh wave reduces FP8 LM-head loss, native-EP combine, DeepSeek-V3 checkpoint conversion, and DR-GRPO runner +coverage from 34 reports to 18. Per-token CE now selects the module and FP32-master paths locally and under TP in one +transaction, while temperature is carried through the primitive and CausalLM consumer together. Importance sampling, +CausalLM FP32 bypass, and TP hidden-gradient reduction remain separate consumer boundaries. + +Native Qwen3.5 combine now reports EP8 admission, exact structural flags, and missing trainer-EP rejection as one policy. +Variable token padding and backward unpadding, invalid expert-ID padding, and maximum-row selection likewise share one +collective contract. FSDP pre-forward routing, the serving fused-gate gradient, and full operand capture remain independent. + +DeepSeek-V3 external merge, internal fused pass-through, save splitting, and multimodal filtering now form one layout +conversion. Dense and packed EP slicing run through one policy; default and requested packed dtypes share one load +transaction; official layout recognition carries through text quantization-config parsing. DR-GRPO legacy input, +temperature, disabled per-token output, and K3-forced output now share one runner option contract without absorbing its full +dispatch, sampler-boundary, or forward-backward integrations. All 18 focused tests pass. + +## Forty-eighth wave: remove analytical echoes and report optimizer/router lifecycles + +The forty-eighth wave reduces SignSGD, CausalLM Z-loss, and MoE train-router coverage from 21 reports to 11. One SignSGD +step now proves sign updates, decoupled decay, zero-sign behavior, and missing-gradient preservation. Multiple steps and +state-dict restoration form one stateless persistence lifecycle, while sparse-gradient rejection and optimizer-factory +grouping remain distinct. + +The random Z-loss reference transaction now proves both forward values and backward gradients. A zero-logit example that +only restated the test's analytical formula is removed, as is a duplicate CausalLM temperature report already retained at +the LM-head consumer. Coefficient-zero behavior, compiled CUDA parity, and TP rejection remain separate. Train-router +all-to-all gradients and DeepEP rejection now share one dispatch policy; argument/model defaults and balanced +forward/replay routing each report their full policy rather than individual examples. All 11 focused tests pass. + +## Forty-ninth wave: join repack and numerical forward/backward contracts + +The forty-ninth wave reduces shared-prefix repacking, multi-part optimization, and batch-invariant fused LM-head coverage +from 20 reports to 13. Shared-prefix detection, exact token and position repacking, decoded loss-field preservation, and +output remapping now execute as one end-to-end transaction. No-sharing and one-token-prompt boundaries remain independent +because they select different backend outcomes. + +MultiOptimizer construction now proves DCP model mapping and complete parameter coverage together. Its live lifecycle +updates every virtual part, clears gradients, and decays every learning-rate group. Single-part selection and invalid custom +groups remain separate admission outcomes. Fused LM-head forward values and backward gradients now compare against eager +from the same graph for both default and non-unit temperature paths. Determinism, unsupported options, unit-temperature +bytes, and the probability-one clamp remain distinct numerical contracts. All 13 focused tests pass, including all six +live CUDA cases. Repository collection is now 2,303 items and the static inventory is 2,261 definitions with 188 curated +decisions. + +## Fiftieth wave: compile configuration and ownership as matrices + +The fiftieth wave reduces model-runner builder propagation, exact GLM gradient ownership, batch conversion, and batch-slice +selection from 18 reports to 9. Fail-closed FP8 defaults, every QARL calibration field, and sharded LM-head loss now pass +through one runner initialization policy. Exact GLM block-FP8 QLoRA remains separate because it also resolves the runtime +target-module set, and raw CausalLM token-sum behavior remains a distinct loss boundary. + +Exact gate-up, dense-MLP, and absorbed-KV leaves now compile through one module-managed ownership matrix with their +component-specific canonical factor names. The EP16 routed leaf rejects missing managed-FSDP ownership and mutated DeepEP +dispatch in one runtime-admission contract. DR-GRPO logprobs and teacher hidden states share one FP32 batch-conversion +transaction, while ragged padding and sequence sharding stay separate. Finally, FSDP, TP, EP, and legacy duplicated-EP +slices are expressed as one topology selector instead of four examples. All 9 focused tests pass. + +## Fifty-first wave: follow LoRA generations and model dtype lifecycles + +The fifty-first wave reduces fused-GDN LoRA, DeepSeek-V4 attention, stochastic rounding, and fused MoE expert coverage from +26 reports to 20. Canonical sliced LoRA folding now carries through the GDN output-projection consumer and its gradients. +Cache slice bounds and release of the previous serialized-request generation execute in one parameter-version lifecycle. + +DeepSeek-V4's attention sink now proves initial FP32 storage and FSDP marking, BF16 module conversion, and FP32 promotion at +the TileLang call boundary together. Stochastic rounding reports its output metadata and non-FP32 rejection as one API +contract while retaining all statistical and deterministic numerical properties. Base and LoRA expert modules share one +fused gate/up registration contract, and Qwen3/Qwen3.5 handlers share one deferred-QLoRA skip policy while preserving their +family-specific key layouts. All 20 focused tests pass. +Repository collection is now 2,288 items and the static inventory is 2,246 definitions with 196 curated decisions. + +## Fifty-second wave: pair numerical values with gradients and policies with consumers + +The fifty-second wave reduces batch-invariant GDN, Nemotron-H, exact GLM synchronization, deferred QLoRA loading, trainer +model-program selection, and P2P transfer coverage from 36 reports to 23. GDN gating and gated RMSNorm each compare forward +values and every input gradient with an independent PyTorch composition from the same graph. Nemotron's output shape, +optional router logits, labeled CausalLM loss, and gradients through Mamba, attention, routed/shared MoE, latent, and +embedding paths now form one model transaction. + +Exact dense MLP and attention projections reject adapter preparation, collective merging, and raw extraction through one +factor-only synchronization matrix. Exact LM-head ordinary and prepacked sparse-delta publication share one pre-side-effect +guard. QKV and gate/up deferred loader key selection now report one merged-projection policy rather than two parameterized +examples. Exact/ordinary server and non-server trunk engagement, structural numerical-family selection, and P2P size cutoff +selection likewise execute as policies. All 23 focused tests pass, including all five live CUDA GDN tests. + +## Fifty-third wave: remove generic serializer examples and report numerical matrices + +The fifty-third wave reduces FWHT, runner message protocol, DeepSeek-V3 auxiliary routing, FlashMLA admission, RL +primitives, and families-v2 coverage from 33 reports to 23. A known Hadamard row, orthonormal roundtrip, and norm +preservation now share one width matrix. Typed runner messages, tensor payloads, JSON conversion, and pickle rejection share +one wire-format contract; generic large nested-list and nested-dictionary serializer examples are removed while IDs, +timestamps, optional fields, and ACK behavior remain. + +DeepSeek-V3 auxiliary router logits now prove both all-MoE output and omission of a dense prefix together. FlashMLA rejects +CPU dispatch, an unproven head shape, and flattened-address overflow before backend import in one admission policy. All KL +estimator modes still compare against independent Slime formulas but no longer publish four parameterized reports. Exact +and nonexact families-v2 selection likewise share one environment policy. All 23 focused tests pass, including the live +CUDA family-byte and batch-invariance cases. + +## Fifty-fourth wave: delete dead placeholders and collapse selector fragments + +The fifty-fourth wave reduces routing regather, tensor collation, runner session/load-state behavior, FLOP counting, and +Kimi target resolution from 22 reports to 10. Sqrt-softplus replay regather now proves eager parity, routed scaling, +requested dtype, and cached expert identity together while keeping the softmax regression separate. Tensor container, +scalar, dtype, empty, and string inputs form one conversion policy; variable-length and packed layouts remain independent. + +LoRA registry synchronization now follows optimizer and checkpoint-load mutations in one lifecycle. DistSignSGD scaling, +clip suppression, and optional CUDA-cache suppression share one optimizer policy. Load-state preparation carries rank-zero +errors and artifact-root containment together, while multi-adapter and single-tenant dispatch share one routing contract. +Two permanently skipped GLM5 FLOP tests are removed because sparse-MLA/DSA accounting is not implemented; executable CP +length invariance remains. Kimi wrapper defaults, explicit targets, and strict manifests now form one target-source +precedence contract. All 10 focused tests pass. +Repository collection is now 2,253 items and the static inventory is 2,215 definitions with 214 curated decisions. + +## Fifty-fifth wave: delete mirrored implementations and follow real lifecycles + +The fifty-fifth wave reduces distributed state, model registration/LoRA, side-payload, scheduler, batch-conversion, and +DR-GRPO coverage from 39 reports to 23. Most importantly, two MoE auto-merge files are removed in full: they copied the +parser, buffer, format detector, and simulated loader into the tests and never imported XoRL. Production buffer +transposition/output and family checkpoint loading remain covered by real handlers. + +Parallel-state publication now carries automatic shard inference and reinitialization rejection through one singleton +lifecycle. Kimi wrapper conversion includes official auxiliary-loss defaults in one config policy, while tokenizer and +processor fallback share one remote-code security contract. DeepSeek-V4 attention-LoRA type/freeze checks now live in its +forward-backward transaction; the stale generic expert-LoRA claim is gone because DeepSeek-V4 expert semantics are +explicitly unsupported by that generic wrapper. + +FIFO defaults and queue operations now share one scheduler policy, and a sleep-based direct state-object smoke is removed +because completion, failure, and abort already run through Scheduler. Missing side-payload keys join typed store +roundtrips, ragged teacher-state padding joins float conversion, and DR-GRPO legacy/options join its loss dispatch. All 30 +focused production tests pass, including the retained expert-buffer and family checkpoint-handler consumers. + +## Fifty-sixth wave: make codec and topology matrices report once + +The fifty-sixth wave reduces six suites from 39 reports to 15. NF4 accounts for the largest change: twenty codebook, +shape, dtype, group-width, accuracy, zero, and layout reports become three codebook, flat-codec, and GKN-codec contracts. +Every 32/64/128 group width still runs, while two 16M-element/14M-element allocation smokes that repeated the same codec +behavior are removed. + +A local GKN matrix/reference-MoE proof and a future sparse-MLA KV-major reference scaffold are removed because neither +invoked production. Production expert buffering and eager/native/Triton backends remain, as do sparse-MLA forward, +backward, deterministic, and combined-kernel contracts. EP supported/unknown backends now share one fail-closed table; +four-rank CP, DP, and HSDP lm-head meshes share one topology matrix; and zero/nonzero cast-once cases share one transaction +for each linear and grouped-expert implementation. + +All seven live CUDA codec/merge/backend tests pass. The real two-rank EP reduction and all three four-rank lm-head +topologies pass as well; the four untouched production sparse-MLA reports remain collected. +Repository collection is now 2,213 items and the static inventory is 2,183 definitions with 229 curated decisions. + +## Fifty-seventh wave: turn router and loss fragments into policies + +The fifty-seventh wave reduces TopKRouter, OPD parity, and reducer coverage from 42 reports to 19. Softmax, balanced, +hash, sqrt-softplus, scaling, invalid-input, and config behavior now execute as selector policies while FP32 selection and +the MoEBlock consumer remain distinct precision and integration boundaries. OPD full-vocabulary modes, estimators, +dispatch, policy-gradient admission, and task weighting likewise report complete policies; clamp behavior, stable metric +keys, and ignored-label handling stay independent. + +TokenPartial now proves denominator and microbatch composition together. SequencePartial covers dense, packed, and +context-parallel layouts as one composition policy, with empty-input zero behavior retained separately. All 19 focused +tests pass. + +## Fifty-eighth wave: execute numerical case matrices without multiplying reports + +The fifty-eighth wave reduces seven shared-loss, RoPE, adapter, and shared-prefix suites from 42 collected items to 25 in +this environment without dropping their numerical cases. A fully provisioned FlashAttention 3 environment also avoids +seven extra parameter reports from the shared-prefix dtype/head-shape product; here that whole optional module already +collects as one dependency skip. Legacy reducer identities, paired KL/K3 implementations, and all importance-sampling and +policy-loss microbatch variants now loop inside one semantic contract apiece. Dense and MoE Qwen half-rotate/cast behavior, +Class-B shapes, supported dtypes, and installed EP adapters use the same pattern. + +All 14 CPU numerical reports pass, and all 10 EP-adapter reports pass on CUDA. The shared-prefix module is lint- and +collection-safe but its optional FlashAttention 3 interface is unavailable in this environment, so that module remains a +dependency skip rather than a claimed runtime pass. + +## Fifty-ninth wave: move machine-specific throughput out of correctness tests + +The fifty-ninth wave removes three Qwen3-8B TFLOPS threshold reports. They were two-to-three-minute benchmark jobs gated +by local model-directory presence and hard-coded H100 baselines, but they did not verify that the executing device was an +H100. Their secondary loss-decrease assertion duplicated retained Qwen LoRA/FSDP end-to-end training coverage. Hardware +performance regression belongs in a controlled benchmark lane with explicit machine admission, not the portable pytest +correctness inventory. + +Repository collection is now 2,170 items and the static inventory is 2,157 definitions with 236 curated decisions. + +## Sixtieth wave: replace existence smokes with production policies + +The sixtieth wave reduces eight data, QLoRA, model, distributed, serialization, weight-sync, and optimizer suites from 72 +collected items to 56. A direct `pack_parallel` report is removed because it asserted only that output was non-empty; +`PackingDataset` still exercises the production path for sequential and multipack operation, while capacity, coverage, and +allocation are checked independently. A QLoRA package `hasattr` smoke is also removed: trainer/model-builder imports and +real adapter consumers cover the public package, and the clean-interpreter dependency-cycle contract remains. + +Default and explicit MLA targets now form one partition policy. CP16 first-rank and padded-tail behavior, both directions +of Muon EP checkpoint resharding plus same-EP identity, replicated/sharded/padded DTensor copies, 1-D/2-D save +materialization, empty/nonempty PP NCCL transfers, and ordinary/Kahan cautious-decay behavior likewise execute as matrices +instead of one report per example. The private `_prod([])` example is replaced by scalar reconstruction through the real +PP receive path. + +All 55 focused CPU reports pass, including the two four-rank DTensor materialization workers. The consolidated Muon +checkpoint report separately passes all three four-rank save/load transitions. + +Repository collection is now 2,154 items and the static inventory is 2,143 definitions with 244 curated decisions. + +## Sixty-first wave: join architecture admission and configuration policies + +The sixty-first wave reduces twelve architecture, loader, optimizer, diagnostics, distributed-policy, and timing suites +from 78 reports to 55. DeepSeek-V4 now carries one standard snapshot through Transformers AutoConfig, AutoModel +registration, and XoRL model construction. DeepSeek-V3, Nemotron-H, and Qwen3.5 registry assertions join their real +configuration conversions. A direct `ModelArguments` field-default test is removed because omitted/explicit YAML +serialization, trainer propagation, and runtime resolution already prove the production path. + +Gradient-checkpoint method propagation, single/multi-part optimizer selection, and token-diagnostic disabled/ignored/top-k +boundaries now report complete policies. DeepSeek-V4 APE, FP8, and MXFP4 examples become one contract per codec, with +unknown-key behavior moved from the private mapper to the checkpoint handler and its unmapped ledger. DeepEP default and +unsafe-opt-in combine behavior likewise form one admission contract. + +FSDP policy coverage accounts for the largest reduction: singleton/sharded/overridden reduce dtypes, supported/rejected +boolean spellings, and backward-only/forward-only/bidirectional/disabled prefetch settings all still run, but thirteen +reports become six policies. Disabled and unrecorded component-timer behavior join one lifecycle while live CUDA hooks stay +separate. All 55 focused tests pass. + +Repository collection is now 2,131 items and the static inventory is 2,120 definitions with 253 curated decisions. + +## Sixty-second wave: make server configuration admission policy-shaped + +The sixty-second wave reduces server argument and weight-sync quantization coverage from 42 reports to 29. R3 payload +transport success and directory rejection now form one admission matrix, exact GLM rank-1 topology covers accepted and +rejected shapes together, and QARL/FP8 full-weight conflicts plus MTP, Mamba, and Nemo source rejection execute as one +fail-closed policy. Nested train and Nemo FP8 aliases now share one normalization contract, while the two unsupported +multi-adapter modes share one rejection policy. + +Weight-sync FP8 configuration now reports normalization once across explicit values, defaults, and module-name cleanup. +Unsupported formats, activation schemes, scale storage, module exclusions, and internal unsupported markers likewise +form one invalid-configuration matrix. All original configuration inputs and exact rejection boundaries remain exercised; +the reduction removes report fragmentation rather than behavior. + +All 29 focused server configuration reports pass. Repository collection is now 2,118 items and the static inventory is +2,107 definitions with 260 curated decisions. + +## Sixty-third wave: consolidate admission and lifecycle policies across subsystems + +The sixty-third wave reduces four independent CLI, simulator, numerical-contract, and weight-sync clusters by 16 reports. +Training CLI FP8 aliases now share one normalization contract, while adapter conflicts, missing QARL calibration, mutual +QARL/FP8 exclusion, MTP metadata, and Mamba configuration form one full-weight low-precision rejection policy. Every YAML +payload and exact rejection boundary still passes through `parse_args`. + +Simulator filesystem admission now exercises built-in traversal, relative escape, missing defaults, symlink escape, and +unapproved model-metadata paths as one security policy. Exact Qwen3.5 topology admission joins three certified shapes with +ten rejected mutations; its model-scope policy joins dense, MoE, and Hugging Face snapshots with wrong-layer and nearby- +geometry rejection. P2P completion now reports two lifecycle policies covering pending-transfer failure, receiver +suppression, best-effort cleanup, completion payload metadata, completion failure, deregistration, and default cache +behavior. + +All 119 focused reports pass, including all 64 remaining P2P backend protocol reports. Repository collection is now 2,102 +items and the static inventory is 2,091 definitions with 266 curated decisions. The heuristic candidate count falls from +63 to 61 without adding parse errors or duplicate bodies. + +## Sixty-fourth wave: turn server lifecycle examples into endpoint policies + +The sixty-fourth wave reduces three session API, request-processing, and adapter-management suites from 91 reports to 66. +LoRA registration now covers full and rank-only overrides plus existing-session refresh as one policy. Reserved checkpoint +creation, per-session isolation, stale replacement, and existing preservation share one persistence contract. Full-weight +default admission and multitenancy/override rejection likewise report once, while kill, checkpoint URI, default-session +protection, and weights-info modes are grouped by endpoint lifecycle. + +Request routing cleanup now covers Mooncake success, Mooncake backend failure, and filesystem payloads together. Token +diagnostics cover packed splitting, empty input, and malformed lengths as one decoder policy; packed-row batching covers +global grouping, rank-local deferral, and replay rejection as one batching policy. Adapter management joins abort states, +direct checkpoint-plan admission, malformed checkpoint structures, PEFT filename/sharding compatibility, and dirty/clean/ +failed/multi-rank eviction outcomes into their respective policies. + +All 66 focused reports pass. Repository collection is now 2,077 items and the static inventory is 2,066 definitions with +279 curated decisions. The candidate inventory remains 61, with no parse errors and the one intentional duplicate-body +group unchanged. + +## Sixty-fifth wave: join inference, GLM5, and optimizer-resume policy matrices + +The sixty-fifth wave reduces three inference API, GLM5 model, and adapter optimizer-resume suites from 83 reports to 58. +Inference endpoint port selection, FP8 KV-cache admission, sync-pool filtering, quantization admission, cache invalidation, +receiver detection, and receiver enrichment now each report one complete policy rather than one report per input example. +All health routes, payload flags, endpoint epochs, HTTP failures, normalization, and unsupported-receiver outcomes remain. + +GLM5 local config security now joins non-JSON and dunder-key rejection. Blocked indexer selection and sparse-MLA reference +coverage each join full-query and query-offset parity, while CPU fallback and unknown sparse-MLA backends form one dispatch +policy. Adapter optimizer resume now joins canonical identity, checkpoint artifact admission, four successful logical +reshards, and the complete invalid-source/no-mutation matrix into four policy reports. Bitwise moment, squared-moment, step, +and resident-state assertions are unchanged. + +All 58 focused reports pass. Repository collection is now 2,052 items and the static inventory is 2,041 definitions with +294 curated decisions. The candidate inventory remains 61, with no parse errors and the intentional duplicate group +unchanged. + +## Sixty-sixth wave: make handler, FP8 linear, and export coverage policy-shaped + +The sixty-sixth wave reduces weight-sync handler, FP8 linear, and quantized-export suites from 73 reports to 44. Receiver +postprocessing, direct-EP ownership, tied parameters, Nemotron conversion, expert gating, compile-wrapper normalization, +and sparse-delta sync now each report one complete handler policy. Every environment override, emitted backend flag, +receiver namespace, tensor split, transport event, rejection, and cache result remains exercised. + +FP8 injection now reports core replacement, recipes, and exclusions as three policies. CPU fallback covers numerical, +output-dtype, and fail-fast behavior together; profiler selection joins call caps, explicit rows, and module-specific calls; +block-FP8 GEMM joins block, rowwise, and torch-scaled-mm reference parity. The live CUDA operand profiler, automatic +fallback, padding, correction, and training-step gates remain independent. Offline export similarly joins quantization, +fused-QKV, MLA-A, linear-attention, and QARL fold admission cases without merging the trained-logprob proof. + +All 44 focused reports pass, including the available CUDA FP8 gates. Repository collection is now 2,023 items and the +static inventory is 2,012 definitions with 312 curated decisions. The candidate inventory remains 61, with no parse errors +and the intentional duplicate group unchanged. + +## Sixty-seventh wave: cross below two thousand with packing and numerical policies + +The sixty-seventh wave reduces packing strategies, core packing, Muon, and FP8-MoE suites from 78 reports to 38. Packing +strategy admission now covers invalid settings and every oversized mode together. Cross-strategy document/token/position/ +capacity/utilization invariants, balanced-DP behavior, and deterministic datum ordering each report one complete policy. +Core token metadata joins OPD, OPRD, HF shift, vector padding, RL padding, cache views, and nested schemas, while disabled +packing joins all target-preservation and warning behavior. + +Muon grouping now executes equal, flattened, transposed, fused, and chunked shapes as one policy; fused gate-up identity +joins gated/non-gated, post-FSDP, and model-family classification. FP8 MoE injection, same-NK forward, same-MN gradient, +scalar-Quack forward, and expert training now report once per semantic boundary while retaining every CUDA numerical case. +The DeepGEMM subprocess remains an independent opt-in backend gate. + +Focused validation reports 37 passed and the existing DeepGEMM opt-in skip. Repository collection is now 1,983 items and +the static inventory is 1,972 definitions with 325 curated decisions. The candidate inventory remains 61, with no parse +errors and the intentional duplicate group unchanged. + +## Sixty-eighth wave: collapse execution examples into end-to-end policies + +The sixty-eighth wave reduces merged-LoRA, API compatibility, and training-simulator suites from 63 reports to 27. +Canonical folding, straight-through gradients, merged linear selection, MoE cache/admission, native-EP routing, and trunk +wrapping now report once per semantic boundary while retaining every exact tensor, gradient, cache-identity, and rejection +check. API session lifecycle, Tinker weights compatibility, worker registration, and optimizer payload/LR resolution use +the same policy shape; heartbeat activity is now deterministic instead of relying on a wall-clock sleep. + +The simulator now has eleven durable boundaries rather than 25 narrow scenario reports: topology accounting, observed-log +planning, model metadata, calibration evaluation, calibrated scenarios, topology what-ifs, path security, built-in packs, +analytical ledgers, kernel admission, and whole-simulator validation. Repeated fixtures now feed complete ingestion-to- +decision paths, and all original measured rows, OOM boundaries, extrapolation flags, topology candidates, and analytical +terms remain asserted. + +All 27 focused reports pass. Repository collection is now 1,947 items and the static inventory is 1,936 definitions with +344 curated decisions. The candidate inventory remains 61, with no parse errors and the intentional duplicate group +unchanged. + +## Sixty-ninth wave: make cache transport and trainer utilities policy-shaped + +The sixty-ninth wave reduces teacher-head/cache, Mooncake transport, and trainer utility suites from 41 reports to 14. +Teacher-head persistence now covers direct files, tied embeddings, sharded manifests, and cross-shard views as one policy; +manager residency covers teacher replacement, dtype replacement, and prefetch. Activation selection joins rank-2, rank-3, +layer-slice, host/device, reuse, dtype-reload, async, and bounds cases into selection and admission policies without losing +any tensor comparisons. + +Mooncake now reports tensor codecs, hidden transport, activation-cache consumption, metadata admission, and store lifecycle +once each. All four dtypes, rank-2/rank-3 layouts, multi-teacher routing, malformed payloads, missing and mismatched objects, +cleanup keys, and environment precedence remain exercised. Trainer utilities similarly join clipping modes, token/voter +accounting, SP and adapter-owned gradient synchronization, and lm-head TP synchronization while leaving PP chunked CE as +its own numerical boundary. + +All 14 focused reports pass. Repository collection is now 1,920 items and the static inventory is 1,909 definitions with +357 curated decisions. The candidate inventory remains 61, with no parse errors and the intentional duplicate group +unchanged. + +## Seventieth wave: consolidate checkpoint transport, optimizer, and GDN contracts + +The seventieth wave reduces checkpoint loading, cautious optimizer, and GDN convolution suites from 49 reports to 22. +Checkpoint behavior now reports object transport, rank-zero loading, source resolution, expert routing, group fallback, +and strict postprocessing as complete policies. DTensor copying/materialization and expert-key classification remain +independent boundaries; every source name, handler call, transfer, dispatch, fallback, and strict diagnostic remains. + +Cautious decay now reports primitive/SignSGD behavior, AnyPrecision decay, AnyPrecision state strategies, Muon behavior, +and builder admission once each. The environment-sensitive DTensor state-offload lifecycle stays separate. All ordinary, +masked, chunked, Kahan, gradient-reuse, Newton-Schulz, fallback, optimizer-family, and kwarg cases retain their numerical +references or exact rejection messages. + +GDN CUDA coverage now reports forward, backward, and end-to-end block policies across fixed, variable-length, batched, +deterministic, eager-parity, and gradient cases. CPU coverage reports packed construction/routing, unsupported-input +admission, and exact-contract lifecycle; the optional SGLang tree comparison remains an independent dependency gate. +Focused validation reports 21 passed and that existing optional skip. Repository collection is now 1,893 items and the +static inventory is 1,882 definitions with 373 curated decisions. The candidate inventory remains 61, with no parse +errors and the intentional duplicate group unchanged. + +## Seventy-first wave: join checkpoint state, numerical programs, and OPD runner policies + +The seventy-first wave reduces model-state compatibility, RoPE/numerical configuration, and OPD runner suites from 52 +reports to 21. Checkpoint state now reports reference/QARL buffers, pipeline key unions, LoRA compatibility, load metadata, +load groups, save groups, and optimizer-state filtering as seven complete policies. Every persistent buffer, union key, +load mode, process-group identity, DCP flag, and per-optimizer state selection remains asserted. + +RoPE and numerical-program coverage now reports Class-B selection, canonical GLM resolution, and exact Qwen3.5 resolution +once each while retaining independent MoE, topology, and model-scope admission. All ordinary defaults, certified fields, +explicit opt-outs, incompatible overrides, CE modes, and family-specific RMSNorm rejection cases remain. + +OPD runner coverage now reports metric aggregation, packed shaping, loss execution, cache contributors, distributed cache +assembly, Mooncake producer/consumer integration, and debug artifacts. Empty-rank collective alignment, per-teacher cache +masking, lm-head anchoring, CP/EP ownership, SP and DP gathering, valid-label trimming, cache indices, gradients, timings, +and JSONL provenance remain exercised. All 21 focused reports pass. Repository collection is now 1,862 items and the +static inventory is 1,851 definitions with 390 curated decisions. The candidate inventory remains 61, with no parse +errors and the intentional duplicate group unchanged. + +## Seventy-second wave: consolidate adapter ownership, RMSNorm, and PP profiling + +The seventy-second wave reduces adapter coordination, RMSNorm family contracts, and pipeline profiling from 49 reports to +17. Adapter coordination now reports auto-load, explicit load/path admission, rank-zero routing, sharded restore, +transactional load failure, registration, and save admission. Checkpoint/fresh paths, synchronized errors, EP shard bytes, +session/topology mismatches, optimizer rejection, PP rejection, worker failure, and every rollback remain exercised. + +RMSNorm now reports family admission, Qwen site declarations, declaration tripwires, family-funnel parity, and module +dispatch. All CPU rejection cases and live CUDA warning, required-family, legacy parity, zero-centered, vitality, trunk, +and three-mode bitwise cases remain. Pipeline profiling reports interval union, schedule formulas, P2P byte accounting, +patch lifecycle, and live GPipe events as five boundaries instead of fourteen examples. + +All 17 focused reports pass, including the CUDA RMSNorm policies and the NCCL single-stage GPipe profiler. Repository +collection is now 1,830 items and the static inventory is 1,819 definitions with 404 curated decisions. The candidate +inventory remains 61, with no parse errors and the intentional duplicate group unchanged. + +## Seventy-third wave: consolidate arguments, pipeline planning, and Mamba2 + +The seventy-third wave reduces argument parsing, pipeline-parallel planning, and Mamba2/SSD suites from 41 reports to 12. +Argument parsing now reports optimizer/numerical controls, checkpoint compatibility, FP8 configuration, and low-precision +mode admission. Every packing field, Muon kwarg, legacy alias, resume flag, FP8 alias/default, vLLM rejection, GLM block- +FP8 QLoRA field, QARL normalization, and full conflict matrix remains exercised. + +Pipeline planning now reports FQN partitioning, stage placement, and schedule metadata/admission instead of thirteen narrow +examples. Default and Qwen names, single/virtual stages, pinned and weighted splits, Torch-reference ownership, loop/V +placement, all six schedules, and every infeasible layout remain. Mamba2 now reports HF mixer parity, SSD recurrence, +packed sequence behavior, missing-kernel admission, and optional live-kernel parity while retaining all output and gradient +comparisons. + +Focused validation reports 11 passed and the existing optional `mamba_ssm` CUDA skip. Repository collection is now 1,801 +items and the static inventory is 1,790 definitions with 415 curated decisions. The candidate inventory remains 61, with +no parse errors and the intentional duplicate group unchanged. + +## Seventy-fourth wave: remove diagnostic, fused-loss, and DistSignSGD example noise + +The seventy-fourth wave reduces token diagnostics, fused selected-logprob, and DistSignSGD suites from 41 reports to 21. +Token diagnostics now report selection/boundaries, KL mapping, log-probability cross-checks, hidden summaries, tensor dumps, +dense hooks, MoE hooks, and trusted override loading. Redundant all-index/component-width examples and a native-MoE hook +case that only pinned internal order constants were removed. + +Fused selected-logprob coverage now reports eager numerical parity, input-gradient policy, irregular tail handling, +dispatcher parity, production vocabularies, causal-LM peak memory, and RL-loss integration. Single-row and repeated large- +vocabulary shapes, the unnamed 100000-vocabulary repetition, and the weaker duplicate peak-memory probe were removed. +DistSignSGD now reports update math, sign communication, local/FSDP hook ownership, topology admission, and optimizer +construction; inherited state-dict behavior and an unsupported sparse-gradient example no longer inflate the suite. + +All 21 focused reports pass, including seven CUDA fused-loss policies. Repository collection is now 1,781 items and the +static inventory is 1,770 definitions with 429 curated decisions. The candidate inventory remains 61, with no parse errors +and the intentional duplicate group unchanged. + +## Seventy-fifth wave: consolidate MoE parity, trunk kernels, and MiniMax support + +The seventy-fifth wave reduces SGLang fused-MoE, batch-invariant trunk-linear, and MiniMax M3 suites from 56 reports to 20. +MoE coverage now reports automatic resolution, block dispatch, admission, trainable dispatch, stock-gradient parity, +masked gradients, weight/layout policy, strided adapter behavior, runtime context, strided numerical parity, and real auto- +mode parity. An install-state-dependent import test and a mocked CUDA auto-dispatch example were removed; the retained +real-kernel gate is the stronger automatic-resolution proof. + +Trunk-linear coverage now reports selection/admission, forward parity, backward parity, and global-interpose gradient +policy. Projection selection, idempotence, exclusions, both bias modes, persistent/global bitwise parity, batch invariance, +cuBLAS gradients, dtype rejection, and inference-safe interposition remain. MiniMax now reports configuration/registration, +activation/routing, text runtime/admission, checkpoint ownership, and MSA paging/admission rather than twelve narrow cases. + +Focused validation reports 16 passed and four existing optional SGLang-kernel skips; every available CUDA trunk policy +passes. Repository collection is now 1,745 items and the static inventory is 1,735 definitions with 444 curated decisions. +The candidate inventory is now 57, including 45 conditional-runtime skips and 12 no-observable outcomes, with no parse +errors and the intentional duplicate group unchanged. + +## Seventy-sixth wave: consolidate dispatcher and weight-sync protocol policies + +The seventy-sixth wave reduces runner-dispatcher, P2P backend, and FP8 synchronization suites from 109 reports to 68. +Dispatcher coverage now reports DP/EP/CP distribution, diagnostics, routing-payload transport/security, packing/dummies, +per-token merging, completion rendezvous, routing-weight slicing, and row-batch provenance. All rank ownership, dummy, +filesystem/Mooncake, trust-boundary, logical-order, CP-replica, and provenance assertions remain. + +P2P preparation now reports payload construction, fanout, cached preparation, completion, Qwen3.5 slicing, and engine +construction as complete policies. The large production transfer-layout and direct-EP matrices were deliberately retained +as independent boundaries. FP8 synchronization now reports core numerical behavior, adapter merging, selector policy, +stack/dtype behavior, CPU expert formatting, workspace lifecycle, GDN folding, and GPU parity rather than 24 examples. + +All 68 focused reports pass, including the available GPU FP8 policy. Repository collection is now 1,704 items and the +static inventory is 1,694 definitions with 462 curated decisions. The candidate inventory remains 57, with 45 conditional- +runtime skips, 12 no-observable outcomes, no parse errors, and the intentional duplicate group unchanged. + +## Seventy-seventh wave: consolidate adapter lifecycle, GLM5, and expert QLoRA + +The seventy-seventh wave reduces adapter-manager, GLM5 support, and expert-adapter QLoRA suites from 85 reports to 43. +Adapter coverage now reports optimizer construction/persistence, gradient-plan admission, capture staging/atomicity, +coordinator checkpoint materialization, and session compatibility as complete policies while retaining independent raw- +capture, epoch, publication, optimizer/collective failure, checkpoint structure, eviction, and mixed-adapter boundaries. + +GLM5 now reports configuration/construction, indexer construction/selection, DSA masks, sparse-MLA reference/wrapping, +sparse attention, kv_b adapters, dispatch, checkpoint filtering, LoRA/MoE integration, HF parity, and forward/recompute. +The optional live TileLang and HF-reference gates remain independent. Expert QLoRA now reports backend identity, factor +ownership, semantic preservation, target injection, model-family construction, and fail-closed admission. + +All 43 focused reports pass. Repository collection is now 1,662 items and the static inventory is 1,652 definitions with +482 curated decisions. The candidate inventory remains 57, with 45 conditional-runtime skips, 12 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## Seventy-eighth wave: consolidate server configuration, exact GLM5.2, and EP serving-kernel contracts + +The seventy-eighth wave reduces server-argument, exact GLM5.2, and SGLang EP suites from 72 reports to 34. Server +configuration now reports removed fields, shipped adapters, runtime round trips, R3 transport, quantized training, +parallel topology, unsupported combinations, optimizer/runner compatibility, and model-specific controls. Every prior +YAML/CLI rejection, shipped file, numerical field, alias, topology, conflict, and serialized value remains exercised. + +Exact GLM5.2 coverage now reports layer planning, logical selection, Hadamard transport, fused projection, sampler key +preparation, sparse codecs, native-selector runtime, kernel loading, and canonical routing as complete contracts. The +independent topology, IndexShare lifecycle/FSDP identity, checkpoint bias, two semantic MoE-stack depths, and fail-closed +native boundary remain. EP serving-kernel coverage now reports admission, disabled-mode dispatch, enabled presentation, +compute guards, slot combine, trainable dispatch, and weight presentation; the missing-package and live GPU-gradient gates +remain independent. + +Focused validation reports 32 passed and the two existing environment-dependent skips. Repository collection is now +1,624 items and the static inventory is 1,614 definitions with 505 curated decisions. The candidate inventory remains 57, +with 45 conditional-runtime skips, 12 no-observable outcomes, no parse errors, and the intentional duplicate group +unchanged. + +## Seventy-ninth wave: consolidate exact RMSNorm and native block-FP8 lifecycles + +The seventy-ninth wave reduces SGLang fused-RMSNorm, native block-FP8 linear, and GLM5.2 native-FP8 suites from 38 reports +to 16. Fused RMSNorm now reports CPU fallback, forward bit-exactness, backward parity, model integration, and trunk-family +behavior. Residual/no-residual calls, BF16/FP32, packed 3D shapes, module dispatch, dense Qwen, the pre-summed final norm, +the serving residual tree, the aten interpose lane, and every hidden/residual/weight gradient remain exercised. + +Native block-FP8 linear coverage now reports encoding, execution/partitioning, admission, checkpoint lifecycle, and FSDP +precision policy. Exact FP8 and scale bytes, protected dtype application, lazy imports, hook traversal, gradient and shape +failures, state/DCP metadata, adversarial apply restoration, and EP-restored global shapes remain. GLM5.2 native-FP8 now +reports configuration, dense pair buffering, model construction, canonical routing, frozen expert execution, and expert +checkpoint ownership rather than twelve narrow steps. + +All 16 focused reports pass, including the available CUDA RMSNorm policies. Repository collection is now 1,602 items and +the static inventory is 1,592 definitions with 518 curated decisions. The candidate inventory is now 56, with 45 +conditional-runtime skips, 11 no-observable outcomes, no parse errors, and the intentional duplicate group unchanged. + +## Eightieth wave: consolidate Qwen3.5 norms, generic QLoRA, and OPD objectives + +The eightieth wave reduces Qwen3.5 norm-family, generic QLoRA, and OPD loss suites from 35 reports to 14. Qwen3.5 now +reports family dispatch, site assignment, bit-exact integration, and the family-two residual equation. Exact/ordinary +coexistence, v1/v2 selection, every zero-centered site, layer-zero/final-norm forcing, module/trunk/layer bytes, and the +sampler BI-mean composition remain exercised. + +Generic QLoRA now reports quantized execution, NVFP4 scale/merge behavior, injection, prequantized block-FP8, and optimizer +reset as five lifecycles. Both formats, memory, dequantization, EMA, delta folding, fused QKV, target ownership, gradients, +state rebuild, non-LoRA preservation, and scheduled merge remain. OPD now reports numerical backends, gradients/reduction, +output edges, the hidden-only objective, and OPRD hidden distance. Chunking, streaming/TileLang/low-memory paths, teacher +shards, ignored tokens, per-token results, zero-KL behavior, and fetched layer slices remain checked. + +All 14 focused reports pass after correcting a helper-name collision caught by the first focused run. Repository collection +is now 1,581 items and the static inventory is 1,571 definitions with 531 curated decisions. The candidate inventory +remains 56, with 45 conditional-runtime skips, 11 no-observable outcomes, no parse errors, and the intentional duplicate +group unchanged. + +## Eighty-first wave: consolidate quantized export, model construction, and DSV4 loading + +The eighty-first wave reduces quantized-export, FP8/QARL model-builder, and DSV4 checkpoint-loader suites from 34 reports +to 16. Export now reports primitive quantization, CLI behavior, base-directory output, projection layouts, MoE layouts, +admission, and trained-QARL logprob preservation. Subprocess execution, BF16 islands, sharding, fused QKV, MLA-A, linear +attention, GKN/fused experts, every emitted tensor/scale, preflight failure, and exact post-fold logprobs remain exercised. + +Training-model construction now reports sharded lm-head threading, full-model FP8 construction, GLM5.2 block-FP8 QLoRA, +quantized-mode admission, and QARL injection/calibration. DSV4 loading now reports translation, FP8/MXFP4 codecs, handler +ownership, and synthetic model loading. All name families, APE inversion, EP filtering/fusion, MTP/unmapped accounting, +window/C4/hash variants, nonpersistent RoPE buffers, and representative bytes remain checked. + +All 16 focused reports pass. Repository collection is now 1,563 items and the static inventory is 1,553 definitions with +544 curated decisions. The candidate inventory remains 56, with 45 conditional-runtime skips, 11 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## Eighty-second wave: consolidate exact GLM5.2 QLoRA subsystem contracts + +The eighty-second wave reduces the general GLM5.2 block-FP8 QLoRA, fused gate-up, and TP16 lm-head suites from 30 reports +to 15. General coverage now reports the full inventory, EP-local routed banks, exact dense component, fail-closed admission, +and product-mode selection. A standalone partial-edge scale test was removed because its only assertion was already present +verbatim in the full 700-target inventory contract. + +Fused gate-up now reports state/loading, effective numerics, gradients, CPU admission, and the live Hopper kernel gate. +TP16 lm-head now reports topology, operands, presentation bytes, surrogate gradients, and its live Hopper gate. All factor +and FP8 bytes, rank/alpha/topology checks, gate/up order, one-time BF16 rounding, branch composition, base-free factor VJP, +mutation safety, TP16 ranges/order/NCCL, sampler strides, vocabulary assembly, custom-autograd saves, and gradients remain. + +Focused validation reports 13 passed and the two existing optional SGLang/Hopper skips. Repository collection is now +1,548 items and the static inventory is 1,538 definitions with 555 curated decisions. The candidate inventory remains 56, +with 45 conditional-runtime skips, 11 no-observable outcomes, no parse errors, and the intentional duplicate group +unchanged. + +## Eighty-third wave: consolidate sparse deltas, EP adapters, and DSV4 runtime + +The eighty-third wave reduces sparse-delta files, EP backend adapters, and DSV4 model suites from 34 reports to 15. +Sparse-delta coverage now reports source capture, single-file encoding, contiguous sharding, ranked files, and translation +futures. Rank/global manifests, empty shards, trust-root filenames, deterministic indices, malformed inputs, encoded/raw +writers, future tags, rank order, and every output statistic remain exercised. + +EP adapters now report registry signatures and native/Triton/Triton-MoE-act FP8 boundaries while retaining independent +score-forwarding and Quack activation checks. DSV4 now reports construction/topology, runtime variants, precision +preservation, and outer gradient checkpointing. C128, ordinary and hash-routed backward, PP rejection, CP wiring, HC +gradient ownership, FP32 carve-outs, registry/direct casts, complex RoPE, and per-layer checkpoint calls remain checked. + +All 15 focused reports pass. Repository collection is now 1,529 items and the static inventory is 1,519 definitions with +567 curated decisions. The candidate inventory is now 50, with 39 conditional-runtime skips, 11 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## Eighty-fourth wave: consolidate routing replay, top-k policies, and OPD payloads + +The eighty-fourth wave reduces routing-replay, top-k router, and OPD pipeline-payload suites from 28 reports to 16. +Routing replay now reports sequence-parallel layout, RingAttention zigzag layout, weight tensors, and wire decoding. +Padded/unpadded positions, unpacked rows, truncation, all ring ranks, packed-document boundaries, FP32 padding, CP slicing, +shaped/inferred base64, and materialized Python/NumPy/Tensor representations remain exercised. + +Top-k routing now reports synthetic balancing, softmax behavior, layer-scoped FP32, sqrt-softplus/noaux, hash routing, and +configuration/scaling. Every tie policy, V4-input isolation, bias/weight distinction, tid2eid ownership, selection, scale, +and config default remains. OPD payload coverage retains chunking, endpoint matching, weight-version verification, and +prepare-worker transitions while reporting shifted payloads and Mooncake teacher-cache transport once each. + +All 16 focused reports pass. Repository collection is now 1,517 items and the static inventory is 1,507 definitions with +576 curated decisions. The candidate inventory remains 50, with 39 conditional-runtime skips, 11 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## Eighty-fifth wave: consolidate adapter persistence and MoE-LoRA semantics + +The eighty-fifth wave reduces checkpoint-manager saves, LoRA checkpoint roundtrips, and MoE-LoRA suites from 35 reports +to 20. Checkpoint-manager coverage now reports rank-zero write atomicity, +factor-only snapshot admission, live adapter targets, and collective MoE export. Internal dtype and export-format keyword +probes were removed in favor of real saved-tensor and SGLang roundtrip coverage; the duplicate strict-manifest artifact +test was removed in favor of the public adapter-manager save/validate lifecycle. + +Checkpoint formats now report runtime-rank export, PEFT hybrid-shared roundtrip, SGLang shared-outer roundtrip/admission, +both expert-ownership modes, and quantized projection subsets. The unrelated cautious-optimizer example was removed +because the dedicated optimizer suite already owns that routing contract. MoE-LoRA now reports construction/runtime-rank +layout, eager execution/hybrid ownership, zero-delta behavior, cross-backend numerics, MoE injection, and EP router-score +semantics. A redundant nonzero-output example, Qwen subclass repetition, and generic linear-injection errors were removed. + +All 20 focused reports pass. Repository collection is now 1,502 items and the static inventory is 1,492 definitions with +589 curated decisions. The candidate inventory remains 50, with 39 conditional-runtime skips, 11 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## Eighty-sixth wave: consolidate target manifests and quantization admission + +The eighty-sixth wave reduces LoRA target-manifest, FP8 compatibility, and NVFP4 fake-quant suites from 25 reports to 11. +Target-manifest coverage now reports the successful injection/coverage lifecycle and one fail-closed schema/runtime policy. +Count, rank, configured-target, unlisted-module, Boolean, and exact scalar-type failures all remain checked without seven +separate reports. + +FP8 compatibility now reports NeMo translation, external configuration rejection, Blackwell admission, and BF16 layer- +island resolution/injection. The thirteen vLLM and ModelOpt rejection paths are table-driven, while first/last overlap, +invalid topology, real replacement boundaries, and summary metadata form one layer-island lifecycle. NVFP4 now reports +2D reference/dispatch/STE behavior, input admission, MoE projection STE, expert independence, and per-half gate/up scaling. + +All 11 focused reports pass. Repository collection is now 1,488 items and the static inventory is 1,478 definitions with +596 curated decisions. The candidate inventory remains 50, with 39 conditional-runtime skips, 11 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## Eighty-seventh wave: consolidate packing, schedules, and API tracking + +The eighty-seventh wave reduces data-packing, learning-rate scheduler, checkpoint-path, and API-type suites from 31 +reports to 22. Packing now reports allocation primitives, sample preprocessing, dataset preprocessing, and the +PackingDataset lifecycle. FFD, binning, rank allocation, position metadata, label filtering, Hugging Face dataset filtering, +and optional eval handling retain all previous assertions without one report per helper. + +Learning-rate coverage now reports constant, linear, cosine, and invalid-configuration policies. Both default and custom +linear/cosine traces retain their warmup, monotonicity, midpoint, endpoint, decay-ratio, and floor checks. API coverage now +reports removed/future session fields as one compatibility boundary, sampler reconciliation as one stale/query-failure +policy, and model-scoped tracking plus failed-load atomicity as one session-tracking contract. + +All 22 focused reports pass. Repository collection is now 1,479 items and the static inventory is 1,469 definitions with +604 curated decisions. The candidate inventory remains 50, with 39 conditional-runtime skips, 11 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## Eighty-eighth wave: consolidate routing position, QARL MoE, and RoPE precision + +The eighty-eighth wave reduces routing-weight-position, NVFP4 QARL-MoE, and RoPE precision suites from 23 reports to 9. +Routing coverage now reports numerical behavior and configuration resolution as two contracts. Both before/after-down +trees, FP64 outputs and gradients, no-score-gradient execution, mutable/env settings, auto regimes, parity opt-in, +explicit Boolean/string values, and invalid input remain checked. + +QARL-MoE now reports conversion, eager execution, and injection as three lifecycles. Parameter identity, idempotence, +non-expert rejection, lossy/disabled execution, restoration, gradients, target selection, metadata, and FP8 rejection all +remain. RoPE now reports registry/frequency precision, contract-lane bytes, CPU cache lifecycle, and exact-architecture +device construction. The optional CUDA device gate remains independent. + +All 9 focused reports pass. Repository collection is now 1,465 items and the static inventory is 1,455 definitions with +611 curated decisions. The candidate inventory remains 50, with 39 conditional-runtime skips, 11 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## Eighty-ninth wave: consolidate QARL injection, Nemotron EP, and endpoint routing + +The eighty-ninth wave reduces generic QARL, Nemotron-H checkpoint, and weight-sync endpoint suites from 22 reports to 14. +Generic QARL retains independent stateful fake-quant, folded-export, injection/admission, and configuration policies. Dense +target/exclusion behavior, parameter names, summary counters, MTP rejection, and Mamba rejection now report as one public +injection lifecycle. + +Nemotron-H retains independent HF parity, published-layout roundtrip, and stacked-HF loading reports while topology +validation, skip-key ownership, local expert slicing, skip accounting, and EP-plan classification form one ownership +policy. Endpoint coverage now reports health fallback/failure, init/direct port routing, two-phase transfer, mixed/chunked +flattened transfer, hybrid receiver fencing, and multi-rank direct-format rejection as six protocol boundaries. + +All 14 focused reports pass. Repository collection is now 1,457 items and the static inventory is 1,447 definitions with +616 curated decisions. The candidate inventory remains 50, with 39 conditional-runtime skips, 11 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## Ninetieth wave: consolidate runner lifecycles, attention paths, and sequence sharding + +The ninetieth wave reduces runner session-operation, attention, and sequence-shard collator suites from 28 reports to 18. +Runner coverage now reports save, registration, fatal optimizer exit, publication, gradient abort, and forward-backward +lifecycles. Cross-rank registration failure, successful/failing publication tails, completion ordering, uniform rejection, +and asymmetric failure promotion all remain exercised. + +Attention now reports registry/resolution, repeat-KV, fixed and varlen FlashAttention, SGL page-size-one KV cache, +alternate paged/FA3/FA4 selection, and cross-attention rejection. The two existing optional FlashAttention gates remain +environment-dependent. Collation now reports SP primitives, packed-label boundaries, full SP splitting/metadata, and +teacher/DRGRPO side-channel alignment across CP2 and CP16. + +Focused validation reports 16 passed and two existing optional FlashAttention skips. Repository collection is now 1,447 +items and the static inventory is 1,437 definitions with 624 curated decisions. The candidate inventory remains 50, with +39 conditional-runtime skips, 11 no-observable outcomes, no parse errors, and the intentional duplicate group unchanged. + +## Ninety-first wave: consolidate sparse-delta and BI norm lifecycles + +The ninety-first wave re-audits merged-LoRA execution, sparse-delta transport, and families-v2 normalization, reducing the +three suites from 27 reports to 21. The nine merged-LoRA reports remain independent: canonical folding, gradient dtype, +straight-through autograd, dense selection/cache behavior, MoE folding/admission/native-EP routing, and trunk composition +exercise distinct numerical or backend boundaries. + +Sparse-delta now reports full, changed, and unchanged streaming updates as one lifecycle, retaining TP path replication, +exact changed-byte indices and values, endpoint metadata, and skip accounting. Per-rank prepacked paths and FP8 KV-cache +metadata now share one publication transaction while preserving unique-file byte accounting. Baseline priming, +initialization, runtime loading, and receiver-failure rollback remain independent. + +Families-v2 keeps correctness, batch invariance, run-to-run determinism, fused/split parity, dispatch, and strided qk-norm +as six contracts. The split-kernel guard now executes inside the parity matrix, and shipped-size, threshold, row cutoff, +and tile-basis checks report as one dispatch policy. + +All 21 focused reports pass, including the CUDA normalization kernels. Repository collection is now 1,441 items and the +static inventory is 1,431 definitions with 628 curated decisions. The candidate inventory remains 50, with 39 +conditional-runtime skips, 11 no-observable outcomes, no parse errors, and the intentional duplicate group unchanged. + +## Ninety-second wave: consolidate clipping, packing, and P2P transfer policies + +The ninety-second wave reduces EP gradient clipping, orchestrator packing, and P2P backend protocol suites from 74 reports +to 62. EP clipping now reports local norm behavior, skip-FSDP ownership, public dispatch, and mixed-mesh foreach handling +as four complete policies. Infinity norm, empty and missing gradients, raw EP-local gradients, ordinary fallback, safe +per-tensor clipping, and explicit foreach rejection all remain. The real two-rank reduction/non-finite gate and three-rank +participation-mask gate remain independent. + +Packing now reports ordinary, overflow, mixed-length, exact-fit, and off-by-one cases as one capacity lifecycle. NumPy input +normalization joins the existing empty, single, oversized, and missing-input boundary. Packed metadata, disabled packing, +position and label generation, validation, unpacking, and the full pack-to-unpack roundtrip remain separate contracts. + +P2P transfer admission now rejects unknown parameters, incompatible receiver shapes, and unsupported source ranks in one +side-effect-free policy. Receiver-memory coalescing covers both distinct and missing handle metadata in one report, while +failure diagnostics cover named tensors and handles, capped samples, omitted counts, and default redaction together. Every +FP8 dequantization topology, direct-EP path, staged CPU/GPU transfer, alignment, completion, and cleanup contract remains. + +All 62 focused reports pass, including the live distributed clipping workers and the complete retained P2P topology matrix. +Repository collection is now 1,429 items and the static inventory is 1,419 definitions with 637 curated decisions. The +candidate inventory remains 50, with 39 conditional-runtime skips, 11 no-observable outcomes, no parse errors, and the +intentional duplicate group unchanged. + +## Ninety-third wave: consolidate runner compiler and Muon policies + +The ninety-third wave re-audits runner LoRA-head compilation, GLM5 support, and Muon optimization, reducing the three suites +from 42 reports to 32. Runner coverage now reports effective LM-head selection, replica topology, unquantized expert +admission, and quantized expert contracts as complete policies. Canonical and legacy head formulas, valid SP/output +replicas, every coverage failure, general-TP rejection, hybrid metadata, all quantized formats, declared-shape drift, and +unsupported EP/eFSDP regimes remain checked. Registration, exact TP16 VJP ownership, staged-capture abort, session-rank +specialization, block-FP8 DeepEP, and the authoritative analytical optimizer step remain independent. + +The 13 GLM5 reports are retained unchanged: they map to distinct construction, selector, mask, sparse-attention, adapter, +checkpoint, HF-parity, and recompute boundaries rather than narrow literal variations. Muon now reports builder validation, +Quack backend selection, and fused/Nemotron parameter classification as three complete policies. Algorithm updates, +grouping geometry, standard Newton-Schulz batching, CUDA FP32 compute preservation, SGD fallback, and the tiny Nemotron +training step remain separate numerical or end-to-end gates. + +All 32 focused reports pass, including the available TileLang, HF-reference, and CUDA paths. Repository collection is now +1,419 items and the static inventory is 1,409 definitions with 644 curated decisions. The candidate inventory remains 50, +with 39 conditional-runtime skips, 11 no-observable outcomes, no parse errors, and the intentional duplicate group +unchanged. + +## Ninety-fourth wave: consolidate layout, optimizer-resume, and FP8 execution policies + +The ninety-fourth wave reduces weight-sync handler, adapter optimizer-resume, and FP8 linear suites from 45 reports to 33. +Weight-sync extraction now reports dense filtering and tied aliases as one ownership policy. DeepSeek and Kimi MLA fusion, +contiguous FP8 views, Nemotron-H conversion, and gated stacked-expert splitting now report as one inference-layout policy; +all name, value, storage, transpose, prefix, and rejection assertions remain. + +Optimizer resume now reports canonical identity with live binding, bitwise continuation with its weights-only control, +scheduled and overridden public LR restoration, and successful or rejected logical resharding as four complete policies. +Moment restoration, full checkpoint writing, manifest identity, artifact admission, incomplete-restore atomicity, recursive +snapshotting, and post-mutation collective failure remain independent. Every one- and two-dimensional topology, replica, +hole, overlap, dtype, shape, step, structure, empty-rank, and resident-state case still executes. + +FP8 linear injection now combines replacement, recipes, and exclusions. Backend parity now includes automatic warn-once +fallback, and CUDA training now includes float32-output dispatch. CPU fallback, profiler behavior, padded matmul parity, +residual correction, activation2 correction, operand diagnostics, and CUDA numerical execution remain distinct gates. + +All 33 focused reports pass, including every retained CUDA FP8 path. Repository collection is now 1,407 items and the +static inventory is 1,397 definitions with 653 curated decisions. The candidate inventory remains 50, with 39 +conditional-runtime skips, 11 no-observable outcomes, no parse errors, and the intentional duplicate group unchanged. + +## Ninety-fifth wave: consolidate endpoint, distributed autograd, and fused-MoE policies + +The ninety-fifth wave reduces inference endpoints, distributed adapter autograd, and SGLang fused-MoE suites from 35 +reports to 26. Endpoint registration now reports discovered TP size and configured sync method as one auto-sync policy; +single-endpoint forwarding and default, named, or unmatched pools now form one routing policy. Port selection, FP8 KV-cache +admission, health fallback, quantization, cache invalidation, receiver detection/enrichment, and method validation remain +independent. + +Distributed adapter autograd now reports dense plus sequence-parallel ownership, unquantized EP2 all-to-all, unquantized +four-rank eFSDP, and quantized EP2 all-to-all as four policies. Every original subprocess still runs: dense, SP, all four +unquantized backends, shared-owner and all-owner layouts, projection subsets, Triton NF4, native NVFP4, Quack block-FP8, +structural-zero handling, analytical clipping, and public optimizer parity. Direct-output and the two dependency-gated +DeepEP compositions remain separate topology gates. + +Fused-MoE now reports stock and masked trainable gradients together, and strided/transient plus auto/explicit real-kernel +parity together. Mocked dispatch, admission, weight cache/layout, adapter layout, runtime context, and auto resolution remain +separate. Focused validation reports 22 passed and four existing dependency skips; every executable distributed worker +passes. Repository collection is now 1,398 items and the static inventory is 1,388 definitions with 661 curated decisions. +The candidate inventory is now 48, with 37 conditional-runtime skips, 11 no-observable outcomes, no parse errors, and the +intentional duplicate group unchanged. + +## Ninety-sixth wave: consolidate canonical-MoE reporting and remove an FP8 plumbing mock + +The ninety-sixth wave reduces canonical-MoE, FP8-MoE, and training-simulator suites from 32 reports to 26. Canonical-MoE +now reports adjacent-tree widths as one numerical policy, transport resolution and direct executor rejection as one +admission policy, both world-32 group layouts as one topology policy, and both distributed contributor widths as one +execution policy. Every original contributor width, group layout, transport guard, distributed worker, and byte-exact +packed-EP16 comparison still executes. + +FP8-MoE drops a positional-argument mock of the internal Quack autograd call. The retained CUDA TP lifecycle now exercises +the same Triton-grouped backend and non-default block size through real forward, backward, finite-gradient checks, and a +master-weight update. Kernel forward and weight-gradient parity, scalar fallback, DeepGEMM isolation, expert training, and +full injected-model training remain separate numerical or lifecycle gates. The simulator's 11 reports remain unchanged +because ingestion, configuration, Qwen calibration, topology extrapolation, path trust, built-in packs, analytical ledgers, +correctness-gated ranking, and consolidated validation are distinct behavioral boundaries. + +Focused validation reports 25 passed and one existing opt-in DeepGEMM skip. Repository collection is now 1,392 items and +the static inventory is 1,384 definitions with 665 curated decisions. The candidate inventory remains 48, with 37 +conditional-runtime skips, 11 no-observable outcomes, no parse errors, and the intentional duplicate group unchanged. + +## Ninety-seventh wave: consolidate prequantized loading, sequence side fields, and Quack safety + +The ninety-seventh wave audits three previously untouched suites and reduces them from 18 reports to 12. Prequantized +checkpoint coverage now reports NVFP4 plus block-FP8 detection as one format policy and dense plus MoE exclusion behavior +as one handler policy. Every nested, flat, config, index, precedence, malformed, missing, wrong-size, passthrough, skip, +shared-expert, and auxiliary-key case still executes. + +Packing-concat coverage now reports teacher hidden states and hidden-match weights as one sequence-side-field policy while +retaining their different ranks, exact concatenated values, shapes, and padding. Quack process safety now reports silent +timeouts plus truncated frames as one receive-protocol policy and structural hashes, unsafe objects, and Cutlass dtype +classes as one cache-key policy. PTXAS output isolation and entry-name selection remain independent filesystem/compiler +boundaries. + +All 12 focused reports pass. Repository collection is now 1,386 items and the static inventory is 1,378 definitions with +670 curated decisions. The candidate inventory remains 48, with 37 conditional-runtime skips, 11 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## Ninety-eighth wave: reuse PP baselines and consolidate NVFP4 and retry lifecycles + +The ninety-eighth wave reduces Qwen3 pipeline-parallel, NVFP4 export, and data-retry suites from 16 reports to 10. Pipeline +schedule parity now trains the 1F1B baseline once and compares all three retained schedules against that trajectory. This +removes two entire duplicate baseline training runs while preserving Interleaved1F1B, InterleavedZeroBubble, and +ZBVZeroBubble convergence plus every per-step tolerance check. The distinct PP/FSDP, Muon, and server E2E gates remain. + +NVFP4 now reports packed layout, dequantization error, and shared global scale as one tensor policy, while weight-only, +W4A4, and already-quantized directory behavior form one export lifecycle. All tensor dtypes and shapes, fused scales, BF16 +islands, metadata, numerical errors, input-scale rules, and re-export rejection remain. Data retry coverage now reports +success, retryable and terminal failures, and exponential, linear, or constant backoff as one policy. Its validation also +exposed and repaired the Hugging Face Hub 1.27 import seam by importing `HfHubHTTPError` from the public errors module. + +All four executable CPU reports pass; the six retained GPU E2E reports collect successfully without launching the costly +training jobs in this audit pass. Repository collection is now 1,380 items and the static inventory is 1,374 definitions +with 674 curated decisions. The candidate inventory remains 48, with 37 conditional-runtime skips, 11 no-observable +outcomes, no parse errors, and the intentional duplicate group unchanged. + +## Ninety-ninth wave: consolidate model-family lifecycles and remove fake data-loader checks + +The ninety-ninth wave reduces OLMo2, Qwen2, cu-seqlen, orchestrator-client, and distributed data-loader suites from 22 +reports to 16. OLMo2 and Qwen2 now each report construction plus TP unfusing as one architecture-layout policy and save +plus load as one bidirectional checkpoint policy. Every family-specific norm, bias, fused/split key, strict-load, +hidden-state, and logits assertion remains; OLMo2's independent TP-plan contract also remains. + +Server cu-seqlen alignment now includes the SP-owned metadata case, while one ZeroMQ lifecycle covers initial health, +repeated requests, and interleaved operations without starting a duplicate client/engine fixture. Forward-backward, +optimizer, serialization, errors, and lifecycle edges remain separate. The data-loader suite keeps all four reports but +drops two genuinely non-observing blocks: a literal `4 * 3 == 12` assertion and an alleged epoch-consistency check that +only compared two list lengths hard-coded to three. Real partitioning, microbatching, sharding, padding, drop-last, +packed-data, multi-DP, and variable-length behavior remains. + +All 16 focused reports pass. Repository collection is now 1,374 items and the static inventory is 1,368 definitions with +681 curated decisions. The candidate inventory remains 48, with 37 conditional-runtime skips, 11 no-observable outcomes, +no parse errors, and the intentional duplicate group unchanged. + +## One-hundredth wave: eliminate false no-outcome signals and make acceptance explicit + +The one-hundredth wave audits every one of the 11 remaining no-observable candidates. None was an inert test: three joined +`torch.multiprocessing.start_processes` wrappers propagate child assertions, the Muon matrix delegates to an +`assert_success` helper, six H100 frozen-bit gates delegate to a SHA-256 assertion helper, and the exact LM-head optimizer +case is an intentional must-not-raise acceptance boundary. + +The audit now recognizes joined multiprocessing as an observable outcome. Muon and BI helpers use assert-prefixed names, +and the scalar optimizer-state acceptance test explicitly asserts the validator's successful result. This removes false +triage noise without weakening or consolidating distinct distributed and numerical gates. The sequence-parallel, exact-DCP, +DTensor materialization, Muon transition, and scalar-state wrappers all pass; the six H100-only golden reports collect +unchanged. + +Repository collection remains 1,374 items and the static inventory remains 1,368 definitions with 684 curated decisions. +The candidate inventory falls from 48 to 37: all 11 no-observable signals are gone, leaving only conditional-runtime-skip +reviews, with no parse errors and the intentional duplicate group unchanged. + +## One-hundred-and-first wave: stop masking supported-kernel failures as skips + +The one-hundred-and-first wave reviews every remaining conditional-runtime-skip candidate and separates honest optional +runtime gates from failure-masking guards. Grouped-GEMM and MoE kernel suites now skip unsupported CPU hosts at their +declared CUDA boundary, but failures importing this repository's own kernel modules fail supported GPU runs. The non-gated +MoE suite likewise imports its core expert implementation normally instead of turning every import-time exception into a +skip. + +The environment-dependent SGLang negative test now simulates a missing SGLang package deterministically, so it executes on +both installed and uninstalled environments. Doing so exposed and corrected a diagnostic that named the unrelated +TP-simulation flag. The QARL CPU suite drops one narrow, usually-skipped registry probe that only checked two internal +dictionary entries; its fake-quant numerics, STE gradients, shadow selection, and restoration contracts remain. + +All 16 focused CPU and CUDA reports pass without skips. Repository collection is now 1,373 items and the static inventory +is 1,367 definitions with 688 curated decisions. The candidate inventory falls from 37 to 29, all remaining candidates +being explicit optional-library, GPU-capacity, distributed-topology, or backend-availability gates. There are no parse +errors, and the intentional duplicate group is unchanged. + +## One-hundred-and-second wave: remove duplicate guards and compose lifecycle seams + +The one-hundred-and-second wave removes two cross-file duplicates and consolidates five tiny seam pairs, reducing ten +reports to three. The examples-only personal-path scan is absorbed into the stronger repository-wide hygiene guard, which +now also recognizes Mac home directories and personal data workspaces. A standalone runner-dispatcher model-id test is +removed because the retained request-processor policy checks the same rank-zero forward path plus routed ids, routed +logits, adapter auto-load, and returned session identity. + +Weight-version forwarding now runs as one composed handler-to-NCCL-synchronizer path instead of two mocked handoffs. HSDP +deferral reports its enabled transitions and non-replicated rejection together. The emitted MoE inference buffer subsumes a +direct one-element runtime-scaling unit, while SGLang RMSNorm mode selection and exact forced-residual numerics form one +policy. Offline and server index-share failures now share one cleanup lifecycle, retaining both caller-specific assertions. + +All eight focused reports pass, including the two stronger retained dispatcher policies. Repository collection is now +1,366 items and the static inventory is 1,360 definitions with 695 curated decisions. The candidate inventory remains 29 +legitimate runtime gates, with no parse errors and the intentional duplicate wrapper group unchanged. + +## One-hundred-and-third wave: consolidate thin server wrapper reports + +The one-hundred-and-third wave reduces three two-report server suites to three complete policies. Importance-sampling +metrics now cover default ratio aggregation and custom TIS extrema together, including weighted means, minima, maxima, +valid-token aggregation, and Python-scalar output. GLM LoRA target resolution now checks raw-HF defaults and explicit +precedence from one model fixture. Remote RPC wrappers now preserve weight-sync timeout and optimizer sparse-delta payloads +in one operation matrix rather than separate single-field reports. + +All three focused reports pass. Repository collection is now 1,363 items and the static inventory is 1,357 definitions +with 698 curated decisions. The candidate inventory remains 29 conditional runtime gates, with no no-outcome candidates, +no parse errors, and the intentional duplicate wrapper group unchanged. + +## One-hundred-and-fourth wave: consolidate mode, precedence, and session lifecycles + +The one-hundred-and-fourth wave reduces four untouched three-report suites from 12 reports to seven. DSv4 RoPE cache +length now reports config default and environment precedence together, while the independent context-parallel short-cache +rejection remains. SGLang JIT and kernel RMSNorm CPU fallbacks now share one numerical matrix covering both exact residual +paths and packed tensors. + +A nonresident LoRA session now demonstrates missing-checkpoint preservation followed by evicted-checkpoint promotion on +the same runner; traversal rejection remains separate. QARL calibration input loading now covers valid truncation and +malformed token shapes together, while its persistent calibration-state lifecycle remains independent. + +All seven focused reports pass. Repository collection is now 1,358 items and the static inventory is 1,352 definitions +with 702 curated decisions. The candidate inventory remains 29 legitimate runtime gates, with no parse errors and the +intentional duplicate wrapper group unchanged. + +## One-hundred-and-fifth wave: merge cross-file modes and delete test-local tests + +The one-hundred-and-fifth wave merges the remaining SGLang RMSNorm CPU checks into one mode policy and removes the redundant +one-test JIT module. Global selection, forced-residual FP32 multiplication, JIT and kernel fallback, packed shape, exact +residual values, and global-state restoration all remain. NCCL rendezvous coverage now reports sticky ephemeral rotation +and explicit port pinning in one lifecycle while keeping bind-failure admission separate. + +FutureStore drops assertions for response builders defined inside the test file itself. Its production entry defaults, +expiry, terminal states, queue transitions, concurrency, deletion, status, error, and TTL behaviors remain. + +All six focused reports pass. Repository collection is now 1,356 items and the static inventory is 1,350 definitions with +705 curated decisions. The candidate inventory remains 29, with no parse errors and the intentional duplicate wrapper +group unchanged. + +## One-hundred-and-sixth wave: consolidate utility matrices and expose FP8 imports + +The one-hundred-and-sixth wave reduces FQN matcher and block-FP8 suites from eight reports to five. Single, all, and any FQN +matching now form one policy containing every exact, wildcard, grouped-number, indexed, prefixed, empty, first-match, and +invalid-input case. Module path get/set remains independent. + +Block-FP8 contiguity and divisibility rejection now execute with shapes, dtypes, scales, and block sizes in one +quantization policy. Imports of this repository's block-FP8 module are no longer caught and labeled as optional feature +absence, and two unused imports are gone. Numerical dequantization and edge/determinism coverage remain separate. + +All five focused CPU and CUDA reports pass. Repository collection is now 1,353 items and the static inventory is 1,347 +definitions with 707 curated decisions. The candidate inventory remains 29, with no parse errors and the intentional +duplicate wrapper group unchanged. + +## One-hundred-and-seventh wave: consolidate source admission and remove an inert compatibility suite + +The one-hundred-and-seventh wave combines dataset source resolution across local files, saved directories, hub datasets, +URLs, missing sources, and string or list data files. Exact DCP skip mode now reports valid exact-model deferral and +non-exact rejection together while retaining FSDP deregistration and the no-HF-read guard. + +The SGLang sparse-delta compatibility module is removed. Its module-level zstd guard skipped every report because the +production writer has no disk-compression API, while its ordinary receiver apply, checksum, validate-only, and parameter +parity contracts are owned by the stronger retained trainer-to-request-processor-to-SGLang E2E. The retained sparse file +and transport backend suites provide the lower-level encoding and posting contracts. + +Focused validation reports 19 passed and one retained external-dependency skip. Repository collection is now 1,351 items, +the static inventory is 1,340 definitions, and there are 710 curated decisions across 357 Python test files. The candidate +inventory remains 29 conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional +duplicate distributed wrapper group unchanged. + +## One-hundred-and-eighth wave: compose launcher precedence and DistSign ownership + +The one-hundred-and-eighth wave reduces launcher address and override fragments plus DistSign local-hook ownership from +six reports to three. Remote rank-zero discovery and explicit connect-host precedence now share one address policy; +schema-agnostic override parsing and removed-field migration validation share one parsing policy. DistSign hook +registration now verifies both local installation and FSDP-managed exclusion in one ownership lifecycle. + +All nine focused launcher and optimizer reports pass. Repository collection is now 1,348 items and the static inventory is +1,337 definitions with 713 curated decisions. The candidate inventory remains 29 legitimate conditional runtime gates. + +## One-hundred-and-ninth wave: replace registry and config fragments with composed behavior + +The one-hundred-and-ninth wave reduces NVFP4 normalization, non-gated MoE, and server Adam initialization by four reports. +NVFP4 alias defaults, activation override, and every invalid block size now form one normalization policy. A direct relu2 +registry-membership probe is removed because exact non-gated MoE forward and gradient parity exercises the production +activation. ServerArguments Adam values now feed ModelRunner initialization in the same test, retaining non-default and +default parameter groups plus malformed-beta rejection. + +All 11 focused CPU and CUDA reports pass. Repository collection is now 1,344 items and the static inventory is 1,333 +definitions with 716 curated decisions. The candidate inventory remains 29 legitimate conditional runtime gates. + +## One-hundred-and-tenth wave: test protocol behavior instead of model field echoes + +The one-hundred-and-tenth wave reduces API and runner protocol coverage from 11 reports to six. Two Pydantic reports that +primarily echoed constructor fields and automatic required-field behavior are removed; real training, checkpoint, sampler, +and session endpoint tests exercise those models, while the unique forward session-id alias remains in the compatibility +policy. Runner protocol constructor, UUID, timestamp, and optional-field fragments are replaced by one full typed-wire +equality contract covering every retained payload, success and error responses, tensors, JSON, ACK correlation, and pickle +rejection. A second API-orchestrator flow that repeated the adjacent roundtrip, builder, validator, and streaming checks is +also removed. + +All six focused reports pass. Repository collection is now 1,339 items and the static inventory is 1,328 definitions with +719 curated decisions across 357 Python test files. The candidate inventory remains 29 conditional runtime gates, with no +no-outcome candidates, no parse errors, and the intentional duplicate distributed wrapper group unchanged. + +## One-hundred-and-eleventh wave: exercise policies through their production consumers + +The one-hundred-and-eleventh wave reduces checkpoint, DeepEP, weight-sync quantization, and request-processor coverage by +eight reports. Legacy and PP-parent EP meshes are now selected through the production restore operation; malformed named +dimensions remain a separate rejection boundary. DeepEP overflow, no-RDMA allowance, byte alignment, and default sizing +form one buffer policy. BF16 no-ops and valid FP8 normalization form one supported policy, while all unsupported or +malformed forms share one rejection policy. + +The direct teacher-sort helper probe is replaced by an OPD model-pass transaction. Nested teacher ids and top-level +precedence drive actual sorting, then compose with packer datum order and Mooncake routing-payload order. All 27 focused +reports pass. Repository collection is now 1,331 items and the static inventory is 1,320 definitions with 723 curated +decisions. + +## One-hundred-and-twelfth wave: remove weaker helper probes and relocate hardware timing + +The one-hundred-and-twelfth wave reduces SignSGD, SGLang fused-MoE, and sparse-MLA coverage by four reports. SignSGD now +reports dense updates, decay, missing gradients, and sparse rejection as one step policy; a separate assertion of base +PyTorch Optimizer state-dict behavior is removed. The direct FP32 routing-scale helper probe is removed because the retained +fused-MoE backward oracle uses the same discriminating values and compares every gradient exactly. + +Sparse-MLA's production-shape H100 speed ratio moves out of pytest into +`certification/glm52/benchmark_sparse_mla_backward.py`. The explicit benchmark retains combined and split warmups, median +timing, environment restoration, and a configurable speedup gate. All three forward/backward numerical kernel reports +remain and pass; nine focused reports pass in total. Repository collection is now 1,327 items and the static inventory is +1,316 definitions with 726 curated decisions. + +## One-hundred-and-thirteenth wave: stop running an unasserted benchmark after correctness + +The vocab-parallel CE distributed test previously continued after eager and compiled value/gradient parity to run 100 +production-scale warmup and timed iterations. Those iterations only printed latency and peak-memory tables and could not +fail the test. They now live in the explicit `certification/benchmark_vocab_parallel_ce.py` torchrun script, while the +two-rank pytest worker ends after numerical correctness. + +The retained distributed report passes, and both certification scripts pass lint and compilation checks. Repository +collection remains 1,327 items, the static inventory remains 1,316 definitions, and there are 727 curated decisions across +357 Python test files. The candidate inventory remains 29 legitimate conditional runtime gates, with no no-outcome +candidates, no parse errors, and the intentional duplicate distributed wrapper group unchanged. + +## One-hundred-and-fourteenth wave: replace smokes and protocol fragments with stronger behavior + +The one-hundred-and-fourteenth wave removes two expensive training smokes already contained in stronger E2E paths. Dense +one-GPU FP8 training remains covered by the checkpoint-and-resume lifecycle, including the original two-step metrics and +module-usage assertions. The basic two-GPU DistSign run is subsumed by the retained FSDP2, Ulysses, DP2, and accumulation +composition. That survivor now uses eager attention, matching the repository's other tiny Ulysses topology tests and +avoiding an unrelated FlashAttention compiler failure on the synthetic shape. + +Nemotron-H packed boundaries now flow through the all-mixer loss and backward contract instead of a finite-output smoke. +Qwen3 tensor-parallel unfusing drops a comparison between independently initialized MLP output shapes and retains direct +plus model-wide projection ownership. The PP NCCL mocks now form an actual sender-to-receiver roundtrip rather than two +manually disconnected halves. + +All eight retained covering reports pass, including FP8 checkpoint/resume and the four-GPU DistSign E2E. Repository +collection is now 1,322 items and the static inventory is 1,311 definitions with 732 curated decisions. + +## One-hundred-and-fifteenth wave: remove duplicate norm proof and synthetic one-rank shards + +The families-v2 dispatch suite drops its second fused-versus-split bitwise comparison. The retained norm contract forces +both implementations, verifies that split execution actually occurred, and covers more hidden sizes plus residual, plain, +and zero-centered forms; dispatch selection at shipped and deep shapes remains independently checked. + +Three exact-GLM component suites no longer launch an additional one-rank FSDP2 process before their real two-rank shard. +The dense MLP, absorbed kv_b, and generic TP1 QLoRA workers execute every common lifecycle, byte-parity, ownership, and +gradient assertion at world size two, where they additionally verify genuinely sharded factor storage. The retained norm +gate passes. The three distributed GLM wrappers collect and reach their explicit optional-dependency gate, then skip +because SGLang is not installed in the repository venv; the removed wrappers had the same gate. + +Repository collection is now 1,318 items and the static inventory is 1,307 definitions with 734 curated decisions across +357 Python test files. The candidate inventory falls from 29 to 26 conditional runtime gates because the three redundant +optional-SGLang wrappers are gone. There are no no-outcome candidates, no parse errors, and the intentional duplicate +distributed wrapper group is unchanged. + +## One-hundred-and-sixteenth wave: replace helper probes with production consumers + +The one-hundred-and-sixteenth wave removes 13 reports across numerical ops, model loading, checkpoint loading, shared +prefix packing, and server optimizer steps. Families-v2 dispatch now spies on actual fused-versus-split runtime selection +inside the retained norm policy; its redundant helper-only dispatch module and duplicate environment-variable rollback +report are gone. Dense Qwen3.5 RMSNorm construction and site assignment are expressed as two complete policies instead of +seven fragments, and a direct rotary-helper comparison is removed because retained dense and MoE attention projections +exercise the same reference while rejecting the wrong rotation. + +Qwen3.5 dense and MoE config conversion now runs through local config files and the production auto-config loader. Its MTP +skip patterns now run inside grouped dense/expert checkpoint loading instead of a direct regex probe. Optimizer-step +learning-rate precedence drives actual API payloads, including both legacy fallbacks, rather than a private resolver. +Finally, the one-line no-shared-prefix fallback joins the complete repack/remap policy. + +All 15 focused CPU, server, and CUDA reports pass. Repository collection is now 1,305 items and the static inventory is +1,294 definitions with 742 curated decisions across 356 Python test files. The audit still contains only 26 legitimate +conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate distributed +wrapper group unchanged. + +## One-hundred-and-seventeenth wave: replace field echoes and helper branches with complete policies + +The one-hundred-and-seventeenth wave removes 13 reports across active-LoRA admission, BI routing and loss, trainer timing, +API models, optimizers, collators, OPD scripting, and sequence parallelism. Active-LoRA truth-table cases and topology +admission now form two complete policies. Router batch invariance now runs through `MoEBlock.route`, including logits, +selections, and weights, while exact and ordinary dispatch share one production policy. Standard and temperature BI fused +LM-head forward/backward comparisons now use one common eager oracle. + +Local phase and memory summaries plus multi-part optimizer selection and updates each become coherent lifecycles instead +of narrow reports. The optimizer/weights Pydantic field-echo report is removed; its unique legacy session-id behavior now +drives the real optimizer endpoint. Tensor-collator coverage replaces a flat pseudo-packed example with the actual +already-batched dict and nested packed-dataset structures. A slicing-comprehension OPD probe and the one-line +sequence-parallel no-group identity probe are removed, while the behavioral pipeline and distributed contracts remain. + +All 28 focused CPU, server, distributed, and CUDA reports pass. Repository collection is now 1,292 items and the static +inventory is 1,281 definitions with 751 curated decisions across 356 Python test files. The audit still contains only 26 +legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate +distributed wrapper group unchanged. + +## One-hundred-and-eighteenth wave: remove tautologies and follow data through real consumers + +The one-hundred-and-eighteenth wave removes 14 reports across numerical MoE backward, router training, data preparation, +BI GEMM configuration, EP adapters, routing replay, OPD loss, checkpointing, and distributed state. A direct FP32 +grouped-GEMM helper probe is removed because exact local and EP custom-autograd oracles already exercise and discriminate +that accumulator. Train-router defaults now feed `MoEBlock.from_config` and prove detached-gate behavior by backward, +alongside the enabled all-to-all path and DeepEP rejection. + +Packing cache identity moves from a string-interpolation test into the real `PackingDataset` cache-path lifecycle and now +covers the ring-attention alignment suffix that the old test missed. A one-line `hashlib` wrapper report is removed. +Two BI GEMM reports are also removed: one set environment variables only after module import and therefore could not test +its claim, while the other asserted that lookup returned the constant it directly injects. The retained CUDA policies +prove table bit neutrality, cross-bucket row invariance, and available-backend parity instead. + +Native, Triton, Triton MoE-act, and Quack adapter argument boundaries now form one capability-aware matrix. Routing wire +decode composes with float-weight construction and sequence-parallel slicing. OPD output edges, DCP process-group +selection, ParallelState construction, and EP LoRA initialization-to-slicing each become coherent policies rather than +adjacent fragments. + +Focused validation reports 30 passed and one legitimate unavailable-DeepGEMM skip. Repository collection is now 1,278 +items and the static inventory is 1,267 definitions with 762 curated decisions across 356 Python test files. The audit +still contains only 26 legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and the +intentional duplicate distributed wrapper group unchanged. + +## One-hundred-and-nineteenth wave: compose quantization configuration into execution + +The one-hundred-and-nineteenth wave removes 10 reports across FP8 configuration, QARL calibration and execution, +stochastic rounding, and API response projection. Supported NeMo FP8 translation and every unsupported external runtime +form now share one compatibility policy. NVFP4 normalization no longer ends at dictionary-field assertions: aliases, +activation selection, and invalid group sizes feed the retained `QARLLinear` forward and STE contract. + +Activation fake-quant forward values and straight-through gradients form one autograd policy. The QARL MoE shadow context +now covers backend admission, no-op modes, and exception restoration together, while activation-quant overrides cover +mixed prior state, both directions, exceptions, and nesting in one lifecycle. Calibration follows parsed and truncated +input through persistent metadata and state restoration. QARL sync configuration similarly flows from an injected model +through selective quantization into the production weight-sync handler. + +The real Triton QARL MoE report now compares enabled lossy quantization and gradients with the exact disabled passthrough +using the same seeded model and routing. Seeded stochastic-rounding reproducibility joins its call contract while unbiased +expectation and adjacent-neighbor properties remain independent. API auto-load information and executor timing fields now +share one response-projection policy; duplicate legacy optimizer-payload coverage is removed from the focused telemetry +test because the broader optimizer policy already owns aliases, Adam fields, defaults, and precedence. + +All 17 focused CPU, server, and CUDA reports pass. Repository collection is now 1,268 items and the static inventory is +1,257 definitions with 772 curated decisions across 356 Python test files. The audit still contains only 26 legitimate +conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate distributed +wrapper group unchanged. + +## One-hundred-and-twentieth wave: follow requests and batches through complete lifecycles + +The one-hundred-and-twentieth wave removes eight reports across request processing, scheduling, packing, distributed data +loading, and orchestrator integration. RequestProcessor readiness and counters now surround three real model passes in one +start-to-stop lifecycle. Scheduler coverage now uses production requests for FIFO, capacity, terminal states, statistics, +clear, and bounded history; constructor constants, repr strings, and raw deque probes are gone. This consolidation also +corrects a misleading running-abort check that had dispatched an older FIFO request and then aborted the named request +while it was still pending. + +A GPU-marked micro-batch report is removed because it had no rank or distributed behavior and duplicated the retained CPU +contract. Packed label generation and position resets now sit inside the pack-to-unpack roundtrip, while exact sequential +legacy layout begins the all-strategy correctness policy. A second mixed-oversized skip assertion is removed from generic +edge cases because admission policy already owns it. + +Finally, orchestrator statistics now describe the retained forward, optimizer, and health transactions. A separate report +that created requests only to move counters, called getters repeatedly to prove they were read-only, and checked that +private methods were callable is gone. Its abort fragment is also removed: the immediate dummy backend usually finished +before abort arrived, and the test asserted only that the original request emitted some output. + +All 35 focused CPU, server, data-loader, and orchestration reports pass. Repository collection is now 1,260 items and the +static inventory is 1,249 definitions with 781 curated decisions across 356 Python test files. The audit still contains +only 26 legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional +duplicate distributed wrapper group unchanged. + +## One-hundred-and-twenty-first wave: make transfer and restore modes observable + +The one-hundred-and-twenty-first wave removes 11 reports across weight synchronization and checkpoint restore. P2P warm +mode no longer ends at a private boolean selector: cached prepare drives the fake Mooncake engine's async API while cold +prepare with the same setting stays synchronous. Small-entry chunking and persistent registration likewise move from +direct helper calls into repeated GPU-direct `transfer_bucket` transactions. The engine observes two-plus-one chunking, +one registration across both buckets, and deregistration during destroy. Direct-EP capability field echoes are removed; +the retained suite uses implicit and explicit sender maps through filtered scatter, dense partitioning, collective +failure, prewarm, and rank-owned transfer policies. + +Weight-sync bucket defaults and environment precedence now sit beside actual byte-cap splitting. Cache metadata is +asserted after complete streaming FP8 and sparse-delta syncs instead of through a standalone dictionary normalizer. +Compile-wrapper name cleanup joins the inference-layout unfusion policy and verifies emitted names and tensors across +experts, broadcast, and Qwen linear attention. + +Checkpoint restoration now follows one EP state through mesh admission, the `ModelState` caller, and final DTensor +construction for legacy and pipeline-parent layouts. CheckpointManager materialization includes successful counter and +optimizer policy plus both zero-meta failure boundaries. ModelRunner's completion flag now surrounds the real initial +checkpoint load for optimizer-enabled, optimizer-disabled, and failing cases rather than a separately mocked wrapper. + +All 58 retained weight-sync and checkpoint reports pass. Repository collection is now 1,249 items and the static +inventory is 1,238 definitions with 791 curated decisions across 356 Python test files. The audit still contains only 26 +legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate +distributed wrapper group unchanged. + +## One-hundred-and-twenty-second wave: keep one authoritative API transaction + +The one-hundred-and-twenty-second wave removes eight reports across session, sampler, inference, and future APIs. Three +copies of canonical-LoRA sampler export become one authoritative checkpoint-path transaction that checks normalized +session metadata, `save_lora_only`, model identity, output path, and returned URI. Create-model admission now handles +conflicting recreation, a distinct base repository, and cross-rank registration rollback together. HF snapshot identity +is no longer tested by calling a private canonicalizer: a bare client repository ID successfully registers against a +cache-resolved server path, while a different repository is rejected through the same endpoint. + +Sampler listing, path resolution, recency tracking, and last-receiver removal now form one adapter lifecycle. Receiver +quantization similarly follows `config.json` detection through name normalization, per-call skip-list enrichment, and +accepted default configuration; MTP, static activation, UE8M0, compressed-tensors, and BF16 boundaries remain in that +policy. + +Finally, a `FutureEntry` report that constructed fields and manually assigned every terminal enum is removed. Real store +jobs already reach pending, processing, completed, failed, and expired states, and queue pause state now follows actual +concurrent processing and statistics. + +All 32 retained API and FutureStore reports pass. Repository collection is now 1,241 items and the static inventory is +1,230 definitions with 797 curated decisions across 356 Python test files. The audit still contains only 26 legitimate +conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate distributed +wrapper group unchanged. + +## One-hundred-and-twenty-third wave: preserve scenarios and collapse trainer helper echoes + +The one-hundred-and-twenty-third wave removes five reports across DeepSeek construction, Trainer bootstrap, manual CUDA +timing, and LoRA dtype selection. DeepSeek router rejection and successful freezing now form one builder policy, while a +direct tensor-parallel validator probe is removed because the retained `build_parallelize_model` policy reaches that same +guard through its production consumer. Trainer bootstrap now covers eager and Quack-linear causal-loss configuration in +one setup rather than rebuilding the full fixture solely to omit `lm_head_fp32`. + +Manual CUDA timing now follows one disabled-to-enabled lifecycle through invalid-mode rejection, forward and recompute +recording, unrecorded-event omission, and draining. The mixed-precision LoRA builder already observes BF16 base weights, +FP32 adapters, and generic-upcast suppression, so the adjacent helper policy keeps only its distinct QLoRA, +explicit-skip, and dense-default branches. The experiment simulator suite is unchanged: its reports exercise separate +scenario ingestion, calibration, trust, topology, ledger, and correctness-gate behavior rather than scale-only variants. + +All six retained focused trainer and timing reports pass. Repository collection is now 1,236 items and the static +inventory is 1,225 definitions with 802 curated decisions across 356 Python test files. The audit still contains only 26 +legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate +distributed wrapper group unchanged. + +## One-hundred-and-twenty-fourth wave: replace P2P scale certification with protocol branches + +The one-hundred-and-twenty-fourth wave removes nine reports from the largest remaining weight-sync file. Expert FP8 +receiver coverage no longer performs production-sized local-shard allocation plus an eight-rank sweep merely to enumerate +all 256 global expert indices. Global names use a size-independent EP-rank formula, so two retained transactions now cover +the actual boundaries: partial blocks, block-128 quantization, one and two receivers, a nonzero EP offset, exact bytes and +scales, and dequantized parity. This reduces eleven expert transfer transactions to two. A separate 2048-by-512 Qwen3.6 +shared-expert report is also removed because the retained shared-expert transaction already covers both namespace +spellings, fused gate/up placement, down placement, scales, and the passthrough gate. + +Receiver-name compatibility now follows compiled-name stripping, `language_model` prefix fallback, and a missing tied +`lm_head` through one transfer policy. Direct-EP dense ownership applies one assignment to both receiver manifests and +outgoing buffers, including fused/split projection aliases and expert exclusion. Nonzero-sender initialization covers +default and prewarmed engine ordering in one policy, and rank filtering covers both owning-rank routing and a fully +filtered non-owner bucket. + +All 34 retained P2P protocol reports pass. Repository collection is now 1,227 items and the static inventory is 1,216 +definitions with 808 curated decisions across 356 Python test files. The audit still contains only 26 legitimate +conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate distributed +wrapper group unchanged. + +## One-hundred-and-twenty-fifth wave: move norm arithmetic into model consumers + +The one-hundred-and-twenty-fifth wave removes seven static reports across Qwen3-MoE residual handling and cross-engine +RMSNorm. Layer-zero, later-layer, and final-model norm declarations now form one Qwen construction policy. A direct +TP-shard materialization probe is gone: discriminating BF16 shard values now flow through the decoder layer and verify +sum-before-residual association at the materialized input, norm call, returned residual, and diagnostic sites. Likewise, +the two O-projection residual association modes now run through `_pre_mlp_forward`; the former private-helper report is +subsumed by exact norm-input, output, residual, capture, and bit-difference assertions. + +Cross-engine RMSNorm funnels now execute inside their supported site policies. Q/K, pre-summed residual-tree, and fused +post-attention paths each compare the real XoRL module with the corresponding SGLang kernel and SGLang family funnel over +every retained adversarial shape. The rare family-difference discriminator is part of the Q/K policy instead of a +standalone test of the test. A nominal trunk-flag report is removed because its fixture explicitly declared the +no-residual family, which selected the same wrapper before the trunk flag could affect dispatch. + +All five retained Qwen3-MoE reports pass. The five-report cross-engine file passes lint and collection parsing but remains +module-skipped because the optional SGLang package is unavailable in the repository venv. Repository collection is now +1,224 items and the static inventory is 1,209 definitions with 813 curated decisions across 356 Python test files. The +audit still contains only 26 legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and +the intentional duplicate distributed wrapper group unchanged. + +## One-hundred-and-twenty-sixth wave: distinguish distributed outcomes from mocked rank labels + +The one-hundred-and-twenty-sixth wave removes three GPU- and distributed-marked data-loader reports that never launch a +process, use a GPU, or observe a distinct distributed result. The partitioning report injected a mocked sampler but proved +non-overlap over rank-index lists constructed by the test itself. Its micro-batch and sequence-parallel fragments duplicated +the retained CPU policies that check exact split values, sampler ownership arguments, and collator insertion. + +The standalone sequence-sharding report asserted only equal output lengths and a loose padding range across mocked ranks. +The retained `TextSequenceShardCollator` policies discriminate exact rank slices, non-divisible padding, labels, attention +metadata, and token-aligned side channels. The packed report likewise changed mocked DP ranks but asserted the same shape; +real packed and variable-length samples already traverse the production data loader, while the collator policy owns exact +values, position resets, padding, extra fields, and flash-attention metadata. Its one distinct boundary, dropping an +incomplete loader batch, now runs inside the retained production data-loader lifecycle. + +All 11 focused data-loader and collator reports pass. Repository collection is now 1,221 items and the static inventory is +1,206 definitions with 816 curated decisions across 355 Python test files. The audit still contains only 26 legitimate +conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate distributed +wrapper group unchanged. + +## One-hundred-and-twenty-seventh wave: keep adapter transitions and GLM composition, not repetitions + +The one-hundred-and-twenty-seventh wave first audits the remaining data-preparation reports and leaves them intact: each +maps to a distinct packing algorithm, preprocessing, persistence, source-routing, file-lock, retry, or cache-identity +lifecycle. It then removes two redundant reports from adapter management and the GLM semantic stack. + +A direct manager report that loaded a SignSGD checkpoint into an AdamW-default manager asserted only the resulting optimizer +type and session field. The retained real `AdapterCoordinator` lifecycle already performs the same checkpoint-driven +selection through both explicit load and eviction auto-load, and the multi-adapter lifecycle reloads mixed optimizer types. +The GLM semantic stack now keeps only its four-layer transaction: the one-layer row changed repetition count but no branch. +Four layers still prove every canonical boundary, final-logprob parity, batch permutation, per-row composition, and a +negative discriminator where skipping canonicalization at the first layer changes the final logprobs. + +All three focused adapter and GLM reports pass. Repository collection is now 1,219 items and the static inventory is 1,205 +definitions with 818 curated decisions across 355 Python test files. The audit still contains only 26 legitimate conditional +runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate distributed wrapper group +unchanged. + +## One-hundred-and-twenty-eighth wave: replace shape smoke and synthetic architecture stubs + +The one-hundred-and-twenty-eighth wave audits every remaining pytest parameter matrix. Those rows now mostly switch real +backends, checkpoint APIs, tensor families, or compiled kernel specializations, so they remain. A thin-wrapper and +shape-only scan instead removes two weak model reports. + +The Qwen Triton expert smoke initialized random weights, ran one forward, and asserted only that output shape equaled input +shape. The retained eager-versus-Triton MoE transaction uses the same backend while comparing numerical outputs and every +LoRA factor gradient. A separate synthetic DeepSeek-like module duplicated the real model's default MLA LoRA targets. Its +unique explicit-target partition case now runs through `inject_lora_into_model_with_moe` on a real tiny DeepSeek model and +also proves the untargeted output projection is left untouched; the five-linear stub module is gone. + +Both retained covering reports pass. Repository collection is now 1,217 items and the static inventory is 1,203 definitions +with 820 curated decisions across 354 Python test files. The audit still contains only 26 legitimate conditional runtime +gates, with no no-outcome candidates, no parse errors, and the intentional duplicate distributed wrapper group unchanged. + +## One-hundred-and-twenty-ninth wave: make configuration change execution + +The one-hundred-and-twenty-ninth wave removes two constructor-only reports. A base MoE expert test instantiated Qwen wrappers +and four backend variants but asserted only stored strings, parameter shapes, LoRA mapping identity, and registry membership. +Retained eager, native, Triton, non-gated, injection, and model-construction policies execute those same registrations and +layouts. The separate LoRA initialization policy remains because frozen base weights and trainable zero-initialized factors +are not observable from forward parity alone. + +DSv4 KV-QAT coverage no longer stops at a helper boolean and the private `_kv_qat_enabled` field. The retained C0 attention +forward/backward transaction now supplies an FP8 quantization configuration, observes the QAT call on the exact no-RoPE KV +slice with block size 64, and then checks finite forward output plus input and parameter gradients. The C128 branch remains +in the same shape-and-gradient policy. + +All five focused MoE and DSv4 reports pass. Repository collection is now 1,215 items and the static inventory is 1,201 +definitions with 822 curated decisions across 354 Python test files. The audit still contains only 26 legitimate conditional +runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate distributed wrapper group +unchanged. + +## One-hundred-and-thirtieth wave: make cache and checkpoint metadata prove execution + +The one-hundred-and-thirtieth wave removes three reports that stopped at shapes, private fields, or constructor topology. +DSv4 RoPE cache precedence now runs through real consumers: the C0/C128 attention forward-backward policy constructs and +uses the environment-sized cache, while the context-parallel C128 compressor forward requires the config-sized fallback to +cover its nonzero-rank slice. The standalone cache-builder shape report is gone. + +DeepSeek packed-checkpoint configuration likewise no longer ends at the selected handler type and its private bit-width and +group-size fields. The retained loader transaction obtains the handler through the model, parses the official nested +compressed-tensors configuration, and dequantizes actual 8-bit/group-64 expert payloads to exact gate, up, and down values. +Default packed loading and requested BF16 output remain in the same transaction. + +Finally, DSv4 component selection is observed through complete paths instead of a presence-only report. The retained C0 and +C128 attention transactions execute their respective topologies, and the C4 synthetic checkpoint load constructs both the +compressor and indexer and validates their separately translated APE tensors. + +All six focused attention, compressor, checkpoint-loader, and cache-boundary reports pass. Repository collection is now +1,212 items and the static inventory is 1,198 definitions with 825 curated decisions across 354 Python test files. The audit +still contains only 26 legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and the +intentional duplicate distributed wrapper group unchanged. + +## One-hundred-and-thirty-first wave: remove dependency canaries and scale repetitions + +The one-hundred-and-thirty-first wave removes four static reports across linear attention and batch-invariant reductions. +The Hopper Gated Delta Rule report imported and executed the optional FLA package directly without reaching an XoRL module, +wrapper, or integration seam. Its illegal-memory/autotuner result was therefore an upstream dependency canary rather than a +repository regression, so the file is gone. + +Full-reduce mean now owns its FP32-output dtype spelling in the same contract as the one- and two-dimensional BF16/FP32 +mean-versus-sum boundaries. Head-v2 likewise keeps one authoritative focused suite. The removed production-hidden and +production-vocabulary reports repeated exact v1 projection bits, the shared decode/scoring statistics tree, and prefix batch +invariance without selecting another launch branch. The retained focused policies additionally prove arbitrary-slice +invariance, selected-logprob composition, fused-loss gradients, and the family-v1 rollback. + +All six retained family-selection, head-v2, and mean reports pass. Repository collection is now 1,209 items and the static +inventory is 1,194 definitions with 828 curated decisions across 353 Python test files. The audit still contains only 26 +legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate +distributed wrapper group unchanged. + +## One-hundred-and-thirty-second wave: stop testing orphan compatibility helpers + +The one-hundred-and-thirty-second wave removes four reports whose private targets have no production callers. The eager +OPRD layer-cache gather wrapper was called only by its test; the live loss path uses the retained streaming slice fetcher, +which verifies selected cache rows, multiple layer ranges, total layer count, and returned shapes. + +The legacy `_accumulate_is_metrics` and `_finalize_is_metrics` pair is likewise absent from runtime call sites. Forward- +backward now routes loss metrics through `_accumulate_loss_metrics` and `_finalize_loss_metrics`, whose retained OPD policy +covers mean, extrema, empty-rank, and loss-specific behavior. Two distributed reports and a CPU report for the dead pair are +gone; the still-live per-micro-batch `_sp_allreduce_kl_metrics` collective remains covered by its two-rank NCCL gate. + +All three retained OPRD fetcher, current metric-aggregation, and live NCCL reduction reports pass. Repository collection is +now 1,205 items and the static inventory is 1,190 definitions with 830 curated decisions across 352 Python test files. The +audit still contains only 26 legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and the +intentional duplicate distributed wrapper group unchanged. + +## One-hundred-and-thirty-third wave: prefer live paths over duplicate wrappers and fabricated states + +The one-hundred-and-thirty-third wave removes seven reports across EP kernels, QARL, DSv4, server launch, routing replay, +and activation offload. A forward-only Triton/Quack routing-score report duplicated the same test stubs and torch reference +used by a retained report that additionally proves routing-score gradients. A direct `QARLLinear` smoke likewise repeated +gradient and persistence observations already covered by retained injected-model calibration and full optimizer/checkpoint +lifecycles. + +DSv4 fallback coverage now enters through public `rotate_activation`: one transaction disables the optional kernel and +proves the known transform, self-inverse behavior, and norm preservation across supported widths. Separate private-helper +algebra, impossible-width rejection, and shape-only dispatch reports are gone. The launcher report that manually combined a +valid flat config with missing parsed server arguments and the routing-replay report that overwrote a private global with a +made-up stage both manufactured states normal callers cannot create. Finally, the standalone activation-offload report +passed `None` outside the typed trainer/server argument boundary and asserted only that an unrelated square operation had a +gradient; it observed no offloading behavior. + +All 17 executed focused reports pass, with one legitimate optional-backend skip. Repository collection is now 1,198 items +and the static inventory is 1,183 definitions with 836 curated decisions across 351 Python test files. The audit now contains +25 legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional duplicate +distributed wrapper group unchanged. + +## One-hundred-and-thirty-fourth wave: stop preserving test-only support surfaces + +The one-hundred-and-thirty-fourth wave removes the three-report `FileLockLoader` suite. The class is neither exported from +`xorl.data.prepare` nor referenced anywhere else under `src`; only its own tests named it. Live dataset preparation remains +covered through packing-cache persistence, preprocessing, source loading, hashing, retries, and collator/data-loader +lifecycles. + +Two broader server policies were also narrowed to live interfaces. Runner messages still round-trip every payload and tensor +through the production transport codec and reject pickle input, but no longer preserve unused `BaseMessage` JSON helpers. +`FutureStore` scheduling, result/failure storage, deletion, model cleanup, and expiry now inspect the live `FutureEntry` +returned by `get`; assertions for six convenience accessors with no production caller are gone. + +All 14 focused server-protocol, future-store, and live data-preparation reports pass. Repository collection is now 1,195 +items and the static inventory is 1,180 definitions with 839 curated decisions across 350 Python test files. The audit still +contains 25 legitimate conditional runtime gates, with no no-outcome candidates, no parse errors, and the intentional +duplicate distributed wrapper group unchanged. + +## One-hundred-and-thirty-fifth wave: fold helper checks into consumers and drop a dormant campaign + +The one-hundred-and-thirty-fifth wave removes two direct helper reports from distributed and FP8 coverage. PyTorch's wrapped +reduce operation now remains qualified by the retained real two-rank FSDP2 custom-reduce-scatter lifecycle rather than a +four-line private canonicalizer truth table. Full-precision expert FSDP kwargs are now checked on an actual exact GLM shared +expert in the retained topmost-unit topology policy, replacing a fake class with one boolean field. + +Four SM90 P5 reports are also gone. They certified an opt-in GDN decode-prep candidate that is not exported, documented, or +called by production; its only source consumer is another unused decode-solve candidate. Those bitwise and graph-capture +campaign gates therefore protected no reachable XoRL execution path. + +All 11 retained focused topology, FP8 checkpoint, and FSDP2 reports pass, including the real two-rank FSDP lifecycle. +Repository collection is now 1,189 items and the static inventory is 1,174 definitions with 842 curated decisions across 349 +Python test files. The audit still contains 25 legitimate conditional runtime gates, with no no-outcome candidates, no parse +errors, and the intentional duplicate distributed wrapper group unchanged. + +## One-hundred-and-thirty-sixth wave: remove support surfaces that only their tests consume + +The one-hundred-and-thirty-sixth wave removes an idealized pipeline-bubble formula with no caller and keeps the measured +`PPBubbleProfiler` transaction used by `Trainer`. Data preparation no longer carries a separate first-fit-decreasing +feasibility checker that never participates in packing; the retained report exercises `pack_group`, sequential allocation, +and `PackingDataset`. Likewise, the grouped and any-FQN matchers are gone because only their truth tables called them, while +the retained matcher remains used by parallel plans and sharded adapter state. + +EP synchronization now has one production model: the coalesced optimizer-boundary reducer. A standalone per-parameter hook, +its single-gradient reducer, and a test-only statistics alias were removed from the real multi-rank report. That report still +proves two-rank clipping, non-finite rejection, bucket accounting, and three-rank participation masks. + +The larger removal is an unreachable DeepSeek-V4 indexer autograd campaign. Production `V4Indexer` calls the forward score +kernel directly and returns discrete top-k indices; it never imported the separate autograd wrapper or backward kernel. The +two parameterized reports for those modules contributed eight collected cases. The retained seven V4 reports all pass on +GPU and cover the reachable forward scores, causal mask, production geometry, numerical range, and zero input. Together with +16 focused CPU and real multi-rank passes, repository collection is now 1,180 items and the static inventory is 1,171 +definitions with 847 curated decisions across 349 Python test files. The audit still contains 25 legitimate conditional +runtime gates, no parse errors, and the intentional duplicate distributed wrapper group unchanged. + +## One-hundred-and-thirty-seventh wave: test behavior through consumers, not synthetic compositions + +The one-hundred-and-thirty-seventh wave deletes the manual CUDA timing module and its sole lifecycle report. No production +module, package export, documentation, example, or script could enable or drain that instrumentation; mocked CUDA events +were testing an isolated subsystem with no XoRL consumer. + +KV repetition is now qualified through eager GQA attention. The retained transaction compares full attention weights and +outputs with an independent `torch.repeat_interleave` reference, replacing a standalone helper smoke that checked shapes, +one tiny pattern, and device preservation. The synthetic “full MoE pipeline” report is also gone: identity experts reduced it +to the same scatter-gather round trip already checked numerically, while the retained kernel and model suites cover routing, +real expert computation, and gradients. Finally, the test-only pipeline single-stage predicate duplicated the live schedule +style table and was removed without weakening schedule construction or admission coverage. + +All 11 focused attention, MoE, and pipeline reports pass; the two skips are the expected unavailable FA3 interface. Repository +collection is now 1,177 items and the static inventory is 1,168 definitions with 851 curated decisions across 348 Python test +files. The audit still contains 25 legitimate conditional runtime gates, no parse errors, and the intentional duplicate +distributed wrapper group unchanged. + +## One-hundred-and-thirty-eighth wave: stop tests from designing unused APIs + +The one-hundred-and-thirty-eighth wave removes six groups of production support surfaces that existed only to make direct +unit assertions convenient. GLM sparse selection no longer carries an unused physical-page translator or selected-value +gatherer, and its inventory and layer plan no longer precompute role and schedule views solely for tests. The retained +contracts derive those observations from the canonical target and layer tuples and still execute the live logical-index +selector across ties, short rows, dead rows, and the production boundary tail. + +NVFP4 now exposes and tests its real format-specific fake quantizer rather than a private one-format string dispatcher. +Native FP8 checkpoint protection likewise keeps the `DistributedCheckpointer`-used real-DCP preflight and removes a +metadata-dictionary adapter whose only caller fabricated that dictionary in a test. `ModelState` no longer advertises a +future lightweight safetensors reference collector with no exporter or save-path consumer; the QARL report retains actual +checkpoint metadata and compatibility rejection. + +Finally, routing replay follows its real transaction lifecycle. Tests no longer preserve cursor-reset methods absent from +the trainer and R3 handler; replay advances the backward cursor and `clear_all` performs teardown. The six focused suites +report 40 passes and one expected optional-SGLang skip. Collection and the static test inventory are unchanged at 1,177 +items and 1,168 definitions because this wave removes test-only support and narrow assertions inside retained behavioral +reports. The ledger now contains 857 curated decisions across 348 Python test files. + +## One-hundred-and-thirty-ninth wave: retire diagnostic router policy matrices + +The one-hundred-and-thirty-ninth wave removes two process-wide router diagnostics that only their tests selected. +`XORL_MOE_ROUTER_TOPK_POLICY` injected stable-sort, artificial tie bias, or raw-logit selection into every ordinary router, +but had no configuration, documentation, script, example, or production consumer. The live softmax and DSv4 paths now call +`torch.topk` directly; retained reports continue to cover softmax weighting and normalization, correction bias, hash +routing, balanced profiling, exact batch-invariant routing, and rejection boundaries. + +The layer-list `XORL_MOE_ROUTER_FP32_LAYERS` parser is also gone. Its standalone report is folded into the real model +configuration contract, which drives `_router_fp32` through `MoEBlock` and observes FP32 hidden and gate operands. This +preserves the shipped configuration behavior while deleting one test report and an environment-only parallel interface. + +Finally, dense, fused-delta, and MoE LoRA modules no longer expose manual merged-weight cache invalidators with no caller. +Their caches already key entries on tensor versions, storage pointers, active rank, and alpha; retained tests prove +automatic invalidation across optimizer steps and runtime rank changes and ensure old fused generations are released. All +20 focused router, merged-LoRA, and Qwen integration reports pass. Repository collection is now 1,176 items and the static +inventory is 1,167 definitions with 860 curated decisions across 348 Python test files. + +## One-hundred-and-fortieth wave: retire test-driven numerical experiments + +The one-hundred-and-fortieth wave removes two environment-only numerical campaigns whose tests never established the +claimed behavior. `XORL_MOE_FP64_ACCUM` selected a slow, inference-only expert detour with no serving counterpart, +configuration, launcher, documentation, or example. Its sole assertion replaced that detour with a mock and proved only +that the switch outranked the fused-SGLang switch. The normal eager, Triton, fused-SGLang, EP, and TP-simulation paths and +their real forward, gradient, layout, determinism, and admission coverage remain. + +Qwen3-MoE no longer contains the delayed-residual tuple, TP-shard-carry, alternate post-attention residual formulas, forced +RMSNorm flags, or candidate-capture matrix. Those branches were reachable only through undocumented process-wide +variables, and four reports exercised them by directly passing private tuple states or attaching attributes to tensors in +test doubles. The retained Qwen report instead executes the normal decoder and final-norm consumers and checks the shipped +no-residual versus residual-tree family declaration. Removing the experiment also removes its generic delayed-output hook +and stale diagnostic-site registrations. + +The primary focused gate reports 10 passes and two expected skips because SGLang is unavailable; the adjacent real TP, +diagnostic-capture, and module-utility suites add 26 passes. Repository collection is now 1,172 items and the static +inventory is 1,163 definitions with 862 curated decisions across 348 Python test files. The audit still contains 25 +legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-forty-first wave: remove helpers whose tests were their product + +The one-hundred-and-forty-first wave deletes the obsolete stacked-LoRA helper module. Its initialization, delta, merge, and +unmerge functions had no production, documentation, example, or script consumer; a package re-export and one arithmetic +truth table were the entire interface. The scaling formula is the exception: it remains directly at the group-GEMM package +boundary because dense, MoE, and quantized adapters use it. Twenty-two retained LoRA, QLoRA, loading, construction, gradient, +and optimizer reports pass through those real consumers after the move. + +Dense op parity also loses a tautology. The private eager SwiGLU wrapper was one `torch.nn.functional.silu` expression and +had no source caller; its report compared it to the identical expression labeled as a serving reference. The independent +RoPE implementation comparison remains, while real fused SwiGLU forward and backward behavior continues to run through the +Triton operator, exact GLM MLP composition, and model-level adapter suites. + +Repository collection is now 1,170 items and the static inventory is 1,161 definitions with 864 curated decisions across +347 Python test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, +and no parse errors. + +## One-hundred-and-forty-second wave: delete a tested kernel with no execution path + +The one-hundred-and-forty-second wave removes the specialized families-v2 Q/K-normalization kernel, its standalone report, +and its frozen golden row. Despite the production-contract language in the test, no model, trainer, dispatcher, +configuration, documentation, example, or script called `qk_norm_v2`; only those two test invocations reached it. The actual +Qwen3.5 and Qwen3.5-MoE attention paths instantiate their declared RMSNorm modules. + +The live families-v2 hidden-state RMSNorm remains intact. Its fused and split realizations, dispatch threshold, exact-model +family selection, cross-structure equivalence, and frozen numerical trees all remain covered. All 25 focused families-v2, +RMSNorm, exact-model selection, and golden-tree reports pass. Repository collection is now 1,169 items and the static +inventory is 1,160 definitions with 865 curated decisions across 347 Python test files. The audit still contains 25 +legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-forty-third wave: retire a diagnostic backend matrix + +The one-hundred-and-forty-third wave removes the `XORL_SGLANG_MOE_TP_SIM` campaign. Its direct, cache, Triton, +alternate-reduce, DeepGEMM, fused-kernel, and runner modes were an environment-only K3 diagnostic with no repository +configuration, launcher, example, documentation, or ordinary execution-path consumer. Commit history also labels later +changes to the lane as experiments and diagnostics. Nine reports exercised it with tiny full-local tensors and Python +fakes for every optional backend, so the tests largely specified a parallel simulation product rather than validating a +supported training topology. + +The supported SGLang fused-expert implementation remains intact: local and EP dispatch, serving weight layouts, autograd, +cache invalidation, runtime-context admission, and failure boundaries retain their real suites. The simulation-only runner +loader, shard attributes, diagnostic capture names, and MoEBlock bypass conditions are gone with the test file. Fifty-one +focused MoE, LoRA, checkpointing, and diagnostic-capture reports pass, with three expected optional-SGLang skips. +Repository collection is now 1,160 items and the static inventory is 1,151 definitions with 866 curated decisions across +346 Python test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, +and no parse errors. + +## One-hundred-and-forty-fourth wave: stop claiming unexecuted backend coverage + +The one-hundred-and-forty-fourth wave removes a CPU-only SGLang RMSNorm report. Although it selected the `sglang_jit` and +`sglang_kernel` diagnostic modes, it never ran their CUDA implementations, optional-package loaders, ABI boundaries, or +serving arithmetic; both cases simply followed an eager CPU formula. The report therefore advertised backend coverage +that it did not provide. + +Those explicit diagnostic modes remain available. The retained RMSNorm suites cover ordinary and fused CPU arithmetic, +fused CUDA forward and backward, exact Qwen model-site integration, family admission, global configuration forwarding, +and real kernels when the optional runtime is present. Repository collection is now 1,159 items and the static inventory +is 1,150 definitions with 867 curated decisions across 345 Python test files. The audit still contains 25 legitimate +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-forty-fifth wave: remove speculative weight-sync and legacy payload surfaces + +The one-hundred-and-forty-fifth wave moves into server and weight-sync coverage. Sparse-delta receiver sharding, +per-rank raw and encoded writers, translation-future collection, and terminal future serialization formed a closed +source/test-only subgraph: no runtime, configuration, documentation, example, or script called it. Three reports built +fake receiver shards and replaced the complete optional `delta_encoding` future API. That speculative layer is gone. +The live source-capture path remains, including optimizer-step snapshots, rank/global manifests, validated single-file +packing, and sparse-delta backend consumption. + +The same wave removes four deprecated R3 payload names: `externalize_r3_payloads`, `keep_r3_payloads`, +`routing_payload_dir`, and `keep_routing_payloads`. They existed only in compatibility branches and two test inputs; +launchers and runtime constructors already use the canonical transport, directory, retention, and namespace fields. +Canonical Mooncake and filesystem reports continue to cover creation, slicing, cleanup, retention, validation, and +serialization. All 34 focused server and weight-sync reports pass. + +Repository collection is now 1,156 items and the static inventory is 1,147 definitions with 869 curated decisions across +345 Python test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, +and no parse errors. + +## One-hundred-and-forty-sixth wave: replace fake bootstrap confidence with real boundaries + +The one-hundred-and-forty-sixth wave removes a trainer bootstrap report that instantiated `Trainer` through `__new__`, +constructed two large fake argument trees, and mocked every bootstrap dependency to observe one direct `ep_intranode` +keyword assignment plus one loss-dictionary branch. The retained distributed suite executes both EP mesh geometries; +loss reports exercise `quack_linear` computation and admission; argument and model-builder reports cover the meaningful +configuration boundaries without reproducing bootstrap implementation details. + +Muon also loses the undocumented `XORL_MUON_QUACK_TUNED` override. Its only assertion replaced Quack GEMMs with lambdas +and checked a keyword, while no configuration, documentation, example, or script exposed the switch. The qualified +`tuned=False` default remains, as do backend import, architecture/dtype dispatch, real optimizer update, grouped Gram +Newton-Schulz, and CUDA dtype reports. All 30 focused optimizer, loss, distributed, and trainer reports pass. + +Repository collection is now 1,155 items and the static inventory is 1,146 definitions with 871 curated decisions across +344 Python test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, +and no parse errors. + +## One-hundred-and-forty-seventh wave: keep data tests at the lifecycle boundary + +The one-hundred-and-forty-seventh wave removes the two-report `CollatePipeline` unit suite. It composed fake collators that +added and multiplied token IDs, then asserted the arithmetic and permissive single-collator, tuple, and empty-list constructor +forms. Production constructs the pipeline only from a non-empty collator list, and the retained dataloader integration runs +that real sequence through tensor conversion, flattening, token shifting, packing, micro-batch splitting, and optional +sequence sharding. The live pipeline now exposes only the sequence constructor it actually consumes. + +This wave also finishes source cleanup identified by earlier test decisions. The unexported `FileLockLoader` and the unused +SHA256 string wrapper had no runtime, documentation, or example caller after their standalone reports were removed. Retry +configuration loses linear and constant policies that only its test selected; the sole production decorator call uses the +retained exponential policy. Its remaining report now observes requested sleeps deterministically instead of measuring real +wall-clock intervals. + +All 26 focused data and checkpoint reports pass. Repository collection is now 1,153 items and the static inventory is 1,144 +definitions with 874 curated decisions across 343 Python test files. The audit still contains 25 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-forty-eighth wave: keep protocol tests on the live transport + +The one-hundred-and-forty-eighth wave removes a closed server-protocol convenience layer. Twelve response builders and four +validation or introspection helpers were exported and tested, but no orchestrator, API server, scheduler, dispatcher, +documentation example, or other runtime caller used them. Their tests manually populated response dictionaries and then +asserted those same fields. The retained protocol report now round-trips the actual typed request and output dataclasses +through msgpack, including operation-payload reconstruction and terminal result fields used by live ZMQ communication. + +The same wave removes another fully mocked Trainer forwarding report. It constructed `Trainer` via `__new__`, supplied a +large fake argument tree, replaced foundation-model construction, and asserted direct keyword copies for numerical flags and +LoRA scalars. Argument parsing already covers the configuration values, while model-policy suites exercise resolution, +admission, and numerical behavior; the forwarding report could not discriminate any of those outcomes. + +All 41 executed focused protocol, orchestrator, argument, and model-policy reports pass. Repository collection is now 1,151 +items and the static inventory is 1,142 definitions with 876 curated decisions across 342 Python test files. The audit still +contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-forty-ninth wave: remove shallow copies of lifecycle contracts + +The one-hundred-and-forty-ninth wave removes a dense NVFP4 injection smoke that wrapped two `Linear` modules and checked +their type, format string, and output shape. The retained QARL contracts already exercise targeted dense injection and +exclusions, wrapper counters and summaries, lossy NVFP4 arithmetic, straight-through gradients, a real optimizer update, +changed log-probabilities, and checkpoint restoration. The extra sequential-model shape check added no distinct failure +boundary. + +The same wave removes a mock-only weight-version forwarding report. Its fake handler/backend chain duplicated the retained +handler policy assertion that verifies `flush_cache` and `weight_version` at `transfer_bucket`, while the P2P protocol suite +verifies the version in the actual completion request body. Seven focused QARL, handler, and protocol reports pass. +Repository collection is now 1,149 items and the static inventory is 1,140 definitions with 878 curated decisions across +341 Python test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, +and no parse errors. + +## One-hundred-and-fiftieth wave: stop testing configuration plumbing with fake neighbors + +The one-hundred-and-fiftieth wave removes two builder-plumbing reports. One replaced foundation-model construction and +parallelization solely to assert that `fsdp_sharded_lm_head_loss=True` crossed one function call. The other replaced the +server's training-model builder and asserted unchanged FP8, QARL, and sharded-loss dictionary values. Neither report +constructed the claimed distributed loss, FP8 model, or calibrated QARL model. Retained suites execute sharded LM-head +loss under FSDP, perform real FP8 and QARL injection, run calibration before parallelization, update parameters, and +restore QARL checkpoints. + +This wave also removes a Dr.GRPO outer-loop report that replaced `_forward_loop` and then asserted that the string +`"drgrpo"` and its inputs reached that fake. The retained runner report executes the actual Dr.GRPO loss branch, including +clipping, KL, temperature, legacy fields, per-token output policy, and K3 output; independent runner and dispatcher +lifecycle suites cover completion, failure, model identity, routing setup, and step accounting. + +All 16 focused behavioral reports pass, including a real sharded-loss FSDP transaction. Repository collection is now +1,146 items and the static inventory is 1,137 definitions with 881 curated decisions across 341 Python test files. The +audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-fifty-first wave: prefer transformed payloads and constructed models + +The one-hundred-and-fifty-first wave removes a request-processor report that supplied raw routed-expert arrays, replaced +both backend methods with `AsyncMock`, and asserted that the same Python objects appeared in the mock kwargs. Retained +routing transactions exercise the behavior that can actually fail: Mooncake and filesystem encoding, datum reordering, +slice loading, cleanup on success and failure, wire decoding, rank-zero selection, and model identity. + +The same wave removes the remaining fake-builder report from the model-runner FP8 file. It replaced +`build_training_model`, asserted direct copies of block-FP8 QLoRA settings, and checked a target set already exercised by +the dedicated GLM target-resolution policy. Real builder suites retain foundation/injection configuration, precondition +rejection, adapter inventory ownership, target selection, quantized construction, and QLoRA execution. + +All seven focused routing and GLM construction reports pass. Repository collection is now 1,144 items and the static +inventory is 1,135 definitions with 883 curated decisions across 341 Python test files. The audit still contains 25 +legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-fifty-second wave: test module-path helpers through real plans + +The one-hundred-and-fifty-second wave removes the standalone distributed-utils suite. Its two reports restated recursive +`getattr`/`setattr`, regex wildcard matching, and exception behavior on a toy `Sequential`/`ModuleDict`. The retained +`ParallelPlan` suites invoke the same helpers through exact and wildcard FQNs while replacing meta parameters, slicing +global expert banks, preserving dtype/trainability, assigning replicated gradient domains, and materializing real GLM +adapter layouts. Sharded adapter-state tests exercise the same matcher on production ownership plans. + +The caller audit also removes the singular `find_free_port` launcher helper, which had no source, test, documentation, +example, or script caller. The live launcher uses the distinct `find_free_ports` allocator for its three- and four-port +rendezvous layouts. + +All 11 focused plan, sharded-state, and launcher reports pass. Repository collection is now 1,142 items and the static +inventory is 1,133 definitions with 885 curated decisions across 340 Python test files. The audit still contains 25 +legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-fifty-third wave: remove unintegrated helper islands + +The one-hundred-and-fifty-third wave removes the mock-only `RemoteBackend` suite. It replaced `_execute`, invoked two thin +wrappers, and asserted operation names, request IDs, timeouts, and fields copied into payload constructors without +serializing or transporting anything. Retained protocol, request-processor, dispatcher, sparse-delta, and weight-sync +suites exercise typed reconstruction and the downstream behavior of those fields. + +This wave also removes the standalone `xorl.rl` package and its five formula-repetition reports. Its six exported +Slime-style tensor helpers had no trainer, runner, loss, CLI, example, documentation, or other source caller. Integrated +loss suites continue to cover the actual policy, KL, importance-sampling, OPD, and Dr.GRPO paths used by training. + +The reachability audit found two smaller test-driven APIs as well. Runner acknowledgement and response factories had no +runtime caller; live transport constructs the dataclasses directly, and the retained protocol report round-trips them +through MessagePack. The sparse source-delta translation-input loader had no translation engine or receiver caller, and +its only report installed a synthetic `delta_encoding` package to fabricate empty shards. Source capture, packed-file +validation, backend upload, receiver application, and trainer-to-SGLang coverage remain. + +All 30 focused protocol, request-processing, sparse-delta, and integrated-loss reports pass. Repository collection is now +1,136 items and the static inventory is 1,127 definitions with 889 curated decisions across 338 Python test files. The +audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-fifty-fourth wave: make reference checks independent + +The one-hundred-and-fifty-fourth wave removes `canonical_moe_reduce_reference` from production code and drops the report +that tested that oracle itself. The oracle shared `_adjacent_pairwise_bf16` with the implementation it was meant to +validate, so actual-versus-expected comparisons were circular. Distributed and GLM model contracts now compute their +expected adjacent BF16 tree independently in test code. The real multi-process transport gate passes with permutation, +chunking, output-distribution, padding, and backward checks intact. + +The same reachability pass removes `SequencePartial` and its synthetic dense, packed, and hand-sliced context-parallel +matrix. No loss, trainer, runner, CLI, example, documentation, or other source caller selected it. `TokenPartial` is the +sole production reducer and remains covered through integrated causal-LM, policy, importance-sampling, OPD, Dr.GRPO, TP, +and FSDP paths. + +Two narrow server reports also go away. An older Tinker compatibility report duplicated focused session creation and +weights-info lifecycles; the retained endpoint report now includes its only distinct flat `lora_rank` assertion. A +sampler-prefill report converted one literal list to a tensor and captured it on a fake model, while retained GLM indexer +and sparse-attention contracts cover the actual prefill-boundary behavior. + +Eighteen focused reducer, runner, API, distributed, and GLM reports pass; one CUDA-only shared-expert report skips on this +host. Repository collection is now 1,132 items and the static inventory is 1,123 definitions with 893 curated decisions +across 338 Python test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## One-hundred-and-fifty-fifth wave: separate telemetry presentation from behavior + +The one-hundred-and-fifty-fifth wave removes two trainer reports that instantiated no trainer lifecycle. One called a +private activation-offload consumer on `SimpleNamespace` objects and asserted byte-to-GB field names, call counts, and an +empty dictionary. The other called private exception-path summarizers and asserted key ordering plus four copies of each +local float. The retained component-timer suite runs real forward and backward hooks on CUDA across GLM- and Qwen-shaped +layers and retains the unrecorded-event recovery policy; broader trainer suites exercise optimization and synchronization. + +Two fake forwarding seams are consolidated into one runner lifecycle report. The old pair stopped once at +`_execute_and_gather` and once at `_execute_compute`. The replacement runs the real rank-zero handler, gather wrapper, and +compute dispatch through to the trainer, while isolating only distributed side effects, so session identity and both R3 +payloads are checked across the complete chain. + +This wave also removes a source-text lint that opened `bi_families_v2.__file__` and searched for banned import strings. The +retained families-v2 suites exercise model-program selection, rollback, numerical dispatch, independent fused/split +realizations, cross-engine bytes, and real CUDA bit gates. The audit now recognizes module-source reads through +`Path(module.__file__).read_text()` as source inspection, preventing this pattern from hiding behind ordinary file I/O. + +All nine focused trainer, dispatcher, and families-v2 reports pass, including the CUDA component-hook path. Repository +collection is now 1,128 items and the static inventory is 1,119 definitions with 896 curated decisions across 336 Python +test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no +parse errors. + +## One-hundred-and-fifty-sixth wave: put repository policy in lint + +The one-hundred-and-fifty-sixth wave relocates the repository-wide private-reference scan out of pytest. It did not +exercise XoRL behavior: it ran `git ls-files`, decoded every tracked file, and matched policy regexes for private paths, +cluster identifiers, and authoring metadata. The same zero-dependency check now lives in `scripts/check_public_tree.py` +and is a local pre-commit hook, so the existing lint workflow still enforces it on every pull request without presenting +it as a product test. + +This wave also removes a duplicate fake register-session dispatcher report from the request-processor file. The dedicated +session-ops suite already forwards the exact typed payload and response through `_handle_register_session` and additionally +tests cross-rank rejection. The retained request-processor report separately runs registration through the processor and +`DummyBackend`, so both meaningful boundaries remain without a third fake coordinator. + +The public-tree lint passes, as do both focused registration paths. Repository collection is now 1,126 items and the +static inventory is 1,117 definitions with 898 curated decisions across 335 Python test files. The audit still contains +25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-fifty-seventh wave: distrust answers supplied by the mock + +The one-hundred-and-fifty-seventh wave removes an orchestrator-client report whose mock engine counted the samples it +had just received and echoed the request learning rate. Those assertions were properties of the test double, not the +server. Retained real-ZMQ reports cover interleaved forward and optimizer traffic plus exact wire payloads, and the real +orchestrator lifecycle covers empty-batch rejection. + +The same pass removes two duplicate API paths and one duplicate validation fragment. Focused endpoint reports already +exercise normalized LoRA worker registration and full-weight admission; the latter now supplies legacy empty optional +configs so that compatibility behavior remains explicit. The focused training-ops report already checks explicit +learning-rate forwarding, while the retained compatibility bundle still covers legacy Adam payloads and the effective +learning-rate fallback priority. At the request processor, only the distinct nonempty batch without valid targets +remains; the empty-list case belongs to the orchestrator end-to-end report. + +All eight focused socket, orchestrator, endpoint, and training-operation reports pass. Repository collection is now +1,124 items and the static inventory is 1,115 definitions with 901 curated decisions across 335 Python test files. The +audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-fifty-eighth wave: make a component contract a test-sized unit + +The one-hundred-and-fifty-eighth wave consolidates microscopic wrapper reports across exact GLM-5.2 LM-head QLoRA, +fused gate-up QLoRA, native FP8 model construction, exact QLoRA admission, and native block-FP8. Twelve collected items +did nothing beyond invoking one to three neighboring assertion helpers for one CPU component. The replacement reports +group topology, operand admission, byte presentation, gradients, construction, checkpointing, and failure behavior at +the component level, explicitly resetting monkeypatch state between independent seams. + +No behavioral assertion was removed. Separate Hopper gates stay separate, as do the native router and expert reports +whose runtime boundaries differ from buffer construction. All nine resulting focused contracts pass. Repository +collection is now 1,112 items and the static inventory is 1,103 definitions with 902 curated decisions across 335 Python +test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no +parse errors. + +## One-hundred-and-fifty-ninth wave: verify checkpoint translation at its destination + +The one-hundred-and-fifty-ninth wave removes a DeepSeek-V4 report that directly asserted private checkpoint-name string +rewrites and round-tripped the APE inverse against a test-defined forward transform. The retained synthetic checkpoint +transaction already drives the real loader across window, C4, hash, shared-expert, and routed-expert families. + +That transaction is now the stronger oracle: it checks loaded values at the embedding/head, norm, attention, HC, +router-bias, shared-expert, fused routed-expert, C4 APE, and renamed indexer destinations. The quantized codec and EP +handler ownership reports remain separate because the synthetic load intentionally uses unquantized, single-rank input. +All three retained DeepSeek-V4 loader reports pass. Repository collection is now 1,111 items and the static inventory is +1,102 definitions with 903 curated decisions across 335 Python test files. The audit still contains 25 legitimate +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-sixtieth wave: a fake injector cannot certify construction + +The one-hundred-and-sixtieth wave removes a GLM-5.2 block-FP8 QLoRA builder report that replaced both foundation-model +construction and QLoRA injection. Its result was six captured argument values and an inventory object supplied by the +test itself; it never built or adapterized the claimed model. + +The real GLM-5.2 suites retain the meaningful boundary: they construct all 700 targets and 1,700 factors, enforce exact +component and product-mode admission, and verify trainable ownership. The builder's distinct fail-closed rule requiring +both LoRA and QLoRA remains in the consolidated quantized-mode admission report. All five focused builder and real-model +reports pass. Repository collection is now 1,110 items and the static inventory is 1,101 definitions with 904 curated +decisions across 335 Python test files. The audit still contains 25 legitimate conditional runtime gates, one intentional +duplicate group, and no parse errors. + +## One-hundred-and-sixty-first wave: prefer an executed topology over its dictionary + +The one-hundred-and-sixty-first wave removes an OLMo-2 report that inspected tensor-parallel plan dictionary entries and +missing keys. Retained two-rank CPU reports apply that production plan to a real OLMo-2 model, execute forward and backward +through local-axis QK RMSNorm, rowwise and colwise projections, post-norm residual flow, and the vocab-sharded LM head, and +compare the custom local-axis norm with an independent numerical reference. + +The GLM sparse-MLA suite no longer compares CPU `auto` dispatch directly with the same torch reference implementation. Its +full-model sparse-versus-dense report already reaches `auto` through `Glm5Model` and checks numerical parity. The distinct +unknown-backend rejection is preserved in that integration report. + +Finally, the standalone LM-head topology matrix is removed. It launched the same CP, DP, and HSDP layouts used by the +retained four-rank FSDP end-to-end cases, but asserted only group membership and mesh labels. The retained cases build a +real sharded LM head and compare parameter synchronization, vocab ranges, global loss, full weight gradients, and local +hidden gradients with eager references. The separate EP-overlay topology report remains because no equivalent EP execution +gate exists. + +All eight focused model, OLMo-2 distributed, and LM-head distributed reports pass. Repository collection is now 1,107 items +and the static inventory is 1,098 definitions with 907 curated decisions across 334 Python test files. The audit still +contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-sixty-second wave: parse checkpoint paths through the API + +The one-hundred-and-sixty-second wave removes a checkpoint URI report that called `_to_xorl_uri` and `_from_xorl_uri` +directly. The adjacent save-and-load lifecycle already creates real checkpoint directories and drives all five documented +spellings through the public API: xorl URI, explicit `weights/model/checkpoint`, `model/checkpoint`, checkpoint-only, and +legacy `weights/checkpoint`. The direct report also treated an undocumented arbitrary raw path as a compatibility contract. + +The retained lifecycle now pins the exact public xorl URI returned by save, instead of merely checking that its model and +checkpoint substrings appear. All seven checkpoint-path lifecycle reports pass. Repository collection is now 1,106 items +and the static inventory is 1,097 definitions with 908 curated decisions across 334 Python test files. The audit still +contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-sixty-third wave: exercise collator primitives through collators + +The one-hundred-and-sixty-third wave removes direct reports for FlashAttention metadata construction and the private +sequence-shard slicing and padding primitives. Those reports supplied synthetic tensors directly to implementation +helpers, including a zero-length padding no-op, even though the live packing and sequence-shard collators own every +production call. + +The retained collator reports are now stronger consumer-level oracles. Packing checks exact cumulative sequence lengths +and maximum lengths after both multi-document and single-document concatenation. Sequence sharding checks exact rank-zero +and rank-one token and label slices, then drives a nondivisible sequence through the last CP rank to verify constant token +padding, ignored-label padding, and sequential position padding together. All six collator reports pass. Repository +collection is now 1,104 items and the static inventory is 1,095 definitions with 910 curated decisions across 334 Python +test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no +parse errors. + +## One-hundred-and-sixty-fourth wave: cross the real boundary + +The one-hundred-and-sixty-fourth wave deletes a standalone API-orchestrator MessagePack report. The retained ZMQ +client-engine lifecycle already serializes and deserializes the same request and output types across the production +sockets; it now pins the request payload, sequence id, timestamp, response identity, type, payload, and terminal flag +after that roundtrip. + +The same pass deletes a DSv4 report that called the RoPE CP-slice helper with a fabricated group and arbitrary cache. +The retained compressor report now constructs an undersized real CP compressor and reaches the fail-loud cache-capacity +guard through `forward_raw`, alongside its successful C128 path. All four focused communication and compressor reports +pass. Repository collection is now 1,102 items and the static inventory is 1,093 definitions with 912 curated decisions +across 332 Python test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## One-hundred-and-sixty-fifth wave: report component decisions, not helper branches + +The one-hundred-and-sixty-fifth wave consolidates thirteen reports without dropping a behavioral assertion. Deferred +QLoRA key planning now reports dense, EP16, missing-pair, and cache-residency branches at the loader-policy level. +DeepSeek router admission groups the foundation and training-builder entry points while leaving its TP guard separate. +Teacher-cache selection already covered through the Mooncake consumer is no longer reported again; unique async, bounds, +device, dtype, and layer-slice behavior remains in the cache lifecycle. + +The same component-level rule removes narrow reporting around optimizer snapshots and immediate checkpoint reload, +request-processor registration and invalid-target branches, create-model normalized registration, folded-LoRA gradient +dtype, checkpoint model-key inventory, and grouped-load expert-name classification. Those assertions now run in the +transaction, resume, processor, endpoint, autograd, checkpoint-compatibility, and grouped-load lifecycles that consume +them. The consolidation also exposed and fixed order-sensitive monkeypatch state between checkpoint helpers. + +The audit explicitly retained three dense suites after semantic review: QARL's fifteen reports cover numerical export, +STE, calibration, optimizer and checkpoint behavior; FP8 training's twenty cover configuration, injection, correction, +profiling, grouped kernels, and optimizer boundaries; and P2P weight sync's thirty-four cover handshake, slicing, +placement, failure propagation, and teardown. Their size reflects distinct contracts rather than branch-level reporting. + +All 56 changed focused reports pass. Repository collection is now 1,089 items and the static inventory is 1,080 +definitions with 921 curated decisions across 332 Python test files. The audit still contains 25 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-sixty-sixth wave: join branch reports at the component boundary + +The one-hundred-and-sixty-sixth wave consolidates sixteen reports across data, loss, model, server, optimizer, +checkpoint, and weight-sync areas without removing an assertion. Packing allocation invariants now run with the real +`PackingDataset` lifecycle, and the empty-mask `TokenPartial` case is part of its denominator and composition contract. +Legacy softmax and configuration selection now form one TopK-router policy matrix, while distinct balanced, +sqrt-softplus, and hash modes remain separate. + +Server runtime configuration now owns R3 transport and adapter-gradient bucket roundtrips and rejection. Exact GLM +construction reports one attention admission matrix and one complete-MoE admission matrix instead of separate reports +for dependency flags, EP16, lm-head TP16, sparse MLA, all-to-all, and rank-one alpha-one branches. Successful inventory, +post-EP ownership, and numerical execution reports remain independent. + +The same rule folds path validation into kill-session checkpoint promotion, custom-group rejection into multi-part +optimizer behavior, and declaration authority errors into the adapter-ownership compiler's fail-closed policy. EP +checkpoint dimension restore and drop now form one mesh contract. Empty PP-NCCL transfer is part of its tensor-roundtrip +protocol, while sparse-delta priming and receiver-failure retry run in the baseline state-machine lifecycle; post-packed +transfer and initialization stay separate. + +All 36 resulting focused reports pass. Repository collection is now 1,073 items and the static inventory is 1,064 +definitions with 931 curated decisions across 332 Python test files. The audit still contains 25 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-sixty-seventh wave: group policy matrices around one production entry point + +The one-hundred-and-sixty-seventh wave consolidates eight branch-level reports without removing an assertion. Packing +and sequence-shard collators now report generic fields, token-aligned side fields, pre-shifted labels, packed boundaries, +padding, and FlashAttention metadata at their respective component boundaries. The large CP16 side-channel contract +remains separate because it exercises a materially different topology. + +R3 reference validation now runs with put, sliced-load, and cleanup ownership, while the low-level Mooncake tensor codec +remains independent. Sync-quantization acceptance and rejection are one configuration policy matrix. NCCL store-bind +failure is part of rendezvous initialization and port lifecycle, and P2P status timeout is part of asynchronous size and +cutoff dispatch; the distinct prepare-request timeout stays separate. Finally, cross-attention cu-seqlens rejection now +runs with the page-size-one SGL KV-cache adapter instead of appearing as a standalone attention behavior. + +The focused selection collects 15 reports: all 13 runnable reports pass and two FlashAttention-dependent reports retain +their existing environment skips. Repository collection is now 1,065 items and the static inventory is 1,056 definitions +with 937 curated decisions across 332 Python test files. The audit still contains 25 legitimate conditional runtime gates, +one intentional duplicate group, and no parse errors. + +## One-hundred-and-sixty-eighth wave: make lifecycles the reporting boundary + +The one-hundred-and-sixty-eighth wave consolidates nine reports around seven production lifecycles without dropping an +assertion. Rank-zero ready handling now reports acknowledgements, early requests, client identity, unexpected messages, +and receive failure together. Runner load-state admission and artifact-root confinement run with multi-adapter and +single-tenant routing. FIFO admission, dispatch, capacity, terminal transitions, statistics, clearing, and bounded history +now form one scheduler report, using fresh instances to keep scenarios isolated. + +ModelRunner's multi-adapter Adam override is now part of the same full, partial, omitted, and non-Adam optimizer-step +policy. ParallelState defaults and validation run with singleton initialization, automatic DP sharding, access, and +reinitialization protection; EP mesh helpers remain separate. DeepEP internode preflight now reports its skip gates, +transport diagnostics, identity roundtrip, and corruption detection as one outcome matrix, while topology discovery and +buffer sizing remain independent. + +Finally, generic ParallelPlan indivisibility is the rejection branch of successful meta slicing, and exact-GLM malformed +singleton dispositions are part of the exact meta EP plan. The materialized real-tensor shard stays separate because it +executes DTensor redistribution. All 15 resulting focused reports pass. Repository collection is now 1,056 items and the +static inventory is 1,047 definitions with 944 curated decisions across 332 Python test files. The audit still contains +25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-sixty-ninth wave: report codecs and pair buffers as transactions + +The one-hundred-and-sixty-ninth wave consolidates nine static reports across checkpoint loading, quantization, and +shared-prefix execution without removing an assertion. Exact absorbed-KV-B checkpoint loading now treats arrival order, +duplicate members, incomplete pairs, dtype mismatch, and shape mismatch as outcomes of one native-FP8 pair-buffer +transaction. The exact-attention source inventory stays separate because it validates construction rather than loader +state. + +Block-FP8 quantization and dequantization now form one CUDA codec contract covering geometry, scales, input admission, +roundtrip error, determinism, storage, magnitude edges, signs, and dimensional consistency. The GKN codec similarly +reports output layout, aligned and tail blocks, zero blocks, large matrices, output dtype, contiguity, and rank admission +together instead of splitting quantize and dequantize reports. + +For shared-prefix attention, singleton membership and a one-token prompt are one edge-layout report; the general dtype, +head-size, GQA, forward, and backward matrix remains separate. The CPU repacker now includes the one-token empty-shared +block in its full detection, repack, and remap lifecycle. All five runnable resulting reports pass, while the FA3-only +module retains its existing collection skip. Repository collection is now 1,048 items and the static inventory is 1,038 +definitions with 949 curated decisions across 332 Python test files. The audit still contains 25 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-seventieth wave: keep admission with the component it admits + +The one-hundred-and-seventieth wave consolidates seven reports without dropping an assertion. Exact dense gate-up +checkpoint loading now reports out-of-order successful emission together with missing, duplicate, invalid-dtype, and +non-finite members as one pair-buffer transaction. The fused native bytes, scale order, model installation, and loaded +state remain checked. + +NVFP4 fake quantization now has one two-dimensional contract for independent reference parity, STE behavior, and input +admission, plus one three-dimensional expert contract for projection STE, expert-isolated scaling, and fused gate-up +per-half scales. DSV4 tensor-parallel rejection now runs with attention storage and backend-call dtype behavior, while +window-only and C128 forward-backward variants remain parameterized separately. + +Eager-versus-native MoE determinism and the all-tokens-to-one-expert edge now run with the same forward and backward +parity matrix. Non-gated MoE constructor rejection similarly belongs to the CPU eager-reference contract; its optional +GPU Triton and native comparisons remain separate. All 10 resulting focused reports pass. Repository collection is now +1,041 items and the static inventory is 1,031 definitions with 954 curated decisions across 332 Python test files. The +audit still contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-seventy-first wave: keep configuration with runtime selection + +The one-hundred-and-seventy-first wave consolidates four CPU reporting boundaries without removing an assertion. +Gradient-checkpoint method defaults and overrides now run with the layer's training, enabled, and full-recompute gate +matrix. Ordinary and MoE method propagation remain covered alongside the exact checkpoint-call truth table. + +Kimi wrapper conversion, official auxiliary defaults, DeepSeek-V3 registry resolution, and local text-config unwrapping +now form one configuration-loading lifecycle. The tokenizer auto-loader similarly reports the dedicated local TikToken +path together with generic tokenizer and processor fallback, retaining token IDs, text roundtrip, right padding, and the +rule that fallback loaders never gain implicit remote-code trust. + +Finally, sqrt-softplus scaling and requested dtype now run with unchanged softmax gather and renormalization as one +`MoEBlock._regather_routing` mode matrix. All four resulting focused reports pass. Repository collection is now 1,037 +items and the static inventory is 1,027 definitions with 958 curated decisions across 332 Python test files. The audit +still contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-seventy-second wave: keep zero and masked branches with the operation + +The one-hundred-and-seventy-second wave consolidates seven reports without removing an assertion. FlashMLA's all-invalid +case now runs with valid-row compaction, TileLang backward, and zero-scatter behavior as one autograd transaction. The +separate dispatch-envelope report remains because it validates device and production geometry admission. + +Causal-LM Z-loss now reports positive coefficient, zero coefficient, and tensor-parallel rejection in one CPU policy, +retaining reference CE and Z-loss values, finite gradients, absent zero-coefficient metrics, and failure before any +collective. Compiled CUDA parity remains separate. + +Streaming forward-KL now reports dense-reference gradients, chunk invariance, ignore-index loss and gradient masking, +and low-memory parity as one kernel contract. OPD backend parity and unsupported logprob clamping form one dispatch +contract, while the independent FP64 gradcheck remains its own numerical oracle. All eight resulting focused reports +pass. Repository collection is now 1,030 items and the static inventory is 1,020 definitions with 961 curated decisions +across 332 Python test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## One-hundred-and-seventy-third wave: report loss objectives, not plumbing branches + +The one-hundred-and-seventy-third wave consolidates nine reports without removing an assertion. Per-token CE now reports +local and TP LM-head module selection, FP32 bypass, and temperature as one policy matrix. Importance sampling and +causal-LM loss wrappers now share one LM-head dispatch contract covering module use, FP32 bypass, TP collectives, and +finite hidden and head gradients. + +Dr.GRPO's zero-advantage, all-ignored, empty-sequence, positive-advantage, missing-reference, and KL-penalty branches now +run with its forward, backward, and metric contract. Temperature-driven behavior K3 and microbatch composition remain +separate numerical reports. + +Fused selected-logprob parity now includes frozen-output input gradients and irregular tails. Per-token CE, causal-LM, +quack-linear, and importance-sampling selection form one dispatcher integration report, while production-vocabulary +finiteness and the no-full-logits memory bound remain independent heavy regressions. All nine resulting focused reports +pass. Repository collection is now 1,021 items and the static inventory is 1,011 definitions with 964 curated decisions +across 332 Python test files. The audit still contains 25 legitimate conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## One-hundred-and-seventy-fourth wave: report invariant matrices at the kernel boundary + +The one-hundred-and-seventy-fourth wave consolidates eight reports without removing an assertion. Families-v2 RMSNorm +now reports FP64 tree agreement, batch-composition invariance, and repeat determinism as one numerical contract. Its +fused-versus-split realization and runtime dispatch remain separate because they exercise different implementations. + +The BI fused LM-head contract now keeps forward and backward parity together with determinism, batch invariance, and +input guards. Unit-temperature identity and the near-one probability clamp form one edge-policy report. TileLang indexer +causal masking now covers both large finite values and zero inputs in the same edge matrix, while the parameterized +forward geometries remain independent kernel-shape reports. + +Finally, NF4 codebook exactness runs with the flat quantize-dequantize codec transaction; the GKN layout remains separate +because it has a different storage boundary. All 12 resulting focused reports pass. Repository collection is now 1,013 +items and the static inventory is 1,003 definitions with 968 curated decisions across 332 Python test files. The audit +still contains 25 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-seventy-fifth wave: collapse branches at their production boundary + +The one-hundred-and-seventy-fifth wave consolidates 17 reports without removing an assertion. OPD diagnostics, reward +weighting, clamps, stable metric keys, and top-k rejection now form one full-vocabulary policy. GDN convolution forward +and backward parity form one numerical contract. Exact TP1 QLoRA now reports configuration and runtime admission, +forward and surrogate-backward reference parity, and backward safety as three transactions rather than six fragments. + +FP8 linear padding and correction modes now share one matmul numerical matrix. TileLang sparse-MLA reports attention-sink +parity and effect together, and checks partially invalid indices across forward and backward in one masking transaction. +Exact routed experts similarly report global and owner-local factor banks as one sampler-buffer policy, with zero-token +and all-sentinel gradients in one routed edge policy. Canonical GLM52 MoE configuration and runtime selection now share +one mode report. + +Exact shared-expert construction and runtime admission now form one component gate, while logical FP32 masters and +immutable checkpoint binding form one persistent-state policy. Optional SGLang factor-view parity remains separate from +native base views so missing SGLang cannot suppress the native report. Finally, generic nesting and the exact shared +expert now exercise the topmost mixed-precision FSDP selector together; reduce dtype, sequence-parallel folding, and +prefetch remain separate policies. + +All 43 runnable resulting reports pass, with eight existing optional or platform skips. Repository collection is now 996 +items and the static inventory is 986 definitions with 977 curated decisions across 332 Python test files. The audit now +contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-seventy-sixth wave: report policy matrices instead of individual fields + +The one-hundred-and-seventy-sixth wave consolidates 11 reports without removing an assertion. Model-runner token +diagnostics now report selection, empty and top-k boundaries, loss-logprob cross-checks, raw-weight references, and +hidden-state summaries as one callable contract. KL diagnostics, tensor persistence, component hooks, and trusted-input +resolution remain separate because they cross distinct implementation boundaries. + +Stochastic BF16 rounding now reports API admission, seeded repeatability, expectation, and neighboring-value bounds in +one numerical contract. Constant, linear, and cosine learning-rate schedules form one builder mode matrix, while invalid +configuration remains an independent admission report. DistSignSGD local hooks, FSDP-managed exclusion, and unsupported +parallel-topology rejection similarly form one configuration transaction; reduce-scatter arithmetic and optimizer +construction remain separate. + +The session API now reports ordinary LoRA teardown, checkpoint URI return, re-registration, and default-session kill and +unload protection as one termination lifecycle. Inference endpoint registration now covers default and explicit worker +ports, adapter routing, auto-sync, discovered topology, and FP8 KV-cache admission together. Weight-sync quantization +admission and FP8 KV-cache invalidation form one sync request policy, while endpoint listing and receiver enrichment stay +separate. + +All 19 resulting focused reports pass. Repository collection is now 985 items and the static inventory is 975 definitions +with 983 curated decisions across 332 Python test files. The audit still contains 24 legitimate conditional runtime gates, +one intentional duplicate group, and no parse errors. + +## One-hundred-and-seventy-seventh wave: report the P2P protocol by state machine + +The one-hundred-and-seventy-seventh wave consolidates 15 reports without removing an assertion. P2P FP8 transfer now +has one receiver-layout matrix covering fused and unfused attention, partial blocks, Qwen3.6 QKVZ and full-attention +layouts, nonexpert namespaces, mixed FP8 and passthrough entries, shared experts, and routed experts. Every source slice, +receiver byte, scale tensor, dequantized value, endpoint, and expert-coverage assertion remains. + +Multi-sender initialization now reports rank-zero filtered scatter, nonzero-rank adoption, explicit sender process groups, +peer-failure propagation, and optional engine prewarm ordering as one state-machine contract. Locator copy modes, dense +sharding, and rank-filtered transfer stay separate because they determine data partitioning rather than initialization. + +Transfer manifest rejection and compatible-name resolution now form one receiver-manifest policy. Small CPU pooling, +GPU-direct persistent registration and chunking, and aligned mixed-dtype scratch views form one source-staging policy; +receiver-handle coalescing and failure diagnostics remain independent scheduling and observability reports. Finally, +pending failures, optional receiver completion, cleanup draining, endpoint results, deregistration, and completion errors +now form one P2P teardown lifecycle. + +All 19 reports in the complete P2P protocol module pass. Repository collection is now 970 items and the static inventory +is 960 definitions with 988 curated decisions across 332 Python test files. The audit still contains 24 legitimate +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-seventy-eighth wave: keep topology matrices and lifecycle outcomes together + +The one-hundred-and-seventy-eighth wave consolidates 11 reports without removing an assertion. The lm-head TP FSDP +end-to-end suite now has one distributed matrix for CP-replica DP1, CP plus DP2, no-CP DP, no-CP HSDP, and the matching +OPD DP and HSDP modes. Each of the six four-process programs still runs through the same embedded eager loss and gradient +oracle, and failures retain a case identifier. + +Adapter gradient epochs now report empty-step rejection, idempotent abort, scratch reset, publication state, and poisoned +or pending abort rejection as one pre-mutation lifecycle. A successful authoritative optimizer step now includes its +analytical clipping, AdamW parameter and moment updates, scratch reuse, global-step commit, and exact single logical-norm +collective. Semantic rejection, partial optimizer failure, and collective failure form one outcome policy retaining each +branch's recoverability, poison, mutation, and publication assertions. + +Adapter checkpoint save policy now keeps trusted-root confinement with strict target-manifest persistence and mismatch +validation. Authoritative restore now reports lifecycle reset, fingerprint restoration, compatible plan replacement, +direct topology mismatch rejection, and atomic nonmutation together. Coordinator materialization, general session +compatibility, and checkpoint structure stay separate because they cross different orchestration or input boundaries. + +The distributed topology matrix and all 17 reports in the complete adapter-manager module pass. Repository collection is +now 959 items and the static inventory is 949 definitions with 994 curated decisions across 332 Python test files. The +audit still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-seventy-ninth wave: keep resume and backend matrices intact + +The one-hundred-and-seventy-ninth wave consolidates 10 reports without removing an assertion. Adapter optimizer shard +emission and manifest identity rejection now form one save contract. Bitwise uninterrupted resume, immediate moment +restore, weights-only divergence, evicted public reload, scheduled learning-rate restoration, and explicit LR override +form one resume lifecycle. Legacy or incomplete artifact admission now also proves that a failed restore leaves an +already-resident adapter unchanged. + +Grouped FP8 same-NK forward and same-MN weight-gradient kernels now form one training arithmetic report. Block-loop and +Triton references, empty groups, irregular tails, nondefault block sizes, precomputed sequence offsets, and scalar-Quack +dispatch all remain covered. SGLang fused-expert EP activation now reports DeepEP and FP8 exclusion, missing runtime, +flag-off stock routing, score dtype, empty ranks, and compute guards together; happy-path compute, slot combination, +autograd ownership, and weight presentation stay separate. + +Finally, the model-runner expert-factor compiler now has one matrix for eager and fused unquantized backends, registered +session-rank specialization, block-FP8 DeepEP, NF4, NVFP4, and generic quantized contracts. Producer families, +quantization guards, metadata mismatch, factor-shape drift, and uncertified parallelism remain asserted for every branch. + +All 19 runnable focused reports pass, with one existing optional-platform skip. Repository collection is now 949 items +and the static inventory is 939 definitions with 998 curated decisions across 332 Python test files. The audit still +contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-eightieth wave: keep calibration and loading workflows whole + +The one-hundred-and-eightieth wave consolidates six reports without removing an assertion. Qwen-235B simulator Markdown +ingestion, leave-one-out calibration evaluation, exact and extrapolated gradient-accumulation scenarios, observed OOM +boundaries, topology what-if cases, and automatic topology sweeps now form one calibration workflow. Built-in pack replay +also runs the consolidated validator across every shipped pack; portable analytical ledgers, path security, and kernel +correctness-gated ranking remain separate simulator boundaries. + +Grouped checkpoint loading now reports dense and expert routing, fused and FFN source formats, local dense-group +fallback, missing EP-group fallback, strict fallback rejection, and persistent-buffer filtering together. State-dict +resolution, transport, DTensor materialization, and strict post-processing stay separate because they exercise different +callables. + +FP8 weight-sync projection inclusion, module exclusions, receiver skip lists, broad selector mode, stacked quantization, +and already-FP8 passthrough now form one input and layout policy. CPU expert projection now includes zero padding, +deferred formatting, exclusions, reusable workspace staging, and workspace quantization in one expert pipeline. + +All 21 reports in the three complete focused suites pass. Repository collection is now 943 items and the static inventory +is 933 definitions with 1,001 curated decisions across 332 Python test files. The audit still contains 24 legitimate +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-eighty-first wave: report component and side-payload lifecycles + +The one-hundred-and-eighty-first wave consolidates nine reports without removing an assertion. Weight-sync quantization +rejection and FP8 BF16-island enrichment now form one `handle_sync_inference_weights` admission policy. Request-processor +Mooncake externalization, datum-order preservation, and normal, exceptional, and default cleanup now form one R3 +side-payload lifecycle; NCCL synchronization, dispatcher forwarding, optimizer and checkpoint operations, and token +unpacking remain separate request boundaries. + +GLM5 indexer geometry, exact FP32 projection, sentinel and padding masking, and sorted and blocked selection now form one +component contract. GLM52 logical selection, Hadamard transport, fused projection, sampler key preparation, portable +codecs, runtime dispatch, and dependency loading similarly form one sparse-selector pipeline. Production-shape SGLang +CUDA codec parity is now an independent optional report, so its runtime skip cannot mask the portable assertions. + +The four complete focused suites resolve 35 reports as passing with one existing optional-platform skip. Repository +collection is now 934 items and the static inventory is 924 definitions with 1,005 curated decisions across 332 Python +test files. The audit still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no +parse errors. + +## One-hundred-and-eighty-second wave: qualify complete optimizer and model programs + +The one-hundred-and-eighty-second wave consolidates nine reports without removing an assertion. Muon's builder now keeps +Gram-Newton-Schulz configuration, invalid grouped-byte admission, and state-free SGD fallback execution together. Fused +gate-up detection, FSDP parameter replacement, gated and non-gated model classification, and a real Nemotron-H optimizer +step now form one parameter-ownership policy; backend dispatch and matrix-group arithmetic remain separate reports. + +Distributed checkpoint metadata admission and synchronous, no-dist, custom-group, and asynchronous load and save routing +now form one I/O policy. Optimizer metadata-key selection moved to the optimizer-state filtering report, while model-key +and pipeline-LoRA compatibility remain independent schema contracts. + +Canonical GLM52 numerical resolution and official geometry now form one exact model program. Qwen3.5 numerical, MoE, +topology, and model-scope admission similarly form one exact training-program contract; the family-independent RoPE +selector remains separate. Finally, quantized-export YAML and CLI precedence, size parsing, module invocation, BF16 +islands, output configuration, and sharded indexing now report one end-to-end command workflow. + +All 20 reports in the four complete focused suites pass. Repository collection is now 925 items and the static inventory +is 915 definitions with 1,010 curated decisions across 332 Python test files. The audit still contains 24 legitimate +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-eighty-third wave: report API and payload lifecycles end to end + +The one-hundred-and-eighty-third wave consolidates eight reports without removing an assertion. Create-model session +normalization, recreation admission, registration rollback, reserved-checkpoint initialization, and full-weight admission +now form one endpoint lifecycle. Direct adapter loading now keeps success, trusted-path admission, synchronized failure, +pipeline rejection, and auto-registration rollback together. Rank-zero broadcast routing, sharded restore, session-spec +mismatch, and transactional optimizer rejection similarly form one load-mode policy. + +Packing now reports empty, missing, oversized, NumPy, valid, and malformed input handling with output validation. Its +per-token unpack modes run with the full pack, metadata, simulated-forward, and sample-boundary round trip. Teacher-cache +CP and EP contributor selection, legacy duplicate mode, sequence-parallel trimming, and cross-rank writer gathering now +form one distributed producer policy; Mooncake storage remains a separate transport boundary. + +Finally, OPD teacher causal shifting, cache-index alignment and rejection, and Mooncake metadata admission form one +pipeline payload contract. Endpoint reuse, student-version verification, and preparation-worker queueing remain separate +orchestration reports. + +All 25 reports in the five complete focused suites pass. Repository collection is now 917 items and the static inventory +is 907 definitions with 1,015 curated decisions across 332 Python test files. The audit still contains 24 legitimate +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-eighty-fourth wave: keep checkpoint and trust boundaries whole + +The one-hundred-and-eighty-fourth wave consolidates eight reports without removing an assertion. Runtime-rank LoRA +export, PEFT hybrid-shared layout, SGLang shared-outer layout, and adapter-manager loading for hybrid-shared and all-owner +experts now form one unquantized checkpoint workflow. Low-level EP slicing and the quantized projection-subset round trip +remain separate representation boundaries. + +Successful optimizer publication, commit and handler-tail poisoning, and rank-zero fatal termination now form one +post-mutation lifecycle. Empty logical packing, discovery, replica classification, and ownership compilation form one +empty-shard layout policy. Coordinate, replica, LoRA-B, session, and FQN-order invariance similarly form one deterministic +adapter initialization contract; real Gloo and explicit EP composition remain separate topology gates. + +Server artifact-root confinement, symlink rejection, and private diagnostic-input admission now form one filesystem trust +boundary. Compile-target allowlisting, safe IPC type round trips, and oversized-frame rejection form one compile-worker +trust boundary, while outbound endpoint validation remains independent. + +All 16 reports in the four complete focused suites pass, including the real two-rank Gloo layout check. Repository +collection is now 909 items and the static inventory is 899 definitions with 1,019 curated decisions across 332 Python +test files. The audit still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no +parse errors. + +## One-hundred-and-eighty-fifth wave: report versioned numerical and manager contracts + +The one-hundred-and-eighty-fifth wave consolidates 11 reports without removing an assertion. Adapter-manager optimizer +construction, hyperparameter persistence, learning-rate updates, and mixed-rank multi-optimizer reload now form one +configuration lifecycle. Session compatibility, weights-only behavior, checkpoint structure, PEFT suffixes, sharded +indices, and rank-capacity admission form one load compatibility policy. + +FSDP topmost protected-module selection, expert policy stripping, mesh-dependent reduction dtype, explicit overrides, +and dtype admission now form one mixed-precision contract. MiniMax M3 clamped SwiGLU, biased sigmoid routing, text +forward and backward, multimodal-token rejection, and parallel-mode admission similarly form one runtime program. + +All v1 BI family, normalization, mean, softmax, matrix, and LM-head frozen hashes now report as one versioned golden-tree +gate; v2 normalization and head hashes form another. FlashQLA packed-versus-individual rows, total-token invariance, and +block-DV tile invariance now form one Gate 2 contract. Auto-CP and Gate 4 state handoff remain separate decisions. Six +inert capability decorators were also removed from helper functions; capability admission remains on the collected +versioned parent reports. + +All 28 reports in the five complete focused suites pass on the available GPU. Repository collection is now 898 items and +the static inventory is 888 definitions with 1,024 curated decisions across 332 Python test files. The audit still +contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-eighty-sixth wave: keep numerical oracles and transport codecs together + +The one-hundred-and-eighty-sixth wave consolidates ten report definitions without removing an assertion. QK, pre-summed +residual-tree, post-attention residual, zero-centered family-1, and families-v2 RMSNorm site classes now form one +cross-engine bitwise matrix. The module-level SGLang dependency gate remains unchanged. Qwen3.5 Class-B dispatch and +fail-closed admission form one rotary policy, while dense and MoE half-rotate behavior and post-RoPE BF16 casting form one +attention projection policy. + +Mooncake byte codecs, canonical dtype strings, metadata emission, suffixed keys, and rank-2 and rank-3 hidden fetches now +form one transport contract. Fused GDN canonical LoRA folding, slice-local gradients, exact projection, cache reuse, +bounded generations, and old-generation release form one merged-weight lifecycle. + +Exact fused gated-RMSNorm routing, unsupported residual rejection, and full GatedDeltaNet routing now form one model +program dispatch contract. BI router GEMM FP32 agreement, empty and dtype admission, and hidden and weight gradients form +one numerical contract. Their lower-level gating, normalization, solve, top-k, and model-integration boundaries remain +separate. + +The six runnable consolidated reports pass; all 19 reports in the runnable affected modules pass, with the cross-engine +module producing one expected dependency skip. Repository collection is now 892 items and the static inventory is 878 +definitions with 1,030 curated decisions across 332 Python test files. The audit still contains 24 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-eighty-seventh wave: report component lifecycles instead of implementation fragments + +The one-hundred-and-eighty-seventh wave consolidates nine report definitions without removing an assertion. RMSNorm family +name admission, residual-shape rejection, and Qwen site declarations now form one structural API policy; funnel equivalence, +family vitality, zero-centered folding, and module dispatch form one numerical-routing contract. Undeclared-family +enforcement remains a separate CUDA tripwire. SGLang-fused residual and no-residual forward and backward comparisons now +form one fused numerical contract, while CPU fallback, model integration, and trunk dispatch remain separate boundaries. + +DeepSeek V4 now reports shared-MLP and routed-expert SwiGLU-limit propagation together. C128 execution, causal-LM forward +and backward, hash-layer input threading, and decoder gradient-checkpoint wrapping form one model runtime contract; +construction, topology admission, and precision preservation remain separate. Nemotron H strict loading, ignored MTP +admission, HF parity, and exact save reconstruction now form one published-layout codec transaction with the supported save +key set kept explicit. + +Qwen3 per-expert and Qwen3.5 stacked fused-expert layouts now share one checkpoint-handler policy, including deferred QLoRA +expert loading. QLoRA quantized storage, forward and backward, dequantization, prequantized NVFP4 loading, EMA scale +convention, and merge-requantization now form one quantized-weight lifecycle; injection, block-FP8 representation, and +optimizer reset remain distinct. + +All 24 reports in the seven affected modules pass. Repository collection is now 883 items and the static inventory is 869 +definitions with 1,037 curated decisions across 332 Python test files. The audit still contains 24 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-eighty-eighth wave: join inverse codec paths and configured execution + +The one-hundred-and-eighty-eighth wave consolidates eight report definitions without removing an assertion. Copying full +checkpoint tensors into existing replicated or sharded DTensors and materializing multi-axis DTensors for a selected writer +now form one load/save tensor codec. Object payload transport, NCCL device selection, weight-load group routing, and +handler-filtered rank-zero loading similarly form one broadcast loading transaction. State-dict discovery and grouped expert +loading remain independent resolution boundaries. + +Muon builder keyword propagation, byte-limit admission, SGD fallback, a real Gram-Newton-Schulz update, and restart +autotuning now form one configured optimizer lifecycle. Quack selection, grouped scheduling, standard Newton-Schulz, CUDA +compute dtype, and model classification remain separate algorithm or platform branches. BI trunk-linear persistent-GEMM +forward bits, batch invariance, dtype admission, and cuBLAS input, weight, and bias gradients now form one wrapped-linear +numerical contract; wrapper selection and global-interpose admission remain separate. + +DeepSeek V3 external per-expert and internal fused checkpoint layouts, dense and packed EP slicing, requested device and +dtype, and quantization-config discovery now form one expert codec policy. Nemotron H router output, loss backward through +Mamba, attention, MoE, and shared experts, and full-layer gradient checkpointing now form one training runtime contract; +packed variable-length equivalence remains its own state-propagation boundary. RoPE registry-wide FP32 frequency +construction, BF16 consumption, and unchanged exact-lane cosine and sine bits now form one precision policy, while lazy +cache growth and architecture-specific CUDA placement remain separate. + +All 20 reports in the six affected modules pass, including the four-rank Gloo materialization and CUDA BI gradient checks. +Repository collection is now 875 items and the static inventory is 861 definitions with 1,043 curated decisions across 332 +Python test files. The audit still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no +parse errors. + +## One-hundred-and-eighty-ninth wave: make state transitions own their reports + +The one-hundred-and-eighty-ninth wave consolidates eight report definitions without removing an assertion. DistSignSGD +FSDP2 builder admission, decay grouping, hook configuration, and a preaggregated update now form one configured optimizer +contract; reduce-scatter sign timing and local-versus-FSDP gradient ownership remain independent communication boundaries. +SignSGD builder construction, decay grouping, dense updates, decoupled decay, and sparse-gradient rejection similarly form +one local optimizer contract. + +Trainer sequence-parallel sums, DTensor skipping, adapter-finalization exclusions, LM-head replica gradient handling, and +marked-parameter broadcast now form one explicit synchronization policy. Dispatcher forward-backward rendezvous and commit, +explicit gradient-epoch abort, uniform rejection, and rank-asymmetric failure conversion form one gradient-epoch lifecycle; +session, save, and post-optimizer publication operations remain separate mutation boundaries. + +Dense residual, attention, normalization, and MLP hooks plus routed-MoE callbacks and shared-expert components now report as +one hidden-component capture pipeline. Summary formatting, ranked tensor dumps, and trusted diagnostic overrides remain +separate output and filesystem boundaries. DeepSeek V3 auxiliary router-logit emission and replay recording of selected +indices and weights similarly form one router observability contract, while full-model backward, router freezing, and LoRA +injection remain independent. + +Local and pipeline checkpoint key discovery, metadata unions, QARL buffer mismatch, base-to-LoRA loading, and LoRA-only +loading now form one model-compatibility policy; distributed transport and optimizer payload filtering remain separate. +Finally, MoE expert-sorted routing weights and scatter-add reconstruction now form the encode and decode directions of one +memory-efficient token permutation codec, while all-to-all ordering and hidden chunking remain separate transports. + +All 26 reports in the eight affected modules pass. Repository collection is now 867 items and the static inventory is 853 +definitions with 1,051 curated decisions across 332 Python test files. The audit still contains 24 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-ninetieth wave: include class-method suites in semantic ownership + +The one-hundred-and-ninetieth wave expands the report-density audit to class-method suites and consolidates nine report +definitions without removing an assertion. Qwen3.5 families-v2 effective folded-weight gradients and the residual twin's +output and residual gradient paths now form one zero-centered backward contract. Legacy FP8 config normalization, +incompatible external runtime rejection, and explicit Blackwell validation-artifact admission form one configuration +boundary; BF16 layer-island transformation remains separate. + +Quack unique PTX outputs, bounded ptxas execution, cleanup, and exact entry discovery now form one compilation process- +safety contract, while worker framing and cache hashing remain separate trust boundaries. NVFP4 independent 2D reference +agreement, shape admission, linear STE, 3D expert STE, expert isolation, and fused gate-up scale ownership now form one +fake-quant contract. + +Merged-LoRA canonical linear and expert folds plus straight-through factor gradients now form one low-level numerical +contract. LoraLinear merged selection, gradient parity, optimizer-step invalidation, and active-rank cache invalidation form +one linear lifecycle. MoE canonical merged views, parameter-version caches, and fused-expert admission form one expert +lifecycle; native EP execution and trunk wrapping remain separate integrations. + +MoE-LoRA backend initialization, frozen and trainable ownership, runtime rank slicing, from-module conversion, model +injection, and block injection now form one construction policy. Zero-delta base equivalence and eager-versus-native or +Triton output and gradient agreement form one GPU numerical policy under the same capability gate. CPU eager execution, +zero-token structural gradients, and EP router-score application remain distinct runtime boundaries. + +All 17 reports in the six affected modules pass, including the GPU cross-backend MoE-LoRA checks. Repository collection is +now 858 items and the static inventory is 844 definitions with 1,057 curated decisions across 332 Python test files. The +audit still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-ninety-first wave: join transport transactions and kernel contract states + +The one-hundred-and-ninety-first wave consolidates eight report definitions without removing an assertion. P2P cold and +cached prepare behavior now form one initialization handshake. Successful completion, pending-transfer failure, receiver +notification, deregistration, and destroy cleanup form one terminal synchronization lifecycle; fanout, slicing, coalescing, +diagnostics, and multi-sender routing remain separate transport boundaries. + +NCCL initialization and transfer endpoint ports plus its optional two-phase receiver protocol now form one endpoint +transaction policy. Flat, chunked-flat, and receiver-fenced hybrid buckets form one flattened load-format contract, while +endpoint health and invalid multi-rank direct format remain independent admission boundaries. + +GDN packed convolution weights, armed routing, fail-closed admission, call-scoped contract state, and checkpoint +recomputation now form one exact-contract lifecycle. Low-level CUDA parity, full-block integration, and optional SGLang +tree-kernel parity remain distinct numerical boundaries. Fixed-length and variable-length FlashAttention calls now form one +API behavior contract. SGL, paged FlashAttention, flags-off, and FA4 selection form one page-size-one KV-cache routing +policy; backend resolution and eager head-layout numerics remain separate. + +The four affected modules collect 29 reports: 27 pass and two have expected capability or optional-dependency skips. +Repository collection is now 850 items and the static inventory is 836 definitions with 1,061 curated decisions across 332 +Python test files. The audit still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and +no parse errors. + +## One-hundred-and-ninety-second wave: make stateful APIs own their complete lifecycle + +The one-hundred-and-ninety-second wave consolidates seven report definitions without removing an assertion. Empty and +aborted adapter epochs, a successful clipped update, nonfinite input, optimizer failure, and collective failure now form one +authoritative optimizer lifecycle. Capture ownership, publication admission, exact LM-head coherence, and checkpoint restore +remain separate. Sharded optimizer-manifest emission, identity rejection, moment restoration, and bitwise continuation now +form one checkpoint codec; artifact admission and logical cross-layout resharding remain independent compatibility gates. + +Sampling-session stale-state reconciliation, transient query failure, model-scoped tracking, and failed-load atomicity now +form one adapter lifecycle. Sampler checkpoint storage and adapter-only export remain separate. Inference weight-sync +endpoint forwarding, pool selection, quantization admission, and cache invalidation now form one API transaction, while +endpoint registration, health refresh, and receiver-capability detection remain separate. + +Model-session normalized registration, duplicate and topology admission, reserved checkpoints, kill, optional final save, +and default-session protection now form one create-to-destroy lifecycle. The lightweight create-session alias remains a +separate endpoint. Disk-backed weights information, path admission, full-weight metadata, and legacy SignSGD upgrade now +form one checkpoint session-spec decoding policy. + +All 30 reports in the five affected modules pass. Repository collection is now 843 items and the static inventory is 829 +definitions with 1,067 curated decisions across 332 Python test files. The audit still contains 24 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-ninety-third wave: report complete artifact and distributed contracts + +The one-hundred-and-ninety-third wave consolidates six report definitions without removing an assertion. Fused QKV, MLA, +linear-attention, GKN expert, and fused gate-up transformations now form one exported-model tensor-layout contract. Exporter +CLI and directory behavior, source admission, QARL folding, and the low-level FP8 quantizer remain separate boundaries. + +Expert-adapter backend capability and plan identity plus factor ownership, reduction domains, and checkpoint persistence +now form one structural contract. Exact target-subset forwarding and GLM, Qwen3, and Qwen3.5 wrapper construction form one +injection policy. Supported SiLU preservation and rejection of incompatible activations, biases, quantization groups, +target sets, and model-family semantics form one fail-closed semantic contract; runtime numerical parity remains separate. + +Canonical MoE trainer and sampler plan identity, topology admission, logical ordinals, and world-32 CP, EP, and expert-FSDP +group aliases now form one topology contract. Its 2- and 8-contributor dense transport and 16-contributor packed and +CP-sharded parity now form one distributed numerical contract; every subprocess still runs. + +All 12 reports in the three affected modules pass, including the 2-, 8-, and 16-process transport checks. Repository +collection is now 837 items and the static inventory is 823 definitions with 1,073 curated decisions across 332 Python test +files. The audit still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no parse +errors. + +## One-hundred-and-ninety-fourth wave: join objective branches and dispatch transactions + +The one-hundred-and-ninety-fourth wave consolidates seven report definitions without removing an assertion. Evicted +auto-load, fresh materialization, explicit path load, rollback, and all-rank or rank-zero-broadcast restoration now form one +adapter load lifecycle; registration and save remain separate mutations. Token selection, boundary behavior, raw-weight +cross-checking, hidden summaries, and CP-sharded KL position mapping now form one token-diagnostics policy, while capture, +tensor-dump output, and trusted override input remain separate. + +OPD reference, streaming, low-memory, and sharded-store forward agreement plus backward, partial reduction, and output dtype +now form one numerical backend contract. Fused selected-logprob dtype, bias, temperature, frozen-head, irregular-tail, Qwen, +and GPT-OSS vocabulary cases similarly form one forward and backward numerical contract; loss dispatch and the no-full- +logits memory gate remain separate. + +DRGRPO forward values, gradients, metrics, zero boundaries, advantage direction, KL penalty, and logprob-temperature behavior +now form one objective contract, with microbatch reducer composition kept independent. Packing concatenation, capacity +splitting, mixed lengths, empty and single input, oversize admission, missing fields, NumPy conversion, and microbatch +validation now form one core packing policy; metadata, disabled mode, and full roundtrip remain separate. + +All 19 reports in the six affected modules pass, including the production-vocabulary and no-full-logits GPU gates. +Repository collection is now 830 items and the static inventory is 816 definitions with 1,079 curated decisions across 332 +Python test files. The audit still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and +no parse errors. + +## One-hundred-and-ninety-fifth wave: make tensor preparation and save formats whole + +The one-hundred-and-ninety-fifth wave consolidates six report definitions without removing an assertion. Receiver +postprocess selection, FP8 KV-cache requirements, unsupported-format admission, and generated BF16 islands now form one +quantized weight-sync configuration contract. Adapter materialization, parameter filtering and tied aliases, compile-name +normalization, and architecture-specific unfusing form one sync-source tensor preparation pipeline; bucket sizing, +transport routing, and sparse-delta selection remain separate. + +P2P sender selection, direct-EP collection admission, and gated or nongated local expert projection collection now form one +EP synchronization-source policy, while transport remains independently tested. Factor-only admission, rank-zero artifact +failures, LoRA-only failures, and pre-barrier error surfacing now form one fail-closed checkpoint save policy. Live dense +target resolution and collective stacked-MoE factor slicing form one LoRA checkpoint export contract; optimizer artifacts +remain in the resume suite. + +All seven reports in the two affected modules pass. Repository collection is now 824 items and the static inventory is 810 +definitions with 1,084 curated decisions across 332 Python test files. The audit still contains 24 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-ninety-sixth wave: join configuration and backend branches + +The one-hundred-and-ninety-sixth wave consolidates six report definitions without removing an assertion. Removed fields, +incompatible quantized modes, vLLM-only runtime knobs, broadcast loading, and unsupported multi-adapter modes now form one +fail-closed server configuration boundary. Canonical defaults, nested runtime controls, receiver cache dtype, R3 transport, +gradient buckets, Muon options, runner compatibility, sparse MLA, and MoE routing controls form one runtime configuration +roundtrip; quantized-training and parallel-topology configuration remain separate. + +Teacher-cache contributor selection, CP and DP assembly, valid-label trimming, Mooncake metadata emission, byte roundtrip, +and activation-cache consumption now form one hidden-cache lifecycle. OPD loss execution and debug artifacts remain +separate. Muon builder options, fallback behavior, restart autotuning, grouped shapes, transpose equivalence, fused halves, +and byte-limit chunking now form one configured Gram-Newton-Schulz contract, while Quack selection, standard Newton-Schulz, +and CUDA compute dtype remain separate. + +Block-loop and Triton-grouped forward and weight gradients plus scalar-Quack per-expert scaling now form one grouped FP8 +GEMM numerical contract. DeepGEMM subprocess isolation and model train-step integration remain separate gates. + +The four affected modules collect 21 reports: 20 pass and one DeepGEMM capability gate skips as expected. Repository +collection is now 818 items and the static inventory is 804 definitions with 1,089 curated decisions across 332 Python test +files. The audit still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no parse +errors. + +## One-hundred-and-ninety-seventh wave: make dispatcher input and completion transactions whole + +The one-hundred-and-ninety-seventh wave consolidates four report definitions without removing an assertion. DP, EP, CP, +and legacy rank selection, routing-payload slicing, rank-local row grouping, and source provenance now form one dispatcher +input-distribution policy. Packing strategy and R3 payload storage remain separate upstream boundaries. + +Local payload trimming, rank rendezvous, CP replica deduplication, disagreement rejection, and rank-zero per-token merging +now form one dispatcher completion transaction; diagnostic dumping remains a separate output. Processor readiness, forward +and backward execution, timing propagation, invalid-target rejection, shutdown, model identity, auto-load, and R3 forwarding +now form one request-to-runner compute lifecycle. NCCL sync, optimizer and checkpoint RPCs, and payload storage remain +separate operations. + +All 11 reports in the two affected modules pass. Repository collection is now 814 items and the static inventory is 800 +definitions with 1,092 curated decisions across 332 Python test files. The audit still contains 24 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-ninety-eighth wave: make adapter capture, restore, and residency transactional + +The one-hundred-and-ninety-eighth wave consolidates five report definitions without removing an assertion. Raw-numerator +accumulation, model-gradient clearing, FP32 scratch reuse, staged commit, direct-DTensor preservation, and atomic +prevalidation now form one gradient-capture transaction. Ownership-plan compilation and optimizer mutation remain separate +lifecycle boundaries. + +Scalar-tensor LM-head optimizer-state coherence now runs as a distributed validation branch of the authoritative optimizer +lifecycle. Trusted paths, strict target manifests, lifecycle reset, ownership-plan admission, optimizer compatibility, +learning-rate rules, checkpoint structure, missing tensors, rank capacity, and PEFT filename and shard compatibility now +form one adapter checkpoint restore and admission policy. Coordinator-driven materialization remains independently tested. + +Mixed ranks and optimizers, adapter switching, training, checkpoint reload, capacity eviction, dirty-state protection, +multi-rank rejection, and save-failure rollback now form one multi-adapter lifecycle. + +All seven reports in the affected module pass. Repository collection is now 809 items and the static inventory is 795 +definitions with 1,096 curated decisions across 332 Python test files. The audit still contains 24 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## One-hundred-and-ninety-ninth wave: close model-loading transactions end to end + +The one-hundred-and-ninety-ninth wave consolidates two report definitions without removing an assertion. Local directory +resolution and distributed shard-list broadcast now form the input phase of rank-zero checkpoint loading alongside tensor +and metadata transport, process-group selection, and handler-filtered prefetch. + +Dense and expert routing, supported expert-key formats, fused and FFN source conversion, process-group fallback, and strict +parameter and persistent-buffer coverage now form one grouped checkpoint-loading transaction. Four-process replicated, +sharded, and target-rank DTensor materialization remains a separate save-side correctness gate. + +All three reports in the affected module pass, including both four-process DTensor checks. Repository collection is now +807 items and the static inventory is 793 definitions with 1,098 curated decisions across 332 Python test files. The audit +still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundredth wave: align GLM-5.2 reports with canonical contracts + +The two-hundredth wave consolidates three report definitions without removing an assertion. Certified world-16, EP16, and +CP16 topology, the official producer schedule, the 38/40 pipeline split, malformed-plan rejection, and full-indexer +allocation now form one canonical layer-plan contract. + +Index publication, reuse, concurrency rejection, exception cleanup, and identity preservation across FSDP mixed-precision +input casting now form one index-share lifecycle. Correction-bias FP32 preservation, meta materialization, strict checkpoint +ingestion, routing-replay rejection, internal transport selection, and exact canonical router and indexer dispatch now form +one canonical MoE configuration and selection contract. Native sampler codec parity and end-to-end semantic logprob +composition remain independent gates. + +The affected module collects six reports: five pass and the SGLang-dependent native-codec capability gate skips as +expected. Repository collection is now 804 items and the static inventory is 790 definitions with 1,101 curated decisions +across 332 Python test files. The audit still contains 24 legitimate conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-first wave: report fused-MoE choices as complete policies + +The two-hundred-and-first wave consolidates three report definitions without removing an assertion. Automatic and explicit +SGLang fused-MoE enablement, supported-module admission, one-time logging, flag-off preservation, and block and +experts-only entrypoint selection now form one resolution and dispatch policy. + +Unsupported expert semantics, pre-import clamp rejection, trainable guards, and the gradient-sensitive choice between the +autograd function and plain kernel now form one fused-expert admission and trainable-dispatch contract. Transient, cached, +and zero-copy strided modes, cache reuse and invalidation, serving tensor order, and split gate-up adapter layout now form +one kernel weight-layout contract. Runtime-context admission, gradient numerics, and real-kernel parity remain separate. + +The affected module collects six reports: four pass and two GPU capability gates skip as expected. Repository collection +is now 801 items and the static inventory is 787 definitions with 1,104 curated decisions across 332 Python test files. The +audit still contains 24 legitimate conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-second wave: make simulator accounting and calibration whole + +The two-hundred-and-second wave consolidates three report definitions without removing an assertion. Topology resolution, +balanced routing, sequence-parallel local shapes, FLOPs, activation storage, and communication bytes now form one +analytical accounting contract. + +Config fingerprinting, cached and known-model metadata resolution, calibration-pack containment, built-in prefix +validation, symlink rejection, and restricted local reads now form one input resolution and admission policy. Qwen +markdown ingestion, leave-one-out calibration evaluation, scenario and topology planning, built-in pack replay, +fit-and-OOM feasibility, and consolidated pack validation now form one calibration lifecycle. Generic observed-run +ingestion and correctness-gated kernel ranking remain separate policies. + +All five reports in the affected module pass. Repository collection is now 798 items and the static inventory is 784 +definitions with 1,107 curated decisions across 332 Python test files. The audit still contains 24 legitimate conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-third wave: join distributed transport and FP8 sync branches + +The two-hundred-and-third wave consolidates six report definitions without removing an assertion or subprocess. Direct +output, shared-owner, and all-owner layouts now form one four-GPU FSDP adapter-gradient ownership contract. Eager, Triton, +native, Quack, NF4, NVFP4, block-FP8, projection-subset, and all-owner cases now form one AllToAll contract. Hybrid-shared, +all-owner, and quantized Quack cases now form one optional DeepEP contract; its dependency gate remains explicit. + +SGLang EP top-k-one pair presentation, FP32 routing weights, local weight layout, flag-off behavior, empty-rank handling, +and semantic admission now form one dispatch policy. BF16 islands, block scale and zero-padding semantics, projection +selection, skip lists, stack handling, and existing FP8 values now form one CPU sync-quantization contract. Dense LoRA, +QLoRA, quantized MoE factors, and fused-GDN factors now form one adapter-folding sync-source policy. + +The three affected modules collect 13 reports. Ten pass, including the live FP8 GPU policy, and two optional capability +gates skip. The four-GPU FSDP report is not currently verified because rank 3 fails at `torch.cuda.set_device(3)` with an +out-of-memory error before any test assertion or model operation; a direct rerun reproduced that external admission +failure. Repository collection is now 792 items and the static inventory is 778 definitions with 1,113 curated decisions +across 332 Python test files. The audit still contains 24 legitimate conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-fourth wave: make runner compilation one ownership policy + +The two-hundred-and-fourth wave consolidates three report definitions without removing an assertion. Generic module and +direct-output ownership, exact TP16 LM-head VJP masks, managed shard capture, replica divisors, group-family coverage, and +fail-closed topology admission now form one runner gradient-ownership compiler policy. Expert-factor compilation remains +a separate specialized contract. + +Merged and legacy effective LM-head selection now run as input branches of the direct-output analytical capture and +optimizer step, which checks selected bytes, factor gradients, logical norm, and parameter mutation end to end. Failed +staged-capture rollback remains independent. + +All four reports in the affected module pass. Repository collection is now 789 items and the static inventory is 775 +definitions with 1,115 curated decisions across 332 Python test files. + +## Two-hundred-and-fifth wave: report optimizer state and routed banks as wholes + +The two-hundred-and-fifth wave consolidates four report definitions without removing an assertion. Denominator chunking, +Kahan compensation, gradient reuse, CPU state offload, DTensor local-shard wrapping, and device restoration now form one +AnyPrecision AdamW state-strategy policy. Cautious decay math and optimizer construction remain separate. + +EP16 and MoE-TP1 admission, all 16-by-16 owner-slot remaps, global and owner-local factor banks, sampler buffer shapes and +dtypes, and unused-rank zero padding now form one routed-bank layout policy. Owned zero-base gradients, all-sentinel +structural zeros, input-layout rejection, and top-k-eight mixed-owner VJPs now form one routed-gradient edge policy. The +full 256-slot literal sampler numerical gate remains separate. + +The two affected modules collect eight reports: six pass and two Hopper/SGLang capability gates skip. Repository +collection is now 785 items and the static inventory is 771 definitions with 1,118 curated decisions across 332 Python +test files. The static audit surfaces 22 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-sixth wave: compile ownership as one transaction + +The two-hundred-and-sixth wave consolidates three report definitions without removing an assertion. Generic topology +declarations, authority masks, rank-local tensor and geometry fingerprint invariance, missing and structurally false +declaration rejection, tensor-parallel admission, group identity and membership validation, and complete orthogonal +replica coverage now form one ownership-plan compilation transaction. + +Fullgraph module-producer execution and bucketed residual gradient transport remain separate runtime contracts. All three +reports in the affected module pass. Repository collection is now 782 items and the static inventory is 768 definitions +with 1,119 curated decisions across 332 Python test files. The static audit surfaces 22 conditional runtime gates, one +intentional duplicate group, and no parse errors. + +## Two-hundred-and-seventh wave: join exact-QLoRA construction and runtime policies + +The two-hundred-and-seventh wave consolidates four report definitions without removing an assertion. Shared-expert +construction, runtime admission, logical state, and checkpoint-state policy now form one shared-expert contract. Physical +SGLang views remain a separate optional-dependency report so the CPU structural policy is never hidden by capability +admission. + +Exact TP1 configuration, runtime admission, dtype moves, packed-state master dtype, and parameter identity now form one +configuration and state-lifecycle policy. Forward values, surrogate VJPs, and backward safety now form one numerical and +autograd policy. Request-scoped NCCL group naming now runs inside the orchestrator's optimizer, checkpoint, synchronization, +registration, and lifecycle control report. + +The three affected modules collect 13 reports: nine pass and four optional GPU or SGLang capability gates skip. Repository +collection is now 778 items and the static inventory is 764 definitions with 1,123 curated decisions across 332 Python test +files. + +## Two-hundred-and-eighth wave: make GLM-5 support reports policy-complete + +The two-hundred-and-eighth wave consolidates three report definitions without removing an assertion. GLM-5 indexer +construction and DSA masking now form one indexer-selection policy. Sparse-MLA reference behavior, wrapper semantics, and +attention integration now form one sparse-attention policy. Sparse-KV adapter weights, adapter dispatch, and MoE dispatch +now form one adapter-and-routing policy. + +TileLang CUDA execution, checkpoint filtering, Hugging Face parity, and full forward/recompute behavior remain separate +capability or end-to-end gates. All eight reports in the affected module pass. Repository collection is now 775 items and +the static inventory is 761 definitions with 1,126 curated decisions across 332 Python test files. + +## Two-hundred-and-ninth wave: report native dispatch and optimizer recovery as transactions + +The two-hundred-and-ninth wave consolidates three report definitions without removing an assertion. Entry through the +experts module and its FSDP pre-forward hook now runs inside the native-combine diagnostic report, which also verifies the +actual gathered, routed, gated, local, and combined operands. EP8 admission, variable-row collectives, and fused-gate +gradient parity remain independent contracts. + +Canonical optimizer parameter identity, wrapper-insensitive fingerprints, live binding validation, recursive state +snapshots, and failed-collective commit behavior now form one optimizer transaction policy. Successful sharded save and +bitwise resume, manifest identity checks, legacy-pickle rejection, missing-artifact rejection, and resident-state +preservation now form one checkpoint recovery and artifact-admission policy. Logical resharding remains separate. + +All seven reports in the two affected modules pass. Repository collection is now 772 items and the static inventory is 758 +definitions with 1,129 curated decisions across 332 Python test files. The static audit surfaces 22 conditional runtime +gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-tenth wave: align P2P reports with transfer lifecycle boundaries + +The two-hundred-and-tenth wave consolidates three report definitions without removing an assertion. Per-receiver slice +placement and reuse of one staged source across replicated receiver locators now form one receiver-placement policy. +Typed CPU/GPU staging, registration lifetime, alignment, and receiver-handle-aware coalescing now form one staging policy. +Requested flush-cache and weight-version propagation now runs inside sync completion alongside cache retention, tied-weight +aliases, pending-transfer failure handling, and cleanup. + +FP8 receiver layouts, invalid-manifest admission, transfer diagnostics, direct-EP transport, initialization, and slicing +remain separate contracts. All 14 reports in the affected module pass. Repository collection is now 769 items and the +static inventory is 755 definitions with 1,132 curated decisions across 332 Python test files. + +## Two-hundred-and-eleventh wave: join adapter state and quantized exclusion lifecycles + +The two-hundred-and-eleventh wave consolidates three report definitions without removing an assertion. Sampling-adapter +reconciliation, transient endpoint-query handling, model-scoped atomic tracking, listing, deletion, resolution, and +receiver-removal invalidation now form one sampler adapter-state policy. Sampler-weight export remains separate. + +Zero-token structural gradients for every local LoRA factor now run as the empty-input edge of eager expert forward, +backward, and MoE-block behavior. Cross-backend numerics and routing-score application remain separate. Prequantized +exclude-module metadata parsing, precedence, malformed input, dense and MoE handler skip behavior, and auxiliary-key +passthrough now form one checkpoint exclusion policy. + +All 12 reports in the three affected modules pass. Repository collection is now 766 items and the static inventory is 752 +definitions with 1,135 curated decisions across 332 Python test files. + +## Two-hundred-and-twelfth wave: make exact construction reports fail-closed + +The two-hundred-and-twelfth wave consolidates three report definitions without removing an assertion. Exact GLM-5.2 +attention construction now reports its complete 780-factor canonical inventory and rejects invalid rank, alpha, component, +dispatch, and sparse-MLA configurations in one fail-closed construction policy. + +Exact GLM-5.2 MoE construction now reports its complete shared and routed inventory, source metadata, shapes, ownership, +and invalid dependency, EP16, LM-head-TP16, rank, and alpha branches together. Post-EP layout and selected-logprob LM-head +specialization remain separate. DeepSeek-V4 construction, pipeline-parallel rejection, parallel-group wiring, FP32 marker +propagation, dtype casts, and complex RoPE preservation now form one construction, topology, and precision policy. Full +forward, backward, hash routing, and checkpoint recomputation remain a separate runtime report. + +All six reports in the three affected modules pass. Repository collection is now 763 items and the static inventory is 749 +definitions with 1,138 curated decisions across 332 Python test files. The static audit surfaces 22 conditional runtime +gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-thirteenth wave: join configuration and numerics at kernel boundaries + +The two-hundred-and-thirteenth wave consolidates three static report definitions and four collected items without removing +an assertion. Before-down and after-down routing-weight numerics, no-router-gradient behavior, lazy configuration, environment +override, automatic regime selection, parity opt-out, explicit settings, and invalid values now form one routing-position +contract. + +Zero and nonzero cast-once LoRA folding now reports the same FP32-add invariant across dense linear and MoE expert layouts. +FlashQLA forward output, final state, and all input gradients now form one parity report per head shape against FLA; the two +head shapes and Hopper/TileLang capability gate remain unchanged. + +All four affected reports pass, including both live FlashQLA Hopper cases. Repository collection is now 759 items and the +static inventory is 746 definitions with 1,141 curated decisions across 332 Python test files. + +## Two-hundred-and-fourteenth wave: report quantized formats by shared invariant + +The two-hundred-and-fourteenth wave consolidates three report definitions without removing an assertion. NVFP4 and block-FP8 +expert loading now form one prequantized expert-load policy covering packed bytes, scales, global factors, amax values, +projection layout, and dequantization. + +Flat and GKN NF4 codebook, packing, scale, zero, shape, and error behavior now form one codec contract. Block-FP8 and NVFP4 +GNK-to-GKN transpose equivalence, direct and transposed dequantization, non-square shapes, expert stacking, and global-scale +absorption now form one prequantized layout-conversion policy. + +All three live GPU reports pass. Repository collection is now 756 items and the static inventory is 743 definitions with +1,144 curated decisions across 332 Python test files. + +## Two-hundred-and-fifteenth wave: group geometry and topology variants by claim + +The two-hundred-and-fifteenth wave consolidates three report definitions without removing an assertion or subprocess. +Same-NK and same-MN grouped GEMM numerics, transpose handling, uneven and empty groups, and input admission now form one +kernel-family contract. + +Dense 2-D Shard(0) and expert 3-D Shard(1) Muon layouts now form one distributed full-gradient oracle-parity report. The +same two layouts form one shard-local negative-control report. All four two-GPU subprocesses remain and still independently +exercise their original mode and layout. + +All three affected reports pass. Repository collection is now 753 items and the static inventory is 740 definitions with +1,147 curated decisions across 332 Python test files. The static audit surfaces 22 conditional runtime gates, one intentional +duplicate group, and no parse errors. + +## Two-hundred-and-sixteenth wave: report Qwen RMSNorm resolution by capability domain + +The two-hundred-and-sixteenth wave consolidates three report definitions without removing an assertion. Dense Qwen3.5 +RMSNorm structural selection, v1 and v2 dispatch, invalid-mode admission, every zero-centered site, GDN exclusion, layer +input policy, and final norm policy now form one CPU resolution contract. + +Qwen3.5-MoE v1 and v2 dispatch, ordinary-mode preservation, every zero-centered site, layer input policy, and final norm +policy likewise form one CPU resolution contract. Its family-1 interpose and full-layer bit parity and family-2 residual +composition now form one GPU bit-exact integration contract. The generic RMSNorm structure, undeclared-family tripwire, +kernel funnel, and dense/MoE capability separation remain independent. + +All three affected reports pass. Repository collection is now 750 items and the static inventory is 737 definitions with +1,150 curated decisions across 332 Python test files. + +## Two-hundred-and-seventeenth wave: make conversion, compilation, and routing reports end to end + +The two-hundred-and-seventeenth wave consolidates four static definitions and five collected items without removing an +assertion. DeepSeek-V4 converter meta-model dtype preservation, ordinary HF-to-DCP roundtrip, legacy sidecar-free LoRA load, +and cross-shard weight/scale deferral now form one conversion policy. AutoModel loading remains a distinct Transformers +integration contract. + +MoE-block and decoder-layer compilation across every available expert backend and both compiler backends now form one +lower-level compile-compatibility report. Full-model per-layer composition remains separate. Local unfiltered, local +filtered, and EP SGLang fused-expert backward paths now form one FP32 routing-gradient oracle report; the former local +parameterization is an internal two-case loop so the EP branch runs once rather than twice. + +All five affected reports pass, including the live compiler paths. Repository collection is now 745 items and the static +inventory is 733 definitions with 1,154 curated decisions across 332 Python test files. + +## Two-hundred-and-eighteenth wave: join registry, wrapping, and native-FP8 state policies + +The two-hundred-and-eighteenth wave consolidates three report definitions without removing an assertion. Optional Quack +registration, registry signatures, common backend arguments, activation forwarding, and explicit FP8 admission now form +one EP adapter boundary contract. + +Exact-hook merged-LoRA preparation and hybrid-model trunk selection now form one Qwen3.5 structural wrapping policy. The +stale synthetic fixture now declares the v1 family and a resolved norm, restores the nonexact numerical family afterward, +and asserts the current documented behavior that wrapping arms the RMSNorm/trunk contract lane. The isolated live GPU +forward remains a separate execution gate and passes after the state leak is removed. + +GLM-5.2 native-FP8 configuration roundtrip and rejection, model replacement, sparse-MLA materialization, pair-buffer bytes +and admission, expert fusion, and grouped checkpoint ownership now form one configuration, model, buffer, and checkpoint +contract. Router dispatch and frozen expert scoring remain separate. + +All six reports in the three affected modules pass. Repository collection is now 742 items and the static inventory is 730 +definitions with 1,157 curated decisions across 332 Python test files. The static audit surfaces 22 conditional runtime +gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-nineteenth wave: join wrapper layers around one execution claim + +The two-hundred-and-nineteenth wave consolidates three report definitions without removing an assertion. Direct Quack MoE +TP-FP8 execution and the MoEExperts Quack wrapper now form one train-step policy across TP reduction, Triton-grouped and +scalar-Quack backends, clamped SwiGLU biases, finite gradients, and master-weight mutation. + +Scoring and trainable SGLang EP dispatch, flag and dependency admission, empty-rank behavior, and autograd guards now form +one dispatch policy. Slot combination, weight presentation/cache modes, and live stock-Triton gradient parity remain +separate. Fused RMSNorm residual and no-residual forward/backward kernels, packed shapes, module dispatch, dense Qwen layer +parity, and serving-tree force calls now form one numerical and model-integration report. CPU fallback and trunk-specific +dispatch remain separate. + +The three affected modules collect 12 reports: ten pass and two optional DeepGEMM or SGLang capability gates skip. +Repository collection is now 739 items and the static inventory is 727 definitions with 1,160 curated decisions across 332 +Python test files. + +## Two-hundred-and-twentieth wave: report rotary, MoE activation, and GDN math as whole policies + +The two-hundred-and-twentieth wave consolidates four report definitions without removing an assertion. Pairwise-interleaved +rotary numerics, Class-B fused admission and CUDA rejection, dense and MoE attention half-rotate semantics, mRoPE behavior, +and post-RoPE BF16 casting now form one Qwen3.5 rotary policy. + +DeepSeek-V4 shared-MLP and routed-expert SwiGLU clamping now run inside the non-hash MoE structure, shared-contribution, +forward, backward, router, and selection-bias report. Hash-table routing and record/replay remain separate. GDN gating and +gated RMSNorm forward, backward, dtype, reference, and row-invariance behavior now form one primitive numerical contract; +exact-model module dispatch and the pinned triangular-solve geometry remain separate. + +All seven reports in the three affected modules pass. Repository collection is now 735 items and the static inventory is +723 definitions with 1,164 curated decisions across 332 Python test files. + +## Two-hundred-and-twenty-first wave: join state, backward, and dispatch phases + +The two-hundred-and-twenty-first wave consolidates three report definitions without removing an assertion. Native block-FP8 +byte packing, dtype-preserving state, strict state-dict and DCP metadata, apply rollback, CPU fail-closed execution, +partition-hook entry, scoring-only admission, range validation, and pair validation now form one encoding, checkpoint, +execution, and admission contract. + +Sparse-MLA production backward parity to the torch reference and deterministic-versus-atomic gradient parity now form one +backward policy; forward reference parity remains separate. RMSNorm v2 fused-versus-split forced-realization bit identity +and the production dispatch heuristic now form one realization and dispatch policy; one-ULP reference correctness, batch +composition invariance, and run-to-run determinism remain separate. + +All five affected reports pass, including the live H100 TileLang paths. Repository collection is now 732 items and the +static inventory is 720 definitions with 1,167 curated decisions across 332 Python test files. The static audit surfaces 22 +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-twenty-second wave: join primitive admission and topology branches + +The two-hundred-and-twenty-second wave consolidates six report definitions without removing an assertion or subprocess. +Class-B RoPE dtype admission, shape and partial-rotary backward behavior, and unique-half table layout now form one primitive +contract. EP backend declarations and parameter-metadata domains now form one fail-closed admission policy, while the real +two-rank reduction remains separate. Exact dense, projection, and LM-head components now share one legacy merged-sync +rejection report; separate-factor publication remains distinct. + +Qwen3.5 Ulysses execution and ring-plus-FLA rejection now run as two subprocess branches of one context-parallel contract. +NVFP4 and block-FP8 loading, execution, merge, and requantization now form one quantized QLoRA format lifecycle. The affected +reports pass, including live CUDA and two-GPU execution. Repository collection is now 726 items and the static inventory is +714 definitions with 1,173 curated decisions across 332 Python test files. + +## Two-hundred-and-twenty-third wave: make builder and artifact reports end to end + +The two-hundred-and-twenty-third wave consolidates five report definitions without removing an assertion. Full-model FP8 +construction, tensor-parallel LM-head selection, incompatible-mode admission, QARL construction, and calibration order now +form one quantized model-builder lifecycle; scoped monkeypatch contexts preserve the former report isolation. Sequence and +ring routing-replay layouts now form one context-parallel layout contract. Sparse-delta encoding, source capture, path +admission, and rank/global manifests now form one artifact lifecycle. + +DeepSeek-V4 FP8 and MXFP4 decoding and EP-aware handler ownership now form one checkpoint conversion policy, while synthetic +end-to-end loading remains separate. All six affected reports pass. Repository collection is now 721 items and the static +inventory is 709 definitions with 1,178 curated decisions across 332 Python test files. + +## Two-hundred-and-twenty-fourth wave: report model support as complete integrations + +The two-hundred-and-twenty-fourth wave consolidates six report definitions without removing an assertion. Active-LoRA atomic +flag mutation and composite admission now form one state policy; server derivation, topology rejection, and cached indexer +and MoE activation now form one propagation policy. Nemotron-H published per-expert and Transformers stacked expert layouts +now run through one bidirectional checkpoint-handler contract. LoRA manifest target selection and fail-closed schema and +runtime validation now form one manifest policy. + +Qwen2 and OLMo2 configuration conversion, architecture construction, tensor-parallel unfusing, checkpoint translation, and +HF parity now each form one architecture-support report. All seven affected reports pass. Repository collection is now 715 +items and the static inventory is 703 definitions with 1,184 curated decisions across 332 Python test files. + +## Two-hundred-and-twenty-fifth wave: join process, operator-edge, and packing policies + +The two-hundred-and-twenty-fifth wave consolidates six static report definitions and five collected items without removing +an assertion. Quack worker framing, PTXAS timeout and temporary-output handling, entry selection, and safe deterministic +cache hashing now form one compilation process-safety policy. Shared-prefix multi-member and singleton behavior now form one +forward and backward equivalence report; its unchanged optional FA3 module gate means neither former definition contributed +to repository collection in this environment. + +OPD ignored-token, per-token, hidden-only, full-materialization, and chunked-fetcher behavior now form one edge contract. +Packing strategy validation, oversized handling, document and token preservation, capacity, utilization, determinism, and +datum order now form one generic packing policy; balanced-DP scheduling remains separate. The five runnable affected reports +pass, and shared-prefix attention skips at its unchanged optional dependency gate. Repository collection is now 710 items +and the static inventory is 697 definitions with 1,190 curated decisions across 332 Python test files. The static audit +surfaces 22 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-twenty-sixth wave: join public service lifecycles + +The two-hundred-and-twenty-sixth wave consolidates six report definitions without removing an assertion. Tinker session +schema publication, creation, follow-up requests, heartbeat activity, and canonical LoRA aliases now form one public +endpoint lifecycle. Mooncake tensor storage, R3 reference publication, selective loading, validation, and cleanup now form +one side-payload lifecycle. Adapter checkpoint fail-closed admission, write failures, and successful dense and MoE live +factor publication now form one save policy with scoped patch isolation. + +Launcher address discovery and worker readiness now form one control lifecycle. P2P size-based dispatch, status timeout, +prepare timeout, and request payload now form one async API policy. K3 tail metrics and temperature-matched behavior logprobs +now form one observability policy across both loss implementations. All nine affected reports pass. Repository collection +is now 704 items and the static inventory is 691 definitions with 1,196 curated decisions across 332 Python test files. + +## Two-hundred-and-twenty-seventh wave: report storage and quantization end to end + +The two-hundred-and-twenty-seventh wave consolidates ten report definitions without removing an assertion. Mooncake hidden +tensor codecs, metadata, teacher consumption, malformed and legacy admission, removal, and configuration precedence now form +two transport and store policies. Dense QARL codec parity, injection, summaries, configuration normalization, and model +admission now form one policy. NVFP4 MoE identity-preserving conversion, eager execution, gradients, passthrough, injection, +target selection, and format admission likewise form one expert lifecycle. + +QARL sync configuration success and mismatch reporting now form one sync policy. Generic expert backend capabilities, +factor ownership, and preserved semantics now form one adapter contract. Teacher-head discovery, sharded storage, cross-shard +views, residency, dtype reload, and prefetch now form one head lifecycle. Exact-server trunk wrapping and numerical-family +selection now form one pre-parallelization program policy. All ten affected reports pass. Repository collection is now 694 +items and the static inventory is 681 definitions with 1,206 curated decisions across 332 Python test files. + +## Two-hundred-and-twenty-eighth wave: join export, pipeline, and trainer admission layers + +The two-hundred-and-twenty-eighth wave consolidates eight report definitions without removing an assertion. NVFP4 packed +codec correctness and full directory export now form one exporter contract. FP8 CLI configuration, base-directory behavior, +architecture-specific layouts, preflight, and QARL-fold rejection now form one command contract; primitive quantization and +trained-logprob preservation remain separate. + +OPD endpoint identity and student-version verification now form one endpoint-admission policy, while chunk queueing, causal +payload shifting, cache-index alignment, and Mooncake transport form one preparation policy. DeepSeek-V3 router-freeze +construction and downstream TP rejection now form one training-admission policy. Class-B selection and the canonical GLM-5.2 +numerical program now form one configuration policy, while exact Qwen3.5 remains separate. Simulator topology, shapes, +analytical ledgers, configuration fingerprints, model metadata, and path admission now form one trusted-input policy. All 13 +affected reports pass. Repository collection is now 686 items and the static inventory is 673 definitions with 1,214 curated +decisions across 332 Python test files. The static audit surfaces 22 conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-twenty-ninth wave: join distributed planning and control policies + +The two-hundred-and-twenty-ninth wave consolidates nine report definitions without removing an assertion or subprocess. +DeepEP node-span detection, preflight skip and roundtrip behavior, failure diagnostics, NVL alignment, and RDMA byte admission +now form one internode transport policy with scoped patches. Generic EP meta slicing and replicated-gradient metadata now +form one plan application; exact GLM meta and materialized already-local factor dispositions form another. Pipeline FQN +partitioning, rank placement, schedule metadata, and microbatch admission now form one layout policy. + +Busy-interval union and P2P-byte estimation now form one PP accounting policy, while patching and live CUDA execution remain +separate. Muon full-gradient oracle parity and shard-local divergence now form one distributed policy while preserving all +four two-GPU subprocesses. DeepSeek-V4 C128 and C4 compression admission now form one context-parallel regime policy. All +nine affected reports pass, including PP CUDA events and the Muon subprocesses. Repository collection is now 677 items and +the static inventory is 664 definitions with 1,223 curated decisions across 332 Python test files. + +## Two-hundred-and-thirtieth wave: report GPU kernel variants as whole contracts + +The two-hundred-and-thirtieth wave consolidates five report definitions without removing an assertion. BI fused LM-head +loss parity, gradients, determinism, batch invariance, guards, unit-temperature identity, and near-one logprob clamping now +form one selected-logprob contract. Full and dimension means now form one batch-invariant reduction policy. Head-v2 +projection and statistics bits, batch invariance, fused CE, gradients, and rollback now form one lifecycle. + +GDN primitive forward and backward numerics now run with exact-model module dispatch, while triangular-solve geometry remains +separate. Quack EP parity against Triton now runs with its independent half-concatenated activation reference; CPU gradient +arity remains separate. All seven affected live GPU reports pass. Repository collection is now 672 items and the static +inventory is 659 definitions with 1,228 curated decisions across 332 Python test files. + +## Two-hundred-and-thirty-first wave: join model construction and runtime boundaries + +The two-hundred-and-thirty-first wave consolidates ten report definitions without removing an assertion. Runner side-channel +conversion, ragged padding, and sequence sharding now form one batch-materialization policy. Exact dense managed factors and +routed fail-closed ownership now form one adapter-gradient policy. DeepSeek-V4 codec, handler ownership, synthetic loading, +construction, topology, precision, forward, backward, routing, and recomputation now form two complete checkpoint and model +contracts with scoped patch isolation. + +Exact attention source inventory and native pair-state behavior now form one checkpoint contract. Canonical routed and shared +MoE partials now form one boundary policy. Native-FP8 serving routing and frozen scoring-only experts now form one runtime +policy. Nemotron-H EP ownership, layout conversion, HF parity, and save behavior now form one checkpoint contract. Fused +RMSNorm ordinary and trunk GPU integration now form one policy while CPU fallback remains separate. RoPE registry recipes and +lazy native caches now form one CPU precision policy while exact serving-device execution remains separate. All 13 affected +reports pass, including fused RMSNorm CUDA execution. Repository collection is now 662 items and the static inventory is 649 +definitions with 1,238 curated decisions across 332 Python test files. The static audit surfaces 22 conditional runtime +gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-thirty-second wave: join residual failure and lifecycle phases + +The two-hundred-and-thirty-second wave consolidates eight report definitions without removing an assertion or subprocess. +Adapter-gradient pre-rendezvous, ModelRunner tail, and publication-commit failures now form one bounded CPU failure policy; +the asymmetric post-mutation GPU boundary remains separate. Live two-rank clipping and nonfinite behavior now run with the +three-rank participation topology as one distributed clipping policy. + +FutureStore creation, processing, concurrency, model operations, expiration, and cleanup now form one async lifecycle; a +fresh identical store preserves the former fixture isolation between its phases. Orchestrator client roundtrip, interleaving, +edge errors, and shutdown likewise form one communication lifecycle. Sparse-delta initialization and runtime loading, FP8 +LM-head selection and loss dispatch, and sequence-shard core and side-channel materialization each now form one policy. All +15 affected reports pass, including the five distributed subprocesses. Repository collection is now 654 items and the static +inventory is 641 definitions with 1,246 curated decisions across 332 Python test files. + +## Two-hundred-and-thirty-third wave: report CPU boundary policies end to end + +The two-hundred-and-thirty-third wave consolidates five report definitions without removing an assertion. Server and CLI +sequence boundaries, int32 metadata, original-position preservation, stale-metadata replacement, LCM padding, post-shard +divisibility, and padded unpacking now form one sequence-metadata and padding policy. AnyPrecision AdamW cautious numerics, +chunked state updates, gradient reuse, and DTensor offload now form one optimizer lifecycle. Inference endpoint registration, +worker and FP8 KV-cache admission, auto-sync, and health-aware listing now form one public endpoint lifecycle. + +All eight resulting reports in the three affected modules pass. Repository collection is now 649 items and the static +inventory is 636 definitions with 1,251 curated decisions across 332 Python test files. The static audit surfaces 22 +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-thirty-fourth wave: join transport, input, and support lifecycles + +The two-hundred-and-thirty-fourth wave consolidates nine report definitions without removing an assertion. P2P prepare, +cached-map behavior, fanout, cleanup, receiver completion, source staging, invalid manifests, and transfer diagnostics now +form three initialization, staging, and failure policies. FP8 receiver-layout handling remains separate from generic staging. + +Microbatch splitting, epoch delegation, builder batch sizing, sampler configuration, sequence-parallel insertion, and custom +collators now form one data-loader construction policy. Dataset name and shard expansion, type inference, splitting, and +merging now form one dataset-composition policy; raw loading and preprocessed persistence remain separate. MiniMax-M3 +configuration, registration, text forward/backward, and unsupported-input admission now form one architecture-support +report. CPU attention backend selection now runs with eager head-layout numerics, while optional FlashAttention paths remain +separate. FP8 training, block-FP8 QLoRA, QARL, aliases, defaults, and incompatible combinations now form one low-precision +argument policy. All 23 runnable affected reports pass and the unchanged FlashAttention capability report skips. Repository +collection is now 640 items and the static inventory is 627 definitions with 1,260 curated decisions across 332 Python test +files. + +## Two-hundred-and-thirty-fifth wave: join direct-EP, FP8, and SSD execution branches + +The two-hundred-and-thirty-fifth wave consolidates five report definitions without removing an assertion. Direct-EP +multi-sender initialization, scatter-copy ownership, dense and expert manifest partitioning, and rank-filtered transfers now +form one lifecycle across rank-zero, nonzero, failure, process-group, engine-order, and empty-rank behavior. FP8Linear padded +matmul recipes and correction numerics now culminate in live CUDA forward, backward, and master-weight mutation under their +shared capability gate. Dense and packed SSD recurrence, boundary-safe convolution, and packed mixer behavior now form one +CPU recurrence policy; unavailable-kernel admission and live GPU kernel parity remain separate. + +All 16 runnable affected reports pass and the unchanged optional SSD kernel report skips. Repository collection is now 635 +items and the static inventory is 622 definitions with 1,265 curated decisions across 332 Python test files. The static audit +surfaces 22 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-thirty-sixth wave: join runner, storage, optimizer, and FP8 lifecycles + +The two-hundred-and-thirty-sixth wave consolidates seven report definitions without removing an assertion. Adapter +coordinator materialization and auto-load now run with checkpoint restore, optimizer compatibility, trusted paths, lifecycle +reset, overrides, and structural admission using a fresh temporary subtree. OPD microbatch loss, gradients, FSDP LM-head +anchoring, metric aggregation, reductions, extrema, and empty-rank key alignment now form one execution-and-metrics policy. +RequestProcessor forward/backward now includes global packed-row batching and its rank-local and routing-replay boundaries. + +Model-scoped and sampler-scoped checkpoint listing, deletion, isolation, resolution, adapter reconciliation, tracking, and +normalized adapter-only export now form one public storage lifecycle with an explicitly fresh APIServer for the sampler +phase. Muon Gram-Newton-Schulz configuration and grouping now include Quack import, dispatch, and dtype selection. Direct +Quack FP8 expert variants now culminate in injected dense-expert-dense forward, backward, and master-weight mutation. All 24 +runnable affected reports pass and the unchanged opt-in DeepGEMM report skips. Repository collection is now 628 items and the +static inventory is 615 definitions with 1,272 curated decisions across 332 Python test files. + +## Two-hundred-and-thirty-seventh wave: join component primitives through their integrations + +The two-hundred-and-thirty-seventh wave consolidates 13 report definitions without removing an assertion. Canonical LoRA +folding, straight-through gradients, LoraLinear selection and cache invalidation, and MoE gate-up/down merged-weight caches +now form one CPU fold policy; native EP and trunk integrations remain separate. GDN delta-linear product correctness now +feeds sliced canonical folding, gradients, projection, and bounded cache behavior. + +Exact LM-head per-token and causal-loss routing, weight and server module selection, and FSDP replicated-factor admission now +form one complete loss policy. Absorbed-KV native state, logical masters, dtype moves, identity, and fail-closed direct +projection form one CPU component policy, while its official CUDA Q/V program remains separate. Canonical MoE capacity +metadata, transport admission, trainer/sampler hashes, topology, and group layouts now form one planning policy; the +distributed reduction subprocess remains separate. BI router GEMM, leading-dimension linear behavior, top-k normalization, +and exact-versus-stock MoEBlock dispatch now form one live CUDA routing contract. All 11 runnable affected reports pass and +the unchanged absorbed-KV CUDA capability report skips. Repository collection is now 615 items and the static inventory is +602 definitions with 1,285 curated decisions across 332 Python test files. The static audit surfaces 22 conditional runtime +gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-thirty-eighth wave: join server preparation, transfer, and diagnostic lifecycles + +The two-hundred-and-thirty-eighth wave consolidates 11 report definitions without removing an assertion. NCCL endpoint +initialization, two-phase completion, flattened and chunked buckets, hybrid receiver fencing, and multi-rank direct-format +admission now form one transfer policy. Weight-sync adapter selection, parameter extraction, inference layout, bucket sizing, +direct-EP sender mapping, and tensor-collection gating now form one source-preparation policy. Shipped MoE LoRA and QLoRA +examples now culminate the quantized server-configuration contract while retaining clean-process parsing. + +Sharded adapter packing, empty-shard ownership, deterministic initialization, and explicit-EP layout discovery now form one +CPU state policy; the real two-rank Gloo DTensor boundary remains separate. Dispatcher registration, nonresident saves, +gradient-epoch completion and abort, optimizer publication, poisoning, and fatal failures now form two session and mutation +lifecycles. Dense and MoE diagnostic hooks now feed the token diagnostic computation they support, while ranked tensor dumps +and trusted override loading form one artifact-boundary policy. All 15 resulting reports in the six affected modules pass. +Repository collection is now 604 items and the static inventory is 591 definitions with 1,296 curated decisions across 332 +Python test files. + +## Two-hundred-and-thirty-ninth wave: join packing, publication, and orchestration lifecycles + +The two-hundred-and-thirty-ninth wave consolidates 10 report definitions without removing an assertion. P2P hostname +selection and fallback engine construction now begin the existing initialization, prepare, fanout, cache, completion, and +cleanup lifecycle. Clean mid-epoch weight admission and strict checkpoint rejection now run inside the authoritative adapter +optimizer lifecycle. Packed teacher, cache, weight, hidden-state, and RL metadata now traverse the full pack, forward, and +unpack pipeline, and token-diagnostic unpacking now closes the RequestProcessor model-pass R3 payload lifecycle. + +Optimizer, forward, and forward-backward API response shaping now form one public training-operation policy. Runner +gradient-ownership compilation now includes staged-capture abort on forward-backward failure. OPD cache-row and last-k +weight shaping now feed live loss, gradient, and profiling execution. Balanced packing and sequential dummy behavior now run +inside dispatcher batch distribution, while routing references and microbatch diagnostic dumps form one side-payload +artifact and security policy. Orchestrator initialization, successful operations, errors, concurrency, statistics, and +shutdown now form one end-to-end lifecycle. All 28 resulting reports in the nine affected modules pass. Repository +collection is now 594 items and the static inventory is 581 definitions with 1,306 curated decisions across 332 Python test +files. + +## Two-hundred-and-fortieth wave: join optimizer, routing, FP8, and exact-component policies + +The two-hundred-and-fortieth wave consolidates 13 report definitions without removing an assertion. Cautious decay +primitives, SignSGD, AnyPrecisionAdamW state strategies, Muon post-Newton-Schulz masking, AdamW fallback, and builder routing +now form one optimizer feature policy. Synthetic balanced, sqrtsoftplus noaux, token-hash, legacy softmax, scaling, +configuration, and FP32 MoEBlock routing now form one TopK router contract. Optional boolean admission now runs with parallel +mixed-precision and reduce-dtype configuration, while sequence-parallel folding and manual prefetch remain separate. + +FP8 injection, recipes, exclusions, CPU fallback, output dtype, and fail-fast behavior now form one CPU module policy. CUDA +operand profiling now culminates the live FP8 matmul, correction, backward, and master-weight mutation report under the same +hardware gate; the CPU profiler remains separately runnable. GLM52 canonical MoE configuration now closes the sparse +selector and codec pipeline, while layer-plan allocation and semantic parity remain separate. Exact dense MLP factor +ownership, runtime admission, forward composition, XoRL load, and PEFT export now form one component lifecycle. Fourteen +runnable reports pass and the unchanged optional GLM52 capability report skips. Repository collection is now 581 items and +the static inventory is 568 definitions with 1,319 curated decisions across 332 Python test files. + +## Two-hundred-and-forty-first wave: join packing, loading, profiling, and loss execution + +The two-hundred-and-forty-first wave consolidates 12 report definitions without removing an assertion. Sample positions, +trainable-token filtering, dataset preprocessing, allocation, PackingDataset construction, caching, and missing-column +admission now form one data-packing lifecycle. Pipeline interval merging, P2P byte accounting, instance patching, +restoration, and patch admission now form one CPU profiling policy; the live CUDA GPipe report remains separate. + +Requested-key-only shard reads, exact merged and expert key plans, missing-pair admission, per-module deferred loads, and +bounded cache release now form one prequantized QLoRA loader lifecycle. NVFP4 and block-FP8 detection, quantized-key skipping, +QKV and bias merging, exclusion parsing, and dense and MoE handler behavior now form one checkpoint policy. Fused selected +logprob numerics now flow through CE, causal LM, Quack, and importance-sampling dispatch while the memory-bound regression +remains separate. Streaming forward-KL dense parity, chunking, masking, low-memory execution, OPD dispatch, and clamp +admission now form one execution policy while fp64 gradcheck remains separate. DistSign communication, hook ownership, +topology admission, construction, grouping, and stepping now form one optimizer lifecycle. Batch-invariant trunk forward and +backward now culminate in global-interpose gradient rejection and no-grad admission under the same CUDA gate. All 12 +resulting reports pass. Repository collection is now 569 items and the static inventory is 556 definitions with 1,331 +curated decisions across 332 Python test files. + +## Two-hundred-and-forty-second wave: join clipping, MoE, QLoRA, and replay lifecycles + +The two-hundred-and-forty-second wave consolidates 16 report definitions without removing an assertion. Shared-replica and +skip-FSDP ownership, norm modes, empty gradients, raw local clipping, EP-aware dispatch, ordinary fallback, and mixed-mesh +foreach behavior now form one CPU clipping policy; real two- and three-rank reductions remain separate. MoE histogram, +expert-slot indexing, deterministic ordering, escape-hatch behavior, scatter, gather, add-gather, and roundtrip execution now +form one CUDA kernel policy. BI GEMM table neutrality now culminates in row invariance across M buckets, while optional +DeepGEMM parity remains separate. + +QLoRA injection, NVFP4 and block-FP8 execution, format loading, merging, optimizer-state reset, and interval integration now +form one CUDA lifecycle. Exact DCP key projection now feeds an official base checkpoint into runtime state, while four-rank +staging and skip-mode admission remain separate. Exact shared-expert construction now includes native TP16 base slicing, +while optional SGLang factor views and Hopper execution remain separate. MoE LoRA construction, ownership, injection, eager +execution, zero-token gradients, and mocked EP score application now form one CPU component policy; cross-backend CUDA +numerics remain separate. Routing replay now joins asynchronous record ordering, MoEBlock replay, router gradients, +multi-layer and 1F1B schedules, base-model checkpoint enabling, and R3 preload under one CUDA lifecycle; its CPU registry +unit report remains separate. Thirteen runnable reports pass and three unchanged optional capability reports skip. Repository +collection is now 553 items and the static inventory is 540 definitions with 1,347 curated decisions across 332 Python test +files. + +## Two-hundred-and-forty-third wave: join checkpoint, exact-construction, and component lifecycles + +The two-hundred-and-forty-third wave consolidates 18 report definitions without removing an assertion. DTensor copy and +four-rank materialization, object transport, rank-zero filtered prefetch, local resolution, grouped expert routing, and +strict coverage now form one checkpoint-load lifecycle. Pipeline key unions, QARL buffer admission, base-to-LoRA +compatibility, optimizer-key filtering, multi-optimizer loading, metadata, load groups, and save groups now form one model +state lifecycle. Exact MoE global inventory now proceeds through EP placement, logical owner shapes, and selected-logprob +head attachment as one construction policy. + +Fused GDN manifest geometry now feeds low-rank products, canonical folding, gradients, bounded caches, export, and sharded +PEFT restore. FlashMLA flattening, invalid-index normalization, valid-row backward compaction, all-invalid behavior, and +production-envelope admission now form one hermetic policy. Exact TP1 construction and admission now run with CPU forward, +surrogate backward, and safety while the CUDA direct program remains separate. RMSNorm family tripwires now precede the +bitwise CUDA funnel while CPU structure remains separate. Fused MoE registration, Qwen checkpoint roundtrips, deferred +expert skipping, QKV unfusing, and QARL filtering now form one export policy. OPD full-vocab modes, VERL estimators, +policy-gradient behavior, and compiled sampled logprobs now form one loss policy. GDN convolution primitive parity now +culminates in end-to-end block output and gradients while optional SGLang parity and CPU admission remain separate. Twelve +runnable reports pass and three unchanged optional capability reports skip. Repository collection is now 535 items and the +static inventory is 522 definitions with 1,365 curated decisions across 332 Python test files. + +## Two-hundred-and-forty-fourth wave: join fused-kernel and architecture-support policies + +The two-hundred-and-forty-fourth wave consolidates 14 report definitions without removing an assertion. SGLang fused-MoE +resolution, block dispatch, admission, trainable dispatch, weight presentation, cache behavior, kernel layout, and runtime +context now form one CPU policy, while stock-Triton and masked gradients culminate in the existing optional real-SGLang +parity gate. Sparse-MLA attention-sink arithmetic and effect coverage now run in the representative first forward +specialization without multiplying the four compiled top-k cases. + +DeepSeek-V3 forward, backward, router freezing, default and explicit LoRA targets, router observability, and routing replay +now form one tiny-model lifecycle. DeepSeek-V4 non-hash routing, shared experts, SwiGLU clamps, hash-table routing, and +record/replay backward now form one architecture MoE policy. MiniMax-M3 configuration, registry, text execution, checkpoint +mapping, EP expert ownership, paged-KV layout, and CPU MSA admission now form one support report. GLM52 full-block QLoRA +inventory now includes EP-local routed banks plus exact-component and training-mode admission, while routed-expert literal +owner-slot coverage now culminates in sentinel and mixed-owner VJPs under its existing hardware gate. + +Six CPU/meta reports, four sparse-MLA forward specializations, both isolated sparse backward edge reports, and the FP8 +grouped-kernel report pass. Optional SGLang and SM100 reports skip on this H100 lane. The larger sparse-MLA backward +specializations and the Quack FP8 optimizer-step report terminate the current pytest process; both are retained as real +failure boundaries, and the passing edge/kernel reports remain separate so those terminations cannot erase their signal. +Repository collection is now 521 items and the static inventory is 508 definitions with 1,381 curated decisions across 332 +Python test files. The static audit surfaces 19 conditional runtime gates, one intentional duplicate group, and no parse +errors. + +## Two-hundred-and-forty-fifth wave: join model-support and EP-combine transactions + +The two-hundred-and-forty-fifth wave consolidates 10 report definitions without removing an assertion. GLM5 configuration, +registry loading, unsafe-value admission, indexer construction and selection, sparse-MLA reference and Ulysses integration, +checkpoint filtering, default adapter targets, EP MoE dispatch, absorbed-KV LoRA execution, tiny-model forward, and +recompute-before-dispatch now form one hermetic architecture-support policy. The real TileLang indexer fast path and HF +logit reference remain separate because they exercise external implementations. + +Qwen3.5 native EP8 admission now proceeds through variable-row token and ID collectives, trainer gradients through the +serving fused gate, FSDP module entry, routed and shared partials, chain summation, diagnostic actual operands, and final +output as one mocked transaction. SGLang EP flag and backend admission, empty-rank and trainable dispatch, slot-ordered +combine, pair-count guards, and transient, cached, and strided weight presentation now form one CPU policy; the optional +real stock-Triton gradient report remains separate. FlashQLA Gate 2 and Gate 4, the four unrelated training-utility APIs, +and GLM52's five CPU, CUDA, topology, and semantic boundaries were reviewed and retained as distinct contracts. + +All five runnable affected reports pass and the unchanged paired-SGLang gradient report skips. Repository collection is now +511 items and the static inventory is 498 definitions with 1,394 curated decisions across 332 Python test files. The static +audit surfaces 19 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-forty-sixth wave: join dispatch, export, optimizer, and sparse-sync lifecycles + +The two-hundred-and-forty-sixth wave consolidates eight report definitions without removing an assertion. MoE token and +routing-weight permutation now flows through mocked pre-dispatch all-to-all ordering, expert cumsums and score gradients, +then chunked and unchunked post-dispatch output and gradient parity. Packing-on capacity, batching, OPD and RL metadata, +generated labels, simulated forward output, and per-sample unpacking now form one roundtrip; packing-disabled remains a +separate supported mode with its own shifting, warning, and loss-mask behavior. + +Trained QARL state and exact dequantized target logprobs now precede the quantized directory and CLI artifact matrix, while +the low-level FP8 block quantization contract stays separate. Server Adam defaults and validation now feed full, partial, +omitted, adapter, and non-Adam optimizer mutation and finally dispatcher payload forwarding. Sparse-delta encoding, +baseline priming and rollback, prepacked per-rank posting, cache metadata, endpoint accounting, post-only admission, and +runtime helper loading now form one backend lifecycle with isolated temporary directories. + +All seven resulting reports pass. Repository collection is now 503 items and the static inventory is 490 definitions with +1,404 curated decisions across 332 Python test files. The static audit surfaces 19 conditional runtime gates, one +intentional duplicate group, and no parse errors. + +## Two-hundred-and-forty-seventh wave: join parsing, dataset, P2P, and FP8-sync policies + +The two-hundred-and-forty-seventh wave consolidates eight report definitions without removing an assertion. Optimizer, +packing, and numeric argument parsing now proceeds through Muon kwargs, EP checkpoint compatibility, automatic checkpoint +resolution, optimizer-state loading, FP8 aliases, fail-fast fallback, runtime-knob rejection, and low-precision mode +admission as one configuration lifecycle with isolated argv and environment contexts. + +Dataset name and shard expansion plus type inference now feed train-validation splitting, merge modes, local-file and saved +directory loading, hub, URL, and data-files routing, preprocessed persistence, reload, and missing-cache behavior. P2P trainer +IB-device precedence now proceeds through abort marker publication, peer observation and cleanup, then distributed +success/failure status gathering. FP8 weight-sync selection and block layout now feed dense and MoE adapter folding, +projection and skip-list policy, CPU expert transposition and padding, deferred quantization, streaming workspace and flush +behavior; its live GPU parity and device-transfer report stays separate. + +Checkpoint CRUD versus model-ID validation, weight-sync receiver versus source versus sparse transport, and inference +registration versus synchronization versus quantization schema were reviewed and retained as distinct public boundaries. +All five resulting reports, including live GPU FP8 execution, pass. Repository collection is now 495 items and the static +inventory is 482 definitions with 1,415 curated decisions across 332 Python test files. The static audit surfaces 19 +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-forty-eighth wave: join adapter, checkpoint, dispatcher, and endpoint lifecycles + +The two-hundred-and-forty-eighth wave consolidates nine report definitions without removing an assertion. Adapter +registration, session materialization, broadcast, cross-rank rollback, worker failure, fresh and evicted loading, explicit +path admission, rank-zero restore, and save refusal for missing evicted state now form one coordinator lifecycle. Optimizer +parameter identity and transactional collective failure now feed sharded manifest creation, artifact admission, bitwise +resume, one- and multi-dimensional logical resharding, replicated layouts, topology changes, and invalid-source rejection. + +Checkpoint manager materialization and zero-meta gates now proceed through model-runner initial base restore, step counters, +optimizer selection, restore failure publication, and guarded fresh default-adapter initialization. Dispatcher batch +sharding, EP and CP provenance, packing and dummy behavior now feed filesystem and Mooncake routing payloads, security and +diagnostic artifacts, CP replica deduplication, completion rendezvous, disagreement rejection, and rank-zero per-token +merge. Endpoint model discovery and health diagnostics now precede NCCL initialization, endpoint-port routing, and +two-phase receiver completion. + +All five resulting server reports pass. Repository collection is now 486 items and the static inventory is 473 definitions +with 1,424 curated decisions across 332 Python test files. The static audit surfaces 19 conditional runtime gates, one +intentional duplicate group, and no parse errors. + +## Two-hundred-and-forty-ninth wave: join ownership, OPD, and Muon producer-consumer policies + +The two-hundred-and-forty-ninth wave consolidates seven report definitions without removing an assertion. A fullgraph +module-managed adapter producer now feeds ownership compilation across dense, direct-output, EP-replicated, and +owner-sharded topology families, stable fingerprints, fail-closed structure and replica-domain admission, then bucketed +residual reduction, immutable raw accumulators, and logical norm accounting. + +Runner ownership compilation now proceeds from dense and exact LM-head producers through replica topology, unquantized and +quantized expert-factor contracts, registered session-rank specialization, rejected backend combinations, effective +LM-head folding, analytical gradients, capture finalization, and optimizer mutation. OPD teacher contributor selection, CP +gathering, Mooncake publication and cache row consumption now feed packed loss execution, metrics, and ranked +vocab-parallel debug artifacts. Muon configuration, grouping, fallback, Gram-Newton-Schulz stepping, and Quack admission now +include fused gate-up discovery and a tiny Nemotron-H parameter update. + +Attention registry versus FlashAttention versus paged-cache APIs, FSDP dtype versus transformation versus prefetch policy, +and standard Newton-Schulz versus live CUDA dtype preservation were reviewed and retained separately. All six resulting +reports pass. Repository collection is now 479 items and the static inventory is 466 definitions with 1,434 curated +decisions across 332 Python test files. The static audit surfaces 19 conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-fiftieth wave: join batch-invariance, DeepSeek, and indexer lifecycles + +The two-hundred-and-fiftieth wave consolidates four report definitions without removing a behavioral assertion. Dense +batch-invariant matmul, RMSNorm, log-softmax, and mean now establish the primitive contract before a padded Qwen sequence +proves full-model composition invariance under the same CUDA gate. DeepSeek-V4 window and compressed attention execution +now includes sink storage, TileLang call-dtype conversion, complex RoPE state, and TP admission, while grouped `wo_a` LoRA +contribution and gradients now feed the all-target attention-adapter freezing and first-step training policy. + +The TileLang indexer no longer launches a fifth kernel report and walks every masked cell in Python. Each of the four +retained execution geometries now checks its valid numerical scores and all invalid future positions from the same output +with a vectorized assertion; large-value and zero-input cases still run once. Dataset split fingerprints versus complete +configuration hashes, dataloader construction versus packed integration, Mooncake positive transport versus fail-closed +metadata handling, DeepSeek HF-to-DCP conversion versus AutoModel loading, and native-FP8 runtime versus checkpoint +construction were reviewed and retained as distinct boundaries. + +All eight resulting reports pass, including the four TileLang geometries in isolated processes. Repository collection is +now 475 items and the static inventory is 462 definitions with 1,443 curated decisions across 332 Python test files. The +static audit surfaces 19 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-fifty-first wave: join numerical trees and exact-state lifecycles + +The two-hundred-and-fifty-first wave consolidates three static definitions and four collected items without removing an +assertion. Families-v2 RMSNorm FP64 bounds, residual and zero-centered behavior, batch invariance, and repeatability now +precede forced fused-versus-split bit equality and live dispatch selection in one frozen-tree policy. Exact GLM dense, +projection, LM-head, streaming, and sparse-delta sync rejection now culminates in separate adapter-factor checkpoint keys, +bytes, and configuration. + +Index-share reentrant and non-reentrant checkpointing are now an internal Boolean mode matrix rather than separate pytest +IDs. Both modes retain producer recomputation, single payload creation, detached shared consumption, gradients, and closure, +then feed forward-only success, forward failure, backward failure, and idempotent cleanup. The standalone `solve_tril` +two-warp pin was reviewed and retained: warp count fixes a bit-relevant Triton reduction tree, while the GDN runtime report +only enforces tolerant numerical agreement and cannot replace that exact source-level gate. + +All three resulting reports pass, including live CUDA RMSNorm execution. Repository collection is now 471 items and the +static inventory is 459 definitions with 1,447 curated decisions across 332 Python test files. The static audit surfaces 19 +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-fifty-second wave: delete duplicated and fake-authored E2E reports + +The two-hundred-and-fifty-second wave removes four expensive or misleading reports. Pipeline schedule parity already +launches a two-GPU PP2/FSDP1 1F1B baseline before its three virtual-stage variants, so that baseline now owns convergence +and the standalone PP2 convergence process is gone. The LoRA FSDP2 checkpoint transaction now explicitly proves the load +marker, resumes to step 20, and enforces the former standalone convergence threshold; the redundant third Qwen3-8B +training process is removed. + +The CPU OPD suite no longer claims production loss invariance from a fake backend that calculated KL and global +normalization entirely in test code. It also drops a fixed teacher-cache metadata echo whose asserted values were authored +by the fake itself. The retained CPU end-to-end report still routes real packed teacher state through RequestProcessor, +Mooncake metadata, `TeacherActivationCache`, `TeacherHeadManager`, and the production `opd_loss_function`, then checks the +result against an independent grouped reference. + +The retained OPD report passes, and the strengthened two-phase LoRA FSDP2 report passes through step 20 in 162 seconds. +The retained pipeline schedule report exposes a current product failure before convergence: the 1F1B baseline reaches +stage backward with `Output gradient: None`. The seven full-weight FP8 E2E mechanisms, three P2P transfer boundaries, and +five adapter-manager state-machine boundaries were reviewed and retained. Repository collection is now 467 items and the +static inventory is 455 definitions with 1,454 curated decisions across 332 Python test files. The static audit surfaces 19 +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-fifty-third wave: remove dead E2E jobs and internalize implementation matrices + +The two-hundred-and-fifty-third wave removes the plain-AdamW eight-GPU PP2/FSDP4 training job: the surviving Muon job +uses the same pipeline, FSDP, packing, and microbatch topology while adding optimizer-partition coverage. More +importantly, the entire Nemotron-H E2E module is gone. Every report in that module requested +`tiny_nemotron_h_model_dir`, a fixture that does not exist in any repository conftest and was never added with the test; +all three definitions (four collected items) therefore errored at setup without constructing a model. Production +Nemotron model, packed-varlen, checkpoint, gradient, and optimizer-step behavior remains covered by functioning suites. + +Four implementation or shape matrices no longer inflate the product-report count. Native-FP8 linear and expert plain +conversion, FlashQLA four-head and production 32-head parity, PEFT MoE down-A and gate-B EP slicing, and non-gated Triton +and native MoE parity now execute as internal cases of their respective semantic policies. No branch or numerical +assertion was removed. + +The native-FP8 plain and two-rank FSDP2 lifecycle, all LoRA checkpoint reports, eager and both installed non-gated MoE +backends, and both FlashQLA head regimes pass. The retained eight-GPU PP2/FSDP4 Muon report reaches the 1F1B schedule and +then exposes the current pipeline backward product failure, consistent with the retained two-GPU gate. Repository +collection is now 458 items and the static inventory is 451 definitions with 1,461 curated decisions across 331 Python +test files. The static audit surfaces 18 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-fifty-fourth wave: prove reachability and remove historical skip shields + +The two-hundred-and-fifty-fourth wave adds a repository-wide fixture-reachability gate instead of trusting collection +alone. It found one remaining structurally dead report: the FP8 hybrid Ulysses/Ring long-tail E2E job requested +`tiny_agent_context_dense_model_dir_with_weights`, a fixture that has never existed in the repository, and therefore could +not reach model or trainer construction. That report and its now-unused dataset writer are removed; the functioning FP8 +matrix still owns checkpoint-resume, TP, Ulysses, Ring, local-MoE, and DeepEP EP/eFSDP boundaries. + +DeepSeek-V4 window and compressed attention and exact GLM sparse-attention CP modes now run as internal isolated matrices, +preserving their numerical, gradient, QAT, sink, exact-factor, query-offset, and topology assertions without multiplying +pytest report IDs. Triton and Quack routing-score gradients likewise form one backend policy. This review exposed that the +Quack parameter had never executed: its test-owned import stub omitted a required grouped-GEMM symbol and a broad +`ImportError` handler converted the defect into a skip. The stub is complete now, both implementations execute against the +reference, and internal import failures fail closed. + +Three historical skip shields are also gone. GLM5 and Dr.GRPO are shipped parts of this source tree, so their reports no +longer pretend those implementations are optional. The eager-versus-native MoE policy retains lazy imports for lightweight +collection but no longer catches every exception from production modules. All seven focused reports pass, including live +eager/native MoE and the formerly skipped Quack branch. A full `pytest --setup-plan` now succeeds with no missing fixture, +and global collection succeeds at 454 items. The static inventory is 450 definitions with 1,468 curated decisions across +331 Python test files; the audit surfaces 18 conditional runtime gates, one intentional duplicate group, and no parse +errors. + +## Two-hundred-and-fifty-fifth wave: remove dormant opt-in coverage and fail core backends closed + +The two-hundred-and-fifty-fifth wave removes the only environment-opt-in pytest report. The DeepGEMM grouped-FP8 +subprocess skipped unless `XORL_TEST_DEEP_GEMM_FP8=1`, but no repository workflow, script, or configuration ever sets +that flag. It was a dormant manual diagnostic presented as suite coverage. The three functioning FP8-MoE reports still +exercise injection, grouped forward and weight gradients, Triton-grouped and scalar-Quack execution, bias and activation +variants, tensor-parallel reduction, dense-plus-MoE composition, and real optimizer updates. + +The same review removes catch-all import skips from core backend gates. Quack is a pinned dependency; Transformers and +TileLang are core pinned dependencies; and the Quack, GKN, GLM5 indexer, and FlashQLA modules are shipped source. Their +tests retain explicit CUDA, SM90, and TileLang-feature admission gates, but packaging errors and internal import regressions +now fail instead of disappearing. All three remaining FP8-MoE reports, both Quack reports, both GKN reports, the GLM5 +TileLang parity report, all four FlashQLA exact-contract reports, and the FlashQLA-versus-FLA numerical report pass. + +Repository collection is now 453 items and the static inventory is 449 definitions with 1,473 curated decisions across +331 Python test files. The audit surfaces 18 conditional runtime gates, one intentional duplicate group, and no parse +errors. + +## Two-hundred-and-fifty-sixth wave: fold narrow admission branches and delete inert test scaffolding + +The two-hundred-and-fifty-sixth wave folds two narrow reports into their owning behavioral policies. Forcing the SSM +kernel while `mamba_ssm` is unavailable now closes the CPU SSD fallback, chunked recurrence, packed-sequence, and mixer +policy. Invalid learning rate, warmup ratio, and schedule-mode inputs now close the constant, linear, and cosine scheduler +policy. Every RuntimeError and ValueError assertion remains; neither implementation branch needs a separate product ID. + +Thirteen modules no longer carry `if __name__ == "__main__": pytest.main(...)` launch blocks. Repository and CI execution +already use pytest paths and node IDs, so those blocks were an unused second runner surface. The FP8 linear and MoE files +also drop GPU and skip markers from directly called private assertion helpers: pytest does not apply marker selection to +such calls. The seven public FP8 reports retain the actual hardware gates and all seven pass with every helper executing. + +The combined scheduler and SSM run reports three passes and the legitimate optional `mamba_ssm` kernel report skips; all +seven FP8 linear and MoE policies pass on CUDA. Global collection succeeds at 451 items. The static inventory is 447 +definitions with 1,477 curated decisions across 331 Python test files; the audit surfaces 18 conditional runtime gates, +one intentional duplicate group, and no parse errors. + +## Two-hundred-and-fifty-seventh wave: strip inert helper metadata and deduplicate synthetic routing + +The two-hundred-and-fifty-seventh wave removes 120 pytest marker decorations from private helpers across 30 distributed, +model, operator, optimizer, server, and weight-sync files. These `_assert_*` and worker helpers are not collected, and +direct calls do not apply pytest marker selection or skipping; the metadata therefore advertised hardware and async gates +that it did not enforce. Every public `test_*` report keeps its actual CPU, CUDA, architecture, async, distributed, and +optional-dependency markers. + +A private-helper reachability pass also removed four unused `NotImplementedError` overrides from a weight-sync QLoRA fake; +the base methods already fail identically, while the retained `dequantize_expert` method is the only fake behavior consumed +by the production merge-and-sync path. Full test-tree Ruff now passes after deleting one unused packed-dataset fixture +variable and documenting three intentional imports after standalone path bootstrapping. + +Balanced synthetic TopK routing no longer has a duplicate report. The canonical TopK policy already proves balanced expert +selection, uniform weights, count balance, and softmax/hash/bias override precedence. The unique MoEBlock replay-regather +assertion now closes the train-router dispatch policy. Sixteen representative reports pass across router, RMSNorm, fused +LM-head, GLM, pipeline profiling, async server, weight-sync, and future-store boundaries; one explicit optional backend +report skips. Global collection succeeds at 450 items. The static inventory is 446 definitions with 1,481 curated decisions +across 331 Python test files; the audit surfaces 18 conditional runtime gates, one intentional duplicate group, and no +parse errors. + +## Two-hundred-and-fifty-eighth wave: remove dead fixtures and separate certification from regression coverage + +The two-hundred-and-fifty-eighth wave maps every pytest fixture through direct parameters, fixture dependencies, +`usefixtures`, indirect parametrization, and dynamic `getfixturevalue` calls. Only two fixtures have no consumer: +`fake_packed_dataset` and `small_dense_model_dir_with_weights`. They are removed with the dead `FakePackedDataset` and +root-level `SimpleCollator` scaffolding. Full setup planning still reaches every collected report without a missing fixture. + +The tests tree also no longer presents manual workloads as pytest protection. DeepEP-versus-AllToAll parity, uneven +vocab-parallel OPD diagnostics, and the 100-step eight-H100 QLoRA comparison define no pytest report and have no workflow +caller; their direct-run value is preserved under `certification/deepep`, `certification/opd`, and +`certification/qwen3_30b`. A smaller standalone reverse-KL file likewise defined only a `main` function despite instructing +users to run pytest; it is removed because the production gathered path is exercised by the retained four-rank lm-head TP +FSDP/OPD policy across six CP, DP, and HSDP topologies. + +One collected report is removed on semantic grounds. It directly called two custom-autograd `backward` methods with a +fabricated `SimpleNamespace` context and `grad_output=None`, then checked a tuple length derived from that same fake context. +It never entered PyTorch autograd or production dispatch. The surviving Quack grouped-GEMM and DeepEP no-permute policies +run real backward graphs and compare all trainable gradients with trusted implementations. + +The focused data and Quack run reports four passes, and the consolidated four-rank lm-head TP FSDP/OPD owner passes all six +topologies in 140 seconds. Full setup planning, test-tree and relocated-certification Ruff, compileall, collection, and +diff-whitespace gates pass. Repository collection is now 449 items and the static inventory is 445 definitions with 1,485 +curated decisions across 327 Python test files. The audit surfaces 18 conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-fifty-ninth wave: replace zero-valued coverage and absorb narrow QARL reports + +The two-hundred-and-fifty-ninth wave finds a report that could never detect its claimed regression. The context-parallel +FLOPs check constructed an `xorl_glm5` configuration, but `XorlFlopsCounter` has no estimator for that model type and falls +back to zero. Comparing CP1 and CP64 therefore asserted only `0 == 0`. The rewritten report uses the supported Qwen3-MoE +estimator, first proves the baseline is nonzero, and then establishes that a global sequence-length input is not multiplied +by context-parallel size. + +Two narrow QARL report identities are removed without losing their useful behavior. The standalone activation NVFP4 report +used the production internal quantizer as its numerical reference and repeated the shared operator's two-dimensional STE +contract. Its unique leading-dimension reshape, exact value, and gradient assertions now close the independent pure-PyTorch +NVFP4 numerical policy. The W4A4 MoE file did not run a down projection or grouped GEMM; it checked only a temporary backend +name and exception restoration. Those assertions now close the existing CPU MoE conversion, execution, and injection +policy. + +The distillation and teacher-cache suites also share one Mooncake object-store fake instead of maintaining identical byte +APIs; the shared helper records the call keys needed by both consumers. Checkpoint process-group selection versus expert-mesh +restore, experiment ingestion and ranking boundaries, quantization primitives versus exporter disk layout, and OPD endpoint +verification versus payload transport were reviewed and retained as distinct production failure surfaces. + +All seven CPU QARL policies pass, and the focused shared NVFP4, FLOPs, Mooncake transport, and teacher-cache run reports +seven passes. Ruff, formatting, diff whitespace, JSON validation, global collection, and the static audit pass. Repository +collection is now 447 items and the static inventory is 443 definitions with 1,490 curated decisions across 326 Python test +files. The audit surfaces 18 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-sixtieth wave: consolidate server dispatch fragments and strengthen import boundaries + +The two-hundred-and-sixtieth wave removes three report identities that were fragments of broader policies. Launcher CLI +override parsing and removed-ZORL rejection now run with the existing YAML, direct-override, and unsupported-configuration +admission boundary. Launcher worker discovery and readiness remain separate because they exercise live address selection and +process failure behavior rather than configuration validation. + +The standalone ModelRunner causal-LM report used a zero-logit, two-token model only to prove that token losses are summed. +Its exact raw-sum and per-token assertions now precede the DR-GRPO cases in the retained `_compute_micro_batch_loss` +dispatch policy. Temperature, legacy field names, per-token output controls, K3 output forcing, and DR-GRPO metrics remain +covered in the same report. + +The QLoRA clean-interpreter smoke report is also absorbed into expert capability and ownership. It previously checked only +that imports returned zero even though its docstring claimed package decoupling. The retained policy now launches the clean +interpreter and explicitly proves that importing QLoRA utilities and expert modules loads neither `xorl.models` nor any of +its children. The server protocol policy likewise stops treating Torch as optional: the package is a core dependency, so a +broken Torch import now fails rather than silently skipping tensor serialization. + +Server API and security reports were reviewed and retained where they protect different consumers: outbound network +admission, artifact and diagnostic path confinement, compile-worker admission, API configuration validation, TensorData +re-nesting, session publication, optimizer fallback, training metrics, and ready-handshake queuing are separate failure +surfaces. The remaining heuristic candidates are explicit hardware or backend gates, not removal evidence by themselves. + +All nine surviving reports across the five touched server and QLoRA modules pass. Full test-tree and certification Ruff, +formatting, decision-JSON validation, public-tree lint, diff whitespace, global collection, and the static audit pass. +Repository collection is now 444 items and the static inventory is 440 definitions with 1,495 curated decisions across 324 +Python test files. The audit surfaces 18 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-sixty-first wave: attach model registration to real family loaders + +The two-hundred-and-sixty-first wave removes two standalone model registration reports. DeepSeek-V4 AutoConfig resolution +and XoRL meta construction previously fabricated the same tiny standard HF snapshot shape used by the retained AutoModel +loader but stopped before loading a tensor. The surviving snapshot policy now proves AutoConfig class resolution, the HF +AutoModel mapping, XoRL meta construction, actual `from_pretrained` loading, and exact embedding bytes in one sequence. HF +to DCP conversion remains separate because it exercises distributed checkpoint serialization rather than model admission. + +Nemotron-H registry lookup and Ultra-style local configuration normalization likewise now open the real model-family policy +instead of reporting alone. That policy proceeds through mixed Mamba, attention, and MoE construction, router output, +causal loss, backward gradients, router freezing behavior, and full-layer gradient checkpointing. Packed variable-length +execution and published-checkpoint parity retain their own reports because they cover independent runtime and codec paths. + +The registry-wide review retains the Kimi-wrapped DeepSeek-V3 policy because nested text-config aliases and official +auxiliary-loss defaults are absent from the base DeepSeek runtime. Qwen3.5 dense and MoE local normalization likewise has no +general family-construction owner. MiniMax M3 already combines registration with runtime, admission, checkpoint, and paging; +Qwen2 and OLMo2 already combine HF construction with fused/unfused layouts, checkpoint round trips, and numerical HF parity. +The small tokenizer, MTP checkpoint remap, gradient-checkpoint dispatcher, and BI operator reports were retained where each +reaches a distinct loader, execution, or reduction branch. + +Both Nemotron-H reports and the consolidated DeepSeek-V4 standard-snapshot loader pass. Full test-tree and certification +Ruff, formatting, decision-JSON validation, public-tree lint, diff whitespace, global collection, and the static audit pass. +Repository collection is now 442 items and the static inventory is 438 definitions with 1,498 curated decisions across 322 +Python test files. The audit surfaces 18 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-sixty-second wave: absorb data resilience and remove repeated stochastic certification + +The two-hundred-and-sixty-second wave removes the standalone request-retry report. The decorator has one production owner, +the high-level `prepare_datasets` operation, so immediate success, transient request and Hub failures, exponential backoff, +retry exhaustion, and unrelated-exception propagation now close the dataset preparation lifecycle. That lifecycle already +owns dataset expansion, type selection, local and remote loading, splitting, merging, saving, and reloading; the utility +report and file no longer multiply the product boundary. + +Stochastic rounding no longer proves one probability law with nested stress loops. The CPU policy previously accumulated +4,000 separately rounded 64-by-64 tensors, performing more than 16 million element updates to infer unbiasedness from a +generic relative-error maximum. It now creates one seeded 65,536-element population exactly one quarter of the way between +adjacent BF16 values and directly checks the legal neighbors, 25 percent round-up probability, and sample mean. + +The four-rank reduce-scatter report also drops 200 extra all-to-all trials that repeated the same unbiased-expectation claim. +It retains the distinct native-FP32 comparison and per-element BF16 transit error bound; the separate FSDP2 report retains +real compositional backward and optimizer coverage. The revised CPU probability policy passes, and the distributed policy +passes on four GPUs in 18 seconds rather than spending most of its execution re-certifying primitive randomness. + +The remaining optimizer reports were reviewed and retained because each already aggregates construction, grouping, +numerical updates, state strategy, cautious decay, and backend admission by optimizer family. Trainer gradient clipping, +token and microbatch metadata, pipeline chunked CE, explicit gradient synchronization, timer fail-soft handling, live CUDA +hooks, collator layouts, fingerprint identity, and packing reach distinct production consumers. + +The consolidated dataset preparation and stochastic-rounding policies report two CPU passes, and the four-GPU collective +reports one pass. Full test-tree and certification Ruff, formatting, compileall, decision-JSON validation, public-tree lint, +diff whitespace, global collection, and the static audit pass. Repository collection is now 441 items and the static +inventory is 437 definitions with 1,502 curated decisions across 321 Python test files. The audit surfaces 18 conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-sixty-third wave: remove scale-only and pseudo-E2E jobs + +The two-hundred-and-sixty-third wave maps the remaining end-to-end reports to production mechanisms rather than treating +GPU count as an independent behavior. The standalone single-GPU Qwen3-8B LoRA convergence job repeated the retained +two-GPU transaction's real model, rank, alpha, learning rate, 20-step horizon, and exact convergence threshold. The +surviving FSDP2 job additionally proves checkpoint save and load, the load marker, and the final global step; model and +server policies already own non-FSDP LoRA construction, forward, backward, optimizer, and checkpoint behavior. + +The CUDA OPD report was an E2E test in name only. It bypassed `ModelRunner` construction, invoked a private loss helper, +and optimized hidden-state and lm-head tensors directly for eight iterations. Its decreasing loss therefore described a +free-tensor optimization problem rather than a trainer or server lifecycle. The retained runner policy already proves +two-teacher cache loading, metrics, loss, and backward through the helper, while the real GPU server OPD report owns +`ModelRunner` startup, the forward/backward API, and the optimizer step. The pseudo-E2E file is removed. + +The remaining 18 E2E reports select distinct paths. FP8 covers dense resume, tensor parallel, Ulysses, Ring, plain MoE, and +DeepEP expert sharding. Pipeline reports separate direct trainer, schedule parity, server ModelRunner, FSDP, and folded +PP-EP-CP topologies. OPD retains request packing and Mooncake grouping, a real SGLang teacher, and the complete +sampler-teacher-Mooncake-trainer-weight-sync loop. DistSignSGD, hybrid shared-LoRA MoE telemetry, and LoRA checkpoint resume +retain their separate production boundaries. + +The retained runner and CPU OPD policies report two passes, and focused collection exposes the LoRA transaction plus all +three OPD integration layers. Full test-tree Ruff, scoped formatting, compileall, decision-JSON validation, public-tree +lint, diff whitespace, global collection, and the static audit pass. Repository collection is now 439 items and the static +inventory is 435 definitions with 1,505 curated decisions across 320 Python test files. The audit surfaces 18 conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-sixty-fourth wave: remove finite-output and scale-only workload tails + +The two-hundred-and-sixty-fourth wave audits semantic anti-patterns across the four largest surviving areas: models, server, +operators, and distributed integration. Qwen3.5's second trunk-wrap report built a tiny all-full-attention model, wrapped it, +and asserted only BF16 dtype and finite hidden states. The retained policy in the same file already proves the exact +full-attention, linear-attention, dense-MLP, shared-expert, and exclusion inventory. Generic GPU policies additionally prove +bitwise forward and backward, batch invariance, BF16 admission, serving-lane identity, and FSDP2 composition, so the +model-specific finite-output report is removed. + +The Quack DeepEP regression no longer treats token count as an independent behavior. Its 16K-token parity case repeated the +same production hidden size, expert geometry, balanced routing, unchunked no-permute path, trusted reference, and gradient +comparisons already exercised at 4K tokens. Its 32K-token checkpoint-training case repeated the same checkpointed three-step +optimizer path already exercised at 8K tokens. Random routing, skewed routing with empty experts, explicit chunking, and +checkpoint versus non-checkpoint training remain. The optional compiled `deep_ep` module is absent from this environment, so +the reduced report collects but its live two-GPU execution remains capability-skipped. + +Two block-FP8 workload tails are also removed. The generic codec allocated a 1024-by-2048 tensor only to compute storage +bytes from already-asserted dtypes and element counts, making the result tautological. The GKN codec's 4096-square roundtrip +selected the same multi-program two-dimensional kernel and error threshold already exercised by divisible and tail-tile +shapes. Both codec reports retain geometry, dtype, scale, accuracy, admission, determinism, sign, magnitude, and zero-block +coverage. + +The weak-name review retains reports whose behavior is stronger than their label. DeepSeek-V4 attention exercises both +window and compressed-KV forward and backward, every trainable gradient, FP8-QAT dispatch, sink dtype transfer, and TP +rejection. FP8 DeepEP uniquely composes no-permute transport with clamped-SwiGLU, native activation, expert biases, grouped +FP8 backward, and every gradient. BF16 stochastic reduction separately proves custom-hook installation and FSDP2 gradient +agreement. + +The Qwen selection policy reports one CPU pass and the two trimmed FP8 codec policies report two GPU passes in 18 seconds. +Full test-tree Ruff, scoped formatting, compileall, decision-JSON validation, public-tree lint, diff whitespace, global +collection, and the static audit pass. Repository collection is now 438 items and the static inventory is 434 definitions +with 1,509 curated decisions across 320 Python test files. The audit surfaces 18 conditional runtime gates, one intentional +duplicate group, and no parse errors. + +## Two-hundred-and-sixty-fifth wave: attach selector fragments to their production owners + +The two-hundred-and-sixty-fifth wave reviews the smallest surviving operator, distributed, model, server, checkpoint, and +loss reports for mocked metadata and argument-plumbing boundaries. Three standalone reports are fragments of broader +policies. The GatedDeltaNet backend report monkeypatched the FlashQLA chunk function and stopped after one call plus input +and output shapes; those assertions now close the CPU FlashQLA backend-selection and exact-contract precedence policy, +while real CUDA numerical, state-chaining, and batch-invariance reports remain separate. + +Runtime-rank MoE LoRA scaling now closes inference-buffer construction and FP8 sync rather than reporting alone. The +retained policy checks active-rank scaling through the production buffer builder, all three emitted projection names, +values, shapes, dtypes, source cleanup, QLoRA folding, and subsequent quantization. The standalone file and its artificial +one-by-one expert boundary are removed. + +Numerical-family selection likewise belongs to the model programs that set it. Both legacy environment aliases and GLM's +exact-v2 override now close the trainer model-builder policy; Qwen's exact-model hook proves its v1 LM-head pin overrides a +legacy v2 request. The selector-only file is removed. Families-v2 CUDA norm reachability stays separate from numerical-tree +realization because an unreachable correct kernel is a distinct regression. + +The remaining small reports were retained where they protect separate behavior: DeepEP async-combine safety versus +internode transport preflight, DSV4 optional-kernel rotation fallback, GLM4 MTP checkpoint remapping, token-loss +composition, gradient-accumulation process-group routing, server batch slicing, OPD cache streaming, and KKT launch +geometry. All four focused owner policies pass. Ruff, scoped formatting, compileall, decision-JSON validation, public-tree +lint, diff whitespace, global collection, and the static audit pass. Repository collection is now 435 items and the static +inventory is 431 definitions with 1,513 curated decisions across 317 Python test files. The audit surfaces 18 conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-sixty-sixth wave: collapse mock boundaries into behavioral owners + +The two-hundred-and-sixty-sixth wave ranks one-report files by monkeypatching, captured arguments, synthetic namespaces, +and absence of runtime work, then traces the highest-ranked candidates to their production consumers. Most mock-heavy +reports remain justified fault injection: NCCL rendezvous protects ephemeral-port rotation and bind-before-inference +ordering, checkpoint restore protects zero-meta admission and base-before-adapter initialization, and protocol round trips +preserve tensors while rejecting pickle. + +Three standalone boundaries are consolidated. GLM and Kimi ModelRunner target-resolution files previously rebuilt large +model configurations even though production reads only top-level `model_type`; both also repeated the same explicit-target +branch. One compact cross-family policy now retains GLM defaults, Kimi defaults including `lm_head`, explicit targets, and +manifest targets while deleting more than one hundred lines of irrelevant fixture data. + +Distributed-checkpointer process-group selection now closes the existing I/O policy rather than reporting alone. The +retained owner covers NCCL-to-Gloo selection, one-time caching, native-Gloo reuse, PP and non-PP metadata selection, custom +groups, and the actual load/save routing. Server batch-slice arithmetic similarly moves into the dispatcher policy, which +already owns distinct EP slices, CP sharing, EP-FSDP coordinates, padding, rollback, routing payloads, and completion. Its +replicated-DP mapping remains asserted, and the integrated sequence verified that the rollback environment switch is +restored before later cases. + +The three surviving owner policies pass. Full Ruff, scoped formatting, compileall, decision-JSON validation, public-tree +lint, diff whitespace, global collection, and the static audit pass. Repository collection is now 432 items and the static +inventory is 428 definitions with 1,517 curated decisions across 314 Python test files. The audit surfaces 18 conditional +runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-sixty-seventh wave: remove production code that exists only for tests + +The two-hundred-and-sixty-seventh wave reverses the audit direction: it inventories low-reference production helpers, then +checks whether tests are their only real consumer. This distinguishes externally callable or dynamically dispatched APIs +from convenience code that survives solely because an assertion imports it. + +The simulator's `compare_kernel_variants` helper had one caller, a test that divided two literal latencies after the retained +ranker had already ordered the same rows. The wrapper and its arithmetic assertions are removed; the surviving policy still +proves the important behavior that a faster unvalidated candidate cannot displace the validated winner. A second simulator +helper, `reference_counter_total_flops`, described itself as test support but had no callers at all after earlier policy +consolidation, so the orphan and its sole `SimpleNamespace` dependency are removed. + +The NVFP4 exporter likewise no longer ships a dequantizer solely so its test can grade the production quantizer with a +second implementation from the same module. The independent fake-quant policy retains exact numerical-reference coverage. +The exporter owner retains packed layout, scale shapes and dtypes, fused shared scales, BF16 islands, activation scales, +directory metadata, and requantization rejection. Low-reference sparse-delta reset, FP8 profiling, dynamic rank filtering, +teacher-store, QARL export, fused-expert cache, and Mooncake store APIs remain because they have operational, runtime, +public-package, or CLI ownership beyond assertion convenience. + +Both affected policies pass. Full test-tree Ruff, scoped formatting, compileall, decision-JSON validation, public-tree lint, +diff whitespace, global collection, and the static audit pass. Repository collection remains 432 items and the static +inventory remains 428 definitions with 1,521 curated decisions across 314 Python test files. The audit surfaces 18 +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-sixty-eighth wave: remove test-only GLM inspection surfaces + +The two-hundred-and-sixty-eighth wave deepens the source-consumer audit by tracing low-reference methods through dynamic +dispatch, framework registration, exports, documentation, and production callers. Runner command handlers, API endpoints, +documented DataLoader extensions, adapter transactions, and dynamically selected weight-sync methods remain. Four GLM +surfaces exist only to make assertions convenient and are removed. + +The exact LM-head component no longer exposes late TP-group binding when production always supplies the group to its +constructor. Its factor-view method also duplicated the BF16 casts already performed inside both real autograd functions; +the retained custom-boundary policy directly captures and checks those production bytes. Shared- and routed-expert public +factor-view wrappers likewise only cast masters before calling internal builders that the value paths already use, so the +tests now exercise those runtime-owned builders directly. + +The routed expert loses a larger test-only branch: a trace dataclass, clone-heavy hook wrappers, a capture flag production +always disabled, and a diagnostic method called only by its GPU test. Instead of inspecting staged caches from that alternate +path, the policy now compares actual module forwards with zero and live LoRA factors and verifies routing-scale linearity +through production execution. Exact buffer layout, owner remapping, independent hybrid VJPs, zero-slot gradients, and mixed +owner coverage remain. IndexShare similarly drops a context-manager wrapper absent from model execution; its lifecycle +policy now drives the same `begin` and `finish_forward` calls used by the model's try/finally. + +Three focused policies pass and the optional SGLang slice policy capability-skips; all affected GPU policies collect. Full +test-tree Ruff, scoped formatting, compileall, decision-JSON validation, public-tree lint, diff whitespace, global collection, +and the static audit pass. Repository collection remains 432 items and the static inventory remains 428 definitions with +1,526 curated decisions across 314 Python test files. The audit surfaces 18 conditional runtime gates, one intentional +duplicate group, and no parse errors. + +## Two-hundred-and-sixty-ninth wave: join numerical probes to their runtime owners + +The two-hundred-and-sixty-ninth wave scans the remaining reports for shape-only, dtype-only, finiteness-only, literal- +configuration, and duplicated family assertions. The apparent weak reports retain stronger numerical, gradient, cache, +checkpoint, or backend comparisons that a surface assertion classifier misses. Generic fused RMSNorm, explicit family +admission, and dense and MoE Qwen3.5 policies also remain separate: they own different kernel guarantees or execute +separate production implementations and call sites. + +Two genuinely narrow reports move into their behavioral owners without losing assertions. OLMo-2 no longer launches two +independent two-rank Gloo subprocesses over the same tensor-parallel mesh. Its surviving end-to-end policy first compares +plain and sharded `Olmo2QKRMSNorm` against local numerical references, then applies the production TP plan and proves model +forward, vocab-sharded LM-head execution, and gradients for every trainable parameter. The standalone Q/K-norm module is +deleted, saving one subprocess launch. + +Qwen3-MoE diagnostic decode likewise no longer reports the two-value FlashAttention causal flag independently. Both flag +branches now open the retained natural-cache and routing-replay cached-forward parity lifecycle. The two focused owner +policies pass. Full test-tree Ruff, scoped formatting, compileall, decision-JSON validation, public-tree lint, diff +whitespace, global collection, and the static audit pass. Repository collection is now 430 items and the static inventory +is 426 definitions with 1,530 curated decisions across 313 Python test files. The audit surfaces 18 conditional runtime +gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-seventieth wave: replace environment snapshots with selected behavior + +The two-hundred-and-seventieth wave audits exact package, accelerator, and process-state assertions before reviewing the +remaining files with multiple reports. Native-FP8 materialization no longer asserts that the worker happens to use the +literal `2.12.1+cu132` Torch wheel. The package and lock files already own that dependency pin, and a behavior regression +should remain runnable after a compatible upgrade. The worker now explicitly disables +`swap_module_params_on_conversion`, selecting the replacement path that exposed the DTensor corruption instead of merely +asserting its ambient default. Both the ordinary conversion policy and the two-rank FSDP2 materialization transaction pass. + +MoE compilation now forms one block-to-model CUDA policy. The surviving report runs every available MoE backend through +AOT eager and Inductor at block and decoder scope, then continues through native and eager full-model forward and backward. +The split reports shared the same capability gate and behavior owner, so a second pytest identity added no isolation or +failure meaning. Its real CUDA matrix passes in 64.78 seconds. + +The remaining apparent pairs stay separate where capability or oracle ownership differs: CPU coverage must not disappear +behind a CUDA skip; FP64 gradcheck is independent from sampled numerical parity; sparse-MLA forward and backward own +different kernels; and routing wire decode, context layout, and packing modes have distinct failure meanings. Full +test-tree Ruff, scoped formatting, compileall, decision-JSON validation, public-tree lint, diff whitespace, global +collection, and the static audit pass. Repository collection is now 429 items and the static inventory is 425 definitions +with 1,533 curated decisions across 313 Python test files. The audit surfaces 18 conditional runtime gates, one intentional +duplicate group, and no parse errors. + +## Two-hundred-and-seventy-first wave: delete dormant external integration and fail admitted backends closed + +The two-hundred-and-seventy-first wave audits exception-driven availability gates rather than accepting every skip as an +environment fact. The 522-line sparse-delta trainer-to-SGLang report depended on an unpinned `delta-encoding` tree and an +SGLang receiver module absent from the repository's pinned submodule. No workflow, script, or configuration supplies its +two path variables, and catch-all imports converted both absence and implementation breakage into a skip. Its fake trainer, +runner, orchestrator, HTTP endpoint, and receiver therefore never established repository coverage. The dormant report is +removed; retained policies still own XORL's packed artifacts, source capture, sorted indices, malformed-update rejection, +hashes, endpoint payloads, version forwarding, and sparse-delta transport lifecycle. Real cross-project byte compatibility +should return only as a dependency-pinned, scheduled integration. + +The standalone `TokenPartial` component file is also gone. Its caller-scaled denominator, microbatch composition, raw-sum, +sequence-mean-token-sum, and empty-mask assertions now close the shared loss policy that already proves explicit reducers +match the legacy policy and importance-sampling implementations. This preserves every oracle while removing an artificial +component report. + +Finally, supported backend admission fails closed consistently. MoE compilation no longer catches errors importing its +shipped capability helper or wraps an infallible Quack list append. Hard-pinned TileLang is no longer presented as an +optional sparse-MLA dependency after the explicit Hopper gate. Three DeepEP NVSHMEM path helpers likewise stop swallowing +errors after their callers have already admitted the package with `importorskip`. The shared loss policy passes, the full +MoE compiler matrix passes in 63.35 seconds, both live TileLang sparse-MLA policies pass in 25.04 seconds, and all affected +DeepEP reports collect. Repository collection is now 427 items and the static inventory is 423 definitions with 1,537 +curated decisions across 311 Python test files. The audit surfaces 18 conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-seventy-second wave: move private fragments into runtime owners + +The two-hundred-and-seventy-second wave traces the smallest remaining single-report files into the production lifecycles +that consume their behavior. Seven files disappear without dropping an assertion. The runtime FLOPs counter's cp1-versus- +cp64 global-sequence-length invariant now closes the simulator topology, shape, and analytical-ledger policy. Distributed +loss-group forwarding, valid-token normalization, and backward scaling now follow trainer token metadata counting, while +HSDP microbatch all-reduce deferral and restoration close the explicit SP and LM-head gradient-synchronization policy. + +OPD layer-cache indexing no longer tests a private fetcher in isolation. Exact selected indices, streamed layer ranges, +layer counts, and output shapes now precede the retained streaming OPD loss, gradient, cache, metric-reduction, and debug- +artifact lifecycle. ModelRunner's LoRA fragments likewise join their consumers: family defaults, explicit targets, and +manifest precedence close the adapter ownership compiler, while nonresident checkpoint promotion, failed-kill +preservation, registry cleanup, and path rejection close the optimizer, checkpoint-load, and session-registry lifecycle. + +Finally, direct sync-quantization dictionary examples now execute inside receiver detection and API admission. BF16 no-op +aliases, valid FP8 defaults and normalization, module exclusion cleanup, and every malformed or unsupported form remain, +alongside receiver discovery, unsupported-marker propagation, per-call enrichment, and persisted default behavior. The +seven focused owner policies pass. Full test-tree Ruff, scoped formatting, compileall, decision-JSON validation, public- +tree lint, diff whitespace, global collection, and the static audit pass. Repository collection is now 420 items and the +static inventory is 416 definitions with 1,542 curated decisions across 304 Python test files. The audit surfaces 18 +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-seventy-third wave: make feature lifecycles own their completion branches + +The two-hundred-and-seventy-third wave removes four more single-feature files by moving their exact assertions to the +policies that own the rest of each lifecycle. IndexShare checkpointing now continues from producer/shared recomputation and +mode-owned close behavior into both public callers: offline trainer and server forward failures each release retained +context exactly once. The orchestrator-runner wire protocol similarly continues from payload and tensor serialization, +command construction, and pickle rejection into the rank-zero ready handshake, including normal ACK, request-before-ACK +queueing, client identity, unexpected message, and channel-failure behavior. + +Packed sequence alignment now includes its distributed completion. The retained server-versus-CLI policy already owns +boundaries, int32 metadata, SP sharding, stale metadata replacement, LCM padding, and unpacking; it now also simulates an +eight-rank maximum and proves 176-token batches extend to 512 with ignored labels, masked attention, corrected cumulative +sequence lengths, and integer max lengths. + +The lm-head TP plus EP report no longer launches a four-rank process only to inspect groups. The retained FSDP transaction's +DP2 by CP2 case now enables EP2, proves the exact TP and replica memberships and active EP mesh, then continues through +parameter synchronization, vocab-sharded causal-LM loss, eager global-loss parity, and full weight and hidden-gradient +parity. The three CPU owners and this strengthened four-rank transaction pass. Full test-tree Ruff, scoped formatting, +compileall, decision-JSON validation, public-tree lint, diff whitespace, global collection, and the static audit pass. +Repository collection is now 416 items and the static inventory is 412 definitions with 1,546 curated decisions across 300 +Python test files. The audit surfaces 18 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-seventy-fourth wave: run server policies against canonical modules + +The two-hundred-and-seventy-fourth wave removes unrealistic import machinery rather than manufacturing a count reduction. +Seven server-runner reports were loading production files under synthetic module names, creating private copies of +`ModelRunner`, `RunnerDispatcher`, `AdapterCoordinator`, `CheckpointManager`, or `LoRAAdapterManager`. Their monkeypatches +therefore targeted test-only module state and could miss package import interactions. The reports now import and patch the +canonical runtime modules while preserving all adapter coordination, checkpoint failure, LoRA roundtrip, optimizer, +session-registry, load-state, and session-operation behavior. All 11 reports pass. + +The server-argument policy had the same issue at a larger scale: it imported the real launcher for override handling, then +re-executed `launcher.py` with fake API-server, orchestrator, session, QARL, and packing modules solely to obtain +`load_server_arguments`. It now uses the canonical launcher and real dependency graph for YAML admission, shipped-example +subprocess parsing, sparse-MLA propagation, and the rest of its configuration lifecycle. All four reports pass. + +The owner-level scan explicitly retains the smallest adjacent reports where size is not semantic duplication: DSv4 +fallback rotation, GLM4 MTP checkpoint remapping, DeepEP async-combine admission, KKT launch geometry, families-v2 norm +dispatch, BI mean, and Class-B RoPE each protect a separate numerical or fail-closed production boundary. This rewrite +leaves collection at 416 items and the static inventory at 412 definitions, with 1,549 curated decisions across 300 Python +test files. + +## Two-hundred-and-seventy-fifth wave: separate public CPU policy from optional kernel machinery + +The two-hundred-and-seventy-fifth wave removes a hidden implementation-detail report from the FA3-gated ring-attention +module. That report combined direct assertions on private `_get_zigzag_step_section` slices with the public +`zigzag_reorder_packed_sequence` behavior consumed by `TextSequenceShardCollator`; because the whole module imported FA3 +at collection time, none of its CPU-only assertions ran in the default environment. The public single-document, +packed-document, multi-rank, identity, and invalid-length contract now closes the collator's existing SP sharding policy, +and passes without FA3. The private step-section assertions are gone. The remaining ring-attention report is solely the +real CUDA partial-output merge policy. + +The skipped-runtime scan retains the FA3 numerical policies and both GLM exact SGLang joins. Those reports cross real +kernel or adapter-export, parser, and memory-pool boundaries rather than checking import availability. Their dependency +lane must remain explicit: XoRL's default profile is Torch 2.12.1 and has no `sglang-kernel`, while pinned SGLang declares +Torch 2.11.0 with `sglang-kernel==0.4.5`; loading a lazy wrapper is not an ABI smoke test. Exact SGLang-kernel execution +therefore belongs in the isolated Torch 2.11 environment, not in the default profile. Repository collection remains 416 +items while the static inventory falls to 411 definitions, with 1,551 curated decisions across 300 Python test files. The +audit surfaces 18 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-seventy-sixth wave: collapse synthetic MoE module graphs into the runtime owner + +The two-hundred-and-seventy-sixth wave removes the standalone EP routing-score report after moving its complete numerical +oracle into the adapter policy that owns `expert_scores`. The old report re-executed Triton and Quack source files under +test-only module names, after installing fake fused-MoE and grouped-GEMM packages in `sys.modules`. The retained policy now +patches the canonical runtime modules, proves both backend outputs and routing-score gradients against the same eager +reference, then continues through real registry admission, common signatures, FP8 rejection, optional MoE-act separation, +and adapter argument forwarding. The adjacent before-versus-after-down policy reuses only explicit CPU grouped-GEMM +doubles rather than importing helpers from another test report. Both owner policies pass. + +Quack's process and cache safety policy had the same synthetic-copy smell. It source-loaded the worker protocol, ptxas +wrapper, and cache utility under private names and supplied fake `cutlass` and `tvm_ffi` modules. Those components import +through the real package in the supported environment, so timeout, truncated-frame, unique temporary output, PTX entry +selection, and non-executable cache-key checks now exercise canonical module state. The focused policy passes. Repository +collection falls to 415 items and the static inventory to 410 definitions, with 1,553 curated decisions across 300 Python +test files. The audit surfaces 18 conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-seventy-seventh wave: replace synthetic dependency state with supported runtime state + +The two-hundred-and-seventy-seventh wave crosses trainer, attention, and executable-script boundaries. HSDP gradient-sync +deferral no longer replaces Torch's composable-FSDP module in `sys.modules` with a fake class. A lightweight instance of +the real Torch 2.12 `FSDPModule` API type now proves two-microbatch all-reduce deferral, last-microbatch enablement, +exception-safe restoration, and the replicate-size negative branch. + +The attention registry policy likewise stops manufacturing fake `flash_attn` and `flash_attn.cute` packages and reloading +two production modules. The default environment already supplies the FA4-only state this regression protects, so the +policy now checks canonical availability, registration under the compatible FA2/FA3 names, explicit FA4 registration, +eager and native resolution, and fail-closed unavailable-flash handling against the live runtime graph. + +Finally, student endpoint matching and weight-version verification now open the OPD pipeline payload and transport policy +instead of forming a second report. The standalone driver module is cached after one source load rather than re-executed +for every helper; all success, mismatch, endpoint-failure, worker, causal-shift, cache-index, and Mooncake metadata +assertions remain. The focused policies report seven passes and the expected FA3-only skip. + +The remaining dependency doubles are intentionally narrower: absent `delta_encoding`, Mooncake's CUDA-bound import, and +the isolated SGLang/`sgl_kernel` ABI lane are represented only to exercise serialization, fallback, runtime-context, or +slot-combine callers. They are not accepted as compiled-kernel smoke tests. Repository collection falls to 414 items and +the static inventory to 409 definitions, with 1,557 curated decisions across 300 Python test files. The audit surfaces 18 +conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-seventy-eighth wave: make shipped backends fail closed + +The two-hundred-and-seventy-eighth wave reviews every remaining conditional-runtime candidate rather than treating a skip +as evidence of uselessness. Two policies were incorrectly permissive. The EP adapter report could skip the whole policy +midway when a backend was absent, and the non-gated MoE report silently tested whichever subset of Triton and native +registered. Those are shipped default-profile backends, not optional integrations. Both reports now assert their complete +registry surface and execute every expected adapter-forwarding, FP8-rejection, forward, input-gradient, and weight-gradient +branch. The focused default-runtime policies report three passes. + +The full OPD pipeline also repeated its three-GPU admission check inside the test after its module-level marker had already +enforced the same rule before execution. The late branch is removed, so a skipped environment no longer creates model +artifacts before reaching a duplicate decision. + +The 15 remaining candidates are retained deliberately. They cover Hopper-only FlashQLA and exact GLM52 kernels, two-rank +FSDP composition, real CUDA profiler events, the isolated SGLang/`sgl_kernel` runtime, and optional DeepGEMM bit parity. +Each guards a complete numerical, distributed, or compiled-runtime transaction; none hides an ordinary CPU policy. +Repository collection remains 414 items and the static inventory remains 409 definitions, with 1,560 curated decisions +across 300 Python test files. The audit now surfaces 15 conditional runtime gates, one intentional duplicate group, and no +parse errors. + +## Two-hundred-and-seventy-ninth wave: close policy islands inside their runtime owners + +The two-hundred-and-seventy-ninth wave removes a tiny QARL "training smoke" that proved ordinary AdamW changes a tiny +model's weights and logprobs, then saved and reloaded its state dictionary. None of those assertions distinguished QARL +from ordinary PyTorch behavior. QARL injection, fake-quant arithmetic, summaries, persistent quantization state, and exact +reload remain in the dense fake-quant and calibration lifecycles. + +Three other standalone reports held useful behavior but no independent owner. QARL activation-quant enable, disable, +exception, non-QARL exclusion, and nested restoration now close the fake-quant owner. The Qwen3.5 families-v2 effective- +weight and dual residual-gradient wiring now closes the existing norm dispatch and site-assignment policy. LoRA BF16 base +retention, FP32 factors, trainability, dtype resolution, and generic-upcast admission now run with the FP8 and QARL model- +builder construction lifecycle. The three retained owner reports pass. + +Repository collection falls to 410 items and the static inventory to 405 definitions, with 1,564 curated decisions across +296 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-eightieth wave: delete test-only APIs and reunite isolated contract fragments + +The two-hundred-and-eightieth wave removes two source surfaces that existed only to support tests. Canonical MoE no longer +publishes an unused sampler-plan role, launcher alias, single-ordinal accessor, or JSON/digest representation; its retained +trainer policy still proves topology validation and the real distributed collective. Adapter gradient capture likewise no +longer exposes a one-call convenience absent from production. The adapter policies now drive the actual stage, commit, and +abort transaction, including the real multi-rank fatal boundary. + +Two useful but overly isolated reports join their contract owners. DeepEP's unsafe async-combine opt-in now closes the same +policy as internode topology, transport preflight, and buffer admission. GLM52 exact routed-ID, scaling, and shared +contributor forwarding now close the exact-MoE construction and global-inventory policy. Both retained owner reports pass. + +Repository collection falls to 408 items and the static inventory to 403 definitions, with 1,568 curated decisions across +294 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-eighty-first wave: replace isolated units with real owner transactions + +The two-hundred-and-eighty-first wave reunites four standalone reports with the runtime policies that consume them. Kimi's +local TikToken loader and safe generic fallback now close its DeepSeek registry/config lifecycle. SignSGD's core dense, +decay, missing-gradient, and sparse-rejection behavior now closes the optimizer policy; its repeated generic parameter-group +assertions are gone. Qwen3 projection unfusing no longer pretends to be a distributed report: the production model-level +transition, TP plan, checkpoint-handler state, and full layer inventory now close the `torch_parallelize` policy. + +The BF16 communication lane receives the stronger change. Its stochastic rounding primitive has no production consumer +outside `BF16StochasticAllToAllReduceScatter`, while the old transaction required four GPUs and was normally skipped. One +default-runtime report now proves rounding admission, deterministic seeding, neighbor distribution, unbiased expectation, +and a real two-rank CPU/Gloo all-to-all plus FP32 accumulation against native reduce-scatter. The six focused owner reports +pass. A repeated balanced synthetic-routing fragment is also removed; the retained TopKRouter policy already owns its +complete expert-sequence, balance, uniform-weight, bias, and hash-table behavior. + +Repository collection falls to 404 items and the static inventory to 399 definitions, with 1,573 curated decisions across +290 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-eighty-second wave: collapse format and kernel policy islands + +The two-hundred-and-eighty-second wave removes seven standalone reports whose setup and objects were already owned by a +larger policy. NVFP4 normalization, forward quantization, straight-through gradients, and weight-disable behavior now close +the dense QARL fake-quant owner. MoE sqrtsoftplus/softmax replay regathering now closes TopKRouter's selection, scaling, +dtype, and configuration policy. DSV4's pure-Torch rotation fallback now closes the compressor that consumes it. Qwen3-MoE +layer and final-norm declarations now join the existing dense-Qwen and shared-attention RMSNorm family owner. + +The batch-invariant lane loses three artificial boundaries. Families-v2 trainer reachability and its kill switch now close +the v2 norm realization/dispatch report. Families-v2 projection, scoring, invariance, backward, and rollback now close the +fused LM-head transaction. Full-reduce mean and dimensional reductions now close the global Torch-interpose policy that +already owns matmul, RMSNorm, log-softmax, and gradient admission. All nine retained owner reports pass, including their +CUDA execution paths. + +Repository collection falls to 397 items and the static inventory to 392 definitions, with 1,580 curated decisions across +283 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-eighty-third wave: reunite lifecycle keys, folds, and sync configuration + +The two-hundred-and-eighty-third wave removes four more standalone files spanning five reports. LoRA's zero-adapter and +nonzero permanent merge now closes the canonical fold/merged-forward owner for both linear and MoE storage; the same +cast-once contract now runs on CPU rather than sitting behind an unconditional CUDA assumption. Dataset split fingerprints +and preparation hashes now close the dataset loading, splitting, saving, and retry lifecycle that consumes those keys, +retaining determinism, sensitivity, fractional-size, tokenizer, column, and order-independence checks. + +Virtual-stage `MultiOptimizer` construction, delegation, single-part fallback, invalid explicit groups, and scheduler +fanout now close the learning-rate scheduler owner; distributed-checkpoint state filtering remains with its checkpoint +owner. QARL-derived FP8 sync metadata, skip lists, handler defaults, quantized buffers, and incompatible overrides now close +the production FP8 `WeightSyncHandler` policy rather than forming a separate QARL report. All six focused owner reports pass. + +Repository collection falls to 392 items and the static inventory to 387 definitions, with 1,584 curated decisions across +279 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-eighty-fourth wave: retire synthetic gates and reunite exact dispatch policy + +The two-hundred-and-eighty-fourth wave removes a synthetic identity-layer checkpoint report that installed `MagicMock` +checkpoint functions and enumerated a local boolean condition. Real Nemotron-H training already proves the default +full-layer checkpoint path executes and propagates gradients, while the GLM-5 model lifecycle proves +`recompute_before_dispatch` bypasses the outer checkpoint and invokes the layer's pre-dispatch checkpoint. The mock truth +table added no independent runtime contract. + +Two remaining standalone policy islands join their natural owners. Generic train-router true/false gradients were already +covered by routing-replay and real-model lifecycles; the unique DeepEP rejection now closes the TopKRouter/MoEBlock policy, +and the frozen server default closes configuration serialization. KKT's exact BK, warp, stage, safety, and off-lane launch +behavior now closes the GDN contract policy beside its existing solve-tril serving geometry. The three focused owner files +report seven passes. + +Repository collection falls to 389 items and the static inventory to 384 definitions, with 1,587 curated decisions across +276 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-eighty-fifth wave: replace component smoke with initialized lifecycle owners + +The two-hundred-and-eighty-fifth wave removes DSV4's isolated attention shape smoke. That report manually initialized +`torch.empty` parameters after bypassing model `post_init`, then repeated C0 and C128 forward shape, finiteness, and backward +reachability already owned by the fully initialized DSV4 model lifecycle. The only unique FP8-QAT dispatch and TP>1 +admission checks now execute through that full model owner; the direct-layer sink-dtype scenario is gone because production +model casting deliberately preserves the sink in FP32. + +Two small policy fragments also join real consumers. Strict LoRA-manifest count, rank, configured-target, unlisted-module, +schema, Boolean, and integer failures now close the fused-GDN injection/checkpoint lifecycle instead of a fake two-layer +attention tree. GLM4 MTP embedding, norm, and head aliases plus ignored auxiliary tail fields now close the GLM4 family +construction and checkpoint lifecycle that creates its ordinary and prequantized handlers. The three focused owner files +report four passes. + +Repository collection falls to 386 items and the static inventory to 381 definitions, with 1,590 curated decisions across +273 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-eighty-sixth wave: reunite reducer and RoPE parity owners + +The two-hundred-and-eighty-sixth wave first rejects a misleading mock-count heuristic. The remaining orchestrator runner, +client, and API-server reports protect different live transactions: rank-zero readiness and tensor-safe serialization, +API-engine ZMQ request handling, and public response metrics. Likewise, low apparent import counts for the model loader, +cautious decay, and Nemotron parallel plans come from package and model-owned call paths, not test-only APIs. Those +boundaries remain. + +Two genuinely duplicated loss reports are removed. The standalone importance-sampling and policy-loss files each rebuilt +the shared owner's masked tensors, global `TokenPartial` denominator, full-batch call, microbatch calls, and summable-metric +loop. Microbatch composition now runs beside legacy-identity coverage in the parameterized shared loss contract for basic, +KL, TIS, and IcePop modes. The copied one-test modules are gone. + +Dense-Qwen eager RoPE parity also no longer owns an isolated file. Its CUDA Q/K bitwise oracle against the serving +arithmetic now closes the existing frequency-table lifecycle, which already owns fp32 construction, lazy serving caches, +exact-model device recipes, and zero-K3 table bits. All four focused owner reports pass, including the CUDA parity path. + +Repository collection falls to 383 items and the static inventory to 378 definitions, with 1,593 curated decisions across +270 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-eighty-seventh wave: remove negative spies and close the handler owner + +The two-hundred-and-eighty-seventh wave removes a production method that existed only to support negative assertions. +`AdapterCoordinator.broadcast_adapter_optimizer_state` was a deprecated no-op with no runtime caller: topology-specific +optimizer shards already fail closed under rank-zero broadcast and restore through the all-ranks checkpoint path. The dead +method and seven lifecycle spies whose sole claim was that it stayed uncalled are gone. The real transactional optimizer +rejection remains in the adapter-coordinator lifecycle, which passes. + +Two one-test weight-sync reports also join the production owner they exercised. Trainer-side HCA selection, physical-GPU +mapping, abort-marker cleanup, and peer-status failure gathering now close the `WeightSyncHandler` configuration and sender +selection policy. The PP NCCL named-tensor codec's empty, sender, metadata, flattened-BF16, scalar, and receiver roundtrips +close the same handler owner. Both standalone private-method modules are removed; all three handler owner reports pass. + +Repository collection falls to 381 items and the static inventory to 376 definitions, with 1,595 curated decisions across +268 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-eighty-eighth wave: retire stale aliases and close complete construction + +The two-hundred-and-eighty-eighth wave removes two deprecated configuration aliases that had become compatibility-only +branches. No shipped configuration or current documentation uses `ep_outside` or `moe_checkpoint_method`; the native +`ep_intranode` and `gradient_checkpointing_method` fields own those policies. The aliases, parser remapping, compatibility +inputs, and redundant simulator dimension are gone. The active `gradient_checkpointing_method="moe_act"` execution mode +remains supported. + +GLM-5.2's standalone exact-attention constructor report also duplicated the admission matrix already exercised by the +complete exact-MoE constructor. Its unique 780 attention-factor names, projection classes, source FQNs, per-layer trainable +sets, and three dense roots now close the complete 1,700-factor inventory. The isolated rank/alpha, dense-component, +sparse-MLA, and all-to-all cases are removed; the retained owner and argument/simulator reports produce six focused passes. + +Repository collection falls to 380 items and the static inventory to 375 definitions, with 1,597 curated decisions across +267 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-eighty-ninth wave: remove algebra stand-ins for loader and backend owners + +The two-hundred-and-eighty-ninth wave removes a prequantized GNK-to-GKN report that never invoked the loader it claimed to +test. It re-derived transpose equivariance for block-FP8 and NVFP4 tensors using local helpers. The retained QLoRA expert +loader owner exercises the real `_load_experts` byte and scale transformations for both formats and multiple shapes, while +the codec owners retain their quantize/dequantize roundtrips. + +The standalone GKN-format report is also gone. It rebuilt a manual MoE and repeated `ExpertWeightBuffer` conversion already +exercised through the actual DeepSeek-V3 checkpoint handler, then repeated eager/native agreement owned by the backend +parity lifecycle. Grouped-GEMM and combined QuACK/SGLang parity continue to cover the Triton path without the removed +report's optional-import bypasses. The four retained owner reports produce four focused passes. + +Repository collection falls to 377 items and the static inventory to 372 definitions, with 1,599 curated decisions across +265 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-ninetieth wave: unify the Qwen3.5 norm-family contract + +The two-hundred-and-ninetieth wave replaces parallel dense and MoE Qwen3.5 RMSNorm reports with one family-wide owner. +Both files independently enumerated the same copied zero-centered dispatch matrix, v2 admission cases, per-site family +propagation, and layer/final residual selection. The unified CPU contract now drives both production classes through those +cases instead of maintaining two model-specific harnesses. + +The consolidation preserves the genuinely distinct boundaries: dense linear-attention GDN remains outside the ordinary +RMSNorm-family surface, both model constructors must propagate v2 to every zero-centered site, and the custom v2 plain and +residual backward paths still compare with autograd. One representative GPU lifecycle retains real-kernel module parity, +family-1 interpose bits, full MoE-layer parity, and the family-2 serving tree. The retained CPU report passes. + +Repository collection falls to 376 items and the static inventory to 371 definitions, with 1,600 curated decisions across +264 Python test files. The audit still surfaces the same 15 substantive conditional runtime gates, one intentional duplicate +group, and no parse errors. + +## Two-hundred-and-ninety-first wave: make the SGLang ABI boundary executable + +The two-hundred-and-ninety-first wave fixes the environment contract behind the exact-kernel reports. The default XoRL +profile remains on Torch 2.12.1 and contains no `sglang-kernel`; the combined profile and installation docs now match the +pinned SGLang tree at Torch 2.11.0 instead of the stale Torch 2.9.1 instructions. A separately ignored `.venv-sglang` was +materialized from that pin without changing the default environment. + +One test is intentionally added because wrapper imports were not a meaningful gate: SGLang loads its compiled extension +lazily. In optional mode the smoke skips when the wheel is absent from the default profile. With +`XORL_REQUIRE_SGL_KERNEL=1`, absence fails; the test eagerly imports `sgl_kernel`, `hash_topk`, and `LoRABatchInfo`, then +executes a real compiled RMSNorm operation. It skips in the default Torch-2.12 environment and passes in `.venv-sglang` +with Torch 2.11.0 and `sglang-kernel` 0.4.5. + +Repository collection rises deliberately to 377 items and the static inventory to 372 definitions, with 1,601 curated +decisions across 265 Python test files. The audit now surfaces 16 substantive conditional runtime gates, one intentional +duplicate group, and no parse errors. + +## Two-hundred-and-ninety-second wave: retire rollback-only server and kernel branches + +The two-hundred-and-ninety-second wave removes compatibility assertions together with the obsolete behavior that made +them necessary. SGLang fused experts now have one weight-presentation policy: the documented `WEIGHT_MODE` values own +strided, transient, and cached layouts. The test-only `CACHE_WEIGHTS` alias and its precedence cases are gone; cache reuse, +explicit invalidation, and all three actual layouts remain covered. + +Server EP dispatch likewise has one correct topology. The undocumented duplicate-batch rollback switch replicated the same +packed batch across every EP rank, causing `ep_size`-times redundant compute. Per-rank EP slices are now unconditional, while +the retained dispatcher and OPD owners still cover EP/CP slice identity, padding, routed payloads, and teacher-cache ordering. + +Two P2P compatibility branches are also retired. The disabled fused QKV slicer had no matching locator in pinned SGLang and +could only construct the wrong layout for its separate Q/K/V receivers; canonical locator slicing is now the only path. +Cold-prepare cache invalidation now uses the native `cache_invalidation_mode=none` opt-out instead of a second undocumented +environment alias. The fused-expert, dispatcher, OPD, handler-layout, P2P slicing, and prepare-lifecycle owner reports pass. + +Repository collection remains at 377 items and the static inventory at 372 definitions, with 1,604 curated decisions across +265 Python test files. The wave deletes compatibility-only assertion blocks inside retained transaction owners rather than +manufacturing a lower count by splitting or renaming them. The audit still surfaces 16 substantive conditional runtime +gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-ninety-third wave: delete duplicate controls and close native FP8 owners + +The two-hundred-and-ninety-third wave removes configuration branches whose tests existed only to preserve a second way to +select an already-owned behavior. Optimizer cache release now uses the native +`skip_empty_cache_after_optim_step` train field instead of an undocumented environment duplicate. Weight sync likewise has +one MoE bucket override, and both manual legacy-receiver post-process switches are gone. Direct P2P writes already target +receiver-native FP8 storage; endpoint requirements remain the owner of KV-cache finalization. + +Direct-EP scatter no longer offers debug-only shallow/deep locator-copy modes or their legacy boolean alias. Prepared +locators are immutable in this phase and `scatter_object_list` serializes recipient payloads, so the production default is +now the only behavior. The retained direct-EP manifest owner still proves rank filtering, dense ownership, endpoint state, +and locator identity through the real payload construction path. + +The unused NeMo `fp8_cfg` translation is also retired. No shipped XoRL configuration or documentation selected it, while +`enable_fp8_training` and the native `fp8_training_*` fields already own the supported contract. Dataclass aliases, +normalization/extraction APIs, launcher remapping, and compatibility assertions are removed; the shared tombstone rejects +the retired key with an explicit migration message. + +That removal makes the standalone FP8 compatibility report unnecessary. Its exhaustive external-knob rejection matrix was +already exercised through the public train and server parsers. The two real behaviors move to runtime owners: BF16 layer +islands execute through FP8 injection, and Blackwell admission executes through `build_training_model`. The focused native +FP8, argument, server, optimizer-step, handler, and P2P owners produce 23 passes. + +Repository collection falls to 375 items and the static inventory to 370 definitions, with 1,608 curated decisions across +264 Python test files. The audit still surfaces 16 substantive conditional runtime gates, one intentional duplicate group, +and no parse errors. + +## Two-hundred-and-ninety-fourth wave: make numerical programs structural + +The two-hundred-and-ninety-fourth wave removes three process-environment rollback paths whose assertions duplicated model +configuration. `XORL_FAMILIES_V2` and `SGLANG_FAMILIES_V2` could move trainer and sampler processes onto different reduction +families even though the exact model program already owns that choice. Ordinary models now use v2, exact Qwen3.5 selects +its qualified v1 norm and LM-head program, and canonical GLM-5.2 selects v2. + +The retained numerical reports no longer test environment kill switches. V1 RMSNorm and fused-LM-head owners select the +Qwen program explicitly; the v2 norm owner executes the default production dispatcher; model-builder reports prove GLM v2, +Qwen v1, and restoration to the ordinary program. The public LM-head contract now documents structural selection rather +than a coordinated environment rollback. + +`XORL_MOE_ROUTING_WEIGHTS_BEFORE_DOWN` is removed for the same reason. The native +`moe_routing_weights_before_down` model/server field already resolves auto, true, and false before model construction. Its +retained CPU oracle still proves both arithmetic positions against fp64, router-score gradients, no-router-gradient behavior, +dispatch regimes, and the SGLang parity exclusion; only the redundant lazy environment override assertion is gone. + +All 14 focused numerical-family, model-program, routing, and handler owners pass. Repository collection remains at 375 +items and the static inventory at 370 definitions, with 1,610 curated decisions across 264 Python test files. The audit still +surfaces 16 substantive conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Two-hundred-and-ninety-fifth wave: prefer lifecycle owners over private input matrices + +The two-hundred-and-ninety-fifth wave removes two standalone reports whose strongest behavior already had a production +lifecycle owner. QARL calibration remains exercised through `build_training_model`, including real calibration-data loading, +observer population, pre-parallelization ordering, and summary counts. Persistent calibration state is retained in the dense +fake-quant owner, so the deleted report's synthetic JSON/JSONL permutations and private bad-shape case no longer form a +separate contract. + +R3 Mooncake side payloads likewise remain covered at their actual boundaries. The request processor writes references in +packed order and cleans them on both success and failure; the runner dispatcher consumes only its rank-local slice. The +deleted fake-store report repeated the same codec and metadata flow in isolation, plus narrow malformed-metadata cases that +did not add a distinct server guarantee. + +DSV4 RoPE cache capacity now has one authority. `config.max_position_embeddings` owns allocation and the context-parallel +consumer still exercises the too-short-cache failure. The test/profiling-only `XORL_DSV4_ROPE_MAX_SEQ_LEN` override and its +setup assertions are removed. The tensor-collator report is also narrowed semantically: rather than hiding a large scalar, +boolean, string, dimensionality, and batch-size matrix inside one collected test, it now covers the four real pipeline forms +and their type boundaries with one representative sample each. + +All 10 focused QARL, R3, DSV4, and collator owners pass. The isolated Torch-2.11 SGLang ABI smoke also passes its real CUDA +operation. Repository collection falls to 373 items and the static inventory to 368 definitions, with 1,614 curated +decisions across 262 Python test files. The audit still surfaces 16 substantive conditional runtime gates, one intentional +duplicate group, and no parse errors. + +## Two-hundred-and-ninety-sixth wave: execute hidden kernel claims and remove fake contracts + +The two-hundred-and-ninety-sixth wave removes an unshipped SGLang EP slot-combine experiment. The default-off sub-flag had +no configuration or documentation owner, was scoring-only, and its assertions replaced both `moe_sum_reduce` and the +distributed return path with world-size-one fakes. The retained EP path uses the qualified serving-kernel expert compute and +the stock all-to-all combine; real extension loading and the independent FP32-routing backward oracle remain its owners. + +Repairing the isolated Torch-2.11 environment made another previously skipped claim executable. The GPU report asserted +bitwise gradients against stock Triton, but failed because stock Triton and the serving wrapper intentionally use different +routing-rounding programs. The wrapper preserves an FP32 routing boundary, so stock equality was not a valid oracle. That +test is deleted; the retained eager oracle proves local and EP input, routing, and weight gradients against the actual +serving program. + +Server batch preparation is consolidated at its lifecycle boundary. The deleted `test_batch_utils.py` directly invoked the +non-packed sharder with a one-row batch even though production routes that shape through `TextSequenceShardCollator`. +`RunnerDispatcher` now owns a real two-row ragged teacher-state conversion and CP shard, while the packed collator reports +retain their own path. Duplicate diagnostic environment aliases are removed in favor of request parameters, and an +untested environment-only minimal dummy constructor is removed in favor of the single zero-loss padding lifecycle. + +The API tensor report now covers valid rank-1 token IDs, rank-2 teacher states, and rank-3 routing payloads without promising +flat-list fallback semantics for malformed or zero-sized metadata. The isolated SGLang instructions and uv profile also pin +the Quack/CUTLASS versions required to import this XoRL trainer source. A fresh uv resolution succeeds, the real +`sgl_kernel` CUDA operation passes, and the FP32-routing plus DSV4 model, LoRA, and compressor owners pass under Torch 2.11. + +The final focused default-environment run produces 9 passes and one intentional optional-kernel skip; the isolated +environment produces 5 passes. Repository collection falls to 371 items and the static inventory to 366 definitions, with +1,620 curated decisions across 261 Python test files. The audit now surfaces 15 substantive conditional runtime gates, one +intentional duplicate group, and no parse errors. + +## Two-hundred-and-ninety-seventh wave: replace false-confidence mocks and assertions + +The two-hundred-and-ninety-seventh wave keeps `XORL_DCP_LOAD_NO_DIST` because it is a real shared-filesystem recovery mode, +but removes two fake-loader cases that merely echoed `process_group=None` and `no_dist=True`. The GLM exact-DCP owner still +performs a real round trip through that mode. The ModelState owner now also saves and loads a real model-only DCP while +proving that a requested optimizer is not fabricated or mutated; the separate pipeline case still protects custom-group +ordering. + +OPD metric finalization no longer guesses its loss family from key prefixes. Production always supplies the resolved +`loss_fn`; only a direct test omitted it. The private finalizer now requires the production signature, and its retained +aggregation, extrema, loss-group reduction, and empty-rank collective-shape policies pass with an explicit `opd_loss`. + +Dataset preparation loses two false contracts. `shards` plus `preprocess_shards` was documented as invalid and absent from +shipped configurations, yet a test promised silent precedence. The input is now rejected, while preprocessing expansion +consumes `preprocess_shards` into concrete shard coordinates. The retry lifecycle drops a redundant immediate-success row; +its transient success already proves return propagation alongside backoff, exhaustion, Hub errors, and unrelated failures. + +The remaining dataset assertions now test outcomes rather than activity. Merge coverage distinguishes ordered +concatenation, a whole-dataset permutation, and per-dataset permutations instead of checking only row count. Downloaded +`data_files` now honor the documented `ds_type`; JSON-string and Parquet-list cases verify both format routing and resolved +files, closing a bug the former download-count assertions could not detect. + +All three focused checkpoint, data-preparation, and OPD policies pass. Full test-tree Ruff, scoped formatting, compileall, +decision-JSON validation, public-tree lint, diff whitespace, global collection, and the static audit pass. Repository +collection remains 371 items and the static inventory remains 366 definitions, with 1,626 curated decisions across 261 +Python test files. The audit still surfaces 15 substantive conditional runtime gates, one intentional duplicate group, and +no parse errors. + +## Two-hundred-and-ninety-eighth wave: make the isolated kernel gate fail closed + +The two-hundred-and-ninety-eighth wave finishes the SGLang dependency boundary rather than treating lazy wrapper imports as +coverage. The default profile remains on Torch 2.12.1 without `sglang-kernel`; the exact SGLang lane remains isolated in +`.venv-sglang` with Torch 2.11.0 and the pinned `sglang-kernel` 0.4.5 wheel. The last stale server-training instructions no +longer install SGLang into the active default environment or advertise the obsolete Torch 2.9.1 contract. + +Required smoke mode is now literal. Import-loader failures from either Python or the dynamic linker fail with the active +Torch version, and `XORL_REQUIRE_SGL_KERNEL=1` cannot turn missing CUDA into a successful skip. It must import `sgl_kernel`, +`hash_topk`, and `LoRABatchInfo` and execute the compiled RMSNorm operation. The default environment intentionally skips the +absent optional wheel; the Torch-2.11 environment passes the real CUDA operation. + +Repository collection remains 371 items and the static inventory remains 366 definitions, with 1,627 curated decisions +across 261 Python test files. This wave strengthens one retained transaction without adding a new collected test. + +## Two-hundred-and-ninety-ninth wave: replace primitive matrices with production transactions + +The two-hundred-and-ninety-ninth wave removes both standalone block-FP8 codec reports. They enumerated dimensionality, +random scale ranges, block sizes, determinism, tiny and large constants, signs, and internal assertion failures without +entering a model operation. The retained FP8-linear owner already runs activation and tiled-weight quantize/dequantize at +block sizes 64 and 128, consumes their scale layouts in a real GEMM, and covers a non-divisible weight row count. It now +also compares both dequantized operands with their originals, preserving the useful independent numerical oracle. Five +unexported, unreferenced compatibility aliases disappear with the obsolete primitive vocabulary. + +Two test-created implementation contracts are also retired. The real two-GPU sequence-parallel metric report was the only +caller using an obsolete positional order, which forced production to detect a dictionary and swap its arguments. It now +uses the same declared signature as both runtime callers and passes its NCCL partial-sum and extrema transaction. The GLM +absorbed-attention report no longer replaces projections and rotary execution merely to assert two private `einsum` +decompositions; the full-model sparse-versus-dense owner already executes both generic projection directions numerically, +while the retained exact-kv_b owner protects factor-only routing and non-materialization. + +Focused GLM and FP8 owners pass, including the real CUDA GEMM and two-GPU NCCL collective. Repository collection falls to +368 items and the static inventory to 363 definitions, with 1,630 curated decisions across 259 Python test files. The audit +still surfaces 15 substantive conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Three-hundredth wave: move NF4 coverage into shipped QLoRA lifecycles + +The three-hundredth wave removes the standalone NF4 codec report. Its flat and GKN sections enumerated shapes, dtypes, +zeros, codebook constants, three group sizes, and large random tensors without entering a shipped module. Production NF4 +QLoRA admits group size 64, not that synthetic matrix. The report also embedded bandwidth benchmarks whose missed target +became a successful skip, so they could neither enforce correctness nor act as a performance gate. + +The useful reconstruction oracle now lives at both real consumers. The QLoRA linear owner includes NF4 in its quantized +storage, forward, backward, and memory transaction, then compares the dequantized flat weight with its original operand. +The distributed expert owner compares every production-shaped GKN NF4 projection with its original base before running the +real Triton EP path through two optimizer steps. The flat CUDA lifecycle and direct two-GPU NCCL transaction both pass. + +The same size-ordered review explicitly keeps three nearby reports. `ToTensorCollator` uniquely owns conversion and +structural preservation across flat, batched, nested, and empty pipeline forms. The QARL GPU transaction uniquely proves +that fake-quantized expert Parameters reach the real Triton grouped GEMM and receive gradients. The local Qwen3.5 and Kimi +registry reports load actual config and tokenizer artifacts, covering conversion behavior that direct tiny-model builders +do not exercise. + +Repository collection falls to 367 items and the static inventory to 362 definitions, with 1,634 curated decisions across +258 Python test files. The audit still surfaces 15 substantive conditional runtime gates, one intentional duplicate group, +and no parse errors. + +## Three-hundred-and-first wave: keep protocol and security transactions intact + +The three-hundred-and-first wave audits the smallest server and API reports plus the remaining compatibility-labeled +branches. It intentionally removes nothing. API metric mapping, rank-zero wire serialization and readiness, scheduler state +transitions, path containment, SSRF and DNS pinning, checkpoint tenant isolation, and weight-sync receiver fencing are all +distinct runtime boundaries. Their process and socket doubles isolate external systems without replacing the payload, +serialization, queue, filesystem, or byte-layout behavior under test. + +The legacy labels also describe live inputs rather than obsolete test fixtures. DRGRPO accepts both current +`old_logprobs` and rollout `logprobs`; Tinker session payloads still map into current model and optimizer controls; supported +checkpoint URI and on-disk forms remain public while unsafe pickle-backed optimizer state fails closed. Removing those +cases would silently narrow compatibility or security guarantees. + +The adjacent numerical reports remain independent as well. The DSV4 compressor owner uniquely exercises +context-parallel offsets, cache capacity, C4 overlap admission, and the CPU Hadamard fallback. The batch-invariant GEMM table +is checked numerically against the pinned reduction tree and across batch buckets rather than by asserting literal table +entries. Repository collection therefore remains at 367 items and the static inventory at 362 definitions, with 1,637 +curated decisions across 258 Python test files. + +## Three-hundred-and-second wave: replace optimizer fakes with real transactions + +The three-hundred-and-second wave removes two narrow reports from optimizer and training instrumentation. The standalone +Gram Newton-Schulz CUDA test replaced every backend operation with logging fakes and asserted only the dtypes observed by +those fakes. Its real requirement now lives in the two-GPU full-gradient Muon owner: the shipped CUDA orthogonalizer must +match an independent FP32 Newton-Schulz program exactly, then the distributed update must match the single-rank optimizer +oracle. + +The per-component timer loses a disabled no-op test that manually inserted fake objects into private event-pair +dictionaries. Its `last_skipped_event_pair_count` production attribute had no runtime consumer and existed solely for that +assertion, so it is removed. Invalid CUDA event pairs remain safely ignored, while the retained real CUDA lifecycle attaches +hooks to GLM-style and Qwen-style decoder layers and records their present forward and backward phases. + +The remaining optimizer reports survive semantic review. Schedule endpoints, DistSignSGD sign and collective ordering, +cautious-decay math, standard batched Newton-Schulz layout, chunked CE parity, token-voter counting, and explicit gradient +synchronization all have independent numerical or collective oracles that end-to-end liveness cannot replace. Focused +owners pass, including real CUDA timer hooks and the strengthened two-GPU Muon transaction. Repository collection falls to +365 items and the static inventory to 360 definitions, with 1,641 curated decisions across 258 Python test files. + +The learning-rate trace is also brought back to the production lifecycle. It previously advanced `LambdaLR` without an +optimizer step, emitting PyTorch's skipped-first-value warning even though the trainer always steps the optimizer first. +Single and multi-optimizer traces now follow that real ordering and retain the same exact warmup, decay, and floor oracles +without warnings. + +## Three-hundred-and-third wave: make artifact lifecycles own helper semantics + +The three-hundred-and-third wave removes synthetic helper coverage from the quantized exporter. String size parsing and +direct-function sharding were separate setup reports even though the module CLI is the public transaction. The retained +subprocess invocation now consumes a `24B` YAML shard limit, writes and reconciles a multi-shard safetensors index, reloads +every emitted shard, and verifies the converted and preserved tensor layouts. This proves the configured string reaches the +artifact writer while deleting the private parser examples and duplicate direct sharding fixture. + +Model-support aggregates lose three dominated implementation reports. GLM's bare-config constants duplicated its real +official-shaped local-config load, and its monkeypatched FP32-indexer dispatch was a subset of the canonical GLM-5.2 router +and indexer contract. MiniMax's three private expert-key aliases were already a strict subset of the centralized checkpoint +key classifier. The retained owners still cover GLM sparse/dense and Hugging Face logit parity, MiniMax forward/backward and +checkpoint/EP behavior, and Qwen2/OLMo2 Hugging Face checkpoint conversion plus numerical parity. + +The focused exporter and model-support run passes all 6 collected transactions. Scoped Ruff, formatting, compileall, +decision-JSON validation, diff whitespace, global collection, and the static audit pass. Repository collection remains 365 +items and the static inventory remains 360 definitions across 258 Python test files, now with 1,644 curated decisions. The +audit still surfaces 15 substantive conditional runtime gates, one intentional duplicate group, and no parse errors. + +## Three-hundred-and-fourth wave: retire test-created configuration and rollback surfaces + +The three-hundred-and-fourth wave follows typed configuration values into distributed execution instead of preserving a +private conversion vocabulary. FSDP prefetch inputs are booleans at the argument and model-builder boundaries, so their +private parser no longer accepts integer values or yes/no/on/off strings solely for a direct helper matrix. FSDP2 now fails +fast on non-booleans while retaining the forward default, backward inheritance, CP-folding admission, and exact neighbor +prefetch directions. + +The optimizer cache field is narrowed the same way. Its lifecycle already supplies a boolean and executes both cache +policies, so the runner no longer interprets arbitrary objects or undocumented truthy strings. The session, optimizer, +P2P-async, checkpoint-broadcast, and adapter-resume reports survive review because they cross real state, transport, +artifact, collective, or failure boundaries; their doubles do not create the behavior being asserted. + +MoE slot assignment loses a larger rollback-only surface. `XORL_MOE_DETERMINISTIC_SCATTER` was not exposed by shipped +arguments, examples, or documentation and existed only to restore a relaxed-atomic kernel whose within-expert order varies +with CTA scheduling. The old kernel, environment parser, false-value alias matrix, and mocked route checks are removed. +Stable sorting is now the sole production program, while the retained CUDA transaction independently proves exact stable +order, full slot coverage, per-expert cumsum regions, run invariance, and both routing integer widths. + +The focused distributed, runner, and CUDA MoE run produces 6 passes. Repository collection remains at 365 items and the +static inventory remains at 360 definitions across 258 Python test files, now with 1,650 curated decisions. This wave +reduces semantic compatibility and implementation surface inside retained transaction owners rather than manufacturing a +lower collected count. + +## Three-hundred-and-fifth wave: finish on authoritative runtime contracts + +The final wave re-audits all fifteen conditional-runtime candidates and keeps them. They are real ABI, CUDA, distributed, +or optional-backend gates: the SGLang owner executes a compiled operation in its isolated Torch 2.11 environment; FlashQLA +compares a two-rank context-parallel program with a local reference; GLM exact kernels compose with real FSDP ownership; +and the DeepGEMM and SGLang MoE reports exercise documented production dispatch with numerical or gradient oracles. Their +conditional admission reflects unavailable hardware or dependencies, not a missing outcome. + +Three remaining test-maintained compatibility surfaces are removed. Routing-weight position accepts its declared boolean +and `auto`/`true`/`false` forms without undocumented numeric-string aliases. `AdapterState.local_params` is now the sole +parameter store after deleting the deprecated `lora_params` property and test fakes that reproduced it. EP replicated- +gradient synchronization no longer falls back to the broader replicated group: the production classifier always emits +`ep_replicated_gradient_sync`, so missing authoritative metadata now fails fast instead of silently changing the reduction +domain. + +The retained owners continue to exercise automatic and explicit routing selection, adapter optimizer and checkpoint +lifecycles, and real multi-rank replicated-gradient coalescing, missing-gradient materialization, nonfinite rejection, and +clipping. Repository collection remains 365 items and the static inventory remains 360 definitions across 258 Python test +files, now with 1,654 curated decisions. The audit closes with fifteen substantive conditional-runtime gates, one +intentional duplicate group, and no parse errors. diff --git a/certification/benchmark_vocab_parallel_ce.py b/certification/benchmark_vocab_parallel_ce.py new file mode 100755 index 00000000..aa0459c5 --- /dev/null +++ b/certification/benchmark_vocab_parallel_ce.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +"""Benchmark eager and compiled vocab-parallel CE outside the pytest suite. + +Run with: + PYTHONPATH=src torchrun --nproc_per_node=2 certification/benchmark_vocab_parallel_ce.py +""" + +from __future__ import annotations + +import argparse +import json +import time + +import torch +import torch.distributed as dist + +from xorl.ops.loss.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy + + +def _benchmark_once(hidden, weight, labels, *, compiled, iterations, backward): + for _ in range(5): + loss = vocab_parallel_cross_entropy(hidden, weight, labels, dist.group.WORLD, use_compile=compiled) + if backward: + loss.sum().backward() + hidden.grad = None + weight.grad = None + torch.cuda.synchronize() + + torch.cuda.reset_peak_memory_stats() + memory_before = torch.cuda.memory_allocated() + start = time.perf_counter() + for _ in range(iterations): + loss = vocab_parallel_cross_entropy(hidden, weight, labels, dist.group.WORLD, use_compile=compiled) + if backward: + loss.sum().backward() + hidden.grad = None + weight.grad = None + torch.cuda.synchronize() + elapsed_ms = (time.perf_counter() - start) / iterations * 1000 + peak_memory = torch.cuda.max_memory_allocated() + return { + "milliseconds": elapsed_ms, + "peak_activation_mb": (peak_memory - memory_before) / 1024**2, + "peak_total_mb": peak_memory / 1024**2, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--iterations", type=int, default=20) + args = parser.parse_args() + + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(rank) + + torch.manual_seed(42) + tokens, hidden_size, vocabulary = 4096, 4096, 152064 + local_vocabulary = vocabulary // world_size + hidden = torch.randn(tokens, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(local_vocabulary, hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True) + labels = torch.randint(0, vocabulary, (tokens,), device="cuda") + + results = {} + for backward in (False, True): + phase = "forward_backward" if backward else "forward" + results[phase] = {} + for compiled in (False, True): + mode = "compiled" if compiled else "eager" + results[phase][mode] = _benchmark_once( + hidden, + weight, + labels, + compiled=compiled, + iterations=args.iterations, + backward=backward, + ) + dist.barrier() + + if rank == 0: + print(json.dumps(results, indent=2, sort_keys=True)) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/distributed/test_deepep_correctness.py b/certification/deepep/compare_dispatch_alltoall.py similarity index 97% rename from tests/distributed/test_deepep_correctness.py rename to certification/deepep/compare_dispatch_alltoall.py index 91c88fe7..656f1a2e 100644 --- a/tests/distributed/test_deepep_correctness.py +++ b/certification/deepep/compare_dispatch_alltoall.py @@ -4,17 +4,17 @@ for both the forward pass and input gradients (backward pass). Usage (single node, EP=8): - torchrun --nproc_per_node=8 tests/distributed/test_deepep_correctness.py + torchrun --nproc_per_node=8 certification/deepep/compare_dispatch_alltoall.py Usage (2 nodes, EP=16): # Node 0: torchrun --nnodes=2 --nproc_per_node=8 --node_rank=0 \ --master_addr= --master_port=29500 \ - tests/distributed/test_deepep_correctness.py + certification/deepep/compare_dispatch_alltoall.py # Node 1: torchrun --nnodes=2 --nproc_per_node=8 --node_rank=1 \ --master_addr= --master_port=29500 \ - tests/distributed/test_deepep_correctness.py + certification/deepep/compare_dispatch_alltoall.py """ import os diff --git a/certification/glm52/benchmark_sparse_mla_backward.py b/certification/glm52/benchmark_sparse_mla_backward.py new file mode 100755 index 00000000..28ad6696 --- /dev/null +++ b/certification/glm52/benchmark_sparse_mla_backward.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +"""Manually certify combined versus split GLM-5 sparse-MLA backward speed.""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics + +import torch + + +def _make_inputs(sequence, kv_sequence, heads, rank, tail, topk): + generator = torch.Generator(device="cuda").manual_seed(1234) + query = torch.randn((sequence, heads, rank + tail), device="cuda", dtype=torch.bfloat16, generator=generator) + kv = torch.randn((kv_sequence, 1, rank + tail), device="cuda", dtype=torch.bfloat16, generator=generator) + relative = torch.arange(topk, device="cuda", dtype=torch.int64) + query_positions = torch.arange(kv_sequence - sequence, kv_sequence, device="cuda", dtype=torch.int64) + indices = query_positions.unsqueeze(1) - (topk - 1 - relative).unsqueeze(0) + indices = indices.clamp(min=-1, max=kv_sequence - 1).to(torch.int32).unsqueeze(1) + return query, kv, indices + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--minimum-speedup", type=float, default=0.15) + args = parser.parse_args() + + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9: + raise SystemExit("This certification requires an H100-class CUDA device") + try: + import tilelang # noqa: F401, PLC0415 + except ImportError as exc: + raise SystemExit("This certification requires TileLang") from exc + + from xorl.ops.glm5_kernels.sparse_mla import SparseMLA # noqa: PLC0415 + + sequence, kv_sequence, heads, rank, tail, topk = 2048, 32768, 64, 512, 64, 2048 + scale = (rank + tail) ** -0.5 + query, kv, indices = _make_inputs(sequence, kv_sequence, heads, rank, tail, topk) + generator = torch.Generator(device="cuda").manual_seed(7) + grad_output = torch.randn((sequence, heads, rank), device="cuda", dtype=torch.bfloat16, generator=generator) + + def time_backward() -> float: + local_query = query.detach().clone().requires_grad_(True) + local_kv = kv.detach().clone().requires_grad_(True) + output, _ = SparseMLA.apply(local_query, local_kv, indices, scale) + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + output.backward(grad_output) + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) + + previous = os.environ.get("XORL_GLM5_SPLIT_SPARSE_MLA_BWD") + try: + timings = {} + for label, mode in (("combined", "0"), ("split", "1")): + os.environ["XORL_GLM5_SPLIT_SPARSE_MLA_BWD"] = mode + for _ in range(args.warmup): + time_backward() + timings[label] = [time_backward() for _ in range(args.trials)] + finally: + if previous is None: + os.environ.pop("XORL_GLM5_SPLIT_SPARSE_MLA_BWD", None) + else: + os.environ["XORL_GLM5_SPLIT_SPARSE_MLA_BWD"] = previous + + combined_ms = statistics.median(timings["combined"]) + split_ms = statistics.median(timings["split"]) + speedup = 1.0 - combined_ms / split_ms + result = { + "combined_ms": combined_ms, + "split_ms": split_ms, + "speedup": speedup, + "minimum_speedup": args.minimum_speedup, + "timings_ms": timings, + } + print(json.dumps(result, indent=2, sort_keys=True)) + if speedup < args.minimum_speedup: + raise SystemExit("Combined sparse-MLA backward did not meet the requested speedup") + + +if __name__ == "__main__": + main() diff --git a/tests/models/test_glm52_official_fp8_inventory.py b/certification/glm52/test_official_fp8_inventory.py similarity index 96% rename from tests/models/test_glm52_official_fp8_inventory.py rename to certification/glm52/test_official_fp8_inventory.py index c4631b98..9b51e54d 100644 --- a/tests/models/test_glm52_official_fp8_inventory.py +++ b/certification/glm52/test_official_fp8_inventory.py @@ -1,3 +1,9 @@ +"""Opt-in certification against an official GLM-5.2 checkpoint inventory. + +Run explicitly with ``pytest certification/glm52/test_official_fp8_inventory.py`` +after setting ``XORL_GLM52_OFFICIAL_MODEL_PATH``. +""" + import hashlib import json import os diff --git a/tests/ops/loss/test_vp_kl_gathered.py b/certification/opd/vocab_parallel_kl_gathered.py similarity index 99% rename from tests/ops/loss/test_vp_kl_gathered.py rename to certification/opd/vocab_parallel_kl_gathered.py index 1b236bc0..fb18d648 100755 --- a/tests/ops/loss/test_vp_kl_gathered.py +++ b/certification/opd/vocab_parallel_kl_gathered.py @@ -10,7 +10,7 @@ - this rank's weight-shard grad == reference grad for its vocab shard. Run (no GPU needed): - python -m pytest tests/ops/loss/test_vp_kl_gathered.py + PYTHONPATH=src .venv/bin/python certification/opd/vocab_parallel_kl_gathered.py """ from __future__ import annotations diff --git a/tests/e2e/qwen3_30b/compare_lora_qlora.py b/certification/qwen3_30b/compare_lora_qlora.py similarity index 97% rename from tests/e2e/qwen3_30b/compare_lora_qlora.py rename to certification/qwen3_30b/compare_lora_qlora.py index 6b30a7a9..f1e2c404 100755 --- a/tests/e2e/qwen3_30b/compare_lora_qlora.py +++ b/certification/qwen3_30b/compare_lora_qlora.py @@ -11,10 +11,10 @@ import tempfile -ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) sys.path.insert(0, ROOT) -from tests.e2e.e2e_utils import ( +from tests.e2e.e2e_utils import ( # noqa: E402 generate_training_config, run_training, ) diff --git a/docs/k3/LM_HEAD_CONTRACT.md b/docs/k3/LM_HEAD_CONTRACT.md index b295b178..ac81b07e 100644 --- a/docs/k3/LM_HEAD_CONTRACT.md +++ b/docs/k3/LM_HEAD_CONTRACT.md @@ -19,11 +19,11 @@ launch. `HEAD_V2_BLOCK_K` fixes the projection's accumulation order and those tile statistics in one explicit pairwise tree. Changing either constant requires a new cross-engine bitwise gate. -The earlier `bi_lm_head_selected_logprob` path remains available with -`XORL_FAMILIES_V2=0` (equivalently `SGLANG_FAMILIES_V2=0`), the same switch -that rolls back the redefined norm trees. One setting moves both engines, -because the trainer and the sampler have to evaluate the same trees. -It materializes one vocabulary chunk at a time, records +The earlier `bi_lm_head_selected_logprob` path remains the structurally selected +v1 program for exact Qwen3.5-family models. Numerical-family selection belongs +to the model program rather than a process-environment rollback switch, so the +trainer and sampler cannot be moved independently. The v1 path materializes one +vocabulary chunk at a time, records the same maximum, exponential sum, and selected logit, then merges chunks in pinned order. This rollback is exact but uses more launches. diff --git a/docs/k3/LORA_CONTRACT.md b/docs/k3/LORA_CONTRACT.md index 12f97081..153fbe2f 100644 --- a/docs/k3/LORA_CONTRACT.md +++ b/docs/k3/LORA_CONTRACT.md @@ -27,6 +27,18 @@ autograd differentiates the fold into the trainable factors. The resolver owns this choice through module state; there is no public `XORL_LORA_MERGED_FORWARD` launch flag on the architecture-selected path. +Qwen3.5/3.6 training defaults keep the GDN input as four +independent rank-r adapters: `q_proj`, `k_proj`, `v_proj`, and `g_proj` (the +trainer name for serving's `in_proj_z`). The optional River/SGLang-shaped +`in_proj_qkvz` adapter remains available only when explicitly targeted. + +On Qwen3.6 MoE layers, `gate_proj`, `up_proj`, and `down_proj` also cover every +shared expert. Gate and up own independent factors even though their frozen +base remains one fused `gate_up_proj` tensor. The trainer folds those logical +factors into the fused GEMM, the ordered EP combine consumes the same folded +weights, and weight synchronization publishes only the corresponding fused +base-weight bytes. Gate and up must therefore be selected together. + ## GLM-5.2: native-FP8 base plus active rank-1 LoRA Folding adapters into GLM's native-FP8 expert base would require FP8 @@ -56,6 +68,8 @@ forward contract and replay gate. ```bash pytest tests/models/test_lora_merged_forward.py -q +pytest tests/models/test_qwen35_lora_projection_topology.py -q +pytest tests/e2e/qwen3_5/test_lora_projection_topology.py -q # one GPU pytest tests/models/test_glm52_exact_qlora.py -q pytest tests/models/test_glm52_exact_gate_up_qlora.py -q pytest tests/models/test_glm52_exact_shared_expert_qlora.py -q diff --git a/docs/src/content/docs/adapters/lora.mdx b/docs/src/content/docs/adapters/lora.mdx index 9fe815ae..84875592 100644 --- a/docs/src/content/docs/adapters/lora.mdx +++ b/docs/src/content/docs/adapters/lora.mdx @@ -106,11 +106,70 @@ lora: save_lora_only: true # checkpoint only LoRA weights, not base model ``` -For QKV-fused models (default `merge_qkv: true`), target the fused projection names: +## Fused projections + +Many architectures store input-side projections fused: attention as `qkv_proj` and +the dense or shared MLP as `gate_up_proj`. xorl audits the requested target set and +fails injection if any target is unmatched. + +For supported fused modules, split target names automatically create independent +logical adapters while retaining the fused base projection. Dynamic LoRA adds each +adapter to its corresponding output slice. Exact merged-forward mode canonically +folds each adapter into that slice and still runs one fused GEMM. Checkpoint and +weight-sync keys remain the standard split names: + +```yaml + lora_target_modules: [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj] +``` + +This is the default path for Qwen2, Qwen3, Qwen3 MoE, Qwen3.5, Llama, OLMo2, and +GLM-4 MoE projections implemented by xorl. + +**Target the fused names.** Adapts each fused module as one target, at no throughput +cost. One `lora_A` is shared across the halves, so this is a more constrained +parameterization than two independent adapters, and the exported adapter carries fused +key names — which do not line up with a HuggingFace base that stores them split: + ```yaml lora_target_modules: [qkv_proj, gate_up_proj, down_proj, o_proj] ``` +**Unfuse before injection.** This explicit fallback is for an architecture whose +fused modules do not implement logical adapters, such as GPT-OSS attention. It splits +the modules into real `q_proj`/`k_proj`/`v_proj` and +`gate_proj`/`up_proj` before injection. It requires `enable_lora` and is rejected +with `enable_qlora`: + +```yaml +lora: + enable_lora: true + unfuse_for_lora: true + lora_target_modules: [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj] +``` + +Unfusing changes both performance and floating-point operation order. Requalify +trainer-to-sampler equality before using it in a workflow that depends on exact +logprobs. + +The layout is fixed for a run's lifetime. Toggling the flag between runs renames +`mlp.gate_up_proj.*` to `mlp.gate_proj.*` / `mlp.up_proj.*`, which checkpoint +validation rejects with "Checkpoint incompatible with model" — resuming requires the +same setting the checkpoint was written with. + +### Audited architecture defaults + +| Family | Default plain-LoRA scope | Projection handling | +|---|---|---| +| Qwen2, Qwen3, Qwen3 MoE, Llama, OLMo2, GLM-4 MoE | Attention and MLP | Independent logical adapters retain fused qkv and gate/up bases | +| Qwen3.5 dense and MoE | Attention, GDN `g_proj`, MLP, routed experts | Separate GDN/full-attention projections plus fused-base logical gate/up adapters | +| DeepSeek V3, Kimi K2/K2.5, GLM-5 | MLA attention and MLP | Architecture-specific separate projections and expert adapters | +| DeepSeek V4 | Attention only | Audited MLA/DSA projection names; routed-expert LoRA is not a generic default | +| GPT-OSS | Attention only | Set `unfuse_for_lora: true` for split q/k/v targets | +| MiniMax M3, Nemotron-H | Attention only | Separate attention projections; expert LoRA is not a generic default | + +Unknown model families have no guessed default. Set `lora_target_modules` explicitly +after auditing their projection and expert semantics. + ## Key Parameters | Parameter | Default | Description | @@ -119,6 +178,7 @@ For QKV-fused models (default `merge_qkv: true`), target the fused projection na | `lora_rank` | `16` | Rank r of the low-rank decomposition | | `lora_alpha` | `32` | Scaling factor; effective LR scale = alpha/rank | | `lora_target_modules` | `null` | List of module name patterns to inject LoRA into | +| `unfuse_for_lora` | `false` | Explicitly split supported fused projections before injection; use only when fused-base logical adapters are unavailable | | `save_lora_only` | `false` | Save only LoRA weights in checkpoints | ## LoRA with MoE Models @@ -130,8 +190,8 @@ lora: enable_lora: true lora_rank: 16 lora_alpha: 32 - lora_target_modules: [qkv_proj, gate_up_proj, down_proj, o_proj] - # Expert layers are included automatically when targeting gate_up_proj/down_proj + lora_target_modules: [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj] + # gate/up/down also select routed experts on supported MoE families ``` ## Checkpoint Behavior diff --git a/docs/src/content/docs/config-reference/local.md b/docs/src/content/docs/config-reference/local.md index bcfd4332..6ea437c1 100644 --- a/docs/src/content/docs/config-reference/local.md +++ b/docs/src/content/docs/config-reference/local.md @@ -37,7 +37,7 @@ torchrun --nproc_per_node=8 -m xorl.cli.train config.yaml \ | `deepep_num_sms` | `20` | SMs assigned to DeepEP communication kernels. Must be even. Lower values leave more SMs for overlapped compute. | | `deepep_async_combine` | `false` | Overlap DeepEP combine with the next layer's compute (experimental, unsafe). Forced to `false` in code unless `XORL_DEEPEP_UNSAFE_ASYNC_COMBINE=1` is exported; without that env var, deferring the comm-stream sync races the transformer block's read of the combined tensor on the default stream. | | `alltoall_combine_hidden_chunk_size` | `0` | Hidden-dimension chunk size for all-to-all EP combine. `0` disables chunking; use a positive value to reduce long-context MoE combine memory peaks. | -| `merge_qkv` | `true` | Keep Q/K/V projections fused as `qkv_proj`. Set `false` for tensor parallelism or per-projection LoRA. | +| `merge_qkv` | `true` | Keep Q/K/V projections fused as `qkv_proj`. Set `false` when required by tensor parallelism; audited LoRA families support independent q/k/v adapters over the fused base. | | `basic_modules` | `[]` | Additional module names (beyond `_no_split_modules`) to shard as separate FSDP units. | | `foundation` | `{}` | Extra foundation model config (dict). | | `encoders` | `{}` | Multimodal encoder configs, keyed by type (`image`, `video`, `audio`). Each value must have `model_path` and optionally `config_path`. | @@ -232,6 +232,7 @@ Each entry in `datasets` (or `test_datasets`) is a dict: | `lora_rank` | `16` | LoRA rank (`r`). | | `lora_alpha` | `16` | LoRA scaling factor (`alpha`). Effective scale = `alpha / rank`. | | `lora_target_modules` | `null` | Module names to inject LoRA into. `null` = default linear projections for the architecture. | +| `unfuse_for_lora` | `false` | Explicitly split supported fused projections before LoRA injection. Requires `enable_lora`; rejected with `enable_qlora`. Audited common families retain fused bases automatically; use this fallback only when logical fused-base adapters are unavailable. See [LoRA](/adapters/lora/#fused-projections). | | `save_lora_only` | `false` | Only save LoRA adapter weights in HF checkpoints (not the full model). | | `enable_qlora` | `false` | Quantize base weights and train LoRA on top. Implies `enable_lora: true`. | | `quant_format` | `nvfp4` | Quantization format: `nvfp4` (4-bit, Hopper+), `block_fp8` (8-bit blocks), `nf4` (4-bit normal float). | diff --git a/docs/src/content/docs/getting-started/installation.md b/docs/src/content/docs/getting-started/installation.md index 02c80ab9..38c0e15d 100644 --- a/docs/src/content/docs/getting-started/installation.md +++ b/docs/src/content/docs/getting-started/installation.md @@ -51,14 +51,25 @@ The repo ships two git submodules under `submodules/`: | [xorl-client](https://github.com/togethercomputer/xorl-client) | Lightweight Python client for the XoRL training service. Required for server/RL training mode. | | [xorl-sglang](https://github.com/togethercomputer/xorl-sglang) | XoRL's fork of [SGLang](https://github.com/sgl-project/sglang). Used as the inference engine in online RL loops. | -Install individually: +Install the client in the default environment. Keep SGLang in an isolated +Torch-2.11 environment so its compiled kernel wheel never enters the default +Torch-2.12 profile: ```bash pip install -e submodules/xorl-client -pip install -e "submodules/xorl-sglang/python[all]" +uv venv .venv-sglang --python 3.12 +uv pip install --python .venv-sglang/bin/python -e submodules/xorl-sglang/python +uv pip install --python .venv-sglang/bin/python \ + torchdata==0.11.0 nvidia-cutlass-dsl==4.5.2 quack-kernels==0.5.0 +uv pip install --python .venv-sglang/bin/python --no-deps -e . +uv pip install --python .venv-sglang/bin/python pytest +PYTHONPATH=src:submodules/xorl-sglang/python XORL_REQUIRE_SGL_KERNEL=1 \ + .venv-sglang/bin/python -m pytest -q tests/ops/test_sgl_kernel_smoke.py ``` -Alternatively, use the bundled `pyproject.sglang.toml` which pins PyTorch to 2.9.1 (required by sglang) and installs xorl, xorl-client, and xorl-sglang together: +The bundled `pyproject.sglang.toml` provides the same combined profile for uv. +Its dependency overrides retain the Quack/CUTLASS versions required by XoRL's +trainer imports while the exact-kernel smoke validates the SGLang boundary. **uv:** ```bash @@ -67,15 +78,7 @@ uv sync source .venv/bin/activate ``` -**conda:** -```bash -conda create -n xorl-sglang python=3.12 -conda activate xorl-sglang -cp pyproject.sglang.toml pyproject.toml -pip install -e . -``` - -> **Note:** The default `pyproject.toml` uses PyTorch 2.10.0. sglang requires PyTorch 2.9.1, so the two cannot coexist in the same environment unless you use `pyproject.sglang.toml`. +> **Note:** The default `pyproject.toml` uses Torch 2.12.1. Pinned SGLang requires Torch 2.11.0; do not install `sglang-kernel` into the default environment. > These submodules are only needed for **server training / online RL**. If you are only running local SFT or pretraining, you can skip this step. @@ -84,9 +87,9 @@ pip install -e . | Package | Version | Notes | |---|---|---| -| PyTorch | 2.10.0+cu129 | CUDA 12.9 build | -| Flash Attention 3 | custom | FA3 + FA4 wheels | -| Triton | 3.6.0 | MoE fused kernels | +| PyTorch | 2.12.1 | Default XoRL profile; SGLang profile uses 2.11.0 | +| Flash Attention 4 | pinned | Selected by each Torch profile | +| Triton | 3.7.1 | Default profile; SGLang profile uses 3.6.0 | | Transformers | 5.0+ | Model loading | | FastAPI + uvicorn | latest | Server training API | | pyzmq | latest | Worker communication | diff --git a/docs/src/content/docs/server-training/sglang.mdx b/docs/src/content/docs/server-training/sglang.mdx index c8b190ea..d4893a0a 100644 --- a/docs/src/content/docs/server-training/sglang.mdx +++ b/docs/src/content/docs/server-training/sglang.mdx @@ -126,20 +126,20 @@ After a weight sync, xorl-sglang flushes its KV cache. The training server sends ## Installation -xorl-sglang is included as a git submodule under `submodules/xorl-sglang`. If you cloned with `--recurse-submodules`, it's already checked out. +xorl-sglang is included as a git submodule under `submodules/xorl-sglang`. If you cloned with `--recurse-submodules`, it's already checked out. Keep it in the isolated Torch 2.11 environment; the default XoRL profile uses Torch 2.12 and is not ABI-compatible with the pinned `sglang-kernel` wheel. ```bash -pip install -e "submodules/xorl-sglang/python[all]" +uv venv .venv-sglang --python 3.12 +uv pip install --python .venv-sglang/bin/python -e submodules/xorl-sglang/python +uv pip install --python .venv-sglang/bin/python \ + torchdata==0.11.0 nvidia-cutlass-dsl==4.5.2 quack-kernels==0.5.0 +uv pip install --python .venv-sglang/bin/python --no-deps -e . +uv pip install --python .venv-sglang/bin/python pytest +PYTHONPATH=src:submodules/xorl-sglang/python XORL_REQUIRE_SGL_KERNEL=1 \ + .venv-sglang/bin/python -m pytest -q tests/ops/test_sgl_kernel_smoke.py ``` -Or use `pyproject.sglang.toml` to install xorl, xorl-client, and xorl-sglang together (pins PyTorch to 2.9.1): - -```bash -cp pyproject.sglang.toml pyproject.toml -uv sync # or: pip install -e . -``` - -See the [installation guide](/xorl/getting-started/installation/#install-submodules) for full details. +The bundled `pyproject.sglang.toml` pins the same Torch 2.11 combined profile. See the [installation guide](/xorl/getting-started/installation/#install-submodules) for full details. ## Launching xorl-sglang diff --git a/docs/src/content/docs/testing/adding-tests.md b/docs/src/content/docs/testing/adding-tests.md index 7e799799..a0673b0e 100644 --- a/docs/src/content/docs/testing/adding-tests.md +++ b/docs/src/content/docs/testing/adding-tests.md @@ -62,11 +62,6 @@ def test_with_dataset(fake_text_dataset): sample = fake_text_dataset[0] assert sample["input_ids"].shape == (128,) -def test_with_packed_dataset(fake_packed_dataset): - # FakePackedDataset: 100 samples, 3 packed seqs each, position_ids included - sample = fake_packed_dataset[0] - assert "position_ids" in sample - def test_with_collator_input(sample_features): # List of 2 dicts with input_ids, attention_mask, labels (len=5) assert len(sample_features) == 2 diff --git a/docs/src/content/docs/testing/existing-tests.md b/docs/src/content/docs/testing/existing-tests.md index e7a12f9c..4d33ee68 100644 --- a/docs/src/content/docs/testing/existing-tests.md +++ b/docs/src/content/docs/testing/existing-tests.md @@ -10,7 +10,6 @@ Tests the full data pipeline: raw dataset loading, packing, batching, and collat - `test_shared.py` — loading datasets from Hugging Face Hub, local paths, and URLs; train/validation splits; merging multiple datasets - `test_packing.py` — FFD (First-Fit Decreasing) packing, sequential allocation, position ID generation, packed dataset merging - `test_hash.py` — dataset fingerprinting and config hashing for cache invalidation -- `test_file_lock_loader.py` — multi-process safe dataset preparation with file locking; counter management and cleanup - `test_utils.py` — retry strategies (exponential, linear, constant backoff), MD5/SHA256 hashing **`data/collators/`** @@ -130,5 +129,4 @@ E2E tests use small randomly-initialized model variants (no downloads required) | `tiny_dense_model_dir_with_weights` | Random-init dense with weights on disk | | `tiny_moe_model_dir` | Random-init Qwen3-MoE | | `tiny_moe_model_dir_with_weights` | Random-init MoE with weights on disk | -| `small_dense_model_dir_with_weights` | Dense, hidden size 256 | | `small_moe_model_dir_with_weights` | MoE, intermediate size 64 (NF4 group_size compatible) | diff --git a/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml b/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml index 2694a532..601c900a 100644 --- a/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml +++ b/examples/server/configs/lora/qwen3_5_35b_a3b_lora.yaml @@ -1,5 +1,5 @@ # Server-side configuration for XORL Training Server -# Qwen3.5-35B-A3B LoRA (bf16 base + LoRA rank 32) +# Qwen3.5/Qwen3.6-35B-A3B projection-topology LoRA (bf16 base + LoRA rank 16) # # 1 node (8 GPUs): EP=8, Ulysses SP=8, FSDP shard=1 # @@ -41,9 +41,11 @@ sample_packing_sequence_len: 32768 enable_packing: true enable_lora: true -lora_rank: 32 -lora_alpha: 32 -lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] +lora_rank: 16 +lora_alpha: 16 +# GDN q/k/v/z are four independent factors (g_proj is in_proj_z), and +# gate/up/down includes every fused-base shared expert. +lora_target_modules: ["q_proj", "k_proj", "v_proj", "g_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] skip_initial_checkpoint: true # bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) diff --git a/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml b/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml index 61b09c4c..acd0f2cf 100644 --- a/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml +++ b/examples/server/configs/lora/qwen3_5_397b_a17b_lora.yaml @@ -43,7 +43,7 @@ enable_packing: true enable_lora: true lora_rank: 32 lora_alpha: 32 -lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] +lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "g_proj", "down_proj"] skip_initial_checkpoint: true # bi_fused = K3 lm-head contract CE (fp32-class, tp1, no z-loss); requires tree >= 09a5ae3d3 (older trees silently ran eager for bi_fused) diff --git a/pyproject.sglang.toml b/pyproject.sglang.toml index 1bebb45a..aa702910 100644 --- a/pyproject.sglang.toml +++ b/pyproject.sglang.toml @@ -1,10 +1,9 @@ -# Alternative pyproject.toml that pins PyTorch 2.9.1 so that xorl, xorl-client, -# and xorl-sglang can all be installed in the same environment. +# Alternative uv profile that pins PyTorch 2.11.0 so that xorl, xorl-client, +# and the exact SGLang kernel path can run in the same environment. # # Usage: # cp pyproject.sglang.toml pyproject.toml -# uv sync # (uv) -# pip install -e . # (conda/pip) +# uv sync [build-system] requires = ["setuptools>=61.0", "wheel"] @@ -41,15 +40,13 @@ dependencies = [ # P2P / Mooncake weight sync "mooncake-transfer-engine==0.3.9", "xorl-client @ git+https://github.com/togethercomputer/xorl-client.git", - # PyTorch 2.9.1 with CUDA 12.9 (compatible with xorl-sglang) - "torch @ https://download.pytorch.org/whl/cu129/torch-2.9.1%2Bcu129-cp312-cp312-manylinux_2_28_x86_64.whl", - "torchvision @ https://download.pytorch.org/whl/cu129/torchvision-0.24.1%2Bcu129-cp312-cp312-manylinux_2_28_x86_64.whl", - "triton @ https://download.pytorch.org/whl/triton-3.5.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", - # Flash Attention - "flash-attn-3 @ https://github.com/windreamer/flash-attention3-wheels/releases/download/2026.02.17-06dc5e7/flash_attn_3-3.0.0%2B20260216.cu129torch291cxx11abitrue.fec3a6-cp39-abi3-linux_x86_64.whl", - "flash-attn-cute @ git+https://github.com/Dao-AILab/flash-attention.git@5678dd909aca97f925957dab716022046fb1e44f#subdirectory=flash_attn/cute", - # Submodule - "sglang[all] @ file:submodules/xorl-sglang/python", + # Match the pinned SGLang ABI. Its dependency set supplies the compatible + # sglang-kernel, CUDA runtime, and Flash Attention packages. + "torch==2.11.0", + "torchvision==0.26.0", + "triton==3.6.0", + # Submodule (resolved through tool.uv.sources below) + "sglang", ] @@ -79,6 +76,16 @@ test = [ # NOTE 2: When updating this line, make sure to update Dockerfile under docker/ to the same # version and release new docker images. required-version = ">=0.8.14" +override-dependencies = [ + # Pinned XoRL's Quack imports use the pre-4.6 CUTLASS API. Pinned SGLang's + # hash_topk, LoRA, and sgl_kernel paths are validated separately by the ABI + # smoke and do not require its newer Quack package. + "nvidia-cutlass-dsl==4.5.2", + "quack-kernels==0.5.0", +] + +[tool.uv.sources] +sglang = { path = "submodules/xorl-sglang/python", editable = true } [tool.setuptools.dynamic] diff --git a/scripts/audit_tests.py b/scripts/audit_tests.py new file mode 100755 index 00000000..73dcedc6 --- /dev/null +++ b/scripts/audit_tests.py @@ -0,0 +1,283 @@ +#!/usr/bin/env python3 +"""Surface test-audit candidates without deciding whether to remove them.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +import subprocess +from collections import Counter, defaultdict +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Iterable + + +DECISIONS_FILE = "test_audit_decisions.json" +VALID_DECISIONS = {"keep", "consolidate", "rewrite", "relocate", "remove"} +VALID_STATUSES = {"proposed", "accepted", "applied", "rejected"} + + +@dataclass(frozen=True) +class TestCase: + path: str + name: str + line: int + end_line: int + body_hash: str + signals: tuple[str, ...] + + +def _call_name(node: ast.Call) -> str: + parts: list[str] = [] + value: ast.expr | None = node.func + while isinstance(value, ast.Attribute): + parts.append(value.attr) + value = value.value + if isinstance(value, ast.Name): + parts.append(value.id) + return ".".join(reversed(parts)) + + +def _body_hash(node: ast.FunctionDef | ast.AsyncFunctionDef) -> str: + body = list(node.body) + if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant): + if isinstance(body[0].value.value, str): + body.pop(0) + normalized = ast.dump(ast.Module(body=body, type_ignores=[]), include_attributes=False) + return hashlib.sha256(normalized.encode()).hexdigest()[:16] + + +def _contains_call(node: ast.AST, names: set[str]) -> bool: + return any( + isinstance(child, ast.Call) and (_call_name(child) in names or _call_name(child).split(".")[-1] in names) + for child in ast.walk(node) + ) + + +def _has_outcome(node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool: + if any(isinstance(child, (ast.Assert, ast.Raise)) for child in ast.walk(node)): + return True + for child in ast.walk(node): + if not isinstance(child, ast.Call): + continue + name = _call_name(child) + leaf = name.split(".")[-1] + normalized_leaf = leaf.lstrip("_") + if name in {"pytest.raises", "pytest.warns", "pytest.deprecated_call"}: + return True + if leaf == "simplefilter" and child.args: + first_arg = child.args[0] + if isinstance(first_arg, ast.Constant) and first_arg.value == "error": + return True + if normalized_leaf.startswith("assert") or leaf in {"fail", "raises", "warns", "start_processes"}: + return True + return False + + +def _skip_inside_condition(node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool: + for child in ast.walk(node): + if isinstance(child, (ast.If, ast.Match)) and _contains_call(child, {"pytest.skip"}): + return True + return False + + +def _reads_module_source(node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool: + """Detect tests that read an imported module's source instead of behavior.""" + for child in ast.walk(node): + if not isinstance(child, ast.Call) or not isinstance(child.func, ast.Attribute): + continue + if child.func.attr not in {"read_text", "read_bytes"}: + continue + receiver = child.func.value + if not isinstance(receiver, ast.Call) or not receiver.args: + continue + if _call_name(receiver).split(".")[-1] not in {"Path", "open"}: + continue + if any( + isinstance(part, ast.Attribute) and part.attr == "__file__" + for argument in receiver.args + for part in ast.walk(argument) + ): + return True + return False + + +def _signals(node: ast.FunctionDef | ast.AsyncFunctionDef) -> tuple[str, ...]: + signals: list[str] = [] + has_outcome = _has_outcome(node) + has_print = _contains_call(node, {"print", "pprint"}) + + if not has_outcome: + signals.append("no-observable-outcome") + if has_print and not has_outcome: + signals.append("print-only") + if _skip_inside_condition(node): + signals.append("conditional-runtime-skip") + if _contains_call(node, {"inspect.getsource", "getsource"}) or _reads_module_source(node): + signals.append("source-inspection") + + return tuple(signals) + + +class Collector(ast.NodeVisitor): + def __init__(self, path: str) -> None: + self.path = path + self.classes: list[str] = [] + self.tests: list[TestCase] = [] + + def visit_ClassDef(self, node: ast.ClassDef) -> None: + self.classes.append(node.name) + self.generic_visit(node) + self.classes.pop() + + def _visit_test(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None: + if node.name.startswith("test_"): + name = ".".join([*self.classes, node.name]) + self.tests.append( + TestCase( + path=self.path, + name=name, + line=node.lineno, + end_line=node.end_lineno or node.lineno, + body_hash=_body_hash(node), + signals=_signals(node), + ) + ) + self.generic_visit(node) + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + self._visit_test(node) + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: + self._visit_test(node) + + +def _repository_test_files(repo_root: Path) -> list[Path]: + try: + result = subprocess.run( + [ + "git", + "-C", + str(repo_root), + "ls-files", + "--cached", + "--others", + "--exclude-standard", + "--", + "tests", + ], + check=True, + capture_output=True, + text=True, + ) + except (FileNotFoundError, subprocess.CalledProcessError): + return sorted((repo_root / "tests").rglob("*.py")) + paths = [repo_root / item for item in result.stdout.splitlines() if item.endswith(".py")] + return [path for path in paths if path.is_file()] + + +def _collect(repo_root: Path) -> tuple[list[TestCase], list[dict[str, str]]]: + tests: list[TestCase] = [] + parse_errors: list[dict[str, str]] = [] + for path in _repository_test_files(repo_root): + relative = path.relative_to(repo_root).as_posix() + try: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=relative) + except (OSError, SyntaxError) as exc: + parse_errors.append({"path": relative, "error": str(exc)}) + continue + collector = Collector(relative) + collector.visit(tree) + tests.extend(collector.tests) + return tests, parse_errors + + +def _duplicate_groups(tests: Iterable[TestCase]) -> list[list[dict[str, object]]]: + by_hash: dict[str, list[TestCase]] = defaultdict(list) + for test in tests: + by_hash[test.body_hash].append(test) + groups = [] + for matches in by_hash.values(): + locations = {(match.path, match.name) for match in matches} + if len(locations) > 1: + groups.append([asdict(match) for match in sorted(matches, key=lambda item: (item.path, item.line))]) + return sorted(groups, key=lambda group: (group[0]["path"], group[0]["line"])) + + +def _load_decisions(repo_root: Path) -> list[dict[str, object]]: + path = repo_root / DECISIONS_FILE + if not path.exists(): + return [] + payload = json.loads(path.read_text(encoding="utf-8")) + if payload.get("schema_version") != 1 or not isinstance(payload.get("items"), list): + raise ValueError(f"{DECISIONS_FILE} must contain schema_version=1 and an items list") + + seen: set[str] = set() + for item in payload["items"]: + missing = {"id", "scope", "decision", "status", "evidence"} - set(item) + if missing: + raise ValueError(f"decision is missing {sorted(missing)}: {item}") + if item["id"] in seen: + raise ValueError(f"duplicate decision id: {item['id']}") + if item["decision"] not in VALID_DECISIONS: + raise ValueError(f"invalid decision {item['decision']!r} for {item['id']}") + if item["status"] not in VALID_STATUSES: + raise ValueError(f"invalid status {item['status']!r} for {item['id']}") + if not isinstance(item["evidence"], list) or not item["evidence"]: + raise ValueError(f"decision {item['id']} needs at least one evidence item") + seen.add(item["id"]) + return payload["items"] + + +def _report(repo_root: Path) -> dict[str, object]: + tests, parse_errors = _collect(repo_root) + signal_counts = Counter(signal for test in tests for signal in test.signals) + candidates = [asdict(test) for test in tests if test.signals] + return { + "summary": { + "repository_python_test_files": len(_repository_test_files(repo_root)), + "test_definitions": len(tests), + "candidate_definitions": len(candidates), + "signal_counts": dict(sorted(signal_counts.items())), + "exact_duplicate_body_groups": len(_duplicate_groups(tests)), + "parse_errors": len(parse_errors), + }, + "candidates": candidates, + "exact_duplicate_body_groups": _duplicate_groups(tests), + "parse_errors": parse_errors, + "curated_decisions": _load_decisions(repo_root), + } + + +def _print_text(report: dict[str, object]) -> None: + summary = report["summary"] + print(json.dumps(summary, indent=2, sort_keys=True)) + print("\nCandidates (signals are prompts for review, not deletion decisions):") + for item in report["candidates"]: + signals = ", ".join(item["signals"]) + print(f" {item['path']}:{item['line']} {item['name']} [{signals}]") + print("\nExact duplicate test bodies:") + for group in report["exact_duplicate_body_groups"]: + print(" group") + for item in group: + print(f" {item['path']}:{item['line']} {item['name']}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parents[1]) + parser.add_argument("--format", choices=("text", "json"), default="text") + args = parser.parse_args() + + report = _report(args.repo_root.resolve()) + if args.format == "json": + print(json.dumps(report, indent=2, sort_keys=True)) + else: + _print_text(report) + return 1 if report["parse_errors"] else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/check_public_tree.py b/scripts/check_public_tree.py new file mode 100755 index 00000000..37567eca --- /dev/null +++ b/scripts/check_public_tree.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 +"""Reject private identifiers and environment-specific references in tracked files.""" + +from __future__ import annotations + +import re +import subprocess +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] + +PATTERNS = { + "home directory": re.compile(r"/(?:home|Users)/[A-Za-z0-9._-]+"), + "personal data workspace": re.compile(r"/data/[A-Za-z0-9._-]+/(?:WorkingProjects|outputs|miniconda3)/"), + "user or mirror under /shared": re.compile(r"/shared/(?:apanda|qywu|huggingface)\b"), + "internal repository name": re.compile(r"\bxorl(?:-sglang|-client)?-internal\b"), + "kubectl invocation": re.compile(r"\bkubectl\b", re.IGNORECASE), + "cluster service address": re.compile(r"\.svc\.cluster\.local\b"), + "scheduling queue label": re.compile(r"\bteam:\s*(?:turbo|shaping)\b", re.IGNORECASE), + "volcano scheduler key": re.compile(r"scheduling\.volcano\.sh"), + "pointer to an internal-only note": re.compile(r"\bdocs/notes/"), + "internal account name": re.compile(r"\bapanda\b"), + "internal branch name": re.compile(r"\bapanda-dev\b"), + "internal tracker ticket": re.compile(r"\bXORL-\d+\b"), + "authoring-assistant attribution": re.compile(r"(?i)\b(?:claude|anthropic|copilot)\b"), +} + +# Pattern definitions necessarily contain the forbidden strings. The gitignore +# names authoring-tool files only to keep them untracked, which is not a leak. +EXEMPT = { + "scripts/check_public_tree.py", + "src/xorl/sim/calibration_packs.py", + ".gitignore", +} + + +def _tracked_files() -> list[str]: + listing = subprocess.run( + ["git", "ls-files"], + cwd=REPO_ROOT, + capture_output=True, + text=True, + check=True, + ) + return [name for name in listing.stdout.splitlines() if name and name not in EXEMPT] + + +def main() -> int: + violations: list[str] = [] + for name in _tracked_files(): + path = REPO_ROOT / name + if not path.is_file(): + continue + try: + text = path.read_text(encoding="utf-8") + except (UnicodeDecodeError, OSError): + continue + for label, pattern in PATTERNS.items(): + match = pattern.search(text) + if match: + violations.append(f"{name}: {label}: {match.group(0)!r}") + + if violations: + print("Internal references in tracked files:") + print("\n".join(sorted(violations))) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/xorl/arguments.py b/src/xorl/arguments.py index f08f8d47..515be9f9 100644 --- a/src/xorl/arguments.py +++ b/src/xorl/arguments.py @@ -154,6 +154,9 @@ def _validate_loading_method(self): """Validate the dataset loading configuration based on path and other fields""" path = self.path + if self.shards is not None and self.preprocess_shards is not None: + raise ValueError("'shards' and 'preprocess_shards' are mutually exclusive.") + if path == "dummy": if self.type != "tokenized": raise ValueError("Dummy dataset only supports type='tokenized'.") @@ -991,16 +994,6 @@ def optimizer_kwargs(self) -> Dict[str, Any]: ) }, ) - fp8_cfg: Optional[Dict[str, Any]] = field( - default=None, - metadata={ - "help": ( - "Optional compatibility alias for NeMo-style FP8 configs. Supported values are " - "{enabled: true, fp8: e4m3, fp8_recipe: blockwise, fp8_param: false}; " - "TransformerEngine-only recipes are rejected." - ) - }, - ) enable_fp8_training: bool = field( default=False, metadata={ @@ -1211,15 +1204,6 @@ def optimizer_kwargs(self) -> Dict[str, Any]: ) }, ) - moe_checkpoint_method: Optional[str] = field( - default=None, - metadata={ - "help": ( - "Deprecated compatibility alias. Use gradient_checkpointing_method instead. " - "The legacy value 'moe_act' maps to 'recompute_before_dispatch'." - ) - }, - ) @property def moe_recomputed(self) -> bool: @@ -1339,14 +1323,6 @@ def moe_recomputed(self) -> bool: default=True, metadata={"help": "Place EP all-to-all within the node (NVLink). When False, EP spans across nodes."}, ) - ep_outside: Optional[bool] = field( - default=None, - metadata={ - "help": ( - "Deprecated compatibility alias for ep_intranode. When set, ep_intranode is resolved to not ep_outside." - ) - }, - ) ulysses_parallel_size: int = field( default=1, metadata={"help": "Ulysses sequence parallel size."}, @@ -1639,11 +1615,8 @@ def moe_recomputed(self) -> bool: ) def __post_init__(self): - from xorl.fp8_training.config_compat import normalize_fp8_training_config # noqa: PLC0415 from xorl.qarl import normalize_qarl_quant_cfg # noqa: PLC0415 - normalized_fp8_config = normalize_fp8_training_config(vars(self), context="train") - self.enable_fp8_training = bool(normalized_fp8_config.get("enable_fp8_training", self.enable_fp8_training)) if self.enable_qarl: if self.enable_fp8_training: raise ValueError( @@ -1657,23 +1630,6 @@ def __post_init__(self): raise ValueError("qarl_calib_size and qarl_quant_sequence_length require qarl_calib_data") self.qarl_quant_cfg = normalize_qarl_quant_cfg(self.qarl_quant_cfg) - if self.ep_outside is not None: - self.ep_intranode = not self.ep_outside - - if self.moe_checkpoint_method is not None: - if self.moe_checkpoint_method != "moe_act": - raise ValueError( - f"Unknown moe_checkpoint_method: {self.moe_checkpoint_method!r}. " - "The only supported legacy value is 'moe_act'." - ) - if self.gradient_checkpointing_method not in (None, "recompute_before_dispatch"): - raise ValueError( - "moe_checkpoint_method='moe_act' is a legacy alias for " - "gradient_checkpointing_method='recompute_before_dispatch'; " - f"got gradient_checkpointing_method={self.gradient_checkpointing_method!r}." - ) - self.gradient_checkpointing_method = "recompute_before_dispatch" - # Resolve gradient_checkpointing_method into internal fields used by # the model and parallelization code. gcm = self.gradient_checkpointing_method or "recompute_full_layer" @@ -1876,6 +1832,18 @@ class LoRAArguments: "and validates exact runtime module counts/ranks before training." }, ) + unfuse_for_lora: bool = field( + default=False, + metadata={ + "help": ( + "Split fused qkv_proj / gate_up_proj before LoRA injection so q/k/v and " + "gate/up can be adapted. Without this they are stored fused, match no " + "target name, and train unadapted. Costs base-forward throughput: the " + "fused GEMMs split, and the MLP loses its fused SiLU-and-mul kernel. " + "Requires enable_lora; rejected with enable_qlora." + ) + }, + ) save_lora_only: bool = field( default=False, metadata={"help": "Only save LoRA weights (not full model) in HF checkpoints"}, @@ -2237,21 +2205,11 @@ def parse_args(rootclass: T) -> T: input_data: Dict[str, Dict[str, Any]] = json.load(f) if input_data: - from xorl.fp8_training.config_compat import ( # noqa: PLC0415 - extract_nemo_fp8_cfg, - validate_external_fp8_runtime_config, - ) + from xorl.fp8_training.config_compat import validate_external_fp8_runtime_config # noqa: PLC0415 + from xorl.server.removed_config import reject_removed_configuration_fields # noqa: PLC0415 + reject_removed_configuration_fields(input_data, context=input_path or "config") validate_external_fp8_runtime_config(input_data, context=input_path or "config") - nemo_fp8_cfg = extract_nemo_fp8_cfg(input_data) - if nemo_fp8_cfg is not None: - train_data = input_data.setdefault("train", {}) - if not isinstance(train_data, dict): - raise ValueError("train config section must be a mapping") - train_data.setdefault("fp8_cfg", nemo_fp8_cfg) - policy_data = input_data.get("policy") - if isinstance(policy_data, dict) and set(policy_data) == {"megatron_cfg"}: - input_data.pop("policy") for base, arg_dict in input_data.items(): for arg_name, arg_value in arg_dict.items(): diff --git a/src/xorl/checkpoint/checkpointer.py b/src/xorl/checkpoint/checkpointer.py index 048aa6ff..46004392 100644 --- a/src/xorl/checkpoint/checkpointer.py +++ b/src/xorl/checkpoint/checkpointer.py @@ -2,7 +2,6 @@ import json import os from abc import ABC, abstractmethod -from collections import OrderedDict from types import SimpleNamespace from typing import Any, Dict, List, Optional, Set @@ -475,50 +474,6 @@ def state_dict(self): return model_state_dict - @torch.no_grad() - def reference_state_dict(self): - """Collect a lightweight state dict of live params/buffers without DCP materialization. - - This is intended for direct safetensors export paths where we only need - references to the current model tensors and will materialize them one at - a time during save. - """ - model_state_dict: "OrderedDict[str, torch.Tensor]" = OrderedDict() - for part in _as_model_parts(self.model): - modules = dict(part.named_modules(remove_duplicate=False)) - - for name, parameter in part.named_parameters(remove_duplicate=False): - if parameter is not None: - model_state_dict[name] = parameter - - for name, buffer in part.named_buffers(remove_duplicate=False): - if buffer is None: - continue - module_name, _, buffer_name = name.rpartition(".") - parent_module = modules[module_name] if module_name else part - if buffer_name in getattr(parent_module, "_non_persistent_buffers_set", set()): - continue - model_state_dict[name] = buffer - - if self.should_ep_aware: - logger.info_rank0( - "Collecting lightweight model tensor references from ModelState wrapper, " - "restoring EP dim for Experts module" - ) - model_state_dict = self.get_state_dict_with_ep_dim(model_state_dict) - - # Compile-agnostic keys (see state_dict): strip after EP-dim restoration. - model_state_dict = OrderedDict((_strip_compile_prefix(k), v) for k, v in model_state_dict.items()) - - if self.exclude_keys: - model_state_dict = OrderedDict((k, v) for k, v in model_state_dict.items() if k not in self.exclude_keys) - - if self.save_lora_only: - model_state_dict = OrderedDict((k, v) for k, v in model_state_dict.items() if "lora_" in k) - logger.info_rank0(f"LoRA-only save: keeping {len(model_state_dict)} LoRA parameters") - - return model_state_dict - @torch.no_grad() def load_state_dict(self, state_dict): """ diff --git a/src/xorl/cli/export_nvfp4.py b/src/xorl/cli/export_nvfp4.py index 9d7b4ba1..c66bb83f 100644 --- a/src/xorl/cli/export_nvfp4.py +++ b/src/xorl/cli/export_nvfp4.py @@ -29,7 +29,6 @@ from safetensors import safe_open from xorl.ops.quantize.nvfp4_fake_quant import ( - _E2M1_ABS, FP4_E2M1_MAX, FP8_E4M3_MAX, _nvfp4_quantize_blocks, @@ -126,25 +125,6 @@ def quantize_weight_to_nvfp4( return entry -def dequantize_nvfp4_export(entry: dict[str, torch.Tensor]) -> torch.Tensor: - """Reconstruct a bf16 weight from an NVFP4 entry (inverse of quantize).""" - packed = entry[WEIGHT_KEY] - M, half = packed.shape - K = half * 2 - block_scale = entry[BLOCK_SCALE_KEY] - block_size = K // block_scale.shape[1] - flat = packed.reshape(-1) - lo = flat & 0x0F - hi = (flat >> 4) & 0x0F - codes = torch.stack([lo, hi], dim=1).reshape(M, K).to(torch.int64) - grid = torch.tensor(_E2M1_ABS, dtype=torch.float32) - sign = torch.where((codes & 0x8) > 0, -1.0, 1.0) - values = sign * grid[codes & 0x7] - eff = block_scale.float() * entry[GLOBAL_SCALE_KEY].float() # [M, K/bs] - eff = eff.repeat_interleave(block_size, dim=1) # [M, K] - return (values * eff).reshape(M, K).to(torch.bfloat16) - - def write_hf_quant_config(save_dir: Path, *, group_size: int, exclude_modules: list[str]) -> Path: cfg = { "producer": {"name": "xorl-qat", "quant_method": "fake-quant-RTN"}, diff --git a/src/xorl/data/collators/collate_pipeline.py b/src/xorl/data/collators/collate_pipeline.py index 72799d9b..f9b5d129 100644 --- a/src/xorl/data/collators/collate_pipeline.py +++ b/src/xorl/data/collators/collate_pipeline.py @@ -1,16 +1,13 @@ -from typing import Any, Callable, Dict, List, Optional, Sequence, Union +from typing import Any, Callable, Dict, Sequence class CollatePipeline: - def __init__(self, data_collators: Optional[Union[Callable, List[Callable]]] = None): + def __init__(self, data_collators: Sequence[Callable]): """ Args: - data_collators: a list of data collators or a single data collator + data_collators: collators to apply in order """ - - if not isinstance(data_collators, (list, tuple)): - data_collators = [data_collators] - self.data_collators = data_collators + self.data_collators = list(data_collators) def __call__(self, batch: Sequence[Dict[str, Any]]): """ diff --git a/src/xorl/data/prepare/file_lock_loader.py b/src/xorl/data/prepare/file_lock_loader.py deleted file mode 100644 index bec17974..00000000 --- a/src/xorl/data/prepare/file_lock_loader.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Logic for loading / preparing a dataset once over all processes.""" - -import time -from pathlib import Path -from typing import Any, Callable - -from filelock import FileLock - -from ...arguments import Arguments -from .constants import DEFAULT_DATASET_PREPARED_PATH - - -LOCK_FILE_NAME = "datasets_prep.lock" -READY_FILE_NAME = "datasets_ready.flag" -PROCESS_COUNTER_FILE_NAME = "process_counter.txt" - - -class FileLockLoader: - """ - Simple class for abstracting single process data loading / processing. The first - process that creates a lock file does the work; the remaining procesees simply load - the preprocessed dataset once the first process is done. - """ - - def __init__(self, args: Arguments): - self.args = args - self.dataset_prepared_path = args.data.dataset_prepared_path or DEFAULT_DATASET_PREPARED_PATH - self.lock_file_path = Path(self.dataset_prepared_path) / LOCK_FILE_NAME - self.ready_flag_path = Path(self.dataset_prepared_path) / READY_FILE_NAME - self.counter_path = Path(self.dataset_prepared_path) / PROCESS_COUNTER_FILE_NAME - - def load(self, load_fn: Callable[[], Any]) -> Any: - # Ensure directory exists - Path(self.dataset_prepared_path).mkdir(parents=True, exist_ok=True) - - with FileLock(str(self.lock_file_path)): - self._increment_counter() - - if not self.ready_flag_path.exists(): - # First process does the work - result = load_fn() - self.ready_flag_path.touch() - return result - else: - # Other processes wait for the first process to finish - # and then load the already prepared data - while not self.ready_flag_path.exists(): - time.sleep(1.0) # Sleep for 1 second - return load_fn() # Load the prepared data - - def _increment_counter(self): - """Safely increment the process counter.""" - try: - if self.counter_path.exists(): - counter_content = self.counter_path.read_text().strip() - count = int(counter_content) if counter_content else 0 - else: - count = 0 - self.counter_path.write_text(str(count + 1)) - except (ValueError, OSError): - # Handle corrupted counter file or I/O errors - # Reset to 1 for this process - self.counter_path.write_text("1") - - def cleanup(self): - """Clean up ready flag when last process is done.""" - with FileLock(str(self.lock_file_path)): - try: - counter_content = self.counter_path.read_text().strip() - count = int(counter_content) if counter_content else 0 - count -= 1 - - if count <= 0: - # Last process cleans everything up - self.ready_flag_path.unlink(missing_ok=True) - self.counter_path.unlink(missing_ok=True) - else: - # Still have active processes - self.counter_path.write_text(str(count)) - except (ValueError, OSError): - # Handle corrupted counter file or I/O errors - # Force cleanup since we can't determine the count - self.ready_flag_path.unlink(missing_ok=True) - self.counter_path.unlink(missing_ok=True) diff --git a/src/xorl/data/prepare/packing.py b/src/xorl/data/prepare/packing.py index ea894e9b..25e0e8de 100644 --- a/src/xorl/data/prepare/packing.py +++ b/src/xorl/data/prepare/packing.py @@ -26,42 +26,6 @@ LOG = logging.get_logger(__name__) -@numba.njit -def ffd_check(sequence_lengths: np.ndarray, bin_capacity: int, num_bins: int) -> bool: - """First-fit-decreasing bin packing algorithm check. - - Checks if sequences with the given lengths could fit in the specified number of - bins. - - Args: - sequence_lengths: Array of sequence lengths. - bin_capacity: Maximum capacity of each bin. - num_bins: Number of bins available. - - Returns: - `True` if all sequences can be packed, `False` otherwise. - """ - # Sort sequence lengths in descending order for optimal packing - sequence_lengths = np.sort(sequence_lengths)[::-1] - # Initialize all bins with full capacity - bins = np.full((num_bins,), bin_capacity, dtype=sequence_lengths.dtype) - - # Try to place each sequence in the first bin it fits - for size in sequence_lengths: - not_found = True - for idx in range(num_bins): - if bins[idx] >= size: - bins[idx] -= size - not_found = False - break - - # If no bin could fit this sequence, packing failed - if not_found: - return False - - return True - - @numba.njit def pack_group( sequence_lengths: np.ndarray, diff --git a/src/xorl/data/prepare/shared.py b/src/xorl/data/prepare/shared.py index adb9d69c..5ee04956 100644 --- a/src/xorl/data/prepare/shared.py +++ b/src/xorl/data/prepare/shared.py @@ -106,12 +106,13 @@ def datasets_with_name_generator( if config.name and isinstance(config.name, list): for name in config.name: yield replace(config, name=name) - elif config.preprocess_shards and not config.shards: + elif config.preprocess_shards: for shard_idx in range(config.preprocess_shards): yield replace( config, shards=config.preprocess_shards, shards_idx=shard_idx, + preprocess_shards=None, ) else: yield config @@ -371,7 +372,7 @@ def _load_from_data_files( else: raise ValueError("data_files must be either a string or list of strings") - return load_dataset("json", data_files=file_path, **load_dataset_kwargs) + return load_dataset(get_dataset_type(dataset_config), data_files=file_path, **load_dataset_kwargs) def get_prepared_dataset_path(args: Arguments, dataset_hash: str) -> Path: diff --git a/src/xorl/data/prepare/utils.py b/src/xorl/data/prepare/utils.py index 6f0f9ea6..f033c00a 100644 --- a/src/xorl/data/prepare/utils.py +++ b/src/xorl/data/prepare/utils.py @@ -3,35 +3,18 @@ import functools import hashlib import time -from enum import Enum from typing import Callable -import huggingface_hub import requests +from huggingface_hub.errors import HfHubHTTPError -from ...utils import logging - -logger = logging.get_logger(__name__) - - -class RetryStrategy(Enum): - """Enum for retry strategies.""" - - CONSTANT = 1 - LINEAR = 2 - EXPONENTIAL = 3 - - -def retry_on_request_exceptions( - max_retries=3, delay=1, retry_strategy: RetryStrategy = RetryStrategy.EXPONENTIAL -) -> Callable: +def retry_on_request_exceptions(max_retries=3, delay=1) -> Callable: """Decorator that retries function calls on specific request exceptions. Args: max_retries: Maximum number of retry attempts. delay: Base delay between retries in seconds. - retry_strategy: Strategy for calculating retry delays. Returns: Decorated function with retry logic. @@ -47,16 +30,10 @@ def wrapper(*args, **kwargs): requests.exceptions.ReadTimeout, requests.exceptions.ConnectionError, requests.exceptions.HTTPError, - huggingface_hub.errors.HfHubHTTPError, + HfHubHTTPError, ) as exc: if attempt < max_retries - 1: - if retry_strategy == RetryStrategy.EXPONENTIAL: - step_delay = delay * 2**attempt - elif retry_strategy == RetryStrategy.LINEAR: - step_delay = delay * (attempt + 1) - else: - step_delay = delay # Use constant delay. - time.sleep(step_delay) + time.sleep(delay * 2**attempt) else: raise exc @@ -71,8 +48,3 @@ def md5(to_hash: str, encoding: str = "utf-8") -> str: return hashlib.md5(to_hash.encode(encoding), usedforsecurity=False).hexdigest() except TypeError: return hashlib.md5(to_hash.encode(encoding)).hexdigest() # nosec - - -def sha256(to_hash: str, encoding: str = "utf-8") -> str: - """Generate SHA256 hash of a string.""" - return hashlib.sha256(to_hash.encode(encoding)).hexdigest() diff --git a/src/xorl/distributed/__init__.py b/src/xorl/distributed/__init__.py index 89b449e8..0155d18f 100644 --- a/src/xorl/distributed/__init__.py +++ b/src/xorl/distributed/__init__.py @@ -12,7 +12,6 @@ ParallelRole, canonical_moe_reduce_cp_sharded_v3, canonical_moe_reduce_packed_ep16_v2, - canonical_moe_reduce_reference, canonical_moe_reduce_v1, resolve_canonical_moe_transport, ) diff --git a/src/xorl/distributed/canonical_moe.py b/src/xorl/distributed/canonical_moe.py index 92045afd..cc4d8b42 100644 --- a/src/xorl/distributed/canonical_moe.py +++ b/src/xorl/distributed/canonical_moe.py @@ -14,11 +14,8 @@ from __future__ import annotations -import hashlib -import json from dataclasses import dataclass from enum import Enum, IntEnum -from typing import Any import torch import torch.distributed as dist @@ -32,7 +29,6 @@ class ParallelRole(str, Enum): TRAINER = "trainer" - SAMPLER = "sampler" PRIMITIVE_TEST = "primitive_test" @@ -161,7 +157,6 @@ class ParallelPlan: logical_ordinals_by_group: tuple[tuple[int, ...], ...] pipeline_layer_ranges: tuple[tuple[int, int], ...] cp_ep_aliases: tuple[tuple[int, int], ...] - launcher_tp_size: int | None = None contract_version: str = CANONICAL_MOE_REDUCE_VERSION def __post_init__(self) -> None: @@ -239,29 +234,6 @@ def validate(self) -> None: raise ValueError("GLM-5.2 trainer combine groups must be contiguous groups of cp_size physical ranks") if self.logical_ordinals_by_group != (expected_ordinals,) * len(expected_groups): raise ValueError("GLM-5.2 trainer requires identity logical contributor ordinals in every group") - if self.launcher_tp_size is not None: - raise ValueError("Trainer ParallelPlan does not accept a launcher_tp_size alias") - elif self.role is ParallelRole.SAMPLER: - expected = (8, 1, 1, 1, 8, 8, 1) - actual = ( - self.world_size, - self.pp_size, - self.tp_size, - self.dp_size, - self.cp_size, - self.ep_size, - self.effective_dense_tp, - ) - if actual != expected: - raise ValueError(f"GLM-5.2 sampler topology must be {expected}, got {actual}") - if self.pipeline_layer_ranges != ((0, 78),): - raise ValueError("GLM-5.2 sampler must own all 78 layers in one pipeline stage") - if self.combine_groups != (tuple(range(8)),): - raise ValueError("GLM-5.2 sampler combine group must contain physical ranks 0..7") - if self.logical_ordinals_by_group != (expected_ordinals,): - raise ValueError("GLM-5.2 sampler requires identity logical contributor ordinals") - if self.launcher_tp_size != 8: - raise ValueError("GLM-5.2 sampler launcher-level tp_size must be exactly 8") elif self.role is ParallelRole.PRIMITIVE_TEST: if self.world_size != self.cp_size or len(self.combine_groups) != 1: raise ValueError("Primitive plans use one combine group spanning the test world") @@ -299,25 +271,6 @@ def glm52_trainer( cp_ep_aliases=tuple((rank, rank) for rank in identity), ) - @classmethod - def glm52_sampler(cls, *, launcher_tp_size: int) -> ParallelPlan: - identity = tuple(range(8)) - return cls( - role=ParallelRole.SAMPLER, - world_size=8, - pp_size=1, - tp_size=1, - dp_size=1, - cp_size=8, - ep_size=8, - effective_dense_tp=1, - combine_groups=(identity,), - logical_ordinals_by_group=(identity,), - pipeline_layer_ranges=((0, 78),), - cp_ep_aliases=tuple((rank, rank) for rank in identity), - launcher_tp_size=launcher_tp_size, - ) - @classmethod def primitive( cls, @@ -348,34 +301,6 @@ def group_index_for_physical_rank(self, physical_global_rank: int) -> int: raise ValueError(f"Physical rank {physical_global_rank} does not belong to exactly one combine group") return matches[0] - def logical_ordinal(self, physical_global_rank: int) -> int: - group_index = self.group_index_for_physical_rank(physical_global_rank) - group = self.combine_groups[group_index] - return self.logical_ordinals_by_group[group_index][group.index(physical_global_rank)] - - def as_dict(self) -> dict[str, Any]: - return { - "contract_version": self.contract_version, - "role": self.role.value, - "world_size": self.world_size, - "pp_size": self.pp_size, - "tp_size": self.tp_size, - "dp_size": self.dp_size, - "cp_size": self.cp_size, - "ep_size": self.ep_size, - "effective_dense_tp": self.effective_dense_tp, - "launcher_tp_size": self.launcher_tp_size, - "combine_groups": self.combine_groups, - "logical_ordinals_by_group": self.logical_ordinals_by_group, - "pipeline_layer_ranges": self.pipeline_layer_ranges, - "cp_ep_aliases": self.cp_ep_aliases, - } - - @property - def digest(self) -> str: - payload = json.dumps(self.as_dict(), sort_keys=True, separators=(",", ":")).encode() - return hashlib.sha256(payload).hexdigest() - @dataclass(frozen=True) class CanonicalMoEGraphMetadata: @@ -534,23 +459,6 @@ def _adjacent_pairwise_bf16(partials: torch.Tensor) -> torch.Tensor: return level[0] -def canonical_moe_reduce_reference( - partials_by_logical_ordinal: torch.Tensor, - metadata: CanonicalMoEGraphMetadata, -) -> torch.Tensor: - """Executable, non-communicating reference used by independent tests.""" - if partials_by_logical_ordinal.ndim < 3: - raise ValueError("Reference partials must have contributor, row, and payload dimensions") - if partials_by_logical_ordinal.shape[1] != metadata.capacity: - raise ValueError("Reference partial row count must equal metadata capacity") - result = _adjacent_pairwise_bf16(partials_by_logical_ordinal) - return torch.where( - metadata.valid_mask.view(-1, *([1] * (result.ndim - 1))), - result, - torch.zeros_like(result), - ) - - def _transport_and_fold( local_partial: torch.Tensor, absolute_positions: torch.Tensor, @@ -974,7 +882,6 @@ def canonical_moe_reduce_cp_sharded_v3( "OutputDistribution", "ParallelPlan", "ParallelRole", - "canonical_moe_reduce_reference", "canonical_moe_reduce_packed_ep16_v2", "canonical_moe_reduce_cp_sharded_v3", "canonical_moe_reduce_v1", diff --git a/src/xorl/distributed/ep_gradients.py b/src/xorl/distributed/ep_gradients.py index ddfb9d40..5aded36c 100644 --- a/src/xorl/distributed/ep_gradients.py +++ b/src/xorl/distributed/ep_gradients.py @@ -22,12 +22,6 @@ class GradientSyncStats: gradient_bytes: int = 0 reduced_bytes: int = 0 - @property - def parameter_count(self) -> int: - """Backward-compatible alias for the configured parameter count.""" - - return self.configured_parameter_count - def _wait_for_local_tensor(tensor: torch.Tensor) -> torch.Tensor: wait = getattr(tensor, "wait", None) @@ -69,30 +63,6 @@ def _set_parameter_local_gradient(parameter: torch.Tensor, local_gradient: torch parameter.grad = local_gradient -def synchronize_ep_replicated_gradient(gradient: torch.Tensor) -> torch.Tensor: - """Sum one replicated EP gradient across the EP group. - - This hook remains available for small standalone callers. The production - model path uses the coalesced optimizer-boundary reducer below. - """ - - if not torch.distributed.is_available() or not torch.distributed.is_initialized(): - return gradient - - from .parallel_state import get_parallel_state # noqa: PLC0415 - - parallel_state = get_parallel_state() - if not parallel_state.ep_enabled: - return gradient - ep_group = parallel_state.ep_group - if ep_group is None or torch.distributed.get_world_size(ep_group) <= 1: - return gradient - - local_gradient = gradient.to_local() if hasattr(gradient, "to_local") else gradient - torch.distributed.all_reduce(local_gradient, group=ep_group) - return gradient - - @torch.no_grad() def synchronize_ep_replicated_gradients( model: torch.nn.Module, @@ -111,11 +81,11 @@ def synchronize_ep_replicated_gradients( if not getattr(model, "_ep_replicated_gradient_sync_enabled", False): return GradientSyncStats() groups = getattr(model, "_ep_param_groups", {}) - if "ep_replicated_gradient_sync" in groups: - parameters = list(groups["ep_replicated_gradient_sync"]) - else: - # Compatibility for small callers constructing the older group shape. - parameters = list(groups.get("ep_replicated", ())) + if "ep_replicated_gradient_sync" not in groups: + raise RuntimeError( + "EP replicated-gradient synchronization is enabled without the ep_replicated_gradient_sync parameter group" + ) + parameters = list(groups["ep_replicated_gradient_sync"]) if not parameters: return GradientSyncStats() if not torch.distributed.is_available() or not torch.distributed.is_initialized(): @@ -242,15 +212,3 @@ def _all_reduce_gradient_bucket( gradient_bytes=gradient_numel * torch.float32.itemsize, reduced_bytes=flat.numel() * torch.float32.itemsize, ) - - -def register_ep_replicated_gradient_hooks(parameters: Iterable[torch.Tensor]) -> None: - """Install one EP-sum hook for standalone tests and diagnostics.""" - - for parameter in parameters: - if not parameter.requires_grad: - continue - if getattr(parameter, "_xorl_ep_replicated_gradient_hook", None) is not None: - continue - handle = parameter.register_hook(synchronize_ep_replicated_gradient) - parameter._xorl_ep_replicated_gradient_hook = handle diff --git a/src/xorl/distributed/pipeline_parallel.py b/src/xorl/distributed/pipeline_parallel.py index c538c74a..c97507fa 100644 --- a/src/xorl/distributed/pipeline_parallel.py +++ b/src/xorl/distributed/pipeline_parallel.py @@ -53,7 +53,6 @@ "build_pipeline_schedule", "schedule_stage_style", "stage_ids_for_rank", - "is_single_stage_schedule", "schedule_splits_backward", "validate_pp_schedule_config", ] @@ -80,11 +79,6 @@ def schedule_stage_style(schedule_name: str) -> str: return _SCHEDULE_STYLES[key] -def is_single_stage_schedule(schedule_name: str) -> bool: - """True when the schedule requires exactly one stage per rank (GPipe, 1F1B).""" - return issubclass(get_schedule_class(schedule_name), PipelineScheduleSingle) - - _BACKWARD_SPLIT_SCHEDULES = frozenset({"interleavedzerobubble", "zbvzerobubble", "dualpipev"}) diff --git a/src/xorl/distributed/pp_profiling.py b/src/xorl/distributed/pp_profiling.py index 04c83f53..f3fa3e2c 100644 --- a/src/xorl/distributed/pp_profiling.py +++ b/src/xorl/distributed/pp_profiling.py @@ -27,7 +27,6 @@ __all__ = [ "PPBubbleProfiler", - "analytic_bubble_fraction", "estimate_p2p_bytes_per_step", "merge_busy_intervals", ] @@ -57,34 +56,6 @@ def merge_busy_intervals(intervals: Iterable[tuple[float, float]]) -> float: return total -def analytic_bubble_fraction(schedule_name: str, pp: int, virtual_stages: int, n_microbatches: int) -> float: - """Textbook bubble fraction for a schedule at (pp, virtual_stages, n_microbatches). - - GPipe/1F1B: ``(p-1)/(m+p-1)``. Interleaved1F1B: ``(p-1)/(v*m+p-1)`` (interleaving - with v virtual stages divides the bubble by ~v). Zero-bubble schedules - (InterleavedZeroBubble, ZBVZeroBubble, DualPipeV): ~0.0 by construction. - - All values are approximations: they assume uniform per-microbatch compute across - stages, fwd:bwd cost ratios matching the schedule's design assumptions, enough - microbatches to fill the pipeline, and zero exposed communication. Real zero-bubble - runs retain small warmup/comm residues, so measured > 0 is expected. - """ - if pp < 1 or virtual_stages < 1 or n_microbatches < 1: - raise ValueError( - f"pp, virtual_stages, n_microbatches must all be >= 1, got ({pp}, {virtual_stages}, {n_microbatches})" - ) - key = schedule_name.lower() - if key in ("gpipe", "1f1b"): - if virtual_stages != 1: - raise ValueError(f"Schedule '{schedule_name}' is single-stage-per-rank; virtual_stages must be 1") - return (pp - 1) / (n_microbatches + pp - 1) - if key == "interleaved1f1b": - return (pp - 1) / (virtual_stages * n_microbatches + pp - 1) - if key in ("interleavedzerobubble", "zbvzerobubble", "dualpipev"): - return 0.0 - raise ValueError(f"No analytic bubble model for schedule '{schedule_name}'") - - def _nbytes(t: torch.Tensor) -> int: return t.numel() * t.element_size() diff --git a/src/xorl/distributed/torch_parallelize.py b/src/xorl/distributed/torch_parallelize.py index ba2446ba..d7bafa7d 100644 --- a/src/xorl/distributed/torch_parallelize.py +++ b/src/xorl/distributed/torch_parallelize.py @@ -273,22 +273,6 @@ def _sequence_parallel_fully_folded_into_fsdp(parallel_state) -> bool: ) -def _coerce_optional_bool_config(value: Any, *, name: str) -> Optional[bool]: - if value is None: - return None - if isinstance(value, bool): - return value - if isinstance(value, int) and value in (0, 1): - return bool(value) - if isinstance(value, str): - normalized = value.strip().lower() - if normalized in {"1", "true", "yes", "y", "on"}: - return True - if normalized in {"0", "false", "no", "n", "off"}: - return False - raise ValueError(f"{name} must be a boolean value, got {value!r}.") - - def _configure_manual_fsdp_prefetch( blocks: List["nn.Module"], *, @@ -341,17 +325,14 @@ def parallelize_model_fsdp2( 4. Result: Expert params [32,H/fsdp_size,I], regular params use standard FSDP2 """ parallel_state = get_parallel_state() - enable_manual_forward_prefetch = _coerce_optional_bool_config( - kwargs.pop("enable_forward_prefetch", True), - name="enable_forward_prefetch", - ) - if enable_manual_forward_prefetch is None: - enable_manual_forward_prefetch = True + enable_manual_forward_prefetch = kwargs.pop("enable_forward_prefetch", True) + if not isinstance(enable_manual_forward_prefetch, bool): + raise ValueError(f"enable_forward_prefetch must be a boolean value, got {enable_manual_forward_prefetch!r}.") enable_manual_backward_prefetch_arg = kwargs.pop("enable_backward_prefetch", None) - enable_manual_backward_prefetch_arg = _coerce_optional_bool_config( - enable_manual_backward_prefetch_arg, - name="enable_backward_prefetch", - ) + if enable_manual_backward_prefetch_arg is not None and not isinstance(enable_manual_backward_prefetch_arg, bool): + raise ValueError( + f"enable_backward_prefetch must be a boolean value, got {enable_manual_backward_prefetch_arg!r}." + ) enable_manual_backward_prefetch = ( enable_manual_forward_prefetch if enable_manual_backward_prefetch_arg is None @@ -967,7 +948,11 @@ def build_parallelize_model( if any(isinstance(m, LoraLinear) for m in model_part.modules()): raise NotImplementedError("Tensor parallelism + LoRA is not currently supported.") - # Unfuse fused projections (qkv_proj, gate_up_proj) for TP compatibility + # Unfuse fused projections (qkv_proj, gate_up_proj) for TP compatibility. + # Skip when the model is already unfused: unfuse_for_tp deletes the fused + # module, so a second call raises AttributeError. Reachable only for a + # MoE-only adapter set, whose MoEExpertsLoRA the LoraLinear check above + # does not catch; any LoraLinear is rejected before this point. if hasattr(model_part, "unfuse_for_tp") and not getattr(model_part, "_unfused_for_tp", False): if i == 0: logger.info_rank0("Unfusing projections for tensor parallelism...") @@ -1117,7 +1102,11 @@ def _reentrant_ckpt_with_kwargs(fn, *args, **kw): if any(isinstance(m, LoraLinear) for m in model.modules()): raise NotImplementedError("Tensor parallelism + LoRA is not currently supported.") - # Unfuse fused projections (qkv_proj, gate_up_proj) for TP compatibility + # Unfuse fused projections (qkv_proj, gate_up_proj) for TP compatibility. + # Skip when the model is already unfused: unfuse_for_tp deletes the fused + # module, so a second call raises AttributeError. Reachable only for a + # MoE-only adapter set, whose MoEExpertsLoRA the LoraLinear check above + # does not catch; any LoraLinear is rejected before this point. if hasattr(model, "unfuse_for_tp") and not getattr(model, "_unfused_for_tp", False): logger.info_rank0("Unfusing projections for tensor parallelism...") model.unfuse_for_tp() diff --git a/src/xorl/distributed/utils.py b/src/xorl/distributed/utils.py index fc866310..e17a88c6 100644 --- a/src/xorl/distributed/utils.py +++ b/src/xorl/distributed/utils.py @@ -1,5 +1,4 @@ import re -from typing import List import torch.nn as nn @@ -22,67 +21,6 @@ def get_module_from_path(model: nn.Module, path: str): return get_module_from_path(next_obj, ".".join(attrs[1:])) -def check_all_fqn_match(path_patterns: List[str], path_keys: List[str]): - """ - Check - """ - assert isinstance(path_patterns, list), f"path_patterns must be a list, got {type(path_patterns)}" - assert isinstance(path_keys, (list, tuple)), f"path_keys must be a list or tuple, got {type(path_keys)}" - - if len(path_patterns) != len(path_keys): - return False - - regex_list = [] - for pattern in path_patterns: - regex_str = re.escape(pattern).replace(r"\*", r"(\d+)") - regex_str = f"^{regex_str}$" - regex_list.append((pattern, re.compile(regex_str))) - - used_patterns = set() - expected_num = None # the first matched number - - for key in path_keys: - matched = False - for p, regex in regex_list: - if p in used_patterns: - continue - match = regex.match(key) - if match: - current_num = match.group(1) - if expected_num is None: - expected_num = current_num - elif current_num != expected_num: - return False - used_patterns.add(p) - matched = True - break - if not matched: - return False - - return True - - -def check_any_fqn_match(path_patterns: List[str], path_key: str, return_idx: bool = False, prefix: str = None): - assert isinstance(path_patterns, list), f"path_patterns must be a list, got {type(path_patterns)}" - assert isinstance(path_key, str), f"path_key must be a str, got {type(path_key)}" - - if prefix: - path_patterns = [".".join([prefix, pattern]) for pattern in path_patterns] - - regex_list = [] - for pattern in path_patterns: - regex_str = re.escape(pattern).replace(r"\*", r"(\d+)") - regex_str = f"^{regex_str}$" - regex_list.append(re.compile(regex_str)) - - for idx, regex in enumerate(regex_list): - match = regex.match(path_key) - if match: - return idx if return_idx else True - - return -1 if return_idx else False - - def check_fqn_match(fqn_pattern: str, fqn: str, prefix: str = None): assert isinstance(fqn_pattern, str), f"fqn_pattern must be a str, got {type(fqn_pattern)}" assert isinstance(fqn, str), f"fqn must be a str, got {type(fqn)}" diff --git a/src/xorl/fp8_training/__init__.py b/src/xorl/fp8_training/__init__.py index b7c298c3..324e6cd2 100644 --- a/src/xorl/fp8_training/__init__.py +++ b/src/xorl/fp8_training/__init__.py @@ -3,10 +3,8 @@ _CONFIG_EXPORTS = { "UnsupportedFP8ConfigError", "enrich_sync_quantization_with_fp8_bf16_islands", - "extract_nemo_fp8_cfg", "is_blackwell_device", "merge_fp8_bf16_layer_island_excludes", - "normalize_fp8_training_config", "resolve_fp8_bf16_layer_islands", "validate_external_fp8_runtime_config", "validate_fp8_blackwell_training_policy", @@ -60,7 +58,6 @@ def __getattr__(name: str): "UnsupportedFP8ConfigError", "clear_linear_error_profile", "enrich_sync_quantization_with_fp8_bf16_islands", - "extract_nemo_fp8_cfg", "fp8_block_loop_group_gemm_same_mn", "fp8_block_loop_group_gemm_same_nk", "fp8_deep_gemm_group_gemm_same_nk", @@ -75,7 +72,6 @@ def __getattr__(name: str): "inject_fp8_training_into_model", "is_blackwell_device", "merge_fp8_bf16_layer_island_excludes", - "normalize_fp8_training_config", "resolve_fp8_bf16_layer_islands", "summarize_fp8_training_model", "validate_external_fp8_runtime_config", diff --git a/src/xorl/fp8_training/config_compat.py b/src/xorl/fp8_training/config_compat.py index 191a0db2..bbf50765 100644 --- a/src/xorl/fp8_training/config_compat.py +++ b/src/xorl/fp8_training/config_compat.py @@ -29,69 +29,6 @@ def _as_bool(value: Any, *, field_name: str) -> bool: raise UnsupportedFP8ConfigError(f"{field_name} must be a boolean, got {value!r}") -def normalize_fp8_training_config(config: Mapping[str, Any], *, context: str = "train") -> dict[str, Any]: - """Normalize NeMo-style ``fp8_cfg`` onto XoRL-native FP8 training fields. - - XoRL intentionally supports only native block-FP8 compute training with - full-precision master parameters. TransformerEngine-only recipes are - rejected here with targeted messages before model construction starts. - """ - - normalized = dict(config) - fp8_cfg = normalized.get("fp8_cfg") - if fp8_cfg is None: - return normalized - if not isinstance(fp8_cfg, Mapping): - raise UnsupportedFP8ConfigError(f"{context}.fp8_cfg must be a mapping, got {type(fp8_cfg).__name__}") - - enabled = _as_bool(fp8_cfg.get("enabled", False), field_name=f"{context}.fp8_cfg.enabled") - if not enabled: - return normalized - - raw_fp8 = str(fp8_cfg.get("fp8", "e4m3")).strip().lower() - if raw_fp8 != "e4m3": - raise UnsupportedFP8ConfigError( - f"Unsupported {context}.fp8_cfg.fp8={fp8_cfg.get('fp8')!r}. " - "XoRL native FP8 training supports E4M3 block-FP8 only; " - "'hybrid' is a TransformerEngine recipe and is not implemented." - ) - - raw_recipe = str(fp8_cfg.get("fp8_recipe", "blockwise")).strip().lower() - if raw_recipe != "blockwise": - raise UnsupportedFP8ConfigError( - f"Unsupported {context}.fp8_cfg.fp8_recipe={fp8_cfg.get('fp8_recipe')!r}. " - "XoRL native FP8 training supports blockwise FP8 only; " - "TransformerEngine tensorwise and MXFP8 recipes are not implemented." - ) - - fp8_param = fp8_cfg.get("fp8_param", False) - if _as_bool(fp8_param, field_name=f"{context}.fp8_cfg.fp8_param"): - raise UnsupportedFP8ConfigError( - f"Unsupported {context}.fp8_cfg.fp8_param=true. XoRL keeps BF16/FP32 master parameters and does not " - "store trainable parameters, optimizer state, or DCP checkpoints in FP8." - ) - - normalized["enable_fp8_training"] = True - return normalized - - -def extract_nemo_fp8_cfg(config: Mapping[str, Any]) -> dict[str, Any] | None: - """Return ``policy.megatron_cfg.fp8_cfg`` when a NeMo-style config is provided.""" - - policy = config.get("policy") - if not isinstance(policy, Mapping): - return None - megatron_cfg = policy.get("megatron_cfg") - if not isinstance(megatron_cfg, Mapping): - return None - fp8_cfg = megatron_cfg.get("fp8_cfg") - if fp8_cfg is None: - return None - if not isinstance(fp8_cfg, Mapping): - raise UnsupportedFP8ConfigError("policy.megatron_cfg.fp8_cfg must be a mapping") - return dict(fp8_cfg) - - def validate_external_fp8_runtime_config(config: Mapping[str, Any], *, context: str = "config") -> None: """Reject non-XoRL low-precision runtime knobs in XoRL configs.""" diff --git a/src/xorl/lora/modules/delta_linear.py b/src/xorl/lora/modules/delta_linear.py index e913b6cc..e768785c 100644 --- a/src/xorl/lora/modules/delta_linear.py +++ b/src/xorl/lora/modules/delta_linear.py @@ -86,9 +86,6 @@ def _active_scaling(self) -> float: def get_delta_weight(self) -> torch.Tensor: return (self.lora_B[:, : self.active_r] @ self.lora_A[: self.active_r]) * self._active_scaling() - def invalidate_merged_weight_cache(self) -> None: - self._merged_weight_cache = {} - def _merged_weight( self, base_weight: torch.Tensor, diff --git a/src/xorl/lora/modules/linear.py b/src/xorl/lora/modules/linear.py index 43afedf5..d6992b82 100644 --- a/src/xorl/lora/modules/linear.py +++ b/src/xorl/lora/modules/linear.py @@ -175,9 +175,6 @@ def _active_scaling(self) -> float: # Merged-forward exact-model contract lane # ------------------------------------------------------------------ - def invalidate_merged_weight_cache(self) -> None: - self._merged_weight_cache = {} - def _merged_weight_key(self) -> tuple: t = (self.lora_A, self.lora_B, self.weight) return ( diff --git a/src/xorl/lora/utils.py b/src/xorl/lora/utils.py index 5fb4320d..5b2d9b33 100644 --- a/src/xorl/lora/utils.py +++ b/src/xorl/lora/utils.py @@ -11,7 +11,7 @@ import re from dataclasses import dataclass from pathlib import Path -from typing import Dict, Iterable, Iterator, List, Optional, Tuple +from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple import torch import torch.distributed as dist @@ -31,7 +31,28 @@ # Default target modules for common model architectures DEFAULT_TARGET_MODULES = { "llama": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], - "qwen": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + "qwen2": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + "qwen3": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + "qwen3_moe": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + # Qwen3.5/3.6 GDN calls its z projection ``g_proj`` in the trainer. + "qwen3_5": ["q_proj", "k_proj", "v_proj", "g_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + "qwen3_5_moe": [ + "q_proj", + "k_proj", + "v_proj", + "g_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + "olmo2": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + "glm4_moe": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + "gpt_oss": ["q_proj", "k_proj", "v_proj", "o_proj"], + "minimax_m3": ["q_proj", "k_proj", "v_proj", "o_proj"], + "xorl_minimax_m3": ["q_proj", "k_proj", "v_proj", "o_proj"], + "nemotron_h": ["q_proj", "k_proj", "v_proj", "o_proj"], + "deepseek_v4": ["wq_a", "wq_b", "wkv", "wo_a", "wo_b"], "mistral": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], "gemma": ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], "deepseek_v3": [ @@ -117,10 +138,12 @@ def _get_default_target_modules(model: nn.Module) -> List[str]: if model_type in DEFAULT_TARGET_MODULES: return list(DEFAULT_TARGET_MODULES[model_type]) if model_type is not None: - for family, targets in DEFAULT_TARGET_MODULES.items(): + for family in sorted(DEFAULT_TARGET_MODULES, key=len, reverse=True): if family in model_type: - return list(targets) - return ["q_proj", "k_proj", "v_proj", "o_proj"] + return list(DEFAULT_TARGET_MODULES[family]) + raise ValueError( + f"No audited default LoRA targets for model_type={model_type!r}; set lora_target_modules explicitly" + ) def _get_submodule(model: nn.Module, target: str) -> Tuple[nn.Module, str]: @@ -144,6 +167,7 @@ def _get_submodule(model: nn.Module, target: str) -> Tuple[nn.Module, str]: def _find_target_modules( model: nn.Module, target_modules: List[str], + satisfied_targets: Optional[Iterable[str]] = None, ) -> List[str]: """ Find all module paths matching target module names. @@ -158,6 +182,12 @@ def _find_target_modules( The algorithm processes modules top-down and skips children of replaced modules to avoid double-replacement. + Raises for any requested target that matched no module, since it would otherwise + train unadapted. That check is satisfied by a single match anywhere, so a target + some modules carry and others lack stays silent: ``MoEExperts`` exposes + ``gate_proj``/``up_proj`` as properties, so routed experts satisfy those names + for the whole model even when a shared expert beside them is bare. + Args: model: Model to search target_modules: List of module name patterns to match @@ -166,7 +196,8 @@ def _find_target_modules( List of full module paths that match (in top-down order) """ matched_paths = [] - replaced_prefixes = set() # Track replaced module paths to skip their children + replaced_prefixes: Set[str] = set() # Track replaced module paths to skip their children + matched_targets: Set[str] = set(satisfied_targets or ()) for name, module in model.named_modules(): # Skip if this module is under an already-matched parent @@ -184,17 +215,30 @@ def _find_target_modules( if module_name in target_modules: matched_paths.append(name) replaced_prefixes.add(name) + matched_targets.add(module_name) continue # Indirect match: module has attributes/children matching target_modules # This handles MoE experts where user specifies "gate_proj" but the # actual module to replace is "experts" which contains gate_proj weights module_attrs = set(dir(module)) - if any(target in module_attrs for target in target_modules): + indirect_matches = {target for target in target_modules if target in module_attrs} + if indirect_matches: matched_paths.append(name) replaced_prefixes.add(name) + matched_targets |= indirect_matches continue + # Partial coverage is never a valid success: a 2-of-7 match otherwise looks like + # a healthy injection while five requested projections remain unadapted. + unmatched = sorted(set(target_modules) - matched_targets) + if unmatched: + raise ValueError( + f"LoRA targets matched no module: {unmatched} " + f"(adapted: {sorted(matched_targets)}). If this architecture stores them fused, " + "either enable unfuse_for_lora or target the fused names directly." + ) + return matched_paths @@ -271,6 +315,72 @@ def _inject_fused_gdn_delta_lora( return injected +def _inject_fused_projection_delta_lora( + model: nn.Module, + *, + r: int, + lora_alpha: int, + target_modules: List[str], + target_manifest: Optional[dict], +) -> tuple[int, set[str]]: + """Attach independent logical factors while retaining fused base GEMMs.""" + + from xorl.lora.modules.delta_linear import LoraDeltaLinear # noqa: PLC0415 + + injected = 0 + satisfied: set[str] = set() + for module_path, module in list(model.named_modules()): + specs = [] + qkv_proj = getattr(module, "qkv_proj", None) + if getattr(module, "_supports_fused_qkv_lora", False) and isinstance(qkv_proj, nn.Linear): + q_dim = int(getattr(module, "q_dim")) + kv_dim = int(getattr(module, "kv_dim")) + expected_outputs = q_dim + 2 * kv_dim + if qkv_proj.out_features != expected_outputs: + raise ValueError( + f"{module_path}: qkv_proj has {qkv_proj.out_features} outputs, expected {expected_outputs}" + ) + specs.append((qkv_proj, (("q_proj", q_dim), ("k_proj", kv_dim), ("v_proj", kv_dim)))) + + gate_up_proj = getattr(module, "gate_up_proj", None) + intermediate_size = getattr(module, "intermediate_size", None) + if ( + getattr(module, "_supports_fused_gate_up_lora", False) + and isinstance(gate_up_proj, nn.Linear) + and intermediate_size is not None + ): + intermediate_size = int(intermediate_size) + if gate_up_proj.out_features != 2 * intermediate_size: + raise ValueError( + f"{module_path}: gate_up_proj has {gate_up_proj.out_features} outputs, " + f"expected {2 * intermediate_size}" + ) + specs.append((gate_up_proj, (("gate_proj", intermediate_size), ("up_proj", intermediate_size)))) + + for base, projections in specs: + for projection, out_features in projections: + if projection not in target_modules: + continue + path = f"{module_path}.{projection}" + if not _manifest_allows_module_path(path, target_manifest): + continue + if hasattr(module, projection): + raise ValueError(f"{path} already exists before fused-projection LoRA injection") + module.add_module( + projection, + LoraDeltaLinear( + base.in_features, + out_features, + r=r, + lora_alpha=lora_alpha, + device=base.weight.device, + ), + ) + injected += 1 + satisfied.add(projection) + return injected, satisfied + + def inject_lora_into_model( model: nn.Module, r: int = 16, @@ -331,12 +441,22 @@ def inject_lora_into_model( target_modules=target_modules, target_manifest=loaded_manifest, ) + fused_projection_count, fused_projection_targets = _inject_fused_projection_delta_lora( + model, + r=r, + lora_alpha=lora_alpha, + target_modules=target_modules, + target_manifest=loaded_manifest, + ) # Find all matching modules - target_paths = _find_target_modules(model, target_modules) + specially_satisfied = set(fused_projection_targets) + if fused_gdn_count: + specially_satisfied.update({name for name in ("in_proj_qkvz", "out_proj") if name in target_modules}) + target_paths = _find_target_modules(model, target_modules, satisfied_targets=specially_satisfied) target_paths = [path for path in target_paths if _manifest_allows_module_path(path, loaded_manifest)] - if not target_paths and fused_gdn_count == 0: + if not target_paths and fused_gdn_count == 0 and fused_projection_count == 0: raise ValueError( f"No modules found matching target_modules={target_modules}. " f"Please check that the model has modules with these names. " @@ -345,7 +465,8 @@ def inject_lora_into_model( logger.info( f"Injecting LoRA into {len(target_paths)} base modules and " - f"{fused_gdn_count} fused-GDN delta modules with r={r}, alpha={lora_alpha}" + f"{fused_gdn_count} fused-GDN delta modules and {fused_projection_count} " + f"fused-projection delta modules with r={r}, alpha={lora_alpha}" ) # Replace each target module @@ -375,7 +496,7 @@ def inject_lora_into_model( logger.debug(f"Replaced {target_path} with {lora_cls.__name__}") # Check if any modules were actually replaced - if replaced_count == 0 and fused_gdn_count == 0: + if replaced_count == 0 and fused_gdn_count == 0 and fused_projection_count == 0: skipped_info = ", ".join([f"{path} ({typ})" for path, typ in skipped_modules[:5]]) if len(skipped_modules) > 5: skipped_info += f"... and {len(skipped_modules) - 5} more" @@ -393,7 +514,8 @@ def inject_lora_into_model( ) logger.info( - f"Successfully injected LoRA into {replaced_count} base modules and {fused_gdn_count} fused-GDN delta modules" + f"Successfully injected LoRA into {replaced_count} base modules, {fused_gdn_count} " + f"fused-GDN delta modules, and {fused_projection_count} fused-projection delta modules" ) if loaded_manifest is not None and not _defer_manifest_validation: validated = validate_lora_target_manifest(model, loaded_manifest) diff --git a/src/xorl/models/layers/attention/multi_head_attention.py b/src/xorl/models/layers/attention/multi_head_attention.py index 64ddd708..b158182a 100644 --- a/src/xorl/models/layers/attention/multi_head_attention.py +++ b/src/xorl/models/layers/attention/multi_head_attention.py @@ -7,6 +7,7 @@ from xorl.distributed.sequence_parallel.strategy import get_cp_strategy from xorl.models.layers.attention.backend import AttentionKwargs, get_attention_fn +from xorl.models.layers.fused_projection_lora import project_fused_linear_with_lora from xorl.models.layers.normalization import RMS_NORM_FAMILY_NO_RESIDUAL, RMSNorm from xorl.models.layers.rope import apply_rotary_pos_emb @@ -22,6 +23,8 @@ class MultiHeadAttention(nn.Module): ``get_cp_strategy()``. """ + _supports_fused_qkv_lora = True + def __init__(self, config, layer_idx: int): super().__init__() self.config = config @@ -78,7 +81,13 @@ def _project_qkv( hidden_shape = (*input_shape, -1, self.head_dim) if hasattr(self, "qkv_proj"): - qkv = self.qkv_proj(hidden_states) + qkv = project_fused_linear_with_lora( + self, + hidden_states, + base_name="qkv_proj", + projection_names=("q_proj", "k_proj", "v_proj"), + projection_sizes=(self.q_dim, self.kv_dim, self.kv_dim), + ) q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1) else: q = self.q_proj(hidden_states) diff --git a/src/xorl/models/layers/fused_projection_lora.py b/src/xorl/models/layers/fused_projection_lora.py new file mode 100644 index 00000000..0f6852fd --- /dev/null +++ b/src/xorl/models/layers/fused_projection_lora.py @@ -0,0 +1,52 @@ +"""Helpers for independent LoRA factors over fused base projections.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence + +import torch +import torch.nn.functional as F + +from xorl.lora.fold import lora_merged_forward_enabled + + +def project_fused_linear_with_lora( + module, + inputs: torch.Tensor, + *, + base_name: str, + projection_names: Sequence[str], + projection_sizes: Sequence[int], + linear: Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None], torch.Tensor] = F.linear, +) -> torch.Tensor: + """Run one fused base GEMM with optional independent logical adapters. + + Delta-only children live on the parent under their logical projection names. + Dynamic mode adds their outputs to the matching fused slices. Exact merged mode + canonically folds each delta into its base slice and still issues one fused GEMM. + """ + if len(projection_names) != len(projection_sizes): + raise ValueError("projection_names and projection_sizes must have equal length") + + base = getattr(module, base_name) + adapters = [getattr(module, name, None) for name in projection_names] + present = [adapter for adapter in adapters if adapter is not None] + if not present: + return base(inputs) + + merged = [lora_merged_forward_enabled(adapter) for adapter in present] + if any(merged): + if not all(merged): + raise RuntimeError(f"{base_name} logical adapters must select merged forward together") + base_parts = base.weight.split(tuple(int(size) for size in projection_sizes), dim=0) + folded_parts = [ + adapter.merged_weight_for_forward(base_part) if adapter is not None else base_part + for adapter, base_part in zip(adapters, base_parts, strict=True) + ] + return linear(inputs, torch.cat(folded_parts, dim=0), base.bias) + + output_parts = list(base(inputs).split(tuple(int(size) for size in projection_sizes), dim=-1)) + for index, adapter in enumerate(adapters): + if adapter is not None: + output_parts[index] = output_parts[index] + adapter(inputs).to(output_parts[index].dtype) + return torch.cat(output_parts, dim=-1) diff --git a/src/xorl/models/layers/moe/backend/eager.py b/src/xorl/models/layers/moe/backend/eager.py index 3f2f930a..1d4e4df5 100644 --- a/src/xorl/models/layers/moe/backend/eager.py +++ b/src/xorl/models/layers/moe/backend/eager.py @@ -65,40 +65,6 @@ def eager_expert_forward( return out -def eager_expert_forward_fp64( - hidden_states: torch.Tensor, - expert_idx: int, - gate_proj: torch.Tensor, - up_proj: torch.Tensor, - down_proj: torch.Tensor, -) -> torch.Tensor: - """fp64-accumulate expert forward for the K3 train/serve parity mode. - - The cross-engine MoE divergence is reduction-ORDER sensitivity: xorl's and - SGLang's expert GEMMs tile the K reduction differently, so their fp32 - accumulators round differently (~1 bf16 ULP on ~0.1% of elements). fp64 - accumulation of bf16 products is order-invariant far below bf16 resolution - (as verified by the cross-engine parity tests), so engines - that both accumulate in fp64 produce identical bf16 outputs regardless of - tiling. - - Cast contract (must stay matched to the serving-side fp64 mode and to the - offline reference): bf16 inputs/weights upcast losslessly to fp64; gate/up - GEMMs, SiLU (``z * sigmoid(z)``), and the gating product all in fp64; the - activation product is cast to bf16 (the standard fused-MoE intermediate - cast point); the down GEMM re-upcasts and accumulates in fp64. Returns the - UNWEIGHTED fp64 down output — the caller applies routing weights and the - cross-expert combine in fp64 and casts to bf16 once at the end. - - Gated SiLU without biases only (Qwen3-MoE family). - """ - x64 = hidden_states.to(torch.float64) - gate = x64 @ gate_proj[expert_idx].to(torch.float64) - up = x64 @ up_proj[expert_idx].to(torch.float64) - h = ((gate * torch.sigmoid(gate)) * up).to(torch.bfloat16) - return h.to(torch.float64) @ down_proj[expert_idx].to(torch.float64) - - def _counts_from_cumsum(cumsum: torch.Tensor, num_experts: int) -> list[int]: """Convert inclusive cumsum token counts to per-expert counts.""" counts = [] diff --git a/src/xorl/models/layers/moe/experts.py b/src/xorl/models/layers/moe/experts.py index 01a1d459..818ecf5f 100644 --- a/src/xorl/models/layers/moe/experts.py +++ b/src/xorl/models/layers/moe/experts.py @@ -23,15 +23,8 @@ logger = logging.getLogger(__name__) _MOE_SGLANG_FUSED_EXPERTS_ENV = "XORL_MOE_SGLANG_FUSED_EXPERTS" -_MOE_SGLANG_FUSED_EXPERTS_SLOT_COMBINE_ENV = "XORL_MOE_SGLANG_FUSED_EXPERTS_SLOT_COMBINE" -_MOE_SGLANG_FUSED_EXPERTS_CACHE_ENV = "XORL_MOE_SGLANG_FUSED_EXPERTS_CACHE_WEIGHTS" _MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE_ENV = "XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE" _MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODES = ("transient", "cached", "strided") -_SG_LANG_MOE_TP_SIM_ENV = "XORL_SGLANG_MOE_TP_SIM" -_SG_LANG_MOE_TP_SIM_SIZE_ENV = "XORL_SGLANG_MOE_TP_SIM_SIZE" -_SG_LANG_MOE_TP_SIM_LAYERS_ENV = "XORL_SGLANG_MOE_TP_SIM_LAYERS" -_SG_LANG_MOE_TP_SIM_BF16_REDUCE_ENV = "XORL_SGLANG_MOE_TP_SIM_BF16_REDUCE" -_SG_LANG_MOE_TP_SIM_CARRY_SHARDS_ENV = "XORL_SGLANG_MOE_TP_SIM_CARRY_SHARDS" def _flag_enabled(name: str) -> bool: @@ -62,19 +55,6 @@ def _env_float_or_none(name: str) -> float | None: return None -def _env_layer_filter_enabled(name: str, layer_idx: int | None) -> bool: - raw = os.environ.get(name, "").strip() - if not raw or raw.lower() in {"all", "*"}: - return True - if layer_idx is None: - return False - try: - enabled_layers = {int(item.strip()) for item in raw.split(",") if item.strip()} - except ValueError as exc: - raise ValueError(f"Invalid {name}={raw!r}; expected comma-separated layer indices") from exc - return layer_idx in enabled_layers - - _MOE_SGLANG_FUSED_EXPERTS_AUTO_LOGGED = False _MOE_SGLANG_FUSED_EXPERTS_STACK_AVAILABLE: bool | None = None @@ -179,47 +159,6 @@ def moe_sglang_fused_experts_enabled( return True -def moe_sglang_fused_experts_slot_combine_enabled() -> bool: - """Formal-guarantee EP combine variant for the serving-kernel parity mode. - - ``XORL_MOE_SGLANG_FUSED_EXPERTS_SLOT_COMBINE=1`` (sub-flag of - ``XORL_MOE_SGLANG_FUSED_EXPERTS``, only honored while that flag is on and - EP dispatch is ``alltoall``) replaces the alltoall fp32 scatter-add combine - with a slot-order gather + sgl_kernel ``moe_sum_reduce`` — the exact top-k - combine kernel the serving engine runs at topk>1. The default scatter-add - combine is empirically bit-identical on validated data; this variant makes - the combine-tree match structural rather than empirical. - """ - return _flag_enabled(_MOE_SGLANG_FUSED_EXPERTS_SLOT_COMBINE_ENV) - - -def moe_sglang_fused_experts_cache_enabled() -> bool: - """Opt-in cache for the serving-layout weight transposes (default off). - - ``XORL_MOE_SGLANG_FUSED_EXPERTS_CACHE_WEIGHTS=1`` keeps the transposed - ``w13``/``w2`` copies alive on the owning :class:`MoEExperts` module - instead of re-materializing them per forward (the transposes dominate the - parity-mode tax: ~4.5 ms vs ~0.45 ms kernel per q30 layer). - - Invalidation contract: entries are keyed on the source parameter's - ``(data_ptr, _version, shape, dtype)``, so plain in-place optimizer updates - invalidate automatically. Under FSDP2 the unsharded parameter buffer can be - re-materialized at the same address with a fresh version counter, which can - FALSELY HIT after a sharded optimizer step — call - :meth:`MoEExperts.invalidate_sglang_fused_weight_cache` after each step, or - leave the cache off (default) under FSDP2 training. - - Memory: the cache duplicates the full expert weights in the serving layout - (~1.1 GiB per q30 layer, ~54 GiB for all 48 layers per rank at bf16) — size - it against free HBM before enabling model-wide. - - Legacy alias: equivalent to ``XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE=cached``; - an explicit weight mode takes precedence (see - :func:`moe_sglang_fused_experts_weight_mode`). - """ - return _flag_enabled(_MOE_SGLANG_FUSED_EXPERTS_CACHE_ENV) - - def moe_sglang_fused_experts_weight_mode() -> str: """How the parity mode presents xorl's GKN weights to the serving kernel. @@ -236,9 +175,10 @@ def moe_sglang_fused_experts_weight_mode() -> str: layer; the fallback if the vendored orchestration cannot bind to the sglang tree on PYTHONPATH after an upstream restructure); - ``cached``: keep the transposed copies alive on the module (fast, but - ~1.1 GiB per q30 layer — see :func:`moe_sglang_fused_experts_cache_enabled` - for the FSDP2 invalidation contract; that legacy env is an alias for this - mode when no explicit mode is set). + ~1.1 GiB per q30 layer). Cache entries key the source parameter's pointer, + version, shape, and dtype; FSDP2 callers must explicitly invalidate after + optimizer steps because an unsharded buffer can be rematerialized at the + same address with a fresh version counter. """ raw = os.environ.get(_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE_ENV, "").strip().lower() if raw: @@ -248,7 +188,7 @@ def moe_sglang_fused_experts_weight_mode() -> str: f"expected one of {', '.join(_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODES)}" ) return raw - return "cached" if moe_sglang_fused_experts_cache_enabled() else "strided" + return "strided" def _validate_sglang_swiglu_limit(swiglu_limit: float | None) -> None: @@ -968,23 +908,6 @@ def backward(ctx, grad_output): ) -def _sglang_moe_tp_sim_accumulate(output: torch.Tensor, shard_output: torch.Tensor) -> torch.Tensor: - if _flag_enabled(_SG_LANG_MOE_TP_SIM_BF16_REDUCE_ENV): - return output.to(torch.bfloat16) + shard_output.to(torch.bfloat16) - return output + shard_output.to(output.dtype) - - -def _attach_sglang_moe_tp_shards( - output: torch.Tensor, - shard_outputs: list[torch.Tensor], - original_shape: torch.Size, -) -> torch.Tensor: - if not _flag_enabled(_SG_LANG_MOE_TP_SIM_CARRY_SHARDS_ENV) or not shard_outputs: - return output - output._xorl_sglang_moe_tp_shards = tuple(shard.reshape(original_shape) for shard in shard_outputs) - return output - - def _deepep_parity_diagnostic_enabled() -> bool: return _flag_enabled("XORL_DEEPEP_PARITY_DIAGNOSTIC") @@ -1668,9 +1591,6 @@ def forward( if self.fp8_training_enabled and self.moe_implementation != "quack": raise NotImplementedError("FP8 grouped MoE compute currently requires moe_implementation='quack'") - if self.sglang_moe_tp_sim_enabled(parallel_state) and expert_idx is None: - return self._sglang_moe_tp_sim_forward(hidden_states, routing_weights, selected_experts, parallel_state) - if self.moe_implementation == "eager": fn = MOE_EXPERT_BACKENDS[self.moe_implementation] assert expert_idx is not None @@ -1714,566 +1634,6 @@ def forward( gated=self.gated, ) - def sglang_moe_tp_sim_enabled(self, parallel_state=None) -> bool: - """Whether to simulate SGLang's MoE expert kernel/reduce order. - - This is a K3 parity diagnostic path for no-EP Qwen3-MoE replays. Under - TP, xorl currently keeps MoE experts full-local while SGLang shards the - expert intermediate dimension and all-reduces the hidden output. Under - TP1/FSDP, the same env-gated path is useful for testing SGLang's local - fused expert kernel/order without changing production MoE defaults. - """ - if not _flag_enabled(_SG_LANG_MOE_TP_SIM_ENV): - return False - if parallel_state is None: - from xorl.distributed.parallel_state import get_parallel_state # noqa: PLC0415 - - parallel_state = get_parallel_state() - if getattr(parallel_state, "ep_enabled", False): - return False - if not _env_layer_filter_enabled(_SG_LANG_MOE_TP_SIM_LAYERS_ENV, getattr(self, "layer_idx", None)): - return False - return True - - def _sglang_moe_tp_sim_forward( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - parallel_state, - ) -> torch.Tensor: - mode = os.environ.get(_SG_LANG_MOE_TP_SIM_ENV, "0").strip().lower() - if mode == "cache": - return self._sglang_moe_tp_sim_cache_forward( - hidden_states, routing_weights, selected_experts, parallel_state - ) - if mode in {"triton", "backend"}: - return self._sglang_moe_tp_sim_triton_forward( - hidden_states, - routing_weights, - selected_experts, - parallel_state, - ) - if mode in {"triton_sgl_reduce", "triton_sglang_reduce", "sglang_reduce"}: - return self._sglang_moe_tp_sim_triton_sglang_reduce_forward( - hidden_states, - routing_weights, - selected_experts, - parallel_state, - ) - if mode in {"deep_gemm", "deepgemm", "dg"}: - return self._sglang_moe_tp_sim_deep_gemm_forward( - hidden_states, - routing_weights, - selected_experts, - parallel_state, - ) - if mode in {"sglang", "sgl_kernel", "fused"}: - return self._sglang_moe_tp_sim_sglang_forward( - hidden_states, - routing_weights, - selected_experts, - parallel_state, - ) - if mode in {"sglang_runner", "sgl_runner", "runner"}: - return self._sglang_moe_tp_sim_sglang_runner_forward( - hidden_states, - routing_weights, - selected_experts, - parallel_state, - ) - return self._sglang_moe_tp_sim_direct_forward(hidden_states, routing_weights, selected_experts, parallel_state) - - def _validate_sglang_moe_tp_sim_inputs( - self, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - parallel_state, - ) -> tuple[int, int]: - if routing_weights is None or selected_experts is None: - raise ValueError(f"{_SG_LANG_MOE_TP_SIM_ENV}=1 requires routing_weights and selected_experts") - if not self.gated: - raise NotImplementedError(f"{_SG_LANG_MOE_TP_SIM_ENV}=1 currently supports gated MoE experts only") - if self.down_bias is not None: - raise NotImplementedError(f"{_SG_LANG_MOE_TP_SIM_ENV}=1 does not support down_bias") - - tp_size = max(1, int(getattr(parallel_state, "tp_size", 1))) - requested_tp_size = _env_int(_SG_LANG_MOE_TP_SIM_SIZE_ENV, 0) - if requested_tp_size > 0: - tp_size = requested_tp_size - if self.intermediate_size % tp_size != 0: - raise ValueError( - f"Cannot simulate SGLang MoE TP with intermediate_size={self.intermediate_size} " - f"not divisible by tp_size={tp_size}" - ) - return tp_size, self.intermediate_size // tp_size - - def _sglang_moe_tp_sim_direct_forward( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - parallel_state, - ) -> torch.Tensor: - """Compute experts as TP-intermediate shards, then sum shard outputs. - - SGLang's unquantized Qwen3-MoE FusedMoE shards w1/w3 on the expert - intermediate output dimension and w2 on the matching input dimension. - Each TP rank computes a hidden-sized partial output and the model path - all-reduces those partials. Since xorl has full expert tensors on every - TP rank today, this parity path computes each shard locally and adds the - shard outputs in rank order. - """ - tp_size, shard_intermediate = self._validate_sglang_moe_tp_sim_inputs( - routing_weights, - selected_experts, - parallel_state, - ) - original_shape = hidden_states.shape - hidden_flat = hidden_states.reshape(-1, int(hidden_states.shape[-1])) - selected_flat = selected_experts.reshape(hidden_flat.shape[0], -1) - routing_flat = routing_weights.reshape(hidden_flat.shape[0], -1) - hidden_dim = int(hidden_flat.shape[-1]) - compute_dtype = torch.bfloat16 if hidden_flat.dtype in {torch.bfloat16, torch.float16} else hidden_flat.dtype - - from xorl.ops.moe.activations import apply_moe_activation # noqa: PLC0415 - - output = hidden_flat.new_zeros(hidden_flat.shape[0], hidden_dim) - shard_outputs = [] - for tp_rank in range(tp_size): - start = tp_rank * shard_intermediate - end = start + shard_intermediate - shard_output = hidden_flat.new_zeros(hidden_flat.shape[0], hidden_dim) - - for expert_idx in range(self.num_experts): - mask = selected_flat == expert_idx - if not bool(mask.any().item()): - continue - token_rows, topk_slots = mask.nonzero(as_tuple=True) - tokens = hidden_flat.index_select(0, token_rows).to(compute_dtype) - - gate = tokens.matmul(self.gate_up_proj[expert_idx, :, start:end].to(compute_dtype)) - up = tokens.matmul( - self.gate_up_proj[ - expert_idx, - :, - self.intermediate_size + start : self.intermediate_size + end, - ].to(compute_dtype) - ) - if self.swiglu_limit > 0: - gate = gate.clamp(-self.swiglu_limit, self.swiglu_limit) - - activated = apply_moe_activation(self.hidden_act, gate, up) - expert_out = activated.matmul(self.down_proj[expert_idx, start:end, :].to(compute_dtype)) - expert_out = expert_out * routing_flat[token_rows, topk_slots].to(expert_out.dtype).unsqueeze(-1) - shard_output.index_add_(0, token_rows, expert_out.to(shard_output.dtype)) - - output = _sglang_moe_tp_sim_accumulate(output, shard_output) - shard_outputs.append(shard_output) - - result = output.reshape(original_shape) - return _attach_sglang_moe_tp_shards(result, shard_outputs, original_shape) - - def _sglang_moe_tp_sim_cache_forward( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - parallel_state, - ) -> torch.Tensor: - """Diagnostic cache-order variant of the SGLang MoE TP simulation.""" - tp_size, shard_intermediate = self._validate_sglang_moe_tp_sim_inputs( - routing_weights, - selected_experts, - parallel_state, - ) - original_shape = hidden_states.shape - hidden_flat = hidden_states.reshape(-1, int(hidden_states.shape[-1])) - selected_flat = selected_experts.reshape(hidden_flat.shape[0], -1) - routing_flat = routing_weights.reshape(hidden_flat.shape[0], -1) - hidden_dim = int(hidden_flat.shape[-1]) - compute_dtype = torch.bfloat16 if hidden_flat.dtype in {torch.bfloat16, torch.float16} else hidden_flat.dtype - - from xorl.ops.moe.activations import apply_moe_activation # noqa: PLC0415 - - output = hidden_flat.new_zeros(hidden_flat.shape[0], hidden_dim) - shard_outputs = [] - for tp_rank in range(tp_size): - start = tp_rank * shard_intermediate - end = start + shard_intermediate - topk = selected_flat.shape[1] - num_assignments = hidden_flat.shape[0] * topk - gate_up_cache = hidden_flat.new_zeros(num_assignments, 2 * shard_intermediate) - down_cache = hidden_flat.new_zeros(hidden_flat.shape[0], topk, hidden_dim) - - for expert_idx in range(self.num_experts): - mask = selected_flat == expert_idx - if not bool(mask.any().item()): - continue - token_rows, topk_slots = mask.nonzero(as_tuple=True) - assignment_rows = token_rows * topk + topk_slots - tokens = hidden_flat.index_select(0, token_rows).to(compute_dtype) - - gate = tokens.matmul(self.gate_up_proj[expert_idx, :, start:end].to(compute_dtype)) - up = tokens.matmul( - self.gate_up_proj[ - expert_idx, - :, - self.intermediate_size + start : self.intermediate_size + end, - ].to(compute_dtype) - ) - if self.gate_up_bias is not None: - gate = gate + self.gate_up_bias[expert_idx, start:end].to(compute_dtype) - up = up + self.gate_up_bias[ - expert_idx, - self.intermediate_size + start : self.intermediate_size + end, - ].to(compute_dtype) - - gate_up_cache[assignment_rows, :shard_intermediate] = gate.to(gate_up_cache.dtype) - gate_up_cache[assignment_rows, shard_intermediate:] = up.to(gate_up_cache.dtype) - - gate = gate_up_cache[:, :shard_intermediate].to(compute_dtype) - up = gate_up_cache[:, shard_intermediate:].to(compute_dtype) - if self.swiglu_limit > 0: - gate = gate.clamp(-self.swiglu_limit, self.swiglu_limit) - activation_cache = apply_moe_activation(self.hidden_act, gate, up).to(gate_up_cache.dtype) - - for expert_idx in range(self.num_experts): - mask = selected_flat == expert_idx - if not bool(mask.any().item()): - continue - token_rows, topk_slots = mask.nonzero(as_tuple=True) - assignment_rows = token_rows * topk + topk_slots - - activated = activation_cache.index_select(0, assignment_rows).to(compute_dtype) - expert_out = activated.matmul(self.down_proj[expert_idx, start:end, :].to(compute_dtype)) - expert_out = expert_out * routing_flat[token_rows, topk_slots].to(expert_out.dtype).unsqueeze(-1) - down_cache[token_rows, topk_slots, :] = expert_out.to(down_cache.dtype) - - shard_output = down_cache.to(torch.float32).sum(dim=1).to(down_cache.dtype) - - output = _sglang_moe_tp_sim_accumulate(output, shard_output) - shard_outputs.append(shard_output) - - result = output.reshape(original_shape) - return _attach_sglang_moe_tp_shards(result, shard_outputs, original_shape) - - def _sglang_moe_tp_sim_triton_forward( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - parallel_state, - ) -> torch.Tensor: - """Diagnostic TP-shard simulation using xorl's local Triton MoE backend.""" - tp_size, shard_intermediate = self._validate_sglang_moe_tp_sim_inputs( - routing_weights, - selected_experts, - parallel_state, - ) - original_shape = hidden_states.shape - hidden_flat = hidden_states.reshape(-1, int(hidden_states.shape[-1])) - selected_flat = selected_experts.reshape(hidden_flat.shape[0], -1) - routing_flat = routing_weights.reshape(hidden_flat.shape[0], -1) - - from xorl.ops.moe.triton import triton_moe_forward # noqa: PLC0415 - - output = hidden_flat.new_zeros(hidden_flat.shape) - shard_outputs = [] - for tp_rank in range(tp_size): - start = tp_rank * shard_intermediate - end = start + shard_intermediate - gate_proj = self.gate_up_proj[:, :, start:end].contiguous() - up_proj = self.gate_up_proj[ - :, - :, - self.intermediate_size + start : self.intermediate_size + end, - ].contiguous() - gate_up_proj = torch.cat([gate_proj, up_proj], dim=-1).contiguous() - down_proj = self.down_proj[:, start:end, :].contiguous() - - shard_output = triton_moe_forward( - module=None, - num_experts=self.num_experts, - routing_weights=routing_flat, - selected_experts=selected_flat, - hidden_states=hidden_flat, - gate_proj=gate_proj, - up_proj=up_proj, - down_proj=down_proj, - gate_up_proj=gate_up_proj, - hidden_act=self.hidden_act, - swiglu_limit=self.swiglu_limit, - gated=self.gated, - ) - output = _sglang_moe_tp_sim_accumulate(output, shard_output) - shard_outputs.append(shard_output) - - result = output.reshape(original_shape) - return _attach_sglang_moe_tp_shards(result, shard_outputs, original_shape) - - @staticmethod - def _sglang_topk_sum_reduce(per_slot: torch.Tensor) -> torch.Tensor: - """Mirror SGLang's top-k MoE output reduction for topk > 1. - - SGLang's deterministic Triton path accumulates the top-k slot outputs - into fp32 in slot order and casts once to the output dtype. PyTorch - ``sum(dim=1)`` can choose a different reduction tree, so keep this - explicit for parity diagnostics. - """ - topk = int(per_slot.shape[1]) - if topk == 1: - return per_slot[:, 0, :] - accumulator = per_slot[:, 0, :].to(torch.float32) - for topk_idx in range(1, topk): - accumulator = accumulator + per_slot[:, topk_idx, :].to(torch.float32) - return accumulator.to(per_slot.dtype) - - def _sglang_moe_tp_sim_triton_sglang_reduce_forward( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - parallel_state, - ) -> torch.Tensor: - """Use xorl grouped GEMMs but SGLang's deterministic top-k combine order.""" - tp_size, shard_intermediate = self._validate_sglang_moe_tp_sim_inputs( - routing_weights, - selected_experts, - parallel_state, - ) - original_shape = hidden_states.shape - hidden_flat = hidden_states.reshape(-1, int(hidden_states.shape[-1])) - selected_flat = selected_experts.reshape(hidden_flat.shape[0], -1) - routing_flat = routing_weights.reshape(hidden_flat.shape[0], -1) - - from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_nk # noqa: PLC0415 - from xorl.ops.group_gemm.kernel.moe import ( # noqa: PLC0415 - expert_histogram, - moe_index_compute, - moe_scatter, - ) - from xorl.ops.moe.activations import apply_moe_activation # noqa: PLC0415 - - splits = expert_histogram(selected_flat, self.num_experts) - cumsum_t = torch.cumsum(splits, dim=0) - scatter_index = moe_index_compute(selected_flat, cumsum_t) - scatter_output = moe_scatter(hidden_flat, scatter_index) - max_m = scatter_output.shape[0] - - output = hidden_flat.new_zeros(hidden_flat.shape) - shard_outputs = [] - for tp_rank in range(tp_size): - start = tp_rank * shard_intermediate - end = start + shard_intermediate - gate_proj = self.gate_up_proj[:, :, start:end].contiguous() - up_proj = self.gate_up_proj[ - :, - :, - self.intermediate_size + start : self.intermediate_size + end, - ].contiguous() - gate_up_proj = torch.cat([gate_proj, up_proj], dim=-1).contiguous() - down_proj = self.down_proj[:, start:end, :].contiguous() - - gate_up_output = group_gemm_same_nk( - a=scatter_output, - b=gate_up_proj, - cumsum_M=cumsum_t, - max_M=max_m, - ) - gate, up = gate_up_output.split(shard_intermediate, dim=-1) - if self.swiglu_limit > 0: - gate = gate.clamp(-self.swiglu_limit, self.swiglu_limit) - activated = apply_moe_activation(self.hidden_act, gate, up) - - down_output = group_gemm_same_nk( - a=activated, - b=down_proj, - cumsum_M=cumsum_t, - max_M=max_m, - ) - per_slot = down_output[scatter_index.flatten()].reshape( - hidden_flat.shape[0], - selected_flat.shape[1], - -1, - ) - weighted = per_slot * routing_flat.to(per_slot.dtype).unsqueeze(-1) - shard_output = self._sglang_topk_sum_reduce(weighted) - output = _sglang_moe_tp_sim_accumulate(output, shard_output) - shard_outputs.append(shard_output) - - result = output.reshape(original_shape) - return _attach_sglang_moe_tp_shards(result, shard_outputs, original_shape) - - @staticmethod - def _deep_gemm_group_gemm_same_nk( - *, - a: torch.Tensor, - b: torch.Tensor, - cumsum_M: torch.Tensor, - ) -> torch.Tensor: - """Run DeepGEMM BF16 contiguous grouped GEMM using xorl's compact layout. - - xorl's grouped-GEMM helpers keep expert rows compact. DeepGEMM's - contiguous grouped layout requires each expert segment to be padded to - the kernel's M alignment, with invalid pad rows marked by ``m_indices``. - ``b`` follows xorl's [expert, K, N] layout and is transposed for - DeepGEMM's NT contract [expert, N, K]. - """ - if a.dtype != torch.bfloat16 or b.dtype != torch.bfloat16: - raise RuntimeError("DeepGEMM BF16 grouped MoE diagnostic requires bf16 tensors") - if not a.is_cuda or not b.is_cuda: - raise RuntimeError("DeepGEMM BF16 grouped MoE diagnostic requires CUDA tensors") - - try: - import deep_gemm # noqa: PLC0415 - except ImportError as exc: - raise ImportError(f"{_SG_LANG_MOE_TP_SIM_ENV}=deep_gemm requires the optional deep_gemm package") from exc - - if a.shape[0] == 0: - return a.new_empty((0, int(b.shape[2]))) - - alignment = int(deep_gemm.get_mk_alignment_for_contiguous_layout()) - starts = torch.cat([cumsum_M.new_zeros(1), cumsum_M[:-1]]).detach().to(torch.int64).cpu().tolist() - ends = cumsum_M.detach().to(torch.int64).cpu().tolist() - chunks: list[torch.Tensor] = [] - m_indices: list[int] = [] - ranges: list[tuple[int, int, int, int]] = [] - padded_start = 0 - for expert_idx, (start, end) in enumerate(zip(starts, ends)): - count = int(end) - int(start) - if count <= 0: - continue - aligned_count = ((count + alignment - 1) // alignment) * alignment - chunk = a[int(start) : int(end)] - if aligned_count != count: - padded = a.new_zeros((aligned_count, int(a.shape[1]))) - padded[:count].copy_(chunk) - chunk = padded - chunks.append(chunk.contiguous()) - m_indices.extend([expert_idx] * count) - m_indices.extend([-1] * (aligned_count - count)) - ranges.append((int(start), int(end), padded_start, padded_start + count)) - padded_start += aligned_count - - out = a.new_empty((int(a.shape[0]), int(b.shape[2]))) - if not chunks: - return out - - padded_a = torch.cat(chunks, dim=0).contiguous() - padded_out = a.new_empty((int(padded_a.shape[0]), int(b.shape[2]))) - m_indices_tensor = torch.tensor(m_indices, dtype=torch.int32, device=a.device) - deep_gemm.m_grouped_bf16_gemm_nt_contiguous( - padded_a, - b.transpose(1, 2).contiguous(), - padded_out, - m_indices_tensor, - ) - for start, end, padded_valid_start, padded_valid_end in ranges: - out[start:end].copy_(padded_out[padded_valid_start:padded_valid_end]) - return out - - @staticmethod - def _sglang_stable_moe_slot_order( - selected_experts: torch.Tensor, - num_experts: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Return SGLang's stable expert-major compact slot order. - - SGLang's BF16 DeepGEMM diagnostic compacts flattened token/top-k slots - by stable-sorting on expert id. xorl's generic Triton scatter index uses - relaxed atomic reservations, which can place the same valid slot at a - different compact M row. DeepGEMM output can vary at one-ulp scale with - that row placement, so the parity diagnostic must use SGLang's order. - """ - flat_experts = selected_experts.reshape(-1).to(torch.int64) - valid_mask = (flat_experts >= 0) & (flat_experts < int(num_experts)) - if not bool(valid_mask.any().item()): - counts = torch.zeros(int(num_experts), dtype=torch.int32, device=selected_experts.device) - return torch.empty(0, dtype=torch.long, device=selected_experts.device), counts - - valid_positions = torch.nonzero(valid_mask, as_tuple=False).flatten() - sort_keys = flat_experts.index_select(0, valid_positions) - try: - sort_relative = torch.argsort(sort_keys, stable=True) - except TypeError: - sort_relative = torch.argsort(sort_keys) - slot_order = valid_positions.index_select(0, sort_relative) - sorted_experts = flat_experts.index_select(0, slot_order) - counts = torch.bincount(sorted_experts, minlength=int(num_experts)).to(dtype=torch.int32) - cumsum_m = torch.cumsum(counts, dim=0).to(dtype=torch.int32) - return slot_order, cumsum_m - - def _sglang_moe_tp_sim_deep_gemm_forward( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - parallel_state, - ) -> torch.Tensor: - """Diagnostic TP-shard simulation using DeepGEMM BF16 grouped GEMM.""" - tp_size, shard_intermediate = self._validate_sglang_moe_tp_sim_inputs( - routing_weights, - selected_experts, - parallel_state, - ) - if self.gate_up_bias is not None: - raise NotImplementedError(f"{_SG_LANG_MOE_TP_SIM_ENV}=deep_gemm does not support gate_up_bias") - original_shape = hidden_states.shape - hidden_flat = hidden_states.reshape(-1, int(hidden_states.shape[-1])) - selected_flat = selected_experts.reshape(hidden_flat.shape[0], -1) - routing_flat = routing_weights.reshape(hidden_flat.shape[0], -1) - - from xorl.ops.moe.activations import apply_sglang_moe_activation # noqa: PLC0415 - - topk = int(selected_flat.shape[1]) - slot_order, cumsum_t = self._sglang_stable_moe_slot_order(selected_flat, self.num_experts) - if slot_order.numel() == 0: - return hidden_flat.new_zeros(hidden_flat.shape).reshape(original_shape) - - hidden_slots = hidden_flat.repeat_interleave(topk, dim=0) - scatter_output = hidden_slots.index_select(0, slot_order).contiguous() - - output = hidden_flat.new_zeros(hidden_flat.shape) - shard_outputs = [] - for tp_rank in range(tp_size): - start = tp_rank * shard_intermediate - end = start + shard_intermediate - gate_proj = self.gate_up_proj[:, :, start:end].contiguous() - up_proj = self.gate_up_proj[ - :, - :, - self.intermediate_size + start : self.intermediate_size + end, - ].contiguous() - gate_up_proj = torch.cat([gate_proj, up_proj], dim=-1).contiguous() - down_proj = self.down_proj[:, start:end, :].contiguous() - - gate_up_output = self._deep_gemm_group_gemm_same_nk( - a=scatter_output, - b=gate_up_proj, - cumsum_M=cumsum_t, - ) - gate, up = gate_up_output.split(shard_intermediate, dim=-1) - if self.swiglu_limit > 0: - gate = gate.clamp(-self.swiglu_limit, self.swiglu_limit) - activated = apply_sglang_moe_activation(self.hidden_act, gate, up).contiguous() - - down_output = self._deep_gemm_group_gemm_same_nk( - a=activated, - b=down_proj, - cumsum_M=cumsum_t, - ) - per_slot_flat = hidden_flat.new_zeros((hidden_flat.shape[0] * topk, hidden_flat.shape[-1])) - per_slot_flat.index_copy_(0, slot_order, down_output) - per_slot = per_slot_flat.reshape(hidden_flat.shape[0], topk, -1) - weighted = per_slot * routing_flat.to(per_slot.dtype).unsqueeze(-1) - shard_output = self._sglang_topk_sum_reduce(weighted) - output = _sglang_moe_tp_sim_accumulate(output, shard_output) - shard_outputs.append(shard_output) - - result = output.reshape(original_shape) - return _attach_sglang_moe_tp_shards(result, shard_outputs, original_shape) - @staticmethod def _ensure_sglang_server_args() -> None: try: @@ -2281,7 +1641,7 @@ def _ensure_sglang_server_args() -> None: from sglang.srt.server_args import ServerArgs # noqa: PLC0415 except ImportError as exc: raise ImportError( - f"{_SG_LANG_MOE_TP_SIM_ENV}=sglang requires an environment with sglang and sgl_kernel installed" + f"{_MOE_SGLANG_FUSED_EXPERTS_ENV}=1 requires an environment with sglang and sgl_kernel installed" ) from exc try: @@ -2318,7 +1678,7 @@ def _load_sglang_fused_experts_impl(): ) except ImportError as exc: raise ImportError( - f"{_SG_LANG_MOE_TP_SIM_ENV}=sglang requires an environment with sglang and sgl_kernel installed" + f"{_MOE_SGLANG_FUSED_EXPERTS_ENV}=1 requires an environment with sglang and sgl_kernel installed" ) from exc MoEExperts._ensure_sglang_server_args() @@ -2331,29 +1691,6 @@ def _load_sglang_fused_experts_impl(): return fused_experts_impl_strided return fused_experts_impl - @staticmethod - def _load_sglang_moe_runner_stack(): - try: - from sglang.srt.layers.moe.moe_runner import MoeRunner, MoeRunnerConfig # noqa: PLC0415 - from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo # noqa: PLC0415 - from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput # noqa: PLC0415 - from sglang.srt.layers.moe.topk import StandardTopKOutput # noqa: PLC0415 - from sglang.srt.layers.moe.utils import MoeRunnerBackend # noqa: PLC0415 - except ImportError as exc: - raise ImportError( - f"{_SG_LANG_MOE_TP_SIM_ENV}=sglang_runner requires an environment with sglang and sgl_kernel installed" - ) from exc - - MoEExperts._ensure_sglang_server_args() - return ( - MoeRunner, - MoeRunnerBackend, - MoeRunnerConfig, - TritonMoeQuantInfo, - StandardDispatchOutput, - StandardTopKOutput, - ) - def sglang_ep_native_routed_partial( self, hidden_flat: torch.Tensor, @@ -2541,8 +1878,8 @@ def sglang_fused_experts_auto_supported(self) -> bool: def invalidate_sglang_fused_weight_cache(self) -> None: """Drop cached serving-layout weight transposes (see - :func:`moe_sglang_fused_experts_cache_enabled` for the invalidation - contract — call after optimizer steps under FSDP2).""" + :func:`moe_sglang_fused_experts_weight_mode`; call after optimizer + steps under FSDP2).""" cache = getattr(self, "_sglang_fused_weight_cache", None) if cache is not None: cache.clear() @@ -2617,9 +1954,9 @@ def sglang_fused_experts_ep_compute( ``w13 [E_local, 2I, H]`` / ``w2 [E_local, H, I]`` per :func:`moe_sglang_fused_experts_weight_mode`: a zero-copy transpose-view (strided mode, default), a transient transpose-copy (~150 MB per q30 - EP8 layer), or a cached copy (~7 GB/rank model-wide at q30 EP8 — see - :func:`moe_sglang_fused_experts_cache_enabled` for the FSDP2 - invalidation contract) — all three bit-identical. + EP8 layer), or a cached copy (~7 GB/rank model-wide at q30 EP8, with + explicit FSDP2 invalidation after optimizer steps) — all three + bit-identical. Trainable: when gradients are required the same serving forward runs under :class:`_SglangFusedExpertsEPTrainFunction`, whose backward is @@ -2720,238 +2057,6 @@ def _compute(permute_tokens, cumsum, _gate_up_proj, _down_proj, _intermediate_si return _compute - @staticmethod - def _sglang_fused_experts_pair_slot_order(selected_experts: torch.Tensor) -> torch.Tensor: - """(token, slot) -> pair-row order for the slot-combine variant. - - Pair rows return from the combine all-to-all in the local ``permute()`` - order: global-expert-major, token-ascending within each expert — i.e. - (expert, token) pairs sorted by (e, t). This is a deterministic function - of ``selected_experts``: returns ``order`` such that - ``slots_flat.index_copy_(0, order, pair_rows)`` lands arrival row ``r`` - at flat slot index ``token * topk + slot``. - """ - num_tokens, topk = selected_experts.shape - token_idx = torch.arange(num_tokens, device=selected_experts.device).repeat_interleave(topk) - keys = selected_experts.reshape(-1).to(torch.int64) * num_tokens + token_idx - return torch.argsort(keys) - - def _sglang_fused_experts_slot_combine(self, expert_output, ctx, dispatch_kwargs, parallel_state): - """Formal-guarantee combine variant: slot-order gather + sgl_kernel ``moe_sum_reduce``. - - The default alltoall combine (fp32 ``scatter_add_`` over arrival-order - pair rows) is empirically order-invariant on bit-identical per-slot - contributions, but ``scatter_add_``'s addition order is not formally - specified. This variant re-orders the returned pair rows into SGLang's - ``[num_tokens, topk, hidden]`` slot layout and reduces with sgl_kernel's - ``moe_sum_reduce`` — the exact combine kernel the serving engine executes - at topk>1. The return-path re-sort + all-to-all is identical to - ``tokens_post_all2all``; only the final unpermute/reduce differs. - - Scoring-only: ``moe_sum_reduce`` writes through an out-parameter with no - autograd, so a grad-requiring expert output would silently detach the - graph — reject it loudly (train with the default scatter-add combine). - """ - if torch.is_grad_enabled() and expert_output.requires_grad: - raise NotImplementedError( - f"{_MOE_SGLANG_FUSED_EXPERTS_SLOT_COMBINE_ENV}=1 is scoring-only (sgl_kernel moe_sum_reduce " - "has no autograd); train with the default alltoall scatter-add combine instead" - ) - from sgl_kernel import moe_sum_reduce # noqa: PLC0415 - - from xorl.distributed.moe.alltoall import _expert_chunk_unpermute_order # noqa: PLC0415 - from xorl.distributed.moe.comm import all_to_all # noqa: PLC0415 - from xorl.distributed.moe.utils import sort_chunks_by_idxs # noqa: PLC0415 - - ep_group = parallel_state.ep_group - expert_output = sort_chunks_by_idxs( - expert_output, - ctx.num_tokens_per_expert.T.ravel(), - _expert_chunk_unpermute_order(ctx.num_experts, ep_group.size()), - ) - pair_rows = all_to_all(ep_group, expert_output, ctx.input_splits, ctx.output_splits) - - selected_experts = dispatch_kwargs["selected_experts"] - selected_flat = selected_experts.reshape(-1, selected_experts.shape[-1]) - num_tokens, topk = selected_flat.shape - if pair_rows.shape[0] != num_tokens * topk: - raise NotImplementedError( - f"{_MOE_SGLANG_FUSED_EXPERTS_SLOT_COMBINE_ENV}=1 requires unique expert selections per " - f"token (got {pair_rows.shape[0]} pair rows for {num_tokens}x{topk} slots)" - ) - order = self._sglang_fused_experts_pair_slot_order(selected_flat) - slots = pair_rows.new_empty((num_tokens * topk, pair_rows.shape[-1])) - slots.index_copy_(0, order, pair_rows) - slots = slots.reshape(num_tokens, topk, -1) - output = torch.empty(ctx.orig_shape, device=pair_rows.device, dtype=pair_rows.dtype) - moe_sum_reduce(slots.contiguous(), output, 1.0) - return output - - def _sglang_moe_tp_sim_sglang_forward( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - parallel_state, - ) -> torch.Tensor: - """Diagnostic TP-shard simulation using SGLang's fused experts kernel.""" - _validate_sglang_swiglu_limit(self.swiglu_limit) - tp_size, shard_intermediate = self._validate_sglang_moe_tp_sim_inputs( - routing_weights, - selected_experts, - parallel_state, - ) - original_shape = hidden_states.shape - hidden_flat = hidden_states.reshape(-1, int(hidden_states.shape[-1])).contiguous() - selected_flat = selected_experts.reshape(hidden_flat.shape[0], -1).contiguous() - routing_flat = routing_weights.reshape(hidden_flat.shape[0], -1).contiguous() - - fused_experts_impl = self._load_sglang_fused_experts_impl() - activation = "gelu" if self.hidden_act == "gelu_tanh" else self.hidden_act - - output = hidden_flat.new_zeros(hidden_flat.shape) - shard_outputs = [] - for tp_rank in range(tp_size): - start = tp_rank * shard_intermediate - end = start + shard_intermediate - gate_proj = self.gate_up_proj[:, :, start:end].transpose(1, 2).contiguous() - up_proj = ( - self.gate_up_proj[ - :, - :, - self.intermediate_size + start : self.intermediate_size + end, - ] - .transpose(1, 2) - .contiguous() - ) - w1 = torch.cat([gate_proj, up_proj], dim=1).contiguous() - w2 = self.down_proj[:, start:end, :].transpose(1, 2).contiguous() - b1 = None - if self.gate_up_bias is not None: - gate_bias = self.gate_up_bias[:, start:end] - up_bias = self.gate_up_bias[:, self.intermediate_size + start : self.intermediate_size + end] - b1 = torch.cat([gate_bias, up_bias], dim=1).contiguous() - - shard_output = fused_experts_impl( - hidden_flat, - w1, - w2, - routing_flat, - selected_flat, - b1=b1, - b2=None, - inplace=False, - activation=activation, - is_gated=self.gated, - apply_router_weight_on_input=False, - no_combine=False, - routed_scaling_factor=None, - gemm1_limit=None, - gate_up_interleaved=False, - filter_expert=False, - ) - output = _sglang_moe_tp_sim_accumulate(output, shard_output) - shard_outputs.append(shard_output) - - result = output.reshape(original_shape) - return _attach_sglang_moe_tp_shards(result, shard_outputs, original_shape) - - def _sglang_moe_tp_sim_sglang_runner_forward( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - parallel_state, - ) -> torch.Tensor: - """Diagnostic TP-shard simulation through SGLang's MoeRunner wrapper.""" - _validate_sglang_swiglu_limit(self.swiglu_limit) - tp_size, shard_intermediate = self._validate_sglang_moe_tp_sim_inputs( - routing_weights, - selected_experts, - parallel_state, - ) - original_shape = hidden_states.shape - hidden_flat = hidden_states.reshape(-1, int(hidden_states.shape[-1])).contiguous() - selected_flat = selected_experts.reshape(hidden_flat.shape[0], -1).contiguous() - routing_flat = routing_weights.reshape(hidden_flat.shape[0], -1).contiguous() - hidden_dim = int(hidden_flat.shape[-1]) - - ( - MoeRunner, - MoeRunnerBackend, - MoeRunnerConfig, - TritonMoeQuantInfo, - StandardDispatchOutput, - StandardTopKOutput, - ) = self._load_sglang_moe_runner_stack() - - activation = "gelu" if self.hidden_act == "gelu_tanh" else self.hidden_act - output = hidden_flat.new_zeros(hidden_flat.shape) - shard_outputs = [] - for tp_rank in range(tp_size): - start = tp_rank * shard_intermediate - end = start + shard_intermediate - gate_proj = self.gate_up_proj[:, :, start:end].transpose(1, 2).contiguous() - up_proj = ( - self.gate_up_proj[ - :, - :, - self.intermediate_size + start : self.intermediate_size + end, - ] - .transpose(1, 2) - .contiguous() - ) - w13 = torch.cat([gate_proj, up_proj], dim=1).contiguous() - w2 = self.down_proj[:, start:end, :].transpose(1, 2).contiguous() - b13 = None - if self.gate_up_bias is not None: - gate_bias = self.gate_up_bias[:, start:end] - up_bias = self.gate_up_bias[:, self.intermediate_size + start : self.intermediate_size + end] - b13 = torch.cat([gate_bias, up_bias], dim=1).contiguous() - - config = MoeRunnerConfig( - num_experts=self.num_experts, - num_local_experts=self.num_experts, - hidden_size=hidden_dim, - intermediate_size_per_partition=shard_intermediate, - layer_id=0, - top_k=int(selected_flat.shape[1]), - params_dtype=hidden_flat.dtype, - activation=activation, - apply_router_weight_on_input=False, - inplace=False, - no_combine=False, - routed_scaling_factor=None, - gemm1_alpha=None, - gemm1_clamp_limit=None, - is_gated=self.gated, - gate_up_interleaved=False, - ) - topk_output = StandardTopKOutput( - topk_weights=routing_flat, - topk_ids=selected_flat, - router_logits=None, - ) - dispatch_output = StandardDispatchOutput( - hidden_states=hidden_flat, - hidden_states_scale=None, - topk_output=topk_output, - ) - quant_info = TritonMoeQuantInfo( - w13_weight=w13, - w2_weight=w2, - b13=b13, - b2=None, - ) - runner = MoeRunner(MoeRunnerBackend.TRITON, config) - combine_input = runner.run(dispatch_output, quant_info) - shard_output = combine_input.hidden_states - output = output + shard_output.to(output.dtype) - shard_outputs.append(shard_output) - - result = output.reshape(original_shape) - return _attach_sglang_moe_tp_shards(result, shard_outputs, original_shape) - @torch.compiler.disable def _ep_forward( self, @@ -3012,7 +2117,7 @@ def _ep_forward( if sglang_fused_experts: # K3 parity mode: the serving-kernel expert compute overrides the # configured backend; dispatch and combine stay on the stock - # alltoall path (combine optionally swaps under the sub-flag below). + # alltoall path. compute_fn = self._sglang_fused_experts_ep_compute_fn() # Step 1: Dispatch tokens to expert-owning ranks @@ -3197,11 +2302,8 @@ def _ep_forward( ) # Step 3: Combine expert outputs back to original ranks - if sglang_fused_experts and moe_sglang_fused_experts_slot_combine_enabled(): - result = self._sglang_fused_experts_slot_combine(expert_output, ctx, dispatch_kwargs, parallel_state) - else: - combine_kwargs = self._build_combine_kwargs(expert_output, ctx, dispatch_kwargs, parallel_state) - result = combine_fn(**combine_kwargs) + combine_kwargs = self._build_combine_kwargs(expert_output, ctx, dispatch_kwargs, parallel_state) + result = combine_fn(**combine_kwargs) if deepep_diagnostic_id is not None: self._emit_deepep_parity_diagnostic( record_id=deepep_diagnostic_id, diff --git a/src/xorl/models/layers/moe/lora.py b/src/xorl/models/layers/moe/lora.py index e8b715c9..c5947cdd 100644 --- a/src/xorl/models/layers/moe/lora.py +++ b/src/xorl/models/layers/moe/lora.py @@ -333,10 +333,6 @@ def merge_weights(self) -> None: # Merged-forward exact-model contract lane # ------------------------------------------------------------------ - def sglang_moe_tp_sim_enabled(self, parallel_state) -> bool: - """TP-sim is outside the LoRA merged-forward envelope.""" - return False - def sglang_fused_experts_auto_supported(self) -> bool: """Auto-default eligibility mirror of :meth:`MoEExperts.sglang_fused_experts_auto_supported`: under the exact model program the adapted experts fold their delta @@ -348,9 +344,6 @@ def sglang_fused_experts_auto_supported(self) -> bool: and self.swiglu_limit == 0.0 ) - def invalidate_merged_weight_cache(self) -> None: - self._merged_weight_cache = {} - def _merged_weight_key(self) -> tuple: params = ( self.gate_proj_lora_A, diff --git a/src/xorl/models/layers/moe/moe_block.py b/src/xorl/models/layers/moe/moe_block.py index b4e40cb3..2c99be30 100644 --- a/src/xorl/models/layers/moe/moe_block.py +++ b/src/xorl/models/layers/moe/moe_block.py @@ -1,6 +1,5 @@ """MoE block: composes gate + router + experts.""" -import os from typing import Optional import torch @@ -14,19 +13,6 @@ from .routing_replay import RoutingReplay, get_replay_stage -_MOE_FP64_ACCUM_ENV = "XORL_MOE_FP64_ACCUM" - - -def _moe_fp64_accum_enabled() -> bool: - """K3 parity mode: fp64-accumulate eager expert GEMMs + weighted combine. - - Order-invariant MoE reductions (see ``eager_expert_forward_fp64``). Slow — - fp64 GEMMs run at H100 fp64 tensor-core rate (~15x below bf16) — so this is - an opt-in reconciliation mode, never a training default. - """ - return os.environ.get(_MOE_FP64_ACCUM_ENV, "").strip().lower() in {"1", "true", "yes", "on"} - - def _moe_bi_router_enabled(config=None) -> bool: """Whether this caller owns the exact batch-invariant router contract. @@ -71,33 +57,6 @@ def backward(ctx, grad_logits): return grad_hidden, grad_weight -_ROUTER_FP32_LAYERS_ENV = "XORL_MOE_ROUTER_FP32_LAYERS" - - -def _router_fp32_layers_enabled(layer_idx: int | None) -> bool: - value = os.environ.get(_ROUTER_FP32_LAYERS_ENV, "").strip().lower() - if value in {"", "0", "false", "no", "off", "none"}: - return False - if value in {"all", "*"}: - return True - if layer_idx is None: - return False - - for part in value.split(","): - part = part.strip() - if not part: - continue - if "-" in part: - start_raw, end_raw = part.split("-", 1) - start = int(start_raw) - end = int(end_raw) - if start <= layer_idx <= end or end <= layer_idx <= start: - return True - elif int(part) == layer_idx: - return True - return False - - class MoEBlock(nn.Module): """Mixture-of-Experts block. @@ -195,25 +154,6 @@ def _override_diagnostic_component(self, name: str, tensor: torch.Tensor) -> tor self._capture_diagnostic_component(f"{name}_override", tensor) return tensor - def _capture_moe_tp_shards(self, name: str, tensor: torch.Tensor) -> None: - capture = getattr(self, "_diagnostic_capture_component", None) - if not callable(capture): - return - tp_shards = getattr(tensor, "_xorl_sglang_moe_tp_shards", None) - if not tp_shards: - return - - shard_sum = None - for shard_idx, shard in enumerate(tp_shards): - self._capture_diagnostic_component(f"{name}_tp_shard_{shard_idx}", shard) - if shard_sum is None: - shard_sum = shard.clone() - else: - shard_sum = shard_sum + shard.to(shard_sum.dtype) - - if shard_sum is not None: - self._capture_diagnostic_component(f"{name}_tp_shard_sum", shard_sum) - def inject_lora( self, r: int = 16, @@ -330,11 +270,7 @@ def route(self, hidden_states: torch.Tensor): - router_logits: ``(num_tokens, num_experts)`` """ # Route (optionally upcast to fp32 for numerical alignment with SGLang) - router_fp32 = ( - getattr(self, "config", None) is not None - and getattr(self.config, "_router_fp32", False) - or _router_fp32_layers_enabled(getattr(self, "layer_idx", None)) - ) + router_fp32 = getattr(self, "config", None) is not None and getattr(self.config, "_router_fp32", False) if self._exact_batch_invariant_router or _moe_bi_router_enabled(getattr(self, "config", None)): router_logits = self._bi_router_logits(hidden_states) elif router_fp32 and not hasattr(self.gate, "fp8_block_size"): @@ -465,28 +401,15 @@ def forward_experts_only(self, hidden_states, routing_weights, selected_experts) parallel_state = get_parallel_state() if ( - _moe_fp64_accum_enabled() - and not parallel_state.ep_enabled - and not self.experts.sglang_moe_tp_sim_enabled(parallel_state) - ): - hidden_states = self._eager_forward_fp64(hidden_states, routing_weights, selected_experts) - elif ( moe_sglang_fused_experts_enabled(getattr(parallel_state, "ep_size", 1), hidden_states.device, self.experts) and not parallel_state.ep_enabled - and not self.experts.sglang_moe_tp_sim_enabled(parallel_state) ): hidden_states = self.experts.sglang_fused_experts_forward(hidden_states, routing_weights, selected_experts) else: hidden_states = self.experts(hidden_states, routing_weights, selected_experts) self._capture_diagnostic_component("moe_experts_output", hidden_states) - self._capture_moe_tp_shards("moe_experts_output", hidden_states) hidden_states = self._override_diagnostic_component("moe_experts_output", hidden_states) - tp_shards = getattr(hidden_states, "_xorl_sglang_moe_tp_shards", None) hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim) - if tp_shards is not None: - hidden_states._xorl_sglang_moe_tp_shards = tuple( - shard.reshape(batch_size, sequence_length, hidden_dim) for shard in tp_shards - ) return hidden_states def forward(self, hidden_states: torch.Tensor): @@ -513,42 +436,23 @@ def forward(self, hidden_states: torch.Tensor): parallel_state = get_parallel_state() if ( - _moe_fp64_accum_enabled() - and not parallel_state.ep_enabled - and not self.experts.sglang_moe_tp_sim_enabled(parallel_state) - ): - # K3 parity mode overrides the configured backend (eager/triton/...): - # the fp64 loop is the order-invariant reference for all of them. - final_hidden_states = self._eager_forward_fp64(flat_hidden_states, routing_weights, selected_experts) - elif ( moe_sglang_fused_experts_enabled( getattr(parallel_state, "ep_size", 1), flat_hidden_states.device, self.experts ) and not parallel_state.ep_enabled - and not self.experts.sglang_moe_tp_sim_enabled(parallel_state) ): # K3 parity mode: SGLang's serving kernel overrides the configured # backend so trainer and serving share one MoE reduction tree. final_hidden_states = self.experts.sglang_fused_experts_forward( flat_hidden_states, routing_weights, selected_experts ) - elif ( - self.moe_implementation == "eager" - and not parallel_state.ep_enabled - and not self.experts.sglang_moe_tp_sim_enabled(parallel_state) - ): + elif self.moe_implementation == "eager" and not parallel_state.ep_enabled: final_hidden_states = self._eager_forward(flat_hidden_states, routing_weights, selected_experts) else: final_hidden_states = self.experts(flat_hidden_states, routing_weights, selected_experts) self._capture_diagnostic_component("moe_experts_output", final_hidden_states) - self._capture_moe_tp_shards("moe_experts_output", final_hidden_states) final_hidden_states = self._override_diagnostic_component("moe_experts_output", final_hidden_states) - tp_shards = getattr(final_hidden_states, "_xorl_sglang_moe_tp_shards", None) final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) - if tp_shards is not None: - final_hidden_states._xorl_sglang_moe_tp_shards = tuple( - shard.reshape(batch_size, sequence_length, hidden_dim) for shard in tp_shards - ) return final_hidden_states, router_logits @@ -559,8 +463,6 @@ def _eager_forward( selected_experts: torch.Tensor, ) -> torch.Tensor: """Per-expert loop for eager mode.""" - if _moe_fp64_accum_enabled(): - return self._eager_forward_fp64(hidden_states, routing_weights, selected_experts) hidden_dim = hidden_states.shape[-1] final_hidden_states = torch.zeros_like(hidden_states) @@ -577,50 +479,6 @@ def _eager_forward( return final_hidden_states - def _eager_forward_fp64( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - ) -> torch.Tensor: - """K3 parity variant of ``_eager_forward``: order-invariant fp64 reductions. - - Same routing/masking as the bf16 loop, but the three expert GEMMs, the - routing-weight multiply, and the cross-expert combine all accumulate in - fp64; the only intermediate bf16 cast is the activation product inside - ``eager_expert_forward_fp64``, and the result is cast back once at the - end. Inference-only (no autograd through the fp64 detour is needed for - logprob scoring; gradients are unsupported). - """ - from .backend.eager import eager_expert_forward_fp64 # noqa: PLC0415 - - experts = self.experts - if not experts.gated or experts.hidden_act != "silu": - raise NotImplementedError(f"{_MOE_FP64_ACCUM_ENV} supports gated SiLU experts only") - if experts.gate_up_bias is not None or experts.down_bias is not None: - raise NotImplementedError(f"{_MOE_FP64_ACCUM_ENV} does not support expert biases") - if experts.swiglu_limit > 0: - raise NotImplementedError(f"{_MOE_FP64_ACCUM_ENV} does not support swiglu_limit") - - hidden_dim = hidden_states.shape[-1] - gate_proj = experts.gate_proj - up_proj = experts.up_proj - down_proj = experts.down_proj - final_hidden_states = torch.zeros(hidden_states.shape, dtype=torch.float64, device=hidden_states.device) - - expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) - - for expert_idx in range(self.num_experts): - idx, top_x = torch.where(expert_mask[expert_idx]) - if top_x.numel() == 0: - continue - - current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) - expert_out64 = eager_expert_forward_fp64(current_state, expert_idx, gate_proj, up_proj, down_proj) - final_hidden_states.index_add_(0, top_x, expert_out64 * routing_weights[top_x, idx, None].to(torch.float64)) - - return final_hidden_states.to(hidden_states.dtype) - @classmethod def from_config(cls, config, moe_implementation: str = "triton"): """Create from a model config (e.g. ``Qwen3MoeConfig``).""" diff --git a/src/xorl/models/layers/moe/router.py b/src/xorl/models/layers/moe/router.py index 47e87cdf..d3e27abc 100644 --- a/src/xorl/models/layers/moe/router.py +++ b/src/xorl/models/layers/moe/router.py @@ -8,7 +8,6 @@ _SYNTHETIC_ROUTING_ENV = "XORL_MOE_SYNTHETIC_ROUTING" -_ROUTER_TOPK_POLICY_ENV = "XORL_MOE_ROUTER_TOPK_POLICY" def balanced_synthetic_routing( @@ -37,36 +36,6 @@ def _synthetic_routing_mode() -> str | None: raise ValueError(f"{_SYNTHETIC_ROUTING_ENV} must be unset or 'balanced', got {mode!r}") -def _router_topk_policy() -> str: - policy = os.environ.get(_ROUTER_TOPK_POLICY_ENV, "").strip().lower() - if policy in {"", "0", "false", "no", "off", "default", "softmax"}: - return "default" - if policy in {"logits", "stable_low_id", "tie_low_id", "tie_high_id"}: - return policy - raise ValueError( - f"{_ROUTER_TOPK_POLICY_ENV} must be unset/default, logits, stable_low_id, " - f"tie_low_id, or tie_high_id; got {policy!r}" - ) - - -def _topk_indices_with_policy(scores: torch.Tensor, top_k: int, policy: str | None = None) -> torch.Tensor: - policy = _router_topk_policy() if policy is None else policy - if policy == "default": - return torch.topk(scores, top_k, dim=-1).indices - if policy == "stable_low_id": - return torch.argsort(scores.float(), dim=-1, descending=True, stable=True)[:, :top_k] - - if policy == "tie_low_id": - expert_id = torch.arange(scores.shape[-1], dtype=torch.float32, device=scores.device) - scores = scores.float() - expert_id * 1e-7 - elif policy == "tie_high_id": - expert_id = torch.arange(scores.shape[-1], dtype=torch.float32, device=scores.device) - scores = scores.float() + expert_id * 1e-7 - else: - scores = scores.float() - return torch.topk(scores, top_k, dim=-1).indices - - def _balanced_selected_experts(router_logits: torch.Tensor, num_experts: int, top_k: int) -> torch.Tensor: """Deterministically spread routed slots evenly over all experts. @@ -154,7 +123,6 @@ def __init__( # Exact model programs are structural and must not be redirected by # process-wide diagnostic environment variables. self.synthetic_routing_mode = None if exact_batch_invariant else _synthetic_routing_mode() - self.topk_policy = "default" if exact_batch_invariant else _router_topk_policy() # ``tid2eid`` is a frozen post-load buffer — validate bounds once and # remember the storage object so the hot path doesn't pay device->host # sync per forward. @@ -237,12 +205,7 @@ def _forward_softmax(self, router_logits: torch.Tensor, input_dtype: torch.dtype input_dtype, ) routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) - topk_policy = self.topk_policy - if topk_policy == "default": - routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) - else: - selected_experts = _topk_indices_with_policy(router_logits, self.top_k, topk_policy) - routing_weights = torch.gather(routing_weights, 1, selected_experts) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) if self._exact_batch_invariant: # Exact router contract: fixed-order renorm + cast so the top-k weights # are bit-identical to SGLang's batch-invariant path (the stock @@ -291,7 +254,7 @@ def _forward_sqrtsoftplus( scores_for_routing = scores + expert_bias else: scores_for_routing = scores - selected_experts = _topk_indices_with_policy(scores_for_routing, self.top_k, self.topk_policy) + selected_experts = torch.topk(scores_for_routing, self.top_k, dim=-1).indices routing_weights = torch.gather(scores, dim=1, index=selected_experts) # V4 paths always renormalize. diff --git a/src/xorl/models/layers/moe/routing_replay.py b/src/xorl/models/layers/moe/routing_replay.py index 8e21ba16..08e56e2e 100644 --- a/src/xorl/models/layers/moe/routing_replay.py +++ b/src/xorl/models/layers/moe/routing_replay.py @@ -22,7 +22,6 @@ model.forward() # records _pp_forward restores "replay_backward" set("replay_backward") loss.backward() loss.backward() # pop_backward checkpoint recompute -> "replay_backward" - reset_all_backward() -> pop_backward set(None) set(None) clear_all() clear_all() """ @@ -129,12 +128,6 @@ def has_weights(self) -> bool: """Whether routing weights are pre-populated.""" return len(self.top_weights_list) > 0 - def reset_forward(self): - self.forward_index = 0 - - def reset_backward(self): - self.backward_index = 0 - def clear(self): self.forward_index = 0 self.backward_index = 0 @@ -148,16 +141,6 @@ def clear_all(cls): for inst in cls._instances: inst.clear() - @classmethod - def reset_all_forward(cls): - for inst in cls._instances: - inst.reset_forward() - - @classmethod - def reset_all_backward(cls): - for inst in cls._instances: - inst.reset_backward() - # --------------------------------------------------------------------------- # Global stage diff --git a/src/xorl/models/module_utils.py b/src/xorl/models/module_utils.py index 98c96050..ac4945ac 100644 --- a/src/xorl/models/module_utils.py +++ b/src/xorl/models/module_utils.py @@ -3123,8 +3123,6 @@ def _moe_forward(self, hidden_states, output_router_logits=False, **kwargs): materialized_hidden_states = layer_output_override(materialized_hidden_states) self._capture_diagnostic_component("layer_output_override", materialized_hidden_states) hidden_states = materialized_hidden_states - elif getattr(self, "_delay_moe_residual_output", False): - hidden_states = (hidden_states, residual) else: hidden_states = materialized_hidden_states diff --git a/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py b/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py index 4176d2ed..61ea1bab 100644 --- a/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py +++ b/src/xorl/models/transformers/glm4_moe/modeling_glm4_moe.py @@ -35,6 +35,7 @@ is_flash_attention, update_causal_mask, ) +from xorl.models.layers.fused_projection_lora import project_fused_linear_with_lora from xorl.models.layers.moe import MoEBlock from xorl.models.layers.moe.routing_replay import get_replay_stage from xorl.models.layers.normalization import compiled_eager_rms_norm @@ -83,6 +84,8 @@ def extra_repr(self) -> str: class Glm4MoeMLP(nn.Module): + _supports_fused_gate_up_lora = True + def __init__(self, config, intermediate_size=None): super().__init__() self.hidden_size = config.hidden_size @@ -103,10 +106,17 @@ def unfuse_for_tp(self): def forward(self, x): if hasattr(self, "gate_up_proj"): + gate_up = project_fused_linear_with_lora( + self, + x, + base_name="gate_up_proj", + projection_names=("gate_proj", "up_proj"), + projection_sizes=(self.intermediate_size, self.intermediate_size), + ) if self._use_fused_silu: - x = fused_silu_and_mul(self.gate_up_proj(x)) + x = fused_silu_and_mul(gate_up) else: - gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + gate, up = gate_up.chunk(2, dim=-1) x = self.act_fn(gate) * up else: x = self.act_fn(self.gate_proj(x)) * self.up_proj(x) @@ -306,7 +316,13 @@ def _project_qkv( hidden_shape = (*input_shape, -1, self.head_dim) if hasattr(self, "qkv_proj"): - qkv = self.qkv_proj(hidden_states) + qkv = project_fused_linear_with_lora( + self, + hidden_states, + base_name="qkv_proj", + projection_names=("q_proj", "k_proj", "v_proj"), + projection_sizes=(self.q_dim, self.kv_dim, self.kv_dim), + ) q, k, v = qkv.split([self.q_dim, self.kv_dim, self.kv_dim], dim=-1) else: q = self.q_proj(hidden_states) diff --git a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py index cd99844b..40634d30 100644 --- a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py @@ -562,13 +562,6 @@ def __init__( self.shard = expected self.tp_group = tp_group - def bind_tp_group(self, tp_group: dist.ProcessGroup) -> None: - """Bind the already-created XoRL lm-head-only TP process group.""" - - if tp_group is None: - raise ValueError("GLM-5.2 exact LM head requires an explicit TP process group") - self.tp_group = tp_group - def _validate_tp_group(self) -> dist.ProcessGroup: group = self.tp_group if group is None: @@ -687,18 +680,6 @@ def _validate_operands( if require_cuda and hidden_states.device.type != "cuda": raise RuntimeError("GLM-5.2 exact LM-head value forward requires CUDA and pinned S4 kernels") - def effective_factor_views(self, lora_A: Tensor, local_lora_B: Tensor) -> tuple[Tensor, Tensor]: - """Return the exact live BF16 bytes consumed by the S4 A/B kernels.""" - - if lora_A.dtype is not torch.float32 or tuple(lora_A.shape) != (1, GLM52_LM_HEAD_HIDDEN_SIZE): - raise TypeError("lora_A must be the official FP32 [1, 6144] master") - if local_lora_B.dtype is not torch.float32 or tuple(local_lora_B.shape) != ( - GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, - 1, - ): - raise TypeError("local_lora_B must be the official FP32 [9680, 1] master") - return lora_A.to(torch.bfloat16).contiguous(), local_lora_B.to(torch.bfloat16).contiguous() - def _exact_local_logits( self, hidden_2d: Tensor, diff --git a/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py b/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py index f64948cb..a6995267 100644 --- a/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_routed_experts_qlora.py @@ -17,7 +17,7 @@ from __future__ import annotations import math -from dataclasses import dataclass, replace +from dataclasses import replace import torch import torch.nn.functional as F @@ -42,18 +42,6 @@ GLM52_ROUTED_MAX_LORAS_PER_BATCH = 8 -@dataclass(frozen=True) -class Glm52ExactRoutedValueTrace: - """Sampler-owned intermediate bytes captured by a component test.""" - - gate_up_base: Tensor - gate_up_post_lora: Tensor - activated: Tensor - down_base_routed: Tensor - down_post_lora_routed: Tensor - owner_output: Tensor - - def localize_glm52_ep16_expert_ids(global_ids: Tensor, ep_rank: int) -> Tensor: """Map global expert IDs to one EP16 owner's contiguous local slots.""" @@ -94,13 +82,12 @@ def forward( effective = tuple( factor.to(torch.bfloat16).contiguous() for factor in (gate_A, gate_B, up_A, up_B, down_A, down_B) ) - output, _ = module._sampler_value( + output = module._sampler_value( hidden, routing, local_ids, *effective, routed_scaling_factor=float(routed_scaling_factor), - capture_trace=False, ) expected_shape = (hidden.shape[0], module.hidden_size) if output.dtype is not torch.bfloat16 or tuple(output.shape) != expected_shape: @@ -499,10 +486,6 @@ def _physical_factor_buffers( "down_lora_b_weights": down_B_buffer, } - def physical_factor_buffers(self) -> dict[str, Tensor]: - effective = tuple(getattr(self, name).to(torch.bfloat16).contiguous() for name in self.logical_factor_names) - return self._physical_factor_buffers(*effective) - def _lora_info(self, rows: int, physical: dict[str, Tensor]): try: from sglang.srt.lora.lora_moe_runners import LoRAInfo # noqa: PLC0415 @@ -546,8 +529,7 @@ def _sampler_value( down_B: Tensor, *, routed_scaling_factor: float, - capture_trace: bool, - ) -> tuple[Tensor, Glm52ExactRoutedValueTrace | None]: + ) -> Tensor: """Invoke S4's literal base/hook/activation/hook/fold sequence.""" physical = self._physical_factor_buffers(gate_A, gate_B, up_A, up_B, down_A, down_B) @@ -581,27 +563,6 @@ def _sampler_value( block_shape=[128, 128], ) hooks = build_lora_hooks(hidden, self._lora_info(hidden.shape[0], physical), local_ids) - trace_values: dict[str, Tensor] = {} - if capture_trace: - original_gate_up = hooks.after_gate_up - original_down = hooks.after_down - - def traced_gate_up(x, cache, weights, ids): - trace_values["gate_up_base"] = cache.clone() - assert original_gate_up is not None - original_gate_up(x, cache, weights, ids) - trace_values["gate_up_post_lora"] = cache.clone() - - def traced_down(activated, cache, weights, ids): - trace_values["activated"] = activated.clone() - trace_values["down_base_routed"] = cache.clone() - assert original_down is not None - original_down(activated, cache, weights, ids) - trace_values["down_post_lora_routed"] = cache.clone() - - hooks.after_gate_up = traced_gate_up - hooks.after_down = traced_down - output = _fused_moe_kernel_sequence( hidden, w1, @@ -642,10 +603,7 @@ def traced_down(activated, cache, weights, ids): gate_up_interleaved=False, a1_q=None, ) - trace = None - if capture_trace: - trace = Glm52ExactRoutedValueTrace(owner_output=output.clone(), **trace_values) - return output, trace + return output def _dequantized_base(self) -> tuple[Tensor, Tensor]: try: @@ -764,28 +722,6 @@ def _surrogate_vjp( iterator = iter(computed) return tuple(next(iterator) if needed else None for needed in needs) - def sampler_value_trace( - self, - hidden: Tensor, - routing: Tensor, - global_ids: Tensor, - *, - routed_scaling_factor: float = 1.0, - ) -> Glm52ExactRoutedValueTrace: - local_ids = self.localize_global_expert_ids(global_ids) - self._validate_runtime_contract(hidden, routing, local_ids) - effective = tuple(getattr(self, name).to(torch.bfloat16).contiguous() for name in self.logical_factor_names) - _, trace = self._sampler_value( - hidden, - routing, - local_ids, - *effective, - routed_scaling_factor=float(routed_scaling_factor), - capture_trace=True, - ) - assert trace is not None - return trace - def forward( self, hidden_states: Tensor, @@ -831,6 +767,5 @@ def forward( "GLM52_ROUTED_GLOBAL_EXPERTS", "GLM52_ROUTED_LOCAL_EXPERTS", "Glm52ExactEP16BlockFP8QLoRARoutedExperts", - "Glm52ExactRoutedValueTrace", "localize_glm52_ep16_expert_ids", ] diff --git a/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py b/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py index 062c8db0..03b0944c 100644 --- a/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_shared_expert_qlora.py @@ -419,23 +419,6 @@ def _physical_factor_views_from_effective( down_B=effective_down_B.unsqueeze(0).contiguous(), ) - def physical_factor_views(self, contributor_ordinal: int) -> Glm52SharedExpertPhysicalFactors: - """Derive one live SGLang slot view from the FP32 logical masters.""" - - self._validate_factor_state() - effective = tuple( - factor.to(torch.bfloat16).contiguous() - for factor in ( - self.gate_proj.lora_A, - self.gate_proj.lora_B, - self.up_proj.lora_A, - self.up_proj.lora_B, - self.down_proj.lora_A, - self.down_proj.lora_B, - ) - ) - return self._physical_factor_views_from_effective(*effective, contributor_ordinal) - @staticmethod def _partition_base_state( projection: NativeBlockFP8Linear, diff --git a/src/xorl/models/transformers/glm5/index_share.py b/src/xorl/models/transformers/glm5/index_share.py index 056b31ac..7c3bb08e 100644 --- a/src/xorl/models/transformers/glm5/index_share.py +++ b/src/xorl/models/transformers/glm5/index_share.py @@ -2,10 +2,9 @@ from __future__ import annotations -from contextlib import contextmanager from dataclasses import dataclass from enum import Enum -from typing import Callable, Iterator +from typing import Callable import torch @@ -192,16 +191,6 @@ def finish_forward(self, context: IndexShareContext, *, succeeded: bool) -> None if not succeeded or context.mode is IndexShareMode.FORWARD_ONLY: self.end(context) - @contextmanager - def invocation(self, *, mode: IndexShareMode | str) -> Iterator[IndexShareContext]: - context = self.begin(mode=mode) - succeeded = False - try: - yield context - succeeded = True - finally: - self.finish_forward(context, succeeded=succeeded) - __all__ = [ "CanonicalLogicalIndices", diff --git a/src/xorl/models/transformers/glm5/layer_plan.py b/src/xorl/models/transformers/glm5/layer_plan.py index fbb62c7f..79c390c9 100644 --- a/src/xorl/models/transformers/glm5/layer_plan.py +++ b/src/xorl/models/transformers/glm5/layer_plan.py @@ -137,22 +137,6 @@ def _validate_pipeline_ranges(ranges: tuple[tuple[int, int], ...], num_layers: i if expected_start != num_layers: raise ValueError(f"Pipeline layer ranges must cover exactly {num_layers} layers") - @property - def full_indexer_layers(self) -> tuple[int, ...]: - return tuple(layer.layer_index for layer in self.layers if layer.indexer_type is IndexerType.FULL) - - @property - def shared_indexer_layers(self) -> tuple[int, ...]: - return tuple(layer.layer_index for layer in self.layers if layer.indexer_type is IndexerType.SHARED) - - @property - def dense_layers(self) -> tuple[int, ...]: - return tuple(layer.layer_index for layer in self.layers if layer.mlp_type is MLPType.DENSE) - - @property - def sparse_layers(self) -> tuple[int, ...]: - return tuple(layer.layer_index for layer in self.layers if layer.mlp_type is MLPType.SPARSE) - @property def identity(self) -> str: payload = { diff --git a/src/xorl/models/transformers/glm5/qlora.py b/src/xorl/models/transformers/glm5/qlora.py index dd99e1d9..164de8cb 100644 --- a/src/xorl/models/transformers/glm5/qlora.py +++ b/src/xorl/models/transformers/glm5/qlora.py @@ -2,7 +2,6 @@ from __future__ import annotations -from collections import Counter from dataclasses import dataclass import torch @@ -101,10 +100,6 @@ def target_names(self) -> frozenset[str]: def factor_names(self) -> frozenset[str]: return frozenset(factor.name for factor in self.factors) - @property - def role_counts(self) -> dict[str, int]: - return dict(Counter(target.role for target in self.targets)) - def _official_indexer_schedule() -> tuple[str, ...]: return tuple("full" if layer_idx < 3 or (layer_idx - 2) % 4 == 0 else "shared" for layer_idx in range(78)) diff --git a/src/xorl/models/transformers/glm5/sparse_selector.py b/src/xorl/models/transformers/glm5/sparse_selector.py index 2154299c..62ea8cdf 100644 --- a/src/xorl/models/transformers/glm5/sparse_selector.py +++ b/src/xorl/models/transformers/glm5/sparse_selector.py @@ -367,50 +367,6 @@ def select_glm52_logical_indices( return Glm52Selection(CanonicalLogicalIndices(output), valid_counts) -def physical_cache_to_logical_indices( - physical_indices: torch.Tensor, - physical_to_logical: torch.Tensor, -) -> CanonicalLogicalIndices: - """Convert cache/page slots to canonical sequence positions.""" - if physical_indices.dtype not in (torch.int32, torch.int64): - raise TypeError("physical_indices must be int32 or int64") - if physical_to_logical.dtype not in (torch.int32, torch.int64): - raise TypeError("physical_to_logical must be int32 or int64") - if physical_to_logical.ndim not in (1, 2): - raise ValueError("physical_to_logical must be [slots] or [batch, slots]") - - valid = physical_indices >= 0 - clamped = physical_indices.clamp_min(0).long() - if physical_to_logical.ndim == 1: - if bool(torch.any(clamped[valid] >= physical_to_logical.shape[0])): - raise ValueError("physical index exceeds the physical-to-logical table") - logical = physical_to_logical[clamped] - else: - if physical_indices.shape[0] != physical_to_logical.shape[0]: - raise ValueError("Batched physical index and page-map batch dimensions differ") - if bool(torch.any(clamped[valid] >= physical_to_logical.shape[1])): - raise ValueError("physical index exceeds the batched physical-to-logical table") - batch_index = torch.arange(physical_indices.shape[0], device=physical_indices.device) - batch_index = batch_index.view(-1, *([1] * (physical_indices.ndim - 1))).expand_as(clamped) - logical = physical_to_logical[batch_index, clamped] - logical = logical.to(torch.int32).masked_fill(~valid, -1) - return CanonicalLogicalIndices(logical) - - -def gather_selected_logical_values(values: torch.Tensor, indices: CanonicalLogicalIndices) -> torch.Tensor: - """Gather selected values with finite-zero semantics for ``-1`` slots.""" - selected = indices.values - if values.ndim != 3 or selected.ndim != 3: - raise ValueError("Expected values [B,K,D] and selected indices [B,Q,T]") - if values.shape[0] != selected.shape[0]: - raise ValueError("Value and selected-index batch dimensions differ") - valid = selected >= 0 - clamped = selected.clamp_min(0).long() - batch = torch.arange(values.shape[0], device=values.device).view(-1, 1, 1).expand_as(clamped) - gathered = values[batch, clamped] - return torch.where(valid.unsqueeze(-1), gathered, torch.zeros_like(gathered)) - - __all__ = [ "FP8_E4M3_MAX", "GLM52_HADAMARD_IMPORT", @@ -421,8 +377,6 @@ def gather_selected_logical_values(values: torch.Tensor, indices: CanonicalLogic "GLM52_SELECTOR_VERSION", "Glm52Selection", "dequantize_e4m3_blocks", - "gather_selected_logical_values", - "physical_cache_to_logical_indices", "quantize_e4m3_dynamic", "quantize_e4m3_ue8m0", "quantize_sparse_key_cache", diff --git a/src/xorl/models/transformers/llama3/checkpoint_handler.py b/src/xorl/models/transformers/llama3/checkpoint_handler.py index e4c2751a..b65fd8f5 100644 --- a/src/xorl/models/transformers/llama3/checkpoint_handler.py +++ b/src/xorl/models/transformers/llama3/checkpoint_handler.py @@ -31,6 +31,12 @@ class Llama3CheckpointHandler(CheckpointHandler): When ``is_prequantized=True`` and a model is provided, quantized weights are loaded inline via QLoRAWeightBuffer (single-pass I/O). When model is not provided, falls back to skipping quantized keys (deferred loading). + + Unfused handling: + - When ``skip_qkv_merge=True``, QKV keys pass through unmerged + (model has separate q_proj/k_proj/v_proj after unfusing). + - When ``skip_gate_up_merge=True``, gate/up keys pass through unmerged + (model has separate gate_proj/up_proj after unfusing). """ def __init__( @@ -38,12 +44,14 @@ def __init__( num_attention_heads: int, num_key_value_heads: int, head_dim: int, + skip_qkv_merge: bool = False, + skip_gate_up_merge: bool = False, is_prequantized: bool = False, exclude_modules: Optional[Set[str]] = None, model: Optional[nn.Module] = None, ): - self._gate_up_buffer = GateUpMergeBuffer() - self._qkv_buffer = QKVMergeBuffer() + self._gate_up_buffer: Optional[GateUpMergeBuffer] = None if skip_gate_up_merge else GateUpMergeBuffer() + self._qkv_buffer: Optional[QKVMergeBuffer] = None if skip_qkv_merge else QKVMergeBuffer() self._q_dim = num_attention_heads * head_dim self._kv_dim = num_key_value_heads * head_dim self._is_prequantized = is_prequantized @@ -112,37 +120,41 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc if OPROJ_WEIGHT_PATTERN.match(key) or DENSE_DOWN_PROJ_PATTERN.match(key): return [] - # QKV merge - if self._is_prequantized and key.endswith(".weight"): - if self._qkv_buffer.is_qkv_key(key): - return [] - else: - qkv_result = self._qkv_buffer.add(key, tensor) - if qkv_result is not None: - return [qkv_result] - if self._qkv_buffer.is_qkv_key(key): - return [] + # QKV merge (skipped when the model carries separate q/k/v projections) + if self._qkv_buffer is not None: + if self._is_prequantized and key.endswith(".weight"): + if self._qkv_buffer.is_qkv_key(key): + return [] + else: + qkv_result = self._qkv_buffer.add(key, tensor) + if qkv_result is not None: + return [qkv_result] + if self._qkv_buffer.is_qkv_key(key): + return [] - # Gate/up merge - if self._is_prequantized and key.endswith(".weight"): - if self._gate_up_buffer.is_gate_up_key(key): - return [] - else: - merge_result = self._gate_up_buffer.add(key, tensor) - if merge_result is not None: - return [merge_result] - if self._gate_up_buffer.is_gate_up_key(key): - return [] + # Gate/up merge (skipped when the model carries separate gate/up projections) + if self._gate_up_buffer is not None: + if self._is_prequantized and key.endswith(".weight"): + if self._gate_up_buffer.is_gate_up_key(key): + return [] + else: + merge_result = self._gate_up_buffer.add(key, tensor) + if merge_result is not None: + return [merge_result] + if self._gate_up_buffer.is_gate_up_key(key): + return [] return [(key, tensor)] def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: - pending_gu = self._gate_up_buffer.get_pending() - if pending_gu: - warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") - pending_qkv = self._qkv_buffer.get_pending() - if pending_qkv: - warnings.warn(f"Incomplete QKV merge groups after loading: {pending_qkv}") + if self._gate_up_buffer is not None: + pending_gu = self._gate_up_buffer.get_pending() + if pending_gu: + warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") + if self._qkv_buffer is not None: + pending_qkv = self._qkv_buffer.get_pending() + if pending_qkv: + warnings.warn(f"Incomplete QKV merge groups after loading: {pending_qkv}") if self._qlora_buffer is not None: self._qlora_buffer.set_inline_metadata() return [] diff --git a/src/xorl/models/transformers/llama3/modeling_llama3.py b/src/xorl/models/transformers/llama3/modeling_llama3.py index ee1a3b19..d114325e 100644 --- a/src/xorl/models/transformers/llama3/modeling_llama3.py +++ b/src/xorl/models/transformers/llama3/modeling_llama3.py @@ -18,6 +18,7 @@ is_flash_attention, update_causal_mask, ) +from xorl.models.layers.fused_projection_lora import project_fused_linear_with_lora from xorl.models.module_utils import GradientCheckpointingLayer from xorl.models.outputs import BaseModelOutput, CausalLMOutput from xorl.models.transformers.llama3 import parallelize @@ -31,6 +32,8 @@ class LlamaMLP(nn.Module): + _supports_fused_gate_up_lora = True + def __init__(self, config): super().__init__() self.hidden_size = config.hidden_size @@ -54,10 +57,17 @@ def unfuse_for_tp(self): def forward(self, x): if hasattr(self, "gate_up_proj"): + gate_up = project_fused_linear_with_lora( + self, + x, + base_name="gate_up_proj", + projection_names=("gate_proj", "up_proj"), + projection_sizes=(self.intermediate_size, self.intermediate_size), + ) if self._use_fused_silu: - x = fused_silu_and_mul(self.gate_up_proj(x)) + x = fused_silu_and_mul(gate_up) else: - gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + gate, up = gate_up.chunk(2, dim=-1) x = self.act_fn(gate) * up else: x = self.act_fn(self.gate_proj(x)) * self.up_proj(x) @@ -145,8 +155,7 @@ def _init_weights(self, module): module.original_inv_freq = module.inv_freq def get_checkpoint_handler(self, **kwargs): - if getattr(self, "_unfused_for_tp", False): - return None + unfused = getattr(self, "_unfused_for_tp", False) weights_path = kwargs.get("weights_path", None) is_prequantized = detect_prequantized_checkpoint(weights_path) @@ -162,6 +171,10 @@ def get_checkpoint_handler(self, **kwargs): num_attention_heads=self.config.num_attention_heads, num_key_value_heads=self.config.num_key_value_heads, head_dim=head_dim, + # Unfused checkpoint keys already match the parameter names, so only the + # merges are skipped; the handler still carries the pre-quantized paths. + skip_qkv_merge=unfused, + skip_gate_up_merge=unfused, is_prequantized=is_prequantized, exclude_modules=exclude_modules, model=self if is_prequantized else None, diff --git a/src/xorl/models/transformers/olmo2/checkpoint_handler.py b/src/xorl/models/transformers/olmo2/checkpoint_handler.py index c07d8a50..18ba291c 100644 --- a/src/xorl/models/transformers/olmo2/checkpoint_handler.py +++ b/src/xorl/models/transformers/olmo2/checkpoint_handler.py @@ -31,6 +31,12 @@ class Olmo2CheckpointHandler(CheckpointHandler): OLMo-2 layer norm names (``post_attention_layernorm``, ``post_feedforward_layernorm``) already match the model's parameter names, so no key remapping is needed for them. + + Unfused handling: + - When ``skip_qkv_merge=True``, QKV keys pass through unmerged + (model has separate q_proj/k_proj/v_proj after unfusing). + - When ``skip_gate_up_merge=True``, gate/up keys pass through unmerged + (model has separate gate_proj/up_proj after unfusing). """ def __init__( @@ -38,12 +44,14 @@ def __init__( num_attention_heads: int, num_key_value_heads: int, head_dim: int, + skip_qkv_merge: bool = False, + skip_gate_up_merge: bool = False, is_prequantized: bool = False, exclude_modules: Optional[Set[str]] = None, model: Optional[nn.Module] = None, ): - self._gate_up_buffer = GateUpMergeBuffer() - self._qkv_buffer = QKVMergeBuffer() + self._gate_up_buffer: Optional[GateUpMergeBuffer] = None if skip_gate_up_merge else GateUpMergeBuffer() + self._qkv_buffer: Optional[QKVMergeBuffer] = None if skip_qkv_merge else QKVMergeBuffer() self._q_dim = num_attention_heads * head_dim self._kv_dim = num_key_value_heads * head_dim self._is_prequantized = is_prequantized @@ -112,37 +120,41 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc if OPROJ_WEIGHT_PATTERN.match(key) or DENSE_DOWN_PROJ_PATTERN.match(key): return [] - # QKV merge - if self._is_prequantized and key.endswith(".weight"): - if self._qkv_buffer.is_qkv_key(key): - return [] - else: - qkv_result = self._qkv_buffer.add(key, tensor) - if qkv_result is not None: - return [qkv_result] - if self._qkv_buffer.is_qkv_key(key): - return [] + # QKV merge (skipped when the model carries separate q/k/v projections) + if self._qkv_buffer is not None: + if self._is_prequantized and key.endswith(".weight"): + if self._qkv_buffer.is_qkv_key(key): + return [] + else: + qkv_result = self._qkv_buffer.add(key, tensor) + if qkv_result is not None: + return [qkv_result] + if self._qkv_buffer.is_qkv_key(key): + return [] - # Gate/up merge - if self._is_prequantized and key.endswith(".weight"): - if self._gate_up_buffer.is_gate_up_key(key): - return [] - else: - merge_result = self._gate_up_buffer.add(key, tensor) - if merge_result is not None: - return [merge_result] - if self._gate_up_buffer.is_gate_up_key(key): - return [] + # Gate/up merge (skipped when the model carries separate gate/up projections) + if self._gate_up_buffer is not None: + if self._is_prequantized and key.endswith(".weight"): + if self._gate_up_buffer.is_gate_up_key(key): + return [] + else: + merge_result = self._gate_up_buffer.add(key, tensor) + if merge_result is not None: + return [merge_result] + if self._gate_up_buffer.is_gate_up_key(key): + return [] return [(key, tensor)] def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: - pending_gu = self._gate_up_buffer.get_pending() - if pending_gu: - warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") - pending_qkv = self._qkv_buffer.get_pending() - if pending_qkv: - warnings.warn(f"Incomplete QKV merge groups after loading: {pending_qkv}") + if self._gate_up_buffer is not None: + pending_gu = self._gate_up_buffer.get_pending() + if pending_gu: + warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") + if self._qkv_buffer is not None: + pending_qkv = self._qkv_buffer.get_pending() + if pending_qkv: + warnings.warn(f"Incomplete QKV merge groups after loading: {pending_qkv}") if self._qlora_buffer is not None: self._qlora_buffer.set_inline_metadata() return [] diff --git a/src/xorl/models/transformers/olmo2/modeling_olmo2.py b/src/xorl/models/transformers/olmo2/modeling_olmo2.py index 2f39d74c..74879ac7 100644 --- a/src/xorl/models/transformers/olmo2/modeling_olmo2.py +++ b/src/xorl/models/transformers/olmo2/modeling_olmo2.py @@ -18,6 +18,7 @@ is_flash_attention, update_causal_mask, ) +from xorl.models.layers.fused_projection_lora import project_fused_linear_with_lora from xorl.models.layers.normalization import native_rms_norm from xorl.models.layers.rope import apply_rotary_pos_emb from xorl.models.module_utils import GradientCheckpointingLayer @@ -33,6 +34,8 @@ class Olmo2MLP(nn.Module): + _supports_fused_gate_up_lora = True + def __init__(self, config): super().__init__() self.hidden_size = config.hidden_size @@ -52,10 +55,17 @@ def unfuse_for_tp(self): def forward(self, x): if hasattr(self, "gate_up_proj"): + gate_up = project_fused_linear_with_lora( + self, + x, + base_name="gate_up_proj", + projection_names=("gate_proj", "up_proj"), + projection_sizes=(self.intermediate_size, self.intermediate_size), + ) if self._use_fused_silu: - x = fused_silu_and_mul(self.gate_up_proj(x)) + x = fused_silu_and_mul(gate_up) else: - gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + gate, up = gate_up.chunk(2, dim=-1) x = self.act_fn(gate) * up else: x = self.act_fn(self.gate_proj(x)) * self.up_proj(x) @@ -224,8 +234,7 @@ def _init_weights(self, module): module.original_inv_freq = module.inv_freq def get_checkpoint_handler(self, **kwargs): - if getattr(self, "_unfused_for_tp", False): - return None + unfused = getattr(self, "_unfused_for_tp", False) weights_path = kwargs.get("weights_path", None) is_prequantized = detect_prequantized_checkpoint(weights_path) @@ -241,6 +250,10 @@ def get_checkpoint_handler(self, **kwargs): num_attention_heads=self.config.num_attention_heads, num_key_value_heads=self.config.num_key_value_heads, head_dim=head_dim, + # Unfused checkpoint keys already match the parameter names, so only the + # merges are skipped; the handler still carries the pre-quantized paths. + skip_qkv_merge=unfused, + skip_gate_up_merge=unfused, is_prequantized=is_prequantized, exclude_modules=exclude_modules, model=self if is_prequantized else None, diff --git a/src/xorl/models/transformers/qwen2/modeling_qwen2.py b/src/xorl/models/transformers/qwen2/modeling_qwen2.py index 458fdd90..65b358e7 100644 --- a/src/xorl/models/transformers/qwen2/modeling_qwen2.py +++ b/src/xorl/models/transformers/qwen2/modeling_qwen2.py @@ -25,6 +25,7 @@ is_flash_attention, update_causal_mask, ) +from xorl.models.layers.fused_projection_lora import project_fused_linear_with_lora from xorl.models.module_utils import GradientCheckpointingLayer from xorl.models.outputs import BaseModelOutput, CausalLMOutput from xorl.models.transformers.qwen2 import parallelize @@ -70,6 +71,8 @@ def _adapt_qwen2_config(config): class Qwen2MLP(nn.Module): + _supports_fused_gate_up_lora = True + def __init__(self, config): super().__init__() self.hidden_size = config.hidden_size @@ -89,10 +92,17 @@ def unfuse_for_tp(self): def forward(self, x): if hasattr(self, "gate_up_proj"): + gate_up = project_fused_linear_with_lora( + self, + x, + base_name="gate_up_proj", + projection_names=("gate_proj", "up_proj"), + projection_sizes=(self.intermediate_size, self.intermediate_size), + ) if self._use_fused_silu: - x = fused_silu_and_mul(self.gate_up_proj(x)) + x = fused_silu_and_mul(gate_up) else: - gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + gate, up = gate_up.chunk(2, dim=-1) x = self.act_fn(gate) * up else: x = self.act_fn(self.gate_proj(x)) * self.up_proj(x) @@ -184,8 +194,7 @@ def _init_weights(self, module): module.original_inv_freq = module.inv_freq def get_checkpoint_handler(self, **kwargs): - if getattr(self, "_unfused_for_tp", False): - return None + unfused = getattr(self, "_unfused_for_tp", False) weights_path = kwargs.get("weights_path", None) is_prequantized = detect_prequantized_checkpoint(weights_path) @@ -201,6 +210,10 @@ def get_checkpoint_handler(self, **kwargs): num_attention_heads=self.config.num_attention_heads, num_key_value_heads=self.config.num_key_value_heads, head_dim=head_dim, + # Unfused checkpoint keys already match the parameter names, so only the + # merges are skipped; the handler still carries the pre-quantized paths. + skip_qkv_merge=unfused, + skip_gate_up_merge=unfused, is_prequantized=is_prequantized, exclude_modules=exclude_modules, model=self if is_prequantized else None, diff --git a/src/xorl/models/transformers/qwen3/checkpoint_handler.py b/src/xorl/models/transformers/qwen3/checkpoint_handler.py index d52dc922..4065b5c4 100644 --- a/src/xorl/models/transformers/qwen3/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3/checkpoint_handler.py @@ -33,6 +33,12 @@ class Qwen3CheckpointHandler(CheckpointHandler): provided, falls back to the old behavior of skipping quantized keys (deferred loading via _deferred_qlora_quantize). + Unfused handling: + - When ``skip_qkv_merge=True``, QKV keys pass through unmerged + (model has separate q_proj/k_proj/v_proj after unfusing). + - When ``skip_gate_up_merge=True``, gate/up keys pass through unmerged + (model has separate gate_proj/up_proj after unfusing). + Bias keys always flow through the normal merge buffers. """ @@ -41,12 +47,14 @@ def __init__( num_attention_heads: int, num_key_value_heads: int, head_dim: int, + skip_qkv_merge: bool = False, + skip_gate_up_merge: bool = False, is_prequantized: bool = False, exclude_modules: Optional[Set[str]] = None, model: Optional[nn.Module] = None, ): - self._gate_up_buffer = GateUpMergeBuffer() - self._qkv_buffer = QKVMergeBuffer() + self._gate_up_buffer: Optional[GateUpMergeBuffer] = None if skip_gate_up_merge else GateUpMergeBuffer() + self._qkv_buffer: Optional[QKVMergeBuffer] = None if skip_qkv_merge else QKVMergeBuffer() self._q_dim = num_attention_heads * head_dim self._kv_dim = num_key_value_heads * head_dim self._is_prequantized = is_prequantized @@ -132,39 +140,43 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc if OPROJ_WEIGHT_PATTERN.match(key) or DENSE_DOWN_PROJ_PATTERN.match(key): return [] - # QKV merge - if self._is_prequantized and key.endswith(".weight"): - # Packed uint8 QKV weights — skip standard merging - if self._qkv_buffer.is_qkv_key(key): - return [] - else: - qkv_result = self._qkv_buffer.add(key, tensor) - if qkv_result is not None: - return [qkv_result] - if self._qkv_buffer.is_qkv_key(key): - return [] + # QKV merge (skipped when the model carries separate q/k/v projections) + if self._qkv_buffer is not None: + if self._is_prequantized and key.endswith(".weight"): + # Packed uint8 QKV weights — skip standard merging + if self._qkv_buffer.is_qkv_key(key): + return [] + else: + qkv_result = self._qkv_buffer.add(key, tensor) + if qkv_result is not None: + return [qkv_result] + if self._qkv_buffer.is_qkv_key(key): + return [] - # Gate/up merge - if self._is_prequantized and key.endswith(".weight"): - # Packed uint8 gate/up weights — skip standard merging - if self._gate_up_buffer.is_gate_up_key(key): - return [] - else: - merge_result = self._gate_up_buffer.add(key, tensor) - if merge_result is not None: - return [merge_result] - if self._gate_up_buffer.is_gate_up_key(key): - return [] + # Gate/up merge (skipped when the model carries separate gate/up projections) + if self._gate_up_buffer is not None: + if self._is_prequantized and key.endswith(".weight"): + # Packed uint8 gate/up weights — skip standard merging + if self._gate_up_buffer.is_gate_up_key(key): + return [] + else: + merge_result = self._gate_up_buffer.add(key, tensor) + if merge_result is not None: + return [merge_result] + if self._gate_up_buffer.is_gate_up_key(key): + return [] return [(key, tensor)] def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: - pending_gu = self._gate_up_buffer.get_pending() - if pending_gu: - warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") - pending_qkv = self._qkv_buffer.get_pending() - if pending_qkv: - warnings.warn(f"Incomplete QKV merge groups after loading: {pending_qkv}") + if self._gate_up_buffer is not None: + pending_gu = self._gate_up_buffer.get_pending() + if pending_gu: + warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") + if self._qkv_buffer is not None: + pending_qkv = self._qkv_buffer.get_pending() + if pending_qkv: + warnings.warn(f"Incomplete QKV merge groups after loading: {pending_qkv}") # Finalize inline-loaded QLoRA modules if self._qlora_buffer is not None: self._qlora_buffer.set_inline_metadata() diff --git a/src/xorl/models/transformers/qwen3/modeling_qwen3.py b/src/xorl/models/transformers/qwen3/modeling_qwen3.py index 718676a9..77ba5a93 100644 --- a/src/xorl/models/transformers/qwen3/modeling_qwen3.py +++ b/src/xorl/models/transformers/qwen3/modeling_qwen3.py @@ -24,6 +24,7 @@ is_flash_attention, update_causal_mask, ) +from xorl.models.layers.fused_projection_lora import project_fused_linear_with_lora from xorl.models.module_utils import GradientCheckpointingLayer from xorl.models.outputs import BaseModelOutput, CausalLMOutput from xorl.models.transformers.qwen3 import parallelize @@ -37,6 +38,8 @@ class Qwen3MLP(nn.Module): + _supports_fused_gate_up_lora = True + def __init__(self, config): super().__init__() self.hidden_size = config.hidden_size @@ -56,10 +59,17 @@ def unfuse_for_tp(self): def forward(self, x): if hasattr(self, "gate_up_proj"): + gate_up = project_fused_linear_with_lora( + self, + x, + base_name="gate_up_proj", + projection_names=("gate_proj", "up_proj"), + projection_sizes=(self.intermediate_size, self.intermediate_size), + ) if self._use_fused_silu: - x = fused_silu_and_mul(self.gate_up_proj(x)) + x = fused_silu_and_mul(gate_up) else: - gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + gate, up = gate_up.chunk(2, dim=-1) x = self.act_fn(gate) * up else: x = self.act_fn(self.gate_proj(x)) * self.up_proj(x) @@ -165,10 +175,7 @@ def _init_weights(self, module): module.original_inv_freq = module.inv_freq def get_checkpoint_handler(self, **kwargs): - # When unfused for TP, checkpoint keys (q_proj, k_proj, v_proj, gate_proj, - # up_proj) already match the model's parameter names — no merging needed. - if getattr(self, "_unfused_for_tp", False): - return None + unfused = getattr(self, "_unfused_for_tp", False) weights_path = kwargs.get("weights_path", None) is_prequantized = detect_prequantized_checkpoint(weights_path) @@ -186,6 +193,10 @@ def get_checkpoint_handler(self, **kwargs): num_attention_heads=self.config.num_attention_heads, num_key_value_heads=self.config.num_key_value_heads, head_dim=head_dim, + # Unfused checkpoint keys already match the parameter names, so only the + # merges are skipped; the handler still carries the pre-quantized paths. + skip_qkv_merge=unfused, + skip_gate_up_merge=unfused, is_prequantized=is_prequantized, exclude_modules=exclude_modules, model=self if is_prequantized else None, diff --git a/src/xorl/models/transformers/qwen3/parallelize.py b/src/xorl/models/transformers/qwen3/parallelize.py index 7ee99cc9..b9ac3520 100644 --- a/src/xorl/models/transformers/qwen3/parallelize.py +++ b/src/xorl/models/transformers/qwen3/parallelize.py @@ -27,7 +27,7 @@ def unfuse_for_tp(model): decoder layer. After unfusing, checkpoint keys from HuggingFace already match - the model's parameter names — no merging handler is needed. + the model's parameter names, so the checkpoint handler skips its merges. """ for layer in model.model.layers: layer.self_attn.unfuse_for_tp() diff --git a/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py index d487c822..53fd1119 100644 --- a/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_5/checkpoint_handler.py @@ -34,6 +34,13 @@ class Qwen3_5CheckpointHandler(CheckpointHandler): weight_scale_2, input_scale, weight_scale_inv) and linear projection ``.weight`` keys are skipped - they are loaded directly by QLoRA modules. Bias keys still flow through the normal merge buffers. + + Unfused handling: + - When ``skip_qkv_merge=True``, QKV keys pass through unmerged. + - When ``skip_gate_up_merge=True``, gate/up keys pass through unmerged + (MLP layers have separate gate_proj/up_proj after unfusing). + Linear-attention remapping runs regardless of either flag: the GatedDeltaNet + projections are packed on disk independently of how the MLP is stored. """ def __init__( @@ -44,11 +51,12 @@ def __init__( linear_key_dim: int, linear_value_dim: int, skip_qkv_merge: bool = False, + skip_gate_up_merge: bool = False, is_prequantized: bool = False, exclude_modules: Optional[Set[str]] = None, ): - self._gate_up_buffer = GateUpMergeBuffer() - self._qkv_buffer = None if skip_qkv_merge else QKVMergeBuffer() + self._gate_up_buffer: Optional[GateUpMergeBuffer] = None if skip_gate_up_merge else GateUpMergeBuffer() + self._qkv_buffer: Optional[QKVMergeBuffer] = None if skip_qkv_merge else QKVMergeBuffer() self._q_dim = num_attention_heads * head_dim self._kv_dim = num_key_value_heads * head_dim self._linear_key_dim = linear_key_dim @@ -128,7 +136,7 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc if linear_attn_results is not None: return linear_attn_results - # QKV merge (skipped when unfused for TP) + # QKV merge (always skipped here: Qwen3_5 attention stores q/k/v separately) if self._qkv_buffer is not None: if self._is_prequantized and key.endswith(".weight"): if self._qkv_buffer.is_qkv_key(key): @@ -140,24 +148,26 @@ def on_load_weight(self, key: str, tensor: torch.Tensor) -> List[Tuple[str, torc if self._qkv_buffer.is_qkv_key(key): return [] - # Gate/up merge - if self._is_prequantized and key.endswith(".weight"): - # Packed uint8 gate/up weights - skip standard merging - if self._gate_up_buffer.is_gate_up_key(key): - return [] - else: - merge_result = self._gate_up_buffer.add(key, tensor) - if merge_result is not None: - return [merge_result] - if self._gate_up_buffer.is_gate_up_key(key): - return [] + # Gate/up merge (skipped when the model carries separate gate/up projections) + if self._gate_up_buffer is not None: + if self._is_prequantized and key.endswith(".weight"): + # Packed uint8 gate/up weights - skip standard merging + if self._gate_up_buffer.is_gate_up_key(key): + return [] + else: + merge_result = self._gate_up_buffer.add(key, tensor) + if merge_result is not None: + return [merge_result] + if self._gate_up_buffer.is_gate_up_key(key): + return [] return [(key, tensor)] def on_load_complete(self) -> List[Tuple[str, torch.Tensor]]: - pending_gu = self._gate_up_buffer.get_pending() - if pending_gu: - warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") + if self._gate_up_buffer is not None: + pending_gu = self._gate_up_buffer.get_pending() + if pending_gu: + warnings.warn(f"Incomplete gate/up merge pairs after loading: {pending_gu}") if self._qkv_buffer is not None: pending_qkv = self._qkv_buffer.get_pending() if pending_qkv: diff --git a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py index cf07c863..300e67cd 100644 --- a/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py +++ b/src/xorl/models/transformers/qwen3_5/modeling_qwen3_5.py @@ -18,6 +18,7 @@ update_causal_mask, ) from xorl.models.layers.attention.backend import get_attention_fn +from xorl.models.layers.fused_projection_lora import project_fused_linear_with_lora from xorl.models.layers.normalization import ( compiled_zero_centered_rms_norm, eager_zero_centered_rms_norm, @@ -73,6 +74,8 @@ def _raise_if_ring_fla_unsupported(config: Qwen3_5Config, ps) -> None: class Qwen3_5MLP(nn.Module): + _supports_fused_gate_up_lora = True + def __init__(self, config): super().__init__() self.hidden_size = config.hidden_size @@ -92,10 +95,17 @@ def unfuse_for_tp(self): def forward(self, x): if hasattr(self, "gate_up_proj"): + gate_up = project_fused_linear_with_lora( + self, + x, + base_name="gate_up_proj", + projection_names=("gate_proj", "up_proj"), + projection_sizes=(self.intermediate_size, self.intermediate_size), + ) if self._use_fused_silu: - x = fused_silu_and_mul(self.gate_up_proj(x)) + x = fused_silu_and_mul(gate_up) else: - gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + gate, up = gate_up.chunk(2, dim=-1) x = self.act_fn(gate) * up else: x = self.act_fn(self.gate_proj(x)) * self.up_proj(x) @@ -411,10 +421,7 @@ def _init_weights(self, module): module.original_inv_freq = module.inv_freq def get_checkpoint_handler(self, **kwargs): - # When unfused for TP, checkpoint keys (q_proj, k_proj, v_proj, gate_proj, - # up_proj) already match the model's parameter names - no merging needed. - if getattr(self, "_unfused_for_tp", False): - return None + unfused = getattr(self, "_unfused_for_tp", False) weights_path = kwargs.get("weights_path", None) is_prequantized = detect_prequantized_checkpoint(weights_path) @@ -435,6 +442,9 @@ def get_checkpoint_handler(self, **kwargs): linear_key_dim=self.config.linear_num_key_heads * self.config.linear_key_head_dim, linear_value_dim=self.config.linear_num_value_heads * self.config.linear_value_head_dim, skip_qkv_merge=True, + # Only the merges are skipped, never the handler: it also remaps the + # GatedDeltaNet in_proj_qkv packing, regardless of how the MLP is stored. + skip_gate_up_merge=unfused, is_prequantized=is_prequantized, exclude_modules=exclude_modules, ) diff --git a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py index 0683f36b..489940ab 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py +++ b/src/xorl/models/transformers/qwen3_5_moe/checkpoint_handler.py @@ -49,7 +49,8 @@ class Qwen3_5MoeCheckpointHandler(CheckpointHandler): - When ``skip_qkv_merge=True``, QKV keys pass through unmerged (model has separate q_proj/k_proj/v_proj after unfuse_for_tp). - When ``skip_gate_up_merge=True``, gate/up keys pass through unmerged - (dense MLP layers have separate gate_proj/up_proj after unfuse_for_tp). + (dense MLP layers and shared experts have separate gate_proj/up_proj after + unfuse_for_tp). - Expert merging is always active (stacking per-expert HF weights). """ diff --git a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py index 9c79f524..97f5e41a 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -7,6 +7,7 @@ from xorl.distributed.parallel_state import get_parallel_state from xorl.distributed.sequence_parallel.strategy import get_cp_strategy +from xorl.lora.fold import lora_merged_forward_enabled from xorl.models.base import XorlPreTrainedModel from xorl.models.checkpoint_handlers.buffers import ( checkpoint_has_per_expert_weights, @@ -74,6 +75,8 @@ def _raise_if_ring_fla_unsupported(config: Qwen3_5MoeConfig, ps) -> None: class Qwen3_5MoeMLP(nn.Module): + _supports_fused_gate_up_lora = True + def __init__(self, config, intermediate_size=None): super().__init__() self.hidden_size = config.hidden_size @@ -90,12 +93,61 @@ def unfuse_for_tp(self): self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False, device=device, dtype=dtype) del self.gate_up_proj + @staticmethod + def _linear_with_contract(module: nn.Linear, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + if getattr(module, "_xorl_bi_trunk_wrapped", False): + from xorl.ops.batch_invariant_ops import _BatchInvariantTrunkLinearFn # noqa: PLC0415 + + return _BatchInvariantTrunkLinearFn.apply(x, weight, module.bias) + return F.linear(x, weight, module.bias) + + def _gate_up_weights_for_forward(self) -> tuple[torch.Tensor, torch.Tensor]: + """Return independently folded gate/up weights for the fused base GEMM.""" + if not hasattr(self, "gate_up_proj"): + raise RuntimeError("Fused gate/up weights requested after unfuse_for_tp()") + gate_base, up_base = self.gate_up_proj.weight.split(self.intermediate_size, dim=0) + gate_adapter = getattr(self, "gate_proj", None) + up_adapter = getattr(self, "up_proj", None) + gate_weight = ( + gate_adapter.merged_weight_for_forward(gate_base) + if gate_adapter is not None and lora_merged_forward_enabled(gate_adapter) + else gate_base + ) + up_weight = ( + up_adapter.merged_weight_for_forward(up_base) + if up_adapter is not None and lora_merged_forward_enabled(up_adapter) + else up_base + ) + return gate_weight, up_weight + + def _project_gate_up(self, x: torch.Tensor) -> torch.Tensor: + gate_adapter = getattr(self, "gate_proj", None) + up_adapter = getattr(self, "up_proj", None) + if gate_adapter is None and up_adapter is None: + return self.gate_up_proj(x) + + adapters = tuple(adapter for adapter in (gate_adapter, up_adapter) if adapter is not None) + merged = tuple(lora_merged_forward_enabled(adapter) for adapter in adapters) + if any(merged): + if not all(merged): + raise RuntimeError("Qwen shared-expert gate/up adapters must select merged forward together") + gate_weight, up_weight = self._gate_up_weights_for_forward() + return self._linear_with_contract(self.gate_up_proj, x, torch.cat((gate_weight, up_weight), dim=0)) + + gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + if gate_adapter is not None: + gate = gate + gate_adapter(x).to(gate.dtype) + if up_adapter is not None: + up = up + up_adapter(x).to(up.dtype) + return torch.cat((gate, up), dim=-1) + def forward(self, x): if hasattr(self, "gate_up_proj"): + gate_up = self._project_gate_up(x) if self._use_fused_silu: - x = fused_silu_and_mul(self.gate_up_proj(x)) + x = fused_silu_and_mul(gate_up) else: - gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + gate, up = gate_up.chunk(2, dim=-1) x = self.act_fn(gate) * up else: x = self.act_fn(self.gate_proj(x)) * self.up_proj(x) @@ -424,8 +476,12 @@ def _ep_combine_native( ).to(torch.bfloat16) self._capture_diagnostic_component("moe_native_routed", routed) - w_gu = self.shared_expert.gate_up_proj.weight # [2I, H], gate rows first - w_down = self.shared_expert.down_proj.weight # [H, I] + gate_weight, up_weight = self.shared_expert._gate_up_weights_for_forward() + w_gu = torch.cat((gate_weight, up_weight), dim=0) # [2I, H], gate rows first + down_proj = self.shared_expert.down_proj + w_down = ( + down_proj.merged_weight_for_forward() if lora_merged_forward_enabled(down_proj) else down_proj.weight + ) # [H, I] shard = inter // ep_size lo_s = ep_rank * shard # Retain the decomposed gate only when operand diagnostics request it. diff --git a/src/xorl/models/transformers/qwen3_5_moe/parallelize.py b/src/xorl/models/transformers/qwen3_5_moe/parallelize.py index b1e480c5..73ba3308 100644 --- a/src/xorl/models/transformers/qwen3_5_moe/parallelize.py +++ b/src/xorl/models/transformers/qwen3_5_moe/parallelize.py @@ -3,7 +3,6 @@ from torch.distributed._tensor import Shard from ....distributed.parallel_plan import ParallelPlan -from ...layers.moe import MoEBlock # TP plan for the base model (Qwen3_5MoeModel). @@ -30,15 +29,17 @@ def unfuse_for_tp(model): """Unfuse fused projections for tensor parallelism compatibility. For ALL layers: splits ``qkv_proj`` -> ``q_proj / k_proj / v_proj`` in attention. - For DENSE layers only: splits ``gate_up_proj`` -> ``gate_proj / up_proj`` in MLP. - MoE layers (``MoEBlock``) are left untouched - their expert weights - are not TP-sharded. + For DENSE layers and each MoE block's shared expert: splits ``gate_up_proj`` -> + ``gate_proj / up_proj``. Routed expert weights are left fused - they are not + TP-sharded. """ for layer in model.model.layers: if getattr(layer, "self_attn", None) is not None and hasattr(layer.self_attn, "unfuse_for_tp"): layer.self_attn.unfuse_for_tp() - # Only unfuse dense MLP layers, not MoE blocks - if not isinstance(layer.mlp, MoEBlock): + # An MoE block exposes a shared expert; a dense MLP is the thing to unfuse. + if hasattr(layer.mlp, "shared_expert"): + layer.mlp.shared_expert.unfuse_for_tp() + else: layer.mlp.unfuse_for_tp() model._unfused_for_tp = True # Override HF config's TP plan (may contain incompatible styles like diff --git a/src/xorl/models/transformers/qwen3_5_shared.py b/src/xorl/models/transformers/qwen3_5_shared.py index 29efb43c..6ee40005 100644 --- a/src/xorl/models/transformers/qwen3_5_shared.py +++ b/src/xorl/models/transformers/qwen3_5_shared.py @@ -92,7 +92,6 @@ def _apply_qwen35_gdn_exact(model: torch.nn.Module) -> dict[str, int]: module._exact_batch_invariant_router = is_moe module.router._exact_batch_invariant = is_moe module.router.synthetic_routing_mode = None - module.router.topk_policy = "default" if not norm_modules: raise RuntimeError("Exact Qwen model construction produced no resolved zero-centered RMSNorm modules.") diff --git a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py index dec3cd3e..a77d0b9e 100644 --- a/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py +++ b/src/xorl/models/transformers/qwen3_moe/modeling_qwen3_moe.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os from typing import Optional, Unpack import torch @@ -40,6 +39,7 @@ is_flash_attention, update_causal_mask, ) +from xorl.models.layers.fused_projection_lora import project_fused_linear_with_lora from xorl.models.layers.moe import MoEBlock, MoEExperts from xorl.models.module_utils import MoEGradientCheckpointingLayer from xorl.models.outputs import MoeCausalLMOutput, MoeModelOutput @@ -53,214 +53,9 @@ logger = logging.get_logger(__name__) -def _qwen3_moe_delayed_residual_pair_enabled() -> bool: - return os.environ.get("XORL_QWEN3_MOE_DELAYED_RESIDUAL_PAIR", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - - -def _qwen3_moe_delayed_residual_pair_sglang_rms_enabled() -> bool: - return os.environ.get("XORL_QWEN3_MOE_DELAYED_RESIDUAL_PAIR_SGLANG_RMS", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - - -def _qwen3_moe_delayed_residual_pair_sglang_kernel_rms_enabled() -> bool: - return os.environ.get("XORL_QWEN3_MOE_DELAYED_RESIDUAL_PAIR_SGLANG_KERNEL_RMS", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - - -def _qwen3_moe_post_attention_sglang_rms_enabled(layer_idx: int) -> bool: - if os.environ.get("XORL_QWEN3_MOE_POST_ATTENTION_SGLANG_RMS", "").strip().lower() not in { - "1", - "true", - "yes", - "on", - }: - return False - layer_filter = os.environ.get("XORL_QWEN3_MOE_POST_ATTENTION_SGLANG_RMS_LAYERS", "").strip() - if not layer_filter or layer_filter.lower() in {"all", "*"}: - return True - try: - enabled_layers = {int(item.strip()) for item in layer_filter.split(",") if item.strip()} - except ValueError as exc: - raise ValueError( - "Invalid XORL_QWEN3_MOE_POST_ATTENTION_SGLANG_RMS_LAYERS=" - f"{layer_filter!r}; expected comma-separated layer indices" - ) from exc - return layer_idx in enabled_layers - - -def _qwen3_moe_post_attention_o_proj_partial_residual_enabled(layer_idx: int) -> bool: - if os.environ.get("XORL_QWEN3_MOE_POST_ATTENTION_O_PROJ_PARTIAL_RESIDUAL", "").strip().lower() not in { - "1", - "true", - "yes", - "on", - }: - return False - layer_filter = os.environ.get("XORL_QWEN3_MOE_POST_ATTENTION_O_PROJ_PARTIAL_RESIDUAL_LAYERS", "").strip() - if not layer_filter or layer_filter.lower() in {"all", "*"}: - return True - try: - enabled_layers = {int(item.strip()) for item in layer_filter.split(",") if item.strip()} - except ValueError as exc: - raise ValueError( - "Invalid XORL_QWEN3_MOE_POST_ATTENTION_O_PROJ_PARTIAL_RESIDUAL_LAYERS=" - f"{layer_filter!r}; expected comma-separated layer indices" - ) from exc - return layer_idx in enabled_layers - - -def _qwen3_moe_capture_o_proj_partial_residual_candidates_enabled(layer_idx: int) -> bool: - if os.environ.get("XORL_QWEN3_MOE_CAPTURE_O_PROJ_PARTIAL_RESIDUAL_CANDIDATES", "").strip().lower() not in { - "1", - "true", - "yes", - "on", - }: - return False - layer_filter = os.environ.get("XORL_QWEN3_MOE_CAPTURE_O_PROJ_PARTIAL_RESIDUAL_CANDIDATES_LAYERS", "").strip() - if not layer_filter or layer_filter.lower() in {"all", "*"}: - return True - try: - enabled_layers = {int(item.strip()) for item in layer_filter.split(",") if item.strip()} - except ValueError as exc: - raise ValueError( - "Invalid XORL_QWEN3_MOE_CAPTURE_O_PROJ_PARTIAL_RESIDUAL_CANDIDATES_LAYERS=" - f"{layer_filter!r}; expected comma-separated layer indices" - ) from exc - return layer_idx in enabled_layers - - -def _qwen3_moe_delayed_residual_pair_tp_shard_carry_enabled() -> bool: - return os.environ.get("XORL_QWEN3_MOE_DELAYED_RESIDUAL_PAIR_TP_SHARD_CARRY", "").strip().lower() in { - "1", - "true", - "yes", - "on", - } - - -def _is_delayed_residual_pair(value) -> bool: - return ( - isinstance(value, tuple) - and len(value) == 2 - and isinstance(value[0], torch.Tensor) - and isinstance(value[1], torch.Tensor) - ) - - -def _get_delayed_pair_moe_tp_shards(hidden_delta: torch.Tensor): - if not _qwen3_moe_delayed_residual_pair_tp_shard_carry_enabled(): - return None - shards = getattr(hidden_delta, "_xorl_sglang_moe_tp_shards", None) - if not shards: - return None - if not all(isinstance(shard, torch.Tensor) for shard in shards): - raise TypeError("_xorl_sglang_moe_tp_shards must contain tensors") - return tuple(shards) - - -def _sum_delayed_pair_moe_tp_shards(hidden_delta: torch.Tensor, dtype: torch.dtype) -> Optional[torch.Tensor]: - shards = _get_delayed_pair_moe_tp_shards(hidden_delta) - if shards is None: - return None - - shard_sum = shards[0].to(dtype) - for shard in shards[1:]: - shard_sum = shard_sum + shard.to(shard_sum.dtype) - return shard_sum.to(dtype) - - -def _materialize_moe_tp_shards_with_residual(hidden_delta: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: - shard_sum = _sum_delayed_pair_moe_tp_shards(hidden_delta, residual.dtype) - if shard_sum is None: - return hidden_delta + residual - return residual + shard_sum - - -def _materialize_delayed_residual_pair(value): - if _is_delayed_residual_pair(value): - return _materialize_moe_tp_shards_with_residual(value[0], value[1]) - return value - - -def _get_o_proj_tp_partials(hidden_states: torch.Tensor): - partials = getattr(hidden_states, "_xorl_o_proj_tp_partials", None) - if not partials: - return None - if not all(isinstance(partial, torch.Tensor) for partial in partials): - raise TypeError("_xorl_o_proj_tp_partials must contain tensors") - return tuple(partials) - - -def _sum_o_proj_tp_partials(partials: tuple[torch.Tensor, ...], dtype: torch.dtype) -> torch.Tensor: - partial_sum = partials[0].to(dtype) - for partial in partials[1:]: - partial_sum = partial_sum + partial.to(partial_sum.dtype) - return partial_sum.to(dtype) - - -def _materialize_o_proj_partial_residual( - hidden_states: torch.Tensor, - residual: torch.Tensor, - partials: tuple[torch.Tensor, ...], -) -> tuple[torch.Tensor, torch.Tensor]: - mode = os.environ.get("XORL_QWEN3_MOE_POST_ATTENTION_O_PROJ_PARTIAL_RESIDUAL_MODE", "sum_then_residual") - mode = mode.strip().lower().replace("-", "_") - return _materialize_o_proj_partial_residual_mode(mode, hidden_states, residual, partials) - - -def _materialize_o_proj_partial_residual_mode( - mode: str, - hidden_states: torch.Tensor, - residual: torch.Tensor, - partials: tuple[torch.Tensor, ...], -) -> tuple[torch.Tensor, torch.Tensor]: - if mode in {"split_output", "output"}: - partial_sum = hidden_states - return partial_sum, residual + partial_sum.to(residual.dtype) - if mode in {"sum_then_residual", "partial_sum"}: - partial_sum = _sum_o_proj_tp_partials(partials, hidden_states.dtype) - return partial_sum, residual + partial_sum.to(residual.dtype) - if mode in {"residual_then_partials", "sequential_bf16", "seq_bf16"}: - residual_out = residual - for partial in partials: - residual_out = (residual_out + partial.to(residual_out.dtype)).to(residual.dtype) - partial_sum = _sum_o_proj_tp_partials(partials, hidden_states.dtype) - return partial_sum, residual_out - if mode in {"fp32_sum_then_residual", "sum_fp32"}: - partial_sum = partials[0].to(torch.float32) - for partial in partials[1:]: - partial_sum = partial_sum + partial.to(torch.float32) - residual_out = (residual.to(torch.float32) + partial_sum).to(residual.dtype) - return partial_sum.to(hidden_states.dtype), residual_out - raise ValueError( - "Invalid XORL_QWEN3_MOE_POST_ATTENTION_O_PROJ_PARTIAL_RESIDUAL_MODE=" - f"{mode!r}; expected split_output, sum_then_residual, residual_then_partials, or fp32_sum_then_residual" - ) - - -_O_PROJ_PARTIAL_RESIDUAL_CANDIDATE_MODES = ( - "split_output", - "sum_then_residual", - "residual_then_partials", - "fp32_sum_then_residual", -) - - class Qwen3MoeMLP(nn.Module): + _supports_fused_gate_up_lora = True + def __init__(self, config, intermediate_size=None): super().__init__() self.hidden_size = config.hidden_size @@ -282,10 +77,17 @@ def unfuse_for_tp(self): def forward(self, x): if hasattr(self, "gate_up_proj"): + gate_up = project_fused_linear_with_lora( + self, + x, + base_name="gate_up_proj", + projection_names=("gate_proj", "up_proj"), + projection_sizes=(self.intermediate_size, self.intermediate_size), + ) if self._use_fused_silu: - x = fused_silu_and_mul(self.gate_up_proj(x)) + x = fused_silu_and_mul(gate_up) else: - gate, up = self.gate_up_proj(x).chunk(2, dim=-1) + gate, up = gate_up.chunk(2, dim=-1) x = self.act_fn(gate) * up else: x = self.act_fn(self.gate_proj(x)) * self.up_proj(x) @@ -472,68 +274,16 @@ def __init__(self, config: Qwen3MoeConfig, layer_idx: int): self.mlp.experts.layer_idx = layer_idx else: self.mlp = Qwen3MoeMLP(config, intermediate_size=config.intermediate_size) - self._delay_moe_residual_output = ( - _qwen3_moe_delayed_residual_pair_enabled() and is_sparse_layer and layer_idx < config.num_hidden_layers - 1 - ) - - def _capture_o_proj_partial_residual_candidates( - self, - hidden_states: torch.Tensor, - residual: torch.Tensor, - partials: tuple[torch.Tensor, ...], - ) -> None: - if not _qwen3_moe_capture_o_proj_partial_residual_candidates_enabled(self.layer_idx): - return - for mode in _O_PROJ_PARTIAL_RESIDUAL_CANDIDATE_MODES: - partial_sum, partial_residual = _materialize_o_proj_partial_residual_mode( - mode, - hidden_states, - residual, - partials, - ) - self._capture_diagnostic_component(f"post_attention_o_proj_partial_sum_{mode}", partial_sum) - self._capture_diagnostic_component(f"post_attention_partial_residual_{mode}", partial_residual) def _pre_mlp_forward(self, hidden_states, attention_mask=None, position_embeddings=None, **kwargs): - self._capture_diagnostic_component( - "materialized_layer_input", _materialize_delayed_residual_pair(hidden_states) + residual = hidden_states + # Family contract: layer-0 input norm is a no-residual site; at + # layer>0 the pre-summed input sits on the serving residual tree + # (the 2026-07-04 norm-seed fix, now declared explicitly). + hidden_states = self.input_layernorm( + hidden_states, + family=RMS_NORM_FAMILY_RESIDUAL_TREE if self.layer_idx > 0 else RMS_NORM_FAMILY_NO_RESIDUAL, ) - if _is_delayed_residual_pair(hidden_states): - hidden_states, residual = hidden_states - self._capture_diagnostic_component("delayed_pair_delta", hidden_states) - self._capture_diagnostic_component("delayed_pair_residual", residual) - shard_sum = _sum_delayed_pair_moe_tp_shards(hidden_states, residual.dtype) - self._capture_diagnostic_component("delayed_pair_shard_sum", shard_sum) - if shard_sum is not None: - self._capture_diagnostic_component("delayed_pair_shard_materialized", residual + shard_sum) - shards = _get_delayed_pair_moe_tp_shards(hidden_states) - # Delayed-pair diagnostic sites keep their env-gated legacy dispatch - # (undeclared family; single-tensor calls fall to the aten interpose - # when the rms env is unset) — bits frozen, tripwire warns. - if shards is not None: - residual = _materialize_moe_tp_shards_with_residual(hidden_states, residual) - hidden_states = self.input_layernorm( - residual, - force_sglang_residual=_qwen3_moe_delayed_residual_pair_sglang_rms_enabled(), - force_sglang_residual_kernel=_qwen3_moe_delayed_residual_pair_sglang_kernel_rms_enabled(), - ) - else: - hidden_states, residual = self.input_layernorm( - hidden_states, - residual=residual, - prenorm=True, - force_sglang_residual=_qwen3_moe_delayed_residual_pair_sglang_rms_enabled(), - force_sglang_residual_kernel=_qwen3_moe_delayed_residual_pair_sglang_kernel_rms_enabled(), - ) - else: - residual = hidden_states - # Family contract: layer-0 input norm is a no-residual site; at - # layer>0 the pre-summed input sits on the serving residual tree - # (the 2026-07-04 norm-seed fix, now declared explicitly). - hidden_states = self.input_layernorm( - hidden_states, - family=RMS_NORM_FAMILY_RESIDUAL_TREE if self.layer_idx > 0 else RMS_NORM_FAMILY_NO_RESIDUAL, - ) self._capture_diagnostic_component("input_norm_residual", residual) self._capture_diagnostic_component("input_norm", hidden_states) hidden_states, _ = self.self_attn( @@ -542,36 +292,14 @@ def _pre_mlp_forward(self, hidden_states, attention_mask=None, position_embeddin position_embeddings=position_embeddings, **kwargs, ) - o_proj_partials = _get_o_proj_tp_partials(hidden_states) - if o_proj_partials is not None: - self._capture_o_proj_partial_residual_candidates(hidden_states, residual, o_proj_partials) - if o_proj_partials is not None and _qwen3_moe_post_attention_o_proj_partial_residual_enabled(self.layer_idx): - partial_sum, partial_residual = _materialize_o_proj_partial_residual( - hidden_states, residual, o_proj_partials - ) - self._capture_diagnostic_component("post_attention_o_proj_partial_sum", partial_sum) - self._capture_diagnostic_component("post_attention_partial_residual", partial_residual) - self._capture_diagnostic_component("post_attention_norm_input", partial_residual) - norm_output = self.post_attention_layernorm( - partial_residual, - force_sglang_residual=_qwen3_moe_post_attention_sglang_rms_enabled(self.layer_idx), - ) - if isinstance(norm_output, tuple): - hidden_states, returned_residual = norm_output - residual = partial_residual if returned_residual is None else returned_residual - else: - hidden_states = norm_output - residual = partial_residual - else: - self._capture_diagnostic_component("post_attention_norm_input", hidden_states) - self._capture_diagnostic_component("post_attention_norm_residual", residual) - hidden_states, residual = self.post_attention_layernorm( - hidden_states, - residual=residual, - prenorm=True, - force_sglang_residual=_qwen3_moe_post_attention_sglang_rms_enabled(self.layer_idx), - family=RMS_NORM_FAMILY_RESIDUAL_TREE, - ) + self._capture_diagnostic_component("post_attention_norm_input", hidden_states) + self._capture_diagnostic_component("post_attention_norm_residual", residual) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, + residual=residual, + prenorm=True, + family=RMS_NORM_FAMILY_RESIDUAL_TREE, + ) self._capture_diagnostic_component("post_attention_norm", hidden_states) self._capture_diagnostic_component("post_attention_residual", residual) return hidden_states, residual @@ -800,7 +528,7 @@ def forward( if decoder_layer is None: # PP: pruned layer continue if output_hidden_states: - all_hidden_states += (_materialize_delayed_residual_pair(hidden_states),) + all_hidden_states += (hidden_states,) if _grad_ckpt_method == "recompute_full_layer": # Recompute entire layer in backward (including dispatch + combine) @@ -851,7 +579,6 @@ def forward( all_router_logits += (layer_outputs[-1],) # PP support: norm may be None on non-last stages - hidden_states = _materialize_delayed_residual_pair(hidden_states) if self.norm is not None: hidden_states = self.norm(hidden_states) if output_hidden_states: diff --git a/src/xorl/ops/batch_invariant_ops.py b/src/xorl/ops/batch_invariant_ops.py index af77cc8a..4026bd98 100644 --- a/src/xorl/ops/batch_invariant_ops.py +++ b/src/xorl/ops/batch_invariant_ops.py @@ -2028,6 +2028,7 @@ def wrap_trunk_linears_batch_invariant( from xorl.lora.fold import lora_merged_forward_enabled # noqa: PLC0415 from xorl.lora.modules.base import LoraModule # noqa: PLC0415 + from xorl.lora.modules.delta_linear import LoraDeltaLinear # noqa: PLC0415 from xorl.lora.modules.linear import LoraLinear # noqa: PLC0415 if is_batch_invariant_mode_enabled(): @@ -2054,6 +2055,13 @@ def _forward_lora_merged(self, input): continue if ".experts." in f".{module_name}.": continue + if isinstance(module, LoraDeltaLinear) and module_name.endswith( + (".mlp.shared_expert.gate_proj", ".mlp.shared_expert.up_proj") + ): + # These are factor-only children of the fused gate_up_proj. The + # parent base GEMM is wrapped and Qwen3_5MoeMLP folds both logical + # projections into that one contracted call. + continue if type(module) is LoraLinear and lora_merged_forward_enabled(module): # Merged-forward contract lane: the adapted linear serves and trains # through the folded weight, so the trunk contract composes. diff --git a/src/xorl/ops/bi_families_v2.py b/src/xorl/ops/bi_families_v2.py index a93e0d4a..db5f896b 100644 --- a/src/xorl/ops/bi_families_v2.py +++ b/src/xorl/ops/bi_families_v2.py @@ -23,8 +23,6 @@ # recorded under v1 must be re-taken under v2, and both engines must flip # together. -import os - import torch import triton import triton.language as tl @@ -55,21 +53,16 @@ def _select_nonexact_families() -> None: def families_v2_enabled() -> bool: """Return the selected reduction family for the current process. - Exact model programs select their family structurally and ignore the - legacy rollback variables. Without an exact model selection, preserve the - pre-existing non-exact behavior for compatibility. + Exact model programs select their family structurally. Ordinary models use + the current v2 family; there is no process-environment rollback path. """ if _EXACT_FAMILIES_VERSION is not None: return _EXACT_FAMILIES_VERSION == "v2" - return not any(os.getenv(v, "1").lower() in _V2_OFF for v in FAMILIES_V2_ENV_VARS) + return True # Contract constants (bit-relevant; never tuning axes). V2_NORM_BLOCK_H = 4096 # per-chunk tree width for hidden-dim norms -V2_QK_MAX_HEAD_DIM = 256 - -FAMILIES_V2_ENV_VARS = ("XORL_FAMILIES_V2", "SGLANG_FAMILIES_V2") -_V2_OFF = ("0", "false", "no") @triton.jit @@ -172,59 +165,6 @@ def _rms_norm_v2_kernel( tl.store(out_ptr + row * stride_out + cols, y.to(out_ptr.dtype.element_ty), mask=mask) -QK_V2_ROWS_PER_PROG = 16 # head-rows per program (perf-only, NOT bit-relevant) - - -@triton.jit -def _qk_norm_v2_kernel( - x_ptr, - w_ptr, - out_ptr, - n_rows, - n_heads, - head_dim, - stride_x_tok, - stride_x_head, - stride_out_tok, - stride_out_head, - eps, - ZERO_CENTERED: tl.constexpr, - BLOCK_D: tl.constexpr, - ROWS: tl.constexpr, -): - """Family-1' strided qk-norm: reads head rows straight out of the packed - qkv projection (no reshape/contiguous copies). Same tree as - _rms_norm_v2_kernel (single chunk: head_dim <= V2_QK_MAX_HEAD_DIM). - Each program handles ROWS independent head-rows (row batching is grid - shape only — per-row math identical, like BLOCK_M in the GEMM). - In-place safe per row: the full head row is loaded before any store. - """ - pid = tl.program_id(0) - rows = pid * ROWS + tl.arange(0, ROWS) - row_mask = rows < n_rows - rows_safe = tl.where(row_mask, rows, 0) - tok = rows_safe // n_heads - head = rows_safe % n_heads - d = tl.arange(0, BLOCK_D) - col_mask = d < head_dim - mask = row_mask[:, None] & col_mask[None, :] - base = tok * stride_x_tok + head * stride_x_head - x = tl.load(x_ptr + base[:, None] + d[None, :], mask=mask, other=0.0).to(tl.float32) - total = _pairwise_tree_sum_rows(x * x, BLOCK_D) - var = total / head_dim.to(tl.float32) - inv_rms = tl.rsqrt(var + eps) - w = tl.load(w_ptr + d, mask=col_mask, other=0.0).to(tl.float32) - if ZERO_CENTERED: - w = 1.0 + w - y = x * inv_rms[:, None] * w[None, :] - out_base = tok * stride_out_tok + head * stride_out_head - tl.store( - out_ptr + out_base[:, None] + d[None, :], - y.to(out_ptr.dtype.element_ty), - mask=mask, - ) - - def rms_norm_v2( x: torch.Tensor, weight: torch.Tensor, @@ -294,55 +234,6 @@ def _rms_norm_v2_fused(x, weight, eps, residual, zero_centered): return out -def qk_norm_v2( - x: torch.Tensor, - weight: torch.Tensor, - eps: float = 1e-6, - *, - head_dim: int, - out: torch.Tensor | None = None, - zero_centered: bool = False, -): - """v2 per-head qk-norm over strided input. - - ``x``: ``[T, n_heads * head_dim]`` view into the packed qkv output — any - row stride, unit element stride, heads contiguous within a row. ``out`` - defaults to a fresh tensor (trainer); pass ``out=x`` for in-place (serving). - """ - assert x.ndim == 2 and x.stride(1) == 1 - assert x.dtype == torch.bfloat16, "families v2 is bf16-only (contract dtype)" - assert x.shape[1] % head_dim == 0 - assert head_dim <= V2_QK_MAX_HEAD_DIM - assert weight.ndim == 1 and weight.shape[0] == head_dim - assert x.is_cuda - T = x.shape[0] - n_heads = x.shape[1] // head_dim - weight = weight.contiguous() - if out is None: - out = torch.empty_like(x) - else: - assert out.shape == x.shape and out.stride(1) == 1 and out.dtype == x.dtype - if T > 0: - n_rows = T * n_heads - _qk_norm_v2_kernel[(triton.cdiv(n_rows, QK_V2_ROWS_PER_PROG),)]( - x, - weight, - out, - n_rows, - n_heads, - head_dim, - x.stride(0), - head_dim, - out.stride(0), - head_dim, - eps, - ZERO_CENTERED=zero_centered, - BLOCK_D=triton.next_power_of_2(head_dim), - ROWS=QK_V2_ROWS_PER_PROG, - ) - return out - - # --------------------------------------------------------------------------- # Head v2 — online-LSE lm-head (component 3, design note §5) # diff --git a/src/xorl/ops/block_fp8_native.py b/src/xorl/ops/block_fp8_native.py index 4f45d9ea..9fc8c338 100644 --- a/src/xorl/ops/block_fp8_native.py +++ b/src/xorl/ops/block_fp8_native.py @@ -85,35 +85,6 @@ def unpack_float32_as_fp8(packed: torch.Tensor, shape: tuple[int, ...]) -> torch return packed.contiguous().view(torch.uint8).view(_FP8_DTYPE).reshape(shape) -def validate_native_fp8_state_metadata( - module: nn.Module, - metadata: dict[str, tuple[torch.dtype, tuple[int, ...]]], - *, - prefix: str = "", -) -> None: - """Fail before DCP load if serialized dtype/shape metadata can cast bytes. - - DCP callers must build ``metadata`` from the checkpoint reader before - invoking ``set_model_state_dict``. State-dict hooks below cover ordinary - ``load_state_dict``; this preflight covers loaders that copy shards without - calling module hooks. - """ - - expected = { - f"{prefix}{name}": (parameter.dtype, tuple(parameter.shape)) - for name, parameter in module.named_parameters() - if "packed_weight_f32" in name or name.endswith("weight_scale_inv") - } - missing = sorted(set(expected) - set(metadata)) - mismatched = { - name: (metadata[name], contract) - for name, contract in expected.items() - if name in metadata and metadata[name] != contract - } - if missing or mismatched: - raise ValueError(f"Native FP8 DCP metadata mismatch: missing={missing[:8]} mismatched={mismatched}") - - def validate_native_fp8_dcp_checkpoint( checkpoint_path: str, expected_state: dict[str, torch.Tensor], @@ -370,5 +341,4 @@ def forward( "pack_fp8_as_float32", "unpack_float32_as_fp8", "validate_native_fp8_dcp_checkpoint", - "validate_native_fp8_state_metadata", ] diff --git a/src/xorl/ops/dsv4/cp_utils.py b/src/xorl/ops/dsv4/cp_utils.py index ba209920..fee36938 100644 --- a/src/xorl/ops/dsv4/cp_utils.py +++ b/src/xorl/ops/dsv4/cp_utils.py @@ -141,6 +141,6 @@ def get_freqs_cis_for_cp( "DSv4 RoPE cache is too short for this context-parallel slice: " f"need positions [{start}, {stop}) with stride {stride}, " f"but freqs_cis only has {freqs_cis.size(0)} positions. " - "Increase XORL_DSV4_ROPE_MAX_SEQ_LEN or config.max_position_embeddings." + "Increase config.max_position_embeddings." ) return result diff --git a/src/xorl/ops/dsv4/kernel/tilelang_indexer.py b/src/xorl/ops/dsv4/kernel/tilelang_indexer.py deleted file mode 100644 index 29b9a87f..00000000 --- a/src/xorl/ops/dsv4/kernel/tilelang_indexer.py +++ /dev/null @@ -1,98 +0,0 @@ -# ruff: noqa -"""TileLang-based DSA Indexer for DeepSeek-V4. - -Adapts GLM-5's lighting_indexer to V4's SBHD data layout and causal masking. -Provides both a low-level per-sample interface and a batched autograd Function. -""" - -import torch - -from .tilelang_indexer_bwd import batched_indexer_bwd -from .tilelang_indexer_fwd import _make_causal_cu_seqlens, batched_indexer_fwd - - -def pytorch_extract_topk_scores(logits, topk_indices, dim=-1): - valid_mask = topk_indices != -1 - safe_indices = topk_indices.clamp(min=0).to(torch.int64) - scores = torch.gather(logits, dim=dim, index=safe_indices) - scores = torch.where(valid_mask, scores, float("-inf")) - return scores - - -class V4IndexerFunction(torch.autograd.Function): - """Autograd function for V4 tilelang indexer. - - Inputs are in V4's native SBHD layout: - q: [seqlen, batch, heads, dim] bf16 - k: [seqlen_kv, batch, dim] bf16 - weights: [seqlen, batch, heads] fp32 - """ - - @staticmethod - def forward( - ctx, - index_q: torch.Tensor, - index_k: torch.Tensor, - weights: torch.Tensor, - compress_ratio: int, - topk: int, - topk_indices: torch.Tensor | None = None, - ): - seqlen_q = index_q.shape[0] - seq_len_kv = index_k.shape[0] - - cu_seqlen_ks, cu_seqlen_ke = _make_causal_cu_seqlens(seqlen_q, seq_len_kv, compress_ratio, index_q.device) - - # [batch, seqlen, seqlen_kv] - logits = batched_indexer_fwd(index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke) - - if topk_indices is None: - actual_topk = min(topk, seq_len_kv) - # torch.topk on bf16 is ~1.6x faster than fp32 (8.6 vs 13.7 ms at - # B=1, S=32k, S_kv=32k, topk=512 on H100). The cast costs ~1.4 ms - # so net win is ~3 ms per indexer call. Indices are dtype-independent - # so no loss; the fp32 scores below are recomputed via gather. - logits_for_topk = logits.bfloat16() - _, topk_indices = torch.topk(logits_for_topk, actual_topk, dim=-1) - topk_indices = topk_indices.to(torch.int32) - # Use the (already-masked) bf16 -inf as the sentinel detector. - sentinel_scores = torch.gather(logits_for_topk, -1, topk_indices.long()) - topk_indices = topk_indices.masked_fill(sentinel_scores == -torch.inf, -1) - - index_score = pytorch_extract_topk_scores(logits, topk_indices) - - ctx.save_for_backward(index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, topk_indices) - ctx.compress_ratio = compress_ratio - ctx.topk = topk - return index_score, topk_indices - - @staticmethod - def backward(ctx, grad_scores, grad_indices): - index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, topk_indices = ctx.saved_tensors - grad_q, grad_w, grad_k = batched_indexer_bwd(index_q, weights, index_k, topk_indices, grad_scores) - return grad_q, grad_k, grad_w, None, None, None - - -def v4_lighting_indexer( - index_q: torch.Tensor, - index_k: torch.Tensor, - weights: torch.Tensor, - compress_ratio: int, - topk: int, - topk_indices: torch.Tensor | None = None, -): - """Main entry point for V4 tilelang indexer. - - Args: - index_q: [seqlen, batch, heads, dim] bf16 - index_k: [seqlen_kv, batch, dim] bf16 - weights: [seqlen, batch, heads] fp32 - compress_ratio: compression ratio (4 for C4 layers) - topk: number of top-k indices to select - topk_indices: optional pre-computed topk indices [batch, seqlen, topk] int32 - - Returns: - index_score: [batch, seqlen, topk] fp32 - topk_indices: [batch, seqlen, topk] int32 - """ - return V4IndexerFunction.apply(index_q, index_k, weights, compress_ratio, topk, topk_indices) diff --git a/src/xorl/ops/dsv4/kernel/tilelang_indexer_bwd.py b/src/xorl/ops/dsv4/kernel/tilelang_indexer_bwd.py deleted file mode 100644 index 4296bf99..00000000 --- a/src/xorl/ops/dsv4/kernel/tilelang_indexer_bwd.py +++ /dev/null @@ -1,245 +0,0 @@ -# ruff: noqa -# Adapted from miles_plugins/models/glm5/ops/tilelang_indexer_bwd.py for DeepSeek-V4. -import tilelang as tl -import tilelang.language as T -import torch - -BF16 = T.bfloat16 -FP32 = T.float32 -INT32 = T.int32 - -pass_configs = { - tl.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tl.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, -} - - -@tl.jit(pass_configs=pass_configs) -def tl_indexer_bwd_impl( - heads: int, - dim: int, - topk: int, - block_I: int = 32, - num_stages: int = 0, - num_threads: int = 128, -): - assert num_stages == 0 - assert topk == tl.math.next_power_of_2(topk) - assert topk % block_I == 0 - assert heads <= 64 and heads % 8 == 0 - seq_len = T.symbolic("seq_len") - q_seq_len = T.symbolic("q_seq_len") - - dtype: str = BF16 - accum_dtype: str = FP32 - index_q_shape = [q_seq_len, heads, dim] - weights_shape = [q_seq_len, heads] - index_k_shape = [seq_len, dim] - shape_p = [q_seq_len, topk] - topk_indices_shape = [q_seq_len, topk] - - pad_heads = heads - if heads < 16: - pad_heads = 16 - - @T.prim_func - def tl_indexer_bwd_kernel( - IndexQ: T.Tensor(index_q_shape, dtype), - IndexK: T.Tensor(index_k_shape, dtype), - Weights: T.Tensor(weights_shape, FP32), - TopkIndices: T.Tensor(topk_indices_shape, INT32), - OGrad: T.Tensor(shape_p, FP32), - dIndexQ: T.Tensor(index_q_shape, dtype), - dWeights: T.Tensor(weights_shape, FP32), - dIndexK: T.Tensor(index_k_shape, FP32), - ): - with T.Kernel(q_seq_len, threads=num_threads) as (bx): - index_q_shared = T.alloc_shared([pad_heads, dim], dtype=FP32) - weights_shared = T.alloc_shared([pad_heads], dtype=FP32) - index_k_shared = T.alloc_shared([block_I, dim], dtype=FP32) - indices_shared = T.alloc_shared([block_I], dtype=INT32) - d_index_q_frag = T.alloc_fragment([pad_heads, dim], dtype=accum_dtype) - d_weights_frag = T.alloc_fragment([pad_heads], dtype=accum_dtype) - d_index_k_frag = T.alloc_fragment([block_I, dim], dtype=accum_dtype) - logits = T.alloc_fragment((block_I, pad_heads), dtype=accum_dtype) - _logits = T.alloc_shared((block_I, pad_heads), dtype=accum_dtype) - grad = T.alloc_shared([block_I], dtype=FP32) - - num_blocks = T.ceildiv(topk, block_I) - for i, j in T.Parallel(pad_heads, dim): - index_q_shared[i, j] = T.if_then_else(i < heads, IndexQ[bx, i, j], 0) - for i in T.Parallel(heads): - weights_shared[i] = Weights[bx, i] - - T.fill(d_index_q_frag, 0) - T.fill(d_weights_frag, 0) - - for bi_i in T.serial(num_blocks): - for i in T.Parallel(block_I): - if bi_i * block_I + i < topk: - indices_shared[i] = TopkIndices[bx, bi_i * block_I + i] - grad[i] = OGrad[bx, bi_i * block_I + i] - - T.sync_threads() - for i, j in T.Parallel(block_I, dim): - index_k_shared[i, j] = T.if_then_else( - indices_shared[i] > -1 and indices_shared[i] < seq_len, IndexK[indices_shared[i], j], 0 - ) - - T.sync_threads() - T.gemm( - index_k_shared, - index_q_shared, - logits, - transpose_A=False, - transpose_B=True, - clear_accum=True, - ) - for i, j in T.Parallel(block_I, heads): - logits[i, j] = T.max(logits[i, j], 0) - - d_weights_i = T.alloc_fragment((block_I, pad_heads), accum_dtype) - for i, j in T.Parallel(block_I, heads): - d_weights_i[i, j] = grad[i] * logits[i, j] - T.reduce_sum(d_weights_i, d_weights_frag, dim=0, clear=False) - - for i, j in T.Parallel(block_I, pad_heads): - _logits[i, j] = T.if_then_else(logits[i, j] > 0 and j < heads, grad[i] * weights_shared[j], 0) - T.sync_threads() - T.gemm( - _logits, - index_k_shared, - d_index_q_frag, - transpose_A=True, - transpose_B=False, - clear_accum=False, - ) - - T.gemm( - _logits, - index_q_shared, - d_index_k_frag, - transpose_A=False, - transpose_B=False, - clear_accum=True, - ) - - for i, j in T.Parallel(block_I, dim): - if indices_shared[i] > -1 and indices_shared[i] < seq_len: - T.atomic_add(dIndexK[indices_shared[i], j], d_index_k_frag[i, j]) - - T.copy(d_index_q_frag[:heads, :], dIndexQ[bx, :, :]) - T.copy(d_weights_frag[:heads], dWeights[bx, :]) - - return tl_indexer_bwd_kernel - - -def indexer_bwd_interface( - index_q: torch.Tensor, - weights: torch.Tensor, - index_k: torch.Tensor, - topk_indices: torch.Tensor, - grad_scores: torch.Tensor, -): - """Backward interface for a single batch element. - - Args: - index_q: [seq_len, heads, dim] bf16 - weights: [seq_len, heads] fp32 - index_k: [seq_len_kv, dim] bf16 - topk_indices: [seq_len, topk] int32 - grad_scores: [seq_len, topk] fp32 - - Returns: - grad_q: [seq_len, heads, dim] bf16 - grad_w: [seq_len, heads] fp32 - grad_k: [seq_len_kv, dim] fp32 - """ - _, head_num, head_dim = index_q.shape - k_top = topk_indices.shape[1] - - grad_scores = grad_scores.contiguous() - grad_q = torch.empty_like(index_q) - grad_w = torch.empty_like(weights, dtype=torch.float32) - grad_k = torch.zeros_like(index_k, dtype=torch.float32) - - # Pad topk to block_I=32 boundary (kernel requires topk % block_I == 0 and topk >= 32) - padded_topk = max(k_top, 32) - padded_topk = ((padded_topk + 31) // 32) * 32 - if padded_topk != k_top: - pad_size = padded_topk - k_top - topk_indices = torch.cat( - [ - topk_indices, - torch.full((topk_indices.shape[0], pad_size), -1, device=topk_indices.device, dtype=topk_indices.dtype), - ], - dim=1, - ).contiguous() - grad_scores = torch.cat( - [ - grad_scores, - torch.zeros((grad_scores.shape[0], pad_size), device=grad_scores.device, dtype=grad_scores.dtype), - ], - dim=1, - ).contiguous() - - # block_I=16, num_threads=64 is 17% faster than the upstream defaults - # (32/128) on DSv4 indexer shape (heads=64, dim=128, topk=512, - # S=S_kv=32k on H100): 11.09 ms vs 13.32 ms. GLM-5 uses (64/256) - # for their heads=32 shape; that's *slower* on DSv4 (15.03 ms) - # because the 2x head count changes the warp-tile arithmetic. - tl_indexer_bwd_impl( - head_num, - head_dim, - padded_topk, - block_I=16, - num_threads=64, - )( - index_q.contiguous(), - index_k.contiguous(), - weights.squeeze(-1).contiguous(), - topk_indices.contiguous(), - grad_scores, - grad_q, - grad_w.squeeze(-1), - grad_k, - ) - - return grad_q, grad_w, grad_k - - -def batched_indexer_bwd(index_q, weights, index_k, topk_indices, grad_scores): - """Batched backward: loops over batch dim. - - Args: - index_q: [seqlen, batch, heads, dim] bf16 - weights: [seqlen, batch, heads] fp32 - index_k: [seqlen_kv, batch, dim] bf16 - topk_indices: [batch, seqlen, topk] int32 - grad_scores: [batch, seqlen, topk] fp32 - - Returns: - grad_q: [seqlen, batch, heads, dim] bf16 - grad_w: [seqlen, batch, heads] fp32 - grad_k: [seqlen_kv, batch, dim] fp32 - """ - seqlen, batch, heads, dim = index_q.shape - seq_len_kv = index_k.shape[0] - - all_grad_q = torch.empty_like(index_q) - all_grad_w = torch.empty(seqlen, batch, heads, device=index_q.device, dtype=torch.float32) - all_grad_k = torch.zeros(seq_len_kv, batch, dim, device=index_q.device, dtype=torch.float32) - - for b in range(batch): - gq, gw, gk = indexer_bwd_interface( - index_q[:, b, :, :].contiguous(), - weights[:, b, :].contiguous(), - index_k[:, b, :].contiguous(), - topk_indices[b].contiguous(), - grad_scores[b].contiguous(), - ) - all_grad_q[:, b, :, :] = gq - all_grad_w[:, b, :] = gw - all_grad_k[:, b, :] = gk - - return all_grad_q, all_grad_w, all_grad_k diff --git a/src/xorl/ops/dsv4/rope.py b/src/xorl/ops/dsv4/rope.py index 918d7031..7aba15c5 100644 --- a/src/xorl/ops/dsv4/rope.py +++ b/src/xorl/ops/dsv4/rope.py @@ -8,7 +8,6 @@ """ import math -import os from functools import lru_cache import torch @@ -86,11 +85,9 @@ def wrapped_precompute_freqs_cis(config, rope_head_dim: int, base: float, yarn_d beta_slow = float(rope_params.get("beta_slow", 1.0)) # Full-weight V4 configs advertise their extended YaRN context in - # max_position_embeddings. Keep the env override for tests/profiling runs - # that intentionally want a smaller cache than the model maximum. - max_seq_len = int( - os.environ.get("XORL_DSV4_ROPE_MAX_SEQ_LEN", getattr(config, "max_position_embeddings", original_max_pos)) - ) + # max_position_embeddings. Tests and profiling configs should declare the + # smaller cache directly instead of overriding the model contract per process. + max_seq_len = int(getattr(config, "max_position_embeddings", original_max_pos)) original_seq_len = 0 if yarn_disabled else original_max_pos return precompute_freqs_cis( diff --git a/src/xorl/ops/fused_silu_and_mul.py b/src/xorl/ops/fused_silu_and_mul.py index ca2055f0..913f7656 100644 --- a/src/xorl/ops/fused_silu_and_mul.py +++ b/src/xorl/ops/fused_silu_and_mul.py @@ -9,13 +9,6 @@ import triton.language as tl -def _native_silu_and_mul(input_tensor: torch.Tensor) -> torch.Tensor: - """Run SwiGLU with the eager PyTorch operation ordering used by serving.""" - assert input_tensor.shape[-1] % 2 == 0, "Last dimension must be even" - split = input_tensor.shape[-1] // 2 - return torch.nn.functional.silu(input_tensor[..., :split]) * input_tensor[..., split:] - - @triton.jit def _silu_and_mul_kernel( input_ptr, diff --git a/src/xorl/ops/group_gemm/kernel/__init__.py b/src/xorl/ops/group_gemm/kernel/__init__.py index c4fe4487..49a241d2 100644 --- a/src/xorl/ops/group_gemm/kernel/__init__.py +++ b/src/xorl/ops/group_gemm/kernel/__init__.py @@ -1,15 +1,8 @@ +import math + # Group GEMM kernels from .group_gemm import group_gemm_same_mn, group_gemm_same_nk -# LoRA utilities -from .lora_utils import ( - compute_lora_scaling, - get_lora_delta_weight_stacked, - init_lora_weights_stacked, - merge_lora_weights_stacked, - unmerge_lora_weights_stacked, -) - # MoE operations from .moe import ( expert_histogram, @@ -21,6 +14,13 @@ from .quack import quack_group_gemm_same_mn, quack_group_gemm_same_nk +def compute_lora_scaling(lora_alpha: int, r: int, use_rslora: bool = False) -> float: + """Compute the standard or rank-stabilized LoRA scale.""" + if use_rslora: + return lora_alpha / math.sqrt(r) + return lora_alpha / r + + __all__ = [ # Group GEMM "group_gemm_same_mn", @@ -34,9 +34,5 @@ "moe_index_compute", "moe_scatter", # LoRA utilities - "init_lora_weights_stacked", "compute_lora_scaling", - "merge_lora_weights_stacked", - "unmerge_lora_weights_stacked", - "get_lora_delta_weight_stacked", ] diff --git a/src/xorl/ops/group_gemm/kernel/lora_utils.py b/src/xorl/ops/group_gemm/kernel/lora_utils.py deleted file mode 100644 index 9482ab2e..00000000 --- a/src/xorl/ops/group_gemm/kernel/lora_utils.py +++ /dev/null @@ -1,148 +0,0 @@ -"""LoRA utilities for MoE implementation. - -This module provides utilities for initializing and managing LoRA weights -in the stacked tensor format used by group GEMM kernels. -""" - -import math -from typing import Optional, Tuple - -import torch -import torch.nn as nn - - -def init_lora_weights_stacked( - num_experts: int, - r: int, - in_features: int, - out_features: int, - init_method: str = "kaiming", - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Initialize stacked LoRA weights for all experts. - - Creates lora_A and lora_B tensors with appropriate initialization: - - lora_A: Kaiming uniform or Gaussian initialization - - lora_B: Zero initialization (ensures delta_W = 0 at start) - - Args: - num_experts: Number of experts - r: LoRA rank - in_features: Input feature dimension - out_features: Output feature dimension - init_method: Initialization method ("kaiming" or "gaussian") - dtype: Data type for the tensors - device: Device for the tensors - - Returns: - Tuple of (lora_A, lora_B) tensors: - - lora_A: Shape [num_experts, in_features, r] - - lora_B: Shape [num_experts, r, out_features] - """ - # lora_A: projects input to low-rank space - # Shape: [num_experts, in_features, r] - lora_A = torch.empty(num_experts, in_features, r, dtype=dtype, device=device) - - # lora_B: projects from low-rank space to output - # Shape: [num_experts, r, out_features] - lora_B = torch.zeros(num_experts, r, out_features, dtype=dtype, device=device) - - # Initialize lora_A - if init_method == "kaiming": - for i in range(num_experts): - # Initialize each expert's lora_A with kaiming uniform - nn.init.kaiming_uniform_(lora_A[i], a=math.sqrt(5)) - elif init_method == "gaussian": - nn.init.normal_(lora_A, std=1.0 / r) - else: - raise ValueError(f"Unknown init_method: {init_method}") - - # lora_B is already zeros - - return lora_A, lora_B - - -def compute_lora_scaling(lora_alpha: int, r: int, use_rslora: bool = False) -> float: - """Compute the LoRA scaling factor. - - Args: - lora_alpha: LoRA alpha parameter - r: LoRA rank - use_rslora: Whether to use rank-stabilized LoRA scaling - - Returns: - Scaling factor - """ - if use_rslora: - return lora_alpha / math.sqrt(r) - else: - return lora_alpha / r - - -def merge_lora_weights_stacked( - base_weight: torch.Tensor, - lora_A: torch.Tensor, - lora_B: torch.Tensor, - scaling: float, -) -> torch.Tensor: - """Merge LoRA weights into base weights. - - Computes: W' = W + A @ B * scaling - - Args: - base_weight: Base weight tensor [num_experts, in_features, out_features] - lora_A: LoRA A tensor [num_experts, in_features, r] - lora_B: LoRA B tensor [num_experts, r, out_features] - scaling: LoRA scaling factor - - Returns: - Merged weight tensor [num_experts, in_features, out_features] - """ - # A @ B: [num_experts, in_features, r] @ [num_experts, r, out_features] - # = [num_experts, in_features, out_features] - delta_weight = torch.bmm(lora_A, lora_B) * scaling - return base_weight + delta_weight - - -def unmerge_lora_weights_stacked( - merged_weight: torch.Tensor, - lora_A: torch.Tensor, - lora_B: torch.Tensor, - scaling: float, -) -> torch.Tensor: - """Unmerge LoRA weights from merged weights. - - Computes: W = W' - A @ B * scaling - - Args: - merged_weight: Merged weight tensor [num_experts, in_features, out_features] - lora_A: LoRA A tensor [num_experts, in_features, r] - lora_B: LoRA B tensor [num_experts, r, out_features] - scaling: LoRA scaling factor - - Returns: - Base weight tensor [num_experts, in_features, out_features] - """ - delta_weight = torch.bmm(lora_A, lora_B) * scaling - return merged_weight - delta_weight - - -def get_lora_delta_weight_stacked( - lora_A: torch.Tensor, - lora_B: torch.Tensor, - scaling: float, -) -> torch.Tensor: - """Compute the LoRA weight delta. - - Computes: delta_W = A @ B * scaling - - Args: - lora_A: LoRA A tensor [num_experts, in_features, r] - lora_B: LoRA B tensor [num_experts, r, out_features] - scaling: LoRA scaling factor - - Returns: - Delta weight tensor [num_experts, in_features, out_features] - """ - return torch.bmm(lora_A, lora_B) * scaling diff --git a/src/xorl/ops/group_gemm/kernel/moe.py b/src/xorl/ops/group_gemm/kernel/moe.py index 053c86cf..9c49d85b 100644 --- a/src/xorl/ops/group_gemm/kernel/moe.py +++ b/src/xorl/ops/group_gemm/kernel/moe.py @@ -1,7 +1,5 @@ """MoE operations: scatter, gather, histogram, and index computation.""" -import os - import torch import triton import triton.language as tl @@ -12,19 +10,8 @@ ) -def _deterministic_scatter_enabled() -> bool: - """Deterministic scatter is the default: build the token->slot permutation with - a stable sort instead of relaxed atomics, so `moe_index_compute` returns the - same permutation on every run (required for bit-reproducible training). - Escape hatch: XORL_MOE_DETERMINISTIC_SCATTER=0 restores the atomics kernel.""" - value = os.environ.get("XORL_MOE_DETERMINISTIC_SCATTER") - if value is None: - return True - return value.strip().lower() not in {"0", "false", "no", "off", ""} - - def _moe_index_compute_deterministic(experts_for_tokens: torch.Tensor) -> torch.Tensor: - """Run-invariant replacement for `_moe_index_compute_kernel`. + """Build a run-invariant token-to-expert-slot permutation. A stable argsort of the flattened expert ids yields slots where expert regions appear in expert-id order (the same `[cumsum[e-1], cumsum[e])` @@ -346,50 +333,6 @@ def moe_scatter(x: torch.Tensor, index: torch.Tensor, out_dtype=None): return out -@triton.jit -def _moe_index_compute_kernel( - indices_ptr, - experts_for_tokens_ptr, - temp_histogram_cumsum_ptr, - num_elts, - NUM_EXPERTS: tl.constexpr, - BLOCK_SIZE: tl.constexpr, # Unlikely to be aligned, so we don't test for alignment. -): - _OOB_EXPERT_ID: tl.constexpr = 1023 - tl.static_assert(_OOB_EXPERT_ID > NUM_EXPERTS, "Too many experts for me.") - - start_pos = tl.program_id(0) - processing_range = start_pos * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - expert_ids = tl.load( - experts_for_tokens_ptr + processing_range, - processing_range < num_elts, - _OOB_EXPERT_ID, - ) - assert expert_ids < NUM_EXPERTS or expert_ids == _OOB_EXPERT_ID - - indices = tl.zeros([BLOCK_SIZE], dtype=tl.int32) - for expert_id in tl.static_range(NUM_EXPERTS): - mask = expert_ids == expert_id - one_if_expert_id_matches = mask.to(tl.int32) - - # Tokens allocated to this expert. - slots_to_reserve = tl.sum(one_if_expert_id_matches) - slot_ids = ( - # Reserve last `slots_to_reserve` slots for us. - tl.atomic_add(temp_histogram_cumsum_ptr + expert_id, -slots_to_reserve, sem="relaxed") - # `atomic_add` returns old value, so we need to do subtraction again. - - slots_to_reserve - # Local offset for each token in `expert_ids`. - + tl.cumsum(one_if_expert_id_matches) - # Result of `cumsum` is "1-based". - - 1 - ) - assigned_slot_or_zero = tl.where(mask, slot_ids, 0) - indices += assigned_slot_or_zero.to(tl.int32) - - tl.store(indices_ptr + processing_range, indices, processing_range < num_elts) - - def moe_index_compute(experts_for_tokens: torch.Tensor, expert_histogram_cumsum: torch.Tensor) -> torch.Tensor: """Calculate row number into activation passed to MoE fc1 for each token. @@ -413,21 +356,4 @@ def moe_index_compute(experts_for_tokens: torch.Tensor, expert_histogram_cumsum: f"experts_for_tokens.device = {experts_for_tokens.device}, expert_histogram_cumsum.device = {expert_histogram_cumsum.device}" ) - if _deterministic_scatter_enabled(): - return _moe_index_compute_deterministic(experts_for_tokens) - - BLOCK_SIZE = 128 # Faster than 1024, not sure why. May be better occupancy? - - histogram_cumsum_copy = expert_histogram_cumsum.clone().detach() # Temporary workspace. - indices = torch.empty_like(experts_for_tokens, dtype=int) - - _moe_index_compute_kernel[(triton.cdiv(experts_for_tokens.numel(), BLOCK_SIZE),)]( - indices_ptr=indices, - experts_for_tokens_ptr=experts_for_tokens, - temp_histogram_cumsum_ptr=histogram_cumsum_copy, - num_elts=experts_for_tokens.numel(), - NUM_EXPERTS=histogram_cumsum_copy.numel(), - BLOCK_SIZE=BLOCK_SIZE, - ) - - return indices + return _moe_index_compute_deterministic(experts_for_tokens) diff --git a/src/xorl/ops/linear_attention/ops/utils/solve_tril_decode.py b/src/xorl/ops/linear_attention/ops/utils/solve_tril_decode.py index 7bb0468f..f1d9ec9d 100644 --- a/src/xorl/ops/linear_attention/ops/utils/solve_tril_decode.py +++ b/src/xorl/ops/linear_attention/ops/utils/solve_tril_decode.py @@ -25,9 +25,8 @@ # The forward-substitution reduction tree must match the pinned solve_tril # kernels (num_warps=2, see SOLVE_TRIL_NUM_WARPS). The diag kernel spells that -# tree out explicitly (see _sum_rows_16_fla_tree), so its launch config is -# free; bit-invariance across the config sweep is measured in -# tests/ops/test_gdn_decode_prep.py rather than assumed. +# tree out explicitly (see _sum_rows_16_fla_tree), while the launch geometry +# remains pinned below. DIAG_HEAD_GROUP = 16 DIAG_NUM_WARPS = 2 DIAG_NUM_STAGES = 1 diff --git a/src/xorl/ops/loss/__init__.py b/src/xorl/ops/loss/__init__.py index 83ce5bbf..76dd9d9d 100644 --- a/src/xorl/ops/loss/__init__.py +++ b/src/xorl/ops/loss/__init__.py @@ -16,7 +16,7 @@ from xorl.ops.loss.loss_output import LossOutput from xorl.ops.loss.opd_loss import OPDLossMetrics, opd_loss_function, opd_vocab_parallel_loss_function from xorl.ops.loss.policy_loss import policy_loss_function -from xorl.ops.loss.reducers import Reducer, SequencePartial, TokenPartial +from xorl.ops.loss.reducers import Reducer, TokenPartial from xorl.ops.loss.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy @@ -57,7 +57,6 @@ def register_loss_function(name: str, fn: Callable) -> None: "OPDLossMetrics", "LOSS_REGISTRY", "Reducer", - "SequencePartial", "TokenPartial", "get_loss_function", "register_loss_function", diff --git a/src/xorl/ops/loss/reducers.py b/src/xorl/ops/loss/reducers.py index 48b414c8..254debbd 100644 --- a/src/xorl/ops/loss/reducers.py +++ b/src/xorl/ops/loss/reducers.py @@ -34,38 +34,7 @@ def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: return (values * mask).sum() / self.scale.clamp(min=1.0) -@dataclass(frozen=True) -class SequencePartial: - """Sum of per-segment token-means, divided by a caller-supplied ``scale``. - - Segment boundaries are flat across ``(values * mask).reshape(-1)``: - - - ``cu_seqlens_local: (N+1,)`` — shard-local segment extents. Under CP each - rank's slice sums to its segment's local contribution. - - ``seq_lengths_global: (N,)`` — pre-CP-shard token count per segment, used - as the per-segment denominator so partial shares from each CP rank sum - to the correct per-segment mean. - """ - - scale: torch.Tensor - cu_seqlens_local: torch.Tensor - seq_lengths_global: torch.Tensor - - def __call__(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: - flat = (values * mask).reshape(-1) - seg_lengths_local = self.cu_seqlens_local.diff() - n_segments = seg_lengths_local.numel() - seg_ids = torch.repeat_interleave( - torch.arange(n_segments, device=flat.device), - seg_lengths_local, - ) - seg_sums = torch.zeros(n_segments, dtype=flat.dtype, device=flat.device).index_add(0, seg_ids, flat) - seg_means = seg_sums / self.seq_lengths_global.clamp(min=1.0) - return seg_means.sum() / self.scale.clamp(min=1.0) - - __all__ = [ "Reducer", - "SequencePartial", "TokenPartial", ] diff --git a/src/xorl/ops/moe/triton.py b/src/xorl/ops/moe/triton.py index 03802fda..aa58656b 100644 --- a/src/xorl/ops/moe/triton.py +++ b/src/xorl/ops/moe/triton.py @@ -1,5 +1,3 @@ -import os - import torch import torch.nn.functional as F @@ -43,13 +41,8 @@ def set_routing_weights_before_down(enabled: bool) -> None: def routing_weights_before_down() -> bool: - """Whether expert_scores fold into the down-GEMM input instead of its output. - - The ``XORL_MOE_ROUTING_WEIGHTS_BEFORE_DOWN=1`` env var force-enables the - before-down position regardless of config; it is read lazily (per forward) - so it keeps working when set after import. - """ - return _ROUTING_WEIGHTS_BEFORE_DOWN_CONFIG or os.environ.get("XORL_MOE_ROUTING_WEIGHTS_BEFORE_DOWN", "0") == "1" + """Whether expert_scores fold into the down-GEMM input instead of its output.""" + return _ROUTING_WEIGHTS_BEFORE_DOWN_CONFIG def resolve_routing_weights_before_down(setting: bool | str, *, train_router: bool, ep_dispatch: str) -> bool: @@ -64,9 +57,9 @@ def resolve_routing_weights_before_down(setting: bool | str, *, train_router: bo if isinstance(setting, bool): return setting normalized = str(setting).strip().lower() - if normalized in ("true", "1"): + if normalized == "true": return True - if normalized in ("false", "0"): + if normalized == "false": return False if normalized != "auto": raise ValueError(f"Invalid moe_routing_weights_before_down={setting!r}; expected 'auto', true, or false.") diff --git a/src/xorl/ops/quantize/block_fp8_gkn_quantize.py b/src/xorl/ops/quantize/block_fp8_gkn_quantize.py index 54989f78..25d7e2f6 100644 --- a/src/xorl/ops/quantize/block_fp8_gkn_quantize.py +++ b/src/xorl/ops/quantize/block_fp8_gkn_quantize.py @@ -314,12 +314,3 @@ def grid(meta): _block_fp8_dequantize_gkn_rowwise_kernel[grid](x, s, y, M, N, BLOCK_SIZE=block_size) return y - - -# --------------------------------------------------------------------------- -# Backward-compat aliases -# --------------------------------------------------------------------------- - -block_fp8_weight_quant = block_fp8_quantize_gkn -block_fp8_weight_dequant = block_fp8_dequantize_gkn -block_fp8_weight_quant_gkn = block_fp8_quantize_gkn diff --git a/src/xorl/ops/quantize/block_fp8_quantize.py b/src/xorl/ops/quantize/block_fp8_quantize.py index 99340db3..b97c8f29 100644 --- a/src/xorl/ops/quantize/block_fp8_quantize.py +++ b/src/xorl/ops/quantize/block_fp8_quantize.py @@ -83,10 +83,6 @@ def grid(meta): return y, s -# Backward-compat alias -block_fp8_quant = block_fp8_quantize - - # --------------------------------------------------------------------------- # 1D dequantization kernel + wrapper # --------------------------------------------------------------------------- @@ -130,10 +126,6 @@ def grid(meta): return x -# Backward-compat alias -block_fp8_dequant = block_fp8_dequantize - - # --------------------------------------------------------------------------- # Autotuned FP8 GEMM # --------------------------------------------------------------------------- diff --git a/src/xorl/ops/quantize/nvfp4_fake_quant.py b/src/xorl/ops/quantize/nvfp4_fake_quant.py index d5e327ee..85f0f2f2 100644 --- a/src/xorl/ops/quantize/nvfp4_fake_quant.py +++ b/src/xorl/ops/quantize/nvfp4_fake_quant.py @@ -162,13 +162,6 @@ def fake_quantize_nvfp4(w: Tensor, block_size: int = 16) -> Tensor: return w + (w_dq - w).detach() -def fake_quantize(w: Tensor, quant_format: str = "nvfp4", block_size: int = 16) -> Tensor: - """Format-dispatched fake quantization (STE). Currently supports ``nvfp4``.""" - if quant_format == "nvfp4": - return fake_quantize_nvfp4(w, block_size) - raise ValueError(f"Unsupported quant_format={quant_format!r}; supported: {sorted(_SUPPORTED_FORMATS)}") - - def fake_quantize_activation_nvfp4(x: Tensor, block_size: int = 16) -> Tensor: """Fake-quantize an **activation** tensor to NVFP4 with an STE backward. diff --git a/src/xorl/optim/gram_newton_schulz.py b/src/xorl/optim/gram_newton_schulz.py index 64ff55df..b7ae7ce4 100644 --- a/src/xorl/optim/gram_newton_schulz.py +++ b/src/xorl/optim/gram_newton_schulz.py @@ -72,11 +72,6 @@ def _import_quack_gemm_interface(): raise ImportError("Muon Gram Newton-Schulz requires the upstream `quack-kernels` package") from exc -def _muon_quack_tuned_enabled() -> bool: - value = os.getenv("XORL_MUON_QUACK_TUNED", "0").strip().lower() - return value in {"1", "true", "yes", "on"} - - @lru_cache(maxsize=1) def _make_quack_backend(): _ensure_cutlass_arch_for_current_device() @@ -84,13 +79,11 @@ def _make_quack_backend(): gemm = gemm_interface.gemm gemm_add = gemm_interface.gemm_add gemm_symmetric = gemm_interface.gemm_symmetric - tuned = _muon_quack_tuned_enabled() - return SimpleNamespace( sym_mm=lambda A, B: gemm_symmetric(A, B), sym_baddbmm=lambda A, B, C, alpha=1.0, beta=1.0: gemm_symmetric(A, B, C=C, alpha=alpha, beta=beta), - mm=lambda A, B: gemm(A, B, tuned=tuned), - mm_add=lambda A, B, C, beta=1.0: gemm_add(A, B, C=C, beta=beta, tuned=tuned), + mm=lambda A, B: gemm(A, B, tuned=False), + mm_add=lambda A, B, C, beta=1.0: gemm_add(A, B, C=C, beta=beta, tuned=False), ) diff --git a/src/xorl/qlora/modules/moe_experts.py b/src/xorl/qlora/modules/moe_experts.py index 12c91e87..53a4209c 100644 --- a/src/xorl/qlora/modules/moe_experts.py +++ b/src/xorl/qlora/modules/moe_experts.py @@ -42,7 +42,7 @@ validate_gated_silu_expert_adapter_semantics, ) from xorl.lora.modules.base import LoraModule -from xorl.ops.group_gemm.kernel.lora_utils import compute_lora_scaling +from xorl.ops.group_gemm.kernel import compute_lora_scaling from xorl.ops.quantize import ( block_fp8_dequantize_gkn, block_fp8_quantize_gkn, diff --git a/src/xorl/rl/__init__.py b/src/xorl/rl/__init__.py deleted file mode 100644 index 8fffc044..00000000 --- a/src/xorl/rl/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -"""Reusable RL objective primitives. - -These helpers intentionally stop at train-time tensor math. Rollout collection, -reward execution, and request construction live outside the core training -engine. -""" - -from xorl.rl.kl import compute_kl_estimate, compute_sequence_kl -from xorl.rl.normalization import reduce_token_or_sample_mean -from xorl.rl.objectives import compute_gspo_kl, compute_opsm_mask, compute_policy_clip_loss - - -__all__ = [ - "compute_gspo_kl", - "compute_kl_estimate", - "compute_opsm_mask", - "compute_policy_clip_loss", - "compute_sequence_kl", - "reduce_token_or_sample_mean", -] diff --git a/src/xorl/rl/kl.py b/src/xorl/rl/kl.py deleted file mode 100644 index 62dc69ad..00000000 --- a/src/xorl/rl/kl.py +++ /dev/null @@ -1,62 +0,0 @@ -from __future__ import annotations - -from typing import Literal - -import torch - - -KLEstimator = Literal["k1", "k2", "k3", "low_var_kl"] - - -def compute_kl_estimate( - policy_logprobs: torch.Tensor, - base_logprobs: torch.Tensor, - kind: KLEstimator, - importance_ratio: torch.Tensor | None = None, -) -> torch.Tensor: - """Compute Slime-compatible sampled-token KL estimates. - - ``policy_logprobs`` are the current policy log-probabilities and - ``base_logprobs`` are reference/base log-probabilities. ``importance_ratio`` - is optional ``pi_current / pi_old`` weighting, matching Slime's unbiased KL - mode. - """ - log_ratio = policy_logprobs.float() - base_logprobs.float() - - if kind == "k1": - kl = log_ratio - elif kind == "k2": - kl = 0.5 * log_ratio.square() - elif kind in ("k3", "low_var_kl"): - neg_log_ratio = -log_ratio - kl = torch.exp(neg_log_ratio) - 1.0 - neg_log_ratio - else: - raise ValueError(f"Unknown KL estimator: {kind}") - - if importance_ratio is not None: - kl = importance_ratio.float() * kl - - if kind == "low_var_kl": - kl = torch.clamp(kl, min=-10.0, max=10.0) - - return kl - - -def compute_sequence_kl( - current_logprobs: torch.Tensor, - old_logprobs: torch.Tensor, - masks: torch.Tensor, - *, - expand: bool = False, -) -> torch.Tensor: - """Compute per-sequence PPO KL ``mean(old - current)`` over valid tokens. - - When ``expand=True``, the per-sequence value is expanded back to token shape, - which is the policy-path GSPO convention used by Slime before PPO clipping. - """ - mask_f = masks.float() - seq_lengths = mask_f.sum(dim=-1).clamp(min=1.0) - seq_kl = ((old_logprobs.float() - current_logprobs.float()) * mask_f).sum(dim=-1) / seq_lengths - if expand: - return seq_kl.unsqueeze(-1).expand_as(current_logprobs) - return seq_kl diff --git a/src/xorl/rl/normalization.py b/src/xorl/rl/normalization.py deleted file mode 100644 index 19363799..00000000 --- a/src/xorl/rl/normalization.py +++ /dev/null @@ -1,42 +0,0 @@ -from __future__ import annotations - -from typing import Literal - -import torch - - -ReductionMode = Literal["token_mean", "sample_mean", "slime_sum_of_sample_mean", "token_sum"] - - -def reduce_token_or_sample_mean( - values: torch.Tensor, - masks: torch.Tensor, - mode: ReductionMode, -) -> torch.Tensor: - """Reduce token-aligned values with an explicit normalization contract. - - Modes: - - ``token_mean``: global token-weighted mean over all valid tokens. - - ``sample_mean``: arithmetic mean of non-empty per-sample means. - - ``slime_sum_of_sample_mean``: sum of per-sample means, matching Slime's - train-time reducer before its later global-batch divisor. - - ``token_sum``: masked sum, useful when normalization is deferred. - """ - mask_f = masks.float() - masked_sum = (values * mask_f).sum() - - if mode == "token_sum": - return masked_sum - if mode == "token_mean": - return masked_sum / mask_f.sum().clamp(min=1.0) - - sample_denoms = mask_f.sum(dim=-1).clamp(min=1.0) - sample_means = (values * mask_f).sum(dim=-1) / sample_denoms - nonempty = mask_f.sum(dim=-1) > 0 - - if mode == "slime_sum_of_sample_mean": - return sample_means.masked_fill(~nonempty, 0.0).sum() - if mode == "sample_mean": - return sample_means.masked_select(nonempty).sum() / nonempty.sum().clamp(min=1) - - raise ValueError(f"Unknown reduction mode: {mode}") diff --git a/src/xorl/rl/objectives.py b/src/xorl/rl/objectives.py deleted file mode 100644 index 4de23df6..00000000 --- a/src/xorl/rl/objectives.py +++ /dev/null @@ -1,67 +0,0 @@ -from __future__ import annotations - -import torch - -from xorl.rl.kl import compute_sequence_kl - - -def compute_policy_clip_loss( - ppo_kl: torch.Tensor, - advantages: torch.Tensor, - eps_clip: float, - eps_clip_high: float, - eps_clip_c: float | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Compute the PPO clipped policy loss used by Slime and XoRL. - - ``ppo_kl`` is ``old_logprobs - current_logprobs``. The returned clip - fraction tensor marks the first PPO clip and intentionally does not include - dual-clip events, matching Slime's reporting behavior. - """ - ratio = torch.exp(-ppo_kl) - pg_losses1 = -ratio * advantages - pg_losses2 = -torch.clamp(ratio, 1.0 - eps_clip, 1.0 + eps_clip_high) * advantages - clip_pg_losses1 = torch.maximum(pg_losses1, pg_losses2) - clipfrac = torch.gt(pg_losses2, pg_losses1).float() - - if eps_clip_c is not None: - if eps_clip_c <= 1.0: - raise ValueError(f"eps_clip_c must be > 1.0, got {eps_clip_c}") - pg_losses3 = -eps_clip_c * advantages - clip_pg_losses2 = torch.minimum(pg_losses3, clip_pg_losses1) - pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1) - else: - pg_losses = clip_pg_losses1 - - return pg_losses, clipfrac, ratio - - -def compute_gspo_kl( - current_logprobs: torch.Tensor, - old_logprobs: torch.Tensor, - masks: torch.Tensor, -) -> torch.Tensor: - """Compute GSPO sequence-level KL and expand it to token shape.""" - return compute_sequence_kl(current_logprobs, old_logprobs, masks, expand=True) - - -def compute_opsm_mask( - current_logprobs: torch.Tensor, - old_logprobs: torch.Tensor, - advantages: torch.Tensor, - masks: torch.Tensor, - delta: float, -) -> tuple[torch.Tensor, torch.Tensor]: - """Compute Slime-style Off-Policy Sequence Masking. - - Tokens with negative advantages are masked out when their sequence-level - ``mean(old - current)`` KL exceeds ``delta``. ``opsm_clipfrac`` follows - Slime's reported value: sum of each sequence's masked-token fraction. - """ - mask_f = masks.float() - seq_kl = compute_sequence_kl(current_logprobs, old_logprobs, mask_f, expand=False) - masked = (advantages < 0) & (seq_kl.unsqueeze(-1) > delta) & masks.bool() - opsm_mask = torch.ones_like(mask_f).masked_fill(masked, 0.0) - per_sequence_fraction = (masked.float() * mask_f).sum(dim=-1) / mask_f.sum(dim=-1).clamp(min=1.0) - opsm_clipfrac = per_sequence_fraction.sum() - return opsm_mask, opsm_clipfrac diff --git a/src/xorl/server/launcher.py b/src/xorl/server/launcher.py index 3c76005e..0140e641 100644 --- a/src/xorl/server/launcher.py +++ b/src/xorl/server/launcher.py @@ -40,7 +40,7 @@ import uvicorn import yaml -from xorl.fp8_training.config_compat import extract_nemo_fp8_cfg, validate_external_fp8_runtime_config +from xorl.fp8_training.config_compat import validate_external_fp8_runtime_config from xorl.server.api_server.server import APIServer from xorl.server.orchestrator.orchestrator import Orchestrator from xorl.server.removed_config import reject_removed_configuration_fields @@ -97,35 +97,6 @@ def configure_uvicorn_logging(): # ============================================================================ -def find_free_port(start_port: int = 50000, max_attempts: int = 10000) -> int: - """ - Find a free port by randomly picking from a range. - - Args: - start_port: Start of port range to search - max_attempts: Maximum number of ports to try - - Returns: - Free port number - - Raises: - RuntimeError: If no free port found - """ - - end_port = min(start_port + max_attempts, 60000) - ports = list(range(start_port, end_port)) - random.shuffle(ports) - for port in ports: - with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock: - try: - sock.bind(("", port)) - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - return port - except OSError: - continue - raise RuntimeError(f"Could not find free port in range {start_port}-{end_port}") - - def find_free_ports(count: int, start_port: int = 50000) -> List[int]: """ Find multiple free ports. @@ -435,7 +406,6 @@ def load_server_arguments(config_path: str, overrides: Optional[Dict[str, any]] reject_removed_configuration_fields(config, context=f"server config {config_path!r}") validate_external_fp8_runtime_config(config, context=config_path) - nemo_fp8_cfg = extract_nemo_fp8_cfg(config) valid_fields = {f.name for f in fields(ServerArguments)} @@ -453,9 +423,6 @@ def load_server_arguments(config_path: str, overrides: Optional[Dict[str, any]] for section in ("model", "train"): for k, v in config.get(section, {}).items(): flat_config[k] = v - if nemo_fp8_cfg is not None and "fp8_cfg" not in flat_config: - flat_config["fp8_cfg"] = nemo_fp8_cfg - # lora.* keys also map 1:1 except exclude_modules → qlora_exclude_modules for k, v in config.get("lora", {}).items(): if k == "exclude_modules": @@ -499,8 +466,6 @@ def load_server_arguments(config_path: str, overrides: Optional[Dict[str, any]] else: # Flat config (ServerArguments style) flat_config = dict(config) - if nemo_fp8_cfg is not None and "fp8_cfg" not in flat_config: - flat_config["fp8_cfg"] = nemo_fp8_cfg filtered_config = {k: v for k, v in flat_config.items() if k in valid_fields} # Handle None values for Optional fields diff --git a/src/xorl/server/orchestrator/request_processor.py b/src/xorl/server/orchestrator/request_processor.py index 620f615e..c0f00928 100644 --- a/src/xorl/server/orchestrator/request_processor.py +++ b/src/xorl/server/orchestrator/request_processor.py @@ -174,8 +174,6 @@ def __init__( r3_payload_dir: Optional[str] = None, r3_payload_keep: bool = False, r3_payload_namespace_prefix: Optional[str] = None, - routing_payload_dir: Optional[str] = None, - keep_routing_payloads: Optional[bool] = None, routing_payload_store: Optional[MooncakeSidePayloadStore] = None, ): """ @@ -198,8 +196,6 @@ def __init__( r3_payload_dir: Shared directory used only by the explicit filesystem fallback. r3_payload_keep: If True, do not delete side payloads after the backend call. r3_payload_namespace_prefix: Optional Mooncake namespace prefix for R3 payload keys. - routing_payload_dir: Backward-compatible alias for filesystem transport. - keep_routing_payloads: Backward-compatible alias for r3_payload_keep. routing_payload_store: Optional injected Mooncake side-payload store for tests. """ self.backend = backend @@ -210,13 +206,6 @@ def __init__( self.packing_strategy = packing_strategy self.on_oversized = on_oversized self.dp_size = max(1, int(dp_size)) - if routing_payload_dir is not None: - if r3_payload_transport != "inline": - raise ValueError("routing_payload_dir alias cannot be combined with r3_payload_transport") - r3_payload_transport = "filesystem" - r3_payload_dir = routing_payload_dir - if keep_routing_payloads is not None: - r3_payload_keep = bool(keep_routing_payloads) if r3_payload_transport not in {"inline", "mooncake", "filesystem"}: raise ValueError(f"Unsupported r3_payload_transport {r3_payload_transport!r}") if r3_payload_transport == "inline" and r3_payload_keep: @@ -376,9 +365,9 @@ def _cleanup_routing_payloads(self, cleanup: Optional[Union[Path, R3PayloadClean log_fn = logger.info if _r3_verbose_logging_enabled() else logger.debug log_fn("Cleaned external R3 Mooncake routing payload keys") return - self._cleanup_routing_payload_dir(cleanup) + self._cleanup_r3_payload_dir(cleanup) - def _cleanup_routing_payload_dir(self, root: Path) -> None: + def _cleanup_r3_payload_dir(self, root: Path) -> None: try: shutil.rmtree(root) except FileNotFoundError: diff --git a/src/xorl/server/protocol/__init__.py b/src/xorl/server/protocol/__init__.py index 47b7f83d..f68d6d79 100644 --- a/src/xorl/server/protocol/__init__.py +++ b/src/xorl/server/protocol/__init__.py @@ -16,22 +16,6 @@ OrchestratorRequest, OutputType, RequestType, - create_error_output, - create_forward_backward_output, - create_health_check_output, - create_load_adapter_state_output, - create_load_state_output, - create_optim_step_output, - create_save_adapter_state_output, - create_save_lora_only_output, - create_save_state_output, - create_sleep_output, - create_sync_weights_output, - create_wake_up_output, - get_operation_from_request, - is_streaming_output, - validate_output, - validate_request, ) from xorl.server.protocol.operations import ( # noqa: F401 AbortData, @@ -59,8 +43,6 @@ RunnerDispatchCommand, RunnerReady, RunnerResponse, - create_ack_for_request, - create_response_for_request, deserialize_message, serialize_message, ) diff --git a/src/xorl/server/protocol/api_orchestrator.py b/src/xorl/server/protocol/api_orchestrator.py index 9da2fd53..8753f687 100644 --- a/src/xorl/server/protocol/api_orchestrator.py +++ b/src/xorl/server/protocol/api_orchestrator.py @@ -199,316 +199,3 @@ def __repr__(self) -> str: status = "finished" if self.finished else "streaming" error_str = f", error='{self.error}'" if self.error else "" return f"OrchestratorOutputs(id={self.request_id[:8]}..., type={self.output_type.value}, status={status}{error_str})" - - -# ============================================================================ -# Response Builder Functions (Engine → API Server) -# ============================================================================ - - -def _build_output( - request_id: str, - output_type: OutputType, - error: Optional[str] = None, - **fields, -) -> OrchestratorOutputs: - """Build an OrchestratorOutputs with optional fields filtered.""" - outputs_data = {} - for key, value in fields.items(): - if value is not None: - outputs_data[key] = value - return OrchestratorOutputs( - request_id=request_id, - output_type=output_type, - outputs=[outputs_data], - finished=True, - error=error, - ) - - -def create_forward_backward_output( - request_id: str, - loss: float, - valid_tokens: Optional[int] = None, - grads_norm: Optional[float] = None, - additional_metrics: Optional[Dict[str, Any]] = None, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create a forward-backward output response.""" - outputs_data = {"loss": loss} - if valid_tokens is not None: - outputs_data["valid_tokens"] = valid_tokens - if grads_norm is not None: - outputs_data["grads_norm"] = grads_norm - if additional_metrics: - outputs_data.update(additional_metrics) - return OrchestratorOutputs( - request_id=request_id, - output_type=OutputType.FORWARD_BACKWARD, - outputs=[outputs_data], - finished=True, - error=error, - ) - - -def create_optim_step_output( - request_id: str, - step: Optional[int] = None, - learning_rate: Optional[float] = None, - grad_norm: Optional[float] = None, - additional_metrics: Optional[Dict[str, Any]] = None, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create an optimizer step output response.""" - outputs_data = {} - if step is not None: - outputs_data["step"] = step - if learning_rate is not None: - outputs_data["lr"] = learning_rate - outputs_data["learning_rate"] = learning_rate - if grad_norm is not None: - outputs_data["grad_norm"] = grad_norm - if additional_metrics: - outputs_data.update(additional_metrics) - return OrchestratorOutputs( - request_id=request_id, - output_type=OutputType.OPTIM_STEP, - outputs=[outputs_data], - finished=True, - error=error, - ) - - -def create_save_state_output( - request_id: str, - checkpoint_path: str, - success: bool = True, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create a save checkpoint output response.""" - return _build_output( - request_id, - OutputType.SAVE_STATE, - error=error, - success=success, - checkpoint_path=checkpoint_path, - ) - - -def create_save_lora_only_output( - request_id: str, - lora_path: str, - success: bool = True, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create a save LoRA-only output response.""" - return _build_output( - request_id, - OutputType.SAVE_LORA_ONLY, - error=error, - success=success, - lora_path=lora_path, - ) - - -def create_load_state_output( - request_id: str, - checkpoint_path: str, - success: bool = True, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create a load checkpoint output response.""" - return _build_output( - request_id, - OutputType.LOAD_STATE, - error=error, - success=success, - checkpoint_path=checkpoint_path, - ) - - -def create_save_adapter_state_output( - request_id: str, - model_id: str, - path: str, - step: int, - success: bool = True, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create a save adapter state output response.""" - return _build_output( - request_id, - OutputType.SAVE_ADAPTER_STATE, - error=error, - success=success, - model_id=model_id, - path=path, - step=step, - ) - - -def create_load_adapter_state_output( - request_id: str, - model_id: str, - path: str, - step: int, - success: bool = True, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create a load adapter state output response.""" - return _build_output( - request_id, - OutputType.LOAD_ADAPTER_STATE, - error=error, - success=success, - model_id=model_id, - path=path, - step=step, - ) - - -def create_health_check_output( - request_id: str, - status: str = "healthy", - active_requests: int = 0, - total_requests: int = 0, - additional_info: Optional[Dict[str, Any]] = None, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create a health check output response.""" - outputs_data = { - "status": status, - "active_requests": active_requests, - "total_requests": total_requests, - } - if additional_info: - outputs_data.update(additional_info) - return OrchestratorOutputs( - request_id=request_id, - output_type=OutputType.HEALTH_CHECK, - outputs=[outputs_data], - finished=True, - error=error, - ) - - -def create_sleep_output( - request_id: str, - status: str = "sleeping", - offload_time: Optional[float] = None, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create a sleep output response.""" - return _build_output( - request_id, - OutputType.SLEEP, - error=error, - status=status, - offload_time=offload_time, - ) - - -def create_wake_up_output( - request_id: str, - status: str = "awake", - load_time: Optional[float] = None, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create a wake_up output response.""" - return _build_output( - request_id, - OutputType.WAKE_UP, - error=error, - status=status, - load_time=load_time, - ) - - -def create_sync_weights_output( - request_id: str, - success: bool, - message: str, - transfer_time: float = 0.0, - total_bytes: int = 0, - num_parameters: int = 0, - num_buckets: int = 0, - timing_breakdown: Optional[Dict[str, float]] = None, - p2p_rank_summaries: Optional[List[Dict[str, Any]]] = None, - endpoint_results: Optional[List[Dict[str, Any]]] = None, - error: Optional[str] = None, -) -> OrchestratorOutputs: - """Create a sync inference weights output response.""" - return _build_output( - request_id, - OutputType.SYNC_INFERENCE_WEIGHTS, - error=error, - success=success, - message=message, - transfer_time=transfer_time, - total_bytes=total_bytes, - num_parameters=num_parameters, - num_buckets=num_buckets, - timing_breakdown=timing_breakdown or {}, - p2p_rank_summaries=p2p_rank_summaries or [], - endpoint_results=endpoint_results or [], - ) - - -def create_error_output( - request_id: str, - error_message: str, - operation_type: OutputType = OutputType.ERROR, -) -> OrchestratorOutputs: - """Create an error output response.""" - return OrchestratorOutputs( - request_id=request_id, - output_type=operation_type, - outputs=[], - finished=True, - error=error_message, - ) - - -# ============================================================================ -# Validation and Utilities -# ============================================================================ - - -def validate_request(request: OrchestratorRequest) -> bool: - """Validate that a request has required fields.""" - if not request.request_id: - raise ValueError("Request must have request_id") - - if request.request_type not in RequestType: - raise ValueError(f"Invalid request_type: {request.request_type}") - - if request.request_type == RequestType.ADD: - if not request.operation: - raise ValueError("ADD request must have 'operation'") - - if request.request_type == RequestType.ABORT: - if not getattr(request.payload, "target_request_id", None): - raise ValueError("ABORT request must have 'target_request_id' in payload") - - return True - - -def validate_output(output: OrchestratorOutputs) -> bool: - """Validate that an output has required fields.""" - if not output.request_id: - raise ValueError("Output must have request_id") - - if output.output_type not in OutputType: - raise ValueError(f"Invalid output_type: {output.output_type}") - - return True - - -def get_operation_from_request(request: OrchestratorRequest) -> Optional[str]: - """Extract operation type from request.""" - return request.operation or None - - -def is_streaming_output(output: OrchestratorOutputs) -> bool: - """Check if output is a streaming (incomplete) output.""" - return not output.finished and output.error is None diff --git a/src/xorl/server/protocol/orchestrator_runner.py b/src/xorl/server/protocol/orchestrator_runner.py index febfa42d..cf2d767f 100644 --- a/src/xorl/server/protocol/orchestrator_runner.py +++ b/src/xorl/server/protocol/orchestrator_runner.py @@ -332,21 +332,3 @@ def deserialize_message(data: bytes) -> BaseMessage: if envelope["version"] != _WIRE_VERSION: raise ValueError(f"Unsupported runner protocol version: {envelope['version']!r}") return _message_from_mapping(envelope["message"]) - - -def create_ack_for_request(request: RunnerDispatchCommand) -> RunnerAck: - """Create an acknowledgement for a given request.""" - return RunnerAck(request_id=request.message_id, received_at=time.time()) - - -def create_response_for_request( - request: RunnerDispatchCommand, - success: bool, - result: Optional[Dict[str, Any]] = None, - error: Optional[str] = None, - execution_time: Optional[float] = None, -) -> RunnerResponse: - """Create a response for a given request.""" - return RunnerResponse( - request_id=request.message_id, success=success, result=result or {}, error=error, execution_time=execution_time - ) diff --git a/src/xorl/server/removed_config.py b/src/xorl/server/removed_config.py index a53054b6..d7133954 100644 --- a/src/xorl/server/removed_config.py +++ b/src/xorl/server/removed_config.py @@ -11,6 +11,9 @@ "observe, or shadow mode" ) _ZORL_REMOVAL_MIGRATION = "ZORL was removed; remove this field and migrate training to forward_backward plus optim_step" +_FP8_CFG_MIGRATION = ( + "NeMo fp8_cfg translation was removed; set enable_fp8_training and the native fp8_training_* fields explicitly" +) # One inventory is shared by YAML loading, CLI overrides, and public request @@ -28,6 +31,7 @@ "zorl_seed": _ZORL_REMOVAL_MIGRATION, "zorl": _ZORL_REMOVAL_MIGRATION, "zorl_config": _ZORL_REMOVAL_MIGRATION, + "fp8_cfg": _FP8_CFG_MIGRATION, } diff --git a/src/xorl/server/runner/adapters/adapter_coordinator.py b/src/xorl/server/runner/adapters/adapter_coordinator.py index 80967b3d..1e6f111e 100644 --- a/src/xorl/server/runner/adapters/adapter_coordinator.py +++ b/src/xorl/server/runner/adapters/adapter_coordinator.py @@ -130,17 +130,6 @@ def broadcast_adapter_state(self, model_id: str, default_lr: float) -> None: logger.debug(f"Rank {self.rank}: Broadcast adapter state for model_id={model_id}") - def broadcast_adapter_optimizer_state(self, model_id: str) -> None: - """Deprecated no-op: optimizer state is topology-specific and not broadcastable.""" - if self.world_size <= 1: - return - self._validate_pipeline_parallel_broadcast_safe() - logger.debug( - "Rank %s: refusing rank-0 optimizer broadcast for local-shard adapter %s; use all_ranks checkpoint restore", - self.rank, - model_id, - ) - @staticmethod def _strip_optimizer_config(session_spec: Dict[str, Any]) -> Dict[str, Any]: stripped = deepcopy(session_spec) diff --git a/src/xorl/server/runner/adapters/manager.py b/src/xorl/server/runner/adapters/manager.py index f638b0bc..66fef319 100644 --- a/src/xorl/server/runner/adapters/manager.py +++ b/src/xorl/server/runner/adapters/manager.py @@ -608,16 +608,6 @@ class AdapterState: lr: float = 1e-5 last_access_time: float = field(default_factory=time.time) # For LRU eviction - @property - def lora_params(self) -> Dict[str, nn.Parameter]: - """Deprecated local-only view retained for external compatibility. - - This is intentionally not a logical/full tensor API. Internal code - must use ``local_params`` and ``tensor_layouts`` explicitly. - """ - - return self.local_params - class LoRAAdapterManager: """ @@ -2113,28 +2103,6 @@ def commit_gradient_capture(self, model_id: str) -> tuple[int, int]: scratch.staged_parameter_fqns = () scratch.staged_numerators.clear() - def capture_gradient_numerators( - self, - model_id: str, - *, - denominator: float, - numerator_scale: float = 1.0, - backward_completed: bool, - ) -> tuple[int, int]: - """Stage and immediately commit one capture for direct manager callers.""" - - try: - self.stage_gradient_numerators( - model_id, - denominator=denominator, - numerator_scale=numerator_scale, - backward_completed=backward_completed, - ) - return self.commit_gradient_capture(model_id) - except BaseException: - self.abort_gradient_capture(model_id) - raise - def prepare_forward(self, model_id: str) -> None: """ Load adapter weights into model before forward pass. diff --git a/src/xorl/server/runner/model_runner.py b/src/xorl/server/runner/model_runner.py index 741e4654..e76572e0 100644 --- a/src/xorl/server/runner/model_runner.py +++ b/src/xorl/server/runner/model_runner.py @@ -106,7 +106,6 @@ MoeMetricsTracker, RoutingReplayHandler, batch_slice_rank_and_size, - ep_duplicate_batches_enabled, run_self_test, validate_token_ids, ) @@ -149,16 +148,8 @@ logger = logging.getLogger(__name__) -def _truthy_flag(value: Any) -> bool: - if isinstance(value, str): - return value.strip().lower() in {"1", "true", "yes", "on"} - return bool(value) - - def _skip_empty_cache_after_optim_step(train_config: Dict[str, Any]) -> bool: - return _truthy_flag(os.environ.get("XORL_SKIP_EMPTY_CACHE_AFTER_OPTIM_STEP", False)) or _truthy_flag( - train_config.get("skip_empty_cache_after_optim_step", False) - ) + return train_config.get("skip_empty_cache_after_optim_step") is True def _optimizer_effective_hparams(optimizer: Any) -> Dict[str, Any]: @@ -277,11 +268,6 @@ def _sp_allreduce_kl_metrics( unfinalized. ``valid_tokens`` SUM-reduces alongside them; downstream accumulation divides mean metrics by the final token count. """ - # Backward-compatible argument order for older tests/call sites: - # _sp_allreduce_kl_metrics(metrics, metric_ops, sp_group). - if isinstance(sp_group, dict): - metric_ops, sp_group = sp_group, metric_ops - device = torch.device(get_device_type()) local_n = float(metrics.get("valid_tokens", metrics.get("_n_valid_kl", 0)) or 0) metrics["valid_tokens"] = local_n @@ -1359,6 +1345,7 @@ def _initialize_model(self): lora_alpha=self.lora_config.get("lora_alpha", 16), lora_target_modules=construction_target_modules, lora_target_manifest=self.lora_config.get("lora_target_manifest"), + unfuse_for_lora=self.lora_config.get("unfuse_for_lora", False), moe_hybrid_shared_lora=self.lora_config.get("moe_hybrid_shared_lora", False), enable_qlora=enable_qlora, block_fp8_qlora_training=block_fp8_qlora_training, @@ -1524,6 +1511,22 @@ def _resolve_lora_target_modules(self) -> List[str]: train_mlp=train_mlp, train_unembed=train_unembed, ) + elif model_type in { + "qwen3_5", + "qwen3_5_text", + "qwen3_5_moe", + "qwen3_5_moe_text", + "xorl_qwen3_5", + "xorl_qwen3_5_moe", + }: + target_modules = [] + if train_attn: + # GDN's g_proj is the serving in_proj_z surface. + target_modules.extend(["q_proj", "k_proj", "v_proj", "g_proj", "o_proj"]) + if train_mlp: + target_modules.extend(["gate_proj", "up_proj", "down_proj"]) + if train_unembed: + target_modules.append("lm_head") else: target_modules = [] if train_attn: @@ -2192,32 +2195,8 @@ def remove(self) -> None: "router_logits": 35, "router_routing_weights": 36, "router_selected_experts": 37, - "materialized_layer_input": 38, - "delayed_pair_delta": 38, - "delayed_pair_residual": 38, - "delayed_pair_shard_sum": 38, - "delayed_pair_shard_materialized": 38, "input_norm_residual": 39, "attn_output_eager_candidate": 40, - "post_attention_o_proj_partial_sum": 41, - "post_attention_partial_residual": 41, - "post_attention_o_proj_partial_sum_split_output": 41, - "post_attention_partial_residual_split_output": 41, - "post_attention_o_proj_partial_sum_sum_then_residual": 41, - "post_attention_partial_residual_sum_then_residual": 41, - "post_attention_o_proj_partial_sum_residual_then_partials": 41, - "post_attention_partial_residual_residual_then_partials": 41, - "post_attention_o_proj_partial_sum_fp32_sum_then_residual": 41, - "post_attention_partial_residual_fp32_sum_then_residual": 41, - "moe_experts_output_tp_shard_0": 42, - "moe_experts_output_tp_shard_1": 42, - "moe_experts_output_tp_shard_2": 42, - "moe_experts_output_tp_shard_3": 42, - "moe_experts_output_tp_shard_4": 42, - "moe_experts_output_tp_shard_5": 42, - "moe_experts_output_tp_shard_6": 42, - "moe_experts_output_tp_shard_7": 42, - "moe_experts_output_tp_shard_sum": 42, "moe_experts_output_override": 43, "moe_input_override": 44, "final_residual_input": 45, @@ -3131,15 +3110,12 @@ def _teacher_hidden_cache_contributor_key(self, ps) -> Optional[int]: """Return this rank's logical cache slice key, or None for duplicate shards. Must mirror the dispatcher's batch_slice_rank_and_size mapping so cache - rows merge back in client datum order. Under legacy EP batch duplication - (XORL_SERVER_EP_DUPLICATE_BATCHES=1) only ep_rank 0 contributes. + rows merge back in client datum order. """ if getattr(ps, "cp_enabled", False) and int(getattr(ps, "cp_rank", 0)) != 0: return None if getattr(ps, "ep_enabled", False): - if ep_duplicate_batches_enabled() and int(getattr(ps, "ep_rank", 0)) != 0: - return None cp_size = max(1, int(getattr(ps, "cp_size", 1) or 1)) if getattr(ps, "cp_enabled", False) else 1 pp_size = max(1, int(getattr(ps, "pp_size", 1))) slice_rank, _ = batch_slice_rank_and_size(self.rank, self.world_size, ps, cp_size, pp_size) @@ -3427,13 +3403,11 @@ def zero(): accumulated[key] = {"sum": zero(), "op": "sum_max"} @staticmethod - def _finalize_loss_metrics(accumulated, result, loss_fn: Optional[str] = None): + def _finalize_loss_metrics(accumulated, result, loss_fn: str): """All-reduce loss metrics, then add reduced values to result dict.""" if not accumulated: return ps = get_parallel_state() - if loss_fn is None and all(str(k).startswith("opd_") for k in accumulated): - loss_fn = "opd_loss" if loss_fn == "opd_loss": reduce_group = ps.loss_group if ps.loss_parallel_enabled else None diff --git a/src/xorl/server/runner/runner_dispatcher.py b/src/xorl/server/runner/runner_dispatcher.py index b249e147..150e9fa2 100644 --- a/src/xorl/server/runner/runner_dispatcher.py +++ b/src/xorl/server/runner/runner_dispatcher.py @@ -1039,7 +1039,7 @@ def _maybe_dump_microbatch_diagnostic( routed_expert_logits: Optional[List[Any]] = None, ) -> None: params = loss_fn_params or {} - dump_dir = params.get("diagnostic_microbatch_dump_dir") or os.getenv("XORL_MICROBATCH_DIAGNOSTIC_DIR") + dump_dir = params.get("diagnostic_microbatch_dump_dir") if not dump_dir: return @@ -1119,9 +1119,7 @@ def _maybe_dump_microbatch_diagnostic( summary_path = out_dir / f"microbatch_{safe_request_id}_rank{self.rank:05d}.json" summary_path.write_text(json.dumps(summary, sort_keys=True, indent=2) + "\n", encoding="utf-8") - dump_tensors = bool(params.get("diagnostic_microbatch_dump_tensors", False)) or os.getenv( - "XORL_MICROBATCH_DIAGNOSTIC_TENSORS", "0" - ).strip().lower() in {"1", "true", "yes"} + dump_tensors = bool(params.get("diagnostic_microbatch_dump_tensors", False)) if dump_tensors: tensor_path = out_dir / f"microbatch_{safe_request_id}_rank{self.rank:05d}.pt" torch.save( @@ -1144,19 +1142,6 @@ def _create_dummy_batch(src_batch: Dict[str, Any]) -> Dict[str, Any]: target_tokens are set to -100 (IGNORE_INDEX) so cross-entropy loss = 0 and gradient_accumulate_loss produces grad_scale = 0 (local_valid_tokens = 0). """ - min_tokens_raw = os.getenv("XORL_SERVER_MINIMAL_DUMMY_BATCH_TOKENS", "").strip() - if min_tokens_raw: - try: - min_tokens = int(min_tokens_raw) - except ValueError: - min_tokens = 0 - if min_tokens > 0: - return RunnerDispatcher._create_minimal_dummy_batch(src_batch, min_tokens) - - return RunnerDispatcher._create_legacy_dummy_batch(src_batch) - - @staticmethod - def _create_legacy_dummy_batch(src_batch: Dict[str, Any]) -> Dict[str, Any]: _LABEL_KEYS = {"labels", "target_tokens"} dummy_batch = {} for key, value in src_batch.items(): @@ -1174,66 +1159,6 @@ def _create_legacy_dummy_batch(src_batch: Dict[str, Any]) -> Dict[str, Any]: dummy_batch["num_samples"] = 0 return dummy_batch - @staticmethod - def _create_minimal_dummy_batch(src_batch: Dict[str, Any], min_tokens: int) -> Dict[str, Any]: - """Create a short zero-loss dummy batch for empty data-slice ranks. - - This keeps collective participation uniform while avoiding cloned full - OPRD rows on ranks that have no real samples. The OPRD trainer-side - teacher forward falls back to this short student sequence when - teacher_input_ids/teacher_kept_indices are absent. - """ - input_ids = src_batch.get("input_ids") - if not isinstance(input_ids, torch.Tensor) or input_ids.dim() < 1: - return RunnerDispatcher._create_legacy_dummy_batch(src_batch) - - seq_len = int(input_ids.shape[-1]) - if seq_len <= 0: - return RunnerDispatcher._create_legacy_dummy_batch(src_batch) - keep = max(1, min(int(min_tokens), seq_len)) - - label_keys = {"labels", "target_tokens"} - drop_keys = { - "cu_seq_lens_q", - "cu_seq_lens_k", - "max_length_q", - "max_length_k", - "_original_position_ids", - "teacher_input_ids", - "teacher_kept_indices", - "teacher_position_ids", - "teacher_cache_indices", - "teacher_cache_local_indices", - "teacher_cache_base", - } - - dummy_batch: Dict[str, Any] = {} - for key, value in src_batch.items(): - if key in drop_keys: - continue - if isinstance(value, torch.Tensor): - if key in label_keys: - if value.dim() >= 1 and int(value.shape[-1]) == seq_len: - dummy_batch[key] = torch.full_like(value[..., :keep], -100) - else: - dummy_batch[key] = torch.full_like(value, -100) - elif key == "position_ids" and value.dim() >= 1 and int(value.shape[-1]) == seq_len: - shape = value.shape[:-1] + (keep,) - pos = torch.arange(keep, dtype=value.dtype, device=value.device) - dummy_batch[key] = pos.reshape((1,) * (len(shape) - 1) + (keep,)).expand(shape).clone() - elif key == "attention_mask" and value.dim() >= 1 and int(value.shape[-1]) == seq_len: - dummy_batch[key] = torch.ones_like(value[..., :keep]) - elif value.dim() >= 1 and int(value.shape[-1]) == seq_len: - dummy_batch[key] = value[..., :keep].clone() - else: - dummy_batch[key] = value.clone() - else: - dummy_batch[key] = value - - dummy_batch["num_samples"] = 0 - dummy_batch["_r3_sample_lengths"] = [] - return dummy_batch - @staticmethod def _dp_batch_range(dp_rank: int, base_count: int, remainder: int): """Return (start_idx, count) for a DP rank under balanced distribution. @@ -1250,10 +1175,8 @@ def _batch_parallel_rank_and_size(self, parallel_state, cp_size: int, pp_size: i """Return the logical data slice rank/size for request batch dispatch. Every logical data replica gets a distinct slice; FSDP, CP/SP, TP, and - same-stage ranks share that slice. EP groups no - longer duplicate a slice across their ranks unless the legacy - XORL_SERVER_EP_DUPLICATE_BATCHES rollback switch is set — see - batch_slice_rank_and_size for the correctness argument. + same-stage ranks share that slice. EP ranks receive distinct slices; + see batch_slice_rank_and_size for the correctness argument. """ return batch_slice_rank_and_size(self.rank, self.world_size, parallel_state, cp_size, pp_size) diff --git a/src/xorl/server/runner/utils/__init__.py b/src/xorl/server/runner/utils/__init__.py index 0534836d..0e33ac1d 100644 --- a/src/xorl/server/runner/utils/__init__.py +++ b/src/xorl/server/runner/utils/__init__.py @@ -3,7 +3,6 @@ batch_packed_rows, batch_slice_rank_and_size, convert_batch_to_tensors, - ep_duplicate_batches_enabled, positive_int_param, simple_sequence_shard, validate_batch_shapes, @@ -19,7 +18,6 @@ "batch_packed_rows", "batch_slice_rank_and_size", "convert_batch_to_tensors", - "ep_duplicate_batches_enabled", "positive_int_param", "simple_sequence_shard", "validate_batch_shapes", diff --git a/src/xorl/server/runner/utils/batch_utils.py b/src/xorl/server/runner/utils/batch_utils.py index 610678f9..bbd9e1d3 100644 --- a/src/xorl/server/runner/utils/batch_utils.py +++ b/src/xorl/server/runner/utils/batch_utils.py @@ -7,7 +7,6 @@ """ import logging -import os from typing import Any, Callable, Dict, Optional import torch @@ -254,20 +253,6 @@ def batch_packed_rows(batches: list[Dict[str, Any]], row_batch_size: int) -> lis return grouped -def ep_duplicate_batches_enabled() -> bool: - """Whether EP groups receive one duplicated batch slice (legacy dispatch). - - Legacy dispatch keyed the batch slice on the ep_fsdp coordinate, so all - ep_size ranks of an EP group computed the same packed batch — ep_size-times - redundant compute. Per-rank-distinct slices are correct: the MoE all-to-all - routes per-rank-distinct tokens (the local-training path always runs this - way) and the OPD full-vocab KL is rank-local; loss normalization by global - valid tokens makes both regimes produce identical gradients. The duplication - is kept only as a rollback switch: XORL_SERVER_EP_DUPLICATE_BATCHES=1. - """ - return os.getenv("XORL_SERVER_EP_DUPLICATE_BATCHES", "0").strip().lower() in {"1", "true", "yes"} - - def batch_slice_rank_and_size( rank: int, world_size: int, @@ -278,27 +263,15 @@ def batch_slice_rank_and_size( """Return the logical request-batch slice rank and count. Ranks that shard the same sample through FSDP, CP/SP, TP, or pipeline - parallelism share a slice. EP ranks receive distinct slices unless the - legacy duplication switch is enabled. + parallelism share a slice. EP ranks receive distinct slices: the MoE + all-to-all routes their distinct tokens, matching local training without + redundant EP-wide batch replication. """ tp_size = max(1, int(getattr(parallel_state, "tp_size", 1))) if getattr(parallel_state, "ep_enabled", False): ranks_per_pp_stage = max(1, world_size // max(1, pp_size)) local_stage_rank = rank % ranks_per_pp_stage - if ep_duplicate_batches_enabled(): - ep_size = max(1, int(getattr(parallel_state, "ep_size", 1))) - ep_fsdp_size = max(1, int(getattr(parallel_state, "dp_shard_in_ep_size", 1))) - ep_mesh = getattr(parallel_state, "ep_fsdp_device_mesh", None) - if ep_mesh is not None: - try: - ep_fsdp_rank = int(ep_mesh.get_local_rank("ep_fsdp")) - return min(ep_fsdp_rank, ep_fsdp_size - 1), ep_fsdp_size - except Exception: - logger.debug("Could not read ep_fsdp local rank; falling back to rank arithmetic", exc_info=True) - ep_fsdp_rank = min(local_stage_rank // ep_size, ep_fsdp_size - 1) - return ep_fsdp_rank, ep_fsdp_size - denom = max(1, cp_size * tp_size) slice_count = max(1, ranks_per_pp_stage // denom) return min(local_stage_rank // denom, slice_count - 1), slice_count diff --git a/src/xorl/server/server_arguments.py b/src/xorl/server/server_arguments.py index 8fb57efd..2706c635 100644 --- a/src/xorl/server/server_arguments.py +++ b/src/xorl/server/server_arguments.py @@ -412,17 +412,6 @@ class ServerArguments: default=None, metadata={"help": "Optional short names, FQNs, or globs to keep out of QARL fake quantization."}, ) - fp8_cfg: Optional[Dict[str, Any]] = field( - default=None, - metadata={ - "help": ( - "Optional compatibility alias for NeMo-style FP8 configs. Supported values are " - "{enabled: true, fp8: e4m3, fp8_recipe: blockwise, fp8_param: false}; " - "TransformerEngine-only recipes are rejected." - ) - }, - ) - fp8_training_num_first_layers_bf16: int = field( default=0, metadata={"help": "Number of initial decoder layers to keep in BF16 when FP8 training is enabled."}, @@ -882,18 +871,6 @@ class ServerArguments: }, ) - externalize_r3_payloads: bool = field( - default=False, - metadata={ - "help": ("Deprecated alias for r3_payload_transport='mooncake'. Kept only for PR-426 compatibility.") - }, - ) - - keep_r3_payloads: bool = field( - default=False, - metadata={"help": "Deprecated alias for r3_payload_keep."}, - ) - storage_limit: str = field( default="10TB", metadata={ @@ -1058,6 +1035,14 @@ class ServerArguments: }, ) + unfuse_for_lora: bool = field( + default=False, + metadata={ + "help": "Replace supported fused projections with split projections before plain LoRA injection. " + "Use only when the architecture lacks fused-base logical LoRA support." + }, + ) + moe_hybrid_shared_lora: bool = field( default=False, metadata={ @@ -1181,7 +1166,10 @@ def optimizer_kwargs(self) -> Dict[str, Any]: def __post_init__(self): """Validate and set defaults.""" - from xorl.fp8_training.config_compat import normalize_fp8_training_config # noqa: PLC0415 + if self.unfuse_for_lora and not self.enable_lora: + raise ValueError("unfuse_for_lora requires enable_lora=True") + if self.unfuse_for_lora and self.enable_qlora: + raise ValueError("unfuse_for_lora is not supported with QLoRA") from xorl.qarl import normalize_qarl_quant_cfg, qarl_unsupported_scope_reason # noqa: PLC0415 from xorl.server.orchestrator.packing import ON_OVERSIZED_MODES, PACKING_STRATEGIES # noqa: PLC0415 @@ -1196,15 +1184,6 @@ def __post_init__(self): ) if self.pad_to_multiple_of < 1: raise ValueError(f"pad_to_multiple_of must be >= 1, got {self.pad_to_multiple_of}") - if self.externalize_r3_payloads: - if self.r3_payload_transport not in {"inline", "mooncake"}: - raise ValueError( - "externalize_r3_payloads=True is a deprecated alias for " - "r3_payload_transport='mooncake' and cannot be combined with filesystem transport" - ) - self.r3_payload_transport = "mooncake" - if self.keep_r3_payloads: - self.r3_payload_keep = True if self.r3_payload_transport == "inline": if self.r3_payload_dir: raise ValueError("r3_payload_dir requires r3_payload_transport='filesystem'") @@ -1223,8 +1202,6 @@ def __post_init__(self): f"r3_payload_transport must be one of: inline, mooncake, filesystem; got {self.r3_payload_transport!r}" ) - normalized_fp8_config = normalize_fp8_training_config(vars(self), context="server.train") - self.enable_fp8_training = bool(normalized_fp8_config.get("enable_fp8_training", self.enable_fp8_training)) if self.enable_qarl and self.enable_fp8_training: raise ValueError( "enable_qarl cannot be combined with enable_fp8_training; choose one low-precision train path" @@ -1447,7 +1424,6 @@ def to_config_dict(self) -> Dict[str, Any]: "qarl_sync_format": self.qarl_sync_format, "qarl_target_modules": self.qarl_target_modules, "qarl_exclude_modules": self.qarl_exclude_modules, - "fp8_cfg": self.fp8_cfg, "fp8_training_num_first_layers_bf16": self.fp8_training_num_first_layers_bf16, "fp8_training_num_last_layers_bf16": self.fp8_training_num_last_layers_bf16, "fp8_training_allow_blackwell": self.fp8_training_allow_blackwell, @@ -1539,6 +1515,7 @@ def to_config_dict(self) -> Dict[str, Any]: "lora_alpha": self.lora_alpha, "lora_target_modules": self.lora_target_modules, "lora_target_manifest": self.lora_target_manifest, + "unfuse_for_lora": self.unfuse_for_lora, "moe_hybrid_shared_lora": self.moe_hybrid_shared_lora, "lora_export_format": self.lora_export_format, "enable_qlora": self.enable_qlora, diff --git a/src/xorl/server/weight_sync/README.md b/src/xorl/server/weight_sync/README.md index e56ddc7d..f5c8b0c1 100644 --- a/src/xorl/server/weight_sync/README.md +++ b/src/xorl/server/weight_sync/README.md @@ -300,12 +300,10 @@ P2P tuning options: - With P2P and explicit FP8 sync quantization, the handler quantizes supported projection weights on the trainer side, transfers FP8 weights plus `weight_scale_inv` tensors, and skips receiver post-processing by default - because direct P2P writes already target receiver-native FP8 storage. Set - `XORL_WEIGHT_SYNC_RUN_POST_PROCESS_WEIGHTS=1` or - `XORL_P2P_RUN_POST_PROCESS_WEIGHTS=1` only for legacy receivers that still - require finalization after P2P writes. If the receiver is FP8 but the sync - request has no FP8 quantization config, tensor-size validation should fail - instead of silently copying bf16 into FP8 locators. + because direct P2P writes already target receiver-native FP8 storage. If the + receiver is FP8 but the sync request has no FP8 quantization config, + tensor-size validation should fail instead of silently copying bf16 into FP8 + locators. - The SGLang receiver must expose a matching block-FP8 layout. XORL emits block-wise `weight_scale_inv` tensors; a receiver exposing only per-tensor `weight_scale` tensors for FusedMoE is not compatible with this sender path. @@ -348,20 +346,9 @@ P2P tuning options: multi-endpoint P2P. It sends each receiver endpoint through its own serialized sync group, avoiding cross-endpoint Mooncake session reuse at the cost of giving up normal endpoint fanout parallelism. -- `XORL_P2P_SCATTER_COPY_MODE`: controls how rank 0 builds per-sender tensor - map payloads for direct-EP scatter. Default `none` reuses read-only locator - lists/dicts while constructing scatter payloads. Set `list` to shallow-copy - lists or `deep` to copy every locator dict for debugging. -- `XORL_P2P_SCATTER_REUSE_LOCATORS`: legacy boolean alias for the default - scatter copy mode. Set `1` to force locator reuse even when older manifests - still set `XORL_P2P_SCATTER_COPY_MODE=list`; set `0` to force shallow list - copies when `XORL_P2P_SCATTER_COPY_MODE` is unset. - `XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES`: explicit MoE bucket cap override. Without this override, P2P uses a 2 GiB MoE bucket cap to amortize Mooncake fixed costs; non-P2P backends keep the 256 MiB default. -- `XORL_WEIGHT_SYNC_BUCKET_BYTES`: legacy alias for the MoE bucket cap. Prefer - `XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES` so dense/root chunking stays independent - from MoE batching. - `XORL_P2P_USE_ASYNC_API=1`: opt into Mooncake's async write API. The default synchronous API path is the sustained-test path; async status polling has shown repeated-update `status=-1` failures and should remain experimental. diff --git a/src/xorl/server/weight_sync/backends/p2p.py b/src/xorl/server/weight_sync/backends/p2p.py index b2014973..0605c0b2 100644 --- a/src/xorl/server/weight_sync/backends/p2p.py +++ b/src/xorl/server/weight_sync/backends/p2p.py @@ -770,7 +770,6 @@ def __init__(self, config: TransportConfig, **kwargs: Any) -> None: super().__init__(config) # backend_config carries optional Mooncake engine setup overrides. be_cfg = config.backend_config or {} - self._qwen_linear_attention_dims: Dict[str, Any] = dict(be_cfg.get("qwen_linear_attention_dims") or {}) self._engine = None # MooncakeTransferEngine (lazy-imported) # tensor_map[name] -> list of receiver locator dicts. self._tensor_map: Dict[str, List[Dict[str, Any]]] = {} @@ -990,58 +989,18 @@ def _endpoint_indices_for_tensor_map( continue return endpoint_indices - @staticmethod - def _copy_locator_list_for_scatter( - locators: List[Dict[str, Any]], - copy_mode: str, - ) -> List[Dict[str, Any]]: - if copy_mode == "deep": - return [dict(loc) for loc in locators] - if copy_mode == "none": - return locators - # Default: keep an independent list per scatter payload without - # duplicating every immutable locator dict. Nonzero ranks copy the - # dicts again when adopting/merging prepared state. - return list(locators) - - @staticmethod - def _scatter_locator_copy_mode() -> str: - if "XORL_P2P_SCATTER_REUSE_LOCATORS" in os.environ: - if _env_flag("XORL_P2P_SCATTER_REUSE_LOCATORS", False): - return "none" - if "XORL_P2P_SCATTER_COPY_MODE" not in os.environ: - return "list" - raw_mode = os.environ.get("XORL_P2P_SCATTER_COPY_MODE") - if raw_mode is None: - # Fast path: locator lists/dicts are read-only after SGLang prepare, - # and scatter_object_list serializes each recipient payload anyway. - return "none" - raw = raw_mode.strip().lower() - if raw in {"deep", "dict", "dicts"}: - return "deep" - if raw in {"list", "shallow", "lists"}: - return "list" - if raw in {"none", "reuse"}: - return "none" - logger.warning("[P2P] invalid XORL_P2P_SCATTER_COPY_MODE=%r; using reuse", raw_mode) - return "none" - def _filter_tensor_map_for_sender( self, tensor_map: Dict[str, List[Dict[str, Any]]], sender_rank: int, *, experts_per_ep: Optional[int] = None, - locator_copy_mode: str = "list", ) -> Dict[str, List[Dict[str, Any]]]: sender_ep_rank = self._sender_ep_ranks.get(int(sender_rank)) if experts_per_ep is None: experts_per_ep = self._experts_per_ep(tensor_map, self._direct_ep_size) if sender_ep_rank is None or experts_per_ep is None: - return { - name: self._copy_locator_list_for_scatter(locators, locator_copy_mode) - for name, locators in tensor_map.items() - } + return dict(tensor_map) keep_dense = int(sender_rank) == 0 filtered: Dict[str, List[Dict[str, Any]]] = {} @@ -1055,7 +1014,9 @@ def _filter_tensor_map_for_sender( continue elif expert_idx // experts_per_ep != sender_ep_rank: continue - filtered[name] = self._copy_locator_list_for_scatter(locators, locator_copy_mode) + # Prepared locators are immutable here. scatter_object_list + # serializes each recipient payload before a nonzero rank adopts it. + filtered[name] = locators return filtered def should_extract_dense_params_on_rank(self, rank: int) -> bool: @@ -1134,7 +1095,6 @@ def _initialize_payloads_for_sender_order(self) -> List[Any]: kind = "tensor_map_with_infos" experts_per_ep = self._experts_per_ep(self._tensor_map, self._direct_ep_size) - locator_copy_mode = self._scatter_locator_copy_mode() payloads: List[Any] = [] for sender_rank in self.sender_rank_order: if int(sender_rank) == 0: @@ -1144,7 +1104,6 @@ def _initialize_payloads_for_sender_order(self) -> List[Any]: self._tensor_map, sender_rank, experts_per_ep=experts_per_ep, - locator_copy_mode=locator_copy_mode, ) if kind == "merge_tensor_map": payloads.append( @@ -1378,11 +1337,7 @@ def _initialize_single_sender(self) -> bool: # re-arming the receiver's Mooncake buffers (the receiver honors p2p_invalidate_cache # via _invalidate_p2p_cache()). Restored from the pre-merge baseline (baf6dcce). cache_invalidation_mode = str(cfg.backend_config.get("cache_invalidation_mode", "auto") or "auto").lower() - invalidate_receiver_cache = ( - cache_invalidation_mode != "none" - and not request_cached_prepare - and _env_flag("XORL_P2P_INVALIDATE_RECEIVER_CACHE_ON_COLD_PREPARE", True) - ) + invalidate_receiver_cache = cache_invalidation_mode != "none" and not request_cached_prepare self._last_prepare_returned_tensor_map = False self._last_prepare_tensor_map_endpoint_indices = set() num_endpoints = len(cfg.endpoints) @@ -1946,7 +1901,7 @@ def _transfer_bucket_impl( if self._rank_filter is not None and not self._rank_filter(loc): continue locators_for_rank += 1 - src_view = self._slice_source_for_locator(name, tensor, loc, self._qwen_linear_attention_dims) + src_view = self._slice_source_for_locator(name, tensor, loc) if src_view is None: skipped_errors.append(f"{name!r}: receiver locator is incompatible with source tensor") continue @@ -2610,7 +2565,6 @@ def _slice_source_for_locator( name: str, full_tensor: torch.Tensor, loc: Dict[str, Any], - qwen_linear_attention_dims: Optional[Dict[str, Any]] = None, ) -> Optional[torch.Tensor]: """Extract the sub-region of the trainer's full HF tensor that corresponds to a single receiver locator. @@ -2627,17 +2581,6 @@ def _slice_source_for_locator( return P2PTransportBackend._normalize_sliced_source_for_locator(name, full_tensor, loc) full_shape = loc.get("full_shape") - fused_linear_view = P2PTransportBackend._slice_qwen_linear_attention_fused_param( - name, - full_tensor, - loc, - full_shape, - slc, - qwen_linear_attention_dims, - ) - if fused_linear_view is not None: - return P2PTransportBackend._normalize_sliced_source_for_locator(name, fused_linear_view, loc) - if full_shape is not None and list(full_tensor.shape) != list(full_shape): local_view = P2PTransportBackend._slice_qwen35_linear_attention_local_param( name, @@ -2658,85 +2601,6 @@ def _slice_source_for_locator( index: Tuple[slice, ...] = tuple(slice(int(s[0]), int(s[1])) for s in slc) return P2PTransportBackend._normalize_sliced_source_for_locator(name, full_tensor[index], loc) - @staticmethod - def _slice_qwen_linear_attention_fused_param( - name: str, - full_tensor: torch.Tensor, - loc: Dict[str, Any], - full_shape: Any, - slc: Any, - qwen_linear_attention_dims: Optional[Dict[str, Any]], - ) -> Optional[torch.Tensor]: - # This combined Q|K|V cat has no matching receiver locator. SGLang's - # p2p_qwen35_linear_attn_qkvz_locators emits SEPARATE contiguous Q/K/V - # locators for in_proj_qkv.weight (and conv1d.weight), each a plain - # contiguous row-range that the generic full_tensor[slc] path in - # _slice_source_for_locator matches exactly. The fused-cat path is an - # orphaned trainer-side divergence from the glm5 rebase (955191dd) with - # no receiver counterpart — when it fires its output is the wrong shape - # for a per-Q/K/V locator. Bypassed by default; returning None makes the - # caller fall through to the canonical contiguous slice. Set - # XORL_P2P_DISABLE_QWEN_LINEAR_ATTN_FUSED_SLICE=0 only for a receiver that - # genuinely expects the fused combined layout. - if _env_flag("XORL_P2P_DISABLE_QWEN_LINEAR_ATTN_FUSED_SLICE", True): - return None - if not ( - ".linear_attn." in name - and (name.endswith(".in_proj_qkv.weight") or name.endswith(".conv1d.weight")) - and full_shape is not None - and list(full_tensor.shape) == list(full_shape) - and isinstance(qwen_linear_attention_dims, dict) - ): - return None - - try: - key_dim = int(qwen_linear_attention_dims["key_dim"]) - value_dim = int(qwen_linear_attention_dims["value_dim"]) - except (KeyError, TypeError, ValueError): - return None - if key_dim <= 0 or value_dim <= 0: - return None - - total_rows = 2 * key_dim + value_dim - if full_tensor.ndim < 1 or int(full_tensor.shape[0]) != total_rows: - return None - - try: - tp_rank = int(loc.get("tp_rank", 0)) - except (TypeError, ValueError): - return None - - tp_size = 0 - try: - tp_size = int(qwen_linear_attention_dims.get("tp_size") or 0) - except (TypeError, ValueError): - tp_size = 0 - if tp_size <= 0 and slc: - try: - local_rows = int(slc[0][1]) - int(slc[0][0]) - except (TypeError, ValueError, IndexError): - local_rows = 0 - if local_rows > 0 and total_rows % local_rows == 0: - tp_size = total_rows // local_rows - if tp_size <= 0 or tp_rank < 0 or tp_rank >= tp_size: - return None - if key_dim % tp_size != 0 or value_dim % tp_size != 0: - return None - - key_chunk = key_dim // tp_size - value_chunk = value_dim // tp_size - q_start = tp_rank * key_chunk - k_start = key_dim + tp_rank * key_chunk - v_start = 2 * key_dim + tp_rank * value_chunk - return torch.cat( - [ - full_tensor.narrow(0, q_start, key_chunk), - full_tensor.narrow(0, k_start, key_chunk), - full_tensor.narrow(0, v_start, value_chunk), - ], - dim=0, - ).contiguous() - @staticmethod def _normalize_source_for_locator( name: str, diff --git a/src/xorl/server/weight_sync/handler.py b/src/xorl/server/weight_sync/handler.py index 0dfaf702..3b07c4ef 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -91,11 +91,9 @@ def _env_int(name: str, default: int, *, minimum: int = 1) -> int: def _moe_bucket_size_bytes(sync_method: str) -> int: - """Default MoE bucket sizing is backend-specific; env vars are explicit overrides.""" + """Default MoE bucket sizing is backend-specific; the env var is an explicit override.""" default = _DEFAULT_P2P_MOE_BUCKET_BYTES if sync_method == "p2p" else _DEFAULT_MOE_BUCKET_BYTES - if "XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES" in os.environ: - return _env_int("XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES", default) - return _env_int("XORL_WEIGHT_SYNC_BUCKET_BYTES", default) + return _env_int("XORL_WEIGHT_SYNC_MOE_BUCKET_BYTES", default) def _env_bool(name: str, default: bool = False) -> bool: @@ -1119,25 +1117,6 @@ def _sync_weights( ) if weight_version is not None: _backend_config["weight_version"] = weight_version - if has_linear_attention_layers(model.config): - linear_num_key_heads = getattr(model.config, "linear_num_key_heads", None) - linear_key_head_dim = getattr(model.config, "linear_key_head_dim", None) - linear_num_value_heads = getattr(model.config, "linear_num_value_heads", None) - linear_value_head_dim = getattr(model.config, "linear_value_head_dim", None) - if all( - value is not None - for value in ( - linear_num_key_heads, - linear_key_head_dim, - linear_num_value_heads, - linear_value_head_dim, - ) - ): - _backend_config["qwen_linear_attention_dims"] = { - "key_dim": int(linear_num_key_heads) * int(linear_key_head_dim), - "value_dim": int(linear_num_value_heads) * int(linear_value_head_dim), - "tp_size": max(int(ep.get("world_size", 1) or 1) for ep in endpoints) if endpoints else 1, - } ib_device = _select_p2p_ib_device(self.rank, self.world_size) if ib_device: _backend_config["ib_device"] = ib_device @@ -3233,16 +3212,6 @@ def _fp8_cpu_workspace_min_capacity() -> int: def _fp8_cpu_workspace_streaming_enabled() -> bool: return os.environ.get("XORL_P2P_FP8_CPU_WORKSPACE_STREAMING", "1") != "0" - @staticmethod - def _p2p_should_run_post_process_weights(quantization: Optional[Dict[str, Any]]) -> bool: - if not (quantization and quantization.get("quant_method") == "fp8"): - return False - if "XORL_P2P_RUN_POST_PROCESS_WEIGHTS" in os.environ: - return _env_bool("XORL_P2P_RUN_POST_PROCESS_WEIGHTS") - if "XORL_WEIGHT_SYNC_RUN_POST_PROCESS_WEIGHTS" in os.environ: - return _env_bool("XORL_WEIGHT_SYNC_RUN_POST_PROCESS_WEIGHTS") - return False - @staticmethod def _should_run_receiver_post_process_after_fp8_sync( sync_method: str, @@ -3257,10 +3226,7 @@ def _should_run_receiver_post_process_after_fp8_sync( if sync_method == "sparse_delta": return bool(fp8_kv_cache_postprocess_required) if sync_method == "p2p": - return bool( - fp8_kv_cache_postprocess_required - or WeightSyncHandler._p2p_should_run_post_process_weights(quantization) - ) + return bool(fp8_kv_cache_postprocess_required) return True @staticmethod @@ -4150,6 +4116,8 @@ def _extract_params_for_sync( # (LoraDeltaLinear, start, end) row-slice of the fused delta to fold into it. lora_param_names = set() fused_gdn_base_deltas = {} + fused_gate_up_deltas = {} + fused_qkv_deltas = {} for mname, mod in lora_modules.items(): prefix = f"{mname}." if mname else "" if isinstance(mod, QLoRALinear): @@ -4179,6 +4147,22 @@ def _extract_params_for_sync( ) elif gdn_leaf == "out_proj": fused_gdn_base_deltas[f"{gdn_parent_name}.o_proj"] = (mod, 0, mod.out_features) + elif gdn_leaf in {"q_proj", "k_proj", "v_proj"} and hasattr(gdn_parent, "qkv_proj"): + base_name = f"{gdn_parent_name}.qkv_proj" + entry = fused_qkv_deltas.setdefault( + base_name, + { + "sizes": ( + int(getattr(gdn_parent, "q_dim")), + int(getattr(gdn_parent, "kv_dim")), + int(getattr(gdn_parent, "kv_dim")), + ) + }, + ) + entry[gdn_leaf] = mod + elif gdn_leaf in {"gate_proj", "up_proj"} and hasattr(gdn_parent, "gate_up_proj"): + base_name = f"{gdn_parent_name}.gate_up_proj" + fused_gate_up_deltas.setdefault(base_name, {})[gdn_leaf] = mod else: raise RuntimeError(f"Unexpected fused-GDN LoRA leaf {gdn_leaf!r} for {mname}") elif isinstance(mod, LoraLinear): @@ -4229,6 +4213,45 @@ def _extract_params_for_sync( parent_name = ".".join(pname.split(".")[:-1]) param_leaf = pname.split(".")[-1] # e.g. "weight", "gate_proj" + # Fused base projections retain one GEMM while training independent + # logical factors. Publish the same canonical fold used by forward. + fused_qkv = fused_qkv_deltas.get(parent_name) + if fused_qkv is not None and param_leaf == "weight": + sizes = fused_qkv["sizes"] + base_parts = param.data.split(sizes, dim=0) + folded_parts = [] + for leaf, base in zip(("q_proj", "k_proj", "v_proj"), base_parts, strict=True): + delta_module = fused_qkv.get(leaf) + folded = base if delta_module is None else delta_module._merged_weight(base) + folded_parts.append(folded.to(dtype=torch.bfloat16)) + buffer.append((full_name, torch.cat(folded_parts, dim=0).clone())) + continue + + fused_gate_up = fused_gate_up_deltas.get(parent_name) + if fused_gate_up is not None and param_leaf == "weight": + if param.shape[0] % 2: + raise RuntimeError(f"Fused gate/up projection {full_name} has odd output size {param.shape[0]}") + gate_base, up_base = param.data.chunk(2, dim=0) + folded_parts = [] + for leaf, base in (("gate_proj", gate_base), ("up_proj", up_base)): + delta_module = fused_gate_up.get(leaf) + if delta_module is None: + folded_parts.append(base.to(dtype=torch.bfloat16)) + continue + if lora_merged_forward_enabled(delta_module): + folded = delta_module._merged_weight(base).to(dtype=torch.bfloat16) + else: + delta = delta_module.get_delta_weight() + if tuple(delta.shape) != tuple(base.shape): + raise RuntimeError( + f"Fused {leaf} delta for {full_name} has shape " + f"{tuple(delta.shape)}, expected {tuple(base.shape)}" + ) + folded = base.to(dtype=torch.bfloat16) + delta.to(dtype=torch.bfloat16) + folded_parts.append(folded) + buffer.append((full_name, torch.cat(folded_parts, dim=0).clone())) + continue + # Fused-GDN fold: this base projection (q/k/v/g/o_proj) gets the # corresponding row-slice of the fused in_proj_qkvz / out_proj delta. fused_delta_spec = fused_gdn_base_deltas.get(parent_name) diff --git a/src/xorl/server/weight_sync/source_delta_capture.py b/src/xorl/server/weight_sync/source_delta_capture.py index ce58cb34..15e72ae3 100644 --- a/src/xorl/server/weight_sync/source_delta_capture.py +++ b/src/xorl/server/weight_sync/source_delta_capture.py @@ -21,7 +21,6 @@ from xorl.server.weight_sync.sparse_delta_files import ( SparseTensorUpdate, _render_rank_filename, - prepare_delta_encoding_runtime, write_sparse_delta_file, ) @@ -214,79 +213,6 @@ def write_sparse_source_delta_global_manifest( return global_manifest -def load_sparse_source_delta_inputs( - manifest_path: str | Path, - *, - delta_encoding_path: str | None = None, - use_native_extension: bool = False, - tag: str | None = "enc", - include_empty: bool = False, -) -> list[tuple[Any, Any]]: - """Load source-rank packed files as ``(StoreKey, EncodedDelta)`` inputs. - - The returned tensors are cloned out of the mmap so callers can close files - before running the translation engine. - When ``include_empty`` is true, the returned inputs also include encoded - empty tensors for manifest entries with ``nnz == 0``. Translation plans for - sharded source layouts need those empty shards to know the input is - complete. - """ - - prepare_delta_encoding_runtime( - delta_encoding_path=delta_encoding_path, - use_native_extension=use_native_extension, - ) - - from delta_encoding.encoding.compression import encode # noqa: PLC0415 - from delta_encoding.encoding.packed import MmapPackedFile # noqa: PLC0415 - from delta_encoding.encoding.types import EncodedDelta # noqa: PLC0415 - from delta_encoding.ops.types import StoreKey # noqa: PLC0415 - - manifest = json.loads(Path(manifest_path).read_text()) - if manifest.get("format") != GLOBAL_MANIFEST_VERSION: - raise ValueError(f"Unsupported sparse source-delta manifest format: {manifest.get('format')!r}") - - inputs: list[tuple[Any, Any]] = [] - seen: set[tuple[int, str]] = set() - for rank_manifest in manifest.get("ranks", []): - packed_path = rank_manifest.get("packed_path") - if not packed_path: - continue - rank = int(rank_manifest["rank"]) - with MmapPackedFile(packed_path) as packed: - for entry in packed.entries: - key = StoreKey(entry.name, rank=rank) - if tag: - key = key.tag(tag) - encoded = EncodedDelta( - packed.flat_deltas_view(entry).clone(), - packed.values_view(entry).clone(), - tuple(entry.shape), - ) - inputs.append((key, encoded)) - seen.add((rank, entry.name)) - - if include_empty: - for rank_manifest in manifest.get("ranks", []): - rank = int(rank_manifest["rank"]) - for tensor in rank_manifest.get("tensors", []): - name = str(tensor["name"]) - if (rank, name) in seen or int(tensor.get("nnz", 0)) != 0: - continue - key = StoreKey(name, rank=rank) - if tag: - key = key.tag(tag) - dtype = _dtype_from_name(tensor.get("dtype") or rank_manifest.get("capture_dtype") or "bfloat16") - shape = tuple(int(dim) for dim in tensor["shape"]) - encoded = encode( - torch.empty(0, dtype=torch.int32), - torch.empty(0, dtype=dtype), - shape, - ) - inputs.append((key, encoded)) - return inputs - - def resolve_sparse_delta_capture_output_dir( config: Mapping[str, Any], *, diff --git a/src/xorl/server/weight_sync/sparse_delta_files.py b/src/xorl/server/weight_sync/sparse_delta_files.py index 4d4c17e8..0d0152a7 100644 --- a/src/xorl/server/weight_sync/sparse_delta_files.py +++ b/src/xorl/server/weight_sync/sparse_delta_files.py @@ -13,7 +13,7 @@ import sys from dataclasses import dataclass from pathlib import Path -from typing import Any, Iterable, Mapping, Optional, Sequence +from typing import Any, Iterable, Optional import torch @@ -38,16 +38,6 @@ class SparseDeltaFileStats: packed_bytes: int -@dataclass(frozen=True) -class SparseTensorShard: - """Receiver-local shard of a logical sparse tensor.""" - - rank: int - name: str - shape: tuple[int, ...] - slices: tuple[tuple[int, int], ...] - - def _load_delta_encoding( *, delta_encoding_path: Optional[str] = None, @@ -192,206 +182,6 @@ def write_sparse_delta_file( ) -def make_contiguous_shards( - *, - name: str, - shape: Sequence[int], - shard_dim: int, - num_shards: int, - shard_sizes: Sequence[int] | None = None, -) -> list[SparseTensorShard]: - """Describe uniform contiguous receiver shards for one logical tensor.""" - - shape = tuple(int(dim) for dim in shape) - if not shape: - raise ValueError("shape must be non-empty") - if num_shards <= 0: - raise ValueError(f"num_shards must be positive, got {num_shards}") - if shard_dim < 0: - shard_dim += len(shape) - if shard_dim < 0 or shard_dim >= len(shape): - raise ValueError(f"shard_dim {shard_dim} is out of range for shape {shape}") - - full = shape[shard_dim] - if shard_sizes is None: - if full % num_shards != 0: - raise ValueError( - f"shape {shape} is not evenly divisible into {num_shards} shards along dim {shard_dim}; " - "pass explicit shard_sizes" - ) - shard_sizes = [full // num_shards] * num_shards - else: - shard_sizes = [int(size) for size in shard_sizes] - if len(shard_sizes) != num_shards: - raise ValueError(f"expected {num_shards} shard sizes, got {len(shard_sizes)}") - if any(size < 0 for size in shard_sizes): - raise ValueError(f"shard sizes must be non-negative, got {shard_sizes}") - if sum(shard_sizes) != full: - raise ValueError(f"shard sizes sum to {sum(shard_sizes)}, expected {full}") - - shards: list[SparseTensorShard] = [] - start = 0 - for rank, size in enumerate(shard_sizes): - stop = start + size - local_shape = list(shape) - local_shape[shard_dim] = size - slices = [(0, dim) for dim in shape] - slices[shard_dim] = (start, stop) - shards.append( - SparseTensorShard( - rank=rank, - name=name, - shape=tuple(local_shape), - slices=tuple(slices), - ) - ) - start = stop - return shards - - -def split_sparse_update_by_shards( - update: SparseTensorUpdate, - shards: Sequence[SparseTensorShard], -) -> dict[int, SparseTensorUpdate]: - """Convert one logical sparse tensor update into receiver-local updates.""" - - _validate_update(update) - if not shards: - raise ValueError("split_sparse_update_by_shards requires at least one shard") - - logical_shape = tuple(update.shape) - for shard in shards: - if len(shard.shape) != len(logical_shape) or len(shard.slices) != len(logical_shape): - raise ValueError( - f"Shard rank {shard.rank} rank mismatch: update shape={logical_shape}, " - f"shard shape={shard.shape}, slices={shard.slices}" - ) - for dim, ((start, stop), local_dim, full_dim) in enumerate(zip(shard.slices, shard.shape, logical_shape)): - if start < 0 or stop < start or stop > full_dim: - raise ValueError(f"Shard rank {shard.rank} has invalid slice {start, stop} for dim {dim}") - if stop - start != local_dim: - raise ValueError( - f"Shard rank {shard.rank} local shape {shard.shape} does not match slice {shard.slices}" - ) - - sorted_update = _sorted_cpu_update(update) - flat = sorted_update.flat_indices.to(torch.int64) - values = sorted_update.values - coords = _flat_to_coords(flat, logical_shape) - - by_rank: dict[int, SparseTensorUpdate] = {} - for shard in shards: - mask = torch.ones(flat.numel(), dtype=torch.bool) - local_coords: list[torch.Tensor] = [] - for coord, (start, stop) in zip(coords, shard.slices): - mask &= (coord >= start) & (coord < stop) - local_coords.append(coord - start) - - if bool(mask.any().item()): - selected_coords = [coord[mask] for coord in local_coords] - local_flat = _coords_to_flat(selected_coords, shard.shape) - local_values = values[mask] - if local_flat.numel() > 1: - order = torch.argsort(local_flat, stable=True) - local_flat = local_flat[order] - local_values = local_values[order] - else: - local_flat = torch.empty(0, dtype=torch.int32) - local_values = values[:0] - - by_rank[int(shard.rank)] = SparseTensorUpdate( - name=shard.name, - flat_indices=local_flat.to(torch.int32), - values=local_values.contiguous(), - shape=tuple(shard.shape), - ) - return by_rank - - -def split_sparse_update_by_contiguous_shards( - update: SparseTensorUpdate, - *, - shard_dim: int, - num_shards: int, - output_name: str | None = None, - shard_sizes: Sequence[int] | None = None, -) -> dict[int, SparseTensorUpdate]: - """Split a logical update into per-rank contiguous local coordinates.""" - - shards = make_contiguous_shards( - name=output_name or update.name, - shape=update.shape, - shard_dim=shard_dim, - num_shards=num_shards, - shard_sizes=shard_sizes, - ) - return split_sparse_update_by_shards(update, shards) - - -def write_sparse_delta_files_by_rank( - updates_by_rank: Mapping[int, Iterable[SparseTensorUpdate]], - output_dir: str | Path, - *, - filename_template: str = "rank{rank}.packed", - delta_encoding_path: Optional[str] = None, - use_native_extension: bool = False, -) -> dict[int, SparseDeltaFileStats]: - """Write rank-local sparse updates as one packed file per receiver rank.""" - - if not updates_by_rank: - raise ValueError("write_sparse_delta_files_by_rank requires at least one rank") - - output_dir = Path(output_dir) - stats: dict[int, SparseDeltaFileStats] = {} - for rank, updates in sorted(updates_by_rank.items()): - rank_int = int(rank) - filename = _render_rank_filename(filename_template, rank_int) - path = output_dir / filename - stats[rank_int] = write_sparse_delta_file( - list(updates), - path, - delta_encoding_path=delta_encoding_path, - use_native_extension=use_native_extension, - ) - return stats - - -def write_encoded_sparse_delta_files_by_rank( - encoded_by_rank: Mapping[int, Mapping[str, Any]], - output_dir: str | Path, - *, - filename_template: str = "rank{rank}.packed", - delta_encoding_path: Optional[str] = None, - use_native_extension: bool = False, -) -> dict[int, SparseDeltaFileStats]: - """Write per-rank ``delta-encoding`` EncodedDelta outputs as packed files.""" - - if not encoded_by_rank: - raise ValueError("write_encoded_sparse_delta_files_by_rank requires at least one rank") - - _, write_packed_file = _load_delta_encoding( - delta_encoding_path=delta_encoding_path, - use_native_extension=use_native_extension, - ) - - output_dir = Path(output_dir) - stats: dict[int, SparseDeltaFileStats] = {} - for rank, encoded_tensors in sorted(encoded_by_rank.items()): - if not encoded_tensors: - raise ValueError(f"Rank {rank} has no encoded sparse-delta tensors") - rank_int = int(rank) - path = output_dir / _render_rank_filename(filename_template, rank_int) - written = Path(write_packed_file(dict(encoded_tensors), path)) - nnz = sum(int(getattr(encoded, "values").numel()) for encoded in encoded_tensors.values()) - stats[rank_int] = SparseDeltaFileStats( - path=str(written), - tensors=len(encoded_tensors), - nnz=nnz, - packed_bytes=written.stat().st_size, - ) - return stats - - def _render_rank_filename(filename_template: str, rank: int) -> str: """Render one rank filename without allowing directory traversal.""" filename = filename_template.format(rank=rank) @@ -405,108 +195,3 @@ def _render_rank_filename(filename_template: str, rank: int) -> str: ): raise ValueError("Sparse-delta filename_template must render a plain filename") return filename - - -def collect_encoded_sparse_deltas_by_rank( - futures: Iterable[Any], - *, - expected_ranks: int | Sequence[int] | None = None, -) -> dict[int, dict[str, Any]]: - """Drain ``delta-encoding`` TranslationFutures into rank-keyed outputs. - - ``TranslationFuture.key`` is a ``delta_encoding.ops.StoreKey`` with a - receiver rank and tensor name. ``TranslationFuture.wait()`` returns the - terminal EncodedDelta for that rank/name. - """ - - by_rank: dict[int, dict[str, Any]] = {} - for future in futures: - key = getattr(future, "key", None) - if key is None: - raise TypeError("Translation future is missing a key attribute") - strip_tags = getattr(key, "strip_tags", None) - if callable(strip_tags): - key = strip_tags() - - rank = getattr(key, "rank", None) - name = getattr(key, "name", None) - if rank is None: - raise ValueError(f"Refusing to serialize unranked sparse-delta terminal output: {key}") - if not name: - raise ValueError(f"Refusing to serialize sparse-delta terminal output with empty name: {key}") - - rank_int = int(rank) - rank_outputs = by_rank.setdefault(rank_int, {}) - name = str(name) - if name in rank_outputs: - raise ValueError(f"Duplicate sparse-delta terminal output for rank {rank_int}: {name!r}") - rank_outputs[name] = future.wait() - - _validate_expected_ranks(by_rank, expected_ranks) - return by_rank - - -def write_translation_futures_as_sparse_delta_files( - futures: Iterable[Any], - output_dir: str | Path, - *, - expected_ranks: int | Sequence[int] | None = None, - filename_template: str = "rank{rank}.packed", - delta_encoding_path: Optional[str] = None, - use_native_extension: bool = False, -) -> dict[int, SparseDeltaFileStats]: - """Write ``delta-encoding`` TranslationFuture outputs as packed files.""" - - encoded_by_rank = collect_encoded_sparse_deltas_by_rank(futures, expected_ranks=expected_ranks) - return write_encoded_sparse_delta_files_by_rank( - encoded_by_rank, - output_dir, - filename_template=filename_template, - delta_encoding_path=delta_encoding_path, - use_native_extension=use_native_extension, - ) - - -def _validate_expected_ranks( - by_rank: Mapping[int, Mapping[str, Any]], - expected_ranks: int | Sequence[int] | None, -) -> None: - if expected_ranks is None: - return - if isinstance(expected_ranks, int): - expected = set(range(expected_ranks)) - else: - expected = {int(rank) for rank in expected_ranks} - - actual = set(by_rank) - missing = sorted(expected - actual) - extra = sorted(actual - expected) - if missing or extra: - parts = [] - if missing: - parts.append(f"missing ranks {missing}") - if extra: - parts.append(f"unexpected ranks {extra}") - raise ValueError("Sparse-delta terminal rank mismatch: " + ", ".join(parts)) - - -def _flat_to_coords(flat_indices: torch.Tensor, shape: tuple[int, ...]) -> list[torch.Tensor]: - remaining = flat_indices.to(torch.int64) - coords = [torch.empty_like(remaining) for _ in shape] - for dim in range(len(shape) - 1, -1, -1): - coords[dim] = remaining % shape[dim] - remaining = remaining // shape[dim] - return coords - - -def _coords_to_flat(coords: list[torch.Tensor], shape: tuple[int, ...]) -> torch.Tensor: - if not coords: - return torch.empty(0, dtype=torch.int32) - flat = torch.zeros_like(coords[0], dtype=torch.int64) - stride = 1 - for coord, dim in zip(reversed(coords), reversed(shape)): - flat += coord.to(torch.int64) * stride - stride *= dim - if flat.numel() == 0 or stride <= torch.iinfo(torch.int32).max: - return flat.to(torch.int32) - return flat diff --git a/src/xorl/sim/analytical_ledgers.py b/src/xorl/sim/analytical_ledgers.py index 50560039..50eae113 100644 --- a/src/xorl/sim/analytical_ledgers.py +++ b/src/xorl/sim/analytical_ledgers.py @@ -19,7 +19,6 @@ from __future__ import annotations -from types import SimpleNamespace from typing import Any @@ -1312,50 +1311,6 @@ def build_model_analytical_coverage( } -def reference_counter_total_flops( - metadata: ModelMetadata, - topology: Topology, - *, - seq_len: int | None = None, - batch_seqlens: list[int] | None = None, -) -> float | None: - """Total FLOPs from the *actual* trainer ``XorlFlopsCounter`` (transcription ground truth). - - Returns None if xorl is not importable. Used by tests to assert the analytical ledger reproduces - the trainer's convention EXACTLY, which is the only honest validation of FLOPs (they are a logged - convention, not a hardware measurement). - """ - seq = _seq_len(topology, seq_len) - if seq is None and not batch_seqlens: - return None - try: - from xorl.utils.count_flops import XorlFlopsCounter # noqa: PLC0415 (lazy: xorl optional/heavy) - except Exception: # pragma: no cover - xorl not importable in this context - return None - cfg = SimpleNamespace( - model_type="qwen3_moe" if metadata.moe_intermediate_size is not None else "qwen3", - hidden_size=metadata.hidden_size, - vocab_size=metadata.vocab_size, - intermediate_size=metadata.intermediate_size, - moe_intermediate_size=metadata.moe_intermediate_size, - num_hidden_layers=metadata.num_hidden_layers, - num_key_value_heads=metadata.num_key_value_heads, - num_attention_heads=metadata.num_attention_heads, - num_experts=metadata.num_experts, - num_experts_per_tok=metadata.top_k, - head_dim=metadata.head_dim, - ) - counter = XorlFlopsCounter(cfg, gradient_checkpointing_enabled=False) - if batch_seqlens is None: - batch_seqlens = [seq] * topology.global_batch_size - tokens_sum = sum(batch_seqlens) - if cfg.model_type == "qwen3_moe": - tflops = counter._estimate_qwen3_moe_flops(tokens_sum, batch_seqlens, delta_time=1.0) - else: - tflops = counter._estimate_qwen2_flops(tokens_sum, batch_seqlens, delta_time=1.0) - return float(tflops) * 1e12 - - def hardware_flops_ledger(ledger: dict[str, Any], train: dict[str, Any]) -> dict[str, Any]: """Recompute-aware HARDWARE FLOPs (what actually runs on the GPU and determines step time). diff --git a/src/xorl/sim/benchmark_behavior.py b/src/xorl/sim/benchmark_behavior.py index f9af483f..0856ea39 100644 --- a/src/xorl/sim/benchmark_behavior.py +++ b/src/xorl/sim/benchmark_behavior.py @@ -368,9 +368,6 @@ def _result_throughput_point( moe_implementation=_first_non_none( throughput.get("moe_implementation"), topology_defaults.get("moe_implementation") ), - moe_checkpoint_method=_first_non_none( - throughput.get("moe_checkpoint_method"), topology_defaults.get("moe_checkpoint_method") - ), muon_update_dtype=_first_non_none( throughput.get("muon_update_dtype"), topology_defaults.get("muon_update_dtype") ), @@ -1135,9 +1132,6 @@ def _best_by_mfu_point( _trial_moe_implementation(trial), topology_defaults.get("moe_implementation"), ), - moe_checkpoint_method=_first_non_none( - row.get("moe_checkpoint_method"), topology_defaults.get("moe_checkpoint_method") - ), muon_update_dtype=_first_non_none( row.get("muon_update_dtype"), _trial_muon_update_dtype(trial), @@ -1679,7 +1673,6 @@ def _resolved_run_behavior_point( fsdp_reduce_dtype=_config_fsdp_reduce_dtype(raw_config), ce_mode=_config_ce_mode(raw_config), moe_implementation=_config_str(raw_config, "model", "moe_implementation"), - moe_checkpoint_method=_config_str(raw_config, "train", "moe_checkpoint_method"), muon_momentum=_config_float(raw_config, "train", "muon_momentum"), muon_update_dtype=_config_str(raw_config, "train", "muon_update_dtype"), attention_backend=_config_attention_backend(raw_config), @@ -1837,7 +1830,6 @@ def _standalone_log_behavior_point( fsdp_reduce_dtype=_config_fsdp_reduce_dtype(raw_config), ce_mode=_config_ce_mode(raw_config), moe_implementation=_config_str(raw_config, "model", "moe_implementation"), - moe_checkpoint_method=_config_str(raw_config, "train", "moe_checkpoint_method"), muon_momentum=_config_float(raw_config, "train", "muon_momentum"), muon_update_dtype=_config_str(raw_config, "train", "muon_update_dtype"), attention_backend=_config_attention_backend(raw_config), @@ -1988,7 +1980,6 @@ def _flashqla_summary_point( fsdp_reduce_dtype=_config_fsdp_reduce_dtype(raw_config), ce_mode=_config_ce_mode(raw_config), moe_implementation=_config_str(raw_config, "model", "moe_implementation"), - moe_checkpoint_method=_config_str(raw_config, "train", "moe_checkpoint_method"), muon_momentum=_config_float(raw_config, "train", "muon_momentum"), muon_update_dtype=_config_str(raw_config, "train", "muon_update_dtype"), attention_backend=backend, @@ -2374,11 +2365,6 @@ def behavior_point_workload_mismatches(point: BenchmarkBehaviorPoint, raw_config ("fsdp_reduce_dtype", point.fsdp_reduce_dtype, _config_fsdp_reduce_dtype(raw_config)), ("ce_mode", point.ce_mode, _config_ce_mode(raw_config)), ("moe_implementation", point.moe_implementation, _config_str(raw_config, "model", "moe_implementation")), - ( - "moe_checkpoint_method", - point.moe_checkpoint_method, - _config_str(raw_config, "train", "moe_checkpoint_method"), - ), ("muon_momentum", point.muon_momentum, _config_float(raw_config, "train", "muon_momentum")), ("muon_update_dtype", point.muon_update_dtype, _config_str(raw_config, "train", "muon_update_dtype")), ) diff --git a/src/xorl/sim/calibration_evaluator.py b/src/xorl/sim/calibration_evaluator.py index 3c18fb8f..8aefab3d 100644 --- a/src/xorl/sim/calibration_evaluator.py +++ b/src/xorl/sim/calibration_evaluator.py @@ -188,7 +188,6 @@ def _apply_point_runtime_signature(raw_config: dict[str, Any], point: BenchmarkB _set_if_known(train, "fsdp_reduce_dtype", point.fsdp_reduce_dtype) _set_if_known(train, "ce_mode", point.ce_mode) _set_if_known(model, "moe_implementation", point.moe_implementation) - _set_if_known(train, "moe_checkpoint_method", point.moe_checkpoint_method) _set_if_known(train, "muon_momentum", point.muon_momentum) _set_if_known(train, "muon_update_dtype", point.muon_update_dtype) if isinstance(simulator, dict): diff --git a/src/xorl/sim/kernel_variants.py b/src/xorl/sim/kernel_variants.py index 3c281448..92d76fd3 100644 --- a/src/xorl/sim/kernel_variants.py +++ b/src/xorl/sim/kernel_variants.py @@ -67,31 +67,6 @@ def rank_kernel_variants( } -def compare_kernel_variants( - baseline: KernelVariantMeasurement | dict[str, Any], - candidate: KernelVariantMeasurement | dict[str, Any], -) -> dict[str, Any]: - base = _measurement(baseline) - other = _measurement(candidate) - if (base.family, base.workload) != (other.family, other.workload): - raise ValueError("kernel variants must share one family and one workload") - return { - "family": base.family, - "workload": base.workload, - "baseline": base.variant, - "candidate": other.variant, - "latency_delta_ms": round(other.latency_ms - base.latency_ms, 6), - "latency_delta_percent": round((other.latency_ms / base.latency_ms - 1.0) * 100.0, 6), - "speedup": round(base.latency_ms / other.latency_ms, 6), - "peak_memory_delta_gb": ( - round(other.peak_memory_gb - base.peak_memory_gb, 6) - if base.peak_memory_gb is not None and other.peak_memory_gb is not None - else None - ), - "candidate_promotable": other.promotable, - } - - def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("measurements", type=Path, help="JSON list of kernel-variant measurements") diff --git a/src/xorl/sim/scenario_planner.py b/src/xorl/sim/scenario_planner.py index 62f82911..ffd49249 100644 --- a/src/xorl/sim/scenario_planner.py +++ b/src/xorl/sim/scenario_planner.py @@ -232,7 +232,6 @@ def _parse_int_list(raw: str | None) -> list[int] | None: ("train", "fsdp_reduce_dtype"), ("train", "ce_mode"), ("model", "moe_implementation"), - ("train", "moe_checkpoint_method"), ("train", "muon_momentum"), ("train", "muon_update_dtype"), ("model", "deepep_async_combine"), diff --git a/src/xorl/sim/schemas.py b/src/xorl/sim/schemas.py index 0d011b05..77ae69a5 100644 --- a/src/xorl/sim/schemas.py +++ b/src/xorl/sim/schemas.py @@ -272,7 +272,6 @@ class BenchmarkBehaviorPoint: fsdp_reduce_dtype: str | None = None ce_mode: str | None = None moe_implementation: str | None = None - moe_checkpoint_method: str | None = None muon_momentum: float | None = None muon_update_dtype: str | None = None attention_backend: str | None = None diff --git a/src/xorl/trainers/model_builder.py b/src/xorl/trainers/model_builder.py index 4fad5d8e..44983672 100644 --- a/src/xorl/trainers/model_builder.py +++ b/src/xorl/trainers/model_builder.py @@ -1,8 +1,8 @@ """Shared model build + LoRA/QLoRA injection + FSDP parallelization pipeline. -Both the offline Trainer and the server ModelRunner call build_training_model() -so that every feature (QLoRA, TP, DeepEP, …) is supported in both paths -without reimplementation. +The server ModelRunner calls build_training_model() so that every feature +(QLoRA, TP, DeepEP, …) is supported without reimplementation. The offline Trainer +builds its model directly and shares the individual helpers here instead. """ from dataclasses import dataclass, field @@ -106,6 +106,86 @@ def maybe_upcast_trainable_adapter_params( logger.info_rank0("Upcast trainable LoRA params to float32") +def maybe_unfuse_projections( + model: nn.Module, + *, + unfuse_for_lora: bool, + enable_lora: bool, + enable_qlora: bool, +) -> None: + """Split fused ``qkv_proj`` / ``gate_up_proj`` into per-projection modules. + + LoRA selects modules by name, so a projection the model stores fused is invisible + to it: ``q_proj`` / ``gate_proj`` and friends match nothing and are skipped, + leaving those weights adapted by nothing. Unfusing first turns them into real + modules, which also leaves every parameter name matching the HuggingFace + checkpoint, so weights load with no merge buffer. + + Must run before LoRA injection and before weights are loaded: the underlying + ``unfuse_for_tp`` allocates fresh ``nn.Linear``s without copying from the fused + weight, so the checkpoint is what fills them. + + Args: + model: Model to unfuse, in place. + unfuse_for_lora: Whether to unfuse at all. False is a no-op. Defaults off + wherever it is plumbed: weight sync canonicalizes trainer parameters to + the fused names, so serving an unfused model needs its own verification. + enable_lora: Whether plain LoRA is enabled for this run. + enable_qlora: Whether QLoRA is enabled for this run. + + Raises: + ValueError: If unfusing is requested without plain LoRA, where it would only + cost throughput, or together with QLoRA, which targets the fused names. + NotImplementedError: If the architecture cannot unfuse its projections. + """ + if not unfuse_for_lora: + return + + if enable_qlora: + # inject_qlora_into_model defaults its targets to the fused names + # ("qkv_proj", "gate_up_proj"). On an unfused model those match nothing while + # "o_proj"/"down_proj" still do, so injection succeeds with q/k/v/gate/up left + # unquantized and unadapted, and the packed checkpoint weights are then + # dispatched into plain bf16 parameters. Refuse rather than half-apply. + raise ValueError( + "unfuse_for_lora is not supported with QLoRA: QLoRA targets the fused module " + "names, so unfusing would leave q/k/v and gate/up both unquantized and " + "unadapted. Disable one of the two." + ) + + if not enable_lora: + raise ValueError( + "unfuse_for_lora requires enable_lora: without LoRA it only splits the fused " + "GEMMs and gives up the fused SiLU-and-mul kernel, costing throughput for no " + "benefit." + ) + + if not hasattr(model, "unfuse_for_tp"): + raise NotImplementedError( + f"{type(model).__name__} cannot unfuse its projections, so LoRA would silently " + "skip every fused target. Either implement unfuse_for_tp() on the architecture, " + "or set unfuse_for_lora=false and target the fused names (qkv_proj, gate_up_proj) " + "directly." + ) + + # unfuse_for_tp does two things: module surgery, which is what we want, and a + # config.base_model_tp_plan rewrite, which we do not. The Trainer keeps a reference + # to this config object and save_pretrained's it into every exported checkpoint, so + # leaving the plan behind ships a config.json advertising TP styles that are not HF + # ParallelInterface names. The write is always spurious here: this path requires + # LoRA, and TP + LoRA is rejected outright. + had_plan = "base_model_tp_plan" in model.config.__dict__ + previous_plan = model.config.__dict__.get("base_model_tp_plan") + + model.unfuse_for_tp() + + if had_plan: + model.config.base_model_tp_plan = previous_plan + else: + model.config.__dict__.pop("base_model_tp_plan", None) + logger.info_rank0("Unfused qkv_proj / gate_up_proj so LoRA can adapt them") + + def build_training_model( *, # --- Model --- @@ -132,6 +212,7 @@ def build_training_model( lora_target_modules: Optional[List[str]] = None, lora_target_manifest: Optional[dict[str, Any] | str] = None, moe_hybrid_shared_lora: bool = False, + unfuse_for_lora: bool = False, # --- QLoRA --- enable_qlora: bool = False, block_fp8_qlora_training: bool = False, @@ -208,7 +289,7 @@ def build_training_model( the server ModelRunner. The steps mirror the original Trainer lifecycle:: 1. build_foundation_model() - 2. Unfuse QKV (for TP) + 2. Unfuse projections (QKV for TP, or both for LoRA) 3. QLoRA or LoRA injection 4. LoRA + mixed-precision: upcast trainable params to fp32 5. Exact model-program setup (pre-FSDP2) + save optimizer pre-hook @@ -296,9 +377,19 @@ def build_training_model( helper.print_device_mem_info("VRAM usage after building model") # ------------------------------------------------------------------ - # 2. Unfuse QKV if merge_qkv=False (needed for TP) + # 2. Unfuse projections — for LoRA coverage, or QKV-only for TP # ------------------------------------------------------------------ - if not merge_qkv: + # Ordering matters: both forms must precede LoRA injection (step 3) and weight + # loading (inside step 6), because unfusing allocates fresh Linears without + # copying from the fused weight. + maybe_unfuse_projections( + model, + unfuse_for_lora=unfuse_for_lora, + enable_lora=enable_lora, + enable_qlora=enable_qlora, + ) + + if not merge_qkv and not unfuse_for_lora: for layer in model.model.layers: if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "unfuse_for_tp"): layer.self_attn.unfuse_for_tp() diff --git a/src/xorl/trainers/per_component_timer.py b/src/xorl/trainers/per_component_timer.py index 9f0c1df7..8668806f 100644 --- a/src/xorl/trainers/per_component_timer.py +++ b/src/xorl/trainers/per_component_timer.py @@ -54,7 +54,6 @@ def __init__(self, enabled: bool) -> None: self._fwd_pairs: dict[str, list[tuple[torch.cuda.Event, torch.cuda.Event]]] = defaultdict(list) self._bwd_pairs: dict[str, list[tuple[torch.cuda.Event, torch.cuda.Event]]] = defaultdict(list) self._recompute_pairs: dict[str, list[tuple[torch.cuda.Event, torch.cuda.Event]]] = defaultdict(list) - self.last_skipped_event_pair_count = 0 def attach(self, model: nn.Module) -> int: """Register hooks on decoder layers. Returns the number of layers found.""" @@ -93,7 +92,6 @@ def end_step(self) -> dict[str, float]: if not self.enabled: return {} self._mode = "idle" - self.last_skipped_event_pair_count = 0 torch.cuda.synchronize() result: dict[str, float] = {} @@ -105,7 +103,6 @@ def _accumulate(pairs: dict[str, list[tuple[torch.cuda.Event, torch.cuda.Event]] try: total_ms += start.elapsed_time(end) except (RuntimeError, ValueError): - self.last_skipped_event_pair_count += 1 continue valid_count += 1 if valid_count: diff --git a/src/xorl/trainers/trainer.py b/src/xorl/trainers/trainer.py index 60ff4f1f..a948d81b 100644 --- a/src/xorl/trainers/trainer.py +++ b/src/xorl/trainers/trainer.py @@ -68,6 +68,7 @@ ) from xorl.qlora.utils import _deregister_qlora_weights_from_fsdp from xorl.trainers.model_builder import ( + maybe_unfuse_projections, maybe_upcast_trainable_adapter_params, resolve_training_model_dtype, should_skip_generic_param_upcast, @@ -760,8 +761,16 @@ def _build_model(self) -> None: ) helper.print_device_mem_info("VRAM usage after building model") - # Unfuse QKV for tensor parallelism - if not args.model.merge_qkv: + # Unfuse projections — for LoRA coverage, or QKV-only for tensor parallelism. + # Both must precede LoRA injection below and the weight load in _parallelize. + maybe_unfuse_projections( + self.model, + unfuse_for_lora=args.lora.unfuse_for_lora, + enable_lora=args.lora.enable_lora, + enable_qlora=args.lora.enable_qlora, + ) + + if not args.model.merge_qkv and not args.lora.unfuse_for_lora: for layer in self.model.model.layers: if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "unfuse_for_tp"): layer.self_attn.unfuse_for_tp() diff --git a/src/xorl/utils/manual_cuda_timing.py b/src/xorl/utils/manual_cuda_timing.py deleted file mode 100644 index c9594d1f..00000000 --- a/src/xorl/utils/manual_cuda_timing.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Lightweight CUDA-event scopes for model-specific phase timing. - -This module complements trainer-level phase timing and hook-based component -timing. It lets model code add fine-grained CUDA-event scopes without depending -on the Trainer class. Timings are accumulated per process and drained by the -trainer at the end of a step. -""" - -from __future__ import annotations - -from collections import defaultdict -from contextlib import contextmanager -from typing import Dict, Iterator, List, Tuple - -import torch - -from xorl.utils.device import get_device_type - - -_enabled = False -_mode = "idle" -_pairs: Dict[str, List[Tuple[torch.cuda.Event, torch.cuda.Event]]] = defaultdict(list) - - -def set_manual_cuda_timing_enabled(enabled: bool) -> None: - """Enable or disable manual CUDA-event timing for the current process.""" - global _enabled - _enabled = bool(enabled) and get_device_type() == "cuda" - if not _enabled: - reset_manual_cuda_timing() - - -def set_manual_cuda_timing_mode(mode: str) -> None: - """Set the current timing mode: ``fwd``, ``bwd``/recompute, or ``idle``.""" - global _mode - if mode not in ("fwd", "bwd", "idle"): - raise ValueError(f"invalid manual CUDA timing mode: {mode}") - _mode = mode - - -def reset_manual_cuda_timing() -> None: - """Clear all accumulated events.""" - _pairs.clear() - - -def _phase_name(name: str) -> str | None: - if not _enabled or _mode == "idle" or get_device_type() != "cuda": - return None - if _mode == "fwd": - return f"fwd_{name}" - if _mode == "bwd": - # During activation checkpointing, model forward code runs under - # backward. These scopes therefore describe recompute work. - return f"recompute_{name}" - return None - - -@contextmanager -def manual_cuda_timing_scope(name: str) -> Iterator[None]: - """Record a CUDA-event scope under the current manual timing mode.""" - phase = _phase_name(name) - if phase is None: - yield - return - - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - try: - yield - finally: - end.record() - _pairs[phase].append((start, end)) - - -def drain_manual_cuda_timing(*, synchronize: bool = True) -> Dict[str, float]: - """Return accumulated timings in seconds and clear the accumulator.""" - if not _pairs: - return {} - if synchronize and get_device_type() == "cuda": - torch.cuda.synchronize() - - result: Dict[str, float] = {} - for phase, events in _pairs.items(): - total_ms = 0.0 - valid_events = 0 - for start, end in events: - try: - total_ms += start.elapsed_time(end) - except ValueError as exc: - if "Both events must be recorded" not in str(exc): - raise - continue - valid_events += 1 - if valid_events: - result[phase] = total_ms / 1000.0 - reset_manual_cuda_timing() - return result diff --git a/test_audit_decisions.json b/test_audit_decisions.json new file mode 100644 index 00000000..1a6ad9ba --- /dev/null +++ b/test_audit_decisions.json @@ -0,0 +1,16525 @@ +{ + "schema_version": 1, + "items": [ + { + "id": "TA-001", + "scope": "tests/models/test_qwen3_moe_fused_lora.py", + "decision": "remove", + "status": "applied", + "evidence": [ + "All seven contracts are duplicated or subsumed by tests/models/test_moe_experts_lora.py; two test bodies are exact AST duplicates." + ] + }, + { + "id": "TA-002", + "scope": "30 config snapshot tests in tests/server/test_server_arguments.py", + "decision": "remove", + "status": "applied", + "evidence": [ + "The referenced experiment and example YAML files are not tracked, and a local wrapper converted every missing fixture into a skip." + ] + }, + { + "id": "TA-003", + "scope": "six print-only or soft-threshold measurement tests under tests/ops", + "decision": "remove", + "status": "applied", + "evidence": [ + "The routines only print measurements or turn a missed performance target into pytest.skip, so they cannot report a regression." + ] + }, + { + "id": "TA-004", + "scope": "tests/models/test_qwen3_5_apply_rotary.py::test_interleaved_pairwise_rotation_d8", + "decision": "remove", + "status": "applied", + "evidence": [ + "The only token is at position zero, making the rotation the identity under competing rotation conventions." + ] + }, + { + "id": "TA-005", + "scope": "tests/models/test_moe_sglang_fused_experts.py::test_sglang_runtime_api_does_not_regress_to_legacy_globals", + "decision": "remove", + "status": "applied", + "evidence": [ + "Adjacent fake-runtime behavioral tests already fail if the implementation returns to the unavailable legacy global API." + ] + }, + { + "id": "TA-006", + "scope": "tests/models/test_glm52_official_fp8_inventory.py", + "decision": "relocate", + "status": "applied", + "evidence": [ + "Both checks require an external checkpoint selected by XORL_GLM52_OFFICIAL_MODEL_PATH and are official-model certification rather than clean-checkout tests." + ] + }, + { + "id": "TA-007", + "scope": "tests/distributed/test_deepep_async_combine_guard.py", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The synchronous-default contract is valuable, but the tests patch a removed module global instead of the current environment-controlled function." + ] + }, + { + "id": "TA-008", + "scope": "tests/models/test_qwen3_5_apply_rotary.py::test_qwen35_modeling_does_not_pass_interleaved_to_rotary", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The source parser was replaced by dense and MoE attention projections at nonzero positions that distinguish half-rotate from pairwise rotation while mrope_interleaved is enabled." + ] + }, + { + "id": "TA-009", + "scope": "shared identity, temperature, and KL-tail contracts in the importance-sampling and policy loss tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Three pairs had exact duplicate bodies; one parameterized contract now runs the same implementation-specific cases and assertions for both production losses." + ] + }, + { + "id": "TA-010", + "scope": "distributed launcher wrappers in test_olmo2_tp_e2e.py and test_vocab_parallel_ce.py", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The wrapper bodies match, but each module supplies a different SCRIPT_PATH and validates a distinct distributed production path." + ] + }, + { + "id": "TA-011", + "scope": "tests/server/test_server_arguments.py module stubbing", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "patch.dict(sys.modules, ...) rolled back native Triton and CUTLASS imports loaded inside the context, making later imports fail despite an exactly synced repository venv; targeted restoration of only the stubbed module keys makes all 48 tests pass." + ] + }, + { + "id": "TA-012", + "scope": "three non-empty _prod unit tests in tests/server/weight_sync/test_pp_nccl_transfer.py", + "decision": "remove", + "status": "applied", + "evidence": [ + "The sender metadata and receiver reconstruction tests already exercise one- and two-dimensional products through the production protocol; only the otherwise uncovered empty-shape identity test remains." + ] + }, + { + "id": "TA-013", + "scope": "tests/trainers/test_step_phase_timing.py::test_order_step_phases_covers_every_canonical_phase", + "decision": "remove", + "status": "applied", + "evidence": [ + "The test used _STEP_PHASE_TIMING_ORDER as both its input and expected result; adjacent tests already cover canonical ordering, unknown-key ordering, and empty input." + ] + }, + { + "id": "TA-014", + "scope": "FA3 and external-FLA test module imports", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The locked repository environment provides FA4 but not the optional FA3 top-level interface or external FLA package; three backend-specific modules now skip at collection when those comparison backends are absent." + ] + }, + { + "id": "TA-015", + "scope": "exact Qwen3.5 and GLM numerical-program admission tests in tests/trainers/test_rope_class_b_config.py", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The full resolved-program tests already assert the Class-B and RMSNorm defaults for dense and MoE architectures; separate assertions of those same defaults were redundant.", + "Invalid override values and unsupported topology values now remain as tables inside one admission contract instead of producing a separate collected test for every literal." + ] + }, + { + "id": "TA-016", + "scope": "synthetic and subsumed cases in tests/server/weight_sync/test_p2p_backend_protocol.py", + "decision": "remove", + "status": "applied", + "evidence": [ + "The opt-in fused QKV and convolution source layouts have no matching receiver locator in the production protocol and are explicitly bypassed by default.", + "The 40-layer by 256-expert sweep repeats layer-independent transfer logic already covered by a real-size local shard and every global expert index.", + "The single block-128 FP8 receiver case is a strict subset of the retained multi-receiver layout contract." + ] + }, + { + "id": "TA-017", + "scope": "checkpoint expert-key classification cases in tests/models/test_module_utils_broadcast.py", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sixteen key spellings exercise one binary classifier and now form one table-driven contract with contextual failure messages instead of fourteen independently collected tests." + ] + }, + { + "id": "TA-018", + "scope": "weight-sync configuration precedence, endpoint normalization, and direct-EP selection tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Closely related environment-precedence branches, metadata aliases, and direct-EP sender choices now live in one contract per production decision rather than field-by-field helper fragments." + ] + }, + { + "id": "TA-019", + "scope": "exact GLM fixture in tests/trainers/test_rope_class_b_config.py", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Production validates official geometry and then installs _glm52_exact_contract before resolving numerical defaults; the direct config fixture omitted that admission step and therefore exercised the ordinary-model branch." + ] + }, + { + "id": "TA-020", + "scope": "TileLang V4 indexer shape matrices in tests/ops/dsv4/test_v4_tilelang_indexer.py", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Four retained geometries cover the basic, batched, production-head/top-k, and C128 kernel paths; the removed sequence, head, and batch literals do not select additional code.", + "The removed V4 real-config sweep repeated the production geometry already covered by both the score-reference and top-k contracts." + ] + }, + { + "id": "TA-021", + "scope": "families-v2 RMSNorm realization and dispatch matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forced fused/split parity now covers tail, aligned, and deep-tile hidden shapes at row-count extremes; intermediate row counts do not select different code.", + "For shipped hidden sizes the tile count is below V2_NORM_SPLIT_MIN_TILES, making the dispatch result independent of every parametrized row literal." + ] + }, + { + "id": "TA-022", + "scope": "decode GDN triangular-solve shape and launch-configuration sweeps", + "decision": "remove", + "status": "applied", + "evidence": [ + "Production fixes diagonal group size and both warp counts internally; the 12-case direct-kernel sweep exercised configurations no supported caller can select, while the retained wrapper parity test covers the deployed launch.", + "Scaling K before immediately normalizing it is not a distinct production input regime; retained cases cover batch/grid boundaries and structurally padded input." + ] + }, + { + "id": "TA-023", + "scope": "pipeline stage-to-rank mapping matrix", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One single-stage loop case and one multi-stage case for each loop and V formula cover every implementation branch; additional PP sizes only repeat the same index arithmetic against PyTorch's reference helper." + ] + }, + { + "id": "TA-024", + "scope": "optional-boolean coercion and sequence-parallel FSDP truth tables", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Each table describes one parser or one policy predicate and now reports as one semantic contract while retaining every accepted literal and mode combination." + ] + }, + { + "id": "TA-025", + "scope": "fused selected-logprob dtype, bias, temperature, shape, and vocabulary matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dtype, optional bias, and scalar temperature are independent implementation branches; two pairwise cases cover both values instead of an eight-case Cartesian product in both forward and backward.", + "Irregular shapes and production vocabularies remain fully exercised inside their respective numerical contracts." + ] + }, + { + "id": "TA-026", + "scope": "batch-invariant GEMM table bit-neutrality matrix", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The dtype-by-shape matrix enforces one doctrine-level invariant: every generated launch config preserves the dtype-pinned reduction tree; all twelve combinations remain executed within one contract." + ] + }, + { + "id": "TA-027", + "scope": "TileLang sparse-MLA geometry and attention-sink matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Four forward geometries retain every compiled top-k specialization while covering batch and the production 64-head path; repeated sequence/head literals did not select new code.", + "Zero and mixed-sign sinks cover the arithmetic boundary and both signs, and the retained large-sink effect test independently proves the sink is consumed." + ] + }, + { + "id": "TA-028", + "scope": "LoRA-head gradient-ownership backend and replica matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Pairwise quant-format/backend coverage retains every format, backend label, and eager-versus-fused producer branch without the 3 by 4 Cartesian product.", + "Table rows use explicit monkeypatch contexts so grouping them preserves the isolation previously supplied by pytest parametrization." + ] + }, + { + "id": "TA-029", + "scope": "MoE torch.compile probes and benchmarks in tests/ops/test_moe_torch_compile.py", + "decision": "remove", + "status": "applied", + "evidence": [ + "The fullgraph probe swallowed every exception and therefore could not report a graph-break regression.", + "Two bench-prefixed routines were not collected by pytest, caught compiler failures, and only printed measurements; the retained three contracts enforce block, decoder-layer, and full-model compilation." + ] + }, + { + "id": "TA-030", + "scope": "inference-endpoint quantization normalization helper fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The HF-config detection contract already covers language-model prefix removal, vision filtering, and weight-suffix normalization together; two direct private-helper tests were strict subsets.", + "Static and null activation schemes reach the same unsupported receiver boundary and now share one contract." + ] + }, + { + "id": "TA-031", + "scope": "packing-strategy invariant and datum-order tests", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Best-fit utilization was algebraically equivalent to the retained row-count property and used fewer seeds.", + "The side-array test only verified Python list indexing; datum-order coverage now compares reported order with document lengths recovered from actual packed position boundaries." + ] + }, + { + "id": "TA-032", + "scope": "weight-sync quantization-config aliases and rejection tables", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "No-op aliases, unsupported methods, FP8 formats, and activation schemes each define one normalization or rejection boundary and retain every literal inside a single contract." + ] + }, + { + "id": "TA-033", + "scope": "FP8 default projection selection and exclusion tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One buffer-level contract now covers fused MLA, four packed linear-attention projections, both shared-expert spellings, and negative embedding/gate cases.", + "The removed packed-projection batch test exactly duplicated the four independently parametrized projection cases; module exclusions now cover suffixed and unsuffixed names in one contract." + ] + }, + { + "id": "TA-034", + "scope": "distributed expert-adapter autograd backend, topology, and quantization matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP2 retains real optimizer/autograd qualification for every unquantized backend; the four-rank test adds eFSDP topology with one shipped Quack representative instead of repeating backend math.", + "Quantized execution uses three pairwise backend/format cases rather than a nine-case Cartesian product while retaining every backend and format; DeepEP and projection-subset compositions remain separate contracts." + ] + }, + { + "id": "TA-035", + "scope": "expert QLoRA backend, model-family, and invalid-target contract tables", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Each backend capability, unsupported semantic, fail-closed model family, and invalid target list is one behavioral boundary; every original row remains executed inside its owning contract.", + "Grouping removes pytest item inflation without deleting model-family construction or rejection coverage." + ] + }, + { + "id": "TA-036", + "scope": "MoE-LoRA backend smoke tests and zero-initialized gradient comparisons", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The old cross-backend comparison left every LoRA-B factor at zero, making LoRA-A gradients trivially zero; the retained comparison initializes nonzero adapters and checks both fused backends against eager for output and every factor gradient.", + "That stronger MoEBlock contract subsumes separate Triton/native forward-backward smokes; one explicit zero-delta reference and one explicit nonzero-effect reference retain the shared semantic boundaries.", + "Backend-independent construction and injection retain a Quack representative, while backend registration and numerical execution remain covered separately." + ] + }, + { + "id": "TA-037", + "scope": "registry-wide RoPE fp32-table and native-lane matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Every registered RoPE initializer still proves that model-wide bf16 casting resolves the exact fp32 CPU table; shared forward consumption uses default and YaRN representatives for unit and non-unit attention scaling.", + "Only default RoPE can select the native SGLang cache, so native-versus-stock comparisons for linear, dynamic, YaRN, LongRoPE, and Llama3 executed identical stock code and had no distinct failure mode." + ] + }, + { + "id": "TA-038", + "scope": "FP8 grouped-MoE geometry and end-to-end backend matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Two same-NK cases retain the small and large N launch branches, aligned and padded K, empty experts, multi-block M, and output tails; further sizes selected no new kernel control flow.", + "Wgrad adds an all-empty case for the max_K zero early return and otherwise uses the same branch-complete small/large geometry split.", + "Both grouped backends retain optimizer-step qualification; the dense-plus-MoE integration uses the default Triton-grouped representative instead of repeating that established backend dimension." + ] + }, + { + "id": "TA-039", + "scope": "FP8 linear block-size and quantization-recipe matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Block sizes 64 and 128 create different padding and scale grids, and every existing recipe row remains executed.", + "Each block-scale layout and padded-matmul recipe is one numerical contract rather than a separate pytest item per table row." + ] + }, + { + "id": "TA-040", + "scope": "adapter-optimizer resume identity, topology-corruption, and staged-state matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Every corruption remains because rank identity, layout fingerprint, parameter order, holes, overlaps, dtype, logical shape, optimizer step, group metadata, and state shape reach distinct validation boundaries.", + "Rows now execute under separate temporary checkpoint roots within three transactional contracts, preserving the filesystem isolation previously supplied by parametrization." + ] + }, + { + "id": "TA-041", + "scope": "gradient-checkpointing defaults, method propagation, and outer-gate truth table", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The method-propagation test repeated three strings through the same assignment branch; one nondefault value proves propagation while argument validation owns the supported-value table.", + "The outer-gate contract now includes the method predicate and proves that a nondefault selective method suppresses full-layer checkpointing, a behavior the old matrix omitted.", + "Base and MoE class defaults plus enable-time defaults remain covered together." + ] + }, + { + "id": "TA-042", + "scope": "MoE routing-weight-position resolver truth tables and explicit aliases", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All automatic-regime combinations and explicit boolean/string aliases remain executed; each resolver boundary now reports as one contract rather than one item per truth-table row." + ] + }, + { + "id": "TA-043", + "scope": "GLM-5.2 exact MoE construction dependency, EP, and rank-alpha rejection matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Every exact-component dependency and both uninitialized/wrong-size EP states remain fail-closed before adapter mutation.", + "Separate rank-only and alpha-only invalid cases cover the combined rank=1 and alpha=1 predicate; the removed rank=16, alpha=16 row was their strict conjunction and selected no new validation branch." + ] + }, + { + "id": "TA-044", + "scope": "GLM-5.2 exact shared-expert constructor rejection matrix", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Hidden size, intermediate size, TP width, rank, alpha, bias, and adaptive-noise rejections all remain executed inside one fail-closed constructor contract." + ] + }, + { + "id": "TA-045", + "scope": "NVFP4 and Block-FP8 QLoRA expert-loading fragments and embedded benchmark harness", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Each synthetic checkpoint load already materializes packed bytes, scales, and all three projections; two integrated format contracts now validate those outputs together across every retained geometry.", + "NVFP4 additionally retains exact amax, absorbed global-scale, and dequantized shape/dtype checks while reducing fourteen repeated loads to six total loads across both formats.", + "The private timing function and __main__ print harness were not pytest-collected, had no performance threshold, and were unused outside the file." + ] + }, + { + "id": "TA-046", + "scope": "FP8 external-config recipe and QARL nesting tables", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Every Transformer-Engine-only recipe key and ModelOpt QARL nesting remains fail-closed; each receiver boundary reports as one contract." + ] + }, + { + "id": "TA-047", + "scope": "RMSNorm family funnel geometries and mode-by-batch-invariant matrix", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both head-dimension and hidden-dimension kernel geometries remain bitwise-qualified inside one funnel contract.", + "The module parity test enables the trunk contract in every row, which already fixes family dispatch; the independent batch-invariant flag axis could not select different code and was removed while native, SGLang, and fused-SGLang modes remain." + ] + }, + { + "id": "TA-048", + "scope": "dense and MoE Qwen3.5 RMSNorm call-site truth tables", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Every layer-zero/later-layer and native/SGLang/fused-SGLang row remains executed; each call-site predicate now reports as one contract." + ] + }, + { + "id": "TA-049", + "scope": "OPD chunk-count, KL-estimator alias, streaming-backend, and compiled smoke matrices", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Chunk coverage retains disabled, one-chunk, and more-chunks-than-valid-tokens boundaries; intermediate positive counts only changed a performance option and selected no new loss arithmetic.", + "All VERL estimator aliases and both streaming implementations remain checked inside their semantic contracts.", + "Three direct compiled tuple/shape smokes were strict subsets of retained end-to-end reverse/forward KL reference and diagnostic tests." + ] + }, + { + "id": "TA-050", + "scope": "NVFP4 fake-quant geometry, approximation, and layout contracts", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The implementation has no M- or legal-K-shape branch, so three arbitrary shapes crossed with two dtypes were reduced to one legal geometry executed in both supported dtypes.", + "Exact equality to the independent quantization reference already proves the fixed random tensor changes and meets the weaker relative-error threshold; the E2M1 grid invariant remains separate.", + "The retained K-within-row trap covers the same divisibility rejection as the removed generic 17-by-17 case and additionally proves a legal K succeeds." + ] + }, + { + "id": "TA-051", + "scope": "EP expert-compute registry signature matrix", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Every function in both live EP registries is still inspected for the shared explicit parameters and forward-compatible kwargs contract; backend labels no longer create separate test items." + ] + }, + { + "id": "TA-052", + "scope": "canonical MoE reference widths, transport resolution, and distributed contributor counts", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The reference keeps the two-contributor base tree and production sixteen-contributor depth; widths four and eight only repeated the same adjacent-pair loop.", + "The distributed transport keeps two and eight contributors plus the separate EP16 packed gate; four contributors selected no different collective, mapping, chunking, padding, or backward behavior.", + "The removed internal-resolution test repeated the admitted EP16 and fallback EP8 assertions already present in the adjacent auto-transport contract." + ] + }, + { + "id": "TA-053", + "scope": "GLM-5.2 QLoRA rank-alpha and unsupported-construction matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-only and alpha-only invalid inputs exercise both sides of the rank=1 and alpha=1 predicate; the removed rank=16, alpha=16 row was their strict conjunction.", + "Every unsupported construction mode still creates a fresh meta model and fails before adapterization inside one admission contract." + ] + }, + { + "id": "TA-054", + "scope": "GLM-5.2 exact-attention native-FP8 checkpoint pair transactions", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both arrival orders and both weight/scale members still execute with a fresh checkpoint handler for byte-exact completion, duplicate, missing, dtype, and shape boundaries; each transactional behavior now reports one test item." + ] + }, + { + "id": "TA-055", + "scope": "Quack-versus-Triton EP token-distribution and score-scaling Cartesian product", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Three pairwise cases retain balanced, empty-expert, extreme-skew, score-free, and score-scaled behavior; applying both score states to every distribution selected no additional wrapper or kernel branch.", + "Each retained case still compares output plus input, gate-up, and down gradients, while the independent half-concatenation reference remains." + ] + }, + { + "id": "TA-056", + "scope": "eager-versus-native MoE geometry sweep and large-scale duplicate", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Parity retains ordinary routing, top-k one and four, and the single-token boundary with forward and backward comparisons.", + "The removed larger ordinary geometries selected no repository control flow; the E64/H512/I1024 test repeated the same contract with substantially larger allocation and looser tolerances.", + "The embedded __main__ pytest launcher was uncollected and added no pass/fail contract." + ] + }, + { + "id": "TA-057", + "scope": "batch-invariant full-reduce mean shape-by-dtype matrix", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both supported dtypes still cross the one-dimensional mean_dim path and multi-dimensional sum/divide path; a third tensor rank reached the same multi-dimensional branch." + ] + }, + { + "id": "TA-058", + "scope": "families-v2 RMSNorm split-versus-fused hidden-size and row-count matrix", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shipped and deep hidden sizes remain at low and high row counts with residual, no-residual, and zero-centered modes; the intermediate row count selected neither a dispatch boundary nor distinct arithmetic." + ] + }, + { + "id": "TA-059", + "scope": "linear and MoE LoRA fp32 cast-once merge tests", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Every original dtype still executes zero and nonzero production merge methods for linear and MoE weights inside four behavioral contracts.", + "The removed precision test never called a production merge method; it only compared the local _naive_merge and _fp32_merge helper formulas on one random tensor." + ] + }, + { + "id": "TA-060", + "scope": "EP backend gradient-reduction domain table", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All five supported backend names still fail fast against EP_SUM locally and are rechecked inside the retained real two-rank autograd and synchronization worker; backend labels no longer create separate local test items." + ] + }, + { + "id": "TA-061", + "scope": "GLM-5.2 exact-attention construction rank-alpha and execution-mode matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-only and alpha-only failures cover both sides of the rank=1 and alpha=1 predicate; the removed rank=16, alpha=16 case was their strict conjunction.", + "Every incomplete dense-component, all-to-all, and sparse-MLA requirement still builds a fresh meta model and fails before adapter mutation." + ] + }, + { + "id": "TA-062", + "scope": "GLM-5.2 exact active-LoRA component conjunction", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Each of the five component flags is still independently cleared and checked against active-LoRA, exact-forward, exact-model, and BI-router admission; the conjunction reports one behavioral contract." + ] + }, + { + "id": "TA-063", + "scope": "Mamba2 chunk, packed-boundary, and upstream-divergence tests", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Exact and tail chunk sizes plus aligned and unaligned packed boundaries still execute with forward and gradient parity inside their respective contracts.", + "The removed divergence canary asserted that a known Transformers fallback bug remained present; an upstream fix would have failed XoRL despite improving behavior.", + "The retained independent sequential SSD recurrence remains the authoritative multi-chunk oracle." + ] + }, + { + "id": "TA-064", + "scope": "Qwen3-MoE layer input-norm family mode matrix", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The call-site family depends only on layer zero versus later layers; the capture module stored mode but never read it, so native, SGLang, and fused-SGLang labels could not select different code in this test.", + "Both no-residual and residual-tree layer boundaries remain executed." + ] + }, + { + "id": "TA-065", + "scope": "FlashQLA M-invariance and chunk-chaining certification rows", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both total-M geometries and all three total-length/chaining-step cases still run bitwise output and state comparisons; each certification gate now reports one test item." + ] + }, + { + "id": "TA-066", + "scope": "argument parser low-precision adapter conflicts and optimizer-load spellings", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FP8 and QARL conflicts with both LoRA and QLoRA still parse fresh YAML and fail independently inside one admission contract.", + "Omitted, explicit true, and explicit false optimizer-load values still parse and resolve inside one defaulting contract." + ] + }, + { + "id": "TA-067", + "scope": "EP checkpoint named-mesh rejection and restoration matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Missing EP, missing expert-FSDP, and duplicate EP dimensions still fail independently; both legacy and PP-parent mesh shapes still restore through fresh fake meshes." + ] + }, + { + "id": "TA-068", + "scope": "GLM-5.2 native-FP8 nonofficial configuration fields", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quant method, FP8 format, activation scheme, and block geometry are still independently mutated from the official config and rejected inside one validation contract." + ] + }, + { + "id": "TA-069", + "scope": "streaming forward-KL dense, compiled, chunking, and backend references", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The retained backward parity test already asserts forward values plus student gradients against the independent dense oracle, subsuming the forward-only test.", + "The retained end-to-end OPD dispatch compares both streaming aliases with the compiled backend including loss and gradients, subsuming the direct compiled forward-only unit.", + "Chunk invariance now compares multi-chunk size 7 directly with single-chunk size 40; the removed size-40000 row compared the function to itself and exact-vocab versus over-vocab selected the same single-iteration path." + ] + }, + { + "id": "TA-070", + "scope": "loss reducer denominator and empty-mask matrices", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "TokenPartial additivity is denominator-independent, so one active-count denominator proves the contract while the separate scale-one raw-sum test remains.", + "TokenPartial and SequencePartial zero-denominator behavior both still execute inside one empty-mask contract." + ] + }, + { + "id": "TA-071", + "scope": "QLoRA random-target convergence and ReLoRA comparison experiments", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Five removed convergence tests ran roughly 750 optimizer iterations against random or synthetic targets and asserted only loss decrease, relative ranking, or a loose two-times threshold; those outcomes select no repository branch.", + "Both NVFP4 and Block-FP8 still exercise quantized storage, memory, forward, backward, loading, merge, and requantization mechanisms.", + "The rewritten scheduler integration proves an off-boundary no-op and on-boundary packed-weight mutation, LoRA-B reset, and optimizer-state removal directly after one state-populating step.", + "Optimizer reset tests use one step, which is sufficient for Adam state materialization, while preserving exact LoRA-only and non-LoRA-state assertions." + ] + }, + { + "id": "TA-072", + "scope": "Mooncake hidden-store malformed metadata table", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Missing key, shapes, dtypes, and invalid rank metadata still mutate fresh dictionaries and fail independently inside one parser contract." + ] + }, + { + "id": "TA-073", + "scope": "LoRA target-manifest exact scalar type tables", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Top-level schema/allow-unlisted types and Boolean count/rank values still fail independently against fresh manifests; each validation family reports one contract." + ] + }, + { + "id": "TA-074", + "scope": "Nemotron-H checkpoint invalid expert-parallel configurations", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Indivisible expert count, out-of-range EP rank, and invalid EP size still construct and fail independently inside one admission contract." + ] + }, + { + "id": "TA-075", + "scope": "DeepSeek-V3 unsupported training-mode matrix", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unfrozen router, QLoRA, and unmerged-QKV modes still enter the real training builder and fail with their specific errors inside one supported-mode contract." + ] + }, + { + "id": "TA-076", + "scope": "GLM-5 config, indexer ownership, and TileLang reference coverage", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The default-shape contract now also asserts model_type, subsuming the bare-construction smoke; the end-to-end four-layer forward already verifies indexer ownership on every attention layer.", + "The prior supposed torch reference used a pure causal mask, which is eligible for the same TileLang fast path. The retained GPU contract sets blocked-scoring sizes to force the independent torch implementation before comparing row sets." + ] + }, + { + "id": "TA-077", + "scope": "merged LoRA shared-factor straight-through autograd matrix", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unshared, shared-A, and shared-B executions and their independent autograd comparisons remain inside one factor-sharing contract rather than three collected parameter rows." + ] + }, + { + "id": "TA-078", + "scope": "training simulator built-in calibration-pack validation", + "decision": "remove", + "status": "applied", + "evidence": [ + "The consolidated validator already discovers all built-in packs, validates schema and sanitation, checks behavior-point counts, and asserts exact raw and promotable golden throughput.", + "Pack names and report schema version were added to the consolidated assertion before removing the weaker sanitation loop and Qwen3.5 winner-only test." + ] + }, + { + "id": "TA-079", + "scope": "DeepEP internode topology and preflight no-op matrices", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All seven uninitialized, malformed, single-node, contiguous, and strided topology scenarios still execute in one truth table.", + "Environment-disabled, uninitialized-distributed, and intranode preflight exits still independently prove that no DeepEP buffer is created." + ] + }, + { + "id": "TA-080", + "scope": "sequential packer artificial types, batch scale, schema whitelist, and roundtrip repetitions", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The removed ABC/custom-subclass test exercised a test-only implementation rather than production packing, and the exact optional-key whitelist failed on harmless schema extension.", + "The arbitrary 100-sample workload selected no new branch; numpy normalization remains. The roundtrip retains one exact sample-boundary integration while repeated multi-batch and shift-mode checks remain in their dedicated contracts." + ] + }, + { + "id": "TA-081", + "scope": "session base-model cache-path canonicalization rows", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Repository-id identity and both cache-path directions still execute inside one canonicalization contract; distinct models and pass-through values remain separately asserted." + ] + }, + { + "id": "TA-082", + "scope": "launcher tests for removed pre-refactor APIs", + "decision": "remove", + "status": "applied", + "evidence": [ + "Thirteen tests were permanently skip-gated because the launcher refactor removed local-master detection, worker override forwarding, and init-time override validation.", + "The retained seven tests cover the live remote and explicit-host address paths, readiness success/failure, current command behavior, parser behavior, and the removed-ZORL migration error." + ] + }, + { + "id": "TA-083", + "scope": "model-runner token-diagnostic shape, ranking, and disabled-input checks", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Valid positions, target IDs, output shapes, top-k width, and target-versus-top1 consistency now share one output contract.", + "Zero top-k and absent labels still independently execute the same disabled diagnostic boundary inside one test." + ] + }, + { + "id": "TA-084", + "scope": "checkpoint metadata production and compatibility consumption", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "QARL buffer metadata is now written, inspected, and consumed by strict and non-strict compatibility checks in one transaction.", + "Pipeline parameter and buffer unions are now written to disk and validated from that artifact instead of testing writer and reader against separately fabricated metadata." + ] + }, + { + "id": "TA-085", + "scope": "Muon restart helper, Quack tuned mode, and backend dtype selection", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The helper-only restart assertion was removed because the retained real optimizer step invokes autotuning and asserts the chosen reset iteration.", + "Untuned and tuned Quack calls and FP32-versus-BF16 SM90 backend selection still execute as complete truth tables." + ] + }, + { + "id": "TA-086", + "scope": "cautious-weight-decay optimizer factory routing", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "SignSGD, explicit AnyPrecisionAdamW, and AdamW-to-AnyPrecision routing still build fresh optimizers and validate cautious flags; AdamW also retains its FP32 momentum assertion." + ] + }, + { + "id": "TA-087", + "scope": "API server constructor, Pydantic assignment smokes, and heartbeat timing", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Constructor field echo and request-object assignment tests were removed; retained endpoint tests exercise aliases, defaults, serialization, registration, and optimizer payloads at the application boundary.", + "Heartbeat activity is now advanced by a deterministic method stub instead of a real sleep, while still proving the endpoint invokes session refresh." + ] + }, + { + "id": "TA-088", + "scope": "trainer gradient clipping, sign-vote scaling, and target-token preference tables", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Regular and DistSign clipping, both nonpositive disable values, and positive/zero/negative voter totals retain every execution inside three behavioral tables.", + "Token and active-microbatch counters both prove target_tokens precedence in one shared caller contract." + ] + }, + { + "id": "TA-089", + "scope": "router diagnostic tie-policy aliases and invalid-policy boundary", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "stable_low_id, tie_low_id, and tie_high_id still construct fresh routers and validate selected experts and gathered weights inside one policy contract.", + "The invalid-policy test now expects the actual fail-closed construction boundary rather than incorrectly constructing outside the exception assertion." + ] + }, + { + "id": "TA-090", + "scope": "teacher activation-cache host/device rank-3 gathers and index bounds", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Host and cached-device token-axis and layer-slice paths still execute with their original shapes and dtype conversion inside two contracts.", + "Negative and upper-bound index failures remain independently asserted as one bounds contract." + ] + }, + { + "id": "TA-091", + "scope": "DeepSeek-V4 successful checkpoint name mappings", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Top-level, attention, indexer, norm, HC, router-bias, and shared-expert mappings all remain exact assertions in one successful-map contract; unknown and MTP names retain their separate rejection contract." + ] + }, + { + "id": "TA-092", + "scope": "orchestrator model-pass results and sequential-operation smokes", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Backend timing-field preservation now runs inside the retained forward/backward and forward-only operation contract.", + "Five identical successful sequential forwards were removed from the error test; exact operation-counter behavior remains in the dedicated statistics contract." + ] + }, + { + "id": "TA-093", + "scope": "runner-dispatch routing expert and logit slicing", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained shard-and-slice test already uses the same datum offset/count and asserts both routed expert IDs and routed logits plus removal of slicing metadata; the second test repeated that contract." + ] + }, + { + "id": "TA-094", + "scope": "pipeline profiling interval and analytic bubble helpers", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Empty, degenerate, disjoint, overlapping, contained, touching, and unsorted intervals all remain in one union truth table.", + "1F1B, GPipe, interleaved, zero-bubble, PP1, and invalid schedule branches all remain, grouped by nonzero formula, zero result, and rejection behavior." + ] + }, + { + "id": "TA-095", + "scope": "OPD driver prompt chunking and student weight-version verification", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Order/tail preservation and nonpositive chunk rejection remain in one chunking contract.", + "Matching, mismatched, and model-info-error weight versions still execute with their exact profile-row side effects inside one verifier truth table." + ] + }, + { + "id": "TA-096", + "scope": "sparse-delta template traversal and malformed update validation", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed filename and manifest filename traversal are still rejected against fresh capture state inside one trust-root contract.", + "Duplicate, floating-point, length-mismatched, and out-of-range indices still enter the real writer and fail with their specific errors inside one receiver contract." + ] + }, + { + "id": "TA-097", + "scope": "EP-aware gradient clipping classification, norm arithmetic, and mixed-mesh smokes", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The retained classify-then-clip contracts already prove skip-FSDP and ordinary parameter grouping, combined L2 clipping, unchanged below-threshold gradients, and absence of EP double division; separate 3-4-5 arithmetic units were strict subsets.", + "Infinity norm, missing gradients, shared replicas, dispatch, mixed DTensor meshes, explicit foreach behavior, and live two- and three-rank reductions remain independently exercised." + ] + }, + { + "id": "TA-098", + "scope": "DistSignSGD update, state, reduce-scatter, and topology micro-tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The weight-decay update covers the ordinary preaggregated update in the same production step, and the state-dict roundtrip proves the optimizer remains state-free after stepping.", + "Signing and forced-SUM behavior now use one AVG-input communication contract, while HSDP, folded sequence parallelism, and EP still fail against fresh models inside one topology table." + ] + }, + { + "id": "TA-099", + "scope": "constant, linear, and cosine learning-rate schedule examples", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Warmup-to-constant, full linear decay, warmup-decay-floor, cosine floor, and warmup-cosine contracts cover every schedule phase; standalone constant and floor examples repeated phases already asserted by those traces.", + "Every invalid learning rate, warmup ratio, and decay style still enters its production validation branch inside one configuration contract." + ] + }, + { + "id": "TA-100", + "scope": "NVFP4 QARL normalization, fake-quant helpers, dense wrappers, and MoE wrappers", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Non-16 group sizes and target-module directions still run as truth tables, while dense and eager-MoE forward contracts now combine configuration, lossy output, parameter restoration, and gradients.", + "Private MoE fake-quant and shadow tests were removed because exact 3D forward and STE arithmetic belongs to the retained op suite; the wrapper contract proves those helpers feed the inherited production forward.", + "The independent NVFP4 reference subsumes a looser grid property, the linear STE contract subsumes the direct identity test, and registry dispatch subsumes the supported-format boolean smoke." + ] + }, + { + "id": "TA-101", + "scope": "FP8 model-builder tensor-parallel lm-head inclusion examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full FP8 and lm-head-excluded tensor-parallel builds still construct fresh models and validate both projection types inside one inclusion contract." + ] + }, + { + "id": "TA-102", + "scope": "QARL activation-quantization override exit and target examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Enabled and disabled overrides, distinct per-module restoration, exception cleanup, and exclusion of ordinary Linear modules now execute as one state transaction; nested restoration remains a separate reentrancy contract." + ] + }, + { + "id": "TA-103", + "scope": "QARL W4A4 activation STE and MoE backend-shadow helper fragments", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The weighted activation-gradient contract proves exact STE identity with nonuniform upstream gradients, subsuming the all-ones sum example.", + "The exception path proves triton_w4a4 selection and restoration together; activation-off and non-Triton no-op cases still execute in one conjunction-boundary table." + ] + }, + { + "id": "TA-104", + "scope": "QARL synthetic convergence smoke and duplicate sync-configuration failure", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The 16-step loss-decrease heuristic is replaced by one real optimizer step that asserts finite loss, both QARL gradients, parameter mutation, changed logprobs, persistent summary state, and exact checkpoint restoration.", + "The weight-sync handler mismatch contract already enters the same block-size validator and additionally proves the user-facing failure response, subsuming the direct mismatch unit." + ] + }, + { + "id": "TA-105", + "scope": "SignSGD and AnyPrecisionAdamW all-aligned cautious-decay comparisons", + "decision": "remove", + "status": "applied", + "evidence": [ + "Each retained mixed-coordinate production step contains both aligned and misaligned coordinates and checks the exact resulting update; the all-aligned comparisons selected no additional optimizer branch." + ] + }, + { + "id": "TA-106", + "scope": "GLM-5.2 IndexShare identity through the test-local tensor mapper", + "decision": "remove", + "status": "applied", + "evidence": [ + "The removed test only proved that a recursive helper defined in the test file leaves non-tensor Python objects unchanged.", + "The retained dense-producer/shared-consumer model forward runs the same mapper in simulated FSDP pre-hooks and proves one context identity across every layer plus lifecycle cleanup." + ] + }, + { + "id": "TA-107", + "scope": "GLM-5 indexer shape, additive-mask, chunked-head, and padding-mask examples", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Shape, dtype, range, final-row validity, and sorted-sentinel ordering now live in one selector output contract.", + "Dense and one-head-chunk scoring both retain the diagonal-only additive-mask result inside one behavioral table.", + "The padding-mask test now runs on CPU and accurately tests prefix acceptance versus interior-hole rejection instead of claiming GPU fast-path execution it never invoked." + ] + }, + { + "id": "TA-108", + "scope": "GLM-5 sparse-attention output-shape smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained full-model sparse-versus-dense contract executes sparse attention and checks numerical output, while the Ulysses integration checks local query, full KV, top-k, mask, offset, and output shapes; the standalone shape-only forward was dominated." + ] + }, + { + "id": "TA-109", + "scope": "exact Qwen3.5 dense and MoE numerical-program and admitted-topology successes", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense and MoE configs still resolve the complete certified numerical program and CE mode from fresh configs inside one family contract.", + "World16 HSDP-plus-EP, world8 EP, and single-GPU dense topologies still enter the real admission validator inside one successful-topology table." + ] + }, + { + "id": "TA-110", + "scope": "exact Qwen3.5 model-scope accepted snapshots and nearby-geometry rejection rows", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense, MoE, and Hugging Face outer-config snapshots still validate from fresh objects in one accepted-scope contract.", + "Both dense and MoE hidden-size near misses still fail independently inside one rejection contract rather than generating parameterized item inflation." + ] + }, + { + "id": "TA-111", + "scope": "weight-sync adapter preparation, bucket helpers, tied aliases, MoE prefix mapping, and protocol field echo", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Requested-adapter materialization and current-adapter fallback still run in one stateful contract, and chunking plus cap decisions now share one bucket contract.", + "The retained prior-module tied-weight test includes the removed root extraction and alias assertions before proving the duplicate is skipped on the subsequent module.", + "Nemotron-H and ordinary-MoE unfuse transactions already prove prefix remapping at the produced tensor boundary; the direct helper smoke was dominated.", + "Sparse-delta protocol fields remain exercised through remote-backend and end-to-end sync requests, so the local Pydantic assignment echo was removed." + ] + }, + { + "id": "TA-112", + "scope": "sparse-delta path fast-path baseline, config, and FP8 cache-metadata harnesses", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Baseline post-only ordering and accounting, explicit baseline configuration, and FP8 KV-cache postprocess metadata still execute independently through one shared fake transport transaction.", + "The rewrite removes two copied endpoint/backend harnesses while preserving backend config, normalized cache epoch, endpoint results, pause/resume order, posted paths, and weight version assertions." + ] + }, + { + "id": "TA-113", + "scope": "adapter-manager ownership, optimizer, lifecycle, eviction, and checkpoint transactions", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The 48 tests cover distinct compile, capture, staging, atomic commit, abort, clipping, collective, poisoning, publication, trust-root, session-spec, rollback, eviction, mixed-rank, and optimizer-state boundaries rather than literal or shape variations.", + "All 48 pass in the source-tree environment." + ] + }, + { + "id": "TA-114", + "scope": "MiniMax-M3 expert-key aliases and DeepSeek-V4 window-only forward smoke", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "All three MiniMax language-model prefix and w1/w2/w3 expert-key aliases still execute inside one classifier contract.", + "The removed DSv4 base-model C0 shape smoke was dominated by the retained full C0 causal-LM forward/backward, which additionally checks logits, loss, required gradients, and intentionally frozen hyperconnection parameters; the distinct C128 forward remains." + ] + }, + { + "id": "TA-115", + "scope": "RMSNorm cross-engine shape/family matrices and fused-kernel dtype rows", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Every adversarial cross-engine shape, both family funnels, both residual modes, trunk lane, zero-centered twin, and families-v2 candidate execution remains; each numerical invariant now reports one test item.", + "BF16 and FP32 residual and no-residual fused-kernel comparisons still run inside two dtype-complete contracts, both of which pass locally.", + "The SGLang cross-engine module is dependency-skipped in the current venv, so its grouped executions were linted and collection-checked but could not run locally." + ] + }, + { + "id": "TA-116", + "scope": "endpoint-manager and NCCL routing examples repeated across the same control transaction", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The fallback health contract already invokes and verifies the primary host-and-port URL before the v1/models fallback, so the primary-success URL-only example was dominated.", + "Configured direct load-format forwarding now lives in the existing bucket endpoint-routing contract instead of repeating its complete synchronizer and broadcast harness.", + "The hybrid receiver-fence contract now proves deferred work lifetime and release on NCCL-group destruction in one transaction; the copied destruction harness was removed." + ] + }, + { + "id": "TA-117", + "scope": "sparse-delta factory, replicated-path, and initialization-policy micro-tests", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The backend-factory isinstance check only mirrored a single factory branch and was removed.", + "Single-path TP replication remains asserted by streaming transfer, while the retained prepacked-path contract owns distinct per-rank paths and now checks unique-file accounting.", + "Post-only import avoidance, streaming rejection under prepacked-only, and the valid prepacked post-only combination all still execute inside one initialization-policy table.", + "The fake unresolvable hostname was replaced with a validated loopback literal so mocked HTTP tests also pass through production URL safety checks deterministically." + ] + }, + { + "id": "TA-118", + "scope": "trainer P2P IB-device selection examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Global-rank, local-rank, single-device fallback, explicit physical-GPU selection, numeric CUDA visibility, and empty-entry autodiscovery all still execute in an environment-isolated precedence table.", + "Seven mapping outcomes now report one behavior contract rather than six separately named examples." + ] + }, + { + "id": "TA-119", + "scope": "private P2P rank-summary dictionary assembly and three-counter addition examples", + "decision": "remove", + "status": "applied", + "evidence": [ + "The removed tests asserted private dictionary field copies and the arithmetic sum of byte, parameter, and bucket counters without exercising a transfer, collective failure, or user-visible result.", + "Abort-marker lifecycle and distributed peer-failure gathering remain as separate operational contracts." + ] + }, + { + "id": "TA-120", + "scope": "server removed-field shapes and shipped MoE adapter configuration rows", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Flat adapter-ownership, nested removed ZORL field, and removed ZORL-section payloads still enter the real server loader and assert their distinct failures inside one rejection table.", + "Both shipped MoE LoRA and all five shipped Qwen MoE QLoRA configurations still parse in a clean subprocess and compare their source and normalized Quack, target-module, and shared-LoRA values; parametrized item inflation was removed." + ] + }, + { + "id": "TA-121", + "scope": "independent server optimizer, resume, prefetch, HSDP, packing, activation, and adapter-state field tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "SignSGD and DistSignSGD each parse through a composed nested config that also checks checkpoint optimizer policy, forward/backward prefetch, HSDP deferral, packing alignment, activation memory limit, and adapter state-load mode.", + "Both the runtime object and its serialized model, train, and LoRA configs remain asserted at the server boundary." + ] + }, + { + "id": "TA-122", + "scope": "server R3 payload success modes and MoE routing-weight default examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit Mooncake transport, the legacy externalization alias, and explicit filesystem fallback still parse and serialize every transport-specific field inside one mode table; the invalid directory-without-filesystem failure stays separate.", + "Explicit routing-before-down and the automatic default both still execute through one defaulting contract." + ] + }, + { + "id": "TA-123", + "scope": "training YAML optimizer, packing, numerical-alignment, and FSDP scalar acceptance tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both SignSGD values still pass the real CLI parser, while multipack shape, model numerical-alignment flags, FSDP reduction dtype, and parameter-upcast policy are checked together from the same production-shaped YAML.", + "Muon, legacy alias transforms, FP8/QARL configuration, automatic checkpoint resolution, load-optimizer defaults, and all incompatibility failures remain independent contracts." + ] + }, + { + "id": "TA-124", + "scope": "GatedDeltaNet exact-convolution unsupported-mode examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Decode-cache, context-parallel, missing-short-convolution, and convolution-bias failures still construct and invoke their distinct production paths inside one guard table.", + "Forward/backward references, kernel and end-to-end determinism, state scoping, checkpoint recompute, packing order, and SGLang parity remain independent contracts." + ] + }, + { + "id": "TA-125", + "scope": "batch-invariant trunk wrapper type guards and bias parameter rows", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "No matching projection, ordinary LoRA, custom Linear subclass, and FP16 weight failures all still execute inside one wrapper-admission table.", + "Biased and bias-free forward and backward bitwise comparisons still run against their independent persistent-GEMM and cuBLAS references, but no longer inflate pytest items." + ] + }, + { + "id": "TA-126", + "scope": "batch-invariant trunk global-state and duplicate RMSNorm loud-failure assertions", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The autouse fixture now establishes disabled global state before as well as after each contract, preventing order-dependent inheritance.", + "The selection contract now correctly asserts that wrapping arms RMSNorm dispatch, matching the implementation contract; its prior assertion required the opposite behavior.", + "The standalone RMSNorm loud-failure regression was removed because the retained multi-op grad-requiring interpose contract already invokes RMSNorm and requires the same failure." + ] + }, + { + "id": "TA-127", + "scope": "attention registry presence, flash-family detection, and resolver examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The FA4-only reload transaction now directly proves flash_attention_2, flash_attention_3, and flash_attention_4 registration plus mask-family detection, subsuming a standalone equality with the registry-membership expression used by the implementation.", + "Registered eager/native resolution, non-flash eager fallback, and unavailable-flash rejection now form one resolver-boundary contract." + ] + }, + { + "id": "TA-128", + "scope": "quantized-export example snapshot and unsupported-source preflight fragments", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The checked-in example's literal field snapshot was removed because generic config parsing and the retained subprocess CLI export already validate the parser and consumer transaction.", + "Existing FP8 scales, MTP config metadata, MTP tensor namespaces, and unfolded QARL state still build real source directories and fail through one export-preflight table." + ] + }, + { + "id": "TA-129", + "scope": "QARL export eight-step synthetic training loop", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "One real AdamW step now establishes finite loss and changed target logprobs before export, replacing an arbitrary eight-iteration mini-training experiment.", + "The retained contract still requires exact target-logprob equality after folding, block-FP8 export, dequantization, and reload." + ] + }, + { + "id": "TA-130", + "scope": "API request removed-configuration parameter rows", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Create-model ZORL, nested adapter-ownership, and create-session ZORL inputs still enter their actual Pydantic request types and assert both field path and migration message inside one rejection table.", + "Unknown rolling-client and nested LoRA fields remain covered by the separate compatibility contract." + ] + }, + { + "id": "TA-131", + "scope": "TensorData rank-specific to_plain_dict examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-one model and loss fields, valid rank-two and rank-three nesting, mismatched shape fallback, and empty higher-rank fallback all still execute in one conversion contract.", + "The retained assertions preserve exact nested values and therefore the sequence-field classification invariant consumed by packing." + ] + }, + { + "id": "TA-132", + "scope": "SGLang fused-expert DeepEP exclusion and pair-slot helper duplicates", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained DeepEP exclusion contract now asserts the all-to-all path plus the precise unverifiable order-and-rounding mechanism, subsuming a second test that only required the same NotImplementedError.", + "The full slot-combine path independently checks slot-ordered weighted reduction from routed inputs, so a direct private pair-order helper example was dominated." + ] + }, + { + "id": "TA-133", + "scope": "API optimizer current/legacy payloads and sampler tracking model-ID examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Current learning-rate/gradient-clip fields and legacy Adam parameter aliases both still reach the orchestrator payload and response metrics in one optimizer transaction.", + "Embedded xorl URI model IDs and explicit request model IDs both still load and track their actual sampler paths in one cleanup-ownership contract." + ] + }, + { + "id": "TA-134", + "scope": "base-model canonicalization equality and inequality examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Repository IDs and Hugging Face cache paths still normalize in both directions, while distinct repositories, ordinary paths, and None remain distinct or unchanged inside one canonicalization contract." + ] + }, + { + "id": "TA-135", + "scope": "FP8 E2E Ulysses and hybrid context-parallel shape ladder", + "decision": "remove", + "status": "applied", + "evidence": [ + "The short Ulysses case is dominated by the retained longer packed Ulysses transaction, which exercises the same FP8 and Ulysses branches at a stronger shape.", + "Three intermediate hybrid context datasets and the basic hybrid case add only sequence length or sample-shape variation; the retained 4096-token long-tail multipack transaction exercises the same Ulysses-plus-Ring composition with heterogeneous near-full bins." + ] + }, + { + "id": "TA-136", + "scope": "FP8 MoE DeepEP checkpoint-resume cross-product", + "decision": "remove", + "status": "applied", + "evidence": [ + "Dense FP8 checkpoint save/resume retains FP8 serialization and optimizer restoration, while the retained DeepEP EP/eFSDP transaction proves FP8 expert compute through that distributed topology.", + "Combining both mechanisms in a second two-phase four-GPU run selected no additional checkpoint or DeepEP branch." + ] + }, + { + "id": "TA-137", + "scope": "retained FP8 E2E metric oracle and configuration generator", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "A live baseline completed two optimizer steps with eight of nine linears using FP8; the sole unused module was lm_head, which Qwen3's resolved numerical program intentionally executes in FP32.", + "The oracle now requires every eligible linear to use FP8 while naming only lm_head as the canonical FP32 exception, and the retained baseline passes.", + "The shared E2E generator now applies the extra_data and extra_model mappings already supplied by retained packed-context and DeepEP cases instead of failing at Python argument binding." + ] + }, + { + "id": "TA-138", + "scope": "GLM-5 architecture registry and loader-selection snapshots", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The retained local Hugging Face config load now proves both architecture aliases, their class relationship, and the selected GLM loader after constructing the actual Glm5Config.", + "Two standalone registry-membership and loader-description tests added no execution boundary beyond that transaction." + ] + }, + { + "id": "TA-139", + "scope": "GLM-5 checkpoint layer normalization and early-skip key examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Configured, MTP-boundary, far-out-of-range, and non-layer keys still pass through both the normalizer and loader early-skip hook in one routing contract.", + "The separate tests repeated the same four key classifications against two surfaces of the same handler." + ] + }, + { + "id": "TA-140", + "scope": "OPD hidden-cache unpacked filtering and payload-merge helper examples", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained gathered-SP writer filters an interior valid target from the full hidden tensor and asserts its persisted cache index, subsuming the direct unpacked row-split helper example.", + "The retained multi-rank writer gathers local and remote chunks, orders them by logical slice, persists the concatenated tensor, and asserts per-sample indices, subsuming the direct payload-merge helper example." + ] + }, + { + "id": "TA-141", + "scope": "P2P generic source-slicing helper examples", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained full transfer transaction writes exact TP row slices to two receiver pointers, while the retained incompatible-shape transaction rejects transfer before the engine is called.", + "Unsliced tensors flow through many retained transfer transactions, so direct private-helper examples for TP rows, full-shape mismatch, and identity return added no independent boundary; specialized Qwen linear-attention transformations remain." + ] + }, + { + "id": "TA-142", + "scope": "P2P initialization failure, sender capability, and scatter-copy mode fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "HTTP and successful-response remote failures still execute independently inside one initialization failure contract.", + "Implicit all-rank and explicit sender sets now share one capability contract, and list, deep, and forced-reuse locator alias policies share one scatter-copy contract.", + "Repeated dense-owner and filtered-buffer assertions inside two retained tests were removed without dropping any value or branch." + ] + }, + { + "id": "TA-143", + "scope": "P2P transfer metadata and direct-EP private-helper smokes", + "decision": "remove", + "status": "applied", + "evidence": [ + "Flush preservation and weight-version propagation now reach the real completion payload in one transfer transaction instead of stopping at backend configuration.", + "Multi-sender nonzero-rank initialization already adopts and validates scattered tensor maps, subsuming a direct state-assignment helper example.", + "The retained all-filtered direct-EP transfer uses a nonzero source rank with a real locator and proves the engine stays untouched, subsuming the empty-bucket no-exception smoke." + ] + }, + { + "id": "TA-144", + "scope": "inference worker-port registration and adapter URL examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The explicit-worker-port registration transaction now checks control and worker health, then uses the registered endpoint for LoRA load, unload, and loaded-adapter discovery.", + "Two standalone adapter URL tests constructed an endpoint by hand and repeated only the final port choice." + ] + }, + { + "id": "TA-145", + "scope": "inference sync-quantization normalization and unsupported-receiver examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The rich FP8 skip-list detection contract already proves default format, activation scheme, block size, language-model normalization, and vision exclusion, subsuming a minimal default-dictionary snapshot.", + "MTP, invalid activation schemes, UE8M0 scale storage, and BF16 MTP now form one receiver-admission policy contract, while compressed-tensors rejection and explicit BF16 no-op normalization share the setter policy boundary." + ] + }, + { + "id": "TA-146", + "scope": "packing strategy literal snapshot and convenience-wrapper smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "Every supported strategy still executes through document, token, position, capacity, balance, order, and determinism contracts, so an exact tuple-literal assertion adds no behavior.", + "The retained full pipeline calls pack_samples, asserts request and packed boundaries, simulates output, and unpacks every sample; direct packed and unpacked contracts separately cover the wrapper's two mode branches." + ] + }, + { + "id": "TA-147", + "scope": "pipeline schedule style, capability, and admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Every schedule now maps to its stage style, single-stage classification, and split-backward policy in one metadata table.", + "All five admitted configurations and every invalid virtual-stage or microbatch branch execute inside one schedule-admission contract." + ] + }, + { + "id": "TA-148", + "scope": "DeepEP preflight and launcher readiness success-only tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One DeepEP roundtrip contract now accepts identity combine and rejects corrupt combine after the same dispatch setup.", + "One launcher readiness contract now accepts a set ready event and independently fails fast when the worker exits." + ] + }, + { + "id": "TA-149", + "scope": "Blackwell FP8 and exact-Qwen MoE admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Blackwell FP8 now rejects no override, rejects an override without a validation artifact, and accepts the explicit validated override in one policy contract.", + "Exact Qwen3.5 MoE defaults and every noncertified implementation, dispatch, or async-combine override now enter one architecture admission boundary." + ] + }, + { + "id": "TA-150", + "scope": "factor-only exact active-LoRA snapshot guard scope", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The retained snapshot contract rejects both full-weight publication paths before downstream work, then proves the same guard leaves an ordinary model unrestricted." + ] + }, + { + "id": "TA-151", + "scope": "FP8 sync output-shape snapshot and contiguous-slice predicate example", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained independent Slime-reference contract now also asserts CPU placement while already proving emitted names, dtypes, scale shape, exact FP8 bytes, exact scales, and dequantized parity, subsuming the weaker output snapshot.", + "Stack-versus-single quantization proves grouping is numerically transparent, and workspace/streaming transactions exercise grouped expert stacks; a direct storage-offset predicate example asserted no user-visible correctness boundary." + ] + }, + { + "id": "TA-152", + "scope": "FP8 GPU stack target-device and CPU-parity fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One CUDA contract now quantizes the same nontrivial stack to CPU and CUDA targets, checks target-specific copy telemetry, and compares both outputs bitwise against the CPU path.", + "Three separate tests previously rebuilt the stack and reported placement and parity independently." + ] + }, + { + "id": "TA-153", + "scope": "two-dimensional DTensor save materialization process launches", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One four-rank 2D CPU-mesh transaction now proves all-rank materialization and writer-only materialization from the same sharded tensor.", + "The separate writer-rank test repeated process-group setup, mesh construction, shards, and full-tensor reconstruction; the distinct one-dimensional mesh contract remains." + ] + }, + { + "id": "TA-154", + "scope": "grouped-GEMM scale, single-group, and property examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The roughly one-gigabyte aligned operand exercised the same masking and autotune shape family as smaller examples; a compact unaligned FP16 case now reaches the previously uncovered K/N masking path.", + "Single-group execution has no separate implementation branch, while retained numerical contracts already cover FP16, BF16, unequal groups, transpose-B, zero-K, contiguity, and device admission." + ] + }, + { + "id": "TA-155", + "scope": "MoE primitive example ladders and non-gated constructor failures", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Repeated flattened histograms, an invalid overlapping slot map, and separate gather/scatter/add-gather smokes were dominated by retained independent references and a permutation roundtrip.", + "The multi-block gather now uses a compact unaligned hidden width, the full pipeline asserts its exact result, and unsupported backend and activation policies share one constructor contract." + ] + }, + { + "id": "TA-156", + "scope": "adapter-manager checkpoint path, rollback, dtype, and learning-rate fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Save-root and load-root rejection now execute in one containment policy transaction.", + "The missing-tensor failure now also proves fresh-adapter rollback, and the current-learning-rate checkpoint transaction also verifies persisted LoRA tensor dtypes." + ] + }, + { + "id": "TA-157", + "scope": "adapter-coordinator path containment and fresh-state broadcasts", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Load paths, save paths, and evicted model-id traversal now enter one output-root policy contract.", + "Direct adapter registration and materialized session registration now prove their ordered broadcasts through one coordinator transaction." + ] + }, + { + "id": "TA-158", + "scope": "routing replay sequence-length and already-decoded input examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One sequence-parallel transaction now covers both exact-length slicing and route padding at the actual position length.", + "Python lists, NumPy arrays, and tensors share one materialized-input contract, while nested Qwen top-k discovery is exercised by the retained raw-base64 shape-inference path." + ] + }, + { + "id": "TA-159", + "scope": "server Adam hyperparameter argument, initialization, step, and dispatcher fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit and default arguments share one configuration contract, while explicit, default, and malformed optimizer initialization share one admission transaction.", + "Full, partial, omitted, and non-Adam step policies execute together; adapter-manager and dispatcher propagation remain independent boundaries." + ] + }, + { + "id": "TA-160", + "scope": "runner dispatcher save variants and rank-parameterized failure policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Save-state and save-LoRA handlers now prove the same nonresident-adapter checkpoint requirement through one dispatcher transaction.", + "Rank-zero and worker forward/backward paths still exercise uniform rejection and asymmetric fatal promotion without generating duplicate parameterized reports." + ] + }, + { + "id": "TA-161", + "scope": "LoRA checkpoint SGLang layout and expert-ownership roundtrip fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The SGLang shared-outer roundtrip now inspects its exact serialized key set and shapes before reloading, subsuming a save-only layout report.", + "Hybrid-shared and all-owner expert layouts now pass through one adapter-manager roundtrip policy while retaining distinct source, checkpoint, and manager state." + ] + }, + { + "id": "TA-162", + "scope": "native block-FP8 state lifecycle and CPU admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One byte-exact lifecycle contract now proves frozen state, FP8 and scale contents, and dtype-application preservation.", + "Import laziness, CPU forward rejection, and explicit input-free CUDA materialization admission now share one fail-closed CPU policy contract." + ] + }, + { + "id": "TA-163", + "scope": "adapter-gradient ownership configuration, fingerprints, replica coverage, and analytical self-test", + "decision": "remove", + "status": "applied", + "evidence": [ + "Positive and invalid bucket configuration, rank-local identity and geometry invariance, and admitted/rejected replica coverage now execute as policy transactions.", + "The removed analytical test exercised only math helpers defined inside the test; a retained adapter-manager transaction checks the real production step against scale, clip, AdamW parameter, and moment equations." + ] + }, + { + "id": "TA-164", + "scope": "DRGRPO forward, backward, metrics, zero boundaries, and KL fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One seeded numerical contract now checks the exact loss, gradient norm, finite values, and metric schema from the same operation.", + "Zero advantages, ignored labels, and empty sequences share one zero-loss policy, while KL reference admission and effect share one positive-KL policy." + ] + }, + { + "id": "TA-165", + "scope": "server endpoint and artifact-path security examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Allowlist admission, DNS pinning, metadata rejection, and malformed host rejection now form one outbound endpoint policy.", + "Generic escape and symlink rejection, environment-root confinement, and explicit-root authority now form one artifact path policy." + ] + }, + { + "id": "TA-166", + "scope": "batch-invariant router GEMM input and top-k example fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FP32-reference parity, empty-token behavior, and BF16 input admission now share one router kernel contract.", + "Top-k renormalization, cast-only behavior, and FP32 input admission now share one post-processing contract; invariance, gradients, and MoEBlock consumers remain independent." + ] + }, + { + "id": "TA-167", + "scope": "local phase ordering, timing summary, and memory summary examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical, custom, and empty phase ordering now form one ordering contract.", + "Phase-time and memory summaries each verify empty output, canonical ordering, local aggregate fields, and normalization through one complete summary map." + ] + }, + { + "id": "TA-168", + "scope": "per-component decoder discovery and model-style hook fragments", + "decision": "remove", + "status": "applied", + "evidence": [ + "Direct decoder-suffix and nested-attribute helper tests were dominated by the retained attach-and-run consumer transaction.", + "One live CUDA transaction now attaches to GLM and Qwen model styles, records present forward/backward phases, and omits absent indexer and shared-expert phases." + ] + }, + { + "id": "TA-169", + "scope": "manual CUDA timer disabled, mode, recorded, and unrecorded fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Disabled behavior and invalid mode rejection now share one API policy.", + "Forward and recompute timing plus unrecorded-pair omission now execute through one event-drain lifecycle." + ] + }, + { + "id": "TA-170", + "scope": "activation-offload populated and empty metric contexts", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One consume lifecycle now converts populated forward/backward byte counters to gigabytes, checks single consumption, and proves empty or unsupported contexts emit no metrics." + ] + }, + { + "id": "TA-171", + "scope": "DCP synchronization and metadata process-group selector fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One synchronization backend policy now covers cached Gloo creation for NCCL and default-group use for Gloo.", + "One metadata policy now covers disabled pipeline parallelism, caller-supplied groups, and global Gloo fallback." + ] + }, + { + "id": "TA-172", + "scope": "parallel-plan meta slicing properties and exact-GLM malformed disposition reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The already-local meta allocation path now proves local shape, meta device, dtype, requires-grad, shard placement, and unrelated replication in one contract.", + "Already-local and force-shard malformed singleton policies both still run without generating separate parameterized reports." + ] + }, + { + "id": "TA-173", + "scope": "adapter-manager local active-slot shape smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "Retained manager registration and multi-adapter forward transactions already assert rank-specific local factor shapes.", + "The sharded-state suite retains logical pack/unpack, deterministic initialization, layout discovery, and a real two-rank uneven DTensor transaction." + ] + }, + { + "id": "TA-174", + "scope": "checkpoint zero-meta stages and initial optimizer-load modes", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Pre-load materialization and post-restore zero-meta failures now execute in one checkpoint lifecycle while preserving stage-specific assertions.", + "Default optimizer restore and weights-only initial restore now share one runner forwarding and state-synchronization policy." + ] + }, + { + "id": "TA-175", + "scope": "MoE TP simulation admission and eager reduction-mode examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Environment, topology, EP rejection, and layer filtering now share one admission policy.", + "Direct, BF16, cache, and TP-size override modes still run against their corresponding independent references from one shared expert and routing fixture." + ] + }, + { + "id": "TA-176", + "scope": "MoE TP carried-shard and diagnostic-capture transactions", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One MoEBlock execution now proves reshaped carried shards, flat per-shard diagnostic captures, their exact sum, and the final reshaped output.", + "Every backend-specific call-layout contract remains independent." + ] + }, + { + "id": "TA-177", + "scope": "DeepSeek-V4 shared-MLP clamp smoke and non-hash MoE property fragments", + "decision": "remove", + "status": "applied", + "evidence": [ + "The removed shared-MLP smoke explicitly avoided asserting a clamp effect; the retained forced-gate numerical case proves bounded clamped output and shape.", + "Non-hash structure, forward/backward, selection-only bias, and shared-expert contribution now execute through one model transaction." + ] + }, + { + "id": "TA-178", + "scope": "DeepSeek-V4 hash-layer structure, input admission, and forward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One hash-layer transaction now proves table and bias structure, rejects missing input IDs, executes table-driven routing, and checks finite gate gradients.", + "Record-to-replay backward and unknown replay-stage failure remain independent state-machine contracts." + ] + }, + { + "id": "TA-179", + "scope": "FP8 LM-head CE module selection, TP, and temperature fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One per-token CE transaction now selects module and FP32-master paths locally and under identity TP collectives.", + "Temperature propagation is checked once through both the primitive and CausalLM consumer; importance-sampling and TP-gradient consumers remain independent." + ] + }, + { + "id": "TA-180", + "scope": "Qwen3.5 native-EP admission and variable-row collective fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP8 admission, exact structural flags, and missing trainer-EP rejection now form one exact-combine policy.", + "Token padding and backward unpadding, invalid ID padding, and shared maximum-row selection now form one variable-row collective contract." + ] + }, + { + "id": "TA-181", + "scope": "DeepSeek-V3 checkpoint expert layout conversion fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "External dense merge, internal fused pass-through, save splitting, and multimodal filtering now execute as one layout-conversion transaction.", + "Dense and packed EP slicing share one policy, packed dtype behavior shares one load transaction, and model-level packed recognition carries through quant-config parsing." + ] + }, + { + "id": "TA-182", + "scope": "DR-GRPO runner legacy input and per-token output option fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy logprobs, temperature forwarding, disabled per-token output, and K3-forced output now run through one runner option contract.", + "Full loss dispatch, sampler-prefill model forwarding, and the forward-backward loop remain separate integration boundaries." + ] + }, + { + "id": "TA-183", + "scope": "SignSGD update, missing-gradient, and stateless persistence fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sign updates, decoupled weight decay, zero-sign behavior, and missing-gradient preservation now execute in one optimizer step.", + "Multiple stateless steps and hyperparameter-only state-dict restoration now share one persistence lifecycle." + ] + }, + { + "id": "TA-184", + "scope": "CausalLM Z-loss analytical example and duplicate temperature report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The random reference comparison now also proves backward gradients, making a separate finiteness-only gradient report redundant.", + "The zero-logit formula example only restated the test reference, and CausalLM temperature propagation is already retained in the FP8 LM-head suite." + ] + }, + { + "id": "TA-185", + "scope": "MoE train-router dispatch, default, and synthetic replay fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One trained-router transaction proves all-to-all gradients before rejecting DeepEP dispatch.", + "Argument and constructed-model defaults share one policy, while balanced forward routing and replay regather share one uniform-routing contract." + ] + }, + { + "id": "TA-186", + "scope": "shared-prefix detection, repack layout, loss-field, and remap fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One shared-prefix transaction now detects groups, repacks exact token and position layouts, preserves decoded loss fields, and remaps outputs to original order.", + "No-sharing and one-token-prompt boundaries remain independent because they select different backend outcomes." + ] + }, + { + "id": "TA-187", + "scope": "multi-part optimizer structure, coverage, step, and scheduler fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "MultiOptimizer construction now proves DCP model mapping and complete parameter-group coverage together.", + "A single optimizer lifecycle proves every virtual part updates, gradients clear, and every learning-rate group decays." + ] + }, + { + "id": "TA-188", + "scope": "batch-invariant fused LM-head forward and backward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forward values, ignored-token loss, and backward gradients now compare with eager from the same graph for the default temperature.", + "The non-unit temperature path likewise proves forward and backward parity in one transaction; determinism, guards, unit-temperature bytes, and probability clamping remain separate." + ] + }, + { + "id": "TA-189", + "scope": "model-runner FP8, QARL, and sharded-loss builder propagation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One builder-policy transaction now checks fail-closed FP8 defaults, every QARL calibration field, and sharded LM-head loss propagation through the same initialization boundary.", + "Exact GLM block-FP8 QLoRA remains separate because it also resolves the runtime target-module set." + ] + }, + { + "id": "TA-190", + "scope": "exact GLM dense-component and routed-expert gradient ownership fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Gate-up, dense MLP, and absorbed-KV leaves now compile through one module-managed ownership matrix with component-specific canonical factor names.", + "Unwrapped EP16 ownership and mutated DeepEP dispatch now fail inside one routed runtime-admission contract." + ] + }, + { + "id": "TA-191", + "scope": "batch tensor conversion float side-channel fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "DR-GRPO old/reference logprobs and distillation teacher hidden states now prove FP32 preservation in one conversion transaction.", + "Ragged teacher padding and sequence-parallel sharding remain independent shape transformations." + ] + }, + { + "id": "TA-192", + "scope": "server batch-slice topology examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FSDP shards, TP replicas, default EP ranks, and legacy EP duplication now map to their exact slice coordinates in one selector policy.", + "All original rank and topology cases still execute without one report per branch." + ] + }, + { + "id": "TA-193", + "scope": "fused GDN LoRA merged-forward and cache-generation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical sliced folding and the GDN output-projection consumer now share one gradient-preserving merged-forward contract.", + "Cache slice bounds and release of the previous adapter generation now execute in one version-change lifecycle." + ] + }, + { + "id": "TA-194", + "scope": "DeepSeek-V4 attention sink storage and call-dtype fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One sink lifecycle now proves initial FP32 storage and FSDP marking, BF16 module conversion, and FP32 promotion at the TileLang call boundary.", + "Attention variants, quantization admission, compressor structure, and TP rejection remain separate." + ] + }, + { + "id": "TA-195", + "scope": "stochastic-round output metadata and input-dtype fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One API contract now proves BF16 output dtype, shape, and device before rejecting a non-FP32 input.", + "Unbiased expectation, neighbor bounds, and generator determinism remain independent numerical properties." + ] + }, + { + "id": "TA-196", + "scope": "MoE fused gate-up registration and deferred QLoRA skip fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Base and LoRA expert modules now prove the same fused parameter registration and gate/up views in one structure contract.", + "Qwen3 and Qwen3.5 handlers now prove their family-specific deferred expert keys through one QLoRA skip policy." + ] + }, + { + "id": "TA-197", + "scope": "batch-invariant GDN gating and gated-norm forward/backward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Each fused primitive now compares forward values and every input gradient with its independent PyTorch composition from the same graph.", + "Gated norm row invariance remains in the numerical transaction; model routing and pinned solve geometry remain separate." + ] + }, + { + "id": "TA-198", + "scope": "Nemotron-H forward, router-output, raw-backward, and labeled-loss fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One model transaction now proves output shape, optional router logits, labeled CausalLM loss, and gradients through every mixer family.", + "Gradient checkpointing and the two packed-sequence contracts remain independent execution modes." + ] + }, + { + "id": "TA-199", + "scope": "exact GLM dense and LM-head legacy weight-sync guard fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense MLP and attention projections now reject adapter preparation, collective merging, and raw extraction through one factor-only policy matrix.", + "LM-head ordinary and prepacked sparse-delta publication both fail before adapter or backend work in one side-effect guard." + ] + }, + { + "id": "TA-200", + "scope": "merged block-FP8 deferred-loader projection key reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "QKV and gate/up merged projections still derive and validate every exact source key inside one key-selection policy.", + "EP expert slicing, missing pairs, retained caches, and per-module release remain independent." + ] + }, + { + "id": "TA-201", + "scope": "trainer trunk-linear engagement and numerical-family selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Exact server, ordinary server, and exact non-server models now share one pre-FSDP trunk engagement policy with explicit harness-state isolation.", + "Exact GLM family selection and ordinary rollback execute as one structural numerical-program transition." + ] + }, + { + "id": "TA-202", + "scope": "P2P async transfer size-threshold examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One transfer policy now sends a large bucket asynchronously, a medium bucket synchronously, and the same medium bucket asynchronously after lowering the configured cutoff.", + "Status timeout and sender preparation timeout remain separate failure and HTTP boundaries." + ] + }, + { + "id": "TA-203", + "scope": "DeepSeek-V4 FWHT known-pattern, roundtrip, and norm examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One numerical matrix now proves a known Hadamard row, orthonormal roundtrip, and norm preservation across power-of-two widths.", + "Invalid width and rotate_activation fallback dispatch remain independent API boundaries." + ] + }, + { + "id": "TA-204", + "scope": "orchestrator-runner generic serialization examples and message metadata fragments", + "decision": "remove", + "status": "applied", + "evidence": [ + "All typed message classes, tensor payloads, JSON conversion, and pickle rejection now share one wire-format contract.", + "Generic large nested-list and nested-dictionary serializer examples were removed; identity, timestamps, optional fields, and ACK construction remain production-specific." + ] + }, + { + "id": "TA-205", + "scope": "DeepSeek-V3 auxiliary router-logit layer-count fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One auxiliary-loss policy now proves all-MoE emission and omission of the configured dense prefix.", + "Base forward/backward, LoRA targeting, and routing replay remain independent consumers." + ] + }, + { + "id": "TA-206", + "scope": "FlashMLA device and production-shape admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One fail-closed admission contract now rejects CPU dispatch, an unproven head shape, and an overflowing flattened KV address space before backend import.", + "Flattening, compacted backward, and all-invalid rows remain independent numerical paths." + ] + }, + { + "id": "TA-207", + "scope": "RL primitive KL estimator mode parameterization", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "K1, K2, K3, and low-variance KL modes still compare against the independent Slime formulas inside one estimator matrix.", + "Sequence KL, clipping, OPSM, and reduction semantics remain independent primitives." + ] + }, + { + "id": "TA-208", + "scope": "families-v2 exact and nonexact selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One environment policy now proves legacy rollback for nonexact models and structural version pinning for exact GLM and Qwen paths.", + "Vendoring, selected-logit bytes, and batch composition invariance remain separate contracts." + ] + }, + { + "id": "TA-209", + "scope": "sqrt-softplus routing regather value, scaling, and dtype fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One replay-regather transaction now matches eager weights, isolates the routed scaling factor, preserves cached experts, and returns the requested BF16 dtype.", + "The unchanged softmax route remains an independent regression contract." + ] + }, + { + "id": "TA-210", + "scope": "tensor collator scalar, dtype, string, and input-container fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "List, NumPy, tensor, boolean, scalar, empty, known-field dtype, and string cases now execute through one conversion policy.", + "Variable-length and packed-sequence handling remains a separate layout outcome." + ] + }, + { + "id": "TA-211", + "scope": "model-runner session registry and optimizer option fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One LoRA session lifecycle now synchronizes registry state after both optimizer update and checkpoint load.", + "One dense optimizer policy now proves DistSignSGD normalization and clipping before the optional empty-cache suppression branch." + ] + }, + { + "id": "TA-212", + "scope": "runner-dispatcher load-state preparation and tenant-routing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Preparation now preserves rank-zero errors while accepting paths under the server output root and rejecting unrelated roots in one policy.", + "Multi-adapter coordinator delegation and single-tenant trainer delegation share one load-state routing transaction." + ] + }, + { + "id": "TA-213", + "scope": "permanently skipped GLM5 FLOP feature placeholders", + "decision": "remove", + "status": "applied", + "evidence": [ + "The two tests were unconditionally skipped because GLM5 sparse-MLA and DSA FLOP accounting is not implemented and currently reports zero.", + "The executable CP-size invariance regression remains; feature coverage should be added with the implementation rather than as dead placeholders." + ] + }, + { + "id": "TA-214", + "scope": "Kimi wrapper LoRA target-source precedence fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One target-resolution policy now proves wrapper-derived defaults, explicit target override, and strict-manifest override in precedence order.", + "Every original target list still executes against an actual wrapper config." + ] + }, + { + "id": "TA-215", + "scope": "test-local MoE expert auto-merge implementations", + "decision": "remove", + "status": "applied", + "evidence": [ + "Both files copied the parser, buffer, format detector, and simulated loader into tests and invoked no XoRL implementation.", + "Production ExpertWeightBuffer transposition/output remains covered directly by test_moe_gkn_format, while family checkpoint handlers cover real dense, packed, EP-sliced, and roundtrip loads." + ] + }, + { + "id": "TA-216", + "scope": "parallel-state singleton initialization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One singleton lifecycle now proves publication, automatic DP shard inference, device selection, and reinitialization rejection.", + "Mesh layout, construction validation, and requires-mesh behavior remain independent contracts." + ] + }, + { + "id": "TA-217", + "scope": "DeepSeek-V3 Kimi wrapper config mapping fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Wrapper unwrapping and official aux_loss_alpha defaults now execute through one configuration-conversion policy.", + "Registry lookup and local-directory auto-config loading remain separate integration boundaries." + ] + }, + { + "id": "TA-218", + "scope": "tokenizer and processor remote-code fallback fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Tokenizer and processor fallback now share one auto-loader security policy that forbids implicit remote code and retains right padding.", + "The local Kimi TikToken implementation remains an independent functional roundtrip." + ] + }, + { + "id": "TA-219", + "scope": "DeepSeek-V4 LoRA structural smokes and unsupported generic expert claim", + "decision": "remove", + "status": "applied", + "evidence": [ + "Attention adapter type/freeze assertions now live in the retained forward-backward gradient transaction instead of standalone structure tests.", + "The generic expert-LoRA claim was stale: DeepSeek-V4 expert semantics are deliberately rejected by the shared semantic guard, so the retained end-to-end path targets supported attention adapters only." + ] + }, + { + "id": "TA-220", + "scope": "scheduler FIFO and ScheduledRequest implementation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Defaults and FIFO add, dispatch, remove, and clear behavior now form one policy lifecycle.", + "A sleep-based direct ScheduledRequest state smoke was removed because scheduler completion/failure/abort transactions already exercise those transitions." + ] + }, + { + "id": "TA-221", + "scope": "Mooncake side-payload missing-key fragment", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Successful typed tensor roundtrips and missing-key rejection now share one store contract.", + "Reference slicing/cleanup and R3 payload validation remain distinct higher-level lifecycles." + ] + }, + { + "id": "TA-222", + "scope": "batch-utils float conversion and ragged-padding fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Scalar logprob fields, rectangular teacher states, and ragged teacher-state padding now execute through one conversion transaction.", + "Sequence-parallel sharding remains an independent consumer boundary." + ] + }, + { + "id": "TA-223", + "scope": "DR-GRPO runner dispatch and option fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Modern and legacy loss fields, KL metrics, temperature, output suppression, and K3-forced output now share one microbatch-loss dispatch contract.", + "Sampler prefill propagation and the full forward-backward loop remain separate integrations." + ] + }, + { + "id": "TA-224", + "scope": "test-local GKN transpose and reference-MoE proof", + "decision": "remove", + "status": "applied", + "evidence": [ + "The removed report compared local matrix multiplications and a local expert loop without calling XoRL.", + "The retained ExpertWeightBuffer test proves production transposition and output, and the live backend matrix compares eager, native, Triton, and MoEBlock consumers." + ] + }, + { + "id": "TA-225", + "scope": "future sparse-MLA KV-major reference scaffold", + "decision": "remove", + "status": "applied", + "evidence": [ + "The test compared two Torch references for a future atomic-free kernel and invoked no production sparse-MLA path.", + "Production forward, backward, deterministic, and combined-kernel performance contracts remain collected." + ] + }, + { + "id": "TA-226", + "scope": "EP gradient backend support and rejection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Every supported backend and the unknown-backend fail-closed outcome now form one reduction-domain table contract.", + "Malformed metadata and the real two-rank reduction lifecycle remain independent." + ] + }, + { + "id": "TA-227", + "scope": "lm-head tensor-parallel topology wrappers", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CP-sourced, DP-sourced, and HSDP-sourced four-rank meshes now execute as one topology matrix rather than three wrapper reports.", + "Each subprocess still verifies TP groups, replica groups, FSDP groups, and topology-specific mesh structure." + ] + }, + { + "id": "TA-228", + "scope": "LoRA cast-once zero and nonzero merge fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Each linear and expert merge transaction now proves both bit-exact zero-adapter behavior and the nonzero FP32-add-then-cast reference across supported dtypes.", + "Linear and grouped-expert implementations remain separate reports." + ] + }, + { + "id": "TA-229", + "scope": "NF4 per-width, size, dtype, and allocation smokes", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Twenty reports now form three codebook, flat-codec, and GKN-codec contracts while still running every 32/64/128 group width, layout, dtype, scale, accuracy, zero, and cross-layout check.", + "Two large-allocation smokes were removed because they repeated the same codec behavior with 16M-element and 14M-element tensors without a distinct production boundary." + ] + }, + { + "id": "TA-230", + "scope": "TopKRouter selector and configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Softmax, balanced, hash, sqrt-softplus, scaling, tie, invalid-input, and model-config outcomes now execute as selector policies rather than isolated examples.", + "FP32 selection and the MoEBlock consumer remain independent because they protect precision and integration boundaries." + ] + }, + { + "id": "TA-231", + "scope": "OPD estimator, policy-gradient, and task-weight fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full-vocabulary modes, estimator formulas, dispatch, policy-gradient admission, and task weighting now report their complete policies.", + "Clamping, stable metric keys, and ignored-label sampled-logprob behavior remain distinct numerical or API boundaries." + ] + }, + { + "id": "TA-232", + "scope": "loss reducer denominator and layout fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "TokenPartial now proves denominator and microbatch composition together, while SequencePartial covers dense, packed, and context-parallel composition as one layout policy.", + "Empty-input zero behavior remains a separate boundary." + ] + }, + { + "id": "TA-233", + "scope": "shared loss implementation and microbatch parameter reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy TokenPartial identities and paired implementation checks now run explicit case matrices within one contract per semantic behavior.", + "Importance-sampling and policy-loss microbatch composition likewise report once while retaining every loss variant." + ] + }, + { + "id": "TA-234", + "scope": "Qwen Class-B RoPE and EP adapter parameter reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense and MoE half-rotate/cast behavior, supported CPU dtypes, Class-B shapes, and available EP adapters now execute internal matrices instead of publishing one report per equivalent case.", + "All numerical cases and every installed adapter backend are still exercised." + ] + }, + { + "id": "TA-235", + "scope": "shared-prefix attention dtype and head-shape parameter reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The two dtypes, two head dimensions, and two query-to-KV head ratios now form one numerical backend contract.", + "No-sharing and one-token-prompt behavior remain separate backend outcomes." + ] + }, + { + "id": "TA-236", + "scope": "hardware-specific Qwen3-8B TFLOPS pytest thresholds", + "decision": "remove", + "status": "applied", + "evidence": [ + "The three reports hard-coded H100 throughput thresholds without admitting only H100 hardware and normally skipped based on local model-directory presence.", + "They were multi-minute benchmark jobs rather than stable correctness contracts; Qwen LoRA/FSDP E2E suites retain real training and loss-convergence coverage." + ] + }, + { + "id": "TA-237", + "scope": "direct pack_parallel non-empty smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report asserted only that two calls returned at least one bin and did not verify capacity, coverage, allocation, or ordering.", + "PackingDataset still invokes the production parallel packer for both sequential and multipack methods, while dedicated FFD, grouping, and allocation contracts retain the actual invariants." + ] + }, + { + "id": "TA-238", + "scope": "QLoRA package hasattr import smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "The deleted report imported xorl.qlora and checked only two exported attribute names.", + "Trainer and model-builder imports exercise the public exports, real QLoRA suites invoke the implementations, and the clean-interpreter dependency-cycle contract remains." + ] + }, + { + "id": "TA-239", + "scope": "default and explicit MLA target partition examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default and caller-specified MLA projections now execute as one MoE target-partition policy.", + "The real DeepSeek-V3 model integration remains separate and still checks attention, shared-expert, and routed-expert replacement." + ] + }, + { + "id": "TA-240", + "scope": "CP16 first-rank and padded-tail side-channel reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ranks 0 and 15 now run inside one CP16 target-sharding contract.", + "Both full and padded-tail slices still verify labels, target tokens, old logprobs, advantages, reference logprobs, and padding values." + ] + }, + { + "id": "TA-241", + "scope": "Muon optimizer EP checkpoint transition wrappers", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP2-to-EP4, EP4-to-EP2, and same-EP identity now form one checkpoint transition matrix.", + "Each case still launches its own four-rank save, reload, global gather, and exact momentum comparison." + ] + }, + { + "id": "TA-242", + "scope": "DTensor copy and save-materialization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Replicated, sharded, padded-tail, and shape-rejection outcomes now form one copy policy.", + "One-dimensional writer-only and two-dimensional all-rank/optional-writer materialization still launch separate four-rank workers within one topology contract." + ] + }, + { + "id": "TA-243", + "scope": "PP NCCL sender, empty-buffer, and product-helper fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Empty sender and receiver outcomes now share one protocol transaction, and nonempty sender return behavior is asserted where metadata and flattening are verified.", + "The direct private product-helper example was replaced by scalar-tensor reconstruction through the production receive path." + ] + }, + { + "id": "TA-244", + "scope": "cautious-decay helper and denominator-mode fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Zero decay, ordinary decoupled decay, and coordinate masking now form one helper policy.", + "Chunked denominator parity still runs both ordinary and Kahan modes with three optimizer steps and complete state comparison." + ] + }, + { + "id": "TA-245", + "scope": "architecture registration, AutoConfig, and builder fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "DeepSeek-V4 now follows one standard snapshot through AutoConfig, AutoModel registration, and the XoRL foundation-model builder.", + "DeepSeek-V3, Nemotron-H, and Qwen3.5 registry assertions now live with their real local/HF configuration conversions; numerical model and checkpoint contracts remain separate." + ] + }, + { + "id": "TA-246", + "scope": "direct ModelArguments routing-weight default assertion", + "decision": "remove", + "status": "applied", + "evidence": [ + "The deleted report instantiated a dataclass and asserted only that one field equaled auto.", + "Server YAML loading still proves the omitted default and explicit value serialize correctly, while trainer alignment and runtime resolution tests retain the real consumers." + ] + }, + { + "id": "TA-247", + "scope": "gradient-checkpoint default and override method fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default class policy, default enablement, and nondefault propagation across dense and MoE checkpoint layers now form one configuration contract.", + "The independent training/flag/method execution gate remains separate." + ] + }, + { + "id": "TA-248", + "scope": "single-part and multi-part optimizer construction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Singleton plain-optimizer selection and multi-part mapping/parameter coverage now execute as one construction policy.", + "Live multi-part step/scheduler behavior and invalid custom groups remain separate outcomes." + ] + }, + { + "id": "TA-249", + "scope": "token-diagnostic disabled, ignored-label, and top-k boundary fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Disabled inputs, all-ignored labels, and top-k larger than vocabulary now form one boundary policy.", + "Ranking, KL position mapping, loss-logprob comparison, raw-weight reference, and hidden-state summaries remain independent numerical contracts." + ] + }, + { + "id": "TA-250", + "scope": "DeepSeek-V4 private mapper, APE, FP8, and MXFP4 fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "APE valid shapes and invalid shape, FP8 full-block and tail dequantization, and MXFP4 values and block scaling now report complete codec policies.", + "Unknown-key handling moved from the private name mapper to the production checkpoint handler and now also verifies unmapped accounting." + ] + }, + { + "id": "TA-251", + "scope": "DeepEP asynchronous-combine default and opt-in fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default synchronous behavior and the explicit unsafe environment opt-in now execute as one admission policy.", + "Both cases still pass through tokens_post_combine and inspect the actual fused-operation argument." + ] + }, + { + "id": "TA-252", + "scope": "FSDP reduce-dtype, boolean coercion, and prefetch setting fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Singleton, sharded, and overridden expert precision plus supported and rejected dtype names now form one reduction policy; boolean spellings and rejection form one admission policy.", + "Backward-only, forward-only, bidirectional, and not-needed prefetch behavior now execute as one direction matrix with exact module ordering." + ] + }, + { + "id": "TA-253", + "scope": "per-component timer disabled and unrecorded-event fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Disabled lifecycle no-op behavior and enabled handling of recorded versus unrecorded event pairs now share one timer admission contract.", + "Live model-style CUDA hook coverage remains a separate integration." + ] + }, + { + "id": "TA-254", + "scope": "R3 payload transport success and directory rejection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Mooncake, legacy externalization, filesystem directory, and inline defaults now execute as one transport-admission matrix.", + "The missing-directory rejection remains in the same policy and still reaches the production argument loader." + ] + }, + { + "id": "TA-255", + "scope": "exact GLM rank-1 topology admission and rejection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The certified TP16 topology and its non-TP16 rejection now form one exact-GLM topology contract.", + "The accepted values and rejection message are both retained." + ] + }, + { + "id": "TA-256", + "scope": "full-weight QARL and FP8 incompatible scope and source fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "LoRA conflicts, mutual QARL and FP8 exclusion, missing calibration, MTP metadata, Mamba config, and Nemo ModelOpt sources now execute as one fail-closed policy.", + "Every prior YAML shape and error boundary is preserved." + ] + }, + { + "id": "TA-257", + "scope": "FP8 configuration alias entry-point fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Nested train fp8_cfg and Nemo policy megatron_cfg aliases now share one normalization contract.", + "Layer-island and Blackwell fields remain asserted on the train alias path." + ] + }, + { + "id": "TA-258", + "scope": "server multi-adapter unsupported-mode fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Merge-interval and pipeline-parallel rejection now execute as one unsupported multi-adapter policy.", + "The general broadcast load-weights rejection remains separate from adapter-specific admission." + ] + }, + { + "id": "TA-259", + "scope": "weight-sync FP8 normalization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit FP8 normalization, default dynamic-block behavior, and module-name cleanup now form one normalization policy.", + "No-quantization aliases remain a separate BF16 no-op contract." + ] + }, + { + "id": "TA-260", + "scope": "weight-sync FP8 invalid configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unsupported formats, activation schemes, UE8M0 scales, module exclusions, and internal unsupported markers now execute as one rejection matrix.", + "Non-FP8 quantization methods retain a separate method-admission policy." + ] + }, + { + "id": "TA-261", + "scope": "training CLI FP8 configuration alias fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Nested train fp8_cfg and Nemo policy megatron_cfg aliases now execute through parse_args as one CLI normalization contract.", + "The layer-island and Blackwell fields remain asserted on the native train path." + ] + }, + { + "id": "TA-262", + "scope": "training CLI QARL and FP8 incompatible configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Adapter conflicts, missing QARL calibration data, QARL plus FP8, MTP metadata, and Mamba config now execute as one CLI rejection policy.", + "Every former YAML payload and exact error boundary remains exercised through parse_args." + ] + }, + { + "id": "TA-263", + "scope": "training simulator untrusted calibration and model path fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Built-in traversal, relative escape, missing default, symlink escape, and unapproved model metadata paths now form one filesystem admission policy.", + "All five attacks still reach the production calibration-pack or metadata loader." + ] + }, + { + "id": "TA-264", + "scope": "exact Qwen3.5 certified topology admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both certified MoE topologies, the certified dense topology, and ten nearby rejected mutations now execute as one topology policy.", + "The exact topology validator remains the observable boundary." + ] + }, + { + "id": "TA-265", + "scope": "exact Qwen3.5 model-scope admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense, MoE, and Hugging Face outer-config snapshots now share one model-scope policy with wrong layer types and nearby geometry rejection.", + "Canonical GLM model-scope validation remains an independent family contract." + ] + }, + { + "id": "TA-266", + "scope": "P2P completion and failed-transfer cleanup fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Pending-transfer failure, receiver-completion suppression, best-effort drain, successful completion payload, failed completion cleanup, and default cache behavior now form two lifecycle policies.", + "All network mocks, engine deregistration checks, endpoint metadata, and failure messages are retained." + ] + }, + { + "id": "TA-267", + "scope": "session endpoint LoRA registration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full LoRA overrides, rank-only defaults, and existing-session refresh now execute as one registration policy.", + "Worker registration, normalized session specs, materialization, and refresh idempotence remain asserted." + ] + }, + { + "id": "TA-268", + "scope": "session endpoint reserved-checkpoint fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default, per-session, stale, and existing reserved checkpoints now execute as one persistence policy.", + "Save counts, model IDs, reserved paths, overwrite, and preservation behavior are retained." + ] + }, + { + "id": "TA-269", + "scope": "full-weight session endpoint admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default full-weight registration, nondefault multitenancy rejection, and per-session override rejection now form one admission policy.", + "The materialize-false worker contract and both rejection messages remain covered." + ] + }, + { + "id": "TA-270", + "scope": "LoRA session kill, checkpoint, and default-session protection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ordinary kill and re-registration, final-checkpoint URI conversion, default kill preservation, and default unload rejection now form two lifecycle policies.", + "Registry, future-store, worker request, checkpoint, and HTTP status outcomes remain asserted." + ] + }, + { + "id": "TA-271", + "scope": "weights-info endpoint mode and path fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "On-disk LoRA metadata, full-weight metadata, and path-escape rejection now execute as one endpoint policy.", + "The disk-over-memory authority check remains intact." + ] + }, + { + "id": "TA-272", + "scope": "request-processor routing payload cleanup fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Mooncake success cleanup, Mooncake backend-failure cleanup, and filesystem payload cleanup now form one transport lifecycle.", + "Payload existence during execution and removal after execution remain verified for every path." + ] + }, + { + "id": "TA-273", + "scope": "request-processor token diagnostic boundary fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed-sample splitting, empty diagnostics, and mismatched-field rejection now execute as one decoding policy.", + "Position rebasing and all diagnostic field alignments remain asserted." + ] + }, + { + "id": "TA-274", + "scope": "request-processor packed-row batching fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Global row grouping, rank-local deferral, and routed-replay rejection now form one batching policy.", + "Batch counts, sequence boundaries, sample counts, metrics, and error output remain covered." + ] + }, + { + "id": "TA-275", + "scope": "adapter gradient-epoch abort fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Successful idempotent abort and rejection while publication-pending or poisoned now execute as one lifecycle policy.", + "Gradient scratch reset, monotonic counters, and publication state remain asserted." + ] + }, + { + "id": "TA-276", + "scope": "authoritative adapter checkpoint plan admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Identity-label differences with equal direct contracts and direct-contract mismatches now execute as one admission policy.", + "Mismatch rejection still proves parameters, optimizer, session spec, step, and learning rate are unchanged." + ] + }, + { + "id": "TA-277", + "scope": "adapter checkpoint structure rejection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Target-module mismatch, missing LoRA tensors, and rank beyond live capacity now form one structural admission policy.", + "Each malformed checkpoint is still built and rejected through load_adapter_state." + ] + }, + { + "id": "TA-278", + "scope": "adapter PEFT filename and sharding compatibility fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Weight-suffixed tensor names and indexed sharded safetensors now execute as one PEFT compatibility policy.", + "Both restored LoRA factors retain exact value checks." + ] + }, + { + "id": "TA-279", + "scope": "adapter residency and eviction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dirty residency, clean-victim selection, auto-save failure, and multi-rank rejection now form one eviction policy.", + "Every accepted or retained adapter identity remains asserted after the transition." + ] + }, + { + "id": "TA-280", + "scope": "inference endpoint port-selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Single-port and explicit worker-port registration now execute as one endpoint port policy.", + "Health checks, LoRA load/unload routes, and adapter discovery remain asserted." + ] + }, + { + "id": "TA-281", + "scope": "inference endpoint FP8 KV-cache registration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit FP8 metadata, dtype inference with cache-version alias, and required-FP8 rejection now form one admission policy.", + "Postprocess, static-scale, epoch, and rejection details remain covered." + ] + }, + { + "id": "TA-282", + "scope": "inference weight-sync pool filtering fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Eval, default, all-pool, and no-match selection now execute as one endpoint-selection policy.", + "The no-match path still proves no orchestrator request is sent." + ] + }, + { + "id": "TA-283", + "scope": "inference weight-sync quantization admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit null quantization and malformed empty quantization now form one request-admission policy.", + "Default suppression and HTTP 400 behavior remain asserted." + ] + }, + { + "id": "TA-284", + "scope": "inference weight-sync cache invalidation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FP8 auto-flush, cache-version aliasing, BF16 no-flush, and explicit none mode now form one cache policy.", + "Payload flags, response flags, endpoint epochs, postprocess, and static-scale outcomes remain checked." + ] + }, + { + "id": "TA-285", + "scope": "receiver quantization detection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Skip-list normalization and unsupported MTP, activation, and UE8M0 receiver configurations now execute as one detection policy.", + "Every detected configuration still passes through the public normalizer." + ] + }, + { + "id": "TA-286", + "scope": "receiver quantization enrichment fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Receiver skip-list fill, user override preservation, BF16 and null pass-through, and unsupported-reason propagation now form one enrichment policy.", + "Unsupported enriched output remains rejected by normalization." + ] + }, + { + "id": "TA-287", + "scope": "GLM5 local config unsafe-value fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Non-JSON object values and dunder keys now execute as one local-config security policy.", + "Both exact validation paths and messages remain exercised." + ] + }, + { + "id": "TA-288", + "scope": "GLM5 blocked indexer selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full-query blocked-versus-dense parity and query-offset shard parity now form one indexer policy.", + "The same production select_topk path remains used for both shapes." + ] + }, + { + "id": "TA-289", + "scope": "GLM5 sparse-MLA torch reference fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full-top-k dense equivalence and local-query offset equivalence now execute as one numerical reference policy.", + "Both outputs retain exact tolerance checks." + ] + }, + { + "id": "TA-290", + "scope": "GLM5 sparse-MLA dispatch fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CPU auto fallback and unknown-backend rejection now form one backend dispatch policy.", + "The production dispatcher remains the observable boundary." + ] + }, + { + "id": "TA-291", + "scope": "adapter optimizer parameter identity fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical parameter ordering and wrapper-insensitive structure fingerprints now execute as one identity policy.", + "Both public checkpoint identity inputs remain directly asserted." + ] + }, + { + "id": "TA-292", + "scope": "adapter optimizer checkpoint artifact fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy pickle, shard-without-manifest, and declared-state-without-artifacts now form one checkpoint admission policy.", + "Single and multi-rank refusal, weights-only migration, and no-partial-registration outcomes remain covered." + ] + }, + { + "id": "TA-293", + "scope": "adapter optimizer logical reshard fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One-dimensional world-size change, multidimensional disjoint rectangles, replicated slicing, and same-world layout change now form one reshard policy.", + "Moments, squared moments, and optimizer step retain exact reconstruction checks." + ] + }, + { + "id": "TA-294", + "scope": "adapter optimizer invalid reshard source fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "World-size mismatch, divergent replicas, holes, overlap, dtype, shape, step, empty ranks, staged defects, and parameter fingerprints now form one fail-closed policy.", + "Resident optimizer no-mutation checks remain active for every defect that reaches staging." + ] + }, + { + "id": "TA-295", + "scope": "weight-sync receiver postprocess fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Environment precedence, FP8 KV-cache requirements, streaming backend flags, BF16 omission, and quantization-method detection now form one postprocess policy.", + "Both handler decisions and emitted backend configuration remain asserted." + ] + }, + { + "id": "TA-296", + "scope": "P2P direct-EP sender and tensor-collection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Replica sender selection, sender-to-EP ownership, and explicit non-sender collection suppression now form one direct-EP policy.", + "Default, round-robin, direct, and NCCL cases remain covered." + ] + }, + { + "id": "TA-297", + "scope": "weight-sync tied-parameter extraction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Declared identical ties, declared unequal parameters, and shared storage without declaration now form one alias policy.", + "Emitted buffers and tied-weight alias maps remain exact." + ] + }, + { + "id": "TA-298", + "scope": "Nemotron-H inference-unfuse fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Per-expert HF layout emission and fused-dense rejection now execute as one Nemotron-H conversion policy.", + "The production checkpoint handler remains the layout oracle." + ] + }, + { + "id": "TA-299", + "scope": "EP MoE gated and non-gated collection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Non-gated two-projection and gated three-projection expert collection now form one gating policy.", + "Projection names, shapes, splits, and values remain asserted." + ] + }, + { + "id": "TA-300", + "scope": "compiled-module wrapper name normalization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "MoE unfuse, backend broadcast, and Qwen linear-attention fusion now share one _orig_mod normalization policy.", + "Each final receiver namespace and transformed tensor remains checked." + ] + }, + { + "id": "TA-301", + "scope": "sparse-delta weight-sync fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fast-path configuration, cache metadata, prepacked-only rejection, and CPU FP8 targeting now form one sparse-delta policy.", + "Health, pause, post, resume, endpoint, and cache ordering remain covered." + ] + }, + { + "id": "TA-302", + "scope": "FP8 training linear replacement fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default replacement, parameter identity, output dtype, state dict, and fully qualified module tags now form one injection policy.", + "All replaced module types and names remain asserted." + ] + }, + { + "id": "TA-303", + "scope": "FP8 training linear recipe fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Global amax scales, FQN recipe overrides, and unknown override rejection now form one recipe policy.", + "Block size, SmoothQuant, correction, and output dtype propagation remain covered." + ] + }, + { + "id": "TA-304", + "scope": "FP8 training linear exclusion fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit module names and FQN globs now execute as one BF16-island exclusion policy.", + "Replacement counts and every retained or converted module remain asserted." + ] + }, + { + "id": "TA-305", + "scope": "FP8 linear CPU fallback fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Numerical fallback, float32 output, and fail-fast mode now form one CPU admission policy.", + "Production FP8Linear calls remain the tested boundary." + ] + }, + { + "id": "TA-306", + "scope": "FP8 linear error-profiler sampling fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Call caps, explicit flattened rows, and module-specific call-row selectors now form one CPU sampling policy.", + "The independent live-CUDA operand-breakdown gate remains separate." + ] + }, + { + "id": "TA-307", + "scope": "block-FP8 GEMM backend and scale-layout fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Block scales, rowwise weight scales, and torch scaled-mm backend parity now form one CUDA GEMM policy.", + "Both 64 and 128 block widths and explicit dequantized references remain." + ] + }, + { + "id": "TA-308", + "scope": "offline FP8 weight quantization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Weight-sync parity and partial-block zero-padding now form one quantization contract.", + "Exact FP8 bytes and scale tensors remain compared." + ] + }, + { + "id": "TA-309", + "scope": "offline fused-QKV export fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Valid QKV splitting, missing metadata rejection, and duplicate output rejection now form one export policy.", + "All Q, K, and V bytes and scales remain checked." + ] + }, + { + "id": "TA-310", + "scope": "offline MLA-A projection export fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Configured q_lora_rank fusion and metadata-absent split preservation now form one MLA export policy.", + "Fused and independent receiver names remain asserted." + ] + }, + { + "id": "TA-311", + "scope": "offline linear-attention name conversion fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Linear-attention metadata-driven fusion and metadata-absent split preservation now form one name policy.", + "All fused projections, convolution, norm, bias, and A_log tensors remain covered." + ] + }, + { + "id": "TA-312", + "scope": "offline QARL fold admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Selective QARL folding and mismatched block-size rejection now form one fold-admission policy.", + "The independent trained-logprob preservation gate remains separate." + ] + }, + { + "id": "TA-313", + "scope": "packing strategy and oversized-sample admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unknown settings, default error, legacy skip, aligned truncation, and HF-shift truncation now form one admission policy.", + "Every former sample payload and rejection remains exercised." + ] + }, + { + "id": "TA-314", + "scope": "cross-strategy packing correctness and utilization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Document multiset, valid tokens, position resets, capacity, and best-fit row count now form one strategy-invariant policy.", + "All sequential, best-fit, and balanced-DP data sets remain." + ] + }, + { + "id": "TA-315", + "scope": "balanced-DP packing behavior fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Zero dummies, balanced bins, full-row utilization, small-sample fallback, and DP1 equivalence now form one balanced-DP policy.", + "Every load and utilization threshold remains asserted." + ] + }, + { + "id": "TA-316", + "scope": "packing determinism and datum-order fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Repeatability, sequential identity, and reordered permutation/build order now form one ordering policy.", + "Routed-expert realignment evidence remains exact." + ] + }, + { + "id": "TA-317", + "scope": "packed token and teacher metadata fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "OPD alignment, HF shifts, hidden-state padding, RL padding, OPRD cache views, cache-base fallback, and nested RL schema now form one metadata policy.", + "Every token field, vector field, cache index, and ignore-index result remains asserted." + ] + }, + { + "id": "TA-318", + "scope": "disabled packing behavior fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Per-sample output, shifted targets, HF warning, explicit target preservation, and loss masking now form one disabled-mode policy.", + "Both flat and nested input conventions remain covered elsewhere in the same suite." + ] + }, + { + "id": "TA-319", + "scope": "Muon Gram-Newton-Schulz grouping fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Equal shapes, flattened matrix shapes, transpose-equivalent shapes, fused gate-up halves, and byte-limit chunking now form one grouping policy.", + "Every orthogonalizer input shape and exact parameter update remains asserted." + ] + }, + { + "id": "TA-320", + "scope": "Muon fused gate-up classification fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Gated versus non-gated experts, post-FSDP parameter replacement, and DeepSeek versus Nemotron model classification now form one fused-split policy.", + "Parameter identity and optimizer-group membership remain checked." + ] + }, + { + "id": "TA-321", + "scope": "FP8 MoE injection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Weight preservation, per-expert bias preservation, backend enablement, and unused-module summary now form one injection policy.", + "The model parameters remain identical objects after conversion." + ] + }, + { + "id": "TA-322", + "scope": "FP8 grouped same-NK forward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Block-loop reference, Triton shapes, nondefault block width, and precomputed sequence offsets now form one CUDA forward policy.", + "Every BF16 reference and tolerance remains unchanged." + ] + }, + { + "id": "TA-323", + "scope": "FP8 grouped same-MN weight-gradient fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Block-loop reference, scalar dispatch, precomputed sequence offsets, and Triton shape matrix now form one CUDA weight-gradient policy.", + "Dispatch spying and all BF16 comparisons remain active." + ] + }, + { + "id": "TA-324", + "scope": "FP8 scalar-Quack grouped forward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "BF16 fallback parity and per-expert scale selection through explicit sequence offsets now form one scalar-Quack policy.", + "Both numerical outputs remain checked on CUDA." + ] + }, + { + "id": "TA-325", + "scope": "FP8 MoE expert training-step fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Triton and scalar backends plus biased clamped-SwiGLU experts now form one expert training policy.", + "Outputs, gradients, master weights, and expert biases retain finite/update checks." + ] + }, + { + "id": "TA-326", + "scope": "canonical merged-LoRA fold fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Pinned arithmetic, shared-factor expansion, expert sharding, linear orientation, and zero-delta identity now form one fold policy.", + "All exact tensor comparisons and dtypes remain asserted." + ] + }, + { + "id": "TA-327", + "scope": "folded-weight straight-through gradient fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shared-factor GKN cases, fused gate-up, and linear orientation now form one straight-through gradient policy.", + "Reference autograd comparisons remain unchanged; FSDP gradient dtype stays independent." + ] + }, + { + "id": "TA-328", + "scope": "LoraLinear merged-forward selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Exact merged bytes, legacy selection, module isolation, and gradient parity now form one selection policy.", + "Step and runtime cache invalidation remains an independent lifecycle report." + ] + }, + { + "id": "TA-329", + "scope": "MoE merged-weight and cache fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical gate-up/down folds and parameter-version cache identity now form one merged-weight policy.", + "Exact values, object reuse, and invalidation remain checked." + ] + }, + { + "id": "TA-330", + "scope": "fused-expert merged-LoRA admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Automatic support and fail-closed unmerged invocation now form one fused-expert admission policy.", + "Both accepted and rejected modes remain exercised." + ] + }, + { + "id": "TA-331", + "scope": "native-EP merged-LoRA routing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Gradient-bearing masked routing and no-grad canonical-fold filtering now form one native-EP policy.", + "Keyword routing, local expert IDs, folded weights, and output identity remain asserted." + ] + }, + { + "id": "TA-332", + "scope": "merged-LoRA trunk-wrap composition fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fail-closed wrapping and enabled composition now form one trunk-wrap policy.", + "The exact contract and wrapper state remain checked." + ] + }, + { + "id": "TA-333", + "scope": "API session creation and activity fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Usable session IDs, deterministic heartbeat refresh, and canonical LoRA storage now form one lifecycle policy.", + "Follow-up save behavior and session registry state remain asserted." + ] + }, + { + "id": "TA-334", + "scope": "Tinker weights-info compatibility fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy flat rank serialization and create-model dictionary storage now form one weights-info compatibility policy.", + "Both direct and create-model paths retain their response assertions." + ] + }, + { + "id": "TA-335", + "scope": "API create-model worker registration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full-weight empty overrides and LoRA worker registration now form one create-model policy.", + "Materialization payloads, worker session specs, and stored optimizer configuration remain checked." + ] + }, + { + "id": "TA-336", + "scope": "API optimizer payload and learning-rate fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Native and Tinker payloads plus request, session, server, create-model, and missing-default LR resolution now form one optimizer policy.", + "All optimizer fields, metrics, defaults, and failure behavior remain asserted." + ] + }, + { + "id": "TA-337", + "scope": "training-simulator topology and shape-accounting fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Balanced routing, DP topology resolution, and sequence-parallel shape accounting now form one topology policy.", + "Every count, batch size, local token, and routed-slot assertion remains." + ] + }, + { + "id": "TA-338", + "scope": "observed benchmark ingestion and planning fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Structured log parsing, resolved fit/OOM ingestion, and tight-margin scenario planning now form one observed-data policy.", + "The same fixture now flows from logs through calibrated feasibility without duplicate setup." + ] + }, + { + "id": "TA-339", + "scope": "training configuration and model-metadata fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Configuration fingerprints, HF-cache metadata, and known-model fallback now form one metadata policy.", + "Hashes, topology fields, parsed architecture fields, and source labels remain checked." + ] + }, + { + "id": "TA-340", + "scope": "Qwen235 calibration ingestion and evaluation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Markdown fit/OOM extraction and leave-one-out GA evaluation now form one calibration policy.", + "Measured rows, labels, throughput, topology, errors, and OOM status remain asserted." + ] + }, + { + "id": "TA-341", + "scope": "Qwen235 calibrated scenario fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Observed GA selection, asymptotic extrapolation, and exact OOM-boundary rejection now form one calibrated-scenario policy.", + "Raw and risk-adjusted ranking, memory basis, remeasurement flags, and infeasibility remain checked." + ] + }, + { + "id": "TA-342", + "scope": "Qwen235 topology what-if fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "TP extrapolation, automatic parallelism sweep, and long-context CP admission now form one topology what-if policy.", + "Candidate spaces, conservative penalties, calibration scope, and OOM risk flags remain asserted." + ] + }, + { + "id": "TA-343", + "scope": "built-in simulator calibration-pack fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Qwen3.6 report matching and Qwen235 fit/OOM replay now form one built-in pack policy.", + "Correctness gating, support status, timing coverage, accuracy, and recall remain checked." + ] + }, + { + "id": "TA-344", + "scope": "portable analytical-ledger fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Portable FLOP/activation/communication coverage, dense MLP activations, and cross-node expert-FSDP normalization now form one ledger policy.", + "Exact status labels, byte terms, activation size, pass normalization, and positive totals remain asserted." + ] + }, + { + "id": "TA-345", + "scope": "teacher-head loading and storage fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Direct safetensors, tied embeddings, sharded stores, and cross-shard row views now form one head-persistence policy.", + "Full tensors, shard row counts, and sliced ranges remain compared exactly." + ] + }, + { + "id": "TA-346", + "scope": "teacher-head manager residency fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Teacher replacement, dtype replacement, and store prefetch now form one manager-residency policy.", + "Resident teacher identity, dtype, and loaded values remain asserted." + ] + }, + { + "id": "TA-347", + "scope": "teacher activation-cache selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-2 indexing, rank-3 token and layer selection, host/device paths, reuse, and dtype reload now form one selection policy.", + "Every source slice, result shape, dtype, and cache identity remains checked." + ] + }, + { + "id": "TA-348", + "scope": "teacher activation-cache async and admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Async prefetch success and negative/high index rejection now form one cache-admission policy.", + "Both failure messages and the prefetched output remain asserted." + ] + }, + { + "id": "TA-349", + "scope": "Mooncake tensor codec fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Byte round trips for four dtypes and canonical string aliases now form one tensor-codec policy.", + "Exact values, dtypes, accepted aliases, and unsupported dtype rejection remain." + ] + }, + { + "id": "TA-350", + "scope": "Mooncake hidden transport fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Metadata emission, rank-2 retrieval, and rank-3 layer retrieval now form one transport policy.", + "Storage key, schema fields, token counts, shapes, dtypes, and exact tensors remain checked." + ] + }, + { + "id": "TA-351", + "scope": "Mooncake teacher activation-consumer fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-2, rank-3, and multi-teacher cache consumption now form one integration policy.", + "Teacher routing, selected values, output shapes, and cache closure remain exercised." + ] + }, + { + "id": "TA-352", + "scope": "Mooncake metadata admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy entries, missing objects, size mismatches, and malformed metadata now form one fail-closed policy.", + "All original invalid payloads and error boundaries remain." + ] + }, + { + "id": "TA-353", + "scope": "Mooncake store lifecycle and configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Best-effort object removal and explicit-versus-environment configuration now form one store lifecycle policy.", + "The suffixed removal key and both configuration precedence outcomes remain asserted." + ] + }, + { + "id": "TA-354", + "scope": "trainer gradient-clipping fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Regular, DistSignSGD, disabled, and scale-factor cases now form one clipping policy.", + "Norms, gradient values, nonpositive thresholds, and rank-count factors remain checked." + ] + }, + { + "id": "TA-355", + "scope": "trainer metadata-counting fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Valid-token reduction, target-token precedence, batched voter reduction, and empty input now form one counting policy.", + "Reduction count, operation, group, device, active microbatches, and voter totals remain asserted." + ] + }, + { + "id": "TA-356", + "scope": "explicit trainer gradient-synchronization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default SP reduction, adapter-owned exclusions, lm-head exclusions, and optional DTensor skipping now form one synchronization policy.", + "Every selected tensor, group, and SUM operation remains checked." + ] + }, + { + "id": "TA-357", + "scope": "lm-head TP synchronization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Single-rank gradient suppression and multi-rank marked-parameter broadcast now form one lm-head TP policy.", + "No-op reduction, global source rank, group, and broadcast tensor remain asserted." + ] + }, + { + "id": "TA-358", + "scope": "checkpoint object-broadcast transport fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "NCCL tensor serialization, default-device selection, and weight-load group routing now form one transport policy.", + "Payload identity, CUDA device index, source rank, and group remain asserted." + ] + }, + { + "id": "TA-359", + "scope": "rank-zero checkpoint broadcast-loading fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Nonzero-rank shard resolution and leader-side handler-filtered prefetch now form one rank-zero loading policy.", + "Load calls, skipped and loaded keys, dispatch names, and batch metadata remain checked." + ] + }, + { + "id": "TA-360", + "scope": "checkpoint state-dict resolution fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Remote rank-zero path resolution and node-local directory resolution now form one source policy.", + "Broadcast suppression and exact iterator paths remain asserted." + ] + }, + { + "id": "TA-361", + "scope": "grouped checkpoint expert-routing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense/expert fanout, HF fused experts, and FFN source formats now form one grouped-routing policy.", + "Handler inputs, prefetch partitions, converted names, transfers, and dispatch targets remain checked." + ] + }, + { + "id": "TA-362", + "scope": "grouped checkpoint group-fallback fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local dense loading plus permissive and strict missing-EP-group outcomes now form one fallback policy.", + "Collective suppression, rank-zero fallback, and fail-closed strict behavior remain asserted." + ] + }, + { + "id": "TA-363", + "scope": "strict checkpoint postprocessing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Missing, unexpected, duplicate, persistent-buffer, and complete-coverage outcomes now form one strict policy.", + "All diagnostic names and successful buffer dispatch remain checked." + ] + }, + { + "id": "TA-364", + "scope": "cautious decay primitive and SignSGD fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Direct zero, ordinary, and sign-masked decay plus SignSGD integration now form one primitive policy.", + "Aligned, misaligned, and zero-update coordinates retain exact expected values." + ] + }, + { + "id": "TA-365", + "scope": "AnyPrecisionAdamW cautious-decay fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy decoupled decay and cautious coordinate masking now form one AnyPrecision decay policy.", + "Both explicit first-step references remain unchanged." + ] + }, + { + "id": "TA-366", + "scope": "AnyPrecisionAdamW state-strategy fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Chunked denominator with ordinary/Kahan state and gradient-reuse CPU offload now form one state policy.", + "Parameters, moments, compensation, cleared gradients, and device placement remain checked." + ] + }, + { + "id": "TA-367", + "scope": "Muon cautious-decay fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ordinary decay, post-Newton-Schulz masking, and AdamW fallback masking now form one Muon policy.", + "All three explicit update references and tolerances remain." + ] + }, + { + "id": "TA-368", + "scope": "optimizer-builder cautious routing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Supported families, ordinary AdamW, unsupported SGD, kwarg rejection, and AnyPrecision options now form one builder policy.", + "Optimizer classes, group fields, accepted kwargs, and rejection messages remain asserted." + ] + }, + { + "id": "TA-369", + "scope": "GDN convolution forward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Qwen3.5 shapes, variable-length and batched layouts, and repeat determinism now form one CUDA forward policy.", + "All outputs remain bitwise compared to serving invocation or a repeated run." + ] + }, + { + "id": "TA-370", + "scope": "GDN convolution backward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Depthwise-autograd parity and repeated backward determinism now form one CUDA backward policy.", + "Input and convolution-weight gradients remain compared." + ] + }, + { + "id": "TA-371", + "scope": "end-to-end GDN block fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Eager output/gradient parity and exact-contract repeat determinism now form one end-to-end policy.", + "Every named parameter gradient and the bounded forward result remain checked." + ] + }, + { + "id": "TA-372", + "scope": "GDN packed-weight and armed-routing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Bitwise-neutral QKV weight packing and exact-contract dispatch now form one construction/routing policy.", + "All packed slices, invocation count, and output shape remain asserted." + ] + }, + { + "id": "TA-373", + "scope": "GDN exact-contract state-lifecycle fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Exact/ordinary module isolation and checkpoint recomputation now form one contract-state policy.", + "Thread-local state is checked during each call and after both ordinary and recomputed execution." + ] + }, + { + "id": "TA-374", + "scope": "checkpoint reference-state and QARL buffer fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ordinary persistent filtering, QARL buffer inclusion, and strict/non-strict QARL mismatch now form one buffer policy.", + "All parameter, buffer, shape, counter, metadata, and mismatch fields remain asserted." + ] + }, + { + "id": "TA-375", + "scope": "pipeline checkpoint key-contract fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Non-pipeline local keys, pipeline stage union, and metadata compatibility now form one key policy.", + "Collective suppression, union ordering, metadata counts, and compatibility results remain checked." + ] + }, + { + "id": "TA-376", + "scope": "pipeline LoRA checkpoint compatibility fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Base-to-LoRA and LoRA-only pipeline checkpoints now form one compatibility policy.", + "Load modes and exact missing LoRA/non-LoRA key sets remain asserted." + ] + }, + { + "id": "TA-377", + "scope": "distributed-checkpointer metadata admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Missing optimizer state and present optimizer metadata keys now form one load-admission policy.", + "Selected state entries, planner, reader, no-dist flag, and optimizer load keys remain checked." + ] + }, + { + "id": "TA-378", + "scope": "distributed-checkpointer load-group fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "No-dist pipeline validation, no-dist non-pipeline loading, and explicit DCP groups now form one load-group policy.", + "Validation and DCP group identities plus no-dist behavior remain asserted." + ] + }, + { + "id": "TA-379", + "scope": "distributed-checkpointer save-group fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Synchronous pipeline metadata reuse and asynchronous non-pipeline group isolation now form one save-group policy.", + "DCP and metadata process-group identities remain checked." + ] + }, + { + "id": "TA-380", + "scope": "optimizer checkpoint-state filtering fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Checkpoint load-key filtering and per-child multi-optimizer filtering now form one optimizer-state policy.", + "State keys, parameter groups, strict flags, and child optimizer assignments remain asserted." + ] + }, + { + "id": "TA-381", + "scope": "RoPE Class-B selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Serving provenance, selector reset, canonical GLM defaults/opt-out, and non-GLM opt-in now form one selection policy.", + "Resolved modes, global state, float32 tables, and rejection messages remain checked." + ] + }, + { + "id": "TA-382", + "scope": "canonical GLM numerical-program fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Exact program resolution, incompatible override and CE rejection, and non-GLM defaults now form one GLM policy.", + "Every resolved field, override case, CE mode, and default remains asserted." + ] + }, + { + "id": "TA-383", + "scope": "exact Qwen3.5 numerical-program fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense/MoE resolution, non-Qwen v2 rejection, incompatible override rejection, and CE rejection now form one Qwen policy.", + "All certified fields, override cases, and exact rejection messages remain." + ] + }, + { + "id": "TA-384", + "scope": "OPD metric aggregation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Namespace, weighted aggregation, zero-valid extrema, loss-group reduction, and empty-rank key seeding now form one metric policy.", + "All values, operations, groups, profile keys, and debug-key filtering remain checked." + ] + }, + { + "id": "TA-385", + "scope": "OPD packed cache and weight-shaping fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Valid-row splitting, OPRD tail weights, and packed-batch hidden chunks now form one packed-shaping policy.", + "Padding removal, position resets, indices, weights, and chunk boundaries remain asserted." + ] + }, + { + "id": "TA-386", + "scope": "OPD microbatch loss-execution fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Per-teacher cache masking and FSDP lm-head anchoring now form one execution policy.", + "Loss finiteness, metrics, timings, anchor call shape, and student/hidden gradients remain checked." + ] + }, + { + "id": "TA-387", + "scope": "teacher hidden-cache contributor fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CP suppression, distinct EP slices, and legacy duplicate-EP mode now form one contributor policy.", + "All contributor keys and suppression outcomes remain asserted." + ] + }, + { + "id": "TA-388", + "scope": "teacher hidden-cache distributed assembly fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unified SP gathering, gathered-label trimming, and cross-rank batch assembly now form one distributed policy.", + "Groups, unpadding, saved rows, token counts, and per-sample cache indices remain checked." + ] + }, + { + "id": "TA-389", + "scope": "teacher hidden-cache Mooncake integration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Mooncake metadata/storage and downstream activation-cache indexing now form one producer-consumer policy.", + "Schema fields, token counts, stored bytes, and selected consumer rows remain asserted." + ] + }, + { + "id": "TA-390", + "scope": "OPD debug artifact fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Vocab-parallel loss contribution and packed teacher-segment JSONL now form one debug-artifact policy.", + "All local/group metrics, segment provenance, cache statistics, and component tensors remain checked." + ] + }, + { + "id": "TA-391", + "scope": "evicted-adapter auto-load fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Checkpoint restore, fresh broadcast, synchronized materialization failure, fresh-state rejection, and restore rollback now form one auto-load policy.", + "Registration/load calls, checkpoint paths, broadcasts, cross-rank errors, and rollback state remain checked." + ] + }, + { + "id": "TA-392", + "scope": "explicit adapter-state load and path-admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All-rank optimizer restore and output-root confinement for load, evicted lookup, and save now form one admission policy.", + "Payload fields, calls, success result, all rejection messages, and absence of escaped writes remain asserted." + ] + }, + { + "id": "TA-393", + "scope": "rank-zero adapter restore routing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Auto-load and explicit-load rank-zero broadcast modes now form one routing policy.", + "Registration, restore invocation, checkpoint path, and suppression of all-rank loads/broadcasts remain checked." + ] + }, + { + "id": "TA-394", + "scope": "rank-zero sharded adapter restore fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP shard slicing, session-spec mismatch, and topology-specific optimizer rejection now form one restore policy.", + "Local tensor bytes, steps, LR, transactional metadata, and unchanged state on both failures remain asserted." + ] + }, + { + "id": "TA-395", + "scope": "adapter-state load failure fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Cross-rank restore failure, pipeline-parallel rejection, and auto-registered session rollback now form one failure policy.", + "Exact errors and cleanup of adapters/session specs remain checked." + ] + }, + { + "id": "TA-396", + "scope": "adapter registration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Adapter/materialized-session success, cross-rank rollback, and worker registration failure now form one registration policy.", + "Session specs, LR calls, broadcasts, rollback, and exception text remain asserted." + ] + }, + { + "id": "TA-397", + "scope": "RMSNorm family-admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unknown families, residual misuse, forced family flips, unsupported fused-add, and zero-centered family misuse now form one admission policy.", + "Construction, module call, and funnel rejection boundaries remain." + ] + }, + { + "id": "TA-398", + "scope": "Qwen RMSNorm site-declaration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense Qwen decoder sites and shared-attention Q/K norms now form one declaration policy.", + "Layer-zero, later-layer, post-attention, Q, and K family assignments remain checked." + ] + }, + { + "id": "TA-399", + "scope": "RMSNorm family declaration-tripwire fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Undeclared warning, required-family failure, declared success, and legacy-explicit silence now form one CUDA tripwire policy.", + "Warning and error messages plus both admitted families remain exercised." + ] + }, + { + "id": "TA-400", + "scope": "RMSNorm family-funnel fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Two-shape legacy parity, zero-centered folding, and family-difference vitality now form one bitwise CUDA funnel policy.", + "All exact outputs and the rare-difference bound remain asserted." + ] + }, + { + "id": "TA-401", + "scope": "RMSNorm family module-dispatch fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Trunk-contract dispatch/warning and declared-versus-legacy calls across three modes now form one module policy.", + "No-residual, residual-tree, and fused residual outputs remain bitwise checked." + ] + }, + { + "id": "TA-402", + "scope": "pipeline bubble formula fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Nonzero schedules, zero-bubble schedules, single-stage behavior, and invalid inputs now form one analytic policy.", + "Every formula value and rejection case remains." + ] + }, + { + "id": "TA-403", + "scope": "pipeline P2P byte-estimation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Middle/edge stages, forward-only, same-rank adjacency, and unpopulated metadata now form one P2P accounting policy.", + "All flow counts, microbatch multipliers, and the unknown result remain asserted." + ] + }, + { + "id": "TA-404", + "scope": "pipeline profiler patch-lifecycle fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Passthrough/restoration, double-patch rejection, single-stage discovery, and report-before-step failure now form one patch policy.", + "Stage calls, instance attributes, cleanup, and both rejection messages remain checked." + ] + }, + { + "id": "TA-405", + "scope": "optimizer, packing, and numerical argument fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "SignSGD/DistSignSGD controls and complete Muon kwargs now form one optimizer argument policy.", + "Packing, load mode, dtypes, numerical flags, and every Muon option remain asserted." + ] + }, + { + "id": "TA-406", + "scope": "checkpoint argument compatibility fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy EP/checkpoint aliases, automatic checkpoint resolution, and optimizer-resume defaults/overrides now form one policy.", + "Resolved checkpoint path, checkpoint method, EP placement, and all load_optimizer values remain checked." + ] + }, + { + "id": "TA-407", + "scope": "FP8 argument configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Optional overrides, native/Nemo aliases, fail-fast fallback, and unsupported vLLM knobs now form one FP8 policy.", + "All accepted fields and five receiver-side rejection boundaries remain." + ] + }, + { + "id": "TA-408", + "scope": "low-precision training argument fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "ModelOpt rejection, GLM block-FP8 QLoRA, QARL normalization, and the full low-precision conflict matrix now form one mode policy.", + "Accepted quantization fields plus LoRA, QLoRA, FP8, calibration, MTP, and Mamba rejections remain." + ] + }, + { + "id": "TA-409", + "scope": "pipeline FQN partitioning fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default/Qwen names, single/virtual stages, pinned endpoints, weighted splits, and infeasible layouts now form one partition policy.", + "All layer coverage, contiguity, counts, endpoints, and errors remain asserted." + ] + }, + { + "id": "TA-410", + "scope": "pipeline stage-placement fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Torch-reference mappings, exact ownership, single-style admission, and V-style endpoints now form one placement policy.", + "Loop, V, single-stage, and invalid virtual-stage cases remain." + ] + }, + { + "id": "TA-411", + "scope": "pipeline schedule metadata and admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Schedule style/splitting metadata and virtual-stage/microbatch validation now form one schedule policy.", + "All six schedules, unknown names, valid layouts, and invalid constraints remain checked." + ] + }, + { + "id": "TA-412", + "scope": "Mamba2 mixer HF-parity fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full/partial single chunks, masked input, and BF16 finite output now form one mixer policy.", + "HF outputs, input/parameter gradients, dtype, and finiteness remain asserted." + ] + }, + { + "id": "TA-413", + "scope": "SSD recurrence fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Multi-chunk output/gradient parity, missing-D behavior, and complete mixer recurrence now form one SSD policy.", + "Both sequence lengths, every input gradient, and projected mixer output remain checked." + ] + }, + { + "id": "TA-414", + "scope": "packed SSD and Mamba2 fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Boundary-aligned/unaligned SSD, single-sequence identity, causal convolution, mixer parity, full-row identity, and batch rejection now form one packed policy.", + "All outputs, input/parameter gradients, sequence boundaries, and error messages remain." + ] + }, + { + "id": "TA-415", + "scope": "optional SSD kernel-parity fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense and packed-sequence mamba_ssm comparisons now form one optional CUDA kernel policy.", + "Outputs and the applicable dense/packed gradient sets retain their tolerances." + ] + }, + { + "id": "TA-416", + "scope": "token diagnostic selection and boundary fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shape, ranking, disabled, absent-label, all-ignored, and top-k clamping behavior now form one selection policy.", + "Every retained field and boundary result remains asserted." + ] + }, + { + "id": "TA-417", + "scope": "token diagnostic log-probability reference fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Loss-path and raw-weight reference comparisons now form one numerical cross-check policy.", + "Exact zero deltas and deliberately nonzero scaled-head deltas remain asserted." + ] + }, + { + "id": "TA-418", + "scope": "hidden diagnostic summary sampling examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Layer summaries, explicit indices, invalid indices, and ordered component summaries now report as one policy.", + "The redundant all-indices spelling and a second component-width implementation example were removed." + ] + }, + { + "id": "TA-419", + "scope": "hidden component hook implementation examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "MoE callback and shared-expert capture now form one hook-integration policy alongside the independent dense equation policy.", + "A native-MoE example that only pinned internal ordering constants was removed." + ] + }, + { + "id": "TA-420", + "scope": "fused selected-logprob forward and backward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forward, finite-output, and input/weight/bias gradient parity now execute from the same eager-reference cases.", + "BF16/no-bias/default-temperature and FP32/bias/nondefault-temperature branches remain covered." + ] + }, + { + "id": "TA-421", + "scope": "fused selected-logprob input-gradient fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Frozen-head gradient parity and all-frozen output detachment now form one input-gradient policy.", + "Hidden-state gradients and absence of weight/bias gradients remain checked." + ] + }, + { + "id": "TA-422", + "scope": "fused selected-logprob irregular shape repetitions", + "decision": "remove", + "status": "applied", + "evidence": [ + "A non-tile-aligned 37x130x777 case retains tail-shape parity.", + "The single-row example and a second large-vocabulary example duplicated boundaries covered by the production-vocabulary policy." + ] + }, + { + "id": "TA-423", + "scope": "fused loss dispatcher fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Per-token CE, scalar causal-LM, and quack-linear per-token dispatch now form one interface policy.", + "Scalar loss, per-token loss, and per-token log-probability parity remain checked." + ] + }, + { + "id": "TA-424", + "scope": "fused loss production-vocabulary repetitions", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The Qwen integration path and largest GPT-OSS direct path retain the two named production boundaries.", + "An unnamed 100000-vocabulary repetition was removed while finiteness, eager parity, and finite gradients remain." + ] + }, + { + "id": "TA-425", + "scope": "duplicate fused-logit peak-memory probe", + "decision": "remove", + "status": "applied", + "evidence": [ + "The causal-LM integration probe is the stronger regression because it fails on eager dispatcher fallthrough.", + "The direct frozen-weight probe repeated the same less-than-half-full-tile heuristic without covering that dispatch boundary." + ] + }, + { + "id": "TA-426", + "scope": "DistSignSGD framework-behavior unit examples", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained optimizer policy checks the production preaggregated update and decoupled decay ordering.", + "Base-Optimizer state-dict round-tripping and an unsupported sparse-gradient example did not exercise distributed sign behavior." + ] + }, + { + "id": "TA-427", + "scope": "DistSignSGD reduce-scatter fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Direct sign voting, forced SUM, async propagation, and SP-before-sign ordering now form one communication policy.", + "Exact inputs, outputs, operations, and process-group routing remain asserted." + ] + }, + { + "id": "TA-428", + "scope": "DistSignSGD unsupported-topology fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "HSDP, sequence-parallel folding, expert parallelism, and unmanaged DTensor rejection now form one admission policy.", + "Each topology still checks its specific failure message." + ] + }, + { + "id": "TA-429", + "scope": "DistSignSGD optimizer-builder fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FSDP2 construction, hook configuration, decay-group ownership, and non-FSDP2 rejection now form one builder policy.", + "Optimizer type, exact parameters, decay values, and rejection remain checked." + ] + }, + { + "id": "TA-430", + "scope": "SGLang MoE automatic-resolution fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP1, EP8, explicit overrides, device/stack requirements, module eligibility, and log-once behavior now form one resolution policy.", + "Every admitted and rejected regime plus warning and logging behavior remains checked." + ] + }, + { + "id": "TA-431", + "scope": "SGLang MoE block-dispatch fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Disabled, full-forward, experts-only, and FP64-precedence paths now form one dispatch policy.", + "Outputs, shapes, call routing, and precedence remain asserted." + ] + }, + { + "id": "TA-432", + "scope": "SGLang MoE fused-expert admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Expert semantics, missing routing inputs, positive SwiGLU limits, and trainable bias/activation guards now form one admission policy.", + "An install-state-dependent missing-SGLang example that skipped whenever SGLang was present was removed." + ] + }, + { + "id": "TA-433", + "scope": "SGLang masked-training gradient fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Partially masked, all-valid, and all-masked expert routing now form one CUDA gradient policy.", + "Forward outputs and every input, routing-weight, gate-up, and down gradient remain bitwise checked." + ] + }, + { + "id": "TA-434", + "scope": "SGLang MoE weight-mode and kernel-layout fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Cache lifecycle, mode resolution, zero-copy strided views, serving layout, and FP32 routing weights now form one policy.", + "Storage identity, invalidation, environment precedence, tensor layout, and kernel flags remain checked." + ] + }, + { + "id": "TA-435", + "scope": "vendored strided MoE adapter fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Accepted storage layouts and split gate/up delegation now form one adapter policy.", + "Contiguous, transpose-view, sliced-layout, and interleaved-gate rejection cases remain." + ] + }, + { + "id": "TA-436", + "scope": "SGLang runtime-context fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Publication, compatible reuse, nondeterministic rejection, and fused-reduction rejection now form one context policy.", + "Exact published arguments, role, no-op behavior, and both incompatibilities remain asserted." + ] + }, + { + "id": "TA-437", + "scope": "mocked EP1 SGLang auto-dispatch CUDA example", + "decision": "remove", + "status": "applied", + "evidence": [ + "CPU resolution and dispatch policies already cover the branch decision and explicit escape hatch.", + "The retained real-kernel CUDA policy proves automatic resolution, deterministic output, and explicit-mode bitwise parity." + ] + }, + { + "id": "TA-438", + "scope": "batch-invariant trunk-linear wrapping fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Selection counts, exclusions, idempotence, routed-expert skipping, and unsupported-module admission now form one CPU policy.", + "Every projection count, ownership marker, and rejection class remains checked." + ] + }, + { + "id": "TA-439", + "scope": "batch-invariant trunk-linear forward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Persistent-GEMM parity, global-interpose parity, batch invariance, dtype admission, and interpose conflict now form one CUDA forward policy.", + "Bias and no-bias outputs remain bitwise checked." + ] + }, + { + "id": "TA-440", + "scope": "global batch-invariant interpose gradient fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Training rejection, no-grad numerical behavior, and grad-enabled operation on grad-free inputs now form one policy.", + "MM, RMSNorm, BMM, log-softmax, mean, and permitted inference outputs remain exercised." + ] + }, + { + "id": "TA-441", + "scope": "MiniMax M3 configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Top-level HF adaptation, local loading/registration, and native-config round-trip now form one configuration policy.", + "Architecture, attention, MoE, sparse-attention, multimodal metadata, and registry fields remain checked." + ] + }, + { + "id": "TA-442", + "scope": "MiniMax M3 activation and router fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "SwiGLU-OAI clamping and sigmoid-router selection/weighting now form one primitive policy.", + "The analytical activation and bias-for-selection-only routing references remain exact." + ] + }, + { + "id": "TA-443", + "scope": "MiniMax M3 text-runtime fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Tiny forward/backward, text-only admission, reserved-token rejection, and unsupported parallelism now form one runtime policy.", + "Loss, output shapes, lm-head gradient, and all rejection messages remain." + ] + }, + { + "id": "TA-444", + "scope": "MiniMax M3 checkpoint fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Language-weight mapping, multimodal skipping, EP ownership, and grouped-loader aliases now form one checkpoint policy.", + "Dense, routed-expert, router, correction-bias, local-shard, and raw-key behavior remains checked." + ] + }, + { + "id": "TA-445", + "scope": "runner dispatcher batch-distribution fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ordinary DP, EP slicing/padding, CP sharing, and both legacy duplicate modes now form one distribution policy.", + "Exact batches, routing payload slices, dummy labels, and EP-FSDP rank ownership remain checked." + ] + }, + { + "id": "TA-446", + "scope": "runner dispatcher routing-payload fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Filesystem and Mooncake slicing, world-size-one loading, pickle rejection, and symlink rejection now form one transport policy.", + "Loaded tensors, object keys, and trust-boundary errors remain asserted." + ] + }, + { + "id": "TA-447", + "scope": "runner dispatcher packing examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Balanced-DP zero-dummy behavior and sequential underfill now form one packing policy.", + "Lockstep rounds, exact datum accounting, and dummy fallback remain checked." + ] + }, + { + "id": "TA-448", + "scope": "runner dispatcher per-token merge fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Logical ordering, empty-slice removal, coherent CP deduplication, and replica disagreement now form one merge policy.", + "Every merged field and the disagreement diagnostic remain asserted." + ] + }, + { + "id": "TA-449", + "scope": "runner dispatcher row-batching provenance fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-local row batching and unmerged-row provenance now form one policy.", + "Token layout, sequence boundaries, source batch/request IDs, sample counts, and spans remain checked." + ] + }, + { + "id": "TA-450", + "scope": "P2P prepare-payload fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Transport/engine metadata and cache-invalidation suppression now form one prepare policy.", + "Endpoint, group, sender identity, rank, and cache-mode behavior remain asserted." + ] + }, + { + "id": "TA-451", + "scope": "P2P initialize-fanout fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Multi-endpoint map aggregation, HTTP/remote failures, and receiver cleanup after partial preparation now form one fanout policy.", + "Locator ownership, receiver sessions, cleanup endpoints, and completion payloads remain checked." + ] + }, + { + "id": "TA-452", + "scope": "P2P cached-prepare fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Map reuse, session change, partial endpoint refresh, unsupported-flag retry, and endpoint-local retry now form one cache policy.", + "Request flags, retry counts, locator replacement, endpoint retention, and session IDs remain checked." + ] + }, + { + "id": "TA-453", + "scope": "P2P complete-sync fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Cached prepare-state preservation and tied-weight alias forwarding now form one completion policy.", + "Tensor maps, session/debug state, and alias payloads remain asserted." + ] + }, + { + "id": "TA-454", + "scope": "P2P Qwen3.5 linear-attention slicing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Convolution squeezing, generic QKV slicing, local state-vector ownership, and receiver-dtype conversion now form one slicing policy.", + "Exact shapes, values, TP ranges, byte sizes, and dtypes remain checked." + ] + }, + { + "id": "TA-455", + "scope": "P2P engine hostname and fallback fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Hostname precedence/DNS fallback and direct Mooncake construction without the SGLang wrapper now form one engine policy.", + "Every address source and initialized engine field remains checked." + ] + }, + { + "id": "TA-456", + "scope": "FP8 synchronization numerical fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "BF16 islands, Slime blockwise parity, and partial-layout zero padding now form one synchronization contract.", + "Exact FP8 bytes, scales, dequantized values, exclusions, and last-element-padding distinction remain." + ] + }, + { + "id": "TA-457", + "scope": "FP8 adapter-merge fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense LoRA, QLoRA linear, and QLoRA MoE merging now form one adapter synchronization policy.", + "Extraction ownership, merged weights, emitted names, exact quantization, and nonzero dequantized deltas remain checked." + ] + }, + { + "id": "TA-458", + "scope": "FP8 projection-selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Supported families, optional weight suffixes, receiver passthrough entries, and broad-selector activation now form one selection policy.", + "Every quantized, scaled, and identity-preserved tensor remains asserted." + ] + }, + { + "id": "TA-459", + "scope": "FP8 stack and existing-dtype fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Stack-versus-single parity and already-quantized passthrough now form one tensor policy.", + "Per-slice quantized values, scales, and identity passthrough remain checked." + ] + }, + { + "id": "TA-460", + "scope": "FP8 CPU expert-projection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "HF naming/scales, partial-layout padding, deferred quantization, and module exclusions now form one CPU expert policy.", + "Bytes, shapes, dtypes, expert indices, timings, exact scales, and passthrough behavior remain." + ] + }, + { + "id": "TA-461", + "scope": "FP8 CPU workspace fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Staging/reuse, streaming transfer, flush/reset, and empty-final-flush metadata now form one workspace lifecycle policy.", + "Storage identity, bucket contents, flush/version routing, capacity reset, and timing fields remain asserted." + ] + }, + { + "id": "TA-462", + "scope": "FP8 GPU synchronization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CPU/CUDA targets, stack CPU parity, direct-EP expert parity, and module exclusions now form one GPU policy.", + "Output devices, dtypes, timings, names, scales, quantized values, and BF16 passthrough remain checked." + ] + }, + { + "id": "TA-463", + "scope": "adapter optimizer construction and persistence fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "SignSGD construction, current-LR persistence, Adam kwarg normalization, and Muon group-LR preservation now form one policy.", + "Checkpoint location/metadata, optimizer types, state fields, hyperparameters, and LR behavior remain checked." + ] + }, + { + "id": "TA-464", + "scope": "adapter gradient-ownership compile and admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit compilation, pending-gradient rejection, uncompiled failure, invalid scale/gradient rejection, collective freedom, and sync exclusions now form one policy.", + "Plan parameters, unchanged state, errors, capture result, and exclusion ownership remain asserted." + ] + }, + { + "id": "TA-465", + "scope": "adapter capture commit fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Preallocated staging, direct-DTensor completion, and atomic prevalidation now form one commit policy.", + "Storage identity, model-gradient ownership, completed tensors, and all-or-nothing numerator mutation remain checked." + ] + }, + { + "id": "TA-466", + "scope": "adapter coordinator checkpoint materialization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit checkpoint loading and automatic evicted-adapter loading now form one coordinator policy.", + "Checkpoint session specs, optimizer type, learning rate, loaded path, and materialization remain checked." + ] + }, + { + "id": "TA-467", + "scope": "adapter checkpoint session-compatibility fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Optimizer mismatch rejection, explicit LR override, weights-only optimizer mismatch, and same-contract LR restoration now form one policy.", + "Weights, optimizer ownership, session spec, live LR, and rejection behavior remain asserted." + ] + }, + { + "id": "TA-468", + "scope": "GLM5 configuration and construction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "RMSNorm declarations, HF adaptation, local loading, unsafe-value rejection, default shape, and ring-attention rejection now form one policy.", + "MLA, MoE, DSA, MTP, family, architecture, topology, and safety fields remain checked." + ] + }, + { + "id": "TA-469", + "scope": "GLM5 indexer construction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Projection shapes and exact-contract FP32 head projection now form one construction policy.", + "All four projection dimensions, family weights, kernel operands, scaling, and output bits remain checked." + ] + }, + { + "id": "TA-470", + "scope": "GLM5 indexer selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Additive masking, padding-prefix detection, sorted output, blocked scoring, and query offsets now form one selection policy.", + "Dense/chunked parity, sentinels, valid ranges, ordering, and local-query equivalence remain checked; TileLang stays independent." + ] + }, + { + "id": "TA-471", + "scope": "GLM5 DSA mask fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Two/three-dimensional mask construction and Ulysses query-axis gathering now form one mask policy.", + "Shapes, excluded keys, gathered shard rows, and values remain asserted." + ] + }, + { + "id": "TA-472", + "scope": "GLM5 sparse-MLA reference and wrapper fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense-equivalent reference behavior, query offsets, and local-query/full-KV TileLang adaptation now form one policy.", + "Outputs, flattened shapes, globalized indices, and scaling remain checked." + ] + }, + { + "id": "TA-473", + "scope": "GLM5 sparse-attention integration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ulysses local-query routing and full-top-k dense parity now form one integration policy.", + "Query/KV/index shapes, query offsets, mask handling, and end-to-end hidden-state parity remain." + ] + }, + { + "id": "TA-474", + "scope": "GLM5 sparse kv_b adapter fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Distributed LoRA, block-FP8 QLoRA dequantization, and BF16 absorb-compute behavior now form one adapter-weight policy.", + "Analytical weights, split shapes, dtypes, downstream equations, and both factor gradients remain checked." + ] + }, + { + "id": "TA-475", + "scope": "GLM5 checkpoint filtering fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Out-of-range layer filtering and EP-owned FP8 expert weight/scale filtering now form one checkpoint policy.", + "Ordinary, MTP, distant, non-layer, local-expert, and remote-expert keys remain checked." + ] + }, + { + "id": "TA-476", + "scope": "GLM5 adapter and MoE dispatch fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default MLA/MoE targets, EP eager dispatch, and sparse kv_b LoRA engagement now form one policy.", + "Wrapped/unwrapped modules, expert call routing, router shapes, and sparse/dense parity with nonzero deltas remain." + ] + }, + { + "id": "TA-477", + "scope": "GLM5 forward and recompute fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Tiny-model layer construction/forward and recompute-before-dispatch checkpoint routing now form one runtime policy.", + "Dense/MoE split, indexer reachability, output shape, and inner-versus-outer checkpoint calls remain checked." + ] + }, + { + "id": "TA-478", + "scope": "expert-adapter backend capability fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Registered backend capabilities and inactive optional-dispatch identity now form one backend policy.", + "Local/EP support, dispatch methods, reduction domain, zero-token behavior, and guard identity remain." + ] + }, + { + "id": "TA-479", + "scope": "expert-adapter factor ownership fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quantized projection subsets, unquantized structural-zero omission, and eager local-only admission now form one ownership policy.", + "Parameters, buffers, checkpoints, shapes, factor domains, and backend capability remain asserted." + ] + }, + { + "id": "TA-480", + "scope": "generic expert semantics fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quantized and unquantized rejection/acceptance across gated, activation, clamp, and bias semantics now form one policy.", + "Standard SiLU preservation, native activation, implementation, target roles, and all failure modes remain checked." + ] + }, + { + "id": "TA-481", + "scope": "model-family expert adapter construction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "GLM4 QLoRA, Qwen3 Quack-NF4, and Qwen3.5 Quack-LoRA construction now form one family policy.", + "Checkpoint buffering, exact targets, quantization format/group, source FQN, activation, and hybrid semantics remain." + ] + }, + { + "id": "TA-482", + "scope": "expert-adapter fail-closed fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Invalid quantization groups, model-specific semantics, and invalid quantized/unquantized target sets now form one admission policy.", + "MiniMax, GPT-OSS, Nemotron, empty, router, mixed, and duplicate target failures remain checked." + ] + }, + { + "id": "TA-483", + "scope": "removed server-configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "YAML fields, CLI overrides, nested ZORL configuration, and unrelated unknown fields now form one removal boundary.", + "Rejection timing, messages, and forward-compatible unknown-field behavior remain checked." + ] + }, + { + "id": "TA-484", + "scope": "shipped adapter example fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "LoRA and quantized-MoE example parsing now forms one shipped-configuration policy.", + "Clean-process parsing, Quack selection, hybrid mode, and every expert target remain checked." + ] + }, + { + "id": "TA-485", + "scope": "server runtime serialization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Automatic defaults, nested runtime controls, and receiver KV-cache normalization now form one round-trip policy.", + "Optimizer variants, checkpoint fields, prefetch, packing, activation, adapter, and model defaults remain checked." + ] + }, + { + "id": "TA-486", + "scope": "quantized training configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FP8 training, GLM block-FP8 QLoRA, QARL, NeMo aliases, and fail-fast fallback now form one mode policy.", + "All numerical fields, target scopes, calibration fields, aliases, and serialized defaults remain checked." + ] + }, + { + "id": "TA-487", + "scope": "server parallel-topology fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Exact GLM rank-one admission, lm-head sharding, and LoRA tensor-parallel boundaries now form one topology policy.", + "EP, CP, lm-head TP, total GPU count, chunking, and model-TP rejection remain checked." + ] + }, + { + "id": "TA-488", + "scope": "unsupported server configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quantized-mode conflicts, vLLM FP8 knobs, broadcast loading, and unsupported adapter modes now form one rejection policy.", + "LoRA, MTP, Mamba, ModelOpt, KV-cache, pipeline, and merge-interval failures remain checked." + ] + }, + { + "id": "TA-489", + "scope": "server optimizer and runner compatibility fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Muon Gram-Newton-Schulz serialization and runner compatibility fields now form one execution policy.", + "Dtypes, restarts, fallbacks, determinism, routing, reduction, decay, and export fields remain checked." + ] + }, + { + "id": "TA-490", + "scope": "model-specific server configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sparse-MLA threading and routing-weight placement now form one model-specific policy.", + "Explicit and automatic resolution plus serialized model fields remain checked." + ] + }, + { + "id": "TA-491", + "scope": "GLM5.2 layer-plan fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Official schedule, malformed plans, stage starts, parameter allocation, and strict reload now form one layer-plan policy.", + "All 78 layers, producer ownership, dense/sparse split, indexer keys, and rejection cases remain checked." + ] + }, + { + "id": "TA-492", + "scope": "GLM5.2 sparse logical-selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Small-row edge cases and the production 4096 boundary now form one logical-index policy.", + "Stable ties, dead rows, valid counts, cache mapping, gathers, and top-2048 ordering remain checked." + ] + }, + { + "id": "TA-493", + "scope": "GLM5.2 selector Hadamard fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Transport inversion and pre-quantization enabled/disabled paths now form one Hadamard policy.", + "BF16 bytes, normalization tolerance, and exact query/key operands remain checked." + ] + }, + { + "id": "TA-494", + "scope": "GLM5.2 fused indexer projection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fused projection row order and FP32 gate scaling now form one projection policy.", + "Single-call fusion, exact BF16 bytes, promotion order, and sampler scoring formula remain checked." + ] + }, + { + "id": "TA-495", + "scope": "GLM5.2 sampler index-key preparation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fused preparation, prompt/decode mixing, and CP16 boundary mapping now form one index-key policy.", + "Projection stride, literal RoPE cache, suffix selection, positions, and 4096 ownership remain checked." + ] + }, + { + "id": "TA-496", + "scope": "GLM5.2 sparse quantization codec fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Query/key scale domains, paged-cache unpacking, and production-shape sampler bytes now form one codec policy.", + "E4M3 and UE8M0 formats, page layout, shapes, scales, and bitwise values remain checked." + ] + }, + { + "id": "TA-497", + "scope": "GLM5.2 native selector runtime fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Batched dispatch/masking and fail-closed admission now form one native-selector runtime policy.", + "Flattening, unwritten cells, CUDA absence, and non-prefix masks remain checked." + ] + }, + { + "id": "TA-498", + "scope": "GLM5.2 sparse kernel loader fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "DeepGEMM capability admission and shared selector import now form one loader policy.", + "Missing score support and exact shared-kernel identity remain checked." + ] + }, + { + "id": "TA-499", + "scope": "GLM5.2 canonical routing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Replay rejection, internal transport, exact canonical routing, and ordinary noncanonical routing now report as two policies.", + "Environment independence, selector version, router class, and public-knob exclusion remain checked." + ] + }, + { + "id": "TA-500", + "scope": "SGLang EP admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "DeepEP mechanism exclusion and FP8 rejection now form one serving-kernel admission policy.", + "Transport rationale, unsupported-mechanism wording, and FP8 failure remain checked." + ] + }, + { + "id": "TA-501", + "scope": "SGLang EP flag-off fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Stock compute selection and score-dtype behavior now form one disabled-mode policy.", + "No serving-kernel engagement, finite output, BF16 stock scores, and FP32 opt-in scores remain checked." + ] + }, + { + "id": "TA-502", + "scope": "SGLang EP compute guard fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Empty-rank short circuit and unsupported compute forms now form one boundary policy.", + "Kernel suppression, score presence, gating, bias, and activation failures remain checked." + ] + }, + { + "id": "TA-503", + "scope": "SGLang EP slot-combine fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Slot-ordered reduction, collapsed-pair rejection, and autograd rejection now form one combine policy.", + "FP32 reduction, token-slot mapping, unique selections, and scoring-only ownership remain checked." + ] + }, + { + "id": "TA-504", + "scope": "SGLang EP trainable-dispatch fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Autograd dispatch, training admission, and empty-rank gradients now form one trainable policy.", + "Plain no-grad routing, bias/activation failures, zero weight gradients, and input-score gradients remain checked." + ] + }, + { + "id": "TA-505", + "scope": "SGLang EP weight-presentation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Zero-copy strided, transient, and cached weight modes now form one presentation policy.", + "Storage identity, contiguity, reuse, version invalidation, and explicit invalidation remain checked." + ] + }, + { + "id": "TA-506", + "scope": "SGLang fused RMSNorm CPU fallback fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Residual and forced single-input CPU fallbacks now form one eager-fallback policy.", + "Exact residual carry, weighted output, and global-mode restoration remain checked." + ] + }, + { + "id": "TA-507", + "scope": "SGLang fused RMSNorm forward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Residual, no-residual, packed-3D, and module-mode examples now form one bit-exact forward policy.", + "BF16/FP32, shape preservation, residual bytes, forced mode, and native no-force behavior remain checked." + ] + }, + { + "id": "TA-508", + "scope": "SGLang fused RMSNorm backward fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Residual and no-residual closed-form gradients now form one backward policy.", + "Hidden, residual, and FP32 weight gradients remain compared with eager autograd." + ] + }, + { + "id": "TA-509", + "scope": "SGLang fused RMSNorm integration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense Qwen layer parity and the pre-summed final-norm boundary now form one model integration policy.", + "Both norm sites, serving residual-tree bytes, and the one-ULP seed boundary remain checked." + ] + }, + { + "id": "TA-510", + "scope": "SGLang fused RMSNorm trunk fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Interposed forward, contract-off dispatch, and trunk backward now form one no-residual trunk policy.", + "Bitwise family identity, ordinary dispatch, finite gradients, and eager gradient parity remain checked." + ] + }, + { + "id": "TA-511", + "scope": "native block-FP8 encoding fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Raw FP8 packing and protected linear state now form one parameter-encoding policy.", + "Shapes, exact bytes, parameter names, frozen ownership, and dtype-apply protection remain checked." + ] + }, + { + "id": "TA-512", + "scope": "native block-FP8 execution fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CPU/materialization failure and partition-forward hook traversal now form one execution-boundary policy.", + "Lazy SGLang imports, CUDA admission, invalid materialization inputs, ranges, hooks, and results remain checked." + ] + }, + { + "id": "TA-513", + "scope": "native block-FP8 admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Gradient ownership, partition geometry, and prequantized pair validation now form one fail-closed policy.", + "Scoring-only mode, frozen bases, block boundaries, widths, dtypes, shapes, and finite scales remain checked." + ] + }, + { + "id": "TA-514", + "scope": "native block-FP8 checkpoint fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "State round trip, metadata preflight, adversarial apply, real DCP, and EP-restored shapes now form one lifecycle policy.", + "Both byte streams, no implicit casts, atomic parameter restoration, and global expert shapes remain checked." + ] + }, + { + "id": "TA-515", + "scope": "GLM5.2 native-FP8 configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "HF metadata round trip and nonofficial-contract rejection now form one configuration policy.", + "Quant method, format, activation scheme, block size, exclusions, and serialization remain checked." + ] + }, + { + "id": "TA-516", + "scope": "GLM5.2 native-FP8 dense pair-buffer fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Successful emission, incomplete/duplicate/bad-scale failures, and injective target ownership now form one dense-buffer policy.", + "DCP names, exact bytes, FP32 scales, duplicate targets, and aliased modules remain checked." + ] + }, + { + "id": "TA-517", + "scope": "GLM5.2 native-FP8 construction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quantized module replacement and sparse-MLA KV materialization now form one model-construction policy.", + "Dense/shared/expert ownership, exclusions, mixed-precision classes, forward hooks, and split layouts remain checked." + ] + }, + { + "id": "TA-518", + "scope": "GLM5.2 native-FP8 expert checkpoint fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Expert pair fusion, topology rejection, and dense/expert handler separation now form one checkpoint policy.", + "Local expert bytes, gate-up fusion, scales, rank/count admission, load families, and skip ownership remain checked." + ] + }, + { + "id": "TA-519", + "scope": "Qwen3.5 norm-family dispatch fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Exact selection, ordinary SGLang, fused SGLang, and v2 candidate dispatch now form one family-selection policy.", + "Structural ownership, residual/no-residual routing, exact coexistence, v1 default, and v2 admission remain checked." + ] + }, + { + "id": "TA-520", + "scope": "Qwen3.5 norm-site assignment fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "V2 module resolution, per-layer input forcing, and final-norm forcing now form one call-site policy.", + "Every zero-centered site, layer-zero exception, native mode, and both SGLang modes remain checked." + ] + }, + { + "id": "TA-521", + "scope": "Qwen3.5 norm bit-exact integration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Module dispatch, trunk family one, and full decoder-layer parity now form one integration gate.", + "Residual, forced, qk-norm, aten interpose, and model output bytes remain checked." + ] + }, + { + "id": "TA-522", + "scope": "generic QLoRA quantized execution fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quantize/train/memory, dequantization, and prequantized NVFP4 loading now form one execution policy.", + "Both quant formats, storage savings, output shape, round-trip tolerance, and LoRA-only gradients remain checked." + ] + }, + { + "id": "TA-523", + "scope": "generic QLoRA NVFP4 merge fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EMA scale convention and merge/requantization now form one NVFP4 lifecycle policy.", + "Amax updates, global scale, quantization error, delta folding, state export, and adapter reset remain checked." + ] + }, + { + "id": "TA-524", + "scope": "generic QLoRA injection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Target replacement/training and checkpoint-format propagation now form one injection policy.", + "Target exclusion, source format, quant format, model forward, and every adapter gradient remain checked." + ] + }, + { + "id": "TA-525", + "scope": "generic QLoRA block-FP8 fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Prequantized block-FP8 loading, fused QKV assembly, forward/backward, and merge now form one format policy.", + "Packed bytes, round-trip error, shapes, adapter gradients, delta folding, and no-EMA behavior remain checked." + ] + }, + { + "id": "TA-526", + "scope": "generic QLoRA optimizer-reset fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Adapter-state clearing, non-LoRA preservation, and scheduled merge integration now form one optimizer lifecycle.", + "State rebuild, selective ownership, pre-boundary no-op, boundary requantization, and factor reset remain checked." + ] + }, + { + "id": "TA-527", + "scope": "OPD numerical backend fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Chunking, streaming/TileLang, low-memory streaming, and sharded teacher storage now form one backend policy.", + "Reference loss, gradients, teacher detachment, chunk equivalence, and safetensor shard reads remain checked." + ] + }, + { + "id": "TA-528", + "scope": "OPD gradient and reduction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default backward, token-partial reduction, and BF16-to-FP32 loss behavior now form one gradient policy.", + "Student gradients, teacher detachment, valid-token scaling, output dtype, and finiteness remain checked." + ] + }, + { + "id": "TA-529", + "scope": "OPD output edge fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All-ignored and per-token outputs now form one output-shaping policy.", + "Zero finite loss/gradients, masks, FP32 dtype, tensor shape, and reduced reference equality remain checked." + ] + }, + { + "id": "TA-530", + "scope": "OPRD hidden-distance fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Materialized and fetched teacher-layer paths now form one hidden-distance policy.", + "Chunk boundaries, full-reference distance, layer counts, student gradients, and teacher detachment remain checked." + ] + }, + { + "id": "TA-531", + "scope": "OPD hidden-only objective boundary", + "decision": "keep", + "status": "applied", + "evidence": [ + "Zero-KL hidden-only MSE remains an independent objective boundary rather than being folded into backend checks.", + "Weighted token math, per-token output, metrics, hidden gradients, absent head gradients, and KL diagnostics remain checked." + ] + }, + { + "id": "TA-532", + "scope": "quantized-export CLI fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "YAML/CLI precedence and byte-size parsing now form one command-line policy.", + "Block sizes, appended exclusions, BF16 layer counts, shard sizes, and decimal/binary units remain checked." + ] + }, + { + "id": "TA-533", + "scope": "base quantized-export fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Module CLI execution, BF16 islands/config writing, and sharded-index output now form one base-directory policy.", + "Subprocess entrypoint, quantized counts, exact dtypes, exclusions, tokenizer copy, shard map, and total bytes remain checked." + ] + }, + { + "id": "TA-534", + "scope": "quantized-export projection layout fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fused QKV, MLA-A fusion, and linear-attention renaming now form one projection-layout policy.", + "Every emitted name, FP8 tensor, scale, metadata rejection, split fallback, convolution, norm, and bias remains checked." + ] + }, + { + "id": "TA-535", + "scope": "quantized-export MoE layout fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "GKN expert conversion and fused gate-up splitting now form one MoE-layout policy.", + "HF expert names, transposes, gate/up/down partitions, quantized bytes, scales, and output counts remain checked." + ] + }, + { + "id": "TA-536", + "scope": "quantized-export admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unsupported source-state rejection and QARL fold admission now form one export preflight policy.", + "Prequantized inputs, active adapters, missing metadata, fold eligibility, and block-size compatibility remain checked." + ] + }, + { + "id": "TA-537", + "scope": "training-model FP8 construction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full-model FP8 injection and tensor-parallel lm-head inclusion now form one construction policy.", + "Dense/MoE replacement, overrides, correction, fallback, trainability, TP inclusion, and explicit exclusion remain checked." + ] + }, + { + "id": "TA-538", + "scope": "training-model GLM5.2 block-FP8 QLoRA fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Foundation/injection threading and missing-QLoRA rejection now form one GLM5.2 adapter-mode policy.", + "Rank, alpha, format, group size, inventory ownership, and enabling preconditions remain checked." + ] + }, + { + "id": "TA-539", + "scope": "training-model quantized-mode admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FP8/adapters, QARL/adapters-or-FP8, and QARL/MoE conflicts now form one admission matrix.", + "Each full-weight ownership conflict and dense-only QARL boundary remains checked." + ] + }, + { + "id": "TA-540", + "scope": "training-model QARL lifecycle fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense fake-quant injection and pre-parallel calibration now form one QARL lifecycle policy.", + "Targets, activation mode, trainability, calibration samples/sequence, forward counts, and learned scales remain checked." + ] + }, + { + "id": "TA-541", + "scope": "DSV4 checkpoint translation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All supported name families and C4 APE hotfix inversion now form one translation policy.", + "Attention/indexer/HC/router/shared-expert mappings, three head dimensions, and invalid layouts remain checked." + ] + }, + { + "id": "TA-542", + "scope": "DSV4 checkpoint quantized-codec fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Block-FP8 and packed MXFP4 dequantization now form one checkpoint-codec policy.", + "Full blocks, tails, known nibbles, scale blocks, output shapes, and dtypes remain checked." + ] + }, + { + "id": "TA-543", + "scope": "DSV4 checkpoint handler fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP-local expert fusion and MTP/unmapped-key accounting now form one handler-ownership policy.", + "Skip decisions, local expert rows, gate-up/down shapes and values, and unknown-key summaries remain checked." + ] + }, + { + "id": "TA-544", + "scope": "DSV4 synthetic load fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Window-only, strict nonpersistent buffers, C4/indexer, and hash-router models now form one end-to-end load policy.", + "Fill summaries, representative values, APE recovery, expert fusion, tid2eid, correction bias, and strict accounting remain checked." + ] + }, + { + "id": "TA-545", + "scope": "GLM5.2 partial-edge QLoRA scale example", + "decision": "remove", + "status": "applied", + "evidence": [ + "The isolated 6144-by-576 scale-shape assertion was an exact subset of the full 700-target inventory contract.", + "The retained inventory still asserts the identical (5, 192) partial-edge storage shape on the official kv_a projection." + ] + }, + { + "id": "TA-546", + "scope": "GLM5.2 exact dense-component fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Logical inventory/physical fused roots and rank-one pre-mutation rejection now form one dense-component policy.", + "All 1,700 factors, three roots, source FQNs, parameter identity, and rank/alpha failures remain checked." + ] + }, + { + "id": "TA-547", + "scope": "GLM5.2 QLoRA admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Wrong targets, missing indexer exclusions, and unsupported construction modes now form one fail-closed policy.", + "No-mutation guarantees, shapes, DSA BF16 ownership, backend, dispatch, exact-contract, and feature-flag failures remain checked." + ] + }, + { + "id": "TA-548", + "scope": "GLM5.2 training-mode fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ordinary product admission and complete exact-active-LoRA alltoall selection now form one training-mode policy.", + "Format rejection, the certified tuple, exact enablement, and deepep rejection remain checked." + ] + }, + { + "id": "TA-549", + "scope": "exact fused gate-up state fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Native-leaf ownership, explicit gate/up loading, and dtype movement now form one state lifecycle.", + "Four FP32 factors, rank/alpha/topology admission, row order, FP8/scale bytes, master identity, and protected dtypes remain checked." + ] + }, + { + "id": "TA-550", + "scope": "exact fused gate-up numerical fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One-time master rounding/logical order and two-branch surrogate parity now form one effective-numerics policy.", + "Captured factor bytes, gate/up ordering, base composition, activation, and output equality remain checked." + ] + }, + { + "id": "TA-551", + "scope": "exact fused gate-up gradient fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Input/factor parity, factor-only VJP, and master-mutation rejection now form one backward policy.", + "All five gradients, no base materialization, and saved-tensor version safety remain checked." + ] + }, + { + "id": "TA-552", + "scope": "exact TP16 lm-head topology fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Official shards, component range ownership, and process-group validation now form one topology policy.", + "All 16 ranges, padding, statelessness, order, rank, world size, and NCCL admission remain checked." + ] + }, + { + "id": "TA-553", + "scope": "exact TP16 lm-head operand fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local operand validation and CPU-before-import rejection now form one execution-admission policy.", + "BF16/FP32 ownership, rank-one shapes, sampler stride, token range, frozen weight, CUDA, and lazy imports remain checked." + ] + }, + { + "id": "TA-554", + "scope": "exact TP16 lm-head presentation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-order vocabulary assembly and effective factor views now form one presentation-bytes policy.", + "Identity token mapping, collective layout rejection, BF16 bytes, and immutable FP32 masters remain checked." + ] + }, + { + "id": "TA-555", + "scope": "exact TP16 lm-head surrogate fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local FP32 surrogate VJP and the custom autograd boundary now form one gradient policy.", + "Base-plus-LoRA hidden gradients, factor gradients, grad-enabled output, and saved effective bytes remain checked." + ] + }, + { + "id": "TA-556", + "scope": "sparse source-delta capture fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank capture, template security, global manifest, and empty-rank reload now form one source lifecycle.", + "Changed indices/values, paths, totals, ordering, traversal rejection, rank tags, shapes, and empty BF16 payloads remain checked." + ] + }, + { + "id": "TA-557", + "scope": "single sparse-delta file fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packing, deterministic index sorting, and malformed-update rejection now form one file policy.", + "Stats, INT32 CPU indices, BF16 values, shapes, duplicates, dtype, length, and range failures remain checked." + ] + }, + { + "id": "TA-558", + "scope": "sparse-delta contiguous-shard fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Index localization and empty local shards now form one contiguous-sharding policy.", + "All rank shapes, localized sorted indices, values, and explicit empty tensors remain checked." + ] + }, + { + "id": "TA-559", + "scope": "ranked sparse-delta file fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Raw rank writes, template security, and pre-encoded writes now form one rank-file policy.", + "Rank order, paths, nnz counts, traversal rejection, and packed API behavior remain checked." + ] + }, + { + "id": "TA-560", + "scope": "sparse-delta translation-future fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Future draining, unranked rejection, and direct ranked-file publication now form one translation policy.", + "Tag stripping, rank maps, exact encoded objects, expected ranks, and output paths remain checked." + ] + }, + { + "id": "TA-561", + "scope": "EP adapter registry fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Optional Quack registration and live signature inspection now form one registry policy.", + "Base-versus-MoE-act registration, explicit shared parameters, and forward-compatible kwargs remain checked." + ] + }, + { + "id": "TA-562", + "scope": "native EP adapter FP8 fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Common FP8-kwarg consumption and explicit FP8-compute rejection now form one native-backend boundary.", + "Disabled-mode output shape and enabled-mode failure remain checked." + ] + }, + { + "id": "TA-563", + "scope": "Triton EP adapter FP8 fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Common FP8-kwarg consumption and explicit FP8-compute rejection now form one Triton-backend boundary.", + "Kernel dispatch, output shape, optional availability, and enabled-mode failure remain checked." + ] + }, + { + "id": "TA-564", + "scope": "Triton MoE-act EP adapter fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Common activation/FP8 kwargs and explicit FP8-compute rejection now form one MoE-act boundary.", + "Activation-native, gating, bias, clamp, output, optional availability, and rejection remain checked." + ] + }, + { + "id": "TA-565", + "scope": "DSV4 model construction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Pipeline rejection and sequence-parallel group wiring now form one topology policy.", + "Hyperconnection PP exclusion, TP/CP groups, CP size, and attention implementation remain checked." + ] + }, + { + "id": "TA-566", + "scope": "DSV4 model runtime fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "C128 shape, full LM backward, and hash-routed execution now form one runtime policy.", + "Finite initialization/output/loss, hidden/logit shapes, ordinary and hash gradients, and frozen HC ownership remain checked." + ] + }, + { + "id": "TA-567", + "scope": "DSV4 precision-preservation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Keep-FP32 marking, registry dtype construction, and direct dtype movement now form one precision policy.", + "Every HC/attention/compressor carve-out, BF16 ordinary weights, and complex RoPE imaginary components remain checked." + ] + }, + { + "id": "TA-568", + "scope": "routing-replay sequence-parallel fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Padded/unpadded positions, independent unpacked rows, and excess truncation now form one SP layout policy.", + "All CP ranks, actual position lengths, pad values, row boundaries, micro-batch shape, and truncation order remain checked." + ] + }, + { + "id": "TA-569", + "scope": "routing-replay RingAttention fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Four-rank zigzag placement and packed-document boundaries now form one RingAttention layout policy.", + "Every rank slice, shape, position ordering, and cross-document separation remain checked." + ] + }, + { + "id": "TA-570", + "scope": "routing-replay weight tensor fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Float-value/padding preservation and decoded NumPy slicing now form one routing-weight policy.", + "FP32 dtype, uniform pad weights, CP rank slicing, tensor shape, and exact values remain checked." + ] + }, + { + "id": "TA-571", + "scope": "routing-replay decode fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shaped base64, inferred-shape base64, and materialized Python/NumPy/Tensor inputs now form one wire-format policy.", + "INT32 bytes, model top-k inference, expert selections, logits conversion, types, shapes, and values remain checked." + ] + }, + { + "id": "TA-572", + "scope": "top-k router softmax fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy reference behavior and environment-selected tie/logit policies now form one softmax policy.", + "Normalization modes, inert V4 inputs, stable low/high IDs, logits selection, weights, and unknown-policy rejection remain checked." + ] + }, + { + "id": "TA-573", + "scope": "top-k router layer-FP32 fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Layer selector parsing and live MoE-block routing now form one scoped-FP32 policy.", + "Ranges/all mode, gate bypass, FP32 operands/logits, and expert selection remain checked." + ] + }, + { + "id": "TA-574", + "scope": "top-k router configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "V4 scaling admission and V4/legacy config construction now form one configuration policy.", + "Post-renorm scaling, softmax rejection, scoring method, top-k method, expert counts, and defaults remain checked." + ] + }, + { + "id": "TA-575", + "scope": "OPD shifted-payload fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Teacher cache input/target shift, OPD cache-index alignment, and unshifted rejection now form one payload policy.", + "Input tokens, targets, teacher IDs/weights, cache indices, and length failure remain checked." + ] + }, + { + "id": "TA-576", + "scope": "OPD teacher-cache transport fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Mooncake metadata return and legacy file-metadata rejection now form one teacher-cache transport policy.", + "Request payload, absence of file paths, metadata identity, cache indices, and fail-closed backend validation remain checked." + ] + }, + { + "id": "TA-577", + "scope": "checkpoint-manager rank-zero save failures", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Adapter-state and LoRA-only write failures now form one pre-barrier failure policy.", + "Both public error messages and the prohibition on entering the barrier remain checked." + ] + }, + { + "id": "TA-578", + "scope": "checkpoint-manager dtype-preservation forwarding probe", + "decision": "remove", + "status": "applied", + "evidence": [ + "The deleted test only inspected an internal preserve_lora_dtype keyword.", + "Adapter persistence coverage checks the dtype of the actual saved LoRA tensors through the public save path." + ] + }, + { + "id": "TA-579", + "scope": "checkpoint-manager MoE save fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Collective-state slicing and resolved-target detection now form one MoE export policy.", + "Gather ownership, active-rank slices, target discovery, exported keys, shapes, values, and rank remain checked." + ] + }, + { + "id": "TA-580", + "scope": "checkpoint-manager export-format forwarding probe", + "decision": "remove", + "status": "applied", + "evidence": [ + "The deleted mock asserted only that one internal keyword was forwarded.", + "The retained SGLang shared-outer roundtrip validates the emitted format marker, complete tensor layout, and reload." + ] + }, + { + "id": "TA-581", + "scope": "checkpoint-manager strict-manifest artifact duplicate", + "decision": "remove", + "status": "applied", + "evidence": [ + "The adapter-manager suite already saves and reads the strict target manifest through the public checkpoint lifecycle.", + "It additionally validates the saved manifest and rejects a mismatched runtime manifest." + ] + }, + { + "id": "TA-582", + "scope": "PEFT hybrid-shared checkpoint fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "PEFT orientation and hybrid-shared reload now form one checkpoint roundtrip.", + "Config fields, shared/expert tensor orientations, shapes, bytes, key inventory, and restored factors remain checked." + ] + }, + { + "id": "TA-583", + "scope": "SGLang shared-outer checkpoint fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shared-outer save/load and incompatible ownership rejection now form one format contract.", + "Format metadata, all six MoE tensor layouts, exact reload, and fail-closed admission remain checked." + ] + }, + { + "id": "TA-584", + "scope": "checkpoint-roundtrip cautious optimizer construction example", + "decision": "remove", + "status": "applied", + "evidence": [ + "The deleted example did not exercise checkpointing and only inspected optimizer type and parameter-group fields.", + "Dedicated optimizer coverage already validates cautious routing, optimizer selection, kwargs, and group policy." + ] + }, + { + "id": "TA-585", + "scope": "MoE-LoRA initialization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Backend construction and active-rank view layout now form one initialization policy.", + "Frozen/trainable ownership, all factor shapes, zero initialization, repr, rank slicing, and contiguity remain checked." + ] + }, + { + "id": "TA-586", + "scope": "eager MoE-LoRA execution fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Eager forward/backward, block integration, and hybrid-shared construction now form one CPU execution policy.", + "Output shapes, gradient ownership, injection, and every supported shared factor shape remain checked." + ] + }, + { + "id": "TA-587", + "scope": "standalone nonzero MoE-LoRA output example", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained cross-backend contract initializes nonzero LoRA B factors and requires factor gradients on real outputs.", + "The separate eager-only max-difference example added no distinct backend or ownership boundary." + ] + }, + { + "id": "TA-588", + "scope": "generic injection examples in MoE-LoRA suite", + "decision": "remove", + "status": "applied", + "evidence": [ + "Qwen subclass wrapping repeated the generic MoE from_module contract, while unmatched linear targets were not MoE-specific.", + "MoE conversion and both injection APIs remain checked here; generic and model-specific injection boundaries remain elsewhere." + ] + }, + { + "id": "TA-589", + "scope": "EP MoE-LoRA router-score fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Alltoall/deepep score names, missing scores, and score gradients now form one router-score contract.", + "Multiplication values, identity behavior, compute-output gradients, and exact score gradients remain checked." + ] + }, + { + "id": "TA-590", + "scope": "LoRA target-manifest rejection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Coverage, rank, configured-target, unlisted-module, and scalar-schema failures now form one fail-closed manifest contract.", + "Every prior error condition and exact scalar-type boundary remains checked alongside the independent success lifecycle." + ] + }, + { + "id": "TA-591", + "scope": "NeMo FP8 configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The supported blockwise alias and unsupported Transformer Engine recipes now form one translation policy.", + "Native enablement plus hybrid, tensorwise, MXFP8, and FP8-parameter rejection remain checked." + ] + }, + { + "id": "TA-592", + "scope": "external FP8 compatibility rejection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "vLLM receiver knobs and ModelOpt QARL configurations now share one table-driven fail-closed contract.", + "All thirteen configuration paths and their diagnostic categories remain checked without repeated setup." + ] + }, + { + "id": "TA-593", + "scope": "FP8 BF16 layer-island fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "First/last resolution, overlap, invalid topology, and real module injection now form one layer-island lifecycle.", + "Every pattern, count, replacement boundary, model summary, and rejection remains checked." + ] + }, + { + "id": "TA-594", + "scope": "NVFP4 two-dimensional quantization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Reference bytes, public-format dispatch, and straight-through linear gradients now form one numerical contract.", + "Both supported dtypes, output shape/dtype, exact quantization, and upstream weight gradient remain checked." + ] + }, + { + "id": "TA-595", + "scope": "NVFP4 input and format admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank, block-divisibility, and unsupported-format failures now form one API admission policy.", + "The cross-row grouping trap and generic dispatch failure remain explicit." + ] + }, + { + "id": "TA-596", + "scope": "NVFP4 MoE projection STE fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Down and fused gate-up projection paths now share one shape and straight-through-gradient contract.", + "Both tensor layouts, lossy forward behavior, exact gradients, expert independence, and per-half scaling remain checked." + ] + }, + { + "id": "TA-597", + "scope": "data-packing allocation primitives", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FFD feasibility, group packing, and sequential rank allocation now form one allocation contract.", + "Capacity, bin size, safe mode, offsets, rank isolation, token accounting, and full coverage remain checked." + ] + }, + { + "id": "TA-598", + "scope": "packing sample-preprocessing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Position metadata and trainable-label filtering now form one sample preprocessing policy.", + "Single/batched inputs, empty/missing fields, preserved metadata, masks, and missing-label failure remain checked." + ] + }, + { + "id": "TA-599", + "scope": "packing dataset-preprocessing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dataset filtering and train/eval packing preparation now form one dataset pipeline.", + "Retained rows, position/length columns, and optional evaluation data remain checked." + ] + }, + { + "id": "TA-600", + "scope": "linear learning-rate schedule fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default full-range decay and warmup/decay-ratio/floor behavior now form one linear schedule policy.", + "Endpoints, equal decrements, warmup values, decay boundary, and post-decay floor remain checked." + ] + }, + { + "id": "TA-601", + "scope": "cosine learning-rate schedule fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Floor/monotonicity and warmup/midpoint/endpoint behavior now form one cosine schedule policy.", + "Decay ratio, minimum floor, warmup values, half-cosine midpoint, and terminal value remain checked." + ] + }, + { + "id": "TA-602", + "scope": "session API configuration compatibility fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Removed-field rejection and unrelated future-field preservation now form one compatibility policy.", + "All migration diagnostics and nested model-extra behavior remain checked." + ] + }, + { + "id": "TA-603", + "scope": "sampling-adapter reconciliation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Stale-entry pruning and endpoint-query failure preservation now form one reconciliation policy.", + "Eviction avoidance, fresh loading, tracked state replacement, and transient-query safety remain checked." + ] + }, + { + "id": "TA-604", + "scope": "sampling-session tracking fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "URI/plain model scoping and failed-load atomicity now form one sampling-session tracking policy.", + "Embedded and requested model IDs, resolved paths, successful tracking, and absence of stale failed entries remain checked." + ] + }, + { + "id": "TA-605", + "scope": "MoE routing-weight numerical fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Before/after-down parity and the no-router-gradient path now form one numerical contract.", + "FP64 outputs, every input/factor/score gradient, error class, and in-place score-fold safety remain checked." + ] + }, + { + "id": "TA-606", + "scope": "MoE routing-weight configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Mutable config, environment forcing, automatic regimes, parity opt-in, explicit values, and invalid input now form one policy.", + "All train-router/dispatch combinations and Boolean/string spellings remain checked." + ] + }, + { + "id": "TA-607", + "scope": "QARL MoE conversion fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Identity preservation, idempotence, and type admission now form one conversion policy.", + "Parameter objects, class identity, QARL attributes, backend retention, repeat conversion, and non-expert rejection remain checked." + ] + }, + { + "id": "TA-608", + "scope": "QARL MoE eager-execution fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quantized execution and disabled-quantization passthrough now form one eager policy.", + "Lossiness, finite output, shape, parameter restoration, gradients, and exact passthrough remain checked." + ] + }, + { + "id": "TA-609", + "scope": "QARL MoE injection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full conversion, FP8 rejection, and target-specific selection now form one injection policy.", + "Linear/expert wrapping, conversion counts, expert-module metadata, independent targets, and fail-closed format admission remain checked." + ] + }, + { + "id": "TA-610", + "scope": "RoPE registry and frequency-precision fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Registry coverage, model-wide BF16 casting, and forward cos/sin precision now form one frequency policy.", + "Every registered recipe, FP32 CPU references, default scaling, YaRN scaling, and bitwise BF16 consumption remain checked." + ] + }, + { + "id": "TA-611", + "scope": "RoPE default-cache fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Lazy cache materialization and Qwen class-B growth now form one cache lifecycle.", + "Execution device, FP32 dtype, indexed values, prefix stability, growth, and CPU recipe reconstruction remain checked." + ] + }, + { + "id": "TA-612", + "scope": "generic QARL dense injection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense wrapping, target/exclusion behavior, model summary, and unsupported architecture admission now form one lifecycle.", + "Parameter names, forward counts, MTP rejection, and Mamba rejection remain checked alongside successful injection." + ] + }, + { + "id": "TA-613", + "scope": "Nemotron-H EP checkpoint fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Topology validation, skip-key ownership, local expert slicing, skip accounting, and parallel-plan classification now form one EP policy.", + "Invalid sizes/ranks, MTP filtering, local ranges, stacked tensors, expert parameter matching, and no-shard modules remain checked." + ] + }, + { + "id": "TA-614", + "scope": "weight-sync endpoint health fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Model-info fallback and exhaustive health failure now form one endpoint-health policy.", + "Endpoint port routing, v1-model fallback order, and diagnostics naming every attempted route remain checked." + ] + }, + { + "id": "TA-615", + "scope": "NCCL endpoint-port routing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Receiver-group initialization and direct bucket transfer now form one port-routing policy.", + "Initialization result metadata, init/update URLs, and direct load-format payload remain checked." + ] + }, + { + "id": "TA-616", + "scope": "NCCL flattened-bucket fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Mixed-dtype flattening and chunked flattened transfer now form one bucket-format policy.", + "Byte packing, payload metadata, flattened/chunked load formats, broadcast counts, and completion waits remain checked." + ] + }, + { + "id": "TA-617", + "scope": "runner session-registration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Coordinator delegation and rank-zero handling of a worker registration error now form one registration policy.", + "Normalized payload delegation, success results, cross-rank synchronization, and failure response text remain checked." + ] + }, + { + "id": "TA-618", + "scope": "runner optimizer-publication fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Successful publication, commit-tail failure, and optimizer-handler tail failure now form one mutation lifecycle.", + "Commit ownership, fatal error translation, causal exceptions, poisoning, and publication ineligibility remain checked." + ] + }, + { + "id": "TA-619", + "scope": "runner forward-backward completion fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Completion ordering and uniform/asymmetric failure policy now form one forward-backward lifecycle.", + "Metric gather, rendezvous, commit-before-merge order, both ranks, uniform rejection, and fatal promotion remain checked." + ] + }, + { + "id": "TA-620", + "scope": "attention backend registry fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FlashAttention registration and runtime backend resolution now form one registry policy.", + "FA4-only import, all flash aliases, eager/native/flex resolution, and unavailable-flash rejection remain checked." + ] + }, + { + "id": "TA-621", + "scope": "SGL page-size-one attention fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Varlen routing and single-sequence metadata synthesis now form one SGL KV-cache policy.", + "Page table/cache shapes, int32 offsets, sequence lengths, scale, causal mode, num_splits, and output shapes remain checked." + ] + }, + { + "id": "TA-622", + "scope": "alternate FlashAttention path fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Paged-KV-cache, flags-off FA3 varlen, and FA4 selection now form one backend-selection policy.", + "Selected call targets, cache layout, num_splits, scale, causal mode, and disabled SGL dispatch remain checked." + ] + }, + { + "id": "TA-623", + "scope": "sequence-shard collator primitive fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "SP slicing and padding now form one collator primitive policy.", + "Both ranks, uneven lengths, ordinary/sequential padding, zero padding, and initialization state remain checked." + ] + }, + { + "id": "TA-624", + "scope": "sequence-shard collator side-channel fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Teacher hidden states and DRGRPO token side channels now form one shard-alignment policy.", + "CP2/CP16 slicing, first/last ranks, shapes, values, ignore-index padding, and zero padding remain checked." + ] + }, + { + "id": "TA-625", + "scope": "sparse-delta streaming lifecycle fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Initial full encoding, exact-byte delta encoding, and unchanged-bucket suppression now execute as one three-transfer lifecycle.", + "TP path replication, endpoint payload, full and changed indices and values, and skip accounting remain checked." + ] + }, + { + "id": "TA-626", + "scope": "sparse-delta prepacked publication fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Per-rank packed paths and FP8 KV-cache metadata now form one prepacked publication transaction.", + "Both paths, unique-file byte accounting, request flags, normalized cache epoch, and endpoint result metadata remain checked." + ] + }, + { + "id": "TA-627", + "scope": "families-v2 fused-split realization guard", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The proof that forcing split reaches the split kernel now runs inside the fused-versus-split bitwise matrix instead of repeating a standalone kernel call.", + "Tail, aligned, and deep-tile shapes, row-count extremes, residual, plain, and zero-centered modes remain checked." + ] + }, + { + "id": "TA-628", + "scope": "families-v2 norm dispatch fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shipped hidden-size admission, row-depth threshold behavior, and the split-kernel tile basis now form one dispatch policy.", + "Fused decisions, both threshold sides, the row cutoff, and the rejected fused-chunk basis remain checked." + ] + }, + { + "id": "TA-629", + "scope": "EP clip local norm and empty-gradient fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Infinity-norm clipping, parameters without gradients, and empty groups now form one local norm policy.", + "Returned norms, uniform clipping, skipped gradients, and the zero norm remain checked." + ] + }, + { + "id": "TA-630", + "scope": "skip-FSDP EP clipping fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Classification, uniform clipping, and raw local-gradient preservation now execute as one skip-FSDP lifecycle.", + "EP and non-EP ownership, the combined norm, clip coefficients, and absence of EP division remain checked." + ] + }, + { + "id": "TA-631", + "scope": "clip-grad-norm dispatch fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP-aware and ordinary FSDP dispatch outcomes now form one public dispatch policy.", + "Both parameter representations still execute and return the expected norm." + ] + }, + { + "id": "TA-632", + "scope": "mixed-mesh foreach clipping fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Safe default per-tensor clipping and explicit foreach rejection now form one mixed-mesh policy.", + "Both DTensor meshes, returned norm, clipped local values, and the explicit cross-mesh error remain checked." + ] + }, + { + "id": "TA-633", + "scope": "orchestrator packing capacity fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Single-batch packing, overflow splitting, mixed lengths, exact fit, and off-by-one capacity now form one batching policy.", + "Token and label shifts, positions, sample counts, maximum length, and boundary outcomes remain checked." + ] + }, + { + "id": "TA-634", + "scope": "orchestrator packing input-normalization fragment", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "NumPy-to-list normalization now executes inside the existing empty, missing, oversized, and single-sample input policy.", + "The production conversion assertion remains unchanged." + ] + }, + { + "id": "TA-635", + "scope": "P2P transfer source and receiver-manifest rejection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unknown receiver names, incompatible shapes, and unsupported source ranks now form one transfer-admission policy.", + "All error messages and the no-transfer side-effect checks remain." + ] + }, + { + "id": "TA-636", + "scope": "P2P receiver-memory coalescing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Distinct receiver handles and absent handle metadata now execute inside one coalescing policy.", + "Session, peer pointers, transfer lengths, and batch count remain checked for both layouts." + ] + }, + { + "id": "TA-637", + "scope": "P2P transfer failure diagnostic fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Named tensor and handle details, bounded coalesced samples, and diagnostics-disabled behavior now form one failure-reporting policy.", + "Pointer details, the six-entry cap, omitted-entry count, and default redaction remain checked." + ] + }, + { + "id": "TA-638", + "scope": "runner effective LM-head selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical merged-LoRA selection and the legacy unmerged fallback now form one effective-weight policy.", + "Exact folded bytes, adapter gradients, frozen base weight, and legacy formula remain checked." + ] + }, + { + "id": "TA-639", + "scope": "runner compiler replica-topology fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "World-discovered replica counts, malformed coverage, and unsupported general tensor parallelism now form one topology policy.", + "SP2, output4, composed world8, every coverage failure, and the TP rejection remain checked." + ] + }, + { + "id": "TA-640", + "scope": "runner unquantized expert admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Eager, exact-merged, and fused-managed expert admission now includes the hybrid-checkpoint metadata rejection in one policy.", + "Producer, topology, parameter count, and fail-closed metadata behavior remain checked." + ] + }, + { + "id": "TA-641", + "scope": "runner quantized expert contract fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "NF4, NVFP4, and block-FP8 admission now includes declared-shape drift and uncertified EP or eFSDP rejection in one contract policy.", + "All formats, eager and fused producers, guard fields, logical shape validation, and both unsupported parallel regimes remain checked." + ] + }, + { + "id": "TA-642", + "scope": "Muon builder configuration fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Complete Gram-Newton-Schulz option forwarding and rejection of a nonpositive grouping byte limit now form one configuration policy.", + "Parameter groups, dtypes, fallback mode, restart count, byte limit, and the validation error remain checked." + ] + }, + { + "id": "TA-643", + "scope": "Muon Quack backend fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quack operation dispatch, tuned-mode selection, missing-package failure, and SM90 dtype routing now form one backend policy.", + "Every GEMM operation, tuned flags, import error, FP32 Torch fallback, and BF16 Quack selection remain checked." + ] + }, + { + "id": "TA-644", + "scope": "Muon fused-weight and Nemotron classification fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fused gate-up discovery and full Nemotron parameter classification now form one optimizer-classification policy.", + "Gated and non-gated experts, post-FSDP attribute loss, every Muon pattern, and AdamW exclusions remain checked." + ] + }, + { + "id": "TA-645", + "scope": "weight-sync parameter extraction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense ownership filtering and tied-weight alias handling now form one extraction policy.", + "Included names, duplicate suppression, declared aliases, false ties, and undeclared shared storage remain checked." + ] + }, + { + "id": "TA-646", + "scope": "weight-sync inference-unfuse layout fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "DeepSeek and Kimi MLA fusion, contiguous FP8 views, Nemotron-H publication, and gated stacked-expert splitting now form one layout-conversion policy.", + "Names, tensor values, storage aliasing, transposes, architecture prefixes, and fail-closed fused Nemotron input remain checked." + ] + }, + { + "id": "TA-647", + "scope": "adapter optimizer identity and live-binding fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical ordering, wrapper-insensitive fingerprints, and live optimizer parameter identity now form one ownership policy.", + "Exact names, equivalent fingerprints, and rejection before state access remain checked." + ] + }, + { + "id": "TA-648", + "scope": "adapter optimizer bitwise-resume control fragment", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The weights-only divergence control now executes inside the bitwise resumed-versus-uninterrupted trajectory.", + "Parameter and moment equality plus the proving divergence without optimizer state remain checked." + ] + }, + { + "id": "TA-649", + "scope": "public adapter optimizer LR-restore fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Scheduled-LR restoration after eviction and explicit LR override now form one public resume policy.", + "Registration generation, ownership fingerprint, restored metadata, optimizer-group LR, next-step parameters, and moments remain checked." + ] + }, + { + "id": "TA-650", + "scope": "adapter optimizer logical-reshard fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Successful logical resharding and every invalid-source rejection now form one topology-transition policy.", + "One- and two-dimensional slices, replicas, same-world changes, holes, overlaps, dtype and shape drift, step mismatch, empty ranks, and resident-state atomicity remain checked." + ] + }, + { + "id": "TA-651", + "scope": "FP8 linear injection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Core replacement, recipe overrides, and exclusions now form one model-injection policy.", + "Parameter identity, module FQNs, global and per-module recipes, unknown-key rejection, explicit exclusions, and glob exclusions remain checked." + ] + }, + { + "id": "TA-652", + "scope": "block-FP8 GEMM backend fallback fragment", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Automatic scaled-matmul failure and warn-once Triton fallback now execute inside the backend and scale-layout policy.", + "Block and row scales, explicit Torch backend parity, fallback equality, and warning suppression remain checked." + ] + }, + { + "id": "TA-653", + "scope": "FP8 linear CUDA execution fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Master-weight training and float32-output execution now form one CUDA FP8 lifecycle.", + "FP8 dispatch, finite gradients, parameter updates, and requested output dtype remain checked." + ] + }, + { + "id": "TA-654", + "scope": "inference endpoint automatic-sync fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Detected TP size and configured synchronization method now form one endpoint auto-sync policy.", + "Receiver discovery, normalized world size, endpoint payload, successful registration, and P2P method forwarding remain checked." + ] + }, + { + "id": "TA-655", + "scope": "inference endpoint weight-sync routing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Single-endpoint forwarding and default, named, or unmatched pool selection now form one routing policy.", + "Endpoint payloads, model ID, timing, rank summaries, all pool outcomes, and no-dispatch failure remain checked." + ] + }, + { + "id": "TA-656", + "scope": "dense and sequence-parallel adapter autograd launchers", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Two-GPU dense and sequence-parallel ownership workers now report as one foundational autograd policy.", + "Both distributed subprocesses, analytical optimizer comparisons, and certification markers still execute." + ] + }, + { + "id": "TA-657", + "scope": "unquantized expert adapter all-to-all launchers", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP2 backend parity, projection-subset ownership, and the all-owner layout now form one unquantized all-to-all policy.", + "Eager, Triton, native, and Quack backends, full and down-only targets, structural zeros, and public optimizer steps remain checked." + ] + }, + { + "id": "TA-658", + "scope": "unquantized expert adapter eFSDP launchers", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Four-GPU shared-owner and all-owner Quack layouts now form one eFSDP topology policy.", + "Both distributed workers and their distinct certification markers remain checked." + ] + }, + { + "id": "TA-659", + "scope": "quantized expert adapter all-to-all launchers", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Backend-format parity and the quantized projection subset now form one EP2 all-to-all policy.", + "Triton NF4, native NVFP4, Quack block-FP8, down-only NF4, structural zeros, and optimizer parity remain checked." + ] + }, + { + "id": "TA-660", + "scope": "SGLang fused-MoE trainable numerical fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Stock-Triton gradient parity and masked-expert gradient semantics now form one trainable numerical policy.", + "Forward inputs, router weights, all parameter gradients, compacted references, masked zeros, and fully masked behavior remain checked." + ] + }, + { + "id": "TA-661", + "scope": "SGLang fused-MoE real parity fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Strided-versus-transient forward and gradient parity now shares one real-kernel policy with automatic-versus-explicit dispatch parity.", + "All output and gradient tensors, repeat determinism, explicit flag behavior, and stock output shape remain checked." + ] + }, + { + "id": "TA-662", + "scope": "canonical-MoE contributor-width parametrization", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Reference-tree widths and distributed contributor widths now execute inside their respective numerical and transport policies instead of reporting each width as a separate test.", + "The 2- and 16-contributor reference trees and the 2- and 8-process distributed workers all still execute." + ] + }, + { + "id": "TA-663", + "scope": "canonical-MoE transport admission fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Automatic selection, explicit-mode rejection, and direct packed or sharded reducer guards now form one transport-admission policy.", + "Admitted eager EP16 behavior, dense fallbacks, graph and consumer-output restrictions, topology rejection, and direct executor fail-closed behavior remain checked." + ] + }, + { + "id": "TA-664", + "scope": "world-32 canonical-MoE group-alias examples", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP8/eFSDP4 and EP16/eFSDP2 now form one world-32 topology policy rather than two layout-named tests.", + "Context-parallel aliases, expert-parallel aliases, and the first and last expert-FSDP groups remain exact for both layouts." + ] + }, + { + "id": "TA-665", + "scope": "Quack TP FP8 positional-forwarding mock", + "decision": "remove", + "status": "applied", + "evidence": [ + "The test only inspected positional arguments passed to an internal autograd function and returned a mocked tensor.", + "The retained CUDA TP lifecycle now executes the same FP8 backend and non-default block size through forward, backward, finite-gradient checks, and a master-weight update." + ] + }, + { + "id": "TA-666", + "scope": "prequantized checkpoint format-detection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "NVFP4 and block-FP8 detection now form one checkpoint-format policy rather than separate format-named tests.", + "Nested, flat, config, index, precedence, malformed, missing, wrong-block-size, and competing-format cases all remain checked." + ] + }, + { + "id": "TA-667", + "scope": "prequantized checkpoint-handler exclusion fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense and MoE exclusion behavior now forms one checkpoint-handler policy.", + "Weight and auxiliary-key passthrough, nonexcluded skipping, empty exclusions, shared experts, and on-load consistency remain checked." + ] + }, + { + "id": "TA-668", + "scope": "packing-concat sequence side-channel fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Teacher hidden states and hidden-match weights now report as one sequence-side-field collation policy.", + "Rank-two and rank-one concatenation, padding to a multiple of four, exact values, shapes, and zero padding remain checked." + ] + }, + { + "id": "TA-669", + "scope": "Quack compile-worker receive failure fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Silent-worker timeout and truncated-body rejection now form one receive-protocol policy.", + "The real pipe timeout bound and exact malformed-frame exception remain checked." + ] + }, + { + "id": "TA-670", + "scope": "Quack cache-key hashing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Structural determinism, unsafe-object rejection, and Cutlass dtype-class support now form one cache-key policy.", + "Tuple boundaries, type distinctions, disabled pickle hooks, deterministic hashes, and metaclass-independent dtype handling remain checked." + ] + }, + { + "id": "TA-671", + "scope": "Qwen3 pipeline-schedule parity parametrization", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The three candidate schedules now share one 1F1B baseline run instead of training the identical baseline once per parameterized report.", + "Interleaved1F1B, InterleavedZeroBubble, and ZBVZeroBubble still run independently with convergence and per-step loss-parity checks." + ] + }, + { + "id": "TA-672", + "scope": "NVFP4 tensor-quantization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed layout, dequantization error, and shared-global-scale behavior now form one tensor-quantization policy.", + "Packed shapes and dtypes, scale geometry, scalar global scale, relative error, and exact cross-tensor scale reuse remain checked." + ] + }, + { + "id": "TA-673", + "scope": "NVFP4 directory-export fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Weight-only export, W4A4 input-scale stamping, and requantization rejection now form one directory-export lifecycle.", + "Fused scales, BF16 islands, metadata, roundtrip error, calibrated input scales, uncalibrated omission, and fail-closed re-export remain checked." + ] + }, + { + "id": "TA-674", + "scope": "data-preparation retry fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Immediate success, retryable and terminal failures, and all backoff strategies now form one retry policy.", + "ReadTimeout, HfHubHTTPError, retry exhaustion, unrelated exceptions, exponential, linear, and constant timing remain checked; the production helper now imports HfHubHTTPError from its stable public module." + ] + }, + { + "id": "TA-675", + "scope": "OLMo2 construction and TP-layout fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "HF-config construction and unfusing into the HF parameter layout now form one architecture-layout policy.", + "Post-norm structure, full-axis QK norms, fused bias rules, split attention and MLP modules, and checkpoint-handler removal remain checked." + ] + }, + { + "id": "TA-676", + "scope": "OLMo2 checkpoint save and load fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "HF-compatible export and HF-to-fused import now form one bidirectional checkpoint policy.", + "All attention, norm, MLP, fused-key, strict loading, hidden-state, and logits assertions remain checked." + ] + }, + { + "id": "TA-677", + "scope": "Qwen2 construction and TP-layout fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "HF-config construction and TP unfusing now form one Qwen2 architecture-layout policy.", + "Norm absence, bias rules, split attention and MLP modules, and checkpoint-handler removal remain checked." + ] + }, + { + "id": "TA-678", + "scope": "Qwen2 checkpoint save and load fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "HF-compatible export and HF-to-fused import now form one bidirectional checkpoint policy.", + "Weight and bias keys, fused-key construction, strict loading, hidden-state parity, and logits parity remain checked." + ] + }, + { + "id": "TA-679", + "scope": "server cu-seqlen SP admission fragment", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The SP-enabled no-emission case now executes inside the server-versus-CLI alignment policy.", + "Two, three, single, and many-sequence boundaries, int32 dtype, maximum lengths, and SP ownership remain checked." + ] + }, + { + "id": "TA-680", + "scope": "orchestrator client communication fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Basic health roundtrip, repeated requests, and interleaved operation types now share one live ZeroMQ communication lifecycle.", + "Connection state, request IDs, response types, finished state, engine receipt, and output transmission remain checked without a second fixture startup." + ] + }, + { + "id": "TA-681", + "scope": "non-observing distributed data-loader examples", + "decision": "remove", + "status": "applied", + "evidence": [ + "A literal 4 times 3 equals 12 assertion did not invoke data-loader code, and the claimed epoch-consistency block only compared two list lengths fixed to three by construction.", + "Real partitioning, microbatching, sequence sharding, padding, drop-last, packed, multi-DP, and variable-length behaviors remain checked." + ] + }, + { + "id": "TA-682", + "scope": "joined multiprocessing outcome detection", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The audit now recognizes torch multiprocessing start_processes as an observable outcome because joined workers propagate child assertion and process failures.", + "Sequence-parallel gradient reduction, four-rank exact-DCP fusion, and DTensor materialization wrappers no longer appear as assertion-free candidates; all three executable wrappers pass." + ] + }, + { + "id": "TA-683", + "scope": "Muon and BI golden assertion-helper signaling", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Muon transition and frozen-bit golden helpers now use explicit assert-prefixed names so their wrapper outcomes are visible to the audit.", + "All three Muon topology transitions pass, and every H100-specific golden hash assertion remains unchanged." + ] + }, + { + "id": "TA-684", + "scope": "exact LM-head scalar optimizer-state acceptance outcome", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The acceptance test now explicitly asserts the validator's successful None result instead of relying only on absence of an exception.", + "Scalar tensor step state, replicated parameter ownership, LM-head TP group selection, and the coherence call remain unchanged and pass." + ] + }, + { + "id": "TA-685", + "scope": "optional QARL triton-w4a4 registry probe", + "decision": "remove", + "status": "applied", + "evidence": [ + "The CPU-marked probe usually skipped and otherwise checked only that two implementation-detail dictionary entries existed and one was callable.", + "QARL shadow selection, activation fake quantization, restoration, and the real W4A4 execution paths remain covered." + ] + }, + { + "id": "TA-686", + "scope": "group-GEMM and MoE kernel dependency gates", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Unsupported CPU hosts now skip at collection-time markers, while import failures in XORL's own grouped-GEMM and MoE kernel modules fail supported GPU runs instead of being converted into dependency skips.", + "All six retained CUDA kernel policies execute and pass on the audit host." + ] + }, + { + "id": "TA-687", + "scope": "non-gated MoE core import guard", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The suite now imports its core XORL expert class normally instead of catching every import-time exception and skipping affected tests.", + "All five retained CPU and CUDA backend policies execute and pass." + ] + }, + { + "id": "TA-688", + "scope": "SGLang missing-dependency diagnostic", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The negative test now simulates an unavailable SGLang import deterministically instead of skipping whenever SGLang is installed.", + "The test exposed and corrected a diagnostic that named the unrelated TP-simulation flag rather than XORL_MOE_SGLANG_FUSED_EXPERTS." + ] + }, + { + "id": "TA-689", + "scope": "personal-path repository hygiene guards", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The examples-only scanner duplicated the repository-wide home-path check and performed a second git traversal.", + "Mac home and personal data-workspace patterns now live in the stronger repository-wide guard, and the narrower file is removed." + ] + }, + { + "id": "TA-690", + "scope": "standalone runner-dispatcher forward model-id test", + "decision": "remove", + "status": "applied", + "evidence": [ + "The standalone test repeated the rank-zero forward handler scenario already covered in the request-processor suite.", + "The retained test additionally checks routed expert ids, routed logits, auto-load selection, rank ownership, and the returned session id." + ] + }, + { + "id": "TA-691", + "scope": "weight-version forwarding fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Two mock-to-mock forwarding tests are replaced by one composed handler-to-NCCL-synchronizer policy.", + "Bucket accounting, cache mode, and the exact weight version are verified at the final transfer seam." + ] + }, + { + "id": "TA-692", + "scope": "HSDP microbatch all-reduce fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Deferral, last-microbatch restoration, and the non-replicated rejection now form one gradient-sync policy.", + "Every original set_requires_all_reduce transition remains asserted." + ] + }, + { + "id": "TA-693", + "scope": "direct runtime-rank MoE LoRA scaling unit", + "decision": "remove", + "status": "applied", + "evidence": [ + "The direct one-by-one delta check was subsumed by the retained inference-buffer test that invokes the same helper for gate, up, and down projections.", + "The retained policy checks active-rank scaling, emitted names, shapes, dtype, values, and source cleanup." + ] + }, + { + "id": "TA-694", + "scope": "SGLang RMSNorm mode fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Global mode selection and forced-residual FP32 weight multiplication now report as one numerical mode policy.", + "Global state restoration and the exact reference calculation remain checked." + ] + }, + { + "id": "TA-695", + "scope": "index-share caller cleanup fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Offline trainer and server runner failure cleanup now execute in one cross-caller lifecycle policy.", + "Both failure messages, mode handoffs, and retained-context release counts remain asserted." + ] + }, + { + "id": "TA-696", + "scope": "importance-sampling metric reduction fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default ratio aggregation and custom TIS extrema now form one no-DP metric policy.", + "Weighted means, extrema, valid-token aggregation, and Python-scalar output remain checked." + ] + }, + { + "id": "TA-697", + "scope": "GLM LoRA target resolution fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Raw-HF default inference and explicit-target precedence now share one resolution policy and one model config fixture.", + "All five default attention targets and the exact explicit override remain asserted." + ] + }, + { + "id": "TA-698", + "scope": "remote backend RPC wrapper fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Weight-sync timeout forwarding and optimizer sparse-delta capture now execute in one operation-payload policy.", + "Operation names, request ids, timeout, pause and cache modes, endpoints, and sparse-delta fields remain checked." + ] + }, + { + "id": "TA-699", + "scope": "DSv4 RoPE cache-length fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Config-derived length and environment override precedence now execute in one cache-sizing policy.", + "The independent context-parallel short-cache rejection remains separate." + ] + }, + { + "id": "TA-700", + "scope": "SGLang JIT and kernel RMSNorm CPU fallback fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "JIT and kernel mode residual numerics plus packed-shape handling now form one CPU fallback matrix.", + "Both exact residual comparisons, the packed output shape, and global mode restoration remain checked." + ] + }, + { + "id": "TA-701", + "scope": "nonresident LoRA kill-session checkpoint fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Missing-checkpoint preservation and later evicted-checkpoint promotion now form one session lifecycle on the same runner.", + "Path traversal rejection remains an independent security boundary." + ] + }, + { + "id": "TA-702", + "scope": "QARL calibration batch loading fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Truncation, calibration-size limiting, attention-mask alignment, and malformed token-shape rejection now form one input policy.", + "The model calibration and persistent-state roundtrip remains a separate lifecycle." + ] + }, + { + "id": "TA-703", + "scope": "RMSNorm SGLang CPU mode modules", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Global SGLang selection, forced-residual FP32 multiplication, JIT fallback, kernel fallback, and packed inputs now live in one mode policy.", + "The redundant one-test JIT module is removed and the original global mode is restored atomically." + ] + }, + { + "id": "TA-704", + "scope": "NCCL rendezvous port fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sticky ephemeral-port rotation and explicit-port pinning now execute in one training rendezvous lifecycle.", + "Store creation, process-group destruction, active-port state, and bind-failure admission remain checked." + ] + }, + { + "id": "TA-705", + "scope": "FutureStore test-local response helper assertions", + "decision": "remove", + "status": "applied", + "evidence": [ + "The removed assertions exercised response-builder functions defined inside the test file rather than production code.", + "Production FutureEntry defaults, expiry, terminal states, and queue-state transitions remain checked." + ] + }, + { + "id": "TA-706", + "scope": "FQN matcher utility fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Single, all, and any matching now form one pattern-policy report instead of three one-function reports.", + "Exact, wildcard, grouped-number, indexed, prefixed, empty, first-match, and invalid-input cases all remain." + ] + }, + { + "id": "TA-707", + "scope": "block-FP8 quantization imports and input fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Contiguity and divisibility rejection now execute inside the quantization shape, dtype, scale, and block-size policy.", + "Imports of XORL's own block-FP8 module now fail visibly instead of being converted into an unavailable-feature skip; two unused imports were removed." + ] + }, + { + "id": "TA-708", + "scope": "dataset source loading fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local file, saved directory, hub, URL, missing source, and string or list data_files now form one source-resolution policy.", + "Download counts and every original source-selection assertion remain checked." + ] + }, + { + "id": "TA-709", + "scope": "exact DCP skip-mode fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Exact-model deferral and non-exact-model rejection now form one load-mode admission policy.", + "FSDP deregistration and the prohibition on unintended HF shard reads remain asserted." + ] + }, + { + "id": "TA-710", + "scope": "inert sparse-delta SGLang compatibility module", + "decision": "remove", + "status": "applied", + "evidence": [ + "A module-level guard skipped every test because its zstd case targeted disk_compression arguments absent from the production writer.", + "The retained sparse-delta file and backend suites pass, and the stronger trainer-to-request-processor-to-SGLang E2E owns receiver application, checksum, validate-only, and final parameter parity." + ] + }, + { + "id": "TA-711", + "scope": "launcher rank-zero address precedence fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Remote discovery and explicit engine-host precedence are two outcomes of the same address-selection policy.", + "The retained report checks the discovery call and the explicit-host short circuit." + ] + }, + { + "id": "TA-712", + "scope": "launcher server-override parsing fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Schema-agnostic parsing and removed-field validation are consecutive boundaries of one override policy.", + "Arbitrary parsed values and the ZORL migration diagnostic remain asserted." + ] + }, + { + "id": "TA-713", + "scope": "DistSignSGD local hook ownership fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local-parameter registration and FSDP-managed exclusion now execute in one ownership report.", + "The retained policy checks hook installation, configuration state, and absence of a duplicate managed-parameter hook." + ] + }, + { + "id": "TA-714", + "scope": "NVFP4 QARL normalization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Alias defaults, activation override, and every invalid block size are outcomes of one normalization policy.", + "Weight-only defaults, fixed group size, explicit activation, and all original rejection literals remain checked." + ] + }, + { + "id": "TA-715", + "scope": "direct relu2 activation-registry probe", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained non-gated MoE reference test constructs relu2 experts and proves exact forward and gradient behavior against relu squared.", + "A separate assertion that relu2 appears in internal dictionaries provided no stronger production regression signal." + ] + }, + { + "id": "TA-716", + "scope": "server Adam configuration and initializer fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "ServerArguments conversion now feeds the ModelRunner initializer in the same test instead of ending at an intermediate dictionary.", + "Non-default betas and epsilon, defaults, optimizer parameter groups, and malformed-beta rejection remain checked." + ] + }, + { + "id": "TA-717", + "scope": "Pydantic API field-echo and automatic roundtrip reports", + "decision": "remove", + "status": "applied", + "evidence": [ + "Two broad reports primarily asserted that Pydantic returned constructor fields and rejected omitted required fields.", + "Real training, checkpoint, sampler, and session endpoint tests construct these models; the unique forward session-id alias was retained in the compatibility policy." + ] + }, + { + "id": "TA-718", + "scope": "runner protocol constructor and default-factory fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Direct payload-field echoes, UUID uniqueness, timestamp range, and optional-None assertions are replaced by one full typed-wire equality contract.", + "All operation payloads, success and error responses, tensors, JSON, ACK correlation, and pickle rejection remain checked." + ] + }, + { + "id": "TA-719", + "scope": "duplicate API-orchestrator request-response flow", + "decision": "remove", + "status": "applied", + "evidence": [ + "The removed flow repeated subsets of the adjacent request and output roundtrips, builders, validators, and streaming-error checks.", + "The retained orchestrator integration suite separately exercises real queue processing and request identity across the live protocol seam." + ] + }, + { + "id": "TA-720", + "scope": "direct EP checkpoint-mesh selection fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy and PP-parent selection are now exercised through the production restore consumer instead of separate direct helper reports.", + "Both selected shapes, named dimensions, placements, ModelState delegation, and malformed-dimension rejection remain checked." + ] + }, + { + "id": "TA-721", + "scope": "DeepEP buffer-size validation fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "RDMA overflow, no-RDMA allowance, 128-byte alignment, and the default two-GB case are one buffer admission policy.", + "Every original size, RDMA allocation, expected byte count, and overflow diagnostic remains checked." + ] + }, + { + "id": "TA-722", + "scope": "weight-sync quantization normalization fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "BF16 no-ops and valid FP8 forms now share one supported normalization policy; unsupported methods and malformed FP8 forms share one rejection policy.", + "All aliases, defaults, exclusions, invalid formats, activation schemes, scale storage, and contextual diagnostics remain checked." + ] + }, + { + "id": "TA-723", + "scope": "direct OPD teacher-sort helper probe", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Nested teacher_id, nested teacher_ids, and top-level precedence now execute through the real OPD model-pass sorting branch.", + "The retained transaction additionally proves that teacher sorting composes with packer datum order and Mooncake routing-payload order." + ] + }, + { + "id": "TA-724", + "scope": "SignSGD sparse-step and base-optimizer state fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense updates, decoupled decay, missing gradients, and sparse rejection now form one SignSGD step policy.", + "A separate report of PyTorch Optimizer state-dict behavior was removed; builder admission and parameter-group ownership remain independent." + ] + }, + { + "id": "TA-725", + "scope": "direct FP32 routing-scale helper probe", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained local fused-MoE backward oracle uses the same discriminating BF16 gradient and small FP32 routing value.", + "It compares every input, routing, gate-up, and down gradient exactly, so a pre-multiply BF16 routing cast still fails through production autograd." + ] + }, + { + "id": "TA-726", + "scope": "sparse-MLA machine-specific speed assertion", + "decision": "relocate", + "status": "applied", + "evidence": [ + "The production-shape H100 timing ratio is hardware certification rather than repository correctness and no longer runs under pytest.", + "An explicit certification script retains combined-versus-split warmup, median timing, environment restoration, and a configurable speedup gate; all three numerical kernel reports remain." + ] + }, + { + "id": "TA-727", + "scope": "unasserted vocab-parallel CE benchmark inside distributed correctness", + "decision": "relocate", + "status": "applied", + "evidence": [ + "The pytest worker previously ran 100 production-scale warmup and timed forward or backward iterations only to print tables after correctness had passed.", + "The retained two-rank test now ends after eager and compiled value and gradient parity; an explicit torchrun certification script owns timing and peak-memory reporting." + ] + }, + { + "id": "TA-728", + "scope": "single-GPU dense FP8 CLI training smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained checkpoint-and-resume test begins with the same eager one-GPU two-step FP8 training configuration and the same finite-loss, gradient, and module-usage assertions.", + "The stronger retained test passed both training phases and additionally verifies DCP checkpoint creation and resume." + ] + }, + { + "id": "TA-729", + "scope": "basic two-GPU DistSignSGD training smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained four-GPU test exercises the same FSDP2, gradient-accumulation, five-step optimizer path while adding Ulysses exact-sum and replicated data-parallel composition.", + "After selecting eager attention to avoid an unrelated tiny-shape FlashAttention compiler failure, the retained composition test passed; unit tests separately pin both no-SP and SP sign-reduction behavior." + ] + }, + { + "id": "TA-730", + "scope": "Nemotron-H packed all-block-types smoke", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed sequence boundaries now flow through the retained all-mixer loss and backward contract instead of a second finite-output smoke.", + "The retained test checks loss plus finite nonzero gradients for Mamba, attention, routed experts, latent projections, shared experts, and embeddings." + ] + }, + { + "id": "TA-731", + "scope": "Qwen3 unfused forward-shape smoke", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Two independently initialized MLPs only proved that fused and unfused implementations returned their declared hidden shape; they did not compare values.", + "Direct projection replacement and model-wide unfuse ownership now form one CPU policy covering attention and MLP modules in every layer." + ] + }, + { + "id": "TA-732", + "scope": "fragmented PP NCCL sender and receiver protocol tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One composed mocked roundtrip now captures the sender's real metadata and flattened payload and replays both directly through the receiver.", + "Names, matrix, vector, and scalar shapes, BF16 storage, and exact reconstructed values remain checked; the empty-payload protocol remains separate." + ] + }, + { + "id": "TA-733", + "scope": "duplicate families-v2 fused-versus-split norm comparison", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained norm contract explicitly forces both realizations, proves the split implementation was reached, and compares more hidden sizes plus residual, plain, and zero-centered variants.", + "The retained dispatch module still independently proves that shipped sizes select fused execution and deep shapes reach split execution." + ] + }, + { + "id": "TA-734", + "scope": "three one-rank exact-GLM FSDP2 component wrappers", + "decision": "remove", + "status": "applied", + "evidence": [ + "Each worker already supports world sizes one and two; its retained two-rank wrapper executes every shared lifecycle, byte-parity, ownership, and gradient assertion.", + "The two-rank branches additionally prove that LoRA factors are genuinely sharded. Both wrapper variants share the same optional SGLang import gate, which is unavailable in the repository venv." + ] + }, + { + "id": "TA-735", + "scope": "standalone families-v2 dispatch helper suite", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The retained norm policy now spies on the production split implementation while exercising shipped hidden sizes and a deep split shape, so it proves actual dispatch rather than only the private predicate.", + "All direct boundary cases remain in that covering policy, and the separate fused-versus-split numerical contract remains independent." + ] + }, + { + "id": "TA-736", + "scope": "direct optimizer-step learning-rate resolver report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit request, registered session, server training config, create-model registration, missing-value rejection, and both legacy fallbacks now drive the actual optimizer-step payload path.", + "The standalone report only called the private resolver with synthetic namespaces and is no longer needed." + ] + }, + { + "id": "TA-737", + "scope": "direct default Qwen3.5 rotary helper comparison", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained dense-and-MoE attention integration reaches the production QKV projection and compares its rotated Q/K values with the same Hugging Face half-rotation reference.", + "It also rejects the pairwise alternative, so the direct helper probe was a strict subset of actual attention behavior." + ] + }, + { + "id": "TA-738", + "scope": "seven fragmented dense Qwen3.5 RMSNorm reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Two retained policy reports cover exact-family selection, fused serving families, explicit v2 admission and rejection, every v2 norm site, GDN separation, layer-input modes, and final-norm modes.", + "No norm case was discarded; the old reports were narrow fragments of the same construction and site-assignment contracts." + ] + }, + { + "id": "TA-739", + "scope": "direct dense and MoE Qwen3.5 config-conversion reports", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "One retained family policy writes realistic local config files and loads both variants through the production local auto-config entry point.", + "The covering path preserves registry admission, derived layer schedules, head geometry, linear-attention fields, mRoPE extraction, and MoE geometry assertions." + ] + }, + { + "id": "TA-740", + "scope": "direct Qwen3.5 checkpoint skip-regex report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained grouped checkpoint loader now uses the real Qwen3.5 skip-pattern constant while routing dense and expert shards.", + "Both top-level MTP forms and the non-MTP negative case are asserted inside the production prefetch filter transaction." + ] + }, + { + "id": "TA-741", + "scope": "duplicate families-v2 environment kill-switch report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained family-selection policy already iterates every trainer and sampler environment variable and verifies default-on plus per-variable rollback.", + "Runtime dispatcher rollback remains separately covered with a kernel-reachability spy." + ] + }, + { + "id": "TA-742", + "scope": "one-line no-shared-prefix repack report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The no-op fallback is now the opening case of the complete repack and remap policy.", + "Shared grouping, layout, position ids, loss fields, cumulative lengths, and output remapping remain covered in the same report; the P-equals-one edge remains independent." + ] + }, + { + "id": "TA-743", + "scope": "fragmented active-LoRA admission truth table and topology reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One composite-admission policy now covers the complete family, every missing component, and the independent scoring-only marker.", + "The production model build now proves both complete flag derivation and non-TP16 rejection in one topology transaction." + ] + }, + { + "id": "TA-744", + "scope": "direct BI router batch-invariance and default-path reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Batch composition invariance now runs through the production MoEBlock route method and checks logits, selections, and normalized weights.", + "Exact-contract selection and the ordinary BF16 route remain together as the two branches of one production dispatch policy." + ] + }, + { + "id": "TA-745", + "scope": "separate standard and temperature BI fused-LM-head oracle reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One forward-and-backward policy compares fused and eager execution at unit and non-unit temperatures.", + "Per-token log probabilities, ignored-token loss, aggregate loss, hidden gradients, and weight gradients remain asserted." + ] + }, + { + "id": "TA-746", + "scope": "three local step phase and memory summary helper fragments", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical ordering, custom ordering, empty input, numeric coercion, phase aggregates, and memory aggregates now form one local-finalization policy.", + "All prior edge cases and exact aggregate values remain covered." + ] + }, + { + "id": "TA-747", + "scope": "optimizer and weights Pydantic field-echo report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The removed assertions primarily repeated automatic constructor assignment, defaults, schema properties, and serialization.", + "The unique legacy optimizer session-id alias now drives the actual optim-step endpoint and payload; retained create-model, create-session, and weights endpoints cover the request and response models." + ] + }, + { + "id": "TA-748", + "scope": "separate multi-part optimizer selection and update reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One optimizer lifecycle now proves multi-part selection, parameter coverage, updates for every part, gradient clearing, scheduler propagation, and single-part fallback.", + "Custom parameter-group rejection remains an independent admission boundary." + ] + }, + { + "id": "TA-749", + "scope": "flat pseudo-packed tensor-collator report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The old report described equal-length flat examples as packed even though the collator intentionally preserves them as separate samples.", + "The retained policy now exercises the real already-batched dict and nested packed-dataset branches while preserving general conversion and dtype coverage." + ] + }, + { + "id": "TA-750", + "scope": "direct OPD list-chunking helper report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The helper is a guarded slicing comprehension, and the production pipeline only enables chunking for positive sizes.", + "The removed report merely repeated Python list slicing; payload alignment, endpoint registration, version verification, and pipeline lifecycle behaviors remain." + ] + }, + { + "id": "TA-751", + "scope": "sequence-parallel no-group identity report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The no-group branch directly returns the input tuple and the test asserted only object values through that one-line identity path.", + "Padding roundtrips and real two-rank gather backward reduction remain as the behavioral sequence-parallel reports." + ] + }, + { + "id": "TA-752", + "scope": "direct SGLang FP32 grouped-GEMM accumulator probe", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained local and EP custom-autograd reports execute the accumulator through the production backward functions.", + "Their exact independent routing-gradient oracles use values that distinguish FP32 accumulation from BF16 rounding, so the direct helper probe was a strict subset." + ] + }, + { + "id": "TA-753", + "scope": "train-router argument and configuration field echoes", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The server default now feeds MoEBlock.from_config and proves by backward behavior that the gate remains detached.", + "The same policy retains the enabled all-to-all gradient path and the unsupported DeepEP rejection." + ] + }, + { + "id": "TA-754", + "scope": "standalone packing-cache string-format report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The old report only repeated the interpolation fields of generate_packing_hash and missed its document-alignment branch.", + "The retained PackingDataset lifecycle now obtains real cache paths and proves ring-attention alignment creates a distinct cache identity." + ] + }, + { + "id": "TA-755", + "scope": "thin MD5 and SHA256 wrapper report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report retested standard-library digest constants and string encoding through one-line wrappers.", + "Dataset split and configuration fingerprint policies retain higher-level determinism, sensitivity, format, and order-independence coverage for the production MD5 consumer; the SHA256 wrapper has no production caller." + ] + }, + { + "id": "TA-756", + "scope": "BI GEMM post-import environment and pinned-constant tautologies", + "decision": "remove", + "status": "applied", + "evidence": [ + "Setting legacy environment variables after importing the module could not prove import-time independence, and the named variables have no production readers.", + "The block-K report asserted that lookup returned the same constant it directly injects; retained CUDA reports instead prove table bit neutrality, cross-bucket row invariance, and DeepGEMM parity." + ] + }, + { + "id": "TA-757", + "scope": "four fragmented EP adapter backend argument-boundary reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Native, Triton, Triton MoE-act, and Quack variants now form one capability-aware backend matrix.", + "Common-argument consumption, explicit unsupported-FP8 rejection, and Quack activation-native forwarding all remain; registry signatures and numerical expert-score forwarding stay independent." + ] + }, + { + "id": "TA-758", + "scope": "separate routing-replay wire decode and weight tensor reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One wire-to-tensor policy now decodes SGLang base64, inferred-shape, Python, NumPy, and tensor forms before exercising float weights, padding, and sequence-parallel slicing.", + "Ring-attention and general sequence-parallel layout contracts remain independent." + ] + }, + { + "id": "TA-759", + "scope": "separate OPD output-edge and hidden-only reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All-ignored stability, per-token output, and zero-KL hidden-only mode now form one output-edge policy.", + "Numerical backend parity, gradient reduction, and OPRD hidden-distance contracts remain separate." + ] + }, + { + "id": "TA-760", + "scope": "separate DCP synchronization and metadata process-group selector reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One process-group policy now covers cached Gloo creation under NCCL, no-op behavior under Gloo, PP-disabled metadata, caller precedence, and global fallback.", + "All prior selection branches and call-count assertions remain." + ] + }, + { + "id": "TA-761", + "scope": "separate ParallelState default and custom-construction reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Defaults, uninitialized properties, invalid modes, distributed rank discovery, topology validation, and enabled flags now form one construction policy.", + "Mesh construction and singleton initialization lifecycles remain independent." + ] + }, + { + "id": "TA-762", + "scope": "separate EP LoRA initialization and plan-slicing reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One lifecycle now follows global GKN factors from initialization and zero-B state through plan registration into exact rank-zero and rank-one expert slices.", + "The retained CUDA report independently exercises EP and non-EP forward and gradient behavior." + ] + }, + { + "id": "TA-763", + "scope": "separate accepted and rejected external FP8 configuration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One compatibility policy now admits the supported NeMo blockwise form and rejects every unsupported recipe or external runtime configuration.", + "Blackwell hardware admission and BF16 layer-island injection remain independent policies." + ] + }, + { + "id": "TA-764", + "scope": "separate QARL activation-override nesting report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One context-manager lifecycle now covers mixed prior values, both override directions, plain-module exclusion, exception restoration, and nested inner-first restoration.", + "No activation-override branch was discarded." + ] + }, + { + "id": "TA-765", + "scope": "separate QARL calibration loading and persistent-state reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One calibration lifecycle now follows truncated JSON input and malformed-shape rejection through model calibration, metadata population, state serialization, and restoration.", + "The retained assertions cover both input admission and durable quantization state." + ] + }, + { + "id": "TA-766", + "scope": "standalone NVFP4 normalization report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "NVFP4 aliases, activation selection, and invalid group sizes now feed the QARLLinear forward and STE policy instead of ending at normalized dictionary fields.", + "Dense-model injection remains an independent production composition report." + ] + }, + { + "id": "TA-767", + "scope": "separate NVFP4 activation forward and STE reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One autograd contract now proves exact dequantized NVFP4 forward values, actual lossiness, and identity straight-through gradients.", + "The forward and backward halves exercise the same activation fake-quant primitive." + ] + }, + { + "id": "TA-768", + "scope": "separate QARL MoE backend-shadow admission and exception reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One shadow-context policy now covers armed Triton W4A4 selection, exception restoration, activation-disabled Triton, and activation-enabled eager no-ops.", + "Every backend and restoration branch remains asserted." + ] + }, + { + "id": "TA-769", + "scope": "separate QARL sync-configuration and handler-derivation reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One training-to-sync lifecycle derives folded modules, quantizes only the QARL weights, preserves BF16 islands, and passes the same configuration through the production WeightSyncHandler.", + "Malformed caller-supplied QARL sync configuration remains a separate rejection boundary." + ] + }, + { + "id": "TA-770", + "scope": "standalone stochastic-rounding seeded-generator report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Seeded-generator reproducibility now belongs to the same call contract as shape, device, output dtype, and invalid-input admission.", + "Unbiased expectation and adjacent-BF16-neighbor properties remain independent statistical and numerical policies." + ] + }, + { + "id": "TA-771", + "scope": "separate Triton QARL MoE weight-quantization disabled report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One real Triton policy now compares the quantized lossy forward and STE gradients with an exact unquantized passthrough from the same seeded weights and routing.", + "This removes a second backend launch that repeated setup without a distinct failure domain." + ] + }, + { + "id": "TA-772", + "scope": "fragmented API training response reports and duplicate legacy optimizer payload", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forward auto-load information and forward-backward executor timings now form one API response-projection policy.", + "The focused optimizer report retains response telemetry and current-control mapping; the broader optimizer-step policy already owns legacy aliases, Adam fields, defaults, precedence, and missing-value rejection." + ] + }, + { + "id": "TA-773", + "scope": "standalone RequestProcessor lifecycle and statistics reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One live processor now starts ready, carries three real forward and forward-backward operations, exposes their counters, and then stops.", + "The removed reports repeated processor construction or a model pass solely to inspect readiness and counter fields." + ] + }, + { + "id": "TA-774", + "scope": "scheduler default-field, repr, and raw FIFO helper report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report asserted constructor constants, class names, and low-level deque methods without carrying a schedulable request through execution.", + "Retained Scheduler transactions prove FIFO dispatch, capacity, pending removal, and clear behavior through the production boundary." + ] + }, + { + "id": "TA-775", + "scope": "separate scheduler terminal-state and statistics-history reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One lifecycle now completes, fails, aborts pending and actually running requests, checks terminal statistics, clears state, and verifies bounded history.", + "The old running-abort block dispatched an older FIFO item and therefore silently aborted its named request while it was still pending." + ] + }, + { + "id": "TA-776", + "scope": "GPU-marked MicroBatchCollator order and uneven-size report in the distributed data-loader suite", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report created no process group, rank, sampler, or sequence shard and duplicated the CPU micro-batch splitting and error contract.", + "Distributed partitioning, sequence sharding, and packed-pipeline reports remain independent." + ] + }, + { + "id": "TA-777", + "scope": "standalone packed position-id and missing-label report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Generated ignored labels, document-position resets, and caller-position replacement now execute after the retained pack-to-unpack roundtrip.", + "The behavior is metadata for the same packed transaction, not an independent failure domain." + ] + }, + { + "id": "TA-778", + "scope": "standalone sequential-packing legacy-layout report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Exact legacy greedy grouping now begins the all-strategy correctness and utilization policy.", + "Best-fit and balanced-DP document preservation remain checked against that sequential baseline in the same report." + ] + }, + { + "id": "TA-779", + "scope": "duplicate mixed-valid and oversized skip assertion in generic packing edge cases", + "decision": "remove", + "status": "applied", + "evidence": [ + "The packing-strategy admission contract already owns the exact mixed-input legacy-drop behavior.", + "The generic edge policy retains the distinct all-skipped rejection, empty, single, missing-input, and NumPy cases." + ] + }, + { + "id": "TA-780", + "scope": "standalone orchestrator statistics and health-classifier report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Scheduler and processor statistics are now asserted after the retained forward, optimizer, and end-to-end health transactions.", + "Callable-existence checks, repeated read-only getter calls, and private classifier truth tables added no behavior beyond that real health response." + ] + }, + { + "id": "TA-781", + "scope": "orchestrator integration abort block that asserted only the original request produced output", + "decision": "remove", + "status": "applied", + "evidence": [ + "The DummyBackend commonly completed before the abort arrived, and the assertion neither checked abort acknowledgement nor an aborted terminal state.", + "The scheduler lifecycle now deterministically proves pending and active abort semantics." + ] + }, + { + "id": "TA-782", + "scope": "standalone MoE weight-sync bucket-size precedence report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Transport defaults and environment precedence now select the cap used by the retained byte-based bucket split policy.", + "Oversized-first-item behavior remains asserted in the same production helper lifecycle." + ] + }, + { + "id": "TA-783", + "scope": "standalone endpoint cache-metadata normalization report", + "decision": "remove", + "status": "applied", + "evidence": [ + "Current cache_epoch metadata now returns through a complete streaming FP8 sync, including post-process configuration and transfer flags.", + "Legacy cache_version normalization remains covered by the complete sparse-delta sync response." + ] + }, + { + "id": "TA-784", + "scope": "standalone compile-wrapper weight-name normalization report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Compiled expert, broadcast, and Qwen linear-attention names now normalize inside the full inference-layout unfusion policy.", + "The retained assertions verify emitted receiver names and tensor values, not only the string helper." + ] + }, + { + "id": "TA-785", + "scope": "direct P2P warm-mode selector report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Warm-only mode now performs real buckets after cached and cold prepare states.", + "The fake Mooncake engine observes async submission only for cached prepare and synchronous transfer for cold prepare." + ] + }, + { + "id": "TA-786", + "scope": "direct P2P small-entry transfer helper report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Three GPU-direct entries now pass through transfer_bucket and flush twice with a two-entry Mooncake chunk limit.", + "Observed engine calls prove the 2-1 chunk sequence and receiver addresses through the production worker." + ] + }, + { + "id": "TA-787", + "scope": "direct persistent-source interval registration report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Repeated GPU-direct buckets backed by one stable allocation now register it exactly once and reuse it on the second bucket.", + "Backend destruction proves the persistent range is deregistered during cleanup." + ] + }, + { + "id": "TA-788", + "scope": "P2P direct-EP capability and sender-rank field echoes", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report stopped at constructor properties for implicit and explicit sender ranks.", + "Retained direct-EP policies use those configurations for filtered scatter, dense partitioning, process-group collectives, failure propagation, prewarming, and rank-owned transfer." + ] + }, + { + "id": "TA-789", + "scope": "fragmented EP checkpoint mesh selector, ModelState caller, and restore reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One restore policy now rejects ambiguous meshes, routes ModelState through named dimensions, and constructs the final DTensor across legacy and PP parent layouts.", + "Dropping the EP dimension remains an independent reverse-conversion report." + ] + }, + { + "id": "TA-790", + "scope": "separate checkpoint materialization success and zero-meta failure reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One load-state policy now proves successful materialization, missing-optimizer omission, restored counters, and both pre-restore and post-restore meta-storage rejection.", + "Each failure still asserts whether the checkpointer was reached." + ] + }, + { + "id": "TA-791", + "scope": "separate ModelRunner initial-load and restore-completion wrapper reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The production restore wrapper now calls the real initial-load method for optimizer-enabled and optimizer-disabled policies and synchronizes both counters.", + "A failing checkpoint manager proves the completion flag stays false after the actual load call raises." + ] + }, + { + "id": "TA-792", + "scope": "three copies of canonical-LoRA sampler-weight export across API suites", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained checkpoint-path contract starts from a normalized session spec and verifies save_lora_only, model identity, destination path, and returned xorl URI.", + "Two smaller copies asserted subsets of the same APIServer method with equivalent fake responses." + ] + }, + { + "id": "TA-793", + "scope": "separate create-model conflicting-recreate and worker-registration-failure reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One admission and rollback policy now covers an existing incompatible session, a mismatched base repository, and a cross-rank registration failure.", + "Failed registration is still proven not to mutate model configuration or registered IDs." + ] + }, + { + "id": "TA-794", + "scope": "direct base-model canonicalizer truth-table report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Create-model registration now succeeds when the server uses an HF snapshot path and the client uses the equivalent repository ID.", + "The same endpoint policy rejects a snapshot path for a distinct repository." + ] + }, + { + "id": "TA-795", + "scope": "standalone last-inference-endpoint adapter-tracking cleanup report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sampler listing, path resolution, recency tracking, and last-receiver removal now form one adapter-tracking lifecycle.", + "Removal proves both endpoint state and receiver-derived adapter state are cleared." + ] + }, + { + "id": "TA-796", + "scope": "separate receiver quantization detection, enrichment, and set-policy reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One receiver policy now reads config.json, normalizes receiver names, enriches a caller FP8 config, and admits that exact result as the sync default.", + "MTP, static activation, UE8M0, compressed-tensors, and BF16 boundaries remain asserted in the same flow." + ] + }, + { + "id": "TA-797", + "scope": "FutureEntry constructor, manually assigned terminal states, and standalone queue-state report", + "decision": "remove", + "status": "applied", + "evidence": [ + "Real FutureStore jobs already prove pending, processing, completed, failed, and expired states plus results and classified errors.", + "Queue pause state now follows actual concurrent processing and statistics rather than a freshly constructed empty store." + ] + }, + { + "id": "TA-798", + "scope": "separate DeepSeek training-builder router rejection and successful-freeze reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One builder policy now rejects unfrozen, QLoRA, and unmerged-QKV modes before exercising the admitted frozen-router construction.", + "The admitted path proves router parameters are frozen while ordinary attention parameters remain trainable." + ] + }, + { + "id": "TA-799", + "scope": "direct DeepSeek tensor-parallel validator report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained build_parallelize_model policy invokes the same validator through its production consumer.", + "It still proves DeepSeek tensor parallelism fails before unsupported parallelization can proceed." + ] + }, + { + "id": "TA-800", + "scope": "duplicate full Trainer bootstrap setup for causal-loss lm-head mode", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One bootstrap policy now covers parallel-state initialization and both eager and Quack-linear causal-loss parameterization.", + "The Quack branch still proves lm_head_fp32 is omitted rather than passed to the unsupported loss implementation." + ] + }, + { + "id": "TA-801", + "scope": "separate disabled and enabled manual CUDA-timing reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One state-machine lifecycle now starts disabled, rejects an invalid mode, records forward and recompute phases, and drains the accumulator.", + "The fake-event path still proves unrecorded CUDA event pairs are omitted." + ] + }, + { + "id": "TA-802", + "scope": "direct LoRA dtype and generic-upcast helper rows already exercised by the model builder", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The production builder policy observes BF16 base weights, FP32 adapters, and skip_param_upcast for mixed-precision LoRA.", + "Only the distinct QLoRA, explicit-skip, and dense-default helper branches remain alongside that lifecycle." + ] + }, + { + "id": "TA-803", + "scope": "P2P expert FP8 production geometry and all-global-index certification sweeps", + "decision": "remove", + "status": "applied", + "evidence": [ + "Global expert names are produced by the size-independent ep_rank times local expert count plus local index formula, so eight full EP-rank transactions selected no new branch.", + "Two retained transactions cover partial blocks plus block-128 quantization, single and multiple receivers, and a nonzero EP offset while checking exact bytes, scales, and dequantized values." + ] + }, + { + "id": "TA-804", + "scope": "production-sized Qwen3.6 shared-expert P2P parity report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained shared-expert transaction exercises singular and plural receiver namespaces, fused gate/up placement, down projection placement, scales, and the passthrough gate.", + "The removed 2048-by-512 allocation repeated those same locator and transfer branches with larger tensors only." + ] + }, + { + "id": "TA-805", + "scope": "separate P2P compiled-name, language-model-prefix, and tied-lm-head lookup reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One receiver-name compatibility policy now performs real transfers for _orig_mod stripping and language_model prefix fallback.", + "The same policy proves a missing tied lm_head locator is skipped without duplicating the embedding transfer." + ] + }, + { + "id": "TA-806", + "scope": "separate direct-EP dense manifest and outgoing-buffer ownership reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One ownership policy assigns every dense family to exactly one sender, keeps fused and split projection aliases together, and applies that assignment to both receiver manifests and outgoing buffers.", + "Expert entries remain rank-owned in manifests and excluded from the dense buffer filter." + ] + }, + { + "id": "TA-807", + "scope": "duplicate nonzero-sender initialization setup for default and prewarmed P2P engines", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One initialization policy runs both environment modes against the same scattered tensor-map contract.", + "It proves default construction follows broadcast while prewarming constructs the engine before broadcast." + ] + }, + { + "id": "TA-808", + "scope": "separate direct-EP rank-filter routed and all-filtered transfer reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One transfer policy routes each slice only to its owning receiver for two ranks.", + "The same policy covers a non-owning rank whose valid bucket is fully filtered and therefore emits no engine transfer." + ] + }, + { + "id": "TA-809", + "scope": "separate Qwen3-MoE layer and final RMSNorm family-declaration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One construction policy now checks layer-zero, later-layer, and final-model norm families.", + "It still executes the decoder call sites and proves the final norm relies on its module declaration rather than a per-call override." + ] + }, + { + "id": "TA-810", + "scope": "direct delayed-TP-shard residual-association helper report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The real Qwen3-MoE decoder-layer path now receives BF16 shard values for which sum-then-residual differs from sequential residual addition.", + "The retained test observes the materialized input, norm call, residual output, and diagnostics through the production consumer." + ] + }, + { + "id": "TA-811", + "scope": "direct O-projection partial-residual mode helper report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Both sum-then-residual and residual-then-partials modes now run through Qwen3MoeDecoderLayer._pre_mlp_forward with discriminating BF16 values.", + "The layer test verifies the norm input, returned residual, output, diagnostic captures, and the expected bitwise association difference." + ] + }, + { + "id": "TA-812", + "scope": "standalone cross-engine RMSNorm funnel and test-discriminator reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "SGLang single-tensor and fused funnel outputs now sit inside the corresponding qk, pre-summed residual-tree, and post-attention module policies across every adversarial shape.", + "The rare family-difference discriminator now guards the qk site policy directly rather than reporting a test of the test." + ] + }, + { + "id": "TA-813", + "scope": "masked cross-engine RMSNorm trunk-flag report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The fixture explicitly declared the no-residual family, which selects the batch-invariant wrapper before the trunk-flag branch and made the flag unable to change dispatch.", + "The retained qk module policy reaches that wrapper and compares both the serving kernel and serving family funnel across all adversarial shapes." + ] + }, + { + "id": "TA-814", + "scope": "mocked distributed data-loader partitioning, micro-batch, and sequence-parallel report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The partitioning outcome was computed from rank-index lists constructed by the test itself; the injected sampler was never observed through the dataloader.", + "The retained builder policy verifies sampler rank and replica ownership plus SP pipeline insertion, while the MicroBatchCollator policy verifies exact split values and failures." + ] + }, + { + "id": "TA-815", + "scope": "shallow mocked sequence-sharding shape report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report asserted only equal shard lengths and a loose padding range after iterating mocked ranks.", + "The retained TextSequenceShardCollator policies check exact rank slices, non-divisible padding, attention metadata, labels, and token-aligned side channels through the production collator." + ] + }, + { + "id": "TA-816", + "scope": "mock-rank packed data-loader shape examples", + "decision": "remove", + "status": "applied", + "evidence": [ + "Changing mocked DP ranks while asserting the same output shape could not detect incorrect partition ownership or data overlap.", + "The retained real data-loader lifecycle covers packed and variable-length samples, while PackingConcatCollator owns exact concatenated values, position resets, extra fields, padding, and flash-attention metadata." + ] + }, + { + "id": "TA-817", + "scope": "direct fresh-manager checkpoint optimizer-selection report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report loaded a SignSGD checkpoint into an AdamW-default manager and asserted only the selected optimizer type and normalized session field.", + "The retained real AdapterCoordinator lifecycle performs that same transition through both explicit checkpoint loading and eviction auto-loading, while the multi-adapter lifecycle also reloads mixed optimizer types." + ] + }, + { + "id": "TA-818", + "scope": "one-layer row of the GLM semantic MoE stack parameterization", + "decision": "remove", + "status": "applied", + "evidence": [ + "Changing the synthetic stack from four MoE layers to one changed only repetition count and selected no distinct model or canonicalization branch.", + "The retained four-layer transaction checks every boundary, final logprob parity, batch permutation, per-row composition, and a deliberately omitted first-layer canonicalization that changes the final result." + ] + }, + { + "id": "TA-819", + "scope": "shape-only Qwen Triton expert forward report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report initialized a Triton expert, ran one random forward, and asserted only that the output shape matched the input shape.", + "The retained eager-versus-Triton MoE transaction exercises the same backend with numerical output comparison and gradients for every LoRA factor." + ] + }, + { + "id": "TA-820", + "scope": "synthetic DeepSeek-like MLA LoRA target stub report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Its default-target assertions duplicated the retained real DeepSeek model policy over the same MLA projection names.", + "The unique explicit-target partition case now runs through inject_lora_into_model_with_moe on the real DeepSeek model and proves the untargeted output projection remains unchanged." + ] + }, + { + "id": "TA-821", + "scope": "base MoE expert constructor field, shape, mapping, and registry echoes", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report instantiated wrappers and backend variants but asserted only stored strings, parameter shapes, mapping identity, and registry membership.", + "Retained eager, native, Triton, non-gated, injection, and model-construction policies execute those registrations and layouts; the separate LoRA initialization policy still protects frozen bases and trainable factor state." + ] + }, + { + "id": "TA-822", + "scope": "DSv4 KV-QAT helper boolean and private-field report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The removed report stopped after calling dsv4_kv_qat_enabled and reading DeepSeekV4Attention._kv_qat_enabled.", + "The retained C0 attention forward/backward transaction now supplies FP8 quantization configuration and observes the QAT call on the exact no-RoPE KV slice with block size 64." + ] + }, + { + "id": "TA-823", + "scope": "standalone DSv4 RoPE cache-length precedence shape report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The report called the cache builder directly and stopped after checking two tensor shapes.", + "The retained C0/C128 attention forward-backward policy now constructs and consumes the environment-sized cache, while the C128 context-parallel compressor forward requires the config-sized fallback to cover its rank-one slice." + ] + }, + { + "id": "TA-824", + "scope": "packed DeepSeek checkpoint handler type and private quantization-field report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The former report ended after checking the selected handler type and its private group-size and bit-width fields.", + "The retained packed-loader transaction now obtains the handler through the model, parses the official nested config, and proves 8-bit group-64 packed expert tensors dequantize to the expected gate, up, and down values." + ] + }, + { + "id": "TA-825", + "scope": "DSv4 compressor and indexer constructor-presence report", + "decision": "remove", + "status": "applied", + "evidence": [ + "C0 and C128 component selection is already exercised by the retained attention forward-backward transaction, including the C128 compressor path.", + "The retained C4 synthetic checkpoint load constructs both compressor and indexer and validates their separately translated APE tensors, so an incorrect topology cannot pass that lifecycle." + ] + }, + { + "id": "TA-826", + "scope": "external FLA Hopper autotuner regression report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report imported fla.ops.gated_delta_rule directly and never called any XoRL module, wrapper, or integration seam.", + "Its pass/fail result was controlled entirely by an optional dependency's production-shape backward kernel, so an upstream autotuner change was not an XoRL repository regression." + ] + }, + { + "id": "TA-827", + "scope": "standalone batch-invariant full-mean dtype spelling report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The dtype keyword reaches the same full-reduction override already owned by the mean-versus-sum regression contract.", + "That retained contract now checks the BF16 input to FP32 output spelling alongside one- and two-dimensional full reductions in both supported dtypes." + ] + }, + { + "id": "TA-828", + "scope": "production-size head-v2 projection and batch-invariance repetitions", + "decision": "remove", + "status": "applied", + "evidence": [ + "The 1024-hidden and 20480-vocabulary reports selected no additional head-v2 launch or merge branch beyond the retained focused geometry.", + "The retained head-v2 policies prove exact v1 projection bits, shared decode/scoring statistics, selected-logprob composition, arbitrary-slice batch invariance, fused-loss gradients, and rollback behavior." + ] + }, + { + "id": "TA-829", + "scope": "orphan eager OPRD layer-cache gather report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The private _get_opd_teacher_layer_hidden_states wrapper has no production caller; its only call site was the removed test.", + "The live OPRD path uses _get_opd_teacher_layer_fetcher, whose retained policy verifies selected cache rows, multiple layer slices, layer count, and returned shapes." + ] + }, + { + "id": "TA-830", + "scope": "legacy IS metric accumulator CPU and distributed reports", + "decision": "remove", + "status": "applied", + "evidence": [ + "ModelRunner._accumulate_is_metrics and _finalize_is_metrics have no production callers; runtime forward-backward now uses _accumulate_loss_metrics and _finalize_loss_metrics.", + "The retained OPD runner policy covers current mean, extrema, empty-rank, and loss-specific accumulation, while the retained two-rank report exercises the still-live _sp_allreduce_kl_metrics collective." + ] + }, + { + "id": "TA-831", + "scope": "EP kernel forward-only routing-score repetition", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report loaded the Triton and Quack EP kernel classes through test stubs and compared only their forward outputs with a local torch reference.", + "The retained EP routing-score report exercises the same two kernel classes and reference computation while checking both forward values and routing-score gradients." + ] + }, + { + "id": "TA-832", + "scope": "direct QARLLinear gradient and state-dict smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "The direct wrapper smoke repeated gradient presence, dynamic scale updates, and state restoration without exercising model injection or an optimizer step.", + "The retained QARL training lifecycle performs injection, a real AdamW update, changed-logprob checks, model checkpoint restoration, and exact restored logprobs; the calibration lifecycle independently checks persistent activation and block-scale state." + ] + }, + { + "id": "TA-833", + "scope": "private DSv4 FWHT helper reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The former reports separately tested a private helper's algebra, an impossible non-power-of-two DSv4 width, and fallback dispatch with only shape and finiteness assertions.", + "One retained public rotate_activation transaction now disables the optional kernel and proves the known transform, self-inverse behavior, and norm preservation across supported widths." + ] + }, + { + "id": "TA-834", + "scope": "launcher worker command with contradictory missing server arguments", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report bypassed Launcher construction, paired a valid flat config path with server_args=None, and asserted only that model_path was absent from a mocked subprocess command.", + "Normal construction either resolves ServerArguments from that config or raises before worker launch; retained launcher policies exercise live address selection, readiness failure, and override admission." + ] + }, + { + "id": "TA-835", + "scope": "DSv4 routing replay corrupted-global-state report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report directly overwrote the private module global with a fabricated future stage value that no production caller emits.", + "Retained DSv4 and generic routing-replay lifecycles exercise every supported stage through record, forward replay, backward replay, checkpoint recomputation, and R3 preload paths." + ] + }, + { + "id": "TA-836", + "scope": "activation-offload None-limit standalone backward smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "The standalone report passed None directly to a helper even though both trainer and server argument surfaces define activation_gpu_limit as a float, then asserted only that an unrelated square operation produced a gradient.", + "It did not observe offload placement, byte accounting, prefetch, or a production training path, so it could not discriminate activation-offload behavior." + ] + }, + { + "id": "TA-837", + "scope": "unreferenced FileLockLoader support suite", + "decision": "remove", + "status": "applied", + "evidence": [ + "FileLockLoader is not exported from xorl.data.prepare and has no caller or import anywhere else under src; its class name appeared only in its three-test module.", + "The retained data-preparation tests cover the live packing cache, dataset loading, hashing, retries, and preprocessing paths without preserving an unused counter-file abstraction." + ] + }, + { + "id": "TA-838", + "scope": "orphan runner-protocol JSON compatibility assertions", + "decision": "remove", + "status": "applied", + "evidence": [ + "Production runner communication exclusively uses serialize_message and deserialize_message; BaseMessage.to_json and BaseMessage.from_json have no runtime caller.", + "The retained protocol transaction still round-trips every live message payload through the actual transport codec, rejects pickle bytes, preserves tensors, and creates request acknowledgements." + ] + }, + { + "id": "TA-839", + "scope": "FutureStore test-only convenience accessor assertions", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "get_status, get_result, get_error, list_by_model, get_queue_state, and set_queue_state have no production caller; only the test used those convenience surfaces.", + "The retained lifecycle now observes live FutureEntry state through FutureStore.get while continuing to prove scheduling, concurrency, result and failure storage, deletion, model cleanup, and TTL expiry." + ] + }, + { + "id": "TA-840", + "scope": "standalone FSDP reduce-op canonicalizer truth table", + "decision": "remove", + "status": "applied", + "evidence": [ + "The four-line CPU report called the private canonicalizer directly with hand-built wrappers and raw enum values.", + "The retained two-rank FSDP2 lifecycle installs BF16StochasticAllToAllReduceScatter through set_custom_reduce_scatter, so PyTorch supplies the real wrapped reduce operation before finite gradient and numerical-error checks." + ] + }, + { + "id": "TA-841", + "scope": "fake full-precision expert FSDP kwargs report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The standalone report defined a fake object with one boolean class field and asserted that one dictionary key was removed.", + "The retained topology policy now makes the same assertion on an actual Glm52ExactTP16SharedExpertBlockFP8QLoRA after proving it is selected as one topmost FSDP unit." + ] + }, + { + "id": "TA-842", + "scope": "unintegrated GDN decode-prep P5 certification suite", + "decision": "remove", + "status": "applied", + "evidence": [ + "The four SM90 reports certified an opt-in gdn_decode_prep candidate that is not exported, documented, or called by any production module.", + "Its only in-repository source dependency is another unused decode-solve candidate, so graph capture and bitwise campaign gates did not protect a reachable XoRL execution path." + ] + }, + { + "id": "TA-843", + "scope": "test-only analytic pipeline bubble formula", + "decision": "remove", + "status": "applied", + "evidence": [ + "analytic_bubble_fraction was exported only by pp_profiling itself and had no production, documentation, example, or script caller; its only consumer was a hand-written formula truth table.", + "The retained PPBubbleProfiler path is constructed by Trainer and measures actual schedule busy intervals, memory, and P2P estimates instead of predicting an idealized formula." + ] + }, + { + "id": "TA-844", + "scope": "orphan first-fit-decreasing feasibility checker assertions", + "decision": "remove", + "status": "applied", + "evidence": [ + "ffd_check had no source caller and was imported only by the packing test, where six literals exercised an algorithm that dataset preparation never invokes.", + "The retained report exercises pack_group, allocate_sequentially, and PackingDataset, which are the actual packing and rank-allocation paths." + ] + }, + { + "id": "TA-845", + "scope": "unused grouped and any-FQN matcher truth tables", + "decision": "remove", + "status": "applied", + "evidence": [ + "check_all_fqn_match and check_any_fqn_match had no runtime caller; only the distributed-utils test preserved their grouped-number, prefix, and return-index behavior.", + "The retained check_fqn_match report protects the matcher actually used throughout ParallelPlan and sharded adapter state resolution." + ] + }, + { + "id": "TA-846", + "scope": "standalone per-parameter EP gradient hook compatibility path", + "decision": "remove", + "status": "applied", + "evidence": [ + "The hook installer and single-gradient reducer were documented for standalone tests and diagnostics and had no production caller; one branch of the real multi-rank report was their only consumer.", + "The retained two-rank and three-rank lifecycles exercise the production coalesced optimizer-boundary reducer, participation masks, bucket accounting, clipping, and non-finite rejection; the test-only parameter_count alias was removed with the hook." + ] + }, + { + "id": "TA-847", + "scope": "unreachable DeepSeek-V4 indexer autograd and backward campaign", + "decision": "remove", + "status": "applied", + "evidence": [ + "The V4IndexerFunction wrapper and batched_indexer_bwd kernel were imported only by two parameterized test reports; no production, documentation, example, or script path referenced either module.", + "The live V4Indexer invokes batched_indexer_fwd directly and returns discrete top-k indices, so the retained forward-score, causal-mask, numerical-range, and zero-input gates cover the reachable kernel while eight dormant wrapper/backward cases are gone." + ] + }, + { + "id": "TA-848", + "scope": "unintegrated manual CUDA timing module and lifecycle report", + "decision": "remove", + "status": "applied", + "evidence": [ + "No production module, package export, documentation, example, or script imported xorl.utils.manual_cuda_timing; its setter, scope, and drain functions were consumed only by its own mocked-event test.", + "The report therefore certified an instrumentation lifecycle that no XoRL execution path could enable or observe." + ] + }, + { + "id": "TA-849", + "scope": "standalone repeat_kv shape, value, and device smoke", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The standalone helper report separately checked identity, repeated shapes, a tiny value pattern, and CUDA device preservation without entering attention.", + "The retained eager-attention transaction now compares GQA weights and outputs against an independent torch.repeat_interleave reference, proving the repeated KV values are consumed correctly through the live backend while retaining invalid-head-layout rejection." + ] + }, + { + "id": "TA-850", + "scope": "identity-expert synthetic full MoE pipeline report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report composed histogram, index, scatter, and gather around an identity expert and ended at the same hidden_states times topk invariant already exercised by the retained scatter-gather round trip.", + "Retained kernel reports independently compare histogram/index ordering and gather, scatter, and add-gather values, while real MoE model and backend suites cover non-identity expert computation and gradients." + ] + }, + { + "id": "TA-851", + "scope": "test-only pipeline single-stage schedule convenience predicate", + "decision": "remove", + "status": "applied", + "evidence": [ + "is_single_stage_schedule had no runtime caller and its only assertions repeated whether schedule_stage_style returned single for GPipe and 1F1B.", + "The retained schedule policy still validates every supported style, split-backward mode, virtual-stage constraint, and the real schedule-class decision inside build_pipeline_schedule." + ] + }, + { + "id": "TA-852", + "scope": "unconsumed GLM sparse-cache mapping and gather helpers", + "decision": "remove", + "status": "applied", + "evidence": [ + "physical_cache_to_logical_indices and gather_selected_logical_values had no source, documentation, example, or script caller; their only consumer was an appended assertion block in the sparse-selector test.", + "The retained selector contract exercises the live canonical logical-index producer, including ties, short rows, dead rows, sorted unique indices, valid counts, and the production boundary tail." + ] + }, + { + "id": "TA-853", + "scope": "test-only GLM inventory and layer-plan convenience projections", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Glm52AdapterInventory.role_counts and four Glm52LayerPlan filtered-layer properties had no runtime caller and merely precomputed views of their public target and layer tuples for tests.", + "The retained contracts derive those counts and schedules from the authoritative tuples while preserving exact official roles, full/shared indexer allocation, dense/sparse placement, and producer boundaries." + ] + }, + { + "id": "TA-854", + "scope": "test-only NVFP4 format-dispatch wrapper", + "decision": "remove", + "status": "applied", + "evidence": [ + "fake_quantize was neither package-exported nor called by source; tests alone checked that its only supported string forwarded to fake_quantize_nvfp4 and that an unrelated string raised.", + "The retained public fake_quantize_nvfp4 contract compares both supported dtypes with an independent reference, proves straight-through gradients, and rejects invalid tensor geometry." + ] + }, + { + "id": "TA-855", + "scope": "standalone native-FP8 metadata-dictionary preflight", + "decision": "remove", + "status": "applied", + "evidence": [ + "validate_native_fp8_state_metadata had no loader caller; its test constructed an artificial metadata dictionary solely to call the helper directly.", + "DistributedCheckpointer uses the retained validate_native_fp8_dcp_checkpoint path, whose real-DCP tests reject castable payload metadata before load and validate EP-restored expected shapes." + ] + }, + { + "id": "TA-856", + "scope": "unintegrated lightweight ModelState reference export path", + "decision": "remove", + "status": "applied", + "evidence": [ + "ModelState.reference_state_dict described a future direct-safetensors path but had no production, export-script, documentation, or example caller; only test fixtures invoked it.", + "The retained checkpoint policies exercise actual DCP state collection, persistent QARL buffer metadata, compatibility rejection, pipeline key unions, optimizer filtering, and save/load process groups." + ] + }, + { + "id": "TA-857", + "scope": "routing-replay cursor reset API used only by tests", + "decision": "remove", + "status": "applied", + "evidence": [ + "reset_forward, reset_backward, reset_all_forward, and reset_all_backward had no runtime caller; tests alone rewound synthetic cursor values and then expected a checkpoint replay cursor to return to zero.", + "The real trainer and R3 handler clear replay instances at transaction teardown, so retained tests now assert that backward replay advances the cursor and that clear_all performs the production cleanup." + ] + }, + { + "id": "TA-858", + "scope": "manual LoRA merged-weight cache invalidators with no caller", + "decision": "remove", + "status": "applied", + "evidence": [ + "The dense, fused-delta, and MoE invalidate_merged_weight_cache methods had no production caller; one test invoked the dense method only to isolate two executions.", + "Merged caches are keyed by tensor versions, storage pointers, active rank, and alpha, and retained tests prove automatic optimizer-step and runtime-configuration invalidation plus bounded generation release." + ] + }, + { + "id": "TA-859", + "scope": "undocumented router top-k diagnostic policy matrix", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_MOE_ROUTER_TOPK_POLICY and its stable-low-id, tie-bias, and raw-logit branches had no documentation, script, example, or production configuration consumer; only a synthetic test truth table selected them.", + "The retained router contracts use the live torch.topk selection and still cover softmax weighting, normalization, DSv4 correction bias, hash routing, balanced synthetic profiling, and exact batch-invariant routing." + ] + }, + { + "id": "TA-860", + "scope": "layer-list router-FP32 environment diagnostic and standalone report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "XORL_MOE_ROUTER_FP32_LAYERS was absent from runtime configuration, documentation, scripts, and examples; its parser and layer-index test were their own consumers.", + "The retained configuration contract now exercises the real _router_fp32 model setting through MoEBlock and proves that hidden states and gate weights enter the FP32 projection without preserving the parallel environment override." + ] + }, + { + "id": "TA-861", + "scope": "unqualified FP64 MoE parity detour and dispatch assertion", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_MOE_FP64_ACCUM had no configuration, documentation, launcher, example, or serving-side implementation in the repository; its only test replaced the FP64 method with a mock and asserted dispatch precedence rather than numerical parity.", + "The retained MoE reports exercise the shipped eager, Triton, fused-SGLang, EP, and TP-simulation paths with real forward, gradient, layout, determinism, and admission checks." + ] + }, + { + "id": "TA-862", + "scope": "Qwen3-MoE delayed-residual and partial-residual experiment family", + "decision": "remove", + "status": "applied", + "evidence": [ + "The delayed residual pair, TP-shard carry, post-attention partial residual modes, alternate RMSNorm force flags, and candidate-capture matrix were controlled only by undocumented process-wide environment variables; no repository configuration, launcher, example, or end-to-end model test enabled them.", + "Four reports constructed private tuple inputs or attached tensor attributes in test doubles. The retained Qwen3-MoE report executes the normal decoder and final-norm consumers and verifies the explicit no-residual versus residual-tree family contract." + ] + }, + { + "id": "TA-863", + "scope": "unused stacked-LoRA initialization and merge utility surface", + "decision": "remove", + "status": "applied", + "evidence": [ + "The stacked initialization, delta, merge, and unmerge helpers had no production, documentation, example, or script caller; their package exports and one arithmetic truth-table report were their only consumers.", + "The live compute_lora_scaling function now remains at the group-GEMM package boundary used by dense, MoE, and quantized adapters, whose retained construction, loading, gradient, and optimizer reports all exercise it." + ] + }, + { + "id": "TA-864", + "scope": "tautological eager SwiGLU parity wrapper and report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The private _native_silu_and_mul helper had no source caller and was exactly one torch.nn.functional.silu expression; its sole report compared that wrapper with the identical expression labeled as the SGLang reference.", + "Real fused SwiGLU forward and backward behavior remains covered through the Triton operator, exact GLM MLP composition, and model-level LoRA and QLoRA paths; the independent dense RoPE cross-engine report remains in the file." + ] + }, + { + "id": "TA-865", + "scope": "unwired families-v2 specialized QK-norm kernel and golden", + "decision": "remove", + "status": "applied", + "evidence": [ + "qk_norm_v2 and its Triton kernel had no model, trainer, dispatcher, configuration, documentation, example, or script caller; only a standalone strided-view report and one frozen golden invocation reached them.", + "The Qwen3.5 and Qwen3.5-MoE attention paths use their declared RMSNorm modules. Retained families-v2 gates continue to cover the live hidden-state RMSNorm fused and split realizations, dispatch boundary, exact-model selection, and frozen numerical trees." + ] + }, + { + "id": "TA-866", + "scope": "undocumented SGLang MoE TP simulation diagnostic matrix", + "decision": "remove", + "status": "applied", + "evidence": [ + "The XORL_SGLANG_MOE_TP_SIM environment family and its direct, cache, Triton, alternate-reduce, DeepGEMM, fused-kernel, and runner modes had no configuration, launcher, example, documentation, or non-diagnostic consumer; history identifies the lane as a K3 parity diagnostic and later experiment.", + "Nine reports constructed tiny full-local tensors and replaced each optional backend with a Python fake. The retained SGLang fused-expert suites exercise the supported local and EP dispatch, weight layouts, loader admission, autograd, cache invalidation, and failure boundaries without preserving the parallel simulation product." + ] + }, + { + "id": "TA-867", + "scope": "CPU-only SGLang JIT and kernel RMSNorm diagnostic report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report selected sglang_jit and sglang_kernel but ran only CPU tensors, so neither optional CUDA implementation, loader, ABI boundary, nor serving arithmetic was exercised; both modes reduced to an ordinary eager fallback formula.", + "The modes remain available for explicit diagnostics. Retained RMSNorm suites cover CPU fallback arithmetic, fused forward and backward, exact Qwen site integration, family admission, global configuration forwarding, and real CUDA kernels when available." + ] + }, + { + "id": "TA-868", + "scope": "unintegrated sparse-delta receiver sharding and translation-future helpers", + "decision": "remove", + "status": "applied", + "evidence": [ + "The contiguous sharder, per-rank raw and encoded writers, translation-future collector, and terminal future writer formed a closed source/test-only subgraph with no runtime, configuration, documentation, example, or script caller.", + "Three reports fabricated receiver shards and the complete optional delta_encoding future API. The retained source-capture lifecycle exercises the live ModelRunner and dispatcher boundary, manifest aggregation, validated single-file packing, and sparse-delta backend consumption." + ] + }, + { + "id": "TA-869", + "scope": "deprecated R3 payload configuration and RequestProcessor aliases", + "decision": "remove", + "status": "applied", + "evidence": [ + "externalize_r3_payloads, keep_r3_payloads, routing_payload_dir, and keep_routing_payloads appeared only in compatibility branches and two test inputs; all launchers and runtime constructors already use r3_payload_transport, r3_payload_dir, and r3_payload_keep.", + "The canonical Mooncake and filesystem tests remain and still cover payload creation, slicing, cleanup, retention, namespace validation, and server configuration serialization without preserving duplicate names." + ] + }, + { + "id": "TA-870", + "scope": "fully mocked Trainer bootstrap forwarding report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report instantiated Trainer through __new__, supplied two large SimpleNamespace configurations, and replaced every bootstrap dependency to assert one direct ep_intranode keyword assignment and the presence or absence of one loss dictionary key.", + "Retained distributed tests execute both EP mesh layouts, loss tests exercise quack_linear numerics and admission, argument tests validate the settings, and the trainer model-boundary report covers meaningful configuration forwarding without reproducing bootstrap internals." + ] + }, + { + "id": "TA-871", + "scope": "undocumented Muon Quack tuning environment override", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_MUON_QUACK_TUNED had no configuration, documentation, example, or script consumer; its only report replaced Quack GEMMs with lambdas and observed the tuned keyword.", + "Muon retains the qualified tuned=False default used by trainer Quack paths. Retained reports cover backend import failure, architecture and dtype selection, real optimizer updates, grouped Gram Newton-Schulz execution, and CUDA compute dtype." + ] + }, + { + "id": "TA-872", + "scope": "fake-arithmetic CollatePipeline suite and test-only constructor forms", + "decision": "remove", + "status": "applied", + "evidence": [ + "Two reports composed fake collators that added and multiplied token IDs, then asserted those fake operations plus single-callable, tuple, and empty-list constructor forms absent from runtime callers.", + "The retained DataLoaderBuilder integration exercises the live non-empty list pipeline through tensor conversion, flattening, shifting, packing, micro-batch splitting, and optional sequence sharding." + ] + }, + { + "id": "TA-873", + "scope": "orphan data-preparation support implementations after test removal", + "decision": "remove", + "status": "applied", + "evidence": [ + "FileLockLoader remained unexported and unreferenced after its test-only suite was removed, while the SHA256 string wrapper had no caller after its standard-library restatement report was removed.", + "Live dataset preparation retains packing-cache persistence, source loading, split and configuration fingerprints, retries, preprocessing, and dataloader lifecycles." + ] + }, + { + "id": "TA-874", + "scope": "test-only retry strategies and wall-clock timing assertions", + "decision": "remove", + "status": "applied", + "evidence": [ + "Linear and constant retry strategies were selected only by the unit test; the sole production decorator use relies on the exponential default.", + "The retained retry contract covers success, retryable request and Hub failures, exhaustion, unrelated exceptions, and exponential delays by observing requested sleeps without real-time thresholds." + ] + }, + { + "id": "TA-875", + "scope": "fully mocked Trainer numerical-flag forwarding report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report constructed Trainer via __new__, supplied a large fake argument tree, replaced foundation-model construction, and asserted direct keyword copies for numerical flags and LoRA scalars.", + "Retained argument and model-policy suites cover parsing, resolution, admission, construction, and numerical behavior without reproducing the Trainer call signature." + ] + }, + { + "id": "TA-876", + "scope": "test-only orchestrator response builders and validation conveniences", + "decision": "remove", + "status": "applied", + "evidence": [ + "Twelve response builders and four validation or introspection helpers had no orchestrator, API server, scheduler, dispatcher, documentation example, or runtime caller; only their protocol test and package re-exports referenced them.", + "The retained protocol contract round-trips the live OrchestratorRequest and OrchestratorOutputs dataclasses through msgpack and reconstructs the typed operation payload consumed by ZMQ communication." + ] + }, + { + "id": "TA-877", + "scope": "redundant dense NVFP4 wrapper-injection smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report wrapped two Linear modules in a Sequential and asserted the wrapper count, format string, and output shape without checking a distinct numerical or lifecycle boundary.", + "Retained QARL reports cover targeted dense injection and exclusions, counters and summaries, NVFP4 arithmetic and straight-through gradients, a real optimizer update, changed log-probabilities, and checkpoint restoration." + ] + }, + { + "id": "TA-878", + "scope": "mock-only weight-version forwarding report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report bound one handler method to a MagicMock and replaced the NCCL synchronizer, so it observed only direct keyword forwarding through mocked collaborators.", + "The retained handler policy contract verifies flush_cache and weight_version at transfer_bucket, and the P2P protocol contract verifies the requested version in the real completion request body." + ] + }, + { + "id": "TA-879", + "scope": "mocked sharded-LM-head builder keyword-copy report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report replaced foundation-model construction and parallelization, then asserted only that fsdp_sharded_lm_head_loss=True crossed the adjacent function call.", + "Retained distributed and training-utility suites execute the sharded LM-head loss under FSDP and cover its admission, chunking, normalization, and gradients." + ] + }, + { + "id": "TA-880", + "scope": "mocked ModelRunner FP8, QARL, and sharded-loss builder keyword copies", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report replaced build_training_model and compared unchanged train-config values with the fake builder's kwargs; it constructed no FP8, QARL, calibrated, or sharded-loss model.", + "Retained builder and QARL lifecycle suites perform real dense and MoE injection, validate FP8 policy, calibrate before parallelization, update parameters, and restore checkpoints." + ] + }, + { + "id": "TA-881", + "scope": "Dr.GRPO outer-loop dispatch into a fake forward loop", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report replaced ModelRunner._forward_loop and asserted that the drgrpo string, input objects, and generic runner fields reached the fake, without executing Dr.GRPO computation.", + "The retained runner report executes the actual Dr.GRPO branch and its clipping, KL, temperature, legacy-field, output, and K3 policies; independent lifecycle suites cover completion, failure, routing, identity, and step accounting." + ] + }, + { + "id": "TA-882", + "scope": "raw routing-replay object forwarding into mocked backends", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report passed routed-expert lists through forward and forward-backward requests, replaced both backend methods with AsyncMock, and asserted object equality at the immediate kwargs seam.", + "Retained routing lifecycles exercise Mooncake and filesystem encoding, datum ordering, slice loading, cleanup, wire decoding, rank-zero selection, model identity, and failures." + ] + }, + { + "id": "TA-883", + "scope": "fake ModelRunner block-FP8 QLoRA builder report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report replaced build_training_model, asserted direct configuration copies, and constructed no quantized or adapter model; its remaining target-set assertion duplicated the dedicated GLM target-resolution contract.", + "Retained real builder suites cover foundation and injection settings, enabling preconditions, adapter inventory ownership, targets, quantized construction, and QLoRA execution." + ] + }, + { + "id": "TA-884", + "scope": "standalone module-path and FQN-matcher utility truth tables", + "decision": "remove", + "status": "applied", + "evidence": [ + "The two reports exercised recursive getattr and setattr plus regex wildcard examples on a toy Sequential and ModuleDict without reaching a sharding or ownership outcome.", + "Retained ParallelPlan and sharded adapter-state suites invoke the same helpers through exact and wildcard FQNs while slicing parameters, assigning placements and gradient domains, and materializing real adapter layouts." + ] + }, + { + "id": "TA-885", + "scope": "unused singular launcher free-port helper", + "decision": "remove", + "status": "applied", + "evidence": [ + "find_free_port had no source, test, documentation, example, or script caller; similarly named test-local helpers were independent definitions.", + "The live launcher retains find_free_ports, which allocates the three- and four-port rendezvous layouts consumed by worker startup." + ] + }, + { + "id": "TA-886", + "scope": "mock-only RemoteBackend payload wrapper suite", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report replaced RemoteBackend._execute and asserted operation names, request ids, timeouts, and fields copied into two payload constructors; it never serialized or transported a request.", + "Retained protocol, request-processor, dispatcher, sparse-delta, and weight-sync suites exercise typed payload reconstruction and the downstream behavior of the same fields." + ] + }, + { + "id": "TA-887", + "scope": "unintegrated standalone xorl.rl primitive package and self-reference suite", + "decision": "remove", + "status": "applied", + "evidence": [ + "The six exported tensor helpers had no trainer, runner, loss, CLI, example, documentation, or other source caller; their only consumers were five tests reproducing the same formulas.", + "Production loss suites retain the integrated policy, KL, importance-sampling, OPD, and Dr.GRPO numerical contracts used by training." + ] + }, + { + "id": "TA-888", + "scope": "test-only runner protocol acknowledgement and response factories", + "decision": "remove", + "status": "applied", + "evidence": [ + "create_ack_for_request and create_response_for_request were exported but had no runtime, documentation, example, or script caller; only the acknowledgement helper had a direct field-copy assertion.", + "Live transport code constructs RunnerAck and RunnerResponse directly, while the retained protocol report round-trips those actual dataclasses through the bounded MessagePack wire format." + ] + }, + { + "id": "TA-889", + "scope": "orphan sparse-delta translation-input loader and fake dependency report fragment", + "decision": "remove", + "status": "applied", + "evidence": [ + "load_sparse_source_delta_inputs had no translation engine, receiver, CLI, example, documentation, or source caller; its only consumer installed a synthetic delta_encoding package and checked fabricated empty shards.", + "Retained sparse-delta tests cover live source capture, manifests, packed-file validation, backend upload, receiver application, and end-to-end trainer-to-SGLang behavior." + ] + }, + { + "id": "TA-890", + "scope": "production-embedded canonical MoE test oracle and oracle self-test", + "decision": "remove", + "status": "applied", + "evidence": [ + "canonical_moe_reduce_reference had no runtime caller and shared _adjacent_pairwise_bf16 with the implementation it was used to validate; its standalone report then compared that shared helper with a hand-written tree.", + "The distributed and GLM model contracts now compute their expected adjacent BF16 tree independently in test code, and the real multi-process transport, permutation, chunking, output-distribution, and backward gate passes." + ] + }, + { + "id": "TA-891", + "scope": "unconsumed SequencePartial reducer and synthetic layout matrix", + "decision": "remove", + "status": "applied", + "evidence": [ + "SequencePartial had no loss, trainer, runner, CLI, example, documentation, or other source caller; only its dense, packed, and hand-sliced context-parallel unit report instantiated it.", + "TokenPartial remains the sole production reducer and retains integrated causal-LM, policy, importance-sampling, OPD, Dr.GRPO, TP, and FSDP coverage." + ] + }, + { + "id": "TA-892", + "scope": "duplicate Tinker weights-info and create-model compatibility report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The older report repeated session-spec persistence, normalized LoRA and optimizer metadata, and weights-info loading already covered by the focused session-endpoint lifecycle.", + "The retained report now also asserts Tinker's flat lora_rank response field while preserving disk-over-memory metadata, path confinement, and full-weight cases." + ] + }, + { + "id": "TA-893", + "scope": "fake-model sampler prefill-length forwarding report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report supplied one literal list to _compute_micro_batch_loss, captured kwargs on a fake model, and asserted only tensor dtype, device, and unchanged value.", + "Retained GLM indexer and sparse-attention contracts exercise prefill-boundary validation and behavior; the Dr.GRPO runner report remains focused on executing the actual loss branch." + ] + }, + { + "id": "TA-894", + "scope": "private trainer telemetry formatting reports", + "decision": "remove", + "status": "applied", + "evidence": [ + "The two reports instantiated no trainer lifecycle; they asserted byte-to-GB field naming, private key ordering, duplicated local summary floats, call counts, and empty dictionaries on synthetic namespaces.", + "Retained component-timer coverage runs real GLM- and Qwen-shaped forward and backward hooks on CUDA and preserves unrecorded-event recovery, while trainer suites cover optimization and synchronization behavior." + ] + }, + { + "id": "TA-895", + "scope": "split fake seams for runner forward session and R3 propagation", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "One report replaced _execute_and_gather and the other called _execute_compute directly, so neither covered the handler-to-trainer chain.", + "The replacement executes the real rank-zero handler, gather wrapper, and compute dispatch through the trainer while stubbing only distributed side effects." + ] + }, + { + "id": "TA-896", + "scope": "families-v2 source-text import lint", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report opened bi_families_v2.__file__ and banned four import substrings without executing a numerical or cross-engine behavior.", + "Retained suites cover model-program selection, rollback, dispatch, fused-versus-split equality, cross-engine bytes, and CUDA bit gates; the audit now identifies module-source reads as source inspection." + ] + }, + { + "id": "TA-897", + "scope": "repository-wide private-reference pytest scan", + "decision": "relocate", + "status": "applied", + "evidence": [ + "The report exercised no XoRL behavior; it listed every tracked file with Git, decoded source and documentation, and matched repository-policy regexes.", + "The same zero-dependency scan now runs as scripts/check_public_tree.py from the pre-commit lint workflow, preserving pull-request enforcement outside pytest." + ] + }, + { + "id": "TA-898", + "scope": "duplicate fake register-session dispatcher report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report called _handle_register_session on a fake coordinator and repeated payload and response assertions from the dedicated runner session-ops suite.", + "The retained session-ops report also covers cross-rank rejection, while the retained request-processor lifecycle separately reaches DummyBackend registration." + ] + }, + { + "id": "TA-899", + "scope": "self-fulfilling orchestrator-client forward and optimizer socket report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report's mock engine computed num_samples from its own received list and echoed the request learning rate, so both asserted results were authored by the test double rather than XoRL behavior.", + "Retained real-ZMQ reports verify interleaved forward and optimizer requests plus exact serialization, while the orchestrator end-to-end error report rejects empty model-pass batches." + ] + }, + { + "id": "TA-900", + "scope": "duplicate request-processor empty-batch rejection fragment", + "decision": "remove", + "status": "applied", + "evidence": [ + "The processor report repeated empty-list rejection already exercised through the real orchestrator request lifecycle.", + "Its distinct nonempty batch without valid targets remains as a focused processor validation report." + ] + }, + { + "id": "TA-901", + "scope": "duplicate API registration and explicit optimizer forwarding paths", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The legacy create-model report repeated focused full-weight and normalized LoRA worker registration; the retained full-weight report now also supplies empty optional configs to preserve that compatibility boundary.", + "The optimizer bundle's explicit learning-rate forwarding repeated the focused training-ops report, while its legacy Adam payload and learning-rate fallback priority remain covered." + ] + }, + { + "id": "TA-902", + "scope": "microscopic exact GLM and native FP8 contract wrappers", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Twelve collected items existed only to call one to three adjacent assertion helpers for the same CPU component, fragmenting topology, operand, byte, gradient, construction, admission, and checkpoint facets into narrow reports.", + "Component-level contracts now execute every original helper, with monkeypatch state explicitly reset between independent seams; separate Hopper and behaviorally distinct router or expert reports remain separate." + ] + }, + { + "id": "TA-903", + "scope": "DeepSeek-V4 private checkpoint-name and APE helper report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The report asserted private string rewrites and called the APE inverse on a synthetic forward transform even though the retained synthetic checkpoint transaction loads window, C4, hash, shared-expert, and routed-expert families through the real handler.", + "The retained transaction now checks values at every major destination, including norms, attention, HC, router bias, shared and fused experts, C4 APE tensors, and renamed indexer projections." + ] + }, + { + "id": "TA-904", + "scope": "fake GLM-5.2 block-FP8 QLoRA builder plumbing report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report replaced both foundation-model construction and QLoRA injection, then asserted that flags, rank, alpha, quant format, group size, and a fabricated inventory crossed those fakes.", + "Retained GLM-5.2 suites construct the full 700-target, 1700-factor model and exercise exact component admission; the builder's fail-closed requirement for LoRA plus QLoRA remains in the quantized-mode admission report." + ] + }, + { + "id": "TA-905", + "scope": "static OLMo-2 tensor-parallel plan dictionary report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report inspected TP_PLAN and MODEL_TP_PLAN entry types and missing keys without constructing a device mesh or applying tensor parallelism.", + "Retained two-rank CPU reports apply the production plan to OLMo-2, execute forward and backward through local-axis QK RMSNorm, rowwise and colwise projections, post-norm residual flow, and the vocab-sharded LM head, and compare the custom QK norm numerically." + ] + }, + { + "id": "TA-906", + "scope": "direct GLM sparse-MLA auto-dispatch self-comparison", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report compared auto dispatch on CPU with a direct call to the same torch reference implementation, duplicating the retained full-model sparse-versus-dense numerical path.", + "The full GLM attention integration report reaches auto dispatch through Glm5Model and checks dense parity; the distinct unknown-backend rejection was moved into that report rather than removed." + ] + }, + { + "id": "TA-907", + "scope": "standalone LM-head TP CP-DP-HSDP topology matrix", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report launched the same four-rank CP, DP, and HSDP layouts as the retained LM-head FSDP end-to-end suite, then asserted group membership, sizes, and mesh labels without executing a model or loss.", + "The retained end-to-end cases build a real FSDP-sharded LM head on those meshes and compare parameter synchronization, vocab ranges, global loss, full weight gradients, and local hidden gradients with eager references; the distinct EP-overlay topology report remains." + ] + }, + { + "id": "TA-908", + "scope": "private checkpoint URI construction and parsing snapshot", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The report called _to_xorl_uri and _from_xorl_uri directly for the same five documented path spellings already exercised by the adjacent save-and-load lifecycle, plus an undocumented arbitrary raw path.", + "The retained lifecycle creates real checkpoint directories and loads xorl URI, explicit weights path, model/checkpoint, checkpoint-only, and legacy weights/checkpoint inputs; its save assertion now pins the exact public xorl URI." + ] + }, + { + "id": "TA-909", + "scope": "direct FlashAttention metadata helper report", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The report called the metadata helper on synthetic dictionaries even though production reaches it through the packing and sequence-shard collators.", + "The retained packing-collator report now checks exact multi-sequence and single-sequence cumulative lengths and maximum lengths after real concatenation." + ] + }, + { + "id": "TA-910", + "scope": "direct sequence-shard slicing and padding primitive reports", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The reports called sp_slice and sp_padding directly and included a zero-length padding no-op.", + "The retained full-collator report now checks exact CP rank-zero and rank-one slices plus constant, sequential, label, and position padding on the real last-rank path." + ] + }, + { + "id": "TA-911", + "scope": "standalone API-orchestrator message roundtrip", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report serialized a request and response directly while the retained ZMQ client-engine tests send and receive those same message types over the production socket protocol.", + "The real communication lifecycle now pins the request payload, sequence id, timestamp, response id, type, payload, and terminal flag after crossing both sockets." + ] + }, + { + "id": "TA-912", + "scope": "fabricated DSv4 RoPE context-parallel slice guard", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The report invoked get_freqs_cis_for_cp with a fake rank and an arbitrary tensor rather than a model consumer.", + "The retained DSv4 compressor report now constructs a real short-cache CP compressor and reaches the same fail-loud capacity guard through forward_raw after also proving the supported C128 path." + ] + }, + { + "id": "TA-913", + "scope": "fragmented deferred QLoRA key-plan and cache-lifetime reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense merged, EP16 expert, and missing-pair cases are branches of the same prequantized module-key planner and now report as one policy contract.", + "Retained-cache clear and release behavior now runs inside the per-module deferred-loader residency lifecycle; every original assertion remains." + ] + }, + { + "id": "TA-914", + "scope": "standalone DeepSeek router rejection wrapper", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Foundation-model and training-builder router rejection are two entry points to the same DeepSeek router admission rule and now share one report with the admitted freeze behavior.", + "The tensor-parallel guard remains separate because it protects a different topology boundary." + ] + }, + { + "id": "TA-915", + "scope": "duplicate TeacherActivationCache selection report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ordinary rank-2 and rank-3 selection is already exercised through the Mooncake teacher-store consumer and the runner producer-consumer lifecycle.", + "Unique async completion, bounds, device residency, dtype reload, and rank-3 layer-slice checks remain in the cache lifecycle report." + ] + }, + { + "id": "TA-916", + "scope": "optimizer snapshot and immediate-load helper reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Recursive CPU snapshot semantics now run with the optimizer transaction failure policy that consumes those snapshots.", + "Immediate moment and step restoration now runs with the stronger uninterrupted-versus-resumed trajectory contract; all bytewise assertions remain." + ] + }, + { + "id": "TA-917", + "scope": "single-branch request-processor reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "A nonempty batch with no valid targets is the failure branch of the retained forward-backward processor lifecycle.", + "Register-session forwarding now runs with the processor control-operation lifecycle instead of presenting one backend roundtrip as a separate product behavior." + ] + }, + { + "id": "TA-918", + "scope": "standalone create-model normalized registration report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Normalized LoRA and optimizer registration is the successful branch of the same create-model lifecycle that handles recreation and registration rollback.", + "The create-session route remains separate because it has distinct refresh and override behavior." + ] + }, + { + "id": "TA-919", + "scope": "standalone folded-LoRA gradient-dtype helper report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FSDP gradient metadata controls the same folded-weight autograd boundary as the straight-through gradient comparisons.", + "The metadata, dtype, nonzero-gradient, shared-factor, fused gate-up, and linear assertions now execute in one autograd contract." + ] + }, + { + "id": "TA-920", + "scope": "separate QARL checkpoint buffer-key report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "QARL persistent-buffer mismatch detection and pipeline parameter or buffer key unions are facets of the same checkpoint model-key compatibility contract.", + "Strict and non-strict mismatch assertions remain alongside non-pipeline and pipeline metadata behavior." + ] + }, + { + "id": "TA-921", + "scope": "direct checkpoint expert-key classifier report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Expert-key classification exists to route dense and expert tensors in grouped checkpoint loading and is now checked inside that consumer lifecycle.", + "All supported expert, shared-expert, dense-MLP, and attention name cases remain asserted." + ] + }, + { + "id": "TA-922", + "scope": "standalone packing allocation primitive report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Bin capacity, safe-mode, offset, rank coverage, and non-overlap assertions describe the allocation used by PackingDataset.", + "They now run in the dataset lifecycle that also constructs sequential and multipack bins, loads cached bins, and checks cache identity." + ] + }, + { + "id": "TA-923", + "scope": "standalone empty TokenPartial reducer report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Zero scale with an empty mask is the boundary branch of the same denominator and additive-composition policy.", + "The exact zero assertion remains in the TokenPartial component report." + ] + }, + { + "id": "TA-924", + "scope": "standalone legacy TopK router wrapper", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy softmax selection, V4-input isolation, routed scaling, configuration selection, and FP32 gate routing are one router policy matrix.", + "Synthetic balanced, sqrt-softplus, and hash-routing behaviors remain separate because they execute distinct routing modes." + ] + }, + { + "id": "TA-925", + "scope": "scattered server runtime configuration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The one-helper R3 transport report is now part of the server runtime roundtrip contract.", + "Adapter-gradient bucket serialization and zero-value rejection moved from the gradient-math suite into that same user-facing configuration lifecycle." + ] + }, + { + "id": "TA-926", + "scope": "fragmented exact GLM attention and MoE construction rejection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dependency flags, EP16, lm-head TP16, rank-one alpha-one, sparse-MLA, and all-to-all requirements are branches of exact-component construction admission.", + "Every fail-before-mutation assertion remains in one attention and one complete-MoE admission report; successful inventory and post-EP ownership reports stay separate." + ] + }, + { + "id": "TA-927", + "scope": "standalone kill-session path validation report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Path-like model ID rejection is a security branch of the same kill-session checkpoint lifecycle.", + "The lifecycle still proves failure preservation, evicted-checkpoint promotion, metadata cleanup, and path rejection." + ] + }, + { + "id": "TA-928", + "scope": "standalone multi-part optimizer custom-group rejection report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Custom parameter-group rejection is the admission branch of multi-part optimizer construction.", + "It now runs with real multi-part updates, zeroing, scheduler propagation, model mapping, and single-part fallback." + ] + }, + { + "id": "TA-929", + "scope": "separate adapter ownership declaration rejection report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Overlapping and foreign pending authority are declaration-level branches of the compiler's fail-closed ownership policy.", + "They now run with missing-universe, unsupported-TP, managed-FSDP, and false-EP ownership rejection." + ] + }, + { + "id": "TA-930", + "scope": "split EP checkpoint mesh restore and drop reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Restoring the EP dimension and dropping it back to the stage-local expert-FSDP mesh are opposite directions of one checkpoint mesh policy.", + "Legacy and PP-parent mesh validation, placement, target shape, and local tensor assertions all remain." + ] + }, + { + "id": "TA-931", + "scope": "weight-sync protocol edge reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Empty PP-NCCL payload handling now runs with the named-tensor roundtrip instead of reporting a two-call mock branch separately.", + "Sparse-delta baseline priming and receiver-failure retry now run in the transfer state-machine lifecycle; post-packed transfer and initialization remain separate boundaries." + ] + }, + { + "id": "TA-932", + "scope": "fragmented packing and sequence-shard collator reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Generic extras, token-aligned fields, pre-shifted labels, and packed boundaries are field policies of their respective production collators.", + "All concatenation, padding, label, position, and FlashAttention metadata assertions remain; the CP16 token-side-channel contract stays separate." + ] + }, + { + "id": "TA-933", + "scope": "standalone R3 side-payload validation report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Malformed reference, bounds, and count failures are branches of the same R3 put, sliced-load, and cleanup lifecycle.", + "The low-level Mooncake tensor codec remains separate because it exercises serialization rather than R3 reference ownership." + ] + }, + { + "id": "TA-934", + "scope": "split sync-quantization normalization and rejection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Supported defaults and normalization plus unsupported methods, formats, schemes, list shapes, and explicit reasons form one configuration policy matrix.", + "Every accepted output and rejection message remains asserted through normalize_sync_quantization_config." + ] + }, + { + "id": "TA-935", + "scope": "NCCL rendezvous bind-failure branch report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Store-bind failure before inference initialization is the fail-closed branch of NCCL rendezvous initialization.", + "It now runs with ephemeral-port rotation and explicit-port pinning in one training rendezvous lifecycle." + ] + }, + { + "id": "TA-936", + "scope": "P2P async status-timeout branch report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Status polling timeout is the failure branch of asynchronous size and cutoff dispatch.", + "The prepare-request timeout remains separate because it governs a different HTTP phase." + ] + }, + { + "id": "TA-937", + "scope": "standalone SGL cross-attention cu-seqlens rejection report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Cross-attention cu-seqlens rejection belongs to the page-size-one SGL KV-cache adapter policy.", + "Generic fixed-length, variable-length, eager, backend-registry, and alternate FlashAttention paths remain separate contracts." + ] + }, + { + "id": "TA-938", + "scope": "split Rank0 ready-handshake branch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Normal acknowledgement, request-before-acknowledgement, client identity, unexpected message, and receive failure are outcomes of one Rank0Protocol ready handshake.", + "All wire-message, queue, acknowledgement, identity, and request-count assertions remain in the handshake lifecycle." + ] + }, + { + "id": "TA-939", + "scope": "split runner load-state preparation and routing reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Path confinement and preparation failure are admission phases of the same load-state operation that routes multi-adapter and single-tenant restores.", + "Error preservation, artifact-root enforcement, adapter conversion, trainer routing, and step reset remain asserted." + ] + }, + { + "id": "TA-940", + "scope": "split scheduler dispatch and terminal-state reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FIFO admission, capacity, dispatch, completion, failure, abort, statistics, clearing, and bounded history form one scheduler lifecycle.", + "A fresh scheduler isolates the terminal-transition matrix from the capacity scenario without creating a second product report." + ] + }, + { + "id": "TA-941", + "scope": "standalone multi-adapter Adam override report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Multi-adapter Adam beta and epsilon propagation is a branch of ModelRunner.optim_step's full, partial, omitted, and non-Adam override policy.", + "The adapter-manager call and adapter optimizer parameter-group assertions remain in the optimizer-step report." + ] + }, + { + "id": "TA-942", + "scope": "split ParallelState construction and singleton initialization reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Defaults, validation, enabled flags, initialization, automatic DP sharding, access, and reinitialization protection are one ParallelState lifecycle.", + "EP mesh construction and requires-mesh behavior remain separate because they exercise a different helper boundary." + ] + }, + { + "id": "TA-943", + "scope": "fragmented DeepEP internode-preflight outcome reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Skip, uninitialized, intranode, transport failure, identity roundtrip, and corruption are branches of preflight_internode_transport.", + "Topology detection and buffer-size admission remain separate contracts; every node diagnostic and corruption assertion is retained." + ] + }, + { + "id": "TA-944", + "scope": "standalone ParallelPlan rejection branch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Indivisible generic meta slicing now runs with the successful metadata-preserving slice policy.", + "Malformed exact-GLM already-local and force-shard singletons now run with the exact meta EP disposition contract; the materialized real-tensor shard remains separate." + ] + }, + { + "id": "TA-945", + "scope": "fragmented exact absorbed-KV-B checkpoint pair reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Arrival order, duplicate members, incomplete pairs, dtype mismatch, and shape mismatch are outcomes of one NativeBlockFP8PairBuffer transaction.", + "The exact-attention source inventory remains separate because it validates model construction rather than checkpoint pair state." + ] + }, + { + "id": "TA-946", + "scope": "split block-FP8 quantization, dequantization, and edge reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shapes, scales, admission, roundtrip error, determinism, storage, magnitude edges, signs, and dimensional consistency are one block-FP8 codec contract.", + "All CUDA assertions execute in the same component report under the existing GPU gate." + ] + }, + { + "id": "TA-947", + "scope": "split GKN block-FP8 quantize and dequantize reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "GKN output geometry, roundtrip accuracy, tail blocks, zero blocks, scale range, output dtype, contiguity, and rank admission form one two-dimensional codec policy.", + "The large-matrix path and dequantization failure assertions remain under the same CUDA report." + ] + }, + { + "id": "TA-948", + "scope": "split shared-prefix attention edge reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "A singleton member and a one-token prompt are degeneracies of the same shared-prefix attention layout.", + "The general dtype, head-size, GQA, forward, and backward matrix remains separate from this edge-layout report." + ] + }, + { + "id": "TA-949", + "scope": "standalone one-token shared-prefix repack report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The empty shared block for a one-token prompt is an edge branch of shared-prefix detection, repacking, and remapping.", + "Its exact decoded blocks and empty cross-attention indices remain asserted in the full repack lifecycle." + ] + }, + { + "id": "TA-950", + "scope": "split exact dense gate-up pair-buffer reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Out-of-order gate and up emission plus missing, duplicate, invalid-dtype, and non-finite members are outcomes of one exact dense checkpoint transaction.", + "Fused byte layout, scale order, model installation, base-loaded state, and every failure remain asserted." + ] + }, + { + "id": "TA-951", + "scope": "fragmented NVFP4 fake-quant reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Two-dimensional reference parity, STE behavior, rank and block admission form one base fake-quant contract.", + "Three-dimensional projection STE, expert-independent scaling, and fused gate-up per-half scaling form one expert fake-quant contract." + ] + }, + { + "id": "TA-952", + "scope": "standalone DSV4 attention TP rejection report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unsupported tensor-parallel construction is an admission branch of the DSV4 attention storage and backend-call contract.", + "Window-only and C128 forward-backward variants remain separately parameterized because they execute different attention structures." + ] + }, + { + "id": "TA-953", + "scope": "separate eager-versus-native MoE determinism report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Backend determinism and all-tokens-to-one-expert behavior are edge branches of the eager-versus-native forward and backward parity matrix.", + "All expert-count, hidden-size, top-k, batch, sequence, gradient, and edge assertions remain under the same CUDA gate." + ] + }, + { + "id": "TA-954", + "scope": "standalone non-gated MoE constructor rejection report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unsupported backend and gated-activation rejection are admission branches of the CPU eager non-gated expert contract.", + "GPU Triton and native parity remains separate because it crosses optional backend and device boundaries." + ] + }, + { + "id": "TA-955", + "scope": "split gradient-checkpoint configuration and runtime-gate reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default and nondefault method propagation configure the same GradientCheckpointingLayer gate exercised by training, enabled, and method combinations.", + "Both ordinary and MoE layer configuration plus the exact checkpoint-call truth table remain asserted." + ] + }, + { + "id": "TA-956", + "scope": "split Kimi wrapper conversion and local registry-loading reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Kimi wrapper mapping, official auxiliary defaults, registry resolution, and local text-config unwrapping are one DeepSeek-V3 configuration-loading lifecycle.", + "Every MLA, MoE, routing, RoPE, registry, and local-load assertion remains." + ] + }, + { + "id": "TA-957", + "scope": "split local Kimi tokenizer and fallback-loader policy reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The dedicated local TikToken path and generic tokenizer or processor fallback are branches of the public auto-loader policy.", + "Token IDs, text roundtrip, right padding, and absence of implicit remote-code trust remain asserted." + ] + }, + { + "id": "TA-958", + "scope": "split sqrt-softplus and softmax routing-regather reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sqrt-softplus scaling and dtype plus unchanged softmax gather and renormalization are modes of MoEBlock._regather_routing.", + "Both eager-router comparisons and the independent softmax formula remain in one mode matrix." + ] + }, + { + "id": "TA-959", + "scope": "standalone all-invalid FlashMLA autograd report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All-invalid rows are the zero-row branch of the same FlashMLA compaction, TileLang backward, and scatter transaction.", + "Compacted valid rows, zero-scattered invalid gradients, all-zero output, and backward bypass remain asserted; dispatch-envelope admission stays separate." + ] + }, + { + "id": "TA-960", + "scope": "fragmented causal-LM Z-loss CPU reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Positive coefficient, zero coefficient, and tensor-parallel rejection are outcomes of one causal-LM Z-loss policy.", + "Reference CE and Z-loss values, gradients, absent metrics at zero, and fail-before-collective behavior remain; compiled CUDA parity stays separate." + ] + }, + { + "id": "TA-961", + "scope": "fragmented streaming forward-KL reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense-reference gradients, chunk invariance, ignore masking, and low-memory parity are one streaming-kernel contract.", + "OPD backend parity and unsupported logprob clamping form one dispatch contract; the independent FP64 gradcheck remains separate." + ] + }, + { + "id": "TA-962", + "scope": "fragmented LM-head module selection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local and TP module selection, FP32 bypass, and logprob temperature are one per-token CE policy matrix.", + "Importance-sampling module use plus causal-LM FP32 bypass and TP hidden-gradient reduction form one outer loss-dispatch contract." + ] + }, + { + "id": "TA-963", + "scope": "standalone Dr.GRPO objective branch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Zero and empty boundaries, positive-advantage direction, and KL reference admission are branches of the forward, backward, and metric contract.", + "Logprob-temperature behavior and microbatch composition remain separate numerical contracts." + ] + }, + { + "id": "TA-964", + "scope": "fragmented fused selected-logprob reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Frozen-output input gradients and irregular tails are shape and autograd branches of fused selected-logprob forward-backward parity.", + "Per-token, causal-LM, quack-linear, and importance-sampling dispatch now form one integration report; production-vocabulary finiteness and no-full-logits memory remain separate regressions." + ] + }, + { + "id": "TA-965", + "scope": "split families-v2 RMSNorm property reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FP64 proximity, residual and zero-centered behavior, batch composition invariance, and run determinism are properties of one RMSNorm-v2 tree.", + "Forced fused-versus-split equivalence and production dispatch remain separate because they certify realization and selection boundaries." + ] + }, + { + "id": "TA-966", + "scope": "fragmented BI fused LM-head reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Eager parity, gradients, deterministic batch composition, and unsupported-mode guards form one BI fused loss contract.", + "Unit temperature identity and near-probability-one clamping form one kernel edge policy." + ] + }, + { + "id": "TA-967", + "scope": "split TileLang V4 indexer edge reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Causal masking, large-value stability, and zero-input output are edge branches of the same batched indexer kernel.", + "The parameterized forward matrix remains separate because it covers four execution geometries." + ] + }, + { + "id": "TA-968", + "scope": "standalone NF4 codebook roundtrip report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Codebook ordering and exact-value roundtrip are foundational branches of the flat NF4 codec contract.", + "The GKN codec remains separate because it uses a different packed layout and scale geometry." + ] + }, + { + "id": "TA-969", + "scope": "fragmented OPD full-vocabulary policy reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Diagnostics, loss clamping, reward weighting, stable metric keys, and unsupported top-k dispatch are branches of one full-vocabulary OPD policy.", + "Policy-gradient behavior, KL estimators, and compiled sampled-token logprobs remain separate because they exercise different objectives or callables." + ] + }, + { + "id": "TA-970", + "scope": "split GDN convolution forward and backward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forward parity, backward parity, variable-length and batch behavior, and repeat determinism form one causal-convolution numerical contract.", + "End-to-end GDN integration, optional SGLang parity, input admission, and contract-state lifecycle remain separate boundaries." + ] + }, + { + "id": "TA-971", + "scope": "fragmented exact TP1 QLoRA wrapper reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Static configuration and fail-before-import runtime admission are one wrapper policy; rounded forward values and surrogate gradients are one numerical transaction.", + "Factor-only VJP behavior and saved-master mutation rejection form one backward safety policy; packed-state dtype movement and literal CUDA parity remain separate." + ] + }, + { + "id": "TA-972", + "scope": "split FP8 linear numerical-mode reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Padding recipes, full residual correction, and activation-only correction are modes of the same FP8 matmul numerical contract.", + "Backend selection, profiling, CPU fallback, injection, and an optimizer train step remain separate implementation boundaries." + ] + }, + { + "id": "TA-973", + "scope": "fragmented TileLang sparse-MLA feature reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Attention-sink reference parity and observable effect form one sink policy; partial-invalid forward and backward behavior form one masking transaction.", + "Parameterized geometries and deterministic-versus-atomic backward remain independent kernel reports." + ] + }, + { + "id": "TA-974", + "scope": "fragmented exact routed-expert factor-buffer and empty-route reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Global and post-EP local factor banks are two inputs to the same physical sampler-buffer contract.", + "Zero-token bank gradients, all-sentinel zero gradients, and stride admission are edge branches of one routed-gradient policy." + ] + }, + { + "id": "TA-975", + "scope": "split canonical GLM52 MoE configuration and selection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Routing-replay rejection, internal transport resolution, exact canonical selection, and ordinary noncanonical selection form one canonical-MoE mode policy.", + "Sparse-selector arithmetic, codecs, loaders, and runtime dispatch remain separate production boundaries." + ] + }, + { + "id": "TA-976", + "scope": "fragmented exact shared-expert admission and state reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Construction restrictions and fail-before-kernel runtime checks form one component admission policy.", + "FP32 logical masters, dtype movement, canonical checkpoint sources, and immutable binding form one persistent-state policy; optional SGLang views remain separate from native base views." + ] + }, + { + "id": "TA-977", + "scope": "split topmost mixed-precision FSDP selection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Generic nested-module selection and the exact shared-expert specialization exercise the same topmost protected-unit selector.", + "Expert mixed precision, reduce dtype, sequence-parallel folding, and prefetch direction remain separate configuration policies." + ] + }, + { + "id": "TA-978", + "scope": "fragmented model-runner token-diagnostic reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Selection, empty and top-k boundaries, loss-logprob cross-checks, raw-weight references, and hidden summaries are output branches of one token-diagnostic callable.", + "KL token diagnostics, tensor dumps, component hooks, and trusted diagnostic inputs remain separate callables or lifecycle boundaries." + ] + }, + { + "id": "TA-979", + "scope": "split stochastic-rounding property reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shape and dtype, input admission, seeded repeatability, expectation, and neighboring-value bounds are properties of one BF16 stochastic-rounding operation.", + "Every original deterministic and statistical assertion remains in one CPU numerical contract." + ] + }, + { + "id": "TA-980", + "scope": "one-report-per-mode learning-rate scheduler tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Constant, linear, and cosine schedules are branches of the same scheduler builder and now form one mode matrix.", + "Invalid configuration remains separate because it is the builder admission boundary." + ] + }, + { + "id": "TA-981", + "scope": "split DistSignSGD hook and topology-admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local hook registration, FSDP-managed exclusion, and unsupported HSDP, CP, EP, and non-FSDP DTensor rejection are outcomes of configure_distsignsgd.", + "Reduce-scatter arithmetic, optimizer construction, and parameter updates remain separate production boundaries." + ] + }, + { + "id": "TA-982", + "scope": "standalone default-session kill and unload protection report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ordinary LoRA teardown, checkpoint URI return, re-registration, and default-session protection form one API session-termination lifecycle.", + "Model creation, weight metadata, and legacy session-spec loading remain separate endpoints or persistence boundaries." + ] + }, + { + "id": "TA-983", + "scope": "fragmented inference-endpoint registration and sync reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Port routing, explicit worker health, adapter routing, auto-sync, server topology, and FP8 KV-cache admission are branches of endpoint registration.", + "Weight-sync quantization admission and FP8 KV-cache invalidation now form one sync request policy; listing and receiver enrichment remain separate." + ] + }, + { + "id": "TA-984", + "scope": "one-report-per-layout P2P FP8 receiver tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fused and unfused attention, partial blocks, Qwen3.6 QKVZ and full attention, nonexpert namespaces, mixed passthrough entries, shared experts, and routed experts are receiver-layout branches of one transfer protocol.", + "All byte, scale, slice, dequantization, endpoint, and expert-coverage assertions remain in one FP8 receiver-layout matrix." + ] + }, + { + "id": "TA-985", + "scope": "fragmented P2P multi-sender initialization reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-zero scatter, nonzero-rank adoption, explicit sender groups, peer-failure propagation, and engine prewarm ordering are branches of one multi-sender initialization state machine.", + "Locator copy modes, dense sharding, and rank-filtered transfer remain separate data-partitioning policies." + ] + }, + { + "id": "TA-986", + "scope": "split P2P transfer manifest rejection and name compatibility reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Missing locators, incompatible shapes, invalid source ranks, canonical name aliases, orig_mod stripping, and language-model prefix fallback form one receiver-manifest resolution policy.", + "Replicated locator staging and transfer metadata remain separate transport behaviors." + ] + }, + { + "id": "TA-987", + "scope": "split P2P source-staging mode reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Small CPU pooling, GPU-direct persistent registration and chunking, and aligned mixed-dtype scratch views are staging modes of transfer_bucket.", + "Receiver-handle coalescing and failure diagnostics remain separate scheduling and observability boundaries." + ] + }, + { + "id": "TA-988", + "scope": "split P2P pending-transfer and destroy reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Pending failure propagation, skip-completion cleanup, failed-transfer draining, receiver completion, endpoint results, deregistration, and completion errors form one teardown lifecycle.", + "Preparation, transfer, and explicit complete_sync metadata remain separate protocol stages." + ] + }, + { + "id": "TA-989", + "scope": "one-wrapper-per-topology lm-head TP FSDP tests", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CP-replica, DP, no-CP DP, no-CP HSDP, and OPD variants invoke the same four-process launcher and embedded loss-gradient oracle.", + "All six topology and loss-mode programs remain executed with per-case failure identifiers in one distributed matrix." + ] + }, + { + "id": "TA-990", + "scope": "split empty and aborted adapter gradient-epoch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Empty-step rejection, idempotent abort, scratch reset, publication state, and poisoned or pending rejection form one pre-mutation epoch lifecycle.", + "Gradient capture and successful optimizer commit remain separate transactions." + ] + }, + { + "id": "TA-991", + "scope": "split authoritative adapter optimizer success reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Analytical clipping, AdamW parameters and moments, scratch reuse, global-step commit, and the single logical-norm collective are properties of one successful optim_step transaction.", + "Exact LM-head replicated optimizer coherence remains separate because it validates a different topology-specific helper." + ] + }, + { + "id": "TA-992", + "scope": "split authoritative adapter optimizer outcome reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Pre-mutation semantic rejection, partial optimizer failure, and collective failure are outcome branches of optim_step.", + "Recoverability, poison state, publication gates, parameter mutation or preservation, and restart guidance remain asserted for every branch." + ] + }, + { + "id": "TA-993", + "scope": "split adapter checkpoint path and target-manifest reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Trusted-root confinement, strict target-manifest persistence, validation, and mismatch rejection form one checkpoint save and validation policy.", + "Restore compatibility and checkpoint structure remain separate load boundaries." + ] + }, + { + "id": "TA-994", + "scope": "split authoritative adapter checkpoint restore-plan reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Lifecycle reset, plan fingerprint restoration, compatible replacement, direct topology mismatch rejection, and atomic nonmutation form one authoritative restore policy.", + "Coordinator materialization and general session compatibility remain separate orchestration and user-policy boundaries." + ] + }, + { + "id": "TA-995", + "scope": "fragmented adapter optimizer save, resume, and admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shard emission and manifest identity rejection form one save contract; uninterrupted, evicted public, LR override, and weights-only control paths form one resume lifecycle.", + "Legacy and incomplete artifact rejection now includes resident-state atomicity; parameter identity, transactional commit, and logical resharding remain separate boundaries." + ] + }, + { + "id": "TA-996", + "scope": "split FP8 grouped forward and weight-gradient reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Same-NK forward and same-MN weight-gradient kernels are the two arithmetic halves of one grouped FP8 training contract.", + "Block-loop and Triton references, empty groups, tail shapes, block sizes, precomputed sequence offsets, and scalar-Quack dispatch remain asserted." + ] + }, + { + "id": "TA-997", + "scope": "fragmented SGLang fused-expert EP activation reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "DeepEP and FP8 exclusion, missing runtime, flag-off stock dispatch, score dtype, empty-rank behavior, and compute guards are branches of one EP dispatch admission boundary.", + "Happy-path compute, slot combine, trainable autograd, and weight presentation remain separate numerical or ownership contracts." + ] + }, + { + "id": "TA-998", + "scope": "split model-runner expert-factor compiler reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Certified unquantized backends, registered-rank specialization, block-FP8 DeepEP, NF4, NVFP4, and generic quantized admission all compile the same expert-factor ownership plan.", + "Module versus fused producer families, quantization guards, metadata mismatch, shape drift, and uncertified parallelism remain asserted in one compiler matrix." + ] + }, + { + "id": "TA-999", + "scope": "fragmented Qwen-235B simulator calibration and built-in pack reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Markdown ingestion, leave-one-out evaluation, calibrated and extrapolated scenarios, OOM boundaries, topology what-if and auto sweeps form one Qwen-235B calibration workflow.", + "Built-in pack replay now includes consolidated cross-pack validation; portable ledgers, security admission, and kernel ranking remain separate simulator boundaries." + ] + }, + { + "id": "TA-1000", + "scope": "split grouped checkpoint load routing and fallback reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense and expert routing, fused and FFN source formats, local dense fallback, absent EP-group fallback, strict rejection, and persistent-buffer filtering are branches of grouped_load_weights.", + "State-dict resolution, object broadcast, DTensor copying, and strict post-processing remain separate callables." + ] + }, + { + "id": "TA-1001", + "scope": "fragmented FP8 weight-sync selection and CPU expert reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Projection inclusion and exclusion, receiver skip lists, broad selection, stacked quantization, and already-FP8 passthrough form one input and layout policy.", + "CPU expert projection, zero padding, deferred formatting, exclusion, reusable workspace staging, and workspace quantization form one expert CPU pipeline." + ] + }, + { + "id": "TA-1002", + "scope": "split weight-sync quantization admission and enrichment reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unsupported quantization rejection and FP8 BF16-island enrichment are branches of handle_sync_inference_weights.", + "Backend post-processing, adapter materialization, parameter extraction, and sparse-delta sync remain separate transactions." + ] + }, + { + "id": "TA-1003", + "scope": "fragmented request-processor R3 side-payload reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Mooncake externalization, datum-order preservation, and normal, exceptional, and default cleanup form one model-pass side-payload lifecycle.", + "NCCL sync, dispatcher forwarding, optimizer and checkpoint operations, and token unpacking remain separate request boundaries." + ] + }, + { + "id": "TA-1004", + "scope": "split GLM5 indexer construction and selection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Parameter geometry, FP32 projection, sentinel masking, padding detection, and sorted and blocked selection form one indexer component contract.", + "The TileLang fast path and sparse-attention model integration remain separate implementation boundaries." + ] + }, + { + "id": "TA-1005", + "scope": "fragmented GLM52 sparse-selector pipeline reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Logical selection, Hadamard transport, fused projection, sampler key preparation, portable codecs, runtime dispatch, and dependency loading form one sparse-selector pipeline.", + "Production-shape SGLang CUDA codec parity is now a separate optional report so its runtime skip cannot mask the portable pipeline assertions." + ] + }, + { + "id": "TA-1006", + "scope": "split Muon builder configuration and SGD fallback reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Threading the fallback choice through build_optimizer and proving its state-free SGD update are one builder configuration contract.", + "Direct Gram-Newton-Schulz arithmetic, backend dispatch, and grouped update geometry remain separate numerical boundaries." + ] + }, + { + "id": "TA-1007", + "scope": "split Muon fused-expert classification and Nemotron integration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fused gate-up detection, FSDP attribute-loss recovery, gated and non-gated family classification, and a real Nemotron-H optimizer step form one parameter-ownership policy.", + "Matrix grouping and Newton-Schulz implementation details remain separate optimizer arithmetic reports." + ] + }, + { + "id": "TA-1008", + "scope": "fragmented distributed-checkpointer metadata, load, and save reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Metadata admission and synchronous, no-dist, custom-group, and asynchronous load and save routing form one DistributedCheckpointer I/O policy.", + "Optimizer metadata-key filtering now runs with the optimizer-state contract; model-key and LoRA compatibility remain separate schema boundaries." + ] + }, + { + "id": "TA-1009", + "scope": "fragmented exact GLM52 and Qwen3.5 training-program admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical GLM52 numerical resolution and official geometry form one model program; exact Qwen3.5 numerical, MoE, topology, and model-scope admission form another.", + "The family-independent RoPE selector remains separate because it also covers ordinary non-GLM behavior." + ] + }, + { + "id": "TA-1010", + "scope": "split quantized-export CLI parsing and base-directory reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "YAML and CLI precedence, size parsing, module invocation, BF16 islands, configuration output, and shard indexing form one end-to-end export command contract.", + "Projection layouts, MoE layouts, QARL logprob preservation, and source admission remain separate export boundaries." + ] + }, + { + "id": "TA-1011", + "scope": "fragmented create-model endpoint lifecycle reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Session-spec normalization, recreation admission, registration rollback, reserved-checkpoint initialization, and full-weight admission are outcomes of create_model_endpoint.", + "The lower-level create-session endpoint, termination, and checkpoint metadata lookup remain separate API lifecycles." + ] + }, + { + "id": "TA-1012", + "scope": "split adapter load admission, failure, and rank-zero broadcast reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Direct load success, trusted-path admission, synchronized failure, pipeline rejection, and auto-registration rollback form one handle_load_adapter_state outcome policy.", + "Rank-zero routing, sharded restoration, session-spec mismatch, and transactional optimizer rejection form one broadcast-load mode; eviction and registration remain separate lifecycles." + ] + }, + { + "id": "TA-1013", + "scope": "split packing edge, validation, and unpacking reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Empty, missing, oversized, NumPy, valid, and malformed inputs form one packing admission and output-validation policy.", + "Per-token unpacking modes now run with the end-to-end pack, metadata, forward-output, and sample-boundary round trip." + ] + }, + { + "id": "TA-1014", + "scope": "split teacher-cache contributor and distributed assembly reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CP and EP contributor selection, legacy duplicate mode, SP trimming, and cross-rank writer gathering form one distributed teacher-cache producer policy.", + "Mooncake storage round trips, OPD loss execution, and debug artifacts remain separate transport, numerical, and observability boundaries." + ] + }, + { + "id": "TA-1015", + "scope": "split OPD pipeline shift and teacher-cache transport reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Causal teacher shifting, cache-index alignment and rejection, and Mooncake metadata admission form one OPD pipeline payload contract.", + "Endpoint reuse, student-version verification, and preparation-worker queueing remain separate orchestration boundaries." + ] + }, + { + "id": "TA-1016", + "scope": "fragmented unquantized LoRA checkpoint export and round-trip reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Runtime-rank export, PEFT hybrid-shared layout, SGLang shared-outer layout, and adapter-manager loading for both ownership modes form one unquantized checkpoint workflow.", + "Low-level EP slicing and quantized projection-subset round trips remain separate conversion and representation boundaries." + ] + }, + { + "id": "TA-1017", + "scope": "split optimizer publication and fatal dispatcher failure reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Successful publication, commit and handler-tail poisoning, and rank-zero fatal termination are outcomes of one post-mutation optimizer publication lifecycle.", + "Forward-backward completion and explicit epoch abort remain separate gradient-epoch transactions." + ] + }, + { + "id": "TA-1018", + "scope": "split empty-layout and deterministic adapter initialization reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Empty logical packing, discovery, replica classification, and ownership compilation form one empty-shard layout policy.", + "Coordinate, replica, LoRA-B, session, and FQN-order invariance form one deterministic initialization contract; real Gloo and explicit EP composition remain separate topology reports." + ] + }, + { + "id": "TA-1019", + "scope": "split server filesystem and compile-worker security reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Authoritative artifact roots, symlink rejection, and private diagnostic input admission form one server filesystem trust boundary.", + "Compile-target allowlisting, safe protocol type round trips, and oversized-frame rejection form one worker IPC trust boundary; outbound endpoint validation remains separate." + ] + }, + { + "id": "TA-1020", + "scope": "fragmented adapter-manager optimizer and checkpoint-load reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Optimizer construction, hyperparameter persistence, learning-rate updates, and mixed-rank multi-optimizer reload form one manager configuration and persistence lifecycle.", + "Session compatibility, weights-only behavior, structure admission, PEFT suffixes, sharded indices, and rank capacity form one load compatibility policy." + ] + }, + { + "id": "TA-1021", + "scope": "split FSDP mixed-precision selection and reduce-dtype reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Topmost protected-module selection, expert FSDP policy stripping, mesh-dependent reduce dtype, explicit overrides, and dtype admission form one mixed-precision policy.", + "Sequence-parallel folding, optional-boolean parsing, and manual prefetch direction remain separate configuration boundaries." + ] + }, + { + "id": "TA-1022", + "scope": "golden-bit reports split below BI contract-tree version", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All family, normalization, mean, softmax, matrix, and LM-head frozen hashes now form one v1 golden-tree gate; normalization and head hashes form one v2 gate.", + "Every failure retains its case and output label, and both versioned gates share the same H100 capability contract." + ] + }, + { + "id": "TA-1023", + "scope": "split FlashQLA Gate 2 shape-invariance reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed-versus-individual rows, total-token invariance, and block-DV tile invariance are the three shape branches of one Gate 2 bitwise contract.", + "Auto-CP admission and Gate 4 chunk-state handoff remain separate because they exercise different production decisions." + ] + }, + { + "id": "TA-1024", + "scope": "split MiniMax M3 activation, router, and text-runtime reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Clamped SwiGLU, biased sigmoid routing, text forward and backward, multimodal-token rejection, and parallel-mode admission form one MiniMax M3 runtime program.", + "Configuration, checkpoint mapping, and sparse-attention paging remain separate serialization and kernel boundaries." + ] + }, + { + "id": "TA-1025", + "scope": "cross-engine RMSNorm reports split by serving site class", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "QK, pre-summed residual-tree, post-attention residual, zero-centered family-1, and families-v2 cases share one cross-engine bitwise oracle and shape matrix.", + "The module-level SGLang dependency gate remains visible; native family admission, module dispatch, and fused backward stay in separate suites." + ] + }, + { + "id": "TA-1026", + "scope": "split Qwen3.5 Class-B and attention rotary reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Class-B dispatch precedence and fail-closed CPU or dtype admission form one rotary admission policy.", + "Dense and MoE half-rotate behavior plus exact post-RoPE BF16 casting form one attention projection policy; the low-level interleaved reference remains separate." + ] + }, + { + "id": "TA-1027", + "scope": "split Mooncake tensor codec and hidden-transport reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Tensor byte round trips, canonical dtype strings, metadata emission, suffixed object keys, and rank-2 and rank-3 fetches form one hidden transport contract.", + "Teacher cache consumption, malformed metadata, removal, and configuration precedence remain separate consumer, admission, and lifecycle boundaries." + ] + }, + { + "id": "TA-1028", + "scope": "split fused-GDN merged-forward and cache reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical LoRA folding, slice-local gradients, exact output projection, cache reuse, bounded generations, and release after adapter publication form one merged-weight lifecycle.", + "Delta arithmetic, manifest filtering, and sharded checkpoint loading remain separate representation and serialization boundaries." + ] + }, + { + "id": "TA-1029", + "scope": "split exact GDN module-dispatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fused gated RMSNorm routing, unsupported residual rejection, and full GatedDeltaNet exact-program routing form one model-program dispatch contract.", + "Gating and normalization forward and backward numerics plus the solve-tril warp pin remain separate kernel contracts." + ] + }, + { + "id": "TA-1030", + "scope": "split BI router GEMM forward and backward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FP32 reference agreement, empty and dtype admission, and analytical hidden and weight gradients form one router GEMM numerical contract.", + "Leading-dimension linear behavior, top-k weight processing, and MoE block integration remain separate API and model boundaries." + ] + }, + { + "id": "TA-1031", + "scope": "split RMSNorm family structure and numerical-routing reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Family-name admission, residual-shape rejection, and Qwen site declarations form one explicit family API structure policy.", + "Funnel equivalence, family vitality, zero-centered folding, and module dispatch form one GPU numerical-routing contract; undeclared-family enforcement remains a separate tripwire." + ] + }, + { + "id": "TA-1032", + "scope": "split SGLang-fused RMSNorm forward and backward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Residual and no-residual forward bit exactness and their analytical backward comparisons share one fused numerical implementation and CUDA gate.", + "CPU fallback, model integration, and trunk-contract dispatch remain separate portability and integration boundaries." + ] + }, + { + "id": "TA-1033", + "scope": "split DeepSeek V4 shared and routed SwiGLU limit reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The shared MLP and eager routed-expert checks jointly prove that the same configured SwiGLU limit reaches both expert implementations.", + "MoE routing, hash-table admission, and routing replay remain separate structural and lifecycle contracts." + ] + }, + { + "id": "TA-1034", + "scope": "split Nemotron H published-layout load, parity, and save reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Strict load accounting, HF numerical parity, ignored MTP admission, and byte-exact save reconstruction form one published-checkpoint codec transaction.", + "The supported published key set is explicit so ignored MTP input is not incorrectly required in saved output; stacked HF input and EP ownership remain separate representations." + ] + }, + { + "id": "TA-1035", + "scope": "split Qwen MoE fused-expert checkpoint-handler reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Qwen3 per-expert and Qwen3.5 stacked layouts now report through one fused-expert load/save policy, including deferred QLoRA expert loading.", + "Expert parameter registration and model-weight QARL filtering remain separate module and filesystem boundaries." + ] + }, + { + "id": "TA-1036", + "scope": "split QLoRA quantized execution and NVFP4 scale-merge reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quantized storage, forward and backward, dequantization, prequantized loading, EMA scale convention, and LoRA merge-requantization form one quantized-weight lifecycle.", + "Injection, block-FP8 checkpoint representation, and optimizer-state reset remain separate integration and lifecycle contracts." + ] + }, + { + "id": "TA-1037", + "scope": "split DeepSeek V4 decoder runtime and gradient-checkpoint reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "C128 execution, causal-LM forward and backward, hash-layer threading, and decoder checkpoint wrapping form one model runtime contract.", + "Construction and topology admission plus dtype preservation remain separate structural policies." + ] + }, + { + "id": "TA-1038", + "scope": "split DTensor checkpoint and rank-zero broadcast loading reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Copying full checkpoint tensors into existing replicated or sharded DTensors and materializing multi-axis DTensors for one writer are the load and save directions of one tensor checkpoint codec.", + "Object payload transport, NCCL device selection, weight-load group routing, and handler-filtered rank-zero loading form one broadcast loading transaction; state-dict resolution and grouped expert routing remain separate policies." + ] + }, + { + "id": "TA-1039", + "scope": "split Muon Gram-Newton-Schulz construction and execution reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Builder keyword propagation, byte-limit admission, SGD fallback, an actual Gram-Newton-Schulz parameter update, and restart autotuning form one configured optimizer lifecycle.", + "Quack backend selection, grouped matrix scheduling, standard Newton-Schulz, CUDA compute dtype, and model parameter classification remain separate branches." + ] + }, + { + "id": "TA-1040", + "scope": "split BI trunk-linear forward and backward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Persistent-GEMM forward bit exactness, batch invariance, dtype admission, and cuBLAS input, weight, and bias gradients form one wrapped-linear numerical contract under the same CUDA gate.", + "Wrapper selection and global-interpose training admission remain separate structural and global-mode boundaries." + ] + }, + { + "id": "TA-1041", + "scope": "split DeepSeek V3 expert checkpoint layout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "External per-expert and internal fused layouts, dense and packed EP slicing, requested device and dtype, and quantization-config discovery form one expert checkpoint codec policy.", + "Every prior dense, packed, internal, EP-local, and configured quantization assertion remains in the surviving report." + ] + }, + { + "id": "TA-1042", + "scope": "split Nemotron H model runtime and gradient-checkpoint reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Router output, loss backward through every mixer type, and full-layer gradient checkpointing form one model training runtime contract.", + "Packed variable-length equivalence remains separate because it changes document-boundary state propagation." + ] + }, + { + "id": "TA-1043", + "scope": "split RoPE registry precision and exact-lane bit reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Registry-wide FP32 frequency construction, BF16 consumption, and unchanged zero-K3 contract-lane cosine and sine bits form one frequency-precision policy.", + "Lazy cache growth and architecture-specific CUDA construction remain separate cache-lifecycle and device-placement boundaries." + ] + }, + { + "id": "TA-1044", + "scope": "split DistSignSGD builder and direct-step reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FSDP2 builder admission, weight-decay grouping, hook configuration, and a preaggregated update with decoupled decay form one configured optimizer contract.", + "Reduce-scatter sign timing and local-versus-FSDP hook ownership remain separate communication and gradient-ownership boundaries." + ] + }, + { + "id": "TA-1045", + "scope": "split sequence-parallel and LM-head explicit synchronization reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sequence-parallel gradient sums, DTensor skipping, adapter-finalization exclusions, LM-head replica gradient handling, and marked-parameter broadcast form one explicit synchronization policy.", + "Token counting, gradient clipping, pipeline loss chunking, and their independent collectives remain separate trainer contracts." + ] + }, + { + "id": "TA-1046", + "scope": "split dense and MoE hidden-component hook reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Residual, attention, normalization, dense MLP, routed-MoE callback, and shared-expert components are alternate producers for one hidden-component capture pipeline.", + "Summary formatting, ranked tensor dumps, and trusted override loading remain separate output and trust boundaries." + ] + }, + { + "id": "TA-1047", + "scope": "split base, pipeline, and LoRA checkpoint compatibility reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local and pipeline key discovery, metadata unions, QARL buffer mismatch, base-to-LoRA loading, and LoRA-only loading form one checkpoint model-compatibility policy.", + "Distributed checkpoint transport and optimizer-state filtering remain separate IO and payload boundaries." + ] + }, + { + "id": "TA-1048", + "scope": "split gradient-epoch completion, abort, and failure reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forward-backward rendezvous and commit, explicit abort routing, uniform rejection, and rank-asymmetric failure conversion form one dispatcher gradient-epoch lifecycle.", + "Session registration, save operations, and post-optimizer publication poisoning remain separate RPC and mutation boundaries." + ] + }, + { + "id": "TA-1049", + "scope": "split DeepSeek V3 router-output and replay reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Aux-loss-driven router-logit emission across sparse layers and recording selected indices and weights form one router observability and replay contract.", + "Model backward and router freezing plus LoRA target injection remain separate training and adapter policies." + ] + }, + { + "id": "TA-1050", + "scope": "split MoE token permutation and unpermutation reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Expert-sorted routing weights and scatter-add reconstruction are the encode and decode directions of one memory-efficient token permutation codec.", + "All-to-all score ordering and hidden-dimension chunking remain separate transport boundaries." + ] + }, + { + "id": "TA-1051", + "scope": "split SignSGD builder and direct-step reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Builder construction, decay and no-decay grouping, dense sign updates, decoupled weight decay, and sparse-gradient rejection form one configured SignSGD contract.", + "DistSignSGD distributed communication and cautious-decay behavior remain in their dedicated suites." + ] + }, + { + "id": "TA-1052", + "scope": "split Qwen3.5 families-v2 zero-centered backward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Effective folded-weight gradients and the residual twin's output and residual gradient paths are two call shapes of one zero-centered families-v2 backward contract.", + "Both CPU reference comparisons and every input, residual, and weight gradient assertion remain in the surviving report." + ] + }, + { + "id": "TA-1053", + "scope": "split FP8 config translation and Blackwell admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Legacy FP8 config normalization, incompatible external runtime rejection, and explicit Blackwell artifact admission form one FP8 configuration boundary.", + "BF16 layer-island resolution and injection remain a separate model transformation policy." + ] + }, + { + "id": "TA-1054", + "scope": "split Quack PTX compilation output and entry-discovery reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unique temporary outputs, bounded ptxas execution, cleanup, and exact kernel-entry discovery form one PTX compilation process-safety contract.", + "Worker framing and cache-key hashing remain separate IPC and cache-trust boundaries." + ] + }, + { + "id": "TA-1055", + "scope": "split merged-LoRA fold, cache, and fused-admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical linear and expert folding plus straight-through factor gradients now form one low-level fold numerical contract.", + "LoraLinear merged selection, gradient parity, optimizer and active-rank cache invalidation form one linear lifecycle; MoE canonical merged views, versioned caches, and fused-expert admission form one expert lifecycle.", + "Native EP execution and trunk-linear wrapping remain separate integration boundaries." + ] + }, + { + "id": "TA-1056", + "scope": "split MoE-LoRA construction and GPU numerical reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Backend-specific initialization, frozen and trainable ownership, runtime rank slicing, from-module conversion, model injection, and block injection form one construction policy.", + "Zero-delta base equivalence and eager-versus-native or Triton output and gradient agreement form one GPU numerical policy under the same capability gate.", + "CPU eager execution, zero-token gradients, and EP router-score application remain separate runtime boundaries." + ] + }, + { + "id": "TA-1057", + "scope": "split NVFP4 2D and expert fake-quant reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Independent 2D reference agreement, shape admission, linear STE, 3D expert STE, expert isolation, and fused gate-up scale ownership form one NVFP4 fake-quant contract.", + "Every dense, expert, and fused projection representation remains covered in the surviving report." + ] + }, + { + "id": "TA-1058", + "scope": "split P2P prepare, completion, and cleanup reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Cold and cached prepare behavior are states of one P2P initialization handshake under the same backend contract.", + "Successful completion, pending-transfer failure, receiver notification, and destroy cleanup form one terminal synchronization lifecycle.", + "Fanout, slicing, coalescing, diagnostics, and multi-sender routing remain separate transport boundaries." + ] + }, + { + "id": "TA-1059", + "scope": "split NCCL endpoint and flattened transfer reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Initialization and transfer endpoint ports plus the optional two-phase receiver protocol form one NCCL endpoint transaction policy.", + "Flat, chunked-flat, and receiver-fenced hybrid buckets are load-format branches of one flattened transfer contract.", + "Endpoint health and invalid multi-rank direct format remain independent admission boundaries." + ] + }, + { + "id": "TA-1060", + "scope": "split GDN contract packing, routing, admission, and state reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed convolution weights, armed routing, fail-closed admission, call-scoped state, and checkpoint recomputation form one exact GDN contract lifecycle.", + "Low-level CUDA parity, full-block integration, and optional SGLang tree-kernel parity remain distinct numerical boundaries." + ] + }, + { + "id": "TA-1061", + "scope": "split FlashAttention API and page-size-one cache path reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fixed-length and variable-length calls are public shapes of one FlashAttention API behavior contract.", + "SGL, paged FlashAttention, flags-off, and FA4 selection are branches of one page-size-one KV-cache routing policy.", + "Backend registry resolution and eager head-layout numerics remain separate implementation boundaries." + ] + }, + { + "id": "TA-1062", + "scope": "split authoritative adapter optimizer outcome reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Empty and aborted epochs, a successful clipped update, nonfinite input, optimizer failure, and collective failure are states of one authoritative optimizer lifecycle.", + "Capture ownership, publication admission, exact LM-head coherence, and checkpoint restore remain separate component boundaries." + ] + }, + { + "id": "TA-1063", + "scope": "split optimizer checkpoint save and resume reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sharded manifest emission, identity rejection, moment restoration, and bitwise continuation are the write and read sides of one optimizer checkpoint codec.", + "Artifact admission and logical cross-layout resharding remain separate compatibility boundaries." + ] + }, + { + "id": "TA-1064", + "scope": "split sampling adapter reconciliation and scoped tracking reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Stale-state reconciliation, query failure preservation, model-scoped tracking, and failed-load atomicity form one sampling-session adapter lifecycle.", + "Sampler checkpoint storage and adapter-only export remain separate filesystem and orchestrator operations." + ] + }, + { + "id": "TA-1065", + "scope": "split inference weight-sync forwarding and admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Endpoint forwarding, pool selection, quantization admission, and cache invalidation are request branches of one inference weight-sync API transaction.", + "Endpoint registration, health refresh, and receiver capability detection remain separate lifecycle boundaries." + ] + }, + { + "id": "TA-1066", + "scope": "split model-session creation and kill reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Normalized registration, duplicate and topology admission, reserved checkpoint handling, kill, optional final save, and default-session protection form one model-session lifecycle.", + "The lightweight create-session alias remains a separate endpoint registration contract." + ] + }, + { + "id": "TA-1067", + "scope": "split checkpoint weights-info and legacy session metadata reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Disk-backed weights information, path admission, full-weight metadata, and legacy SignSGD metadata upgrade form one checkpoint session-spec decoding policy.", + "No metadata assertion or legacy compatibility case was removed." + ] + }, + { + "id": "TA-1068", + "scope": "split quantized exporter projection and MoE layout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fused QKV, MLA, linear-attention, GKN expert, and fused gate-up layouts are architecture branches of one exported-model tensor layout contract.", + "CLI and directory behavior, source admission, QARL folding, and low-level FP8 quantization remain separate boundaries." + ] + }, + { + "id": "TA-1069", + "scope": "split expert-adapter backend capability and factor ownership reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Backend capability and plan identity plus factor ownership, reduction domains, and checkpoint persistence form one expert-adapter structural contract.", + "Construction and semantic preservation remain independently reported." + ] + }, + { + "id": "TA-1070", + "scope": "split generic and model-family expert-adapter injection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Exact target-subset forwarding and GLM, Qwen3, and Qwen3.5 wrapper construction form one expert-adapter injection policy.", + "Each family, backend, target set, quantization format, and checkpoint-buffer assertion remains covered." + ] + }, + { + "id": "TA-1071", + "scope": "split expert-adapter semantic preservation and fail-closed reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Supported SiLU preservation and rejection of incompatible activations, biases, quantization groups, target sets, and model-family semantics form one fail-closed semantic contract.", + "Runtime numerical parity remains outside this structural suite." + ] + }, + { + "id": "TA-1072", + "scope": "split canonical MoE plan and group-alias topology reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Trainer and sampler plan identity, topology admission, logical ordinals, and world-32 CP, EP, and expert-FSDP group aliases form one canonical MoE topology contract.", + "Graph metadata and transport selection remain separate boundaries." + ] + }, + { + "id": "TA-1073", + "scope": "split canonical MoE distributed transport numerical reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Two- and eight-contributor dense transport plus sixteen-contributor packed and CP-sharded parity are topology cases of one distributed transport numerical contract.", + "All 2-, 8-, and 16-process subprocess checks still execute in the surviving report." + ] + }, + { + "id": "TA-1074", + "scope": "split adapter auto-load, explicit load, and rank-zero broadcast reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Evicted auto-load, fresh materialization, explicit path load, rollback, and all-rank or rank-zero-broadcast restoration are branches of one adapter load lifecycle.", + "Adapter registration and save admission remain separate mutation boundaries." + ] + }, + { + "id": "TA-1075", + "scope": "split token selection and CP-sharded KL position reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Target selection, boundary behavior, raw-weight cross-checking, hidden summaries, and CP-sharded KL position mapping form one token-diagnostics policy.", + "Hidden-component capture, tensor-dump output, and trusted override input remain separate boundaries." + ] + }, + { + "id": "TA-1076", + "scope": "split OPD backend numerics and gradient-reduction reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Reference, streaming, low-memory, and sharded-store forward agreement plus backward, partial reduction, and output dtype form one OPD numerical backend contract.", + "Output-shape edge behavior and hidden-only distance remain separately reported." + ] + }, + { + "id": "TA-1077", + "scope": "split fused selected-logprob small and production-vocabulary numerical reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dtype, bias, temperature, frozen-head, irregular-tail, Qwen, and GPT-OSS vocabulary cases form one fused selected-logprob forward and backward numerical contract.", + "Loss-dispatch agreement and the no-full-logits memory gate remain separate integration and resource boundaries." + ] + }, + { + "id": "TA-1078", + "scope": "split DRGRPO temperature behavior from forward numerical report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forward value, gradients, metrics, zero boundaries, advantage direction, KL penalty, and logprob-temperature behavior form one DRGRPO objective contract.", + "Microbatch reducer composition remains a separate aggregation boundary." + ] + }, + { + "id": "TA-1079", + "scope": "split packing capacity and edge-admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Concatenation, capacity splitting, mixed lengths, empty and single input, oversize admission, missing fields, NumPy conversion, and microbatch validation form one core packing policy.", + "Packed metadata, disabled mode, and full pack-to-unpack roundtrip remain separate contracts." + ] + }, + { + "id": "TA-1080", + "scope": "split receiver postprocess and sync quantization configuration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Receiver postprocess selection, FP8 KV-cache requirements, unsupported-format admission, and generated BF16 islands form one quantized weight-sync configuration contract.", + "Tensor quantization numerics remain in the dedicated FP8 sync suite." + ] + }, + { + "id": "TA-1081", + "scope": "split LoRA preparation, parameter extraction, and inference-layout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Adapter materialization, parameter filtering and tied aliases, compile-name normalization, and architecture-specific unfusing form one sync-source tensor preparation pipeline.", + "Bucket sizing, transport routing, and sparse-delta selection remain separate boundaries." + ] + }, + { + "id": "TA-1082", + "scope": "split EP collection admission and expert-data reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "P2P sender selection, direct-EP collection admission, and gated or nongated local expert projection collection form one EP synchronization-source policy.", + "Actual transport remains independently tested." + ] + }, + { + "id": "TA-1083", + "scope": "split checkpoint save failure and exact-active-LoRA admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Factor-only admission, rank-zero artifact failures, LoRA-only failures, and pre-barrier error surfacing form one fail-closed checkpoint save policy.", + "No downstream conversion or collective is allowed after an admission or write failure." + ] + }, + { + "id": "TA-1084", + "scope": "split dense and MoE LoRA save-format reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Live dense target resolution and collective stacked-MoE factor slicing are format branches of one LoRA checkpoint export contract.", + "Optimizer-state checkpointing remains in the adapter optimizer resume suite." + ] + }, + { + "id": "TA-1085", + "scope": "split removed and unsupported server configuration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Removed fields, incompatible quantized modes, vLLM-only runtime knobs, broadcast loading, and unsupported multi-adapter modes form one fail-closed server configuration boundary.", + "Valid shipped and feature-specific configurations remain independently reported." + ] + }, + { + "id": "TA-1086", + "scope": "split general, optimizer, runner, and model-specific server runtime reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical defaults, nested runtime controls, receiver cache dtype, R3 transport, gradient buckets, Muon options, runner compatibility, sparse MLA, and MoE routing controls form one runtime configuration roundtrip.", + "Quantized-training and parallel-topology configuration remain separate specialized boundaries." + ] + }, + { + "id": "TA-1087", + "scope": "split teacher-cache distributed assembly and Mooncake integration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Contributor selection, CP and DP assembly, valid-label trimming, Mooncake metadata emission, byte roundtrip, and activation-cache consumption form one teacher hidden-cache lifecycle.", + "OPD loss execution and debug artifacts remain separate consumers and outputs." + ] + }, + { + "id": "TA-1088", + "scope": "split Muon Gram-Newton-Schulz configuration and grouping reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Builder options, fallback behavior, restart autotuning, grouped shapes, transpose equivalence, fused halves, and byte-limit chunking form one configured Gram-Newton-Schulz optimizer contract.", + "Quack backend selection, standard Newton-Schulz, and CUDA compute dtype remain separate algorithm and platform boundaries." + ] + }, + { + "id": "TA-1089", + "scope": "split Triton-grouped and scalar-Quack FP8 numerical reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Block-loop and Triton-grouped forward and weight gradients plus scalar-Quack per-expert scaling are backend branches of one grouped FP8 GEMM numerical contract.", + "DeepGEMM subprocess isolation and model train-step integration remain separate gates." + ] + }, + { + "id": "TA-1090", + "scope": "split dispatcher selection, routing slice, and row-provenance reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "DP, EP, CP, and legacy rank selection, routing-payload slicing, rank-local row grouping, and source provenance form one dispatcher input-distribution policy.", + "Packing strategy and R3 payload storage remain separate upstream boundaries." + ] + }, + { + "id": "TA-1091", + "scope": "split completion rendezvous and per-token result merging reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local payload trimming, rank rendezvous, CP replica deduplication, disagreement rejection, and rank-zero per-token merging form one dispatcher completion transaction.", + "Diagnostic dumping remains a separate output boundary." + ] + }, + { + "id": "TA-1092", + "scope": "split request-processor and runner-dispatcher forward lifecycle reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Processor readiness, forward and backward execution, timing propagation, invalid-target rejection, shutdown, model identity, auto-load, and R3 forwarding form one request-to-runner compute lifecycle.", + "NCCL sync, optimizer and checkpoint RPCs, and payload storage remain separate operations." + ] + }, + { + "id": "TA-1093", + "scope": "split raw-numerator accumulation and staged gradient-capture reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Raw numerator accumulation, model-gradient clearing, FP32 scratch reuse, staged commit, DTensor preservation, and atomic prevalidation form one gradient-capture transaction.", + "Ownership-plan compilation and optimizer mutation remain separate lifecycle boundaries." + ] + }, + { + "id": "TA-1094", + "scope": "standalone exact LM-head optimizer coherence report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Scalar-tensor optimizer-state coherence is a distributed validation branch of the authoritative optimizer lifecycle, not an independently meaningful user behavior.", + "Its collective mocks and assertions now run inside the optimizer lifecycle contract." + ] + }, + { + "id": "TA-1095", + "scope": "split adapter checkpoint trust, restore-lifecycle, and compatibility reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Trusted paths, strict target manifests, lifecycle reset, ownership-plan admission, optimizer compatibility, learning-rate rules, checkpoint structure, and PEFT sharding form one checkpoint restore and admission policy.", + "Coordinator-driven checkpoint materialization remains a separate server integration boundary." + ] + }, + { + "id": "TA-1096", + "scope": "split adapter eviction and mixed-adapter manager reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Mixed ranks and optimizers, adapter switching, training, checkpoint reload, capacity eviction, dirty-state protection, multi-rank rejection, and save-failure rollback form one multi-adapter lifecycle.", + "No eviction or failure assertion was removed." + ] + }, + { + "id": "TA-1097", + "scope": "split state-dict resolution and rank-zero loading reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local-versus-broadcast shard resolution is the input phase of rank-zero checkpoint transport and loading.", + "Object transport, group selection, handler-filtered prefetch, and state-dict resolution now form one rank-zero loading policy." + ] + }, + { + "id": "TA-1098", + "scope": "split grouped-load routing and strict postprocess coverage reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense and expert routing, format conversion, group fallback, and strict parameter and persistent-buffer coverage are phases of one grouped checkpoint-loading transaction.", + "Distributed DTensor checkpoint materialization remains a separate save-side correctness gate." + ] + }, + { + "id": "TA-1099", + "scope": "split canonical GLM-5.2 trainer topology and layer-plan reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Certified world, EP, and CP topology plus the official producer schedule, pipeline split, malformed-plan rejection, and indexer allocation form one canonical layer-plan contract.", + "The topology admission assertions still execute before layer-plan validation." + ] + }, + { + "id": "TA-1100", + "scope": "split index-share lifecycle and FSDP identity reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Publication, reuse, concurrency rejection, exception cleanup, and identity preservation across FSDP input casting form one index-share lifecycle.", + "All lifecycle and model-integration assertions remain intact." + ] + }, + { + "id": "TA-1101", + "scope": "split correction-bias checkpoint and canonical MoE configuration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Correction-bias FP32 preservation and checkpoint admission are router branches of the canonical MoE configuration and selection contract.", + "Native sampler codec parity and end-to-end semantic logprob composition remain independent gates." + ] + }, + { + "id": "TA-1102", + "scope": "split SGLang fused-MoE resolution and block-dispatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Automatic and explicit enablement, eligibility, logging, flag-off preservation, and block entrypoint selection form one fused-MoE resolution and dispatch policy.", + "Real-kernel parity remains an independent GPU gate." + ] + }, + { + "id": "TA-1103", + "scope": "split fused-expert admission and trainable-dispatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Semantic admission failures and the gradient-sensitive choice between the autograd function and plain kernel are branches of one fused-expert dispatch contract.", + "Gradient numerics remain separately capability-gated." + ] + }, + { + "id": "TA-1104", + "scope": "split fused-MoE weight-mode and strided-adapter layout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Transient, cached, and zero-copy strided modes, cache invalidation, serving tensor layout, and adapter gate-up layout form one kernel weight-layout contract.", + "No cache, storage-alias, or layout assertion was removed." + ] + }, + { + "id": "TA-1105", + "scope": "split simulator topology, shape, and analytical-ledger reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Topology resolution, local-token shapes, FLOPs, activation storage, and communication bytes form one analytical accounting contract.", + "Observed benchmark ingestion and kernel ranking remain separate empirical policies." + ] + }, + { + "id": "TA-1106", + "scope": "split simulator config and metadata resolution from path-admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Config fingerprinting, cached and known-model metadata resolution, calibration-pack containment, symlink rejection, and restricted local reads form one input resolution and admission policy.", + "All traversal and trust-boundary assertions remain intact." + ] + }, + { + "id": "TA-1107", + "scope": "split ad-hoc Qwen and built-in calibration-pack reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Markdown ingestion, calibration evaluation, scenario planning, built-in pack replay, feasibility boundaries, and consolidated validation form one calibration lifecycle.", + "Generic observed-run ingestion remains an independent input format and planning contract." + ] + }, + { + "id": "TA-1108", + "scope": "split direct-output and expert eFSDP real-autograd reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Direct-output sharding plus shared-owner and all-owner expert layouts are branches of one four-GPU FSDP adapter-gradient ownership contract.", + "Each distributed subprocess and certification marker remains required." + ] + }, + { + "id": "TA-1109", + "scope": "split unquantized and quantized expert AllToAll reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Eager, Triton, native, Quack, NF4, NVFP4, block-FP8, projection-subset, and all-owner cases are backend and representation branches of one AllToAll ownership policy.", + "Every two-GPU subprocess still executes." + ] + }, + { + "id": "TA-1110", + "scope": "split unquantized and quantized DeepEP expert reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Hybrid-shared, all-owner, and quantized Quack cases form one optional DeepEP adapter-gradient ownership contract.", + "The optional dependency and two-GPU capability gate remain on the report." + ] + }, + { + "id": "TA-1111", + "scope": "split positive SGLang EP presentation and dispatch-admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Top-k-one pair presentation, FP32 routing weights, local weight layout, flag-off preservation, empty-rank handling, and semantic rejection form one EP dispatch and admission policy.", + "Slot combine, trainable dispatch, weight modes, and real gradient parity remain separate contracts." + ] + }, + { + "id": "TA-1112", + "scope": "split FP8 sync quantization and projection-selection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "BF16 islands, block scale and zero-padding semantics, projection selection, skip lists, stacked tensors, and existing FP8 values form one CPU sync-quantization contract.", + "Expert workspace and GPU execution remain independently reported." + ] + }, + { + "id": "TA-1113", + "scope": "split dense, quantized, MoE, and fused-GDN LoRA sync-extraction reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense LoRA, QLoRA, quantized MoE factors, and fused GDN factors are architecture branches of one adapter-folding sync-source policy.", + "Every merged-weight and raw-factor exclusion assertion remains intact." + ] + }, + { + "id": "TA-1114", + "scope": "split generic, exact-LM-head, and replica-topology runner compiler reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Generic module and direct-output ownership, exact TP16 LM-head VJP masks, replica divisors, group coverage, and fail-closed topology admission form one runner gradient-ownership compiler policy.", + "Expert-factor compilation remains a separate specialized contract." + ] + }, + { + "id": "TA-1115", + "scope": "split effective LM-head selection and authoritative analytical-step reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Merged and legacy effective-head selection are input branches of the direct-output analytical gradient capture and optimizer step.", + "The surviving contract checks selection bytes, gradients, capture, norm, and parameter mutation end to end." + ] + }, + { + "id": "TA-1116", + "scope": "split AnyPrecision AdamW state-strategy and DTensor-offload reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Denominator chunking, Kahan compensation, gradient reuse, CPU state offload, DTensor local-shard wrapping, and device restoration form one optimizer state-strategy policy.", + "Cautious decay math remains independently reported." + ] + }, + { + "id": "TA-1117", + "scope": "split routed-expert topology, owner remap, and physical-buffer reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP16 and MoE-TP1 admission, all owner-slot remaps, global and owner-local factor banks, sampler buffer shapes, dtypes, and zero padding form one routed-bank layout policy.", + "Literal sampler numerics remain a separate GPU gate." + ] + }, + { + "id": "TA-1118", + "scope": "split zero-token, sentinel, and mixed-owner routed-gradient reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Owned zero-base gradients, all-sentinel structural zeros, input-layout rejection, and top-k-eight mixed-owner VJPs are branches of one routed-gradient edge policy.", + "The Hopper and SGLang capability gate remains on the combined report." + ] + }, + { + "id": "TA-1119", + "scope": "split ownership compiler topology, fingerprint, declaration, and replica reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Topology declarations, authority masks, rank-local fingerprint invariance, fail-closed structural admission, and orthogonal replica coverage are phases of one ownership-plan compilation transaction.", + "Compiled producer execution and residual gradient transport remain separate runtime contracts." + ] + }, + { + "id": "TA-1120", + "scope": "split shared-expert construction and logical checkpoint-state reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Construction, runtime admission, logical state, and checkpoint-state behavior form one shared-expert contract.", + "Physical SGLang views remain separate so optional dependency admission cannot hide the CPU structural report." + ] + }, + { + "id": "TA-1121", + "scope": "split exact TP1 configuration and dtype-state identity reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Runtime admission and dtype moves are lifecycle phases of the exact TP1 configuration contract.", + "Packed-state master dtype and object identity assertions remain intact." + ] + }, + { + "id": "TA-1122", + "scope": "split exact TP1 forward VJP and backward-safety reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Reference forward values, surrogate gradients, and rejected unsafe backward paths form one numerical and autograd policy.", + "All value, gradient, and error assertions remain intact." + ] + }, + { + "id": "TA-1123", + "scope": "split request-scoped NCCL synchronization from orchestrator control operations", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Request-scoped NCCL group naming is a synchronization branch of the optimizer, checkpoint, registration, and lifecycle control report.", + "The exact group-name and operation assertions remain intact." + ] + }, + { + "id": "TA-1124", + "scope": "split GLM5 indexer construction and DSA mask reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Indexer construction and its DSA mask behavior form one indexer-selection policy.", + "The mask branch runs in an isolated monkeypatch context within the surviving report." + ] + }, + { + "id": "TA-1125", + "scope": "split sparse MLA reference wrapper and attention-integration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Reference behavior, wrapper semantics, and attention integration are layers of one sparse-MLA policy.", + "Full forward and recompute behavior remain separate end-to-end gates." + ] + }, + { + "id": "TA-1126", + "scope": "split GLM5 sparse-KV adapter weight and dispatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sparse-KV adapter weights, adapter dispatch, and MoE dispatch form one adapter-and-routing policy.", + "The sparse-KV branch runs in an isolated monkeypatch context." + ] + }, + { + "id": "TA-1127", + "scope": "split native routed-partial module entry and actual-operand capture reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FSDP pre-forward hook entry is part of the native-combine execution boundary whose actual operands are captured.", + "Collective padding and fused-gate gradient parity remain independent contracts." + ] + }, + { + "id": "TA-1128", + "scope": "split optimizer parameter identity and transaction snapshot failure reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical identity, binding validation, recursive snapshots, and failed-collective commit behavior form one optimizer transaction policy.", + "Logical resharding remains independently reported." + ] + }, + { + "id": "TA-1129", + "scope": "split optimizer checkpoint resume and artifact-admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Successful sharded save and bitwise resume and failed artifact admission are complementary branches of one checkpoint recovery policy.", + "Legacy pickle, missing artifact, and resident-state preservation assertions remain intact." + ] + }, + { + "id": "TA-1130", + "scope": "split P2P receiver placement and replicated staged-source reuse reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Per-receiver slices and replicated locators are branches of one receiver-placement policy.", + "The session, pointer, length, and staged-source identity assertions remain intact." + ] + }, + { + "id": "TA-1131", + "scope": "split P2P staging and receiver-handle coalescing reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Typed staging, registration lifetime, alignment, and coalescing are phases of one transfer staging policy.", + "Coalescing with and without receiver handles runs in an isolated environment context." + ] + }, + { + "id": "TA-1132", + "scope": "split P2P flush and weight-version payload from sync completion", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Flush-cache and weight-version propagation are completion-payload branches of the sync lifecycle.", + "Cache retention, tied aliases, failure completion, and cleanup remain in the same lifecycle report." + ] + }, + { + "id": "TA-1133", + "scope": "split sampler adapter reconciliation from adapter tracking and invalidation", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Reconciliation, model-scoped atomic tracking, receiver invalidation, listing, deletion, and resolution form one sampler adapter-state policy.", + "Sampler-weight export remains a separate orchestration contract." + ] + }, + { + "id": "TA-1134", + "scope": "split zero-token LoRA structural gradients from eager expert forward and backward", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Zero-token output is the empty-input branch of eager expert forward and backward behavior.", + "Every local-factor structural-gradient assertion remains intact." + ] + }, + { + "id": "TA-1135", + "scope": "split prequantized exclude metadata parsing from checkpoint-handler exclusion behavior", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Metadata formats, precedence, malformed inputs, dense and MoE skip behavior, and auxiliary-key passthrough form one checkpoint exclusion policy.", + "General prequantized detection and non-excluded loading remain independently reported." + ] + }, + { + "id": "TA-1136", + "scope": "split exact GLM52 attention inventory and construction-admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The canonical attention-factor inventory and fail-closed input branches define one exact construction policy.", + "All rank, alpha, component, dispatch, sparse-MLA, source, dtype, and identity assertions remain intact." + ] + }, + { + "id": "TA-1137", + "scope": "split exact GLM52 MoE inventory and construction-admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The shared and routed inventory and invalid dependency and topology branches define one exact MoE construction policy.", + "Post-EP layout and selected-logprob LM-head specialization remain independent reports." + ] + }, + { + "id": "TA-1138", + "scope": "split DeepSeek-V4 construction topology and precision-preservation reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Pipeline admission, parallel-group wiring, FP32 markers, dtype casts, and complex RoPE buffers are one construction-state policy.", + "Full model forward, backward, hash routing, and checkpoint recomputation remain a separate runtime report." + ] + }, + { + "id": "TA-1139", + "scope": "split routing-weight position numerics and configuration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Before-down and after-down numerical behavior and the setting that selects them define one routing-position contract.", + "Reference gradients, error classes, lazy configuration, environment, auto, explicit, and invalid branches remain intact." + ] + }, + { + "id": "TA-1140", + "scope": "split dense and MoE FP32 cast-once LoRA merge reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense linear and MoE expert layouts exercise the same zero-preserving FP32-add-then-cast invariant.", + "All dtype, fused gate-up, down-projection, zero, and nonzero assertions remain intact." + ] + }, + { + "id": "TA-1141", + "scope": "split FlashQLA forward and backward parity reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forward output, final state, and input gradients form one numerical parity report per head shape.", + "Both parameterized head shapes and the Hopper and TileLang capability gate remain unchanged." + ] + }, + { + "id": "TA-1142", + "scope": "split NVFP4 and block-FP8 expert-load reference reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both formats are branches of one prequantized expert-load policy.", + "Packed bytes, scales, global factors, amax, projection, shape, and dequantization assertions remain intact." + ] + }, + { + "id": "TA-1143", + "scope": "split flat and GKN NF4 codec reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Flat and GKN layouts implement one NF4 codebook and quantization-error contract.", + "Packing, scale, dtype, zero, shape, and error assertions remain intact." + ] + }, + { + "id": "TA-1144", + "scope": "split block-FP8 and NVFP4 GNK-to-GKN conversion reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both prequantized formats exercise one checkpoint-to-runtime layout-conversion invariant.", + "Direct parity, roundtrip, non-square, stacking, scale, and error assertions remain intact." + ] + }, + { + "id": "TA-1145", + "scope": "split same-NK and same-MN grouped GEMM reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Same-NK and same-MN geometries are branches of one grouped GEMM kernel-family contract.", + "Transpose, uneven, empty-group, contiguity, device, shape, and numerical assertions remain intact." + ] + }, + { + "id": "TA-1146", + "scope": "split dense and expert-layout Muon full-gradient oracle reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense 2-D and expert 3-D sharding are layout branches of one full-gradient oracle-parity claim.", + "Both two-GPU subprocesses remain independently executed." + ] + }, + { + "id": "TA-1147", + "scope": "split dense and expert-layout Muon shard-local negative controls", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense 2-D and expert 3-D sharding are layout branches of one shard-local negative-control claim.", + "Both two-GPU subprocesses remain independently executed." + ] + }, + { + "id": "TA-1148", + "scope": "split dense Qwen3.5 RMSNorm dispatch and site-assignment reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Structural selection, family dispatch, and assignment to every model site form one CPU RMSNorm resolution contract.", + "Invalid mode, v1, v2, coexistence, GDN exclusion, layer, and final-norm assertions remain intact." + ] + }, + { + "id": "TA-1149", + "scope": "split Qwen3.5-MoE RMSNorm dispatch and site-assignment reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Family dispatch and assignment to every MoE model site form one CPU RMSNorm resolution contract.", + "Ordinary-mode, exact-mode, v2, layer, and final-norm assertions remain intact." + ] + }, + { + "id": "TA-1150", + "scope": "split Qwen3.5-MoE family-1 integration and family-2 residual GPU reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Family-1 interpose and full-layer parity and family-2 residual composition form one GPU bit-exact integration contract.", + "The CUDA capability gate and every bitwise assertion remain intact." + ] + }, + { + "id": "TA-1151", + "scope": "split DeepSeek-V4 converter meta-dtype and roundtrip reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Meta-model dtype preparation is the input phase of the same HF-to-DCP roundtrip conversion policy.", + "FP32 destinations, BF16 tensors, DCP output, exact loads, and legacy sidecar-free LoRA assertions remain intact." + ] + }, + { + "id": "TA-1152", + "scope": "split ordinary and cross-shard DeepSeek-V4 conversion reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Single-shard roundtrip and deferred cross-shard weight and scale pairing are branches of one conversion policy.", + "Process-group state is explicitly reset between the two converter invocations." + ] + }, + { + "id": "TA-1153", + "scope": "split MoE-block and decoder-layer torch compile reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Block and decoder composition are lower-level phases of one compiler compatibility policy.", + "Every available expert backend, AOT eager, Inductor, forward, backward, and numerical assertion remains intact." + ] + }, + { + "id": "TA-1154", + "scope": "split local and EP fused-expert FP32 routing-gradient reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unfiltered local, filtered local, and EP backward paths exercise one FP32 routing-gradient oracle.", + "The local cases run as an internal loop and the EP case runs once, preserving all gradient assertions without duplicate work." + ] + }, + { + "id": "TA-1155", + "scope": "split EP adapter registry and backend argument FP8-boundary reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Optional registration, live signatures, shared arguments, activation forwarding, and FP8 admission define one adapter boundary contract.", + "The isolated optional-import context and every available backend branch remain intact." + ] + }, + { + "id": "TA-1156", + "scope": "split exact Qwen hook preparation and hybrid trunk selection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Merged-LoRA preparation and selected trunk modules are phases of one exact wrapping policy.", + "The stale fixture now satisfies the resolved-family contract, restores process-wide family state, and checks that wrapping intentionally arms the contract lane." + ] + }, + { + "id": "TA-1157", + "scope": "split GLM52 native-FP8 configuration model and buffer checkpoint reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Configuration admission, module replacement, pair-buffer materialization, and checkpoint ownership form one native-FP8 state policy.", + "Canonical router dispatch and frozen expert scoring remain separate behavioral contracts." + ] + }, + { + "id": "TA-1158", + "scope": "split direct and MoEExperts Quack FP8 train-step reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Direct TP-FP8 execution and the module wrapper are layers of one Quack expert train-step policy.", + "TP reduction, grouped backends, bias, activation, gradient, finiteness, and master-update assertions remain intact." + ] + }, + { + "id": "TA-1159", + "scope": "split scoring and trainable SGLang EP dispatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Scoring and trainable execution are branches of one SGLang EP dispatch and admission policy.", + "Slot combine, weight presentation, and live stock-Triton gradient parity remain separate contracts." + ] + }, + { + "id": "TA-1160", + "scope": "split fused RMSNorm kernel and model-integration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Residual and no-residual kernel numerics and their dense-model call sites form one GPU integration contract.", + "CPU fallback and trunk-specific dispatch remain independently reported." + ] + }, + { + "id": "TA-1161", + "scope": "split Qwen3.5 pairwise rotary numerics from attention rotary behavior", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Pairwise-interleaved reference numerics and the attention modules that select half-rotate semantics form one rotary policy.", + "Dense and MoE attention and mRoPE assertions remain intact." + ] + }, + { + "id": "TA-1162", + "scope": "split Qwen3.5 Class-B rotary admission from attention rotary behavior", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Class-B fused admission and fail-closed CPU behavior are configuration branches of the same rotary policy.", + "Fused-call and CUDA-rejection assertions remain intact." + ] + }, + { + "id": "TA-1163", + "scope": "split DeepSeek-V4 SwiGLU clamp from non-hash MoE behavior", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shared and routed SwiGLU clamping are activation branches of the non-hash MoE runtime policy.", + "Hash routing and record/replay remain separate contracts." + ] + }, + { + "id": "TA-1164", + "scope": "split GDN gating and gated RMSNorm numerical reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Gating and gated normalization are the two numerical primitives of one GDN forward and backward contract.", + "Exact-model module dispatch and triangular-solve geometry remain separate." + ] + }, + { + "id": "TA-1165", + "scope": "split native block-FP8 encoding checkpoint and execution admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Byte encoding, protected state, checkpoint validation, execution entry, and fail-closed admission form one native block-FP8 state machine.", + "Every byte, dtype, identity, rollback, hook, range, and error assertion remains intact." + ] + }, + { + "id": "TA-1166", + "scope": "split sparse-MLA backward reference and deterministic-atomic parity reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Reference correctness and deterministic-versus-atomic parity are complementary branches of one backward policy.", + "The H100 and TileLang gates and all gradient and finiteness assertions remain intact." + ] + }, + { + "id": "TA-1167", + "scope": "split RMSNorm v2 realization parity and dispatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forced fused-versus-split bit identity and the heuristic that selects those realizations form one realization and dispatch policy.", + "Reference accuracy and batch and run invariance remain a separate numerical-tree contract." + ] + }, + { + "id": "TA-1168", + "scope": "split Class-B RoPE dtype admission and shape-backward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fail-closed table dtype admission and partial-rotary forward and backward behavior are branches of one Class-B primitive contract.", + "All shape cases, untouched-tail bytes, gradients, and error assertions remain intact." + ] + }, + { + "id": "TA-1169", + "scope": "split Class-B RoPE table-layout report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Unique-frequency-half table construction is the input-layout edge of the same Class-B primitive contract.", + "The exact output shape and value assertions remain intact." + ] + }, + { + "id": "TA-1170", + "scope": "split EP gradient backend and metadata-domain admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Backend declarations and parameter-metadata declarations select the same gradient-reduction domain contract and must both fail closed.", + "The real two-rank reduction report remains separate." + ] + }, + { + "id": "TA-1171", + "scope": "split exact dense and LM-head legacy weight-sync rejection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense, ordinary projection, and LM-head components share one prohibition on legacy merged-weight publication.", + "Separate-factor checkpoint byte preservation remains a distinct successful-publication report." + ] + }, + { + "id": "TA-1172", + "scope": "split Qwen3.5 context-parallel positive and negative subprocess reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ulysses execution and ring-plus-FLA rejection are positive and negative branches of one context-parallel admission contract.", + "Both two-GPU subprocesses and their independent success checks still run." + ] + }, + { + "id": "TA-1173", + "scope": "split QLoRA NVFP4 and prequantized block-FP8 lifecycle reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "NVFP4 and block-FP8 are format branches of the same quantized QLoRA execution and merge lifecycle.", + "All live CUDA loading, forward, backward, scale, merge, and requantization assertions remain intact." + ] + }, + { + "id": "TA-1174", + "scope": "split training-model FP8 construction and quantized-mode admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Successful FP8 construction and rejected incompatible modes are the positive and negative branches of one quantized model-builder policy.", + "Former monkeypatch isolation is preserved with scoped contexts." + ] + }, + { + "id": "TA-1175", + "scope": "split training-model QARL lifecycle report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "QARL construction, calibration order, and admission belong to the same model-builder quantization policy as FP8.", + "Dense-only restrictions, calibration state, and parallelization-order assertions remain intact." + ] + }, + { + "id": "TA-1176", + "scope": "split routing-replay sequence-parallel and ring-attention layout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sequence slicing and ring zigzag placement are topology branches of one context-parallel routing-layout contract.", + "Wire decoding and weight tensor construction remain separate." + ] + }, + { + "id": "TA-1177", + "scope": "split sparse-delta single-file and source-capture reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Single-file encoding is the artifact boundary consumed by source capture and manifest publication.", + "Scoped monkeypatch contexts preserve the former module-stub isolation." + ] + }, + { + "id": "TA-1178", + "scope": "split DeepSeek-V4 checkpoint codec and handler-ownership reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quantized decoding and EP-aware expert fusion form one checkpoint conversion and ownership policy.", + "End-to-end synthetic model loading remains separate." + ] + }, + { + "id": "TA-1179", + "scope": "split active-LoRA composite admission and atomic flag reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Atomic setting and clearing establishes the composite flag state consumed by exact-family admission.", + "All missing-component and scoring-only branches remain intact." + ] + }, + { + "id": "TA-1180", + "scope": "split active-LoRA server derivation and cached-surface activation reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-one server derivation and indexer and MoE activation jointly describe complete-composite propagation.", + "Topology rejection and partial-composite negative controls remain intact." + ] + }, + { + "id": "TA-1181", + "scope": "split Nemotron-H published and stacked expert checkpoint layout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Per-expert published weights and stacked in-memory weights are two accepted input layouts for one bidirectional checkpoint handler.", + "EP ownership remains separate." + ] + }, + { + "id": "TA-1182", + "scope": "split LoRA target-manifest success and failure reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Runtime target selection and schema or coverage rejection are positive and negative branches of one manifest contract.", + "Every schema, rank, count, target, and unlisted-module assertion remains intact." + ] + }, + { + "id": "TA-1183", + "scope": "split Qwen2 model construction and checkpoint-handler reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "HF configuration conversion, unfusing, and bidirectional checkpoint translation form one architecture-support contract.", + "All model-layout and HF parity assertions remain intact." + ] + }, + { + "id": "TA-1184", + "scope": "split OLMo2 model construction and checkpoint-handler reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "HF configuration conversion, post-norm construction, unfusing, and checkpoint translation form one architecture-support contract.", + "All model-layout and HF parity assertions remain intact." + ] + }, + { + "id": "TA-1185", + "scope": "split Quack worker-protocol and PTXAS process-safety reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Worker response framing and PTXAS timeout and temporary-output handling are process-boundary safety checks for one compilation path.", + "All timeout, truncation, uniqueness, cleanup, and entry-selection assertions remain intact." + ] + }, + { + "id": "TA-1186", + "scope": "split Quack cache-key hashing report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Deterministic safe cache identity is the persistence edge of the same Quack compilation process contract.", + "Collision boundaries and rejection of unsafe pickle hooks remain intact." + ] + }, + { + "id": "TA-1187", + "scope": "split shared-prefix matrix and singleton attention reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Multi-member and singleton groups are shape branches of one shared-prefix forward and backward equivalence contract.", + "The optional FA3 capability gate remains unchanged." + ] + }, + { + "id": "TA-1188", + "scope": "split OPD output-edge and hidden-distance reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ignored tokens, per-token output, hidden-only loss, and chunked hidden distance describe one OPD edge-behavior contract.", + "Backend numerics and gradient reduction remain separate." + ] + }, + { + "id": "TA-1189", + "scope": "split packing-strategy admission and correctness reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Strategy validation, oversized handling, token preservation, capacity, and utilization form one generic packing-policy contract.", + "Balanced-DP-specific scheduling remains separate." + ] + }, + { + "id": "TA-1190", + "scope": "split packing-strategy determinism and datum-order report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Determinism and datum order are output invariants of the same generic packing-policy contract.", + "All strategies, repeated builds, and reordered-index assertions remain intact." + ] + }, + { + "id": "TA-1191", + "scope": "split Tinker session OpenAPI and activity-lifecycle reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Schema publication, creation, follow-up use, heartbeat, and canonical configuration form one public session-endpoint lifecycle.", + "All HTTP-boundary and server-state assertions remain intact." + ] + }, + { + "id": "TA-1192", + "scope": "split Mooncake side-payload store and R3 slice reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Tensor roundtrip is the storage primitive used by R3 reference publication, selective loading, and cleanup.", + "Integer, float, missing-key, slice, validation, and cleanup assertions remain intact." + ] + }, + { + "id": "TA-1193", + "scope": "split checkpoint-save failure and live-adapter-state reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fail-closed admission, write failures, and successful dense and MoE factor publication are branches of one adapter-save policy.", + "Scoped monkeypatch contexts preserve former isolation." + ] + }, + { + "id": "TA-1194", + "scope": "split launcher worker discovery and readiness reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Address resolution and readiness or early-exit handling are successive phases of one launcher worker-control lifecycle.", + "Server override parsing remains separate." + ] + }, + { + "id": "TA-1195", + "scope": "split P2P async dispatch and prepare-timeout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Transfer cutoff, status timeout, and prepare-request timeout are environment-controlled branches of one P2P async API policy.", + "All sync-versus-async, timeout, request payload, and transport assertions remain intact." + ] + }, + { + "id": "TA-1196", + "scope": "split K3 debug metrics and logprob-temperature reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Tail metrics and temperature-matched zero K3 jointly define behavior-logprob observability for both loss implementations.", + "TokenPartial reducer identity remains separate." + ] + }, + { + "id": "TA-1197", + "scope": "split Mooncake hidden transport and teacher-consumer reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Tensor codec, metadata publication, retrieval, and indexed teacher consumption form one hidden-transport lifecycle.", + "Rank-two, rank-three, and multi-teacher cases remain intact." + ] + }, + { + "id": "TA-1198", + "scope": "split Mooncake metadata admission and store-lifecycle reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Malformed or legacy metadata rejection, object removal, and configuration precedence form one store-admission lifecycle.", + "All missing-key, size, schema, cleanup, and environment assertions remain intact." + ] + }, + { + "id": "TA-1199", + "scope": "split dense QARL weight fake-quant and injection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Export-equivalent fake quantization and model injection are numerical and structural phases of one dense QARL policy.", + "Target selection, summary, forward counts, and model admission remain intact." + ] + }, + { + "id": "TA-1200", + "scope": "split dense QARL configuration-normalization report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Recipe normalization and rejection establish the configuration admitted by the same dense fake-quant lifecycle.", + "Static, unsupported-format, and invalid-block failures remain intact." + ] + }, + { + "id": "TA-1201", + "scope": "split NVFP4 QARL MoE conversion and eager-execution reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Identity-preserving class conversion and eager forward and backward behavior are consecutive phases of one expert lifecycle.", + "Parameter identity, quantization effect, gradients, passthrough, and admission remain intact." + ] + }, + { + "id": "TA-1202", + "scope": "split NVFP4 QARL MoE injection report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Model injection selects and invokes the same expert conversion lifecycle.", + "NVFP4 admission, FP8 rejection, and independent target selection remain intact." + ] + }, + { + "id": "TA-1203", + "scope": "split QARL weight-sync success and bad-configuration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Derived FP8 sync configuration and explicit incompatible configuration are positive and negative branches of one QARL sync policy.", + "Folded and excluded module behavior remains intact." + ] + }, + { + "id": "TA-1204", + "scope": "split expert-adapter capability ownership and semantics reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Backend capability, factor ownership, and preserved activation semantics jointly define the generic expert-adapter contract.", + "Injection and model-family construction remain separate." + ] + }, + { + "id": "TA-1205", + "scope": "split teacher-head storage and manager-residency reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Weight discovery, sharded storage, cross-shard views, residency, dtype reload, and prefetch form one teacher-head lifecycle.", + "Teacher activation caching remains separate." + ] + }, + { + "id": "TA-1206", + "scope": "split exact-server trunk and numerical-family selection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Trunk wrapping and v2 family selection are coordinated pre-parallelization model-program choices.", + "Scoped monkeypatch contexts and the autouse global-state reset preserve isolation." + ] + }, + { + "id": "TA-1207", + "scope": "split NVFP4 quantization roundtrip and directory-export reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed codec correctness is the primitive exercised by end-to-end NVFP4 directory export.", + "Shared scales, BF16 islands, W4A4 inputs, requantization rejection, and reconstruction error remain intact." + ] + }, + { + "id": "TA-1208", + "scope": "split FP8 export CLI and model-layout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CLI configuration, output directory construction, and architecture-specific layout transforms form one export command contract.", + "Every transform runs in its own named temporary case." + ] + }, + { + "id": "TA-1209", + "scope": "split FP8 export admission report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Source preflight and QARL fold rejection are fail-closed branches of the same export command contract.", + "Primitive quantization and trained-logprob preservation remain separate." + ] + }, + { + "id": "TA-1210", + "scope": "split OPD endpoint-registration and student-version verifier reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Existing endpoint identity and expected student weight version jointly establish OPD endpoint admission.", + "Matching, mismatch, and endpoint-error branches remain intact." + ] + }, + { + "id": "TA-1211", + "scope": "split OPD prepare-worker and payload-transport reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Chunk preparation, queue completion, causal shifting, cache-index alignment, and Mooncake metadata form one pipeline-preparation contract.", + "All queue and transport assertions remain intact." + ] + }, + { + "id": "TA-1212", + "scope": "split DeepSeek-V3 trainer and parallelizer admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Router-freeze construction and downstream tensor-parallel rejection are successive admission layers of one DeepSeek training policy.", + "The successful router-freeze path and all incompatible configurations remain intact." + ] + }, + { + "id": "TA-1213", + "scope": "split Class-B RoPE selection and canonical GLM numerical-program reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Class-B selection is one required component of the canonical GLM-5.2 numerical program and its fail-closed overrides.", + "Exact Qwen3.5 program admission remains separate." + ] + }, + { + "id": "TA-1214", + "scope": "split simulator topology-ledger and config-metadata admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Topology resolution and analytical ledgers are derived from the same trusted training configuration and model metadata boundary.", + "Observed benchmarking, calibration, and kernel ranking remain separate." + ] + }, + { + "id": "TA-1215", + "scope": "split DeepEP internode topology and preflight reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Node-spanning topology detection directly controls whether the internode transport preflight runs.", + "Scoped patches preserve every skip, intranode, failure-diagnostic, identity, and corruption branch." + ] + }, + { + "id": "TA-1216", + "scope": "split DeepEP buffer-size and RDMA admission report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "NVL alignment and RDMA byte validation are resource-admission edges of the same internode transport policy.", + "The int32 limit and all byte-layout cases remain intact." + ] + }, + { + "id": "TA-1217", + "scope": "split generic parallel-plan meta slicing and gradient-domain reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Parameter slicing, placement, and explicit replicated-gradient metadata are outputs of one generic EP plan application.", + "Shape, dtype, requires-grad, divisibility, and reduction assertions remain intact." + ] + }, + { + "id": "TA-1218", + "scope": "split exact GLM meta and materialized already-local EP plan reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Meta and materialized tensors are allocation branches of the same exact routed-expert disposition policy.", + "Already-local bases, force-sharded banks, replicated factors, and malformed singleton guards remain intact." + ] + }, + { + "id": "TA-1219", + "scope": "split pipeline FQN partition and stage-placement reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Module partitioning and rank ownership are consecutive outputs of one pipeline layout plan.", + "Single, loop, virtual, v-style, pinned, weighted, and infeasible cases remain intact." + ] + }, + { + "id": "TA-1220", + "scope": "split pipeline schedule metadata and admission report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Schedule style and microbatch admission consume the same stage layout plan.", + "Every supported schedule and invalid virtual-stage or microbatch case remains intact." + ] + }, + { + "id": "TA-1221", + "scope": "split PP profiler interval-union and P2P-byte accounting reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Busy interval union and P2P bytes are pure accounting primitives feeding one bubble-profile report.", + "Profiler patch lifecycle and live CUDA schedule execution remain separate." + ] + }, + { + "id": "TA-1222", + "scope": "split Muon full-gradient oracle and shard-local negative-control reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full-gradient parity and shard-local divergence are positive and negative controls of one distributed Muon policy.", + "All four two-GPU subprocesses across dense and MoE layouts still run." + ] + }, + { + "id": "TA-1223", + "scope": "split DeepSeek-V4 C128 and C4 compression admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "C128 ratio divisibility and C4 overlap divisibility are compression-regime branches of one context-parallel policy.", + "Output shape, cache capacity, and overlap failure assertions remain intact." + ] + }, + { + "id": "TA-1224", + "scope": "split BI fused LM-head integration and kernel edge reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Loss forward and backward parity and unit-temperature or near-one probability edges exercise one selected-logprob kernel contract.", + "Determinism, batch invariance, guards, identity bits, and nonpositive logprobs remain intact." + ] + }, + { + "id": "TA-1225", + "scope": "split BI full and dimension mean reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Full reduction and explicit dimension reductions are dispatch branches of one batch-invariant mean policy.", + "FP32 and BF16 accuracy and bitwise dimension behavior remain intact." + ] + }, + { + "id": "TA-1226", + "scope": "split BI head-v2 projection and trainability reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Projection and statistics bits, batch invariance, fused CE, gradients, and rollback form one head-v2 lifecycle.", + "The live CUDA capability gate remains unchanged." + ] + }, + { + "id": "TA-1227", + "scope": "split GDN primitive numerics and exact-model module-dispatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Gating and gated RMSNorm numerics are the kernel behavior selected by the exact-model GDN dispatch program.", + "Triangular-solve geometry remains separate." + ] + }, + { + "id": "TA-1228", + "scope": "split Quack EP Triton parity and half-concat reference reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Backend parity and the independent half-concatenated activation reference jointly establish one Quack EP numerical contract.", + "CPU gradient-arity behavior remains separate." + ] + }, + { + "id": "TA-1229", + "scope": "split runner batch conversion and sequence-shard reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Float side-channel conversion, ragged padding, and sequence sharding form one runner batch-materialization policy.", + "Teacher hidden-state dtype, padding, shape, and shard assertions remain intact." + ] + }, + { + "id": "TA-1230", + "scope": "split exact dense and routed adapter-gradient ownership reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense managed-factor compilation and routed fail-closed ownership are topology branches of one exact adapter-gradient policy.", + "Scoped patches preserve manager and parallel-state isolation." + ] + }, + { + "id": "TA-1231", + "scope": "split DeepSeek-V4 checkpoint codec-handler and synthetic-load reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Codec and EP-handler behavior culminate in the same synthetic end-to-end checkpoint load contract.", + "FP8, MXFP4, window, C4, hash, strict-buffer, and ownership cases remain intact." + ] + }, + { + "id": "TA-1232", + "scope": "split DeepSeek-V4 construction and runtime reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Construction, topology, precision preservation, forward, backward, hash routing, and recomputation form one model contract.", + "A scoped patch restores the former construction-test isolation before runtime execution." + ] + }, + { + "id": "TA-1233", + "scope": "split exact attention checkpoint inventory and pair-state reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Source inventory selects the native absorbed-kv pair state machine exercised by order, byte, completion, dtype, and shape checks.", + "All exact factor inventory and handler assertions remain intact." + ] + }, + { + "id": "TA-1234", + "scope": "split canonical MoE routed and shared boundary reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Routed and shared partials are the two contributor boundaries of one canonical MoE block policy.", + "Global and local IDs, scale, root invocation, and contributor ordinal remain intact." + ] + }, + { + "id": "TA-1235", + "scope": "split canonical native-FP8 router and expert-state reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Serving-compatible routing and frozen scoring-only expert execution form one canonical native-FP8 runtime policy.", + "Configuration, buffer, and checkpoint ownership remain separate." + ] + }, + { + "id": "TA-1236", + "scope": "split Nemotron-H EP ownership and checkpoint-layout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP admission and slicing govern the same published and stacked layouts exercised by the bidirectional handler.", + "HF parity, exact saved bytes, skips, and plan targeting remain intact." + ] + }, + { + "id": "TA-1237", + "scope": "split fused RMSNorm model integration and trunk reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ordinary model integration and trunk-specific no-residual dispatch are branches of one GPU fused RMSNorm policy.", + "CPU fallback remains a separate capability domain." + ] + }, + { + "id": "TA-1238", + "scope": "split RoPE registry precision and native-cache reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Registry FP32 recipes and lazy device-local cache materialization form one CPU RoPE precision and cache policy.", + "Exact architecture serving-device execution remains a separate GPU report." + ] + }, + { + "id": "TA-1239", + "scope": "split adapter-gradient pre-rendezvous and ModelRunner tail-failure reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both subprocesses exercise bounded fail-closed behavior before a capture can commit.", + "The independent two-rank subprocesses and their original assertions remain intact." + ] + }, + { + "id": "TA-1240", + "scope": "split adapter-gradient publication-commit failure report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Publication commit is the final CPU failure phase of the same bounded fatal-gradient lifecycle.", + "The asymmetric post-mutation GPU boundary remains a separate capability report." + ] + }, + { + "id": "TA-1241", + "scope": "split live two-rank clip and three-rank participation reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Norm clipping, nonfinite admission, and all-rank participation form one live distributed clipping policy.", + "Both independent subprocess topologies still execute." + ] + }, + { + "id": "TA-1242", + "scope": "split FutureStore creation-processing and model-expiration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Creation, processing, concurrency, model operations, expiration, and cleanup form one async store lifecycle.", + "A fresh identical store instance preserves the former fixture isolation between phases." + ] + }, + { + "id": "TA-1243", + "scope": "split orchestrator communication roundtrip and edge-lifecycle reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Roundtrip, interleaving, exceptions, and shutdown are phases of one client communication lifecycle.", + "All async assertions and engine interactions remain intact." + ] + }, + { + "id": "TA-1244", + "scope": "split sparse-delta initialization and runtime-helper reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Post-only initialization, prepacked admission, and runtime delta loading govern one backend policy.", + "Encoding and load assertions remain intact." + ] + }, + { + "id": "TA-1245", + "scope": "split FP8 LM-head CE selection and loss-dispatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Per-token LM-head selection and dispatcher routing are consecutive layers of one FP8 loss policy.", + "Numerical selection and dispatch assertions remain intact." + ] + }, + { + "id": "TA-1246", + "scope": "split sequence-shard core and token-side-channel reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sharding, padding, flash-attention metadata, and per-token side channels form one collator materialization policy.", + "The original patched parallel-state contexts remain intact." + ] + }, + { + "id": "TA-1247", + "scope": "split server-CLI sequence-boundary and shard-preservation reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Server-CLI boundary parity culminates in preservation and regeneration by the sequence-shard collator.", + "Dtype, padding, stale-metadata, and original-position assertions remain intact." + ] + }, + { + "id": "TA-1248", + "scope": "split sequence LCM-padding and collator-divisibility report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "LCM padding is the topology branch of the same sequence-metadata alignment policy.", + "RequestProcessor and post-shard divisibility assertions remain intact." + ] + }, + { + "id": "TA-1249", + "scope": "split padded-unpacking boundary report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Padded output unpacking is the consumer boundary of the same packed sequence policy.", + "Padded and unpadded sample-count assertions remain intact." + ] + }, + { + "id": "TA-1250", + "scope": "split AnyPrecision AdamW cautious execution and state-strategy reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Cautious decay, chunked denominators, state reuse, and DTensor offload form one optimizer lifecycle.", + "The temporary-path fixture and all numerical and state assertions remain intact." + ] + }, + { + "id": "TA-1251", + "scope": "split inference-endpoint registration and list-health reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Registration and subsequent health-aware listing form one public endpoint lifecycle.", + "Explicit worker, auto-sync, FP8 KV-cache, and v1-model fallback assertions remain intact." + ] + }, + { + "id": "TA-1252", + "scope": "split P2P prepare and initialize-fanout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Prepare payloads, cached maps, endpoint fanout, and fanout cleanup are phases of one initialization handshake.", + "Scoped monkeypatch contexts preserve environment isolation between initialization modes." + ] + }, + { + "id": "TA-1253", + "scope": "split P2P complete-sync report from initialization handshake", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Receiver completion, tied aliases, cache flushing, version forwarding, and cleanup close the same P2P lifecycle.", + "All completion and failure assertions remain intact." + ] + }, + { + "id": "TA-1254", + "scope": "split P2P receiver placement and source-staging reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Receiver placement, replicated-source reuse, scratch alignment, coalescing, and registration form one staging policy.", + "FP8 receiver layouts remain a separate format contract." + ] + }, + { + "id": "TA-1255", + "scope": "split P2P invalid-manifest and transfer-failure reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Source, receiver, shape, rank, and name admission failures belong with runtime transfer diagnostics.", + "Diagnostic detail, sample caps, and disabled-sampling assertions remain intact." + ] + }, + { + "id": "TA-1256", + "scope": "split microbatch splitting and DataLoaderBuilder configuration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Builder batch sizing and collator selection culminate in microbatch splitting and epoch delegation.", + "Sampler, sequence-parallel, custom-collator, edge, and set-epoch assertions remain intact." + ] + }, + { + "id": "TA-1257", + "scope": "split dataset expansion-type and split-merge reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dataset expansion and type selection feed the same composition policy as train-validation splitting and merging.", + "Raw loading and preprocessed persistence remain separate I/O boundaries." + ] + }, + { + "id": "TA-1258", + "scope": "split MiniMax-M3 configuration-registration and text-runtime reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Configuration conversion and registry admission select the text-only model runtime exercised by forward, backward, and rejection checks.", + "Checkpoint ownership and MSA paging remain separate mechanisms." + ] + }, + { + "id": "TA-1259", + "scope": "split attention backend resolution and eager-layout reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CPU backend selection now executes the eager numerical head-layout contract it resolves.", + "FlashAttention and SGL page-cache paths retain their independent optional capability gates." + ] + }, + { + "id": "TA-1260", + "scope": "split FP8 training and other low-precision argument reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FP8 training, block-FP8 QLoRA, QARL, aliases, defaults, and incompatible combinations form one low-precision parsing policy.", + "Fresh temporary subdirectories preserve the former configuration-file isolation." + ] + }, + { + "id": "TA-1261", + "scope": "split direct-EP multi-sender initialization and scatter-copy reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Scatter-copy ownership is an initialization branch of the same direct-EP multi-sender lifecycle.", + "List, deep, and locator-reuse modes remain intact under a scoped environment." + ] + }, + { + "id": "TA-1262", + "scope": "split direct-EP dense-sharding manifest report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense and expert manifest partitioning determines the buffers published by the same multi-sender lifecycle.", + "QKV and gate-up ownership affinity assertions remain intact." + ] + }, + { + "id": "TA-1263", + "scope": "split direct-EP rank-filter transfer report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-filtered receiver transfers are the execution boundary of the direct-EP manifest policy.", + "Owning-rank and empty-rank transfer assertions remain intact." + ] + }, + { + "id": "TA-1264", + "scope": "split FP8Linear CUDA matmul and train-step reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Padded recipes and residual correction now culminate in live FP8 forward, backward, and master-weight mutation.", + "Both phases retain the same CUDA capability gate and every numerical assertion." + ] + }, + { + "id": "TA-1265", + "scope": "split dense and packed SSD recurrence reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense, multichunk, boundary-crossing packed, convolution, and mixer behavior form one CPU recurrence contract.", + "Unavailable-kernel admission and live GPU kernel parity remain separate capability reports." + ] + }, + { + "id": "TA-1266", + "scope": "split adapter checkpoint materialization and restore-admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Coordinator materialization and auto-load are entry points into the same checkpoint restore and admission lifecycle.", + "A fresh temporary subtree preserves manager and optimizer isolation." + ] + }, + { + "id": "TA-1267", + "scope": "split OPD loss execution and metric aggregation reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Per-microbatch OPD execution now culminates in the aggregation, extrema, reduction, and empty-rank policy consuming its metrics.", + "Scoped patches preserve device and parallel-state isolation." + ] + }, + { + "id": "TA-1268", + "scope": "split request-processor forward and packed-row batching reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed-row batching is an OPD branch of the same live forward-backward processor lifecycle.", + "The helper uses a fresh internal processor while retaining the shared-processor rejection case." + ] + }, + { + "id": "TA-1269", + "scope": "split model-scoped and sampler-scoped checkpoint listing reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Listing, deletion, resolution, and isolation now cover both storage namespaces in one public checkpoint policy.", + "Explicit setup and teardown give the sampler namespace a fresh APIServer and temporary root." + ] + }, + { + "id": "TA-1270", + "scope": "split sampler adapter tracking and normalized export reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Adapter reconciliation, model-scoped tracking, and adapter-only export form one sampler-weight lifecycle.", + "The normalized session-spec request and output URI assertions remain intact." + ] + }, + { + "id": "TA-1271", + "scope": "split Muon Gram-Newton-Schulz configuration and Quack backend reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quack import, dispatch, tuned-mode, and dtype admission are backend branches of the configured Gram-Newton-Schulz optimizer.", + "Nested monkeypatch contexts preserve cache and import isolation." + ] + }, + { + "id": "TA-1272", + "scope": "split FP8 MoE expert and injected full-model train-step reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Direct expert variants now culminate in injected dense-expert-dense forward, backward, and master-weight mutation.", + "Both phases retain the same CUDA capability gate." + ] + }, + { + "id": "TA-1273", + "scope": "split canonical LoRA fold and LoraLinear merged-forward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical folding and straight-through gradients now execute through LoraLinear selection, ordinary isolation, and cache invalidation.", + "All exact and legacy forward assertions remain intact." + ] + }, + { + "id": "TA-1274", + "scope": "split canonical LoRA fold and MoE merged-weight cache reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Expert gate-up and down folding, weight-sync views, versioned caches, and fused admission are the MoE branch of the same fold policy.", + "Native EP execution and trunk wrapping remain separate integrations." + ] + }, + { + "id": "TA-1275", + "scope": "split GDN delta-linear product and merged-forward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The explicit low-rank product now feeds sliced canonical folding, gradient ownership, GDN projection, and bounded cache behavior.", + "Geometry-manifest and checkpoint roundtrip remain separate boundaries." + ] + }, + { + "id": "TA-1276", + "scope": "split exact LM-head per-token and causal-loss routing reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Per-token exact-head selection and causal-loss TP-group admission are consecutive layers of one loss-routing contract.", + "A scoped patch preserves each dispatcher replacement." + ] + }, + { + "id": "TA-1277", + "scope": "split exact LM-head weight and server selector report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Weight and module selection prove that the exact loss route never materializes a merged delta.", + "Identity assertions remain intact in the complete exact-head policy." + ] + }, + { + "id": "TA-1278", + "scope": "split exact LM-head FSDP replicated-factor report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FSDP replicated-factor admission is the parallel ownership boundary of the same exact-head loss policy.", + "The lora-A-only fail-closed assertion remains intact." + ] + }, + { + "id": "TA-1279", + "scope": "split absorbed-KV contract and dtype-move state reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Native base bytes, FP32 logical masters, identities, and dtype moves form one CPU state contract.", + "The official CUDA Q/V program remains a separate capability report." + ] + }, + { + "id": "TA-1280", + "scope": "split absorbed-KV state and direct-projection admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Direct projection, branch, materialization, and factor-dtype rejection define the fail-closed edge of the same CPU component contract.", + "Every negative assertion remains intact." + ] + }, + { + "id": "TA-1281", + "scope": "split canonical MoE graph metadata and transport-admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Capacity metadata and transport resolution jointly define the canonical MoE planning boundary.", + "Dense, packed, CP-sharded, graph, and output-distribution admission cases remain intact." + ] + }, + { + "id": "TA-1282", + "scope": "split canonical MoE metadata and exact parallel-plan reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Trainer and sampler plan hashes, group layouts, topology rejection, and metadata now form one CPU planning policy.", + "The distributed reduction and backward subprocess remains separate." + ] + }, + { + "id": "TA-1283", + "scope": "split BI router GEMM and leading-dimension linear reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The leading-dimension wrapper is a shape branch of the same BF16-input FP32-output router GEMM contract.", + "Forward and backward comparisons remain intact." + ] + }, + { + "id": "TA-1284", + "scope": "split BI router GEMM and top-k weight reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FP32 logits and top-k normalization/casting are consecutive stages of exact router selection.", + "Renormalized, cast-only, and dtype-rejection cases remain intact." + ] + }, + { + "id": "TA-1285", + "scope": "split BI router primitives and MoEBlock dispatch report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "MoEBlock exact and ordinary routing now exercise the complete GEMM and top-k primitive policy under the same CUDA gate.", + "Batch-composition and stock-path assertions remain intact." + ] + }, + { + "id": "TA-1286", + "scope": "split NCCL endpoint transfer and flattened or hybrid bucket reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Endpoint initialization, two-phase completion, and flattened, chunked, and hybrid broadcasts now form one NCCL transfer policy.", + "Scoped environment contexts retain independent load-format setup and receiver-fence assertions." + ] + }, + { + "id": "TA-1287", + "scope": "split NCCL transfer and multi-rank direct-format admission reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Multi-rank direct-load rejection is the admission edge of the same endpoint transfer policy.", + "The health-check lifecycle remains a separate report." + ] + }, + { + "id": "TA-1288", + "scope": "split weight-sync source and bucket-sizing reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Adapter selection, parameter extraction, inference layout, and byte-capped chunking jointly define the sync-source preparation policy.", + "Default, shared override, MoE override, split, and oversize-item assertions remain intact." + ] + }, + { + "id": "TA-1289", + "scope": "split weight-sync source and direct-EP sender-selection reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Direct-EP sender mapping and collection gating determine which prepared source tensors enter the transfer.", + "Scoped environment contexts isolate default and round-robin replica strategies." + ] + }, + { + "id": "TA-1290", + "scope": "split shipped adapter examples and quantized server-configuration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shipped MoE LoRA and QLoRA parsing now culminates the low-precision server-configuration contract.", + "Clean-process parsing and every certified Quack target assertion remain intact." + ] + }, + { + "id": "TA-1291", + "scope": "split sharded-adapter layout and deterministic-initialization reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local packing, empty-shard ownership, and topology-invariant deterministic initialization now form one CPU adapter-state policy.", + "A fresh temporary subtree preserves manager and checkpoint isolation." + ] + }, + { + "id": "TA-1292", + "scope": "split sharded-adapter state and explicit-EP layout-discovery reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit expert sharding and generic replication are layout-discovery branches of the same CPU adapter-state policy.", + "The real two-rank Gloo DTensor report remains separate." + ] + }, + { + "id": "TA-1293", + "scope": "split dispatcher save and session-registration reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Registration, cross-rank failure, nonresident auto-load, and state or adapter saves now form one session checkpoint lifecycle.", + "A scoped patch preserves the rank-zero failure boundary." + ] + }, + { + "id": "TA-1294", + "scope": "split optimizer publication and gradient-epoch completion reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forward-backward completion, abort, uniform rejection, optimizer publication, and fatal tail failures now form one mutation lifecycle.", + "Commit ordering, poisoning, and process-termination assertions remain intact." + ] + }, + { + "id": "TA-1295", + "scope": "split token diagnostics and hidden-component hook reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense and MoE component hooks now feed the same diagnostic computation and summary policy they support.", + "Equation-term tensors, hook cleanup, selection, loss cross-checks, and CP mapping remain asserted." + ] + }, + { + "id": "TA-1296", + "scope": "split diagnostic tensor-dump and trusted-override reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Ranked tensor persistence and trusted diagnostic replay now form one artifact-boundary policy.", + "A fresh subtree preserves file isolation and the missing-root rejection remains intact." + ] + }, + { + "id": "TA-1297", + "scope": "split P2P engine construction and initialization reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Hostname precedence and Mooncake fallback construction now precede the prepare, fanout, cache, completion, and cleanup lifecycle.", + "A scoped environment context preserves every resolution branch." + ] + }, + { + "id": "TA-1298", + "scope": "split adapter weight-publication and authoritative optimizer reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Clean mid-epoch weight admission and strict checkpoint rejection are publication branches of the authoritative optimizer lifecycle.", + "A fresh manager subtree preserves state isolation." + ] + }, + { + "id": "TA-1299", + "scope": "split packed side-metadata and full pipeline reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Teacher, cache, weight, RL, and hidden-state metadata now flow through the same full pack, forward, and unpack pipeline.", + "Capacity and packing-disabled policies remain separate." + ] + }, + { + "id": "TA-1300", + "scope": "split model-pass R3 payload and token-diagnostic unpacking reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed token diagnostics are now unpacked in the model-pass side-payload lifecycle they support.", + "Position rebasing and every aligned diagnostic field remain asserted." + ] + }, + { + "id": "TA-1301", + "scope": "split API optimizer and forward response-metric reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Optimizer, forward, and forward-backward response shaping now form one public training-operation response policy.", + "Each phase constructs a fresh API server." + ] + }, + { + "id": "TA-1302", + "scope": "split runner gradient compiler and staged-capture abort reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forward-backward abort now closes the failure edge of the runner gradient-ownership compilation policy.", + "The abort phase uses a fresh adapter-manager path." + ] + }, + { + "id": "TA-1303", + "scope": "split OPD cache shaping and loss-execution reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed cache rows, teacher boundaries, last-k weights, loss, gradients, and profiling now form one OPD execution policy.", + "Distributed cache assembly and debug artifacts remain separate capabilities." + ] + }, + { + "id": "TA-1304", + "scope": "split dispatcher batch distribution and packing-dummy reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Balanced packing and sequential dummy behavior are input-construction branches of dispatcher sharding and provenance.", + "A scoped parallel-state patch preserves rank isolation." + ] + }, + { + "id": "TA-1305", + "scope": "split routing payload and microbatch diagnostic artifact reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Routing references and diagnostic dumps now form one dispatcher side-payload artifact and security policy.", + "Raw manifests, Mooncake slices, legacy-pickle rejection, symlink rejection, and R3 dump contents remain intact." + ] + }, + { + "id": "TA-1306", + "scope": "split orchestrator success and error or concurrency reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Initialization, operations, errors, end-to-end completion, concurrent requests, statistics, and shutdown now form one lifecycle.", + "The same live fixture carries the successful and failure phases without restarting the capability." + ] + }, + { + "id": "TA-1307", + "scope": "split cautious primitive and SignSGD report from optimizer construction", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Primitive masking and SignSGD execution now begin the cautious weight-decay optimizer policy.", + "Zero decay, ordinary decay, aligned, misaligned, and zero-direction assertions remain intact." + ] + }, + { + "id": "TA-1308", + "scope": "split AnyPrecisionAdamW cautious report from optimizer construction", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "AnyPrecisionAdamW numerics, chunked state, gradient reuse, and DTensor offload now run inside the complete cautious optimizer policy.", + "The temporary path remains isolated within the parent report." + ] + }, + { + "id": "TA-1309", + "scope": "split Muon cautious report from optimizer construction", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Muon and its AdamW fallback now exercise the same cautious feature before builder routing and kwarg admission.", + "Post-Newton-Schulz masking and ordinary-decay equivalence remain asserted." + ] + }, + { + "id": "TA-1310", + "scope": "split synthetic balanced TopK routing report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Synthetic balanced routing is now an environment-selected branch of the complete TopK router contract.", + "A scoped environment context prevents the synthetic mode from leaking into ordinary routing cases." + ] + }, + { + "id": "TA-1311", + "scope": "split sqrtsoftplus noaux TopK routing report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Biased selection, unbiased weights, normalization, and missing-bias rejection now run with the router configuration policy.", + "All DSv4-specific assertions remain intact." + ] + }, + { + "id": "TA-1312", + "scope": "split hash-table TopK routing report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Token-to-expert lookup, bias independence, input admission, softmax, scaling, and from-config behavior now form one router contract.", + "Each routing algorithm still uses independent tensors and router instances." + ] + }, + { + "id": "TA-1313", + "scope": "split optional boolean coercion from parallel policy configuration", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Boolean admission now accompanies mixed-precision and reduce-dtype configuration in one CPU parallel-policy report.", + "Sequence-parallel folding and manual prefetch remain separate topology policies." + ] + }, + { + "id": "TA-1314", + "scope": "split FP8 module injection and CPU fallback reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Injected module identity, recipes, exclusions, CPU execution, output dtype, and fail-fast fallback now form one CPU policy.", + "The CPU profiler remains a separate observability capability." + ] + }, + { + "id": "TA-1315", + "scope": "split CUDA FP8 profiler and live train-step reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CUDA operand-error breakdown now culminates the live FP8 matmul, correction, backward, and master-weight mutation policy.", + "Both phases retain the same CUDA capability gate and scoped profiler environment." + ] + }, + { + "id": "TA-1316", + "scope": "split GLM52 canonical MoE configuration and sparse selector reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical routing selection, transport, configuration rejection, codecs, and sparse selector pipeline now form one selection policy.", + "Layer-plan allocation and semantic logprob parity remain separate contracts." + ] + }, + { + "id": "TA-1317", + "scope": "split exact dense MLP root-state and forward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical factor ownership and immutable checkpoint-source binding now precede exact fused gate-up, activation, and down execution.", + "All unique-path and state-dictionary assertions remain intact." + ] + }, + { + "id": "TA-1318", + "scope": "split exact dense MLP runtime admission and forward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Rank-alpha construction, atomic runtime updates, and pre-forward consistency rejection now bound the exact MLP execution policy.", + "Every fail-closed assertion remains intact." + ] + }, + { + "id": "TA-1319", + "scope": "split exact dense MLP checkpoint roundtrip and forward reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "XoRL load and PEFT export now close the same six-factor exact MLP component lifecycle.", + "The roundtrip uses a fresh pytest temporary directory." + ] + }, + { + "id": "TA-1320", + "scope": "split sample metadata from packing-dataset lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Position metadata and trainable-token filtering now begin the PackingDataset construction policy.", + "Single, batched, missing-field, and rejection assertions remain intact." + ] + }, + { + "id": "TA-1321", + "scope": "split dataset preprocessing from packing-dataset lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Record filtering and optional evaluation preprocessing now feed the dataset packing, allocation, cache, and missing-column behavior they support.", + "All preprocessing assertions execute before PackingDataset construction." + ] + }, + { + "id": "TA-1322", + "scope": "split pipeline-profiler patching from interval and P2P accounting", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Instance patching, restoration, double-patch admission, interval merging, and P2P byte estimation now form one CPU profiler policy.", + "The live CUDA GPipe step remains a separate capability report." + ] + }, + { + "id": "TA-1323", + "scope": "split selected QLoRA shard-cache loading from deferred loader", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Requested-key-only shard reads now begin the bounded deferred prequantized-loader lifecycle.", + "A scoped monkeypatch context preserves fake shard and cache isolation." + ] + }, + { + "id": "TA-1324", + "scope": "split prequantized QLoRA key planning from deferred loader", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Merged projections, EP16 expert slices, missing-pair rejection, per-module loads, and cache release now form one loader policy.", + "Every exact-key and peak-residency assertion remains intact." + ] + }, + { + "id": "TA-1325", + "scope": "split prequantized checkpoint detection from handler policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "NVFP4 and block-FP8 detection now precede the Qwen checkpoint-handler behavior selected by those formats.", + "A dedicated temporary subtree preserves all file-format cases." + ] + }, + { + "id": "TA-1326", + "scope": "split prequantized skip and load behavior from exclusion policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Quantized-key skipping, QKV and bias merges, exclusion parsing, dense handling, and MoE handling now form one checkpoint policy.", + "Normal and prequantized paths retain independent handler instances." + ] + }, + { + "id": "TA-1327", + "scope": "split fused selected-logprob primitive and loss dispatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Primitive forward and backward parity now flow through per-token CE, causal LM, Quack, and importance-sampling dispatch.", + "The full-logits memory-bound regression remains a separate performance contract." + ] + }, + { + "id": "TA-1328", + "scope": "split streaming forward-KL primitive and OPD dispatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense parity, chunking, masking, low-memory execution, compiled OPD dispatch, and clamp rejection now form one forward-KL execution policy.", + "The fp64 autograd gradcheck remains an independent numerical guard." + ] + }, + { + "id": "TA-1329", + "scope": "split DistSign reduce-scatter from optimizer lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "SP summation, post-sum sign, forced SUM reduction, builder selection, parameter grouping, and update numerics now form one optimizer policy.", + "A scoped distributed patch preserves communication isolation." + ] + }, + { + "id": "TA-1330", + "scope": "split DistSign hook configuration from optimizer lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Local and FSDP-managed hook ownership plus unsupported-topology admission now run before optimizer construction and stepping.", + "Each fake parallel topology remains independently asserted." + ] + }, + { + "id": "TA-1331", + "scope": "split batch-invariant global interpose and trunk-linear reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Wrapped trunk forward and backward parity now culminate in global-interpose training rejection and no-grad admission under the same CUDA gate.", + "The trunk contract is explicitly reset before the global-interpose phase." + ] + }, + { + "id": "TA-1332", + "scope": "split EP shared-replica classification from clipping policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Shared-replica ownership classification now begins the complete CPU EP clipping policy.", + "The live multi-rank reduction remains separate." + ] + }, + { + "id": "TA-1333", + "scope": "split EP norm modes and empty-gradient report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Inf norm, empty groups, missing gradients, and mixed-mesh clipping now form one local clipping policy.", + "Every norm and scaling assertion remains intact." + ] + }, + { + "id": "TA-1334", + "scope": "split skip-FSDP classification and local clipping report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Skip-FSDP expert classification, raw local norms, and uniform clipping now execute in the complete EP policy.", + "The no-reduction and no-division assertions remain intact." + ] + }, + { + "id": "TA-1335", + "scope": "split clip-grad dispatch report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP-aware and ordinary FSDP dispatch now precede mixed-mesh foreach behavior in one CPU report.", + "The world-one Gloo fixture still backs real DTensor mesh handling." + ] + }, + { + "id": "TA-1336", + "scope": "split MoE histogram and index kernels from gather-scatter report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Histogram and expert-slot index construction now feed scatter, gather, add-gather, and roundtrip execution.", + "All kernels retain the same CUDA capability gate." + ] + }, + { + "id": "TA-1337", + "scope": "split deterministic MoE scatter report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Stable ordering, stock coverage, dtype handling, and escape-hatch routing now run inside the complete MoE kernel policy.", + "A scoped environment context prevents the injected failure path from leaking." + ] + }, + { + "id": "TA-1338", + "scope": "split BI GEMM row invariance from table neutrality report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Table-entry bit neutrality now culminates in row invariance across M sizes and bucket boundaries.", + "Optional DeepGEMM parity remains an independent capability report." + ] + }, + { + "id": "TA-1339", + "scope": "split QLoRA injection from quantized execution lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Module injection and checkpoint-format selection now precede NVFP4 and block-FP8 execution, loading, and merging.", + "All phases retain the same CUDA gate." + ] + }, + { + "id": "TA-1340", + "scope": "split QLoRA optimizer reset from quantized execution lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "LoRA optimizer-state clearing and interval-triggered merge integration now close the QLoRA lifecycle.", + "Non-LoRA state preservation remains asserted." + ] + }, + { + "id": "TA-1341", + "scope": "split exact base DCP key contract and official-state load reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Projected dense and scale aliases now feed an official base-DCP load into exact runtime state.", + "A fresh checkpoint subtree and scoped load configuration preserve isolation." + ] + }, + { + "id": "TA-1342", + "scope": "split exact shared-expert native base views from construction report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Construction, runtime admission, logical factors, checkpoint sources, and native TP16 base slices now form one CPU component policy.", + "Optional SGLang factor slicing and Hopper execution remain separate." + ] + }, + { + "id": "TA-1343", + "scope": "split MoE LoRA initialization from eager execution report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Backend construction, frozen and trainable ownership, active-rank views, injection, eager forward, backward, and MoEBlock integration now form one CPU lifecycle.", + "Cross-backend CUDA numerics remain separate." + ] + }, + { + "id": "TA-1344", + "scope": "split EP LoRA router-score report from eager component lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All-to-all and DeepEP score application, identity behavior, and gradient flow now close the mocked CPU MoE LoRA policy.", + "Dispatch, compute, and combine mocks remain isolated in local contexts." + ] + }, + { + "id": "TA-1345", + "scope": "split asynchronous routing-replay record report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "CUDA stream ordering now begins the MoEBlock routing-replay integration under the same hardware gate.", + "The CPU registry and stage-management unit report remains independently runnable." + ] + }, + { + "id": "TA-1346", + "scope": "split multi-layer and pipeline routing-replay report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Single-layer record and replay now proceeds through multi-layer, multi-microbatch, and 1F1B scheduling in one CUDA lifecycle.", + "Global replay state is explicitly reset between phases." + ] + }, + { + "id": "TA-1347", + "scope": "split base-model and R3 routing-replay integration report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Checkpoint enabling, non-MoE admission, full backward replay, and R3 forward preload now close the routing-replay lifecycle.", + "Global replay state is explicitly reset before this phase." + ] + }, + { + "id": "TA-1348", + "scope": "split DTensor checkpoint materialization from grouped loading", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Replicated and sharded DTensor copy and four-rank save materialization now begin the grouped checkpoint-load policy.", + "The real CPU process workers remain intact." + ] + }, + { + "id": "TA-1349", + "scope": "split rank-zero checkpoint resolution and transport from grouped loading", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Object transport, rank-zero filtered prefetch, local resolution, grouped expert routing, and strict coverage now form one load lifecycle.", + "Scoped monkeypatch contexts preserve transport isolation." + ] + }, + { + "id": "TA-1350", + "scope": "split checkpoint model-key compatibility from distributed IO", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Pipeline key unions, QARL buffer admission, base-to-LoRA compatibility, metadata, load groups, and save groups now form one state lifecycle.", + "A fresh compatibility subtree preserves file isolation." + ] + }, + { + "id": "TA-1351", + "scope": "split optimizer-state filtering from distributed checkpoint IO", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Metadata-selected optimizer keys and multi-optimizer child filtering now close the distributed checkpoint policy.", + "A fresh optimizer subtree and scoped patches preserve isolation." + ] + }, + { + "id": "TA-1352", + "scope": "split exact MoE post-EP layout from construction inventory", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The 1700-factor inventory now proceeds through EP placement and logical owner-shape discovery in one construction policy.", + "A scoped EP16 rank patch isolates the layout phase." + ] + }, + { + "id": "TA-1353", + "scope": "split exact selected-logprob head attachment from MoE construction", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Selected-logprob LM-head ownership and fail-closed execution now close the complete exact construction policy.", + "The head phase uses an independent world16 patch context." + ] + }, + { + "id": "TA-1354", + "scope": "split fused GDN delta and merged-forward report from geometry", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Manifest selection and geometry now feed low-rank products, canonical folding, gradients, output projection, and bounded caches.", + "Every gradient-slice and cache-release assertion remains intact." + ] + }, + { + "id": "TA-1355", + "scope": "split fused GDN sharded checkpoint load from component lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Two-shard PEFT load now closes the same fused GDN injection, execution, and serialization lifecycle.", + "A fresh sharded checkpoint subtree preserves the original export." + ] + }, + { + "id": "TA-1356", + "scope": "split FlashMLA input flattening from backward compaction", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Batch offsets and invalid-index normalization now feed valid-row backward compaction and zero-gradient scattering.", + "All behavior remains in the hermetic CPU/mock policy." + ] + }, + { + "id": "TA-1357", + "scope": "split FlashMLA production-envelope admission from backward policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Device, shape, and Hopper admission now close the FlashMLA mock execution policy.", + "A scoped CUDA-capability patch isolates the admission phase." + ] + }, + { + "id": "TA-1358", + "scope": "split exact TP1 configuration from CPU forward policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Construction, rank-alpha admission, dtype identity, input rejection, exact forward, surrogate backward, and safety now form one CPU component policy.", + "The literal CUDA direct-program report remains separate." + ] + }, + { + "id": "TA-1359", + "scope": "split RMSNorm family declaration tripwire from funnel execution", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Undeclared warnings and required-family rejection now precede bitwise family funnel and module dispatch under one CUDA gate.", + "CPU structural guards remain independently runnable." + ] + }, + { + "id": "TA-1360", + "scope": "split fused gate-up registration from MoE checkpoint export", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Base and LoRA fused parameter ownership now begins the MoE checkpoint and HF export policy.", + "All registration shapes and aliases remain asserted." + ] + }, + { + "id": "TA-1361", + "scope": "split fused expert checkpoint-handler roundtrip from model export", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Qwen3 and Qwen3.5 fused expert load, save, deferred-skip, QKV unfusing, and QARL-buffer filtering now form one checkpoint policy.", + "Every handler uses an independent instance." + ] + }, + { + "id": "TA-1362", + "scope": "split OPD KL-estimator report from full-vocab policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "VERL estimator values, k3-plus straight-through gradients, and estimator dispatch now run with full-vocab modes and diagnostics.", + "All estimator formulas remain asserted." + ] + }, + { + "id": "TA-1363", + "scope": "split OPD policy-gradient mode report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Input admission, clipping metrics, PPO KL, and finite policy-gradient loss now form a branch of the complete OPD policy.", + "Backward-compatible metrics remain asserted." + ] + }, + { + "id": "TA-1364", + "scope": "split compiled OPD sampled-logprob ignored-label report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Compiled sampled-token logprobs and ignored-label zeroing now close the OPD dispatch policy.", + "Student and teacher output shapes and masks remain asserted." + ] + }, + { + "id": "TA-1365", + "scope": "split GDN convolution primitive and end-to-end block reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Primitive forward, backward, varlen, batching, and determinism now culminate in full GatedDeltaNet output and gradient parity.", + "Optional SGLang parity and CPU admission remain separate reports." + ] + }, + { + "id": "TA-1366", + "scope": "split SGLang fused-expert admission from CPU resolution and dispatch", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Feature resolution and block dispatch now proceed through unsupported activation, bias, clamp, and trainable-dispatch admission.", + "Each phase uses a scoped monkeypatch context while preserving every guard assertion." + ] + }, + { + "id": "TA-1367", + "scope": "split SGLang fused-expert weight mode and layout report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Transient, cached, and strided weight ownership now closes the CPU fused-expert policy.", + "Cache invalidation, zero-copy views, kernel layout, and split gate-up assertions remain intact." + ] + }, + { + "id": "TA-1368", + "scope": "split SGLang runtime-context report from fused-expert policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Runtime creation, compatible reuse, and incompatible-context rejection now form the final CPU configuration phase.", + "The real optional SGLang execution gate remains separate." + ] + }, + { + "id": "TA-1369", + "scope": "split SGLang trainable-gradient numerics from real parity", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Stock-Triton gradients and masked routing gradients now precede strided and auto-mode bitwise parity under one CUDA and SGLang capability gate.", + "CPU policy coverage is not hidden by the optional dependency gate." + ] + }, + { + "id": "TA-1370", + "scope": "split sparse-MLA attention-sink report from forward parity sweep", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Zero and signed sink parity plus the non-ignored effect check now run in the representative first forward specialization.", + "All four compiled top-k forward specializations remain collected." + ] + }, + { + "id": "TA-1371", + "scope": "split DeepSeek-V3 default LoRA target report from tiny model lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default and explicit MLA and MoE adapter targets now follow the tiny forward, backward, and router-freeze transaction.", + "Every projection type assertion remains in a named scenario helper." + ] + }, + { + "id": "TA-1372", + "scope": "split DeepSeek-V3 router observability and replay report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sparse-layer router output counts and recorded routing weights now close the DeepSeek-V3 model lifecycle.", + "Dense-prefix and all-MoE schedules remain covered." + ] + }, + { + "id": "TA-1373", + "scope": "split DeepSeek-V4 hash-layer structure report from non-hash MoE policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Non-hash routing, shared experts, SwiGLU clamps, hash admission, table selection, and gradients now form one CPU MoE policy.", + "Hash and non-hash branches retain independent fixtures inside named helpers." + ] + }, + { + "id": "TA-1374", + "scope": "split DeepSeek-V4 hash routing replay report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Hash-table routing now culminates in record and replay-backward behavior in the same architecture MoE report.", + "Recorded indices and replay gradients remain asserted." + ] + }, + { + "id": "TA-1375", + "scope": "split MiniMax-M3 checkpoint mapping from architecture support", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Configuration, registry, text runtime, checkpoint fusion, EP ownership, and expert aliases now form one architecture-support report.", + "Raw-key skip accounting and all projection mappings remain asserted." + ] + }, + { + "id": "TA-1376", + "scope": "split MiniMax-M3 paging and CPU MSA admission report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Stable paged-KV layout and loud CPU rejection now close the MiniMax-M3 architecture-support report.", + "The helper remains independent of optional CUDA execution." + ] + }, + { + "id": "TA-1377", + "scope": "split GLM52 routed-bank EP checkpoint slice report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The exact full-block QLoRA inventory now includes scoped EP16 local-bank ownership and offset selection.", + "All 75 routed banks retain local-expert and global-offset assertions." + ] + }, + { + "id": "TA-1378", + "scope": "split GLM52 exact component and admission report from QLoRA inventory", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Logical inventory now proceeds through exact dense roots, rank-alpha rejection, target validation, and supported training-mode admission.", + "The meta-device construction keeps the combined policy hermetic." + ] + }, + { + "id": "TA-1379", + "scope": "split GLM52 routed-expert gradient edges from literal sampler coverage", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "All 256 owner-slot remaps now culminate in sentinel, mixed-owner, and top-k logical VJP behavior under one hardware gate.", + "CPU topology and physical-buffer policy remains separately runnable." + ] + }, + { + "id": "TA-1380", + "scope": "sparse-MLA deterministic-dKV and invalid-index backward reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Both reports pass independently and protect distinct deterministic accumulation and invalid-slot gradient behavior.", + "The larger base backward sweep terminates the current H100 pytest process, so folding these passing boundaries into that sweep would destroy failure isolation." + ] + }, + { + "id": "TA-1381", + "scope": "FP8 grouped-kernel numerics and full MoE optimizer-step reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The grouped forward and weight-gradient report passes independently on H100.", + "The full Quack optimizer-step report terminates the current pytest process, which is a real lifecycle failure that must not erase the passing kernel-level result." + ] + }, + { + "id": "TA-1382", + "scope": "split GLM5 indexer construction from architecture support", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "GLM5 configuration and registry loading now proceed through indexer geometry, FP32 head projection, selection, and masking.", + "Each mutable indexer scenario runs in a scoped monkeypatch context." + ] + }, + { + "id": "TA-1383", + "scope": "split GLM5 sparse-MLA wrapper and attention integration report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Reference arithmetic, dispatch shaping, Ulysses integration, dense parity, and backend rejection now form the sparse-attention phase of one GLM5 support policy.", + "The real TileLang fast path remains an independent CUDA report." + ] + }, + { + "id": "TA-1384", + "scope": "split GLM5 checkpoint filtering from architecture support", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Partial-layer and EP expert-key filtering now precede adapterization and model execution in the architecture lifecycle.", + "Layer, non-layer, local-expert, and out-of-range key assertions remain intact." + ] + }, + { + "id": "TA-1385", + "scope": "split GLM5 adapter sparse-KV and MoE dispatch report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Default MLA, shared-expert, and routed-expert targets now feed EP dispatch and sparse absorbed-KV LoRA execution.", + "Indexer exclusions and live delta contribution remain asserted." + ] + }, + { + "id": "TA-1386", + "scope": "split GLM5 forward and recompute report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense and MoE layer construction now culminates in end-to-end hidden-state output and recompute-before-dispatch checkpoint routing.", + "Optional HF-reference logits remain separately runnable." + ] + }, + { + "id": "TA-1387", + "scope": "split native EP-combine variable-row collective report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "EP8 admission now proceeds through token and ID padding, reduce-scatter gradients, and shared maximum-row selection.", + "Mocked collectives run in an isolated monkeypatch context." + ] + }, + { + "id": "TA-1388", + "scope": "split native EP-combine serving fused-gate gradient report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Serving fused-gate output and trainer gradients now form the differentiable gate phase of the native-combine policy.", + "Hidden, gate, shared, and routed gradients still match the eager reference." + ] + }, + { + "id": "TA-1389", + "scope": "split native EP-combine dispatch and actual-operand report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FSDP module entry, local routed partials, shared projection, chain sum, diagnostics, and final output now close one native EP transaction.", + "Every captured exact-combine boundary remains asserted." + ] + }, + { + "id": "TA-1390", + "scope": "split SGLang EP slot-combine report from dispatch and admission", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Flag, backend, empty-rank, and trainable admission now proceed through slot-ordered reduction and pair-count guards.", + "The CPU EP fixture and mock kernel boundary are identical." + ] + }, + { + "id": "TA-1391", + "scope": "split SGLang EP weight-presentation report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Strided zero-copy views plus transient and cached presentation now close the CPU SGLang EP policy.", + "The real optional stock-Triton gradient comparison remains separate." + ] + }, + { + "id": "TA-1392", + "scope": "FlashQLA auto-CP pin, shape invariance, and chunk-chaining reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The reports protect distinct control-flow pinning, packed and batch-composition invariance, and recurrent state-handoff exactness.", + "Gate 2 and Gate 4 are production numerical boundaries rather than repeated shape smoke cases." + ] + }, + { + "id": "TA-1393", + "scope": "training utility clipping, metadata, pipeline loss, and synchronization reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The four reports exercise unrelated gradient clipping, distributed counting, chunked cross-entropy, and explicit synchronization APIs.", + "Sharing a utility module is not evidence that these failure boundaries are equivalent." + ] + }, + { + "id": "TA-1394", + "scope": "GLM52 index sharing, codec parity, semantic MoE, layer plan, and selector reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The five reports separate FSDP identity, optional sampler bytes, end-to-end logprob composition, static topology, and sparse-selector implementation behavior.", + "Only the sampler codec requires the paired CUDA serving stack; folding would hide the four hermetic CPU contracts." + ] + }, + { + "id": "TA-1395", + "scope": "split MoE all-to-all pre-dispatch score-order report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Token permutation and routing-weight order now feed the mocked all-to-all pre-dispatch transaction.", + "Received ordering, expert cumsums, and routing-weight gradients remain asserted." + ] + }, + { + "id": "TA-1396", + "scope": "split MoE post-all-to-all hidden-chunking report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The same permutation lifecycle now closes with chunked and unchunked post-dispatch output and gradient parity.", + "Pre- and post-dispatch collectives use isolated monkeypatch contexts." + ] + }, + { + "id": "TA-1397", + "scope": "split packing full-pipeline roundtrip from packing-on policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Capacity, batching, token metadata, generated labels, simulated forward output, and sample-boundary unpacking now form one packing-on lifecycle.", + "Packing-disabled behavior remains a separate supported mode." + ] + }, + { + "id": "TA-1398", + "scope": "split trained-QARL logprob preservation from quantized export policy", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "A trained QARL directory now proves exact dequantized logprobs before the CLI layout and admission matrix runs.", + "The scenario uses a dedicated temporary export directory." + ] + }, + { + "id": "TA-1399", + "scope": "split optimizer-step Adam override report from server initialization", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Server defaults and validation now feed full, partial, omitted, adapter, and non-Adam optimizer-step overrides.", + "All mutable runner patches are scoped to the step phase." + ] + }, + { + "id": "TA-1400", + "scope": "split optimizer dispatcher payload-forwarding report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Explicit and omitted Adam fields now traverse configuration, runner mutation, and dispatcher forwarding as one server lifecycle.", + "Backward-compatible None values remain asserted." + ] + }, + { + "id": "TA-1401", + "scope": "split sparse-delta prepacked-path posting report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Sparse encoding, baseline updates, receiver failures, per-rank packed paths, cache metadata, and endpoint accounting now form one backend lifecycle.", + "The prepacked phase uses its own temporary directory." + ] + }, + { + "id": "TA-1402", + "scope": "split sparse-delta initialization admission report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Post-only and prepacked-only admission plus runtime helper loading now close the sparse-delta backend policy.", + "Initialization artifacts use a dedicated temporary directory." + ] + }, + { + "id": "TA-1403", + "scope": "packing-on and packing-disabled reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Packing-disabled is a supported control-flow mode with different batching, warning, shifting, and loss-mask behavior.", + "It is not a narrow input variation of the packing-on roundtrip." + ] + }, + { + "id": "TA-1404", + "scope": "FP8 weight quantization primitive and directory export reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The primitive report protects block-scale numerical and layout behavior without filesystem or CLI dependencies.", + "The export report protects model-state transformation, artifact layout, configuration, and admission." + ] + }, + { + "id": "TA-1405", + "scope": "split checkpoint compatibility argument parsing report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Optimizer, packing, and numeric configuration parsing now proceeds through EP checkpoint compatibility, automatic checkpoint resolution, and optimizer-state loading.", + "The checkpoint phase uses a dedicated configuration directory and monkeypatch context." + ] + }, + { + "id": "TA-1406", + "scope": "split low-precision argument parsing report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The same parser lifecycle now closes with FP8 aliases, fail-fast defaults, vLLM knob rejection, and supported low-precision modes.", + "All environment and argv mutations remain isolated." + ] + }, + { + "id": "TA-1407", + "scope": "split dataset local, hub, URL, and data-files loading report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dataset name and shard expansion plus type inference now feed local-file, saved-directory, hub, URL, and data-files loading.", + "The loading phase uses an isolated cache and temporary directory." + ] + }, + { + "id": "TA-1408", + "scope": "split preprocessed dataset save-load report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Expanded and loaded datasets now proceed through split, merge, save, reload, and missing-cache behavior as one preparation lifecycle.", + "The persisted artifact uses a dedicated temporary directory." + ] + }, + { + "id": "TA-1409", + "scope": "split P2P trainer abort-marker report from device selection", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "IB-device mapping precedence now proceeds through sync-abort publication, peer observation, and cleanup.", + "Abort state is scoped to a temporary transfer directory." + ] + }, + { + "id": "TA-1410", + "scope": "split P2P trainer peer-status gather report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The trainer transfer policy now closes by gathering local success and remote failure status across ranks.", + "Distributed mocks run in an isolated monkeypatch context." + ] + }, + { + "id": "TA-1411", + "scope": "split FP8 adapter-merge weight-sync report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "BF16-island selection and block layout now feed dense QLoRA and MoE adapter folding before FP8 quantization.", + "Projection selection, skip lists, existing FP8 state, and stack behavior remain asserted." + ] + }, + { + "id": "TA-1412", + "scope": "split FP8 CPU expert projection and workspace report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The CPU weight-sync policy now proceeds through expert transposition, zero padding, deferred quantization, streaming workspaces, flush reset, and completion metadata.", + "The independent live GPU parity and device-transfer report remains separate." + ] + }, + { + "id": "TA-1413", + "scope": "checkpoint save-load, list-delete, and model-ID validation API reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The reports protect distinct mutating checkpoint I/O, listing and deletion isolation, and request-boundary path validation.", + "Collapsing public endpoint and security failures into one large transaction would reduce actionable failure isolation." + ] + }, + { + "id": "TA-1414", + "scope": "weight-sync receiver, source preparation, and sparse-delta handler reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Receiver post-processing, trainer-side parameter extraction, and sparse prepacked transport are different endpoints and data paths.", + "Each report already consolidates its internal configuration, layout, and admission variants." + ] + }, + { + "id": "TA-1415", + "scope": "inference endpoint registration, weight sync, and quantization normalization reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Registration and health discovery, explicit synchronization, and quantization schema validation expose distinct public API boundaries.", + "Each report contains its own success, rejection, and compatibility lifecycle rather than shape-only variants." + ] + }, + { + "id": "TA-1416", + "scope": "split adapter registration report from coordinated load lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Fresh and evicted adapter loading now proceeds through adapter and session registration, broadcasts, cross-rank failure rollback, and worker exceptions.", + "Registration uses a dedicated temporary checkpoint root." + ] + }, + { + "id": "TA-1417", + "scope": "split adapter save admission report from coordinated load lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The coordinator policy now closes with save admission that refuses to recreate missing evicted state.", + "Load, registration, and save paths retain independent fixtures inside named helpers." + ] + }, + { + "id": "TA-1418", + "scope": "split optimizer checkpoint save-resume report from identity transaction", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Canonical parameter identity and transactional collective failure now feed sharded manifest creation, artifact admission, and bitwise resumed training.", + "The save-resume phase uses an isolated monkeypatch context and artifact root." + ] + }, + { + "id": "TA-1419", + "scope": "split optimizer logical reshard report from checkpoint resume", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Saved optimizer moments now proceed through one-dimensional, multidimensional, replicated, same-world, and invalid-source logical resharding.", + "Reshard fixtures use a dedicated checkpoint directory and patch context." + ] + }, + { + "id": "TA-1420", + "scope": "split model-runner initial checkpoint restore report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Checkpoint materialization and zero-meta admission now feed runner step restoration, optimizer selection, and failure-state publication.", + "Manager-level patches are undone before the runner phase." + ] + }, + { + "id": "TA-1421", + "scope": "split default-adapter initialization report from checkpoint loading", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The base restore lifecycle now culminates in guarded fresh default-adapter initialization and ownership compilation.", + "Uninitialized to-empty storage remains explicitly rejected before registration." + ] + }, + { + "id": "TA-1422", + "scope": "split dispatcher routing-payload transport and security report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Batch sharding and provenance now feed filesystem and Mooncake routing payload slicing, legacy and symlink rejection, and diagnostic artifacts.", + "Routing transport uses an isolated temporary root and monkeypatch context." + ] + }, + { + "id": "TA-1423", + "scope": "split dispatcher completion rendezvous and per-token merge report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The dispatcher lifecycle now closes with CP replica deduplication, disagreement rejection, bounded completion payloads, and rank-zero merge.", + "All completion cases remain in a named scenario helper." + ] + }, + { + "id": "TA-1424", + "scope": "split endpoint health preflight from NCCL transfer report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Endpoint model discovery and all-failure diagnostics now precede NCCL initialization, direct bucket routing, and two-phase receiver completion.", + "The combined report covers one endpoint health-to-transfer lifecycle." + ] + }, + { + "id": "TA-1425", + "scope": "split adapter ownership compiler topology report from producer execution", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "A declared module-managed producer now feeds all four topology families, stable fingerprints, fail-closed admission, and replica-domain coverage.", + "Fullgraph producer gradients and compiler structure remain asserted." + ] + }, + { + "id": "TA-1426", + "scope": "split adapter residual transport report from ownership compilation", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The compiled ownership plan now culminates in bucketed residual reduction, immutable raw accumulators, and logical norm accounting.", + "Distributed finalizer mocks use an isolated monkeypatch context." + ] + }, + { + "id": "TA-1427", + "scope": "split runner expert-factor compilation and admission report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense, exact LM-head, and replica topology compilation now proceeds through unquantized, block-FP8, NVFP4, NF4, and session-rank expert-factor contracts.", + "Certified and rejected backend combinations retain dedicated helper fixtures." + ] + }, + { + "id": "TA-1428", + "scope": "split direct-output analytical-step report from runner ownership compiler", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The runner compiler policy now closes with authoritative analytical gradients, effective LM-head folding, capture finalization, and parameter mutation.", + "The direct-output phase uses a dedicated manager root and monkeypatch context." + ] + }, + { + "id": "TA-1429", + "scope": "split OPD distributed teacher-cache assembly report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Teacher contributor selection, CP gathering, Mooncake metadata, and cache row consumption now feed the OPD loss execution lifecycle.", + "Patch-decorated helpers are invoked with isolated fixture contexts." + ] + }, + { + "id": "TA-1430", + "scope": "split OPD debug-artifact report from loss execution", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Packed cache and weight shaping plus teacher assembly now culminate in loss metrics and ranked vocab-parallel debug artifacts.", + "Loss and debug files use separate temporary directories." + ] + }, + { + "id": "TA-1431", + "scope": "split fused gate-up and Nemotron-H Muon classification report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Muon construction, grouping, fallback, Gram-Newton-Schulz stepping, and Quack backend admission now include fused gate-up discovery and a tiny Nemotron-H update.", + "Standard Newton-Schulz arithmetic and CUDA FP32-compute preservation remain separate." + ] + }, + { + "id": "TA-1432", + "scope": "attention backend registry, FlashAttention API, and SGL page-cache reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Backend resolution, variable-length FlashAttention behavior, and paged KV-cache semantics exercise different public APIs and storage models.", + "Their shared module location does not make their regressions equivalent." + ] + }, + { + "id": "TA-1433", + "scope": "mixed precision, folded sequence parallelism, and manual FSDP prefetch reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The reports protect independent dtype selection, model transformation, and execution-order configuration boundaries.", + "None is an input variation of another policy." + ] + }, + { + "id": "TA-1434", + "scope": "Muon standard Newton-Schulz and CUDA compute-dtype reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Batched leading-dimension preservation tests a different orthogonalization algorithm from the main Gram-Newton-Schulz policy.", + "Live CUDA FP32 backend dtype preservation cannot be replaced by CPU construction or model-classification assertions." + ] + }, + { + "id": "TA-1435", + "scope": "split dense primitive and model batch-composition reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Batch-invariant matmul, RMSNorm, log-softmax, and mean now establish the primitive contract before the padded dense-model consumer runs under the same CUDA gate.", + "The surviving report retains every primitive equality and full-model hidden-state assertion." + ] + }, + { + "id": "TA-1436", + "scope": "split DeepSeek-V4 attention runtime and sink-dtype reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Window and compressed attention forward-backward coverage now includes sink storage, TileLang call-dtype conversion, RoPE state, and TP admission in one attention lifecycle.", + "The TileLang boundary remains mocked and hermetic inside the representative window-attention case." + ] + }, + { + "id": "TA-1437", + "scope": "split DeepSeek-V4 wo_a and full attention-LoRA training reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The grouped wo_a delta and gradient regression now precedes all-target injection, base freezing, autograd reachability, and first-step adapter gradients.", + "Both scenarios retain independent model construction and inputs within one supported attention-LoRA policy." + ] + }, + { + "id": "TA-1438", + "scope": "standalone TileLang indexer causal-mask report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Every forward geometry now verifies both valid numerical scores and all invalid future positions from the same kernel output.", + "A vectorized negative-infinity assertion replaces the redundant kernel invocation and Python element walk; large-value and zero-input cases still run once." + ] + }, + { + "id": "TA-1439", + "scope": "dataset split-fingerprint and configuration-hash reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Split derivation and complete dataset-configuration identity are different production algorithms with different inputs and consumers.", + "Their shared digest representation does not make either report an example of the other." + ] + }, + { + "id": "TA-1440", + "scope": "dataloader builder configuration and packed pipeline integration reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The first report isolates sampler, microbatch, parallel-state, and collator construction; the second executes real loader batches and packed sequence layouts.", + "Keeping the unit boundary separately runnable preserves useful failure localization without one report per input example." + ] + }, + { + "id": "TA-1441", + "scope": "Mooncake hidden transport and metadata-admission reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Successful tensor transport and TeacherActivationCache consumption are independent from malformed metadata, missing objects, removal, and environment configuration failures.", + "The positive data path and fail-closed admission path already consolidate their internal variants." + ] + }, + { + "id": "TA-1442", + "scope": "DeepSeek-V4 HF-to-DCP conversion and AutoModel loading reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Distributed-checkpoint conversion and Transformers AutoModel dispatch are distinct public ingestion paths even though both consume a synthetic HF snapshot.", + "A failure in one does not imply or diagnose a failure in the other." + ] + }, + { + "id": "TA-1443", + "scope": "GLM52 native-FP8 routing-runtime and checkpoint-buffer reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Canonical routing and scoring-only expert execution are runtime behavior, while configuration validation, module replacement, byte packing, and checkpoint ownership are construction and persistence boundaries.", + "Each report already joins the narrow examples inside its own production boundary." + ] + }, + { + "id": "TA-1444", + "scope": "split families-v2 RMSNorm reference and realization reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "FP64 wrongness bounds, residual behavior, zero-centered mode, batch invariance, and repeatability now precede forced fused-versus-split bit equality and dispatch selection.", + "All checks exercise one frozen RMSNorm numerical tree under the same CUDA gate." + ] + }, + { + "id": "TA-1445", + "scope": "split exact factor-only weight-sync and checkpoint-publication reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Exact dense, projection, LM-head, streaming, and sparse-delta preflight rejection now culminate in separate A and B factor persistence and byte verification.", + "The filesystem publication uses an isolated temporary checkpoint inside the same factor-only contract." + ] + }, + { + "id": "TA-1446", + "scope": "parameterized index-share checkpoint modes and separate cleanup report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Reentrant and non-reentrant checkpoint modes are now an internal two-case input matrix rather than two product report IDs.", + "Producer reuse, detached payloads, gradients, forward-only completion, forward failure, backward failure, and idempotent cleanup form one index-share lifecycle." + ] + }, + { + "id": "TA-1447", + "scope": "standalone solve_tril two-warp configuration pin", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The warp count selects the association tree of a Triton reduction and is therefore part of the exact arithmetic contract, not a cosmetic tuning constant.", + "The tolerant GDN runtime parity report cannot detect a bit-level drift in that source configuration." + ] + }, + { + "id": "TA-1448", + "scope": "standalone two-GPU PP 1F1B convergence job", + "decision": "remove", + "status": "applied", + "evidence": [ + "The schedule-parity report already launches the same PP2 FSDP1 1F1B baseline before comparing all virtual-stage schedules.", + "Baseline convergence is now asserted in that surviving transaction, eliminating a separate full training subprocess." + ] + }, + { + "id": "TA-1449", + "scope": "separate LoRA FSDP2 convergence and checkpoint-resume jobs", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The two-phase FSDP2 checkpoint transaction now resumes to the former 20-step convergence horizon and asserts the explicit checkpoint-load marker.", + "The surviving report passed with checkpoint restoration, global step 20, and the original loss-convergence threshold." + ] + }, + { + "id": "TA-1450", + "scope": "test-authored OPD packing-strategy loss invariance report", + "decision": "remove", + "status": "applied", + "evidence": [ + "Its fake backend computed KL and global normalization entirely with test helpers rather than invoking XoRL's OPD loss implementation.", + "Production OPD numerics and reducer composition plus sequential, best-fit, and balanced-DP packing remain covered directly." + ] + }, + { + "id": "TA-1451", + "scope": "fake teacher-cache metadata echo report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The backend returned fixed metadata and timing literals that the assertions merely read back from the request-processor response.", + "Real ModelRunner Mooncake publication, cache-index construction, activation-cache consumption, and OPD pipeline transport remain covered." + ] + }, + { + "id": "TA-1452", + "scope": "remaining full-weight FP8 end-to-end topology matrix", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The seven survivors select distinct checkpoint-resume, tensor-parallel, Ulysses, Ring, hybrid long-tail packing, local MoE, and DeepEP EP/eFSDP mechanisms.", + "Earlier shape-ladder and checkpoint-by-topology cross-products are already removed; no survivor is only a larger input example of another." + ] + }, + { + "id": "TA-1453", + "scope": "P2P staging, FP8 layout, and transfer-failure reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Source staging and receiver placement, fused FP8 byte layouts, and fail-closed manifest and runtime diagnostics exercise different transport boundaries.", + "Each report already consolidates the narrow locator, shape, retry, and layout variants within its boundary." + ] + }, + { + "id": "TA-1454", + "scope": "adapter manager ownership, capture, optimizer, restore, and multi-adapter reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Plan compilation, gradient capture, authoritative mutation, checkpoint admission, and multi-adapter eviction are independent state-machine transitions with different rollback obligations.", + "The reports already aggregate their per-input and per-optimizer branches rather than multiplying product IDs." + ] + }, + { + "id": "TA-1455", + "scope": "duplicate eight-GPU PP2 FSDP4 AdamW convergence job", + "decision": "remove", + "status": "applied", + "evidence": [ + "The retained Muon job exercises the same PP2, FSDP4, two-microbatch, packed training topology while also covering optimizer partitioning.", + "AdamW construction and stepping are covered independently; changing the optimizer does not select a different pipeline transport or loss-normalization mechanism." + ] + }, + { + "id": "TA-1456", + "scope": "dead Nemotron-H end-to-end training module", + "decision": "remove", + "status": "applied", + "evidence": [ + "All three static reports required tiny_nemotron_h_model_dir, a fixture absent from every repository conftest and absent since the module was introduced.", + "Every report therefore errored during pytest setup without constructing a model; real Nemotron runtime, packed-varlen, checkpoint, gradient, and Muon-step policies remain covered." + ] + }, + { + "id": "TA-1457", + "scope": "native-FP8 plain-conversion family parametrization", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Linear and expert modules are two inputs to the same frozen-FP32 ordinary conversion contract and now execute in one internal family matrix.", + "Both families still prove non-DTensor materialization, byte preservation, device placement, dtype, and requires-grad state." + ] + }, + { + "id": "TA-1458", + "scope": "FlashQLA small-head and production-head report IDs", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Four-head and 32-head inputs now run sequentially inside one FlashQLA-versus-FLA numerical policy instead of producing separate pytest reports.", + "Both shapes retain forward, final-state, and every input-gradient cosine and finiteness check on Hopper." + ] + }, + { + "id": "TA-1459", + "scope": "PEFT MoE EP-slice orientation parametrization", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Down-projection A and gate-projection B are transpose orientations of the same global-to-local expert slice conversion and now form one report.", + "Both orientations still construct all eight published experts and verify the exact rank-two local shard." + ] + }, + { + "id": "TA-1460", + "scope": "non-gated MoE backend report parametrization", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Installed Triton and native implementations now feed one eager-reference policy instead of emitting one report per backend.", + "Each available backend still verifies forward values plus input, gate-up, and down-projection gradients." + ] + }, + { + "id": "TA-1461", + "scope": "surviving PP2 FSDP4 Muon topology gate", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The report is the only direct-trainer eight-GPU composition of pipeline parallelism and FSDP4 after duplicate removal.", + "A live run reached the 1F1B schedule and exposed the current pipeline backward product failure, so it is an effective behavioral gate rather than a configuration smoke." + ] + }, + { + "id": "TA-1462", + "scope": "unreachable FP8 hybrid Ulysses Ring long-tail E2E report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report requested tiny_agent_context_dense_model_dir_with_weights, a fixture that has never existed in repository history, so it failed during setup before constructing a trainer.", + "The surviving FP8 E2E matrix separately exercises Ulysses, Ring, checkpoint-resume, local MoE, and DeepEP EP/eFSDP, while packing behavior has functioning lower-level coverage." + ] + }, + { + "id": "TA-1463", + "scope": "DeepSeek-V4 window and compressed attention report parametrization", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Window attention and C128 compression are internal modes of one forward-backward shape policy and now execute as an isolated in-report matrix.", + "Both modes retain output, parameter-gradient, RoPE, QAT, sink-dtype, TileLang call-dtype, and topology-admission assertions." + ] + }, + { + "id": "TA-1464", + "scope": "exact GLM sparse attention CP report parametrization", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Non-CP and Ulysses-CP select two branches of the same exact q and v projection routing contract and now run as an isolated Boolean mode matrix.", + "Both branches still verify call order, factor-only execution, query offsets, tensor shapes, and the absence of frozen-weight materialization." + ] + }, + { + "id": "TA-1465", + "scope": "stale optional GLM5 FLOPs-counter module gate", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The test skipped the entire module on a missing GLM5 configuration that is now a shipped source-tree dependency, not an optional package.", + "Direct imports make a broken GLM5 package or configuration import fail the report instead of silently removing coverage." + ] + }, + { + "id": "TA-1466", + "scope": "stale Dr.GRPO model-runner implementation gate", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "ModelRunner now has both the Dr.GRPO dispatch branch and its loss-exclusion inventory, making the upstream-WIP module skip obsolete.", + "The report now unconditionally covers dispatch, legacy field names, temperature, output suppression, and K3 output requests." + ] + }, + { + "id": "TA-1467", + "scope": "eager versus native MoE broad import skip", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The helper caught every exception while importing shipped internal MoE modules and converted production import regressions into skips.", + "Lazy imports remain collection-safe, but a failure now fails the live forward, backward, determinism, and edge-case policy." + ] + }, + { + "id": "TA-1468", + "scope": "Triton and Quack EP routing-score report parametrization and import skip", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both backends implement the same routing-score forward and gradient boundary and now run inside one isolated backend matrix rather than emitting separate product reports.", + "The Quack case had always skipped because its test-owned module stub omitted a required grouped-GEMM symbol; the repaired stub now executes Quack, and internal import failures are no longer suppressed." + ] + }, + { + "id": "TA-1469", + "scope": "opt-in DeepGEMM grouped-FP8 subprocess report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report skipped unless XORL_TEST_DEEP_GEMM_FP8=1, and no repository workflow, script, or configuration ever sets that flag.", + "A permanently dormant manual diagnostic is not suite coverage; the remaining FP8-MoE reports execute injection, grouped forward and weight gradients, Quack and Triton backends, and full optimizer steps." + ] + }, + { + "id": "TA-1470", + "scope": "Quack EP parity broad internal-import skips", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Both CPU gradient-arity and GPU Quack-versus-Triton reports caught every exception from shipped internal modules and converted implementation regressions into skips.", + "Quack is a pinned project dependency; explicit CUDA gates remain, while internal imports now fail closed and both reports pass." + ] + }, + { + "id": "TA-1471", + "scope": "GKN checkpoint and eager-backend internal-import skips", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The reports described transformers as optional even though it is a core dependency and the imported checkpoint buffer and eager backend are shipped source files.", + "Direct imports preserve the GKN checkpoint, eager, native, Triton, and cross-backend assertions while exposing packaging or import regressions." + ] + }, + { + "id": "TA-1472", + "scope": "GLM5 TileLang indexer broad import skip", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The supported-CUDA report caught every exception importing the shipped TileLang indexer and could hide a broken kernel module behind an availability skip.", + "TileLang is pinned by the project; the explicit CUDA gate remains and the direct-import fast-path parity report passes." + ] + }, + { + "id": "TA-1473", + "scope": "FlashQLA contract and numerical parity broad import skips", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The reports already gate on CUDA architecture and the required TileLang feature, so an additional catch-all around shipped FlashQLA imports only concealed backend regressions.", + "SM90 and prefer_instruction skips remain explicit; all four exact-contract reports and the FlashQLA-versus-FLA numerical report pass with direct imports." + ] + }, + { + "id": "TA-1474", + "scope": "standalone forced-SSM-kernel-unavailable report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Forcing use_kernel=True with the kernel removed is the admission branch of the same CPU SSD fallback and packed-recurrence policy.", + "The RuntimeError assertion now runs after dense, packed, chunking, gradient, and full-mixer recurrence checks without a separate pytest report." + ] + }, + { + "id": "TA-1475", + "scope": "standalone learning-rate scheduler invalid-configuration report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Invalid learning rate, warmup ratio, and mode inputs are builder admission branches adjacent to constant, linear, and cosine scheduling.", + "All five rejection cases now close the scheduler mode policy instead of producing an independent product report." + ] + }, + { + "id": "TA-1476", + "scope": "direct pytest.main launch blocks in test modules", + "decision": "remove", + "status": "applied", + "evidence": [ + "Thirteen test modules carried __main__ blocks that are never reached by repository pytest invocation, collection, or CI.", + "Removing the duplicate direct-launch surface changes no test behavior; individual modules remain runnable through pytest paths and node IDs." + ] + }, + { + "id": "TA-1477", + "scope": "pytest GPU and skip markers attached to private FP8 assertion helpers", + "decision": "remove", + "status": "applied", + "evidence": [ + "Pytest markers on directly called _assert helpers do not enforce selection or skipping and falsely implied that each helper had an independent hardware gate.", + "The seven collected FP8 policies retain their real GPU markers, execute every helper, and pass after the inert helper decorations are removed." + ] + }, + { + "id": "TA-1478", + "scope": "pytest markers attached to private assertion helpers repository-wide", + "decision": "remove", + "status": "applied", + "evidence": [ + "One hundred twenty private helpers across 30 files carried MarkDecorator metadata even though their names prevent collection and callers invoke them directly.", + "All public test reports retain their CPU, GPU, async, distributed, architecture, and optional-dependency markers; representative model, operator, distributed, server, and weight-sync reports pass." + ] + }, + { + "id": "TA-1479", + "scope": "unused abstract-method lookalikes in the FP8 weight-sync QLoRA fake", + "decision": "remove", + "status": "applied", + "evidence": [ + "The fake subclass overrode four base methods only to raise NotImplementedError, but the exercised merge-and-sync path never calls them and the base methods already fail identically.", + "The retained dequantize_expert implementation supplies the only fake behavior consumed by the production weight-sync path, whose focused policy passes." + ] + }, + { + "id": "TA-1480", + "scope": "test-tree Ruff defects in fixture and path-bootstrap scaffolding", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The packed-dataset fixture assigned an unused current_pos variable, while three standalone distributed or E2E modules intentionally imported after path bootstrapping without local E402 annotations.", + "The dead assignment is removed, the intentional imports are explicit, and the complete tests tree now passes Ruff." + ] + }, + { + "id": "TA-1481", + "scope": "duplicate balanced synthetic TopK routing report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The dedicated TopK router policy already proves balanced expert selection, uniform weights, count balance, and override precedence for synthetic routing.", + "Only MoEBlock replay regather behavior was unique in the second report; that assertion now closes the train-router dispatch policy and both surviving router policies pass." + ] + }, + { + "id": "TA-1482", + "scope": "unused repository-wide and E2E fixtures plus their dead helper types", + "decision": "remove", + "status": "applied", + "evidence": [ + "An AST fixture dependency map covering function parameters, fixture-to-fixture dependencies, usefixtures, indirect parametrization, and getfixturevalue found no consumer for fake_packed_dataset or small_dense_model_dir_with_weights.", + "FakePackedDataset and the root SimpleCollator were reachable only from that dead fixture or nowhere; full pytest setup planning succeeds after all four dead setup surfaces are removed." + ] + }, + { + "id": "TA-1483", + "scope": "standalone vocab-parallel reverse-KL pseudo-test", + "decision": "remove", + "status": "applied", + "evidence": [ + "The file instructed users to invoke pytest but defined only main and worker functions, so collection produced zero reports and no repository workflow called it directly.", + "The retained four-rank lm-head TP FSDP policy reaches the gathered vocab-parallel OPD implementation through production code and passes loss plus hidden and weight gradient comparison across six CP, DP, and HSDP topologies." + ] + }, + { + "id": "TA-1484", + "scope": "manual DeepEP, uneven vocab-parallel OPD, and real-model QLoRA campaign workloads under tests", + "decision": "relocate", + "status": "applied", + "evidence": [ + "These modules define no pytest report and are not called by a repository workflow; they require direct torchrun or Python execution, or two 100-step eight-H100 real-checkpoint training runs.", + "Their diagnostic value is preserved under certification/deepep, certification/opd, and certification/qwen3_30b while the tests tree now contains only repository regression coverage and shared test support." + ] + }, + { + "id": "TA-1485", + "scope": "fabricated Quack DeepEP None-gradient backward report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report called custom-autograd backward methods directly with a SimpleNamespace context and grad_output=None, then asserted a tuple length derived from that same fake context; PyTorch autograd and production dispatch were never entered.", + "The retained Quack grouped-GEMM and DeepEP no-permute reports execute real forward and backward graphs, compare outputs and every trainable gradient against trusted implementations, and pass." + ] + }, + { + "id": "TA-1486", + "scope": "context-parallel FLOPs report using an unsupported GLM model type", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "XorlFlopsCounter has no xorl_glm5 estimator, so both sides of the former cp_size comparison were always zero and the report could not detect double-counted sequence lengths.", + "The rewritten report uses the supported qwen3_moe estimator, proves the baseline is nonzero, and then checks that changing cp_size does not alter global-sequence FLOPs." + ] + }, + { + "id": "TA-1487", + "scope": "standalone QARL activation NVFP4 value and STE report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report used the production internal quantizer as its numerical reference and repeated the two-dimensional STE already owned by the independent pure-PyTorch NVFP4 operator policy.", + "Its unique leading-dimension reshape contract now executes inside that independent policy, preserving value, shape, and gradient assertions while removing one pytest report." + ] + }, + { + "id": "TA-1488", + "scope": "standalone W4A4 MoE temporary backend-name report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report exercised only the temporary triton-to-triton_w4a4 name switch and restoration; it did not run a quantized down projection or grouped GEMM.", + "That admission and exception-restoration lifecycle now closes the existing CPU NVFP4 MoE conversion, eager execution, and injection policy, which passes with all former assertions." + ] + }, + { + "id": "TA-1489", + "scope": "duplicate Mooncake byte-store fake in distillation and utility suites", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The distillation file redefined the same in-memory put, get, existence, and removal object API already provided by tests._helpers.opd solely to record call keys.", + "The shared fake now records those calls for both consumers; transport and teacher-cache report boundaries remain unchanged and all four reports pass." + ] + }, + { + "id": "TA-1490", + "scope": "remaining checkpoint, experiment-planning, exporter, and OPD script reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Checkpoint process-group selection is independent of expert-mesh restore; experiment ingestion, path admission, calibration, and correctness-gated ranking enter different production branches.", + "Quantization primitives are independent of on-disk CLI layout, while OPD endpoint-version verification is independent of payload preparation and transport, so merging these reports would hide distinct failure boundaries." + ] + }, + { + "id": "TA-1491", + "scope": "standalone launcher CLI override parsing and removed-field report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report covered one schema-agnostic parse and one removed ZORL override rejection, both part of the existing server removed-configuration admission boundary.", + "Those exact assertions now run with YAML and direct override rejection in test_removed_configuration_boundary; launcher worker discovery and readiness remain a separate report." + ] + }, + { + "id": "TA-1492", + "scope": "standalone two-token ModelRunner causal-LM loss summation report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report invoked _compute_micro_batch_loss on a zero-logit toy model solely to prove that two equal token losses are summed rather than averaged.", + "The same raw-sum and per-token assertions now precede the retained ModelRunner loss-dispatch policy, which exercises DR-GRPO fields, options, temperature, and output controls through the same production method." + ] + }, + { + "id": "TA-1493", + "scope": "standalone QLoRA clean-interpreter import smoke report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report only checked that importing QLoRA utilities and expert modules returned zero, so it did not prove its stated model-package decoupling contract.", + "A strengthened clean-interpreter check now closes the expert capability and ownership policy by asserting that neither xorl.models nor any child module is loaded; the standalone report and file are removed." + ] + }, + { + "id": "TA-1494", + "scope": "optional Torch import inside the server protocol serialization policy", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The protocol report had already imported and exercised XoRL server modules that require Torch, yet used pytest.importorskip before its tensor round trip.", + "Torch is a core dependency and is now imported normally, so a broken installation fails closed instead of silently dropping the tensor serialization branch." + ] + }, + { + "id": "TA-1495", + "scope": "remaining server API, path security, worker protocol, and lifecycle report boundaries", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Outbound endpoint validation, artifact and diagnostic path confinement, and compile-worker function and message admission protect different trust boundaries.", + "API configuration validation, TensorData re-nesting, session publication, optimizer fallback, training response metrics, and ready-handshake queuing reach separate consumers and failure paths rather than input variants of one branch." + ] + }, + { + "id": "TA-1496", + "scope": "standalone DeepSeek-V4 AutoConfig and meta-builder registration report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The standalone report fabricated the same tiny standard HF snapshot shape used by the retained AutoModel from_pretrained loader report, then stopped after AutoConfig and meta construction.", + "AutoConfig class resolution, HF mapping, XoRL meta construction, actual AutoModel weight loading, and tensor equality now form one standard-snapshot loader policy; DCP conversion remains a separate report." + ] + }, + { + "id": "TA-1497", + "scope": "standalone Nemotron-H registry and local-config construction report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report proved registry lookup and Ultra-style config normalization but never constructed or executed the resulting model family.", + "Those loader assertions now open the retained Nemotron-H runtime, backward, router, and gradient-checkpointing policy; packed variable-length behavior and checkpoint parity remain independent reports." + ] + }, + { + "id": "TA-1498", + "scope": "remaining model-family config, registry, checkpoint, and runtime report boundaries", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Kimi-wrapped DeepSeek-V3 conversion covers nested text-config aliases absent from the base runtime model, while Qwen3.5 dense and MoE config normalization has no general model-construction owner to absorb it.", + "MiniMax M3 already owns config, registry, runtime, admission, checkpoint, and paging in one policy; Qwen2 and OLMo2 each combine HF construction, fused and unfused layouts, checkpoint round trips, and numerical HF parity." + ] + }, + { + "id": "TA-1499", + "scope": "standalone request-retry decorator report in data preparation", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "retry_on_request_exceptions has one production consumer: the high-level prepare_datasets operation, so its immediate success, transient HTTP failures, exhaustion, backoff, and unrelated-exception behavior are part of dataset preparation resilience.", + "Those assertions now close the existing dataset expansion, split, loader, merge, save, and reload lifecycle; the standalone utility report and file are removed." + ] + }, + { + "id": "TA-1500", + "scope": "stochastic-rounding unbiasedness stress loop", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The former CPU check accumulated 4,000 separately rounded 64 by 64 tensors, performing more than 16 million element updates to estimate a generic relative-error bound.", + "One seeded 65,536-element sample at exactly one quarter of a BF16 interval now checks the expected 25 percent round-up probability, sample mean, and legal neighbors directly and passes." + ] + }, + { + "id": "TA-1501", + "scope": "two-hundred-trial distributed stochastic-rounding bias tail", + "decision": "remove", + "status": "applied", + "evidence": [ + "After its four-rank reduce-scatter comparison, the report launched 200 additional all-to-all collectives solely to repeat the unbiased-expectation property owned by the CPU primitive policy.", + "The distributed report retains the distinct native-FP32 comparison and per-element BF16 transit error bound and passes on four GPUs in 18 seconds; FSDP2 integration remains separate." + ] + }, + { + "id": "TA-1502", + "scope": "remaining optimizer, trainer-utility, and data-preparation report boundaries", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Optimizer reports already aggregate construction, grouping, numerical updates, state strategy, cautious decay, and backend admission by optimizer family rather than by input size.", + "Gradient clipping, token and microbatch metadata, pipeline chunked CE, explicit gradient synchronization, timer fail-soft handling, live CUDA hooks, collator layouts, fingerprints, and packing each reach separate production consumers." + ] + }, + { + "id": "TA-1503", + "scope": "standalone single-GPU Qwen3-8B LoRA convergence job", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The job repeated the retained two-GPU checkpoint transaction's real model, LoRA rank and alpha, learning rate, 20-step horizon, and exact convergence threshold without selecting another production branch.", + "The surviving FSDP2 job adds checkpoint save and load, an explicit load marker, and final-step validation; model and server LoRA policies own non-FSDP construction, forward, backward, optimizer, and checkpoint behavior." + ] + }, + { + "id": "TA-1504", + "scope": "pseudo-E2E CUDA OPD free-tensor convergence loop", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report bypassed ModelRunner initialization, called a private loss helper directly, and optimized hidden-state and lm-head tensors as free parameters for eight steps, so decreasing loss did not represent a trainer or server lifecycle.", + "The runner policy already owns two-teacher cache loading, metrics, loss, and backward through that helper, while the real GPU OPD server report owns ModelRunner startup, forward_backward, and optim_step; the pseudo-E2E file is removed." + ] + }, + { + "id": "TA-1505", + "scope": "remaining end-to-end topology and integration report boundaries", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The FP8 jobs select dense checkpoint-resume, tensor parallel, Ulysses, Ring, plain MoE, and DeepEP expert-sharded paths; pipeline jobs separately reach direct trainer, schedule parity, server ModelRunner, FSDP, and folded PP-EP-CP topologies.", + "The retained OPD reports own request packing and Mooncake grouping, a real SGLang teacher, and the complete sampler-teacher-Mooncake-trainer-weight-sync loop; DistSignSGD, hybrid shared-LoRA MoE telemetry, and LoRA checkpoint resume exercise distinct production mechanisms." + ] + }, + { + "id": "TA-1506", + "scope": "Qwen3.5 trunk-wrap finite-output GPU report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report wrapped a tiny Qwen3.5 model and asserted only BF16 dtype and finite hidden states, after the same file had already proved the exact full-attention, linear-attention, dense-MLP, shared-expert, and exclusion inventory.", + "Generic trunk policies already prove bitwise forward and backward, batch invariance, BF16 admission, serving-lane equality, and two-rank FSDP2 composition, so the model-specific finite-output report selected no uncovered runtime behavior." + ] + }, + { + "id": "TA-1507", + "scope": "scale-only Quack DeepEP parity and checkpoint-training cases", + "decision": "remove", + "status": "applied", + "evidence": [ + "The 16K-token parity case repeated the same production geometry, balanced routing, unchunked no-permute path, reference comparison, and gradients already exercised at 4K tokens.", + "The 32K-token checkpoint case repeated the same checkpointed three-step optimizer path already exercised at 8K tokens; random routing, empty experts, explicit chunking, and checkpoint versus non-checkpoint training remain." + ] + }, + { + "id": "TA-1508", + "scope": "tautological and scale-only block-FP8 workload tails", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The generic codec allocated a 1024 by 2048 tensor only to recompute storage bytes from dtypes and element counts that the same report had already asserted; the result could not detect a codec defect.", + "The GKN codec's 4096-square roundtrip repeated the same multi-program two-dimensional kernel and error threshold covered by divisible and tail-tile shapes; removing both tails preserves geometry, accuracy, admission, edge, and determinism coverage." + ] + }, + { + "id": "TA-1509", + "scope": "remaining smoke-named and shape-named model, operator, and distributed reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "DeepSeek-V4 attention checks both window and compressed-KV forward and backward, every trainable gradient, FP8-QAT dispatch, sink dtype transfer, and TP rejection rather than shapes alone.", + "FP8 DeepEP uniquely composes no-permute transport with clamped-SwiGLU, native activation, expert biases, grouped FP8 backward, and all gradients; BF16 stochastic reduction separately proves custom-hook installation and FSDP2 gradient agreement." + ] + }, + { + "id": "TA-1510", + "scope": "standalone GatedDeltaNet FlashQLA environment-dispatch report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report monkeypatched the FlashQLA chunk function and asserted one call plus input and output shapes, making it a backend-selection branch rather than a numerical GDN policy.", + "Its assertions now close the CPU FlashQLA selection and exact-contract precedence policy; real CUDA numerical, state-chaining, and batch-invariance gates remain separate." + ] + }, + { + "id": "TA-1511", + "scope": "standalone runtime-rank MoE LoRA scaling report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both reports exercised WeightSyncHandler's inference-buffer construction, while the standalone file stopped at a hard-coded active-rank delta and a three-name buffer.", + "The active-rank scaling, emitted names, values, dtypes, shapes, and source cleanup now close the broader QLoRA merge and FP8 sync lifecycle; the standalone file is removed." + ] + }, + { + "id": "TA-1512", + "scope": "standalone families-v2 selector report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "GLM exact selection and nonexact environment rollback are properties of model-builder structure, while Qwen's v1 pin is applied by its exact-model hook; the selector-only report called private setters without either production consumer.", + "Both legacy environment aliases now close the trainer model-selection policy, and the Qwen hook proves its v1 pin overrides the legacy v2 request; CUDA norm reachability and numerical-tree reports remain distinct." + ] + }, + { + "id": "TA-1513", + "scope": "remaining small BI, DeepEP, DSV4, checkpoint, loss, and server reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The remaining small files protect distinct behavior: DeepEP async-combine safety and internode preflight are different failure modes, DSV4 rotation fallback owns optional-kernel-free orthonormality, and MTP checkpoint remapping has no broader GLM4 loader owner.", + "Token-loss composition, gradient-accumulation group routing, batch-slice topology mapping, OPD layer-cache slicing, KKT contract geometry, and families-v2 norm reachability each exercise a separate production consumer or dispatch boundary rather than construction metadata alone." + ] + }, + { + "id": "TA-1514", + "scope": "family-specific ModelRunner LoRA target-resolution reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The GLM and Kimi files rebuilt large model configurations even though ModelRunner reads only the top-level model_type, and both repeated the same explicit-target precedence branch.", + "One compact cross-family policy now retains GLM defaults, Kimi defaults including lm_head, explicit targets, and manifest targets; one duplicate branch and more than one hundred lines of irrelevant fixture data are removed." + ] + }, + { + "id": "TA-1515", + "scope": "standalone distributed-checkpointer process-group selector report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The selector report called private helpers with synthetic distributed namespaces, while the retained distributed-checkpointer I/O policy already owns load and save process-group routing.", + "NCCL-to-Gloo selection, one-time caching, native-Gloo reuse, non-pipeline omission, and pipeline custom/default metadata groups now close that I/O policy; the standalone file is removed." + ] + }, + { + "id": "TA-1516", + "scope": "standalone server batch-slice rank helper report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The dispatcher policy already exercises distinct EP slices, CP sharing, EP-FSDP coordinates, padding, routing side payloads, and the rollback switch through real batch selection.", + "The remaining replicated-DP and direct helper mappings now close that dispatcher policy, with environment restoration verified by the subsequent cases; the standalone file is removed." + ] + }, + { + "id": "TA-1517", + "scope": "remaining mock-heavy checkpoint, rendezvous, protocol, and model reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "NCCL rendezvous fault injection protects port rotation and bind-before-inference ordering; checkpoint restore protects zero-meta admission and base-before-adapter initialization; protocol round trips reject pickle and preserve tensors.", + "Gradient-checkpoint gating, native-EP combine, LoRA target manifests, sparse-delta artifacts, and exact factor publication each retain numerical, persistence, security, or production-lifecycle outcomes beyond captured argument plumbing." + ] + }, + { + "id": "TA-1518", + "scope": "test-only kernel-variant comparison API and assertion tail", + "decision": "remove", + "status": "applied", + "evidence": [ + "compare_kernel_variants had no runtime, CLI, documentation, or example consumer; its only caller computed a speedup from two literals after the production ranker had already ordered the same rows.", + "The retained rank policy still proves that a faster unvalidated candidate cannot displace the validated winner, which is the simulator's actual promotion contract." + ] + }, + { + "id": "TA-1519", + "scope": "test-only NVFP4 export dequantizer and duplicate numerical assertions", + "decision": "remove", + "status": "applied", + "evidence": [ + "dequantize_nvfp4_export was imported only by the exporter test, so production shipped an inverse implementation solely to grade its own quantizer twice.", + "The independent fake-quant policy retains exact numerical-reference coverage; the exporter policy retains packed bytes, scale shapes and dtypes, fused shared scales, BF16 islands, activation scales, directory metadata, and requantization rejection." + ] + }, + { + "id": "TA-1520", + "scope": "orphaned simulator reference_counter_total_flops helper", + "decision": "remove", + "status": "applied", + "evidence": [ + "The helper described itself as test support but had no callers anywhere in source, tests, documentation, examples, or scripts after simulator-policy consolidation.", + "Removing the 42-line adapter also removes its sole SimpleNamespace dependency without changing the analytical ledger or trainer FLOPs counter." + ] + }, + { + "id": "TA-1521", + "scope": "remaining low-reference cache, profiling, export, and weight-sync APIs", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Sparse-delta baseline reset is an explicit operational hook, the FP8 error recorder is called by FP8Linear, and dense-buffer rank filtering is invoked dynamically by WeightSyncHandler.", + "Teacher-store preparation and QARL export are public package or CLI surfaces; fused-expert cache invalidation and Mooncake hidden-store methods own runtime state rather than serving only as assertion oracles." + ] + }, + { + "id": "TA-1522", + "scope": "test-only GLM exact LM-head group binding and factor-view APIs", + "decision": "remove", + "status": "applied", + "evidence": [ + "Production supplies the TP group to the constructor and converts FP32 factor masters inside both real autograd functions; neither public-looking helper had a runtime caller.", + "The test now constructs the real object with its group, while the retained custom-boundary policy directly proves the actual autograd function saves the BF16 factor bytes." + ] + }, + { + "id": "TA-1523", + "scope": "GLM routed-expert trace-only hook path and public factor-buffer wrapper", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "sampler_value_trace installed clone-heavy wrapper hooks and returned a test-specific dataclass through a branch production always disabled; physical_factor_buffers only converted masters before calling the real internal builder.", + "The trace machinery and wrapper are removed; the GPU policy now compares real production forwards with zero and live LoRA factors and proves routing-scale linearity through the actual module path." + ] + }, + { + "id": "TA-1524", + "scope": "test-only GLM shared-expert physical-factor convenience view", + "decision": "remove", + "status": "applied", + "evidence": [ + "The public wrapper had no production consumer and only duplicated the BF16 conversion performed before the runtime-owned _physical_factor_views_from_effective builder.", + "CPU SGLang slice parity and the official CUDA operand policy now exercise that same builder directly, preserving physical layout and byte assertions." + ] + }, + { + "id": "TA-1525", + "scope": "test-only IndexShare context-manager convenience path", + "decision": "remove", + "status": "applied", + "evidence": [ + "The model uses begin plus finish_forward in its own try/finally and never calls the context-manager wrapper; only the unit helper exercised that alternate lifecycle.", + "The retained policy now drives begin and finish_forward exactly as production does for failed and successful forward-only invocations." + ] + }, + { + "id": "TA-1526", + "scope": "remaining source symbols with test-dominant lexical references", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Runner and RequestProcessor handlers are reached through command dispatch, API endpoint functions are framework-registered, and dense rank filtering is selected dynamically by WeightSyncHandler.", + "DataLoader collator mutation is a documented extension surface, while adapter transactions, exact-factor ownership, and cache invalidation participate in runtime lifecycle contracts despite sparse direct-name references." + ] + }, + { + "id": "TA-1527", + "scope": "generic, dense-Qwen3.5, and MoE-Qwen3.5 RMSNorm policies", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The generic fused policy owns kernel forward, backward, bitwise, packed-shape, and trunk behavior, while the family policy owns explicit dispatch admission and fail-closed structure.", + "Dense and MoE Qwen3.5 use separate module implementations and call sites; each policy executes its own construction and dispatch rather than repeating inputs against one owner." + ] + }, + { + "id": "TA-1528", + "scope": "remaining shape, dtype, finiteness, registry, and configuration reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The apparent weak-output reports also contain numerical reference, gradient, cache, checkpoint, or cross-backend assertions that simple assert-shape classification missed.", + "Compressor capacity, GKN format, ragged batching, RoPE bytes, and model-family construction each protect a distinct runtime consumer." + ] + }, + { + "id": "TA-1529", + "scope": "split OLMo-2 QK-norm and full tensor-parallel subprocess reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both reports initialized the same two-rank Gloo CPU tensor-parallel mesh and exercised LocalAxisRMSNormShard with Olmo2QKRMSNorm.", + "The surviving full-model subprocess now begins with plain and sharded numerical RMSNorm oracles, then applies the production TP plan and proves forward, lm-head, and all-parameter backward execution." + ] + }, + { + "id": "TA-1530", + "scope": "standalone FlashAttention diagnostic decode causal-flag report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The two boolean branches exist only to support diagnostic decode-cache behavior and do not describe an independent product policy.", + "Both causal-flag assertions now open the retained Qwen3-MoE natural and routing-replay cached-forward parity lifecycle." + ] + }, + { + "id": "TA-1531", + "scope": "native-FP8 materialization exact Torch-wheel and ambient-future assertions", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The dependency lock already pins Torch 2.12.1+cu132, so repeating the wheel string inside a behavior test made compatible environment changes fail before materialization ran.", + "The worker now explicitly disables swap_module_params_on_conversion to select the replacement path that exposed the regression, then proves plain and two-rank FSDP2 frozen-state behavior." + ] + }, + { + "id": "TA-1532", + "scope": "split MoE block, decoder, and full-model torch.compile reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both reports shared the same CUDA capability gate and compiler-compatibility owner; the full-model report only continued the lower-level block and decoder sequence.", + "The surviving policy still runs every available MoE backend through AOT eager and Inductor at lower levels, then runs native and eager compiled layers through full-model forward and backward." + ] + }, + { + "id": "TA-1533", + "scope": "remaining multi-report numerical and lifecycle files", + "decision": "keep", + "status": "accepted", + "evidence": [ + "CPU and CUDA RMSNorm, Z-loss, non-gated MoE, and native-FP8 reports require separate capability outcomes so CPU coverage is not hidden behind a GPU skip.", + "FP64 gradcheck, forward and backward sparse-MLA kernels, routing wire decode versus context layout, and enabled versus disabled packing own distinct numerical or runtime failure meanings." + ] + }, + { + "id": "TA-1534", + "scope": "unscheduled sparse-delta trainer-to-SGLang external integration report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report required both an unpinned delta-encoding source tree and an SGLang sparse-delta receiver absent from the pinned submodule; no workflow, script, or configuration supplies either test path.", + "A catch-all converted every missing or broken external implementation into a skip, so the 522-line fake trainer and orchestrator never established repository coverage; retained policies own XORL artifact encoding, source capture, hashing, posting, and transport lifecycles." + ] + }, + { + "id": "TA-1535", + "scope": "standalone TokenPartial component report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Caller-scaled denominator, microbatch additivity, raw-sum, sequence-mean-token-sum, and empty-mask behavior are branches of the shared loss reducer contract rather than an independent product policy.", + "Every direct reducer assertion now closes the retained policy and importance-sampling loss identity report, eliminating one file and pytest identity without losing an oracle." + ] + }, + { + "id": "TA-1536", + "scope": "remaining pinned MoE and TileLang availability exception shields", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The MoE compiler helper broadly suppressed failures importing a shipped availability utility and wrapped an infallible Quack list append; the sparse-MLA reports treated hard-pinned TileLang as optional.", + "Explicit CUDA and Hopper admission remains, while internal dependency failures now fail closed; the real MoE compile matrix and both TileLang sparse-MLA policies pass." + ] + }, + { + "id": "TA-1537", + "scope": "NVSHMEM library-path catch-all helpers", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "All three callers already use pytest.importorskip for nvidia.nvshmem before launching their DeepEP workers, so swallowing every subsequent path-resolution error only hid malformed installations.", + "The helpers now resolve the admitted package directly; all affected reports collect and the installed package path is importable." + ] + }, + { + "id": "TA-1538", + "scope": "standalone runtime FLOPs context-parallel denominator report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The runtime counter and simulator analytical ledger share the same global-sequence-length numerator contract; context parallelism must change placement rather than multiply global work.", + "The exact cp1-versus-cp64 assertion now closes the retained topology, shape, and analytical-ledger policy instead of occupying a one-test utility file." + ] + }, + { + "id": "TA-1539", + "scope": "standalone gradient-accumulation loss and HSDP deferral reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Loss-group forwarding, local/global token normalization, and backward scaling are the consumer half of the trainer metadata-counting policy.", + "HSDP microbatch deferral and restoration are branches of the explicit SP and LM-head gradient-synchronization policy; all assertions remain in those two owner reports." + ] + }, + { + "id": "TA-1540", + "scope": "standalone private OPD layer-cache fetcher report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Index filtering, streamed layer slices, and layer-count reporting exist only to feed the retained OPD microbatch loss lifecycle.", + "The exact requested indices, slice ranges, and output shapes now execute before the real streaming loss, gradient, cache, metric, and debug-artifact assertions." + ] + }, + { + "id": "TA-1541", + "scope": "standalone ModelRunner LoRA target-resolution and kill-session reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Implicit family targets, explicit lists, and manifest precedence now close the runner's adapter ownership compiler policy that consumes the selected targets.", + "Nonresident checkpoint promotion, failed-kill preservation, registry cleanup, and path rejection now close the existing optimizer, checkpoint-load, and session-registry lifecycle." + ] + }, + { + "id": "TA-1542", + "scope": "standalone sync-quantization dictionary normalization report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "BF16 no-ops, valid FP8 normalization, module exclusions, and malformed or unsupported forms are admission branches of receiver quantization detection and enrichment.", + "All direct normalization assertions now run in the API policy that detects receiver configuration, propagates unsupported markers, enriches user input, and persists the default." + ] + }, + { + "id": "TA-1543", + "scope": "standalone IndexShare trainer and server caller-cleanup report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Offline and server failure cleanup are the public-caller completion branches of the same IndexShare lifecycle that owns forward-only, backward-retained, checkpoint-recompute, and idempotent close behavior.", + "Both real caller wrappers and their exact release counts now close the retained checkpointed lifecycle policy." + ] + }, + { + "id": "TA-1544", + "scope": "standalone rank-zero ready-handshake report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "ACK handling, request-before-ACK queueing, client identity, malformed frames, and message preservation are runtime branches of the orchestrator-runner wire protocol.", + "The async handshake now follows serialization, tensor roundtrip, command creation, and pickle rejection in one protocol owner." + ] + }, + { + "id": "TA-1545", + "scope": "standalone cross-rank packed-padding synchronization report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Tensor padding, ignored labels, attention masks, cumulative sequence boundaries, and max-length updates are the distributed completion of server-versus-CLI packed sequence metadata alignment.", + "The exact 176-to-512 cross-rank case now runs alongside local padding, SP sharding, stale-metadata replacement, LCM admission, and unpacking behavior." + ] + }, + { + "id": "TA-1546", + "scope": "standalone lm-head TP plus EP mesh-membership subprocess", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The mesh-only report stopped after constructing four-rank TP, replica, and EP groups, while the retained lm-head FSDP policy already used the same DP2 by CP2 topology for real loss and gradients.", + "That full transaction now enables EP2, asserts exact group membership, and continues through parameter sync, vocab-sharded loss, global loss parity, weight gradients, and hidden gradients." + ] + }, + { + "id": "TA-1547", + "scope": "server runner reports that execute source-loaded production module copies", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Seven reports loaded ModelRunner, RunnerDispatcher, AdapterCoordinator, CheckpointManager, or LoRAAdapterManager from file paths under synthetic module names even though the canonical package modules import successfully.", + "The reports now patch and exercise the actual runtime module objects, preserving all 11 adapter, checkpoint, optimizer, session, and dispatcher assertions while removing duplicate module state." + ] + }, + { + "id": "TA-1548", + "scope": "source-loaded launcher copy and fake dependency graph in server-argument policy", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The report imported the canonical launcher for override parsing but re-executed the same source file with fake API-server, orchestrator, session, QARL, and packing modules solely to obtain load_server_arguments.", + "All four server-argument reports now call the canonical launcher and real dependency graph, including shipped-config subprocess parsing and sparse-MLA propagation." + ] + }, + { + "id": "TA-1549", + "scope": "smallest remaining operator, model, and distributed report boundaries", + "decision": "keep", + "status": "accepted", + "evidence": [ + "DSV4 fallback rotation, GLM4 MTP checkpoint remapping, DeepEP async-combine admission, KKT launch geometry, families-v2 dispatch, BI mean, and Class-B RoPE each reach a different production boundary rather than repeating size-only examples.", + "These reports retain numerical or fail-closed behavior not owned by their adjacent kernel, model, or topology policies; line count alone is not a deletion signal." + ] + }, + { + "id": "TA-1550", + "scope": "private ring-attention zigzag helper report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report mixed assertions on the private _get_zigzag_step_section implementation with the public packed-sequence reorder consumed by TextSequenceShardCollator.", + "The public single-document, packed-document, multi-rank, identity, and invalid-length behavior now closes the collator policy, while the private section-helper assertions are removed." + ] + }, + { + "id": "TA-1551", + "scope": "optional FA3 and SGLang exact-runtime reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "FA3 ring-attention merging and shared-prefix attention still exercise production numerical kernels, while the GLM exact SGLang joins cross real checkpoint export, adapter parsing, and memory-pool boundaries.", + "Their environment contract is intentionally isolated: the default XoRL profile remains Torch 2.12.1 without sglang-kernel, and pinned SGLang declares Torch 2.11.0 with sglang-kernel 0.4.5. Lazy wrapper imports are not ABI validation." + ] + }, + { + "id": "TA-1552", + "scope": "source-loaded EP routing-score report and synthetic backend registration graph", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The routing-score report rebuilt Triton and Quack from source under test-only names with fake kernel packages, while the EP adapter policy already owned expert_scores forwarding and backend registration.", + "Its complete forward and routing-score gradient oracle now closes the canonical EP adapter policy; the routing-position policy shares only explicit CPU kernel doubles, and the standalone report is removed." + ] + }, + { + "id": "TA-1553", + "scope": "source-loaded Quack process, compiler, and cache module copies", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The Quack safety policy executed worker protocol, ptxas, and cache files under synthetic module names with fake cutlass and tvm_ffi modules even though their canonical package imports succeed.", + "Timeout, truncation, temporary-output, entry-point, and safe cache-key assertions now run against the actual runtime module objects." + ] + }, + { + "id": "TA-1554", + "scope": "synthetic Torch FSDP module in HSDP gradient-sync policy", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The policy replaced torch.distributed._composable.fsdp.fully_shard in sys.modules solely to make a fake object pass the FSDPModule type gate.", + "It now uses a lightweight instance of Torch 2.12's real FSDPModule API type and preserves the complete deferral, last-microbatch, restoration, and replicate-size assertions." + ] + }, + { + "id": "TA-1555", + "scope": "fake FA4 package graph and runtime module reloads in attention registry policy", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The default profile already installs FA4 without flash_attn_interface, but the policy replaced flash_attn modules and reloaded both production registry modules to recreate that state.", + "Registry admission and resolution now run against the installed canonical FA3 or FA4 availability and retain the eager fallback and unavailable-flash failure boundaries." + ] + }, + { + "id": "TA-1556", + "scope": "split OPD driver version-verification and payload-transport reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Student endpoint matching and weight-version success, mismatch, and failure branches are setup and publication checks of the same OPD pipeline that owns worker preparation, causal payloads, and Mooncake metadata transport.", + "All assertions now execute in one driver lifecycle, and the standalone executable script is loaded once instead of re-executed for every helper." + ] + }, + { + "id": "TA-1557", + "scope": "remaining synthetic optional-dependency modules", + "decision": "keep", + "status": "accepted", + "evidence": [ + "delta_encoding is absent from the default environment, Mooncake cannot load without its CUDA runtime, and the default Torch 2.12 lane intentionally lacks the pinned SGLang and sgl_kernel runtime.", + "The retained doubles exercise explicit serialization, fallback, runtime-context, and slot-combine contracts only; they are not used as evidence that the compiled dependency ABI or real kernel operation works." + ] + }, + { + "id": "TA-1558", + "scope": "conditional skips for shipped EP and non-gated MoE backends", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The default profile ships Triton, native Torch, and Quack support, but the EP adapter report could skip midway and the non-gated report tested whichever subset happened to register.", + "Both policies now fail closed if a shipped backend disappears and execute every expected forwarding, rejection, forward, and gradient branch." + ] + }, + { + "id": "TA-1559", + "scope": "duplicate three-GPU skip inside OPD full-pipeline report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The module-level marker already skips unless three CUDA devices are visible before fixtures or the test body run.", + "The second in-body device-count branch repeated the same admission rule after model artifacts had already been created." + ] + }, + { + "id": "TA-1560", + "scope": "remaining conditional runtime reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The remaining reports require Hopper-only FlashQLA or exact GLM kernels, two-rank FSDP, CUDA profiler events, a real SGLang and sgl_kernel lane, or optional DeepGEMM execution.", + "Each gate protects a numerical, distributed, or compiled-runtime transaction that has no ordinary CPU branch hidden behind the skip." + ] + }, + { + "id": "TA-1561", + "scope": "tiny dense QARL AdamW training smoke", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report only proved that one ordinary AdamW step changes a tiny model's weights and logprobs, behavior owned by PyTorch rather than QARL.", + "QARL injection and summaries remain in the fake-quant policy, while forward state and exact state-dict restoration remain in the calibration lifecycle." + ] + }, + { + "id": "TA-1562", + "scope": "standalone QARL activation-quant override report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The override changes the same QARLLinear and QARLMoEExperts fake-quant state already constructed by the dense fake-quant owner.", + "Enable, disable, per-module restoration, exception safety, non-QARL exclusion, and nested restoration now close that owner lifecycle without a separate report." + ] + }, + { + "id": "TA-1563", + "scope": "standalone Qwen3.5 families-v2 backward candidate report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report replaced both numerical kernels with CPU references and tested only the zero-centered autograd wrapper wiring.", + "That wiring now closes the Qwen3.5 norm dispatch and site-assignment owner, preserving effective-weight and dual residual-gradient parity." + ] + }, + { + "id": "TA-1564", + "scope": "standalone LoRA mixed-precision model-builder report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report patched the same build_training_model construction boundary as the existing FP8 and QARL model-builder lifecycle.", + "Base BF16 retention, FP32 adapter factors, trainability, dtype resolution, and generic-upcast admission now execute in that owner policy." + ] + }, + { + "id": "TA-1565", + "scope": "test-only canonical-MoE sampler plan and serialization surface", + "decision": "remove", + "status": "applied", + "evidence": [ + "ParallelRole.SAMPLER, glm52_sampler, launcher_tp_size, logical_ordinal, as_dict, and digest had no production, documentation, example, or launcher consumer; only the canonical-MoE test called them.", + "The retained trainer plan still validates exact group membership, logical ordinals, topology rejection, and the real distributed collective contract." + ] + }, + { + "id": "TA-1566", + "scope": "test-only one-call adapter gradient capture convenience", + "decision": "remove", + "status": "applied", + "evidence": [ + "Production ModelRunner owns the two-phase stage_gradient_numerators and commit_gradient_capture transaction; capture_gradient_numerators had only test callers.", + "Tests now drive the production boundary directly, including explicit abort behavior and real multi-rank fatal paths, with setup-only repetition isolated in test helpers." + ] + }, + { + "id": "TA-1567", + "scope": "standalone DeepEP async-combine opt-in report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The unsafe async-combine environment gate is part of the same DeepEP transport admission surface as internode topology, preflight, and buffer sizing.", + "Both default-safe and explicit-opt-in branches now close the existing DeepEP admission policy without a separate report." + ] + }, + { + "id": "TA-1568", + "scope": "standalone GLM52 exact-MoE modeling forwarding report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report constructed exact routed and shared experts solely to inspect arguments forwarded by Glm5MoEBlock.", + "Routed IDs, scaling, and shared contributor-ordinal forwarding now close the exact-MoE construction and inventory owner that already validates those concrete modules." + ] + }, + { + "id": "TA-1569", + "scope": "standalone Kimi tokenizer loader report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Kimi's local tokenizer selection is part of the same model-family loading boundary as its wrapper config and registry resolution.", + "Local TikToken decoding and generic tokenizer/processor fallback behavior now close the Kimi/DeepSeek registry owner without a separate report." + ] + }, + { + "id": "TA-1570", + "scope": "standalone SignSGD builder and step report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Its no-decay parameter grouping repeated the generic builder contract already exercised by DistSignSGD, while the cautious optimizer policy already constructed and stepped SignSGD.", + "Dense sign updates, decoupled decay, missing-gradient behavior, and sparse-gradient rejection now close the retained optimizer policy." + ] + }, + { + "id": "TA-1571", + "scope": "isolated stochastic-rounding report and four-GPU BF16 all-to-all gate", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "stochastic_round_to_bf16 has one production consumer: BF16StochasticAllToAllReduceScatter, so two independent reports overstated confidence in one transaction.", + "One default-runtime policy now proves deterministic admission, neighbor distribution, unbiased expectation, and the real two-rank Gloo all-to-all with FP32 accumulation against reduce-scatter." + ] + }, + { + "id": "TA-1572", + "scope": "standalone Qwen3 projection-unfusing structure report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report carried a distributed marker but initialized no process group and asserted only module replacement shapes.", + "The production model-level unfuse path, TP plan, checkpoint-handler transition, and every layer's projection inventory now close the torch_parallelize policy owner." + ] + }, + { + "id": "TA-1573", + "scope": "duplicate balanced synthetic-routing fragment in train-router report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The fragment checked only that cached expert IDs survived regather and weights became uniform under the balanced environment switch.", + "The retained TopKRouter owner already proves the exact cyclic expert sequence, balance bound, uniform weights, and override of softmax, hash-table, and bias inputs." + ] + }, + { + "id": "TA-1574", + "scope": "standalone NVFP4 QARL linear report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report constructed the same QARLLinear and normalized the same recipe surface as the dense fake-quant owner, differing only by format.", + "NVFP4 group admission, weight-only forward, straight-through gradient, and disabled-weight behavior now close the format-spanning QARL policy." + ] + }, + { + "id": "TA-1575", + "scope": "standalone MoE sqrtsoftplus regather report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Routing replay regather is a continuation of TopKRouter's softmax, sqrtsoftplus, selected-expert, dtype, and scaling contract.", + "Both regather branches now close the router owner that already proves selection bias, hash routing, normalization, and configured FP32 execution." + ] + }, + { + "id": "TA-1576", + "scope": "standalone families-v2 norm reachability report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report invoked the same v2 RMSNorm tree and payload as the numerical, realization, and dispatch policy, then inspected only trainer entry-point reachability.", + "Trainer dispatch and the v1 kill switch now close the v2 norm owner under isolated monkeypatch contexts." + ] + }, + { + "id": "TA-1577", + "scope": "standalone families-v2 LM-head report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Projection, selected-logprob, batch-invariance, backward, and kill-switch assertions target the same final-token probability transaction as the bi_fused LM-head policy.", + "The v2 realization now closes that owner alongside eager parity, temperature, determinism, guards, and probability-boundary behavior." + ] + }, + { + "id": "TA-1578", + "scope": "standalone batch-invariant mean interpose regression", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The regression exercises set_batch_invariant_mode's global Torch interpose, already owned by the trunk-linear forward and gradient-admission policy.", + "Full, typed, one-dimensional, keepdim, and multi-dimensional reductions now close that interpose owner while preserving the former sum-versus-mean bug oracle." + ] + }, + { + "id": "TA-1579", + "scope": "standalone DSV4 rotate-activation fallback report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "rotate_activation is consumed by the DSV4 compressor and indexer, and the standalone report only disabled the optional fast transform to inspect its fallback.", + "Fallback basis mapping, involution, and norm preservation now close the compressor's context-parallel shape and admission policy." + ] + }, + { + "id": "TA-1580", + "scope": "standalone Qwen3-MoE RMSNorm family declaration report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Dense Qwen and shared-attention family declarations already belong to the generic RMSNorm family contract.", + "Qwen3-MoE layer-zero, residual-tree, explicit call-site, and bare final-norm behavior now close that same structural owner." + ] + }, + { + "id": "TA-1581", + "scope": "standalone LoRA permanent cast-once merge report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report rebuilt the same LoraLinear, MoEExpertsLoRA, factor deltas, and canonical FP32 fold owned by the merged-forward contract.", + "Zero-adapter preservation and permanent BF16/FP16 cast-once merge behavior now close the canonical fold owner on CPU instead of requiring CUDA unconditionally." + ] + }, + { + "id": "TA-1582", + "scope": "standalone dataset hash and split-fingerprint reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both helpers exist to key the dataset split, loading, saving, and preparation lifecycle already exercised by the shared-data owner.", + "Determinism, split separation, input sensitivity, fractional sizes, multi-dataset order independence, and tokenizer/column sensitivity now close that lifecycle." + ] + }, + { + "id": "TA-1583", + "scope": "standalone virtual-stage MultiOptimizer report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report's end-to-end assertion was that build_lr_scheduler wraps every child optimizer and decays all parameter groups after delegated steps.", + "Multi-part construction, delegation, model mapping, single-part fallback, invalid explicit groups, and scheduler fanout now close the scheduler owner; checkpoint state filtering remains separately covered." + ] + }, + { + "id": "TA-1584", + "scope": "standalone QARL-to-FP8 weight-sync report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report derived an FP8 sync configuration from QARL, invoked WeightSyncHandler's quantizer, and exercised its request handler.", + "Folded-module metadata, skip-list behavior, derived request configuration, quantized buffers, and incompatible overrides now close the FP8 weight-sync owner." + ] + }, + { + "id": "TA-1585", + "scope": "synthetic identity-layer gradient-checkpoint truth table", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report attached MagicMock checkpoint functions to an identity layer and exhaustively restated the outer Python gate's training, flag, and method condition.", + "Real Nemotron-H training already proves default full-layer checkpoint execution and gradients, while the GLM-5 lifecycle proves recompute-before-dispatch bypasses the outer checkpoint and invokes the layer checkpoint." + ] + }, + { + "id": "TA-1586", + "scope": "standalone generic train-router dispatch report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Trainable and frozen router-gradient behavior is already exercised by routing-replay and real model lifecycles; the standalone tiny MoE repeated those backward assertions.", + "The unique DeepEP rejection now closes TopKRouter's MoEBlock configuration policy, and the frozen server default now closes server configuration serialization." + ] + }, + { + "id": "TA-1587", + "scope": "standalone mocked GDN KKT launch-geometry report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report replaced both KKT kernels with launch recorders solely to inspect the same gdn_contract switch owned by the GDN exact-contract policy.", + "Pinned BK, warp, stage, safety, and autotuned-off-lane behavior now close the existing GDN serving-geometry owner beside its solve-tril geometry." + ] + }, + { + "id": "TA-1588", + "scope": "standalone DSV4 attention shape smoke report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report directly constructed and manually initialized isolated C0 and C128 attention layers, then repeated forward shape, finiteness, and backward reachability already covered by the fully initialized DSV4 model lifecycle.", + "The unique FP8-QAT dispatch and TP rejection now run through the full model owner; direct-component dtype behavior that production model casting deliberately overrides was removed." + ] + }, + { + "id": "TA-1589", + "scope": "standalone synthetic LoRA target-manifest report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report built a fake two-layer attention tree even though the fused GDN lifecycle already consumes strict target manifests for real fused projection paths.", + "Count, rank, configured-target, unlisted-module, Boolean, schema, and integer validation now fail closed against the real fused-GDN manifest owner." + ] + }, + { + "id": "TA-1590", + "scope": "standalone GLM4 MTP tail-remap fragment", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report instantiated only Glm4MoeCheckpointHandler and checked three MTP tail aliases plus three ignored tail fields.", + "Those aliases and exclusions now close the GLM4 model-family construction and checkpoint lifecycle that already creates both ordinary and prequantized handlers." + ] + }, + { + "id": "TA-1591", + "scope": "remaining orchestrator and API communication reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The small reports exercise different live boundaries: runner rank-zero readiness and tensor-safe serialization, API-engine ZMQ request handling, and APIServer response metrics.", + "Their mocks isolate external processes but do not duplicate the protocol transaction or its failure modes, so combining them would obscure ownership rather than remove repeated behavior." + ] + }, + { + "id": "TA-1592", + "scope": "standalone importance-sampling and policy-loss microbatch reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both one-test files rebuilt the same masked tensors, TokenPartial denominator, full-batch call, microbatch calls, and summable-metric loop already parameterized by the shared loss-contract owner.", + "The shared owner now checks legacy identity and microbatch composition together for basic, KL, TIS, and IcePop modes; both copied reports are removed." + ] + }, + { + "id": "TA-1593", + "scope": "standalone dense-Qwen eager RoPE parity report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The report's remaining assertion compared RotaryEmbedding table consumption and Q/K application with an inlined serving arithmetic reference.", + "That exact CUDA oracle now closes the existing RoPE frequency-table, lazy-cache, serving-device, and zero-K3 lifecycle; the standalone owner is removed without dropping its bitwise checks." + ] + }, + { + "id": "TA-1594", + "scope": "deprecated adapter optimizer broadcast no-op and negative spies", + "decision": "remove", + "status": "applied", + "evidence": [ + "AdapterCoordinator.broadcast_adapter_optimizer_state had no production caller and deliberately performed no optimizer transfer; topology-specific optimizer state is restored through the all-ranks checkpoint path.", + "The method is removed together with seven lifecycle spies whose only assertion was that the dead no-op stayed uncalled; the real transactional optimizer-restore rejection remains covered." + ] + }, + { + "id": "TA-1595", + "scope": "standalone trainer-P2P and PP-NCCL handler reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Both files directly bound or called private WeightSyncHandler helpers for HCA selection, abort markers, peer status gathering, and PP named-tensor flatten/reconstruction.", + "Those policies now close the existing handler configuration, sender-selection, bucketing, and inference-layout owner; both one-test modules are removed while their complete cases remain." + ] + }, + { + "id": "TA-1596", + "scope": "deprecated ep_outside and moe_checkpoint_method configuration aliases", + "decision": "remove", + "status": "applied", + "evidence": [ + "Neither alias appears in shipped configurations or current documentation; ep_intranode and gradient_checkpointing_method are the documented native fields.", + "The aliases, parser remapping branches, duplicated simulator dimension, and compatibility-only parser inputs are removed. The active gradient_checkpointing_method='moe_act' execution mode remains available." + ] + }, + { + "id": "TA-1597", + "scope": "standalone GLM52 exact-attention construction lifecycle", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Its rank-alpha, dense-component, sparse-MLA, and all-to-all admission failures duplicated the complete exact-MoE constructor's cases.", + "The unique 780 attention-factor names, projection classes, source FQNs, per-layer trainable sets, and three dense roots now close the complete 1,700-factor constructor; the standalone report is removed." + ] + }, + { + "id": "TA-1598", + "scope": "synthetic prequantized GNK-to-GKN transpose report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report never invoked prequantized loading; it quantized tensors in both orientations and re-derived transpose equivariance for block-FP8 and NVFP4.", + "The retained QLoRA expert-loader owner exercises the real GNK-to-GKN byte and scale transforms for both formats and multiple shapes, while the codec owners cover numerical roundtrips." + ] + }, + { + "id": "TA-1599", + "scope": "standalone synthetic GKN format and backend report", + "decision": "remove", + "status": "applied", + "evidence": [ + "Its checkpoint half rebuilt ExpertWeightBuffer behavior already exercised through the DeepSeek-V3 checkpoint handler's exact load/save layout contract.", + "Its eager/native comparisons duplicate the dedicated forward/backward backend owner, while grouped-GEMM primitives and combined Quack/SGLang parity own the Triton path; the report's optional imports could silently bypass that branch." + ] + }, + { + "id": "TA-1600", + "scope": "parallel dense and MoE Qwen3.5 RMSNorm reports", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The two files independently enumerated the same copied zero-centered dispatch matrix, v2 admission policy, norm-site propagation, and residual-family selection.", + "One family-wide owner now drives both production classes through the shared CPU contract, retains the dense-only GDN and v2 backward boundaries, and uses one GPU model lifecycle for their shared normalization kernels." + ] + }, + { + "id": "TA-1601", + "scope": "SGLang compiled-kernel dependency and smoke contract", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The default Torch-2.12 profile intentionally excludes sglang-kernel, while pinned SGLang requires Torch 2.11.0 and sglang-kernel 0.4.5 in an isolated environment.", + "The retained smoke fails when the exact lane lacks the package, eagerly loads the compiled extension and wrapper symbols, and executes a real RMSNorm operation; it passes in the new .venv-sglang and skips in the default profile." + ] + }, + { + "id": "TA-1602", + "scope": "legacy SGLang fused-expert cache environment alias", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_MOE_SGLANG_FUSED_EXPERTS_CACHE_WEIGHTS appeared only in implementation and tests; no shipped configuration or documentation used it.", + "XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE=cached remains the sole cache policy, with reuse, invalidation, transient, and zero-copy strided coverage retained." + ] + }, + { + "id": "TA-1603", + "scope": "legacy EP-wide duplicate server batch rollback mode", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_SERVER_EP_DUPLICATE_BATCHES was an undocumented rollback switch for ep_size-times redundant server compute; no shipped configuration selected it.", + "The correct per-rank EP slice mapping remains covered across EP, CP, padding, routing payload, and OPD teacher-cache assembly, while compatibility-only duplicate-broadcast assertions are removed." + ] + }, + { + "id": "TA-1604", + "scope": "orphaned Qwen linear-attention P2P compatibility branches", + "decision": "remove", + "status": "applied", + "evidence": [ + "The disabled fused QKV slicer produced a combined layout for which pinned SGLang has no receiver locator; its env gate, dimension plumbing, fake bypass assertion, and dead implementation are removed in favor of canonical locator slices.", + "Cold-prepare invalidation now has one policy boundary: cache_invalidation_mode=none disables it. The undocumented duplicate env override is removed, and cold, cached, and disabled prepare cases remain covered." + ] + }, + { + "id": "TA-1605", + "scope": "duplicate optimizer empty-cache environment override", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_SKIP_EMPTY_CACHE_AFTER_OPTIM_STEP appeared only in ModelRunner and its test, duplicating the native skip_empty_cache_after_optim_step train configuration field.", + "The optimizer-step owner now selects both active cache policies through train_config while retaining gradient scaling, clipping, mutation, synchronization, and result-metric coverage." + ] + }, + { + "id": "TA-1606", + "scope": "legacy and debug-only weight-sync environment policies", + "decision": "remove", + "status": "applied", + "evidence": [ + "Both legacy receiver post-process overrides and XORL_WEIGHT_SYNC_BUCKET_BYTES were compatibility-only; pinned P2P writes already target receiver-native FP8 storage, while KV-cache finalization is selected from endpoint requirements.", + "Direct-EP scatter's shallow/deep copy modes and legacy boolean alias only tested optional copies of immutable prepared locators. The production default reuses locators, and scatter serialization plus the retained manifest owner preserve the recipient boundary." + ] + }, + { + "id": "TA-1607", + "scope": "NeMo fp8_cfg compatibility translation", + "decision": "remove", + "status": "applied", + "evidence": [ + "No shipped XoRL configuration or documentation used fp8_cfg, while native enable_fp8_training and fp8_training_* fields already own the complete supported policy.", + "The dataclass aliases, NeMo policy extraction, normalization API, launcher remapping, and acceptance assertions are removed; the shared configuration tombstone now fails fast with the native-field migration." + ] + }, + { + "id": "TA-1608", + "scope": "tests/fp8_training/test_config_compat.py", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "Its direct external-runtime rejection matrix duplicated the public train and server parser owners after fp8_cfg translation was retired.", + "BF16 island selection now runs inside the FP8 injection owner and Blackwell rejection runs through build_training_model, eliminating the standalone file and two collected tests without losing production execution." + ] + }, + { + "id": "TA-1609", + "scope": "legacy numerical-family environment rollback switches", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_FAMILIES_V2 and SGLANG_FAMILIES_V2 duplicated the model program's structural v1/v2 selection and allowed trainer and sampler processes to drift independently.", + "Ordinary models now use v2, exact Qwen3.5 selects its qualified v1 program, and canonical GLM-5.2 selects v2. The surviving norm and LM-head tests select those production programs structurally instead of testing kill switches." + ] + }, + { + "id": "TA-1610", + "scope": "routing-weight position environment override", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_MOE_ROUTING_WEIGHTS_BEFORE_DOWN duplicated the model and server moe_routing_weights_before_down configuration and could silently override the resolved model program per process.", + "The numerical forward/backward oracle and auto, explicit, router-training, dispatch, and SGLang-parity policies remain in the production configuration owner; only the redundant lazy environment assertion is removed." + ] + }, + { + "id": "TA-1611", + "scope": "tests/qarl/test_calibration.py", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The standalone report mostly enumerated private JSON and JSONL loader shapes around a synthetic model, while the training-model owner already executes the real calibration lifecycle before parallelization.", + "Persistent QARL calibration state is retained in the dense fake-quant owner; the production builder still proves real batch loading, observer population, ordering, and calibrated-module counts." + ] + }, + { + "id": "TA-1612", + "scope": "tests/server/test_side_payloads.py", + "decision": "remove", + "status": "applied", + "evidence": [ + "The standalone fake-store report repeated Mooncake metadata encoding and error permutations outside the request and dispatcher lifecycles that own the feature.", + "The retained request processor writes and cleans R3 references on success and failure, and the retained dispatcher loads only the rank-local packed slice in production order." + ] + }, + { + "id": "TA-1613", + "scope": "DSV4 RoPE cache environment override", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_DSV4_ROPE_MAX_SEQ_LEN was a test and profiling override for a value already owned by config.max_position_embeddings, creating two cache-capacity authorities.", + "DSV4 model, LoRA, compressor, and context-parallel capacity owners now select explicit model configuration; the production capacity guard still fails loudly when that configured cache is too short." + ] + }, + { + "id": "TA-1614", + "scope": "tests/data/collators/test_tensor_collator.py permutation matrix", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "After earlier count consolidation, the single test still hid a broad matrix of scalar, boolean, string, dimensionality, and batch-size permutations for a mechanical converter.", + "One production-shaped contract now covers the four supported pipeline forms: flat features, already-batched dictionaries, nested packed features, and empty input, including dtype and tensor-passthrough boundaries." + ] + }, + { + "id": "TA-1615", + "scope": "SGLang EP slot-combine experiment and fake-kernel assertions", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_MOE_SGLANG_FUSED_EXPERTS_SLOT_COMBINE was an undocumented, default-off, scoring-only branch with no shipped configuration or runtime owner.", + "Its assertions replaced both moe_sum_reduce and distributed transport with world-size-one fakes. The qualified SGLang EP compute, stock all-to-all combine, real extension smoke, and independent FP32-routing backward oracle remain." + ] + }, + { + "id": "TA-1616", + "scope": "tests/server/runner/test_batch_utils.py", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The standalone report called the non-packed sequence sharder with batch_size=1, although production routes batch_size=1 through TextSequenceShardCollator.", + "Its useful float-side-channel and teacher-hidden-state behavior now runs through RunnerDispatcher with a real two-row batch, ragged conversion, and CP sharding; the data-collator owners retain packed behavior." + ] + }, + { + "id": "TA-1617", + "scope": "duplicate and unowned dispatcher diagnostic/dummy environment branches", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_MICROBATCH_DIAGNOSTIC_TENSORS duplicated diagnostic_microbatch_dump_tensors, and XORL_MICROBATCH_DIAGNOSTIC_DIR duplicated diagnostic_microbatch_dump_dir.", + "XORL_SERVER_MINIMAL_DUMMY_BATCH_TOKENS selected a second undocumented dummy constructor with no tests or callers. Diagnostics now use request parameters and padding uses the single retained zero-loss dummy lifecycle." + ] + }, + { + "id": "TA-1618", + "scope": "malformed TensorData compatibility assertions", + "decision": "remove", + "status": "applied", + "evidence": [ + "The API-type report treated mismatched and zero-sized tensor metadata falling back to a flat list as a compatibility promise, despite those inputs being invalid and prone to downstream misclassification.", + "The retained production-shaped contract covers rank-1 token IDs plus rank-2 teacher states and rank-3 routing tensors, which are the shapes the server must preserve." + ] + }, + { + "id": "TA-1619", + "scope": "test_ep_trainable_grads_match_stock_triton", + "decision": "remove", + "status": "applied", + "evidence": [ + "The test was skipped in the default environment and failed when the isolated SGLang environment made it executable: stock Triton and the serving wrapper intentionally use different routing-rounding programs.", + "The custom backward preserves the serving FP32-routing boundary, so stock bitwise equality is the wrong oracle. The retained independent eager oracle proves local and EP input, routing, and weight gradients against that actual contract." + ] + }, + { + "id": "TA-1620", + "scope": "isolated SGLang exact-kernel environment resolution", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Pinned SGLang resolves Quack 0.6 and CUTLASS DSL 4.6, while this XoRL source imports the Quack 0.5 and CUTLASS 4.5 APIs; the prior smoke-only environment therefore could not collect XoRL model tests.", + "The isolated setup and uv profile now override those trainer-side packages after installing SGLang. Fresh uv resolution succeeds, the real sgl_kernel operation passes, and the FP32-routing plus DSV4 model, LoRA, and compressor owners all pass under Torch 2.11." + ] + }, + { + "id": "TA-1621", + "scope": "mocked XORL_DCP_LOAD_NO_DIST routing assertions", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The no-dist environment is a real shared-filesystem recovery mode and remains covered by the GLM exact-DCP round trip, but two ModelState helpers only captured process_group=None and no_dist=True from a fake dcp.load call.", + "Those mocks are removed. The retained ModelState owner now performs a real model-only DCP save/load while proving a requested optimizer remains untouched; pipeline custom-group ordering stays separately covered." + ] + }, + { + "id": "TA-1622", + "scope": "test-only OPD loss-family inference", + "decision": "remove", + "status": "applied", + "evidence": [ + "Every production _finalize_loss_metrics call supplies its resolved loss_fn, while one direct test omitted it and forced runtime code to guess OPD from metric-name prefixes.", + "The private finalizer now requires the production argument and the retained OPD aggregation owner uses that actual call contract." + ] + }, + { + "id": "TA-1623", + "scope": "shards plus preprocess_shards compatibility behavior", + "decision": "remove", + "status": "applied", + "evidence": [ + "Documentation declares shards and preprocess_shards mutually exclusive and no shipped configuration combines them, but the preparation test promised that shards silently won.", + "DatasetConfig now rejects the invalid combination and the generator no longer carries the compatibility-only precedence branch." + ] + }, + { + "id": "TA-1624", + "scope": "dataset merge length-only shuffle assertions", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "Both shuffle modes previously asserted only that concatenation preserved row count, so they passed if production ignored the shuffle flags entirely.", + "The retained lifecycle now distinguishes ordered concatenation, whole-merge permutation, and within-dataset permutation while proving row preservation." + ] + }, + { + "id": "TA-1625", + "scope": "data_files JSON-only routing assertions", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The string/list cases only counted Hub downloads and always selected JSON, masking that production ignored the documented ds_type field for downloaded data_files.", + "The loader now selects the configured dataset type and the retained cases verify JSON string and Parquet list routing plus their resolved files." + ] + }, + { + "id": "TA-1626", + "scope": "immediate-success retry decorator assertion", + "decision": "remove", + "status": "applied", + "evidence": [ + "The one-call success row added no behavior beyond the retained transient-success case, which already proves return-value propagation after the retry boundary.", + "Request and Hub transient classification, exponential backoff, exhaustion, and unrelated-exception propagation remain covered." + ] + }, + { + "id": "TA-1627", + "scope": "required SGLang ABI smoke and remaining default-profile installation path", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The server-training guide still installed pinned SGLang into the active environment and advertised Torch 2.9.1, contradicting the Torch-2.12 default and Torch-2.11 SGLang split.", + "Required smoke mode now fails unless CUDA is available for a real compiled RMSNorm operation and reports both import-loader error forms; the isolated Torch 2.11.0 and sglang-kernel 0.4.5 lane passes while the default profile contains no kernel wheel." + ] + }, + { + "id": "TA-1628", + "scope": "standalone block-FP8 activation and tiled-weight edge matrices", + "decision": "remove", + "status": "applied", + "evidence": [ + "The two component reports enumerated random scale ranges, dimensionality, block sizes, determinism, magnitude and sign examples, plus internal assertion failures without entering a model operation.", + "The retained FP8-linear owner executes activation and tiled-weight quantize/dequantize at block sizes 64 and 128, consumes both scale layouts in a real GEMM, and now checks each codec against its original production-shaped operand; five unexported and unreferenced compatibility aliases are also removed." + ] + }, + { + "id": "TA-1629", + "scope": "test-owned sequence-parallel metric argument-order compatibility", + "decision": "remove", + "status": "applied", + "evidence": [ + "Both production callers pass metrics, process group, and metric operations in the declared order, while the retained distributed test alone used the obsolete metrics, operations, group order.", + "The type-based argument swap is removed and the real two-GPU NCCL reduction now passes through the production signature without weakening its partial-sum or extrema assertions." + ] + }, + { + "id": "TA-1630", + "scope": "private generic GLM absorbed-projection einsum decomposition", + "decision": "remove", + "status": "applied", + "evidence": [ + "The test replaced GLM projections and rotary execution, then asserted the exact private _project_qkv_absorb and _project_absorbed_value einsum decomposition.", + "The retained full-model GLM sparse-versus-dense policy executes both generic absorbed projections numerically, while the exact-kv_b owner separately proves factor-only branch routing and non-materialization." + ] + }, + { + "id": "TA-1631", + "scope": "standalone NF4 codec and bandwidth matrices", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report repeated flat and GKN shapes, dtypes, constants, three group sizes, and large random tensors without entering a shipped module; its bandwidth cases converted a missed performance target into a successful skip.", + "The retained QLoRA linear owner now executes flat NF4 quantization through forward, backward, storage, and reconstruction, while the real two-GPU Triton expert transaction reconstructs each production-shaped GKN base before two optimizer steps." + ] + }, + { + "id": "TA-1632", + "scope": "ToTensorCollator pipeline-shape report", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The report covers the four inputs admitted by the production collator: flat samples, an already-batched dictionary, packed nested samples, and empty input.", + "The larger dataloader owner starts from tensors and therefore does not prove list and NumPy conversion, field-specific integer dtypes, tensor identity, or structural preservation." + ] + }, + { + "id": "TA-1633", + "scope": "QARL Triton expert integration report", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The CPU fake-quant report independently owns rounding and STE math, but cannot prove that QARLMoEExperts exposes quantized Parameters to the production grouped-GEMM implementation.", + "The retained GPU transaction executes the real Triton expert kernel, distinguishes quantized and disabled paths numerically, and proves finite gradients reach both expert projections." + ] + }, + { + "id": "TA-1634", + "scope": "local model config and registry reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The Qwen3.5 and Kimi reports load real local config.json and tokenizer artifacts through the auto-loader rather than merely asserting registry dictionary entries.", + "Tiny model forward/backward reports construct configs directly, so they do not cover wrapper conversion, derived hybrid layer types, official auxiliary-field spelling, or local tokenizer dispatch." + ] + }, + { + "id": "TA-1635", + "scope": "small server API, scheduler, and rank-zero protocol reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The reports cross real boundaries: API response and metric mapping, typed wire serialization, rank-zero ready and acknowledgement handling, and FIFO pending, running, terminal, and bounded-history transitions.", + "Their doubles isolate sockets and worker processes but do not replace the serialization, payload conversion, queue mutation, or failure behavior being asserted." + ] + }, + { + "id": "TA-1636", + "scope": "server security, checkpoint-path, and weight-sync routing reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The security and checkpoint owners enforce traversal, symlink, SSRF and DNS-pinning, tenant isolation, reserved-checkpoint, and destructive-operation boundaries that ordinary success lifecycles cannot replace.", + "The mock-heavy endpoint report owns two-phase receiver fencing, flattened byte layout, chunking, cache metadata, and health fallback rather than merely checking that HTTP helpers were called." + ] + }, + { + "id": "TA-1637", + "scope": "remaining legacy-labeled server compatibility behavior", + "decision": "keep", + "status": "accepted", + "evidence": [ + "DRGRPO still consumes both old_logprobs and rollout logprobs at the live loss boundary, and the public API still maps Tinker session payloads into current model identifiers and optimizer controls.", + "Checkpoint and optimizer compatibility cases reject unsafe pickle state or resolve supported on-disk and URI forms; none exists solely to satisfy a test-only branch." + ] + }, + { + "id": "TA-1638", + "scope": "fake Gram Newton-Schulz CUDA dtype report", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The standalone report replaced every backend operation with logging functions and therefore proved only that its fakes received FP32 tensors, not that the shipped CUDA program preserved FP32 arithmetic.", + "The retained two-GPU full-gradient Muon transaction now runs a real CUDA Gram Newton-Schulz tree and matches an independent FP32 program exactly in addition to its single-rank optimizer oracle." + ] + }, + { + "id": "TA-1639", + "scope": "disabled timer and manually injected unrecorded-event report", + "decision": "remove", + "status": "applied", + "evidence": [ + "The report asserted disabled no-ops and wrote fake objects directly into private event-pair dictionaries; last_skipped_event_pair_count existed only to expose that synthetic path to the test.", + "The test-only counter is removed, invalid event pairs remain safely ignored, and the retained real CUDA lifecycle attaches hooks to GLM-style and Qwen-style decoder layers and records forward and backward phase timings." + ] + }, + { + "id": "TA-1640", + "scope": "remaining optimizer and training utility reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Learning-rate reports own exact warmup, decay, floor, and multi-optimizer behavior, while DistSignSGD owns sign-after-SP-sum ordering, forced SUM reduction, local versus FSDP-managed hooks, and unsupported-topology admission.", + "Cautious decay, standard batched Newton-Schulz, token counting, chunked CE, and explicit gradient synchronization retain independent numerical or collective oracles that a successful training smoke cannot replace." + ] + }, + { + "id": "TA-1641", + "scope": "scheduler-only stepping in learning-rate traces", + "decision": "rewrite", + "status": "applied", + "evidence": [ + "The schedule report advanced LambdaLR repeatedly without optimizer steps, unlike Trainer._clip_and_step, and emitted PyTorch warnings about skipping the first learning-rate value.", + "Single and multi-optimizer traces now use the production optimizer-step then scheduler-step order while preserving the exact warmup, decay, and floor assertions without warnings." + ] + }, + { + "id": "TA-1642", + "scope": "quantized-export size parsing and direct sharding helpers", + "decision": "consolidate", + "status": "applied", + "evidence": [ + "The standalone parser examples and direct-function sharding setup exercised the same implementation path separately without proving that a configured command-line export honored its string-valued shard size.", + "The retained subprocess CLI transaction now consumes a 24B YAML value, emits a multi-shard index, reconciles its total size and shard count, reloads every artifact, and checks the converted tensor dtypes and layouts." + ] + }, + { + "id": "TA-1643", + "scope": "model-support default and private-dispatch assertions", + "decision": "remove", + "status": "applied", + "evidence": [ + "The GLM bare-config field report repeated values already exercised through an official-shaped local config artifact, while its monkeypatched exact-indexer dispatch probe is a subset of the canonical GLM-5.2 contract report's router and indexer selection transaction.", + "MiniMax's three private expert-key aliases are a strict subset of the centralized checkpoint expert-key classifier matrix; its real forward, backward, checkpoint packing, and EP ownership checks remain." + ] + }, + { + "id": "TA-1644", + "scope": "remaining command-line and model-support transactions", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The NVFP4 and FP8 export owners write and reload real safetensors/config artifacts, enforce conversion and fail-closed admission, and preserve trained QARL logprobs after folding rather than merely asserting helper mappings.", + "Qwen2 and OLMo2 load Hugging Face checkpoints into fused models and compare hidden states and logits; OPD exercises causal payload alignment and Mooncake metadata transport; private FSDP policy reports uniquely own CP-folding admission and prefetch direction." + ] + }, + { + "id": "TA-1645", + "scope": "undocumented FSDP prefetch boolean coercion matrix", + "decision": "remove", + "status": "applied", + "evidence": [ + "The authoritative training and model-builder fields are bool or Optional[bool], and every production caller supplies those typed values; integer and yes/no/on/off spellings were accepted only by a direct private-helper test.", + "The compatibility parser and its ten-case matrix are removed. FSDP2 now fails fast on non-boolean values while preserving forward defaults and backward inheritance." + ] + }, + { + "id": "TA-1646", + "scope": "context-parallel folding and manual FSDP prefetch reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The CP folding truth table decides whether expert parameters need a separate FSDP root across Ulysses, Ring Attention, and mixed layouts.", + "The prefetch transaction uniquely checks forward and backward neighbor direction plus the no-op topology; these are live topology and performance policies rather than configuration aliases." + ] + }, + { + "id": "TA-1647", + "scope": "runner session, optimizer, and P2P async reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "Optimizer hyperparameters cross ServerArguments, wire payloads, dispatcher forwarding, dense optimizers, and adapter optimizers; the session registry crosses update, checkpoint load, eviction promotion, and kill.", + "The P2P async report owns size-based transfer selection and timeout failure at the real transport entry point. Its doubles isolate HTTP and Mooncake engines without replacing the production decision." + ] + }, + { + "id": "TA-1648", + "scope": "permissive optimizer cache-policy truthy adapter", + "decision": "remove", + "status": "applied", + "evidence": [ + "The skip_empty_cache_after_optim_step field has one boolean runtime owner; accepting arbitrary truthy objects and undocumented yes/on string spellings widened the contract without a configuration source.", + "The optimizer lifecycle still executes both cache policies using actual booleans and reports whether cache release was skipped." + ] + }, + { + "id": "TA-1649", + "scope": "deterministic MoE scatter rollback aliases and relaxed-atomic kernel", + "decision": "remove", + "status": "applied", + "evidence": [ + "XORL_MOE_DETERMINISTIC_SCATTER was absent from shipped arguments, examples, and documentation; it existed only to restore the older run-variant relaxed-atomic slot assignment.", + "The environment parser, four false-value aliases, route mocks, and obsolete Triton kernel are removed. The retained CUDA transaction proves stable order, full slot coverage, per-expert cumsum regions, and int32/int64 routing inputs." + ] + }, + { + "id": "TA-1650", + "scope": "mock-heavy checkpoint broadcast and adapter optimizer reports", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The broadcast owner includes real four-process Gloo DTensor materialization and checks grouped dense/expert transfer, strict coverage, checkpoint-handler filtering, and fallback admission.", + "Adapter optimizer resume retains bitwise uninterrupted comparison, artifact identity, topology-changing reshard, corruption rejection, atomic nonmutation, and a separate real two-rank rollback transaction; private-helper calls support those lifecycle guarantees rather than form isolated examples." + ] + }, + { + "id": "TA-1651", + "scope": "remaining conditional-runtime backend gates", + "decision": "keep", + "status": "accepted", + "evidence": [ + "The SGLang smoke executes a compiled operation under the pinned ABI, FlashQLA compares a real two-rank context-parallel program with its local reference, and the FSDP/GLM reports compose exact kernels with real distributed ownership.", + "The optional DeepGEMM and real SGLang MoE reports exercise documented production dispatches with independent numerical or gradient oracles; their conditional admission reflects hardware and dependency availability rather than an absent assertion." + ] + }, + { + "id": "TA-1652", + "scope": "numeric-string routing-weight position aliases", + "decision": "remove", + "status": "applied", + "evidence": [ + "The public field is a bool or one of auto, true, and false; configuration loaders can therefore supply actual booleans or those declared strings, but no shipped caller produces the undocumented strings 1 and 0.", + "The resolver and retained policy test preserve typed booleans, declared string spellings, automatic regime selection, explicit overrides, and invalid-value rejection without expanding the interface for a test-only matrix." + ] + }, + { + "id": "TA-1653", + "scope": "deprecated AdapterState.lora_params compatibility property", + "decision": "remove", + "status": "applied", + "evidence": [ + "Production, documentation, and examples use AdapterState.local_params; only tests accessed the deprecated lora_params alias or reproduced it on fake state objects.", + "The alias and fake properties are removed, while optimizer resume, coordination, checkpoint round-trip, gradient ownership, and rollback tests continue through the authoritative local parameter storage." + ] + }, + { + "id": "TA-1654", + "scope": "legacy EP replicated-gradient group fallback", + "decision": "remove", + "status": "applied", + "evidence": [ + "The production EP classifier always emits ep_replicated_gradient_sync when synchronization is enabled; falling back to the broader ep_replicated group masked incomplete model metadata and was sustained only by a test-created group dictionary.", + "Synchronization now fails fast when the authoritative group is missing, and the retained real multi-rank transaction uses the same four-group metadata built by production while preserving coalesced, missing-gradient, nonfinite, and clipping coverage." + ] + } + ] +} diff --git a/tests/_helpers/__init__.py b/tests/_helpers/__init__.py index e69de29b..51990824 100644 --- a/tests/_helpers/__init__.py +++ b/tests/_helpers/__init__.py @@ -0,0 +1,17 @@ +def stage_and_commit_gradient_capture( + manager, + model_id: str, + *, + denominator: float, + numerator_scale: float = 1.0, + backward_completed: bool = True, +): + """Advance the same explicit two-phase capture boundary used by ModelRunner.""" + + manager.stage_gradient_numerators( + model_id, + denominator=denominator, + numerator_scale=numerator_scale, + backward_completed=backward_completed, + ) + return manager.commit_gradient_capture(model_id) diff --git a/tests/_helpers/moe.py b/tests/_helpers/moe.py new file mode 100644 index 00000000..1d43897a --- /dev/null +++ b/tests/_helpers/moe.py @@ -0,0 +1,49 @@ +"""CPU kernel doubles shared by MoE contract tests.""" + +import torch + + +def counts_from_cumsum(cumsum: torch.Tensor) -> list[int]: + counts = torch.empty_like(cumsum) + counts[0] = cumsum[0] + counts[1:] = cumsum[1:] - cumsum[:-1] + return counts.tolist() + + +def _naive_group_gemm_same_nk(a, b, cumsum_M, max_M, transpose_a=False, transpose_b=False, **kwargs): + del max_M, kwargs + assert not transpose_a + + outputs = [] + start = 0 + for expert_idx, count in enumerate(counts_from_cumsum(cumsum_M)): + end = start + count + weight = b[expert_idx] + outputs.append(a[start:end] @ (weight.transpose(0, 1) if transpose_b else weight)) + start = end + return torch.cat(outputs, dim=0) + + +def _naive_group_gemm_same_mn(a, b, c, cumsum_K, max_K, transpose_a=False, transpose_b=False, **kwargs): + del max_K, kwargs + + start = 0 + for expert_idx, count in enumerate(counts_from_cumsum(cumsum_K)): + end = start + count + lhs = a[start:end].transpose(0, 1) if transpose_a else a[start:end] + rhs = b[start:end].transpose(0, 1) if transpose_b else b[start:end] + c[expert_idx].copy_(lhs @ rhs) + start = end + return c + + +def patch_ep_kernels(monkeypatch, module) -> None: + """Replace grouped kernels while preserving the canonical runtime module.""" + if module.__name__.endswith(".triton"): + monkeypatch.setattr(module, "group_gemm_same_nk", _naive_group_gemm_same_nk) + monkeypatch.setattr(module, "group_gemm_same_mn", _naive_group_gemm_same_mn) + else: + monkeypatch.setattr(module, "_group_gemm_same_nk", _naive_group_gemm_same_nk) + monkeypatch.setattr(module, "_group_gemm_same_mn", _naive_group_gemm_same_mn) + monkeypatch.setattr(module, "quack_group_gemm_same_nk", _naive_group_gemm_same_nk) + monkeypatch.setattr(module, "quack_group_gemm_same_mn", _naive_group_gemm_same_mn) diff --git a/tests/_helpers/opd.py b/tests/_helpers/opd.py index b72e5a6e..2366d779 100644 --- a/tests/_helpers/opd.py +++ b/tests/_helpers/opd.py @@ -23,12 +23,17 @@ class FakeMooncakeClient: def __init__(self) -> None: self.objects: dict[str, bytes] = {} + self.put_calls: list[str] = [] + self.get_calls: list[str] = [] + self.removed: list[str] = [] def put(self, key: str, value: bytes) -> int: self.objects[key] = bytes(value) + self.put_calls.append(key) return 0 def get(self, key: str) -> bytes: + self.get_calls.append(key) return self.objects.get(key, b"") def is_exist(self, key: str) -> int: @@ -36,6 +41,7 @@ def is_exist(self, key: str) -> int: def remove(self, key: str) -> int: self.objects.pop(key, None) + self.removed.append(key) return 0 diff --git a/tests/checkpoint/test_checkpointer_process_group.py b/tests/checkpoint/test_checkpointer_process_group.py deleted file mode 100644 index 6cfd5ebb..00000000 --- a/tests/checkpoint/test_checkpointer_process_group.py +++ /dev/null @@ -1,81 +0,0 @@ -from types import SimpleNamespace - -import pytest - -from xorl.checkpoint import checkpointer -from xorl.checkpoint.checkpointer import DistributedCheckpointer - - -pytestmark = pytest.mark.cpu - - -def test_sync_dcp_process_group_uses_gloo_for_nccl_default(monkeypatch): - created = [] - fake_group = object() - - monkeypatch.setattr( - checkpointer, - "dist", - SimpleNamespace( - is_available=lambda: True, - is_initialized=lambda: True, - get_backend=lambda: "nccl", - new_group=lambda backend: created.append(backend) or fake_group, - ), - ) - DistributedCheckpointer._sync_process_group = None - - assert DistributedCheckpointer._get_sync_process_group() is fake_group - assert DistributedCheckpointer._get_sync_process_group() is fake_group - assert created == ["gloo"] - - -def test_sync_dcp_process_group_uses_default_for_gloo(monkeypatch): - monkeypatch.setattr( - checkpointer, - "dist", - SimpleNamespace( - is_available=lambda: True, - is_initialized=lambda: True, - get_backend=lambda: "gloo", - ), - ) - DistributedCheckpointer._sync_process_group = None - - assert DistributedCheckpointer._get_sync_process_group() is None - - -def test_metadata_process_group_is_disabled_without_pipeline_parallelism(monkeypatch): - monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(pp_enabled=False)) - monkeypatch.setattr( - DistributedCheckpointer, - "_get_sync_process_group", - classmethod(lambda cls: (_ for _ in ()).throw(AssertionError("unexpected process group"))), - ) - - assert DistributedCheckpointer._get_metadata_process_group() is None - assert DistributedCheckpointer._get_metadata_process_group(object()) is None - - -def test_metadata_process_group_prefers_pipeline_caller_group(monkeypatch): - custom_group = object() - monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(pp_enabled=True)) - monkeypatch.setattr( - DistributedCheckpointer, - "_get_sync_process_group", - classmethod(lambda cls: (_ for _ in ()).throw(AssertionError("unexpected global process group"))), - ) - - assert DistributedCheckpointer._get_metadata_process_group(custom_group) is custom_group - - -def test_metadata_process_group_uses_global_gloo_for_pipeline_default(monkeypatch): - metadata_group = object() - monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(pp_enabled=True)) - monkeypatch.setattr( - DistributedCheckpointer, - "_get_sync_process_group", - classmethod(lambda cls: metadata_group), - ) - - assert DistributedCheckpointer._get_metadata_process_group() is metadata_group diff --git a/tests/checkpoint/test_ep_checkpoint_mesh.py b/tests/checkpoint/test_ep_checkpoint_mesh.py index 435f1187..e3611845 100644 --- a/tests/checkpoint/test_ep_checkpoint_mesh.py +++ b/tests/checkpoint/test_ep_checkpoint_mesh.py @@ -27,91 +27,70 @@ def __getitem__(self, mesh_dim_names): return _FakeDeviceMesh(shape, names, self.selections) -def test_legacy_2d_ep_checkpoint_mesh_is_unchanged(): - mesh = _FakeDeviceMesh((8, 2), ("ep", "ep_fsdp")) - - selected = _get_ep_checkpoint_mesh(mesh) - - assert selected is mesh - assert selected.shape == (8, 2) - assert mesh.selections == [("ep", "ep_fsdp")] - - -def test_pp2_ep8_cp8_parent_selects_stage_local_checkpoint_mesh(): - # CP8 lives in the primary training mesh. The separate EP mesh represents - # those same eight ranks per PP stage as EP8 x expert-FSDP1. - mesh = _FakeDeviceMesh((2, 8, 1), ("_pp_ep", "ep", "ep_fsdp")) - - selected = _get_ep_checkpoint_mesh(mesh) - - assert selected.mesh_dim_names == ("ep", "ep_fsdp") - assert selected.shape == (8, 1) - assert mesh.selections == [("ep", "ep_fsdp")] - - -@pytest.mark.parametrize( - "mesh_dim_names", - [ +def test_ep_checkpoint_mesh_restore_policy(monkeypatch): + for mesh_dim_names in ( ("_pp_ep", "ep"), ("_pp_ep", "ep_fsdp"), ("ep", "ep", "ep_fsdp"), - ], -) -def test_ep_checkpoint_mesh_rejects_missing_or_ambiguous_dimensions(mesh_dim_names): - mesh = _FakeDeviceMesh((1,) * len(mesh_dim_names), mesh_dim_names) - - with pytest.raises(RuntimeError, match="exactly one 'ep' and one 'ep_fsdp'"): - _get_ep_checkpoint_mesh(mesh) + ): + mesh = _FakeDeviceMesh((1,) * len(mesh_dim_names), mesh_dim_names) + with pytest.raises(RuntimeError, match="exactly one 'ep' and one 'ep_fsdp'"): + _get_ep_checkpoint_mesh(mesh) -@pytest.mark.parametrize( - ("parent_shape", "parent_names"), - [ + # ModelState must route both legacy and PP-parent meshes through the named + # checkpoint dimensions before restoring expert state. + for parent_shape, parent_names in ( ((8, 2), ("ep", "ep_fsdp")), ((2, 8, 1), ("_pp_ep", "ep", "ep_fsdp")), - ], -) -def test_model_state_restores_ep_dim_from_named_legacy_or_pp_mesh(monkeypatch, parent_shape, parent_names): - parent_mesh = _FakeDeviceMesh(parent_shape, parent_names) - state = ModelState.__new__(ModelState) - state.ep_fqn2spec_info = {"experts.weight": SimpleNamespace(placement=Shard(0), ep_fsdp_mesh=parent_mesh)} - restored = object() - - def fake_restore(tensor, mesh): - assert tensor is original - assert _get_ep_checkpoint_mesh(mesh).mesh_dim_names == ("ep", "ep_fsdp") - return restored - - monkeypatch.setattr(checkpointer, "_restore_ep_dim", fake_restore) - original = torch.ones(1) - - result = state.get_state_dict_with_ep_dim({"experts.weight": original}) - - assert result["experts.weight"] is restored - - -def test_restore_ep_dim_uses_only_named_checkpoint_dimensions(monkeypatch): - parent_mesh = _FakeDeviceMesh((2, 8, 1), ("_pp_ep", "ep", "ep_fsdp")) - local_tensor = torch.ones(1) - origin = type("DTensor", (), {"_local_tensor": local_tensor})() - restored = object() - call = {} - - def fake_from_local(tensor, *, device_mesh, placements): - call.update(tensor=tensor, device_mesh=device_mesh, placements=placements) - return restored - - monkeypatch.setattr(checkpointer, "DTensor", SimpleNamespace(from_local=fake_from_local)) - - result = checkpointer._restore_ep_dim(origin, parent_mesh) - - assert result is restored - assert call["tensor"] is local_tensor - assert call["device_mesh"].mesh_dim_names == ("ep", "ep_fsdp") - assert [placement.dim for placement in call["placements"]] == [0, 1] - - -def test_drop_ep_dim_uses_stage_local_expert_fsdp_dimension(monkeypatch): + ): + parent_mesh = _FakeDeviceMesh(parent_shape, parent_names) + state = ModelState.__new__(ModelState) + state.ep_fqn2spec_info = {"experts.weight": SimpleNamespace(placement=Shard(0), ep_fsdp_mesh=parent_mesh)} + restored = object() + original = torch.ones(1) + + def fake_restore(tensor, mesh): + assert tensor is original + assert _get_ep_checkpoint_mesh(mesh).mesh_dim_names == ("ep", "ep_fsdp") + return restored + + with monkeypatch.context() as scoped_patch: + scoped_patch.setattr(checkpointer, "_restore_ep_dim", fake_restore) + result = state.get_state_dict_with_ep_dim({"experts.weight": original}) + + assert result["experts.weight"] is restored + + # The production restore then constructs the expected two-dimensional + # DTensor for each supported parent layout. + for parent_shape, parent_names, checkpoint_shape in ( + ((8, 2), ("ep", "ep_fsdp"), (8, 2)), + ((2, 8, 1), ("_pp_ep", "ep", "ep_fsdp"), (8, 1)), + ): + parent_mesh = _FakeDeviceMesh(parent_shape, parent_names) + local_tensor = torch.ones(1) + origin = type("DTensor", (), {"_local_tensor": local_tensor})() + restored = object() + call = {} + + def fake_from_local(tensor, *, device_mesh, placements): + call.update(tensor=tensor, device_mesh=device_mesh, placements=placements) + return restored + + monkeypatch.setattr(checkpointer, "DTensor", SimpleNamespace(from_local=fake_from_local)) + + result = checkpointer._restore_ep_dim(origin, parent_mesh) + + assert result is restored + assert call["tensor"] is local_tensor + assert call["device_mesh"].mesh_dim_names == ("ep", "ep_fsdp") + assert call["device_mesh"].shape == checkpoint_shape + assert [placement.dim for placement in call["placements"]] == [0, 1] + + _assert_drop_ep_dim_uses_stage_local_expert_fsdp_dimension(monkeypatch) + + +def _assert_drop_ep_dim_uses_stage_local_expert_fsdp_dimension(monkeypatch): parent_mesh = _FakeDeviceMesh((2, 8, 1), ("_pp_ep", "ep", "ep_fsdp")) local_tensor = torch.ones(1) loaded = SimpleNamespace(_local_tensor=local_tensor, placements=(Shard(0), Shard(1))) diff --git a/tests/conftest.py b/tests/conftest.py index d2e2e6bc..684af694 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,20 +1,8 @@ -from typing import Any, Dict, Sequence - import pytest import torch from torch.utils.data import Dataset -class SimpleCollator: - """Simple collator for testing that stacks tensors.""" - - def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, torch.Tensor]: - result = {} - for key in features[0].keys(): - result[key] = torch.stack([f[key] for f in features]) - return result - - class FakeTextDataset(Dataset): """ A fake text dataset for testing purposes. @@ -30,7 +18,7 @@ def __init__(self, num_samples: int = 100, seq_len: int = 128, vocab_size: int = def __len__(self): return self.num_samples - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: # Create deterministic but varied data based on index torch.manual_seed(idx) @@ -45,75 +33,12 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: } -class FakePackedDataset(Dataset): - """ - A fake dataset that returns packed sequences with position_ids. - - Simulates packing multiple sequences together with position IDs. - """ - - def __init__( - self, - num_samples: int = 100, - min_seq_len: int = 64, - max_seq_len: int = 256, - vocab_size: int = 1000, - num_sequences_per_sample: int = 3, - ): - self.num_samples = num_samples - self.min_seq_len = min_seq_len - self.max_seq_len = max_seq_len - self.vocab_size = vocab_size - self.num_sequences_per_sample = num_sequences_per_sample - - def __len__(self): - return self.num_samples - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - torch.manual_seed(idx) - - # Generate multiple sequences and pack them - sequences = [] - position_ids = [] - current_pos = 0 - - for _ in range(self.num_sequences_per_sample): - seq_len = torch.randint(self.min_seq_len, self.max_seq_len, (1,)).item() - seq = torch.randint(1, self.vocab_size, (seq_len,), dtype=torch.long) - sequences.append(seq) - - # Position IDs start from 0 for each sequence - pos_ids = torch.arange(seq_len, dtype=torch.long) - position_ids.append(pos_ids) - - # Concatenate all sequences - input_ids = torch.cat(sequences) - position_ids = torch.cat(position_ids) - attention_mask = torch.ones_like(input_ids) - labels = input_ids.clone() - - return { - "input_ids": input_ids, - "attention_mask": attention_mask, - "labels": labels, - "position_ids": position_ids, - } - - @pytest.fixture def fake_text_dataset(): """Provides a fake text dataset.""" return FakeTextDataset(num_samples=100, seq_len=128, vocab_size=1000) -@pytest.fixture -def fake_packed_dataset(): - """Provides a fake packed dataset with position IDs.""" - return FakePackedDataset( - num_samples=100, min_seq_len=64, max_seq_len=256, vocab_size=1000, num_sequences_per_sample=3 - ) - - @pytest.fixture def sample_features(): """Provides sample features for testing collators.""" diff --git a/tests/data/collators/test_collate_pipeline.py b/tests/data/collators/test_collate_pipeline.py deleted file mode 100644 index 919b1fd1..00000000 --- a/tests/data/collators/test_collate_pipeline.py +++ /dev/null @@ -1,100 +0,0 @@ -from typing import Any, Dict, Sequence - -import pytest -import torch - -from xorl.data.collators import CollatePipeline, DataCollator - - -pytestmark = [pytest.mark.cpu, pytest.mark.collator] - - -class MockCollator1(DataCollator): - """Mock collator that adds a constant to input_ids.""" - - def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, torch.Tensor]: - result = {} - for key in features[0].keys(): - if key == "input_ids": - result[key] = torch.cat([f[key] for f in features]) + 1 - else: - result[key] = torch.cat([f[key] for f in features]) - return result - - -class MockCollator2(DataCollator): - """Mock collator that multiplies input_ids by 2.""" - - def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, torch.Tensor]: - if isinstance(features, dict): - features["input_ids"] = features["input_ids"] * 2 - return features - else: - result = {} - for key in features[0].keys(): - if key == "input_ids": - result[key] = torch.cat([f[key] for f in features]) * 2 - else: - result[key] = torch.cat([f[key] for f in features]) - return result - - -class MockCollator3(DataCollator): - """Mock collator that adds a new field.""" - - def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, torch.Tensor]: - if isinstance(features, dict): - features["new_field"] = torch.tensor([100]) - return features - else: - result = {} - for key in features[0].keys(): - result[key] = torch.cat([f[key] for f in features]) - result["new_field"] = torch.tensor([100]) - return result - - -class TestCollatePipeline: - """Test suite for CollatePipeline.""" - - def test_single_and_sequential_collators(self, sample_features): - """Covers single collator, multiple collators in sequence, three-collator chain, - and adding new fields through pipeline.""" - base_ids = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) - - # Single collator - pipeline1 = CollatePipeline(MockCollator1()) - r1 = pipeline1(sample_features) - assert torch.equal(r1["input_ids"], base_ids + 1) - - # Two collators: add 1 then multiply by 2 - pipeline2 = CollatePipeline([MockCollator1(), MockCollator2()]) - r2 = pipeline2(sample_features) - assert torch.equal(r2["input_ids"], (base_ids + 1) * 2) - - # Three collators: add 1, multiply by 2, add new_field - pipeline3 = CollatePipeline([MockCollator1(), MockCollator2(), MockCollator3()]) - r3 = pipeline3(sample_features) - assert torch.equal(r3["input_ids"], (base_ids + 1) * 2) - assert "new_field" in r3 and torch.equal(r3["new_field"], torch.tensor([100])) - - def test_empty_list_tuple_and_key_preservation(self, sample_features): - """Covers empty collator list, single collator as list vs direct, tuple of collators, - and key preservation through pipeline.""" - # Empty list returns original features - assert CollatePipeline([])(sample_features) == sample_features - - # Single collator direct vs wrapped in list produces same result - c = MockCollator1() - r_direct = CollatePipeline(c)(sample_features) - r_list = CollatePipeline([c])(sample_features) - assert torch.equal(r_direct["input_ids"], r_list["input_ids"]) - - # Tuple of collators works - r_tuple = CollatePipeline((MockCollator1(), MockCollator2()))(sample_features) - expected = (torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) + 1) * 2 - assert torch.equal(r_tuple["input_ids"], expected) - - # All keys preserved - r_keys = CollatePipeline([MockCollator1()])(sample_features) - assert all(k in r_keys for k in ["input_ids", "attention_mask", "labels"]) diff --git a/tests/data/collators/test_packing_concat_collator.py b/tests/data/collators/test_packing_concat_collator.py index 61f1d283..cce53a57 100644 --- a/tests/data/collators/test_packing_concat_collator.py +++ b/tests/data/collators/test_packing_concat_collator.py @@ -3,35 +3,12 @@ import pytest import torch -from xorl.data.collators import PackingConcatCollator, add_flash_attention_kwargs_from_position_ids +from xorl.data.collators import PackingConcatCollator pytestmark = [pytest.mark.cpu, pytest.mark.collator] -class TestAddFlashAttentionKwargs: - """Tests for add_flash_attention_kwargs_from_position_ids function.""" - - def test_kwargs_and_correctness(self): - """Covers all required kwargs added, cu_seq_lens correctness, max_length correctness, and single sequence.""" - # Multiple sequences: [0,1,2] and [0,1,2,3] - batch = { - "input_ids": torch.tensor([[1, 2, 3, 4, 5, 6, 7]]), - "position_ids": torch.tensor([[0, 1, 2, 0, 1, 2, 3]]), - } - cu_q, cu_k, max_q, max_k = add_flash_attention_kwargs_from_position_ids(batch) - - assert all(k in batch for k in ["cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k"]) - assert torch.equal(cu_q, torch.tensor([0, 3, 7], dtype=torch.int32)) - assert max_q == 4 and max_k == 4 - - # Single sequence - batch2 = {"position_ids": torch.tensor([[0, 1, 2, 3, 4]])} - cu_q2, _, max_q2, _ = add_flash_attention_kwargs_from_position_ids(batch2) - assert torch.equal(cu_q2, torch.tensor([0, 5], dtype=torch.int32)) - assert max_q2 == 5 - - class TestPackingConcatCollator: """Tests for PackingConcatCollator.""" @@ -56,6 +33,11 @@ def test_concatenation_and_flash_attn(self, mock_parallel_state, sample_packed_f # Flash attn kwargs present when SP disabled assert all(k in batch for k in ["cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k"]) + expected_cu_seq_lens = torch.tensor([0, 3, 5, 10], dtype=torch.int32) + assert torch.equal(batch["cu_seq_lens_q"], expected_cu_seq_lens) + assert torch.equal(batch["cu_seq_lens_k"], expected_cu_seq_lens) + assert batch["max_length_q"] == 5 + assert batch["max_length_k"] == 5 # Position IDs generated if missing batch_gen = collator(sample_features) @@ -79,9 +61,14 @@ def test_concatenation_and_flash_attn(self, mock_parallel_state, sample_packed_f ] batch_single = collator(single) assert batch_single["input_ids"].shape == (1, 3) + assert torch.equal(batch_single["cu_seq_lens_q"], torch.tensor([0, 3], dtype=torch.int32)) + assert batch_single["max_length_q"] == 3 + + self._assert_extra_fields_and_multiple_seqs() + self._assert_sequence_side_fields_concatenate_and_pad() @patch("xorl.data.collators.packing_concat_collator.get_parallel_state") - def test_extra_fields_and_multiple_seqs(self, mock_parallel_state): + def _assert_extra_fields_and_multiple_seqs(self, mock_parallel_state): """Covers extra field handling and multiple packed sequences per sample.""" mock_ps = Mock() mock_ps.cp_enabled = False @@ -129,7 +116,7 @@ def test_extra_fields_and_multiple_seqs(self, mock_parallel_state): assert torch.equal(batch2["position_ids"], torch.tensor([[0, 1, 2, 0, 1, 0, 1, 2, 3]])) @patch("xorl.data.collators.packing_concat_collator.get_parallel_state") - def test_teacher_hidden_states_concatenate_and_pad_as_sequence_field(self, mock_parallel_state): + def _assert_sequence_side_fields_concatenate_and_pad(self, mock_parallel_state): mock_ps = Mock() mock_ps.cp_enabled = False mock_parallel_state.return_value = mock_ps @@ -159,9 +146,9 @@ def test_teacher_hidden_states_concatenate_and_pad_as_sequence_field(self, mock_ torch.tensor([[[0.25, 0.5], [1.25, 1.5], [2.25, 2.5], [0.0, 0.0]]]), ) - @patch("xorl.data.collators.packing_concat_collator.get_parallel_state") - def test_hidden_match_weights_are_padded_as_sequence_field(self, mock_parallel_state): - mock_parallel_state.return_value = Mock(cp_enabled=False) + self._assert_hidden_match_weights_are_padded_as_sequence_field() + + def _assert_hidden_match_weights_are_padded_as_sequence_field(self): collator = PackingConcatCollator(pad_to_multiple_of=4) features = [ { diff --git a/tests/data/collators/test_sequence_shard_collator.py b/tests/data/collators/test_sequence_shard_collator.py index 6f5a740a..8edddc32 100644 --- a/tests/data/collators/test_sequence_shard_collator.py +++ b/tests/data/collators/test_sequence_shard_collator.py @@ -4,6 +4,7 @@ import torch from xorl.data.collators import TextSequenceShardCollator +from xorl.data.collators.sequence_shard_collator import zigzag_reorder_packed_sequence from xorl.data.constants import IGNORE_INDEX @@ -18,51 +19,11 @@ def _make_mock_ps(cp_size=2, cp_rank=0, ringattn_size=1): return mock_ps -class TestSPSliceAndPadding: - """Tests for sp_slice and sp_padding utility methods.""" - - @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") - def test_sp_slice_across_ranks_and_uneven(self, mock_parallel_state): - """Covers initialization, basic slicing rank 0/1, and uneven split.""" - mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=0) - collator = TextSequenceShardCollator() - assert collator.cp_size == 2 and collator.cp_rank == 0 - - tensor = torch.tensor([[1, 2, 3, 4, 5, 6]]) - assert torch.equal(collator.sp_slice(tensor, dim=-1), torch.tensor([[1, 2, 3]])) - - # Rank 1 - mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=1) - collator1 = TextSequenceShardCollator() - assert torch.equal(collator1.sp_slice(tensor, dim=-1), torch.tensor([[4, 5, 6]])) - - # Uneven split (length 5, cp_size=2 -> chunk_size=3 for rank 0) - mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=0) - collator0 = TextSequenceShardCollator() - assert torch.equal(collator0.sp_slice(torch.tensor([[1, 2, 3, 4, 5]]), dim=-1), torch.tensor([[1, 2, 3]])) - - @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") - def test_sp_padding_basic_sequential_zero(self, mock_parallel_state): - """Covers basic padding, sequential padding, and zero-length padding.""" - mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=0) - collator = TextSequenceShardCollator() - tensor = torch.tensor([[1, 2, 3]]) - - assert torch.equal( - collator.sp_padding(tensor, dim=-1, pad_value=0, pad_length=2), torch.tensor([[1, 2, 3, 0, 0]]) - ) - assert torch.equal( - collator.sp_padding(tensor, dim=-1, pad_value=0, pad_length=2, sequential=True), - torch.tensor([[1, 2, 3, 0, 1]]), - ) - assert torch.equal(collator.sp_padding(tensor, dim=-1, pad_value=0, pad_length=0), tensor) - - class TestCollatorCall: """Tests for the full collator __call__ method.""" @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") - def test_preshifted_labels_and_packed_sequences(self, mock_parallel_state): + def _assert_preshifted_labels_and_packed_sequences(self, mock_parallel_state): """Covers pre-shifted labels pass-through and packed sequence boundary masking with cp_size=1.""" mock_parallel_state.return_value = _make_mock_ps(cp_size=1, cp_rank=0) collator = TextSequenceShardCollator(pad_token_id=0) @@ -92,6 +53,8 @@ def test_preshifted_labels_and_packed_sequences(self, mock_parallel_state): def test_sp_splitting_padding_and_flash_attn_kwargs(self, mock_parallel_state): """Covers SP padding to multiple, splitting across ranks, flash attention kwargs, attention_mask/position_ids preservation, and padding values.""" + self._assert_preshifted_labels_and_packed_sequences() + # SP splitting with cp_size=2 mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=0) collator0 = TextSequenceShardCollator(pad_token_id=0) @@ -104,7 +67,8 @@ def test_sp_splitting_padding_and_flash_attn_kwargs(self, mock_parallel_state): } r0 = collator0(batch) assert r0["input_ids"].shape[-1] == 3 - assert r0["input_ids"][0, 0] == 1 + assert torch.equal(r0["input_ids"], torch.tensor([[1, 2, 3]])) + assert torch.equal(r0["labels"], torch.tensor([[2, 3, 4]])) mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=1) collator1 = TextSequenceShardCollator(pad_token_id=0) @@ -116,7 +80,8 @@ def test_sp_splitting_padding_and_flash_attn_kwargs(self, mock_parallel_state): } r1 = collator1(batch_r1) assert r1["input_ids"].shape[-1] == 3 - assert r1["input_ids"][0, 0] == 4 + assert torch.equal(r1["input_ids"], torch.tensor([[4, 5, 6]])) + assert torch.equal(r1["labels"], torch.tensor([[5, 6, IGNORE_INDEX]])) # Padding to SP multiple (length 5 -> padded to 6, then split to 3) mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=0) @@ -129,6 +94,22 @@ def test_sp_splitting_padding_and_flash_attn_kwargs(self, mock_parallel_state): } r_pad = collator_pad(batch5) assert r_pad["input_ids"].shape[-1] == 3 + assert torch.equal(r_pad["input_ids"], torch.tensor([[1, 2, 3]])) + + # The last rank observes both constant and sequential padding through + # the production collator, rather than through its private primitives. + mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=1) + collator_pad_last = TextSequenceShardCollator(pad_token_id=0) + batch5_last = { + "input_ids": torch.tensor([[1, 2, 3, 4, 5]]), + "attention_mask": torch.tensor([[1, 1, 1, 1, 1]]), + "labels": torch.tensor([[2, 3, 4, 5, IGNORE_INDEX]]), + "position_ids": torch.tensor([[0, 1, 2, 3, 4]]), + } + r_pad_last = collator_pad_last(batch5_last) + assert torch.equal(r_pad_last["input_ids"], torch.tensor([[4, 5, 0]])) + assert torch.equal(r_pad_last["labels"], torch.tensor([[5, IGNORE_INDEX, IGNORE_INDEX]])) + assert torch.equal(r_pad_last["position_ids"], torch.tensor([[0, 1, 2, 3, 4, 0]])) # Flash attention kwargs added (cp_size=1) mock_parallel_state.return_value = _make_mock_ps(cp_size=1, cp_rank=0) @@ -157,8 +138,11 @@ def test_sp_splitting_padding_and_flash_attn_kwargs(self, mock_parallel_state): assert r_single["input_ids"].shape[-1] == 5 assert r_single["labels"].shape[-1] == 5 + self._assert_token_side_channels_follow_sequence_shards() + _assert_zigzag_reorder_policy() + @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") - def test_teacher_hidden_states_are_sharded_with_token_fields(self, mock_parallel_state): + def _assert_token_side_channels_follow_sequence_shards(self, mock_parallel_state): mock_parallel_state.return_value = _make_mock_ps(cp_size=2, cp_rank=1) collator = TextSequenceShardCollator(pad_token_id=0) teacher_hidden_states = torch.tensor( @@ -189,49 +173,103 @@ def test_teacher_hidden_states_are_sharded_with_token_fields(self, mock_parallel torch.tensor([[[3.25, 3.5], [4.25, 4.5], [0.0, 0.0]]]), ) - @pytest.mark.parametrize("cp_rank", [0, 15]) + self._assert_drgrpo_side_channels_follow_cp16_target_shards() + @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") - def test_drgrpo_side_channels_follow_cp16_target_shard(self, mock_parallel_state, cp_rank): + def _assert_drgrpo_side_channels_follow_cp16_target_shards(self, mock_parallel_state): """Canonical and reference logprobs must follow the same padded CP slice as labels.""" cp_size = 16 seq_len = 4099 chunk_size = 257 - mock_parallel_state.return_value = _make_mock_ps(cp_size=cp_size, cp_rank=cp_rank) - collator = TextSequenceShardCollator(pad_token_id=0) - input_ids = torch.arange(seq_len).unsqueeze(0) target_tokens = torch.cat([torch.arange(1, seq_len), torch.tensor([IGNORE_INDEX])]).unsqueeze(0) old_logprobs = torch.arange(seq_len, dtype=torch.float32).add(0.25).unsqueeze(0) advantages = torch.arange(seq_len, dtype=torch.float32).add(0.5).unsqueeze(0) ref_logprobs = torch.arange(seq_len, dtype=torch.float32).neg().sub(0.75).unsqueeze(0) - result = collator( - { - "input_ids": input_ids, - "attention_mask": torch.ones_like(input_ids), - "labels": target_tokens.clone(), - "target_tokens": target_tokens.clone(), - "position_ids": torch.arange(seq_len).unsqueeze(0), - "old_logprobs": old_logprobs, - "advantages": advantages, - "ref_logprobs": ref_logprobs, - } - ) - start = cp_rank * chunk_size - end = min(start + chunk_size, seq_len) - valid_length = end - start - assert result["labels"].shape == (1, chunk_size) - for field in ("target_tokens", "old_logprobs", "advantages", "ref_logprobs"): - assert result[field].shape == result["labels"].shape - - torch.testing.assert_close(result["target_tokens"][0, :valid_length], target_tokens[0, start:end]) - torch.testing.assert_close(result["old_logprobs"][0, :valid_length], old_logprobs[0, start:end]) - torch.testing.assert_close(result["advantages"][0, :valid_length], advantages[0, start:end]) - torch.testing.assert_close(result["ref_logprobs"][0, :valid_length], ref_logprobs[0, start:end]) - if valid_length < chunk_size: - assert torch.equal( - result["target_tokens"][0, valid_length:], - torch.full((chunk_size - valid_length,), IGNORE_INDEX), + for cp_rank in (0, 15): + mock_parallel_state.return_value = _make_mock_ps(cp_size=cp_size, cp_rank=cp_rank) + collator = TextSequenceShardCollator(pad_token_id=0) + result = collator( + { + "input_ids": input_ids, + "attention_mask": torch.ones_like(input_ids), + "labels": target_tokens.clone(), + "target_tokens": target_tokens.clone(), + "position_ids": torch.arange(seq_len).unsqueeze(0), + "old_logprobs": old_logprobs, + "advantages": advantages, + "ref_logprobs": ref_logprobs, + } ) - for field in ("old_logprobs", "advantages", "ref_logprobs"): - assert torch.equal(result[field][0, valid_length:], torch.zeros(chunk_size - valid_length)) + + start = cp_rank * chunk_size + end = min(start + chunk_size, seq_len) + valid_length = end - start + assert result["labels"].shape == (1, chunk_size) + for field in ("target_tokens", "old_logprobs", "advantages", "ref_logprobs"): + assert result[field].shape == result["labels"].shape + + torch.testing.assert_close(result["target_tokens"][0, :valid_length], target_tokens[0, start:end]) + torch.testing.assert_close(result["old_logprobs"][0, :valid_length], old_logprobs[0, start:end]) + torch.testing.assert_close(result["advantages"][0, :valid_length], advantages[0, start:end]) + torch.testing.assert_close(result["ref_logprobs"][0, :valid_length], ref_logprobs[0, start:end]) + if valid_length < chunk_size: + assert torch.equal( + result["target_tokens"][0, valid_length:], + torch.full((chunk_size - valid_length,), IGNORE_INDEX), + ) + for field in ("old_logprobs", "advantages", "ref_logprobs"): + assert torch.equal(result[field][0, valid_length:], torch.zeros(chunk_size - valid_length)) + + +def _assert_zigzag_reorder_policy() -> None: + ringattn_size = 2 + tensor = torch.arange(40).unsqueeze(0) + position_ids = torch.arange(40).unsqueeze(0) + reordered = zigzag_reorder_packed_sequence(tensor, position_ids, ringattn_size, dim=-1) + expected = torch.cat( + ( + torch.arange(0, 10), + torch.arange(30, 40), + torch.arange(10, 20), + torch.arange(20, 30), + ) + ).unsqueeze(0) + assert torch.equal(reordered, expected) + + doc_len = 20 + position_ids = torch.cat((torch.arange(doc_len), torch.arange(doc_len))).unsqueeze(0) + reordered = zigzag_reorder_packed_sequence(torch.arange(2 * doc_len).unsqueeze(0), position_ids, ringattn_size) + expected = torch.cat( + ( + torch.arange(0, 5), + torch.arange(15, 20), + torch.arange(20, 25), + torch.arange(35, 40), + torch.arange(5, 10), + torch.arange(10, 15), + torch.arange(25, 30), + torch.arange(30, 35), + ) + ).unsqueeze(0) + assert torch.equal(reordered, expected) + reordered_position_ids = zigzag_reorder_packed_sequence(position_ids, position_ids, ringattn_size) + assert (reordered_position_ids[0, :doc_len] == 0).nonzero(as_tuple=False).view(-1).numel() == 2 + + for ringattn_size in (2, 4, 8): + sequence_length = 8 * 2 * ringattn_size + tensor = torch.arange(sequence_length).unsqueeze(0) + reordered = zigzag_reorder_packed_sequence(tensor, tensor, ringattn_size) + assert reordered.shape == tensor.shape + assert torch.equal(reordered.sort().values, tensor) + rank_width = sequence_length // ringattn_size + for rank in range(ringattn_size): + rank_slice = reordered[0, rank * rank_width : (rank + 1) * rank_width] + half = rank_width // 2 + assert rank_slice[:half].max() < rank_slice[half:].min() + + tensor = torch.arange(20).unsqueeze(0) + assert zigzag_reorder_packed_sequence(tensor, tensor, 1) is tensor + with pytest.raises(ValueError, match="not divisible"): + zigzag_reorder_packed_sequence(torch.arange(15).unsqueeze(0), torch.arange(15).unsqueeze(0), 2) diff --git a/tests/data/collators/test_tensor_collator.py b/tests/data/collators/test_tensor_collator.py index 974e6d7d..3bd39be4 100644 --- a/tests/data/collators/test_tensor_collator.py +++ b/tests/data/collators/test_tensor_collator.py @@ -10,141 +10,37 @@ pytestmark = [pytest.mark.cpu, pytest.mark.collator] -class TestToTensorCollator: - """Test suite for ToTensorCollator.""" +def test_to_tensor_collator_preserves_its_pipeline_shapes_and_types(): + collator = ToTensorCollator() + labels = torch.tensor([4, 5, 6]) - def test_type_conversion_and_passthrough(self): - """Covers list->tensor, numpy->tensor, tensor passthrough, mixed inputs, and boolean lists.""" - collator = ToTensorCollator() - - # Lists to tensors - result = collator( - [{"input_ids": [1, 2, 3], "labels": [4, 5, 6]}, {"input_ids": [7, 8, 9], "labels": [10, 11, 12]}] - ) - assert isinstance(result, list) and len(result) == 2 - assert isinstance(result[0]["input_ids"], torch.Tensor) and result[0]["input_ids"].shape == (3,) - assert torch.equal(result[0]["input_ids"], torch.tensor([1, 2, 3])) - assert torch.equal(result[1]["input_ids"], torch.tensor([7, 8, 9])) - - # Numpy arrays to tensors - result_np = collator([{"input_ids": np.array([1, 2, 3]), "labels": np.array([4, 5, 6])}]) - assert isinstance(result_np[0]["input_ids"], torch.Tensor) - - # Already tensors pass through - t = torch.tensor([1, 2, 3]) - result_t = collator([{"input_ids": t, "labels": torch.tensor([4, 5, 6])}]) - assert isinstance(result_t[0]["input_ids"], torch.Tensor) - assert torch.equal(result_t[0]["input_ids"], t) - - # Mixed list and tensor - result_mix = collator([{"input_ids": [1, 2, 3], "labels": torch.tensor([4, 5, 6])}]) - assert isinstance(result_mix[0]["input_ids"], torch.Tensor) - assert isinstance(result_mix[0]["labels"], torch.Tensor) - - # Boolean lists - result_bool = collator([{"input_ids": [1, 2, 3], "mask": [True, False, True]}]) - assert isinstance(result_bool[0]["mask"], torch.Tensor) - assert result_bool[0]["mask"].dtype == torch.bool - - def test_scalars_empty_and_dtype_inference(self): - """Covers scalar fields, empty features, dtype inference for known fields, numpy dtypes, and batch_size=1.""" - collator = ToTensorCollator() - - # Scalar fields - result = collator([{"input_ids": [1, 2, 3], "length": 3, "score": 0.95}]) - assert isinstance(result[0]["length"], torch.Tensor) and result[0]["length"].shape == () - - # Empty features - assert collator([]) == {} - - # Dtype inference - result_dtype = collator( - [ - { - "input_ids": [1, 2, 3], - "labels": [4, 5, 6], - "position_ids": [0, 1, 2], - "attention_mask": [1, 1, 1], - "other_field": [1.0, 2.0, 3.0], - } - ] - ) - assert result_dtype[0]["input_ids"].dtype == torch.long - assert result_dtype[0]["labels"].dtype == torch.long - assert result_dtype[0]["position_ids"].dtype == torch.long - assert result_dtype[0]["attention_mask"].dtype == torch.long - assert result_dtype[0]["other_field"].dtype in [torch.float32, torch.float64] - - # Numpy dtype preservation - result_np = collator( - [ - { - "input_ids": np.array([1, 2, 3], dtype=np.int32), - "embeddings": np.array([1.0, 2.0, 3.0], dtype=np.float32), - } - ] - ) - assert result_np[0]["embeddings"].dtype == torch.float32 - - # Batch size one - result_one = collator([{"input_ids": [1, 2, 3], "labels": [4, 5, 6]}]) - assert result_one[0]["input_ids"].shape == (3,) - - def test_2d_lists_and_string_handling(self): - """Covers 2D numeric lists, 2D string lists, single strings, and mixed numeric/string fields.""" - collator = ToTensorCollator() - - # 2D numeric lists - result = collator([{"position_ids": [[0, 1], [2, 3]]}]) - assert isinstance(result[0]["position_ids"], torch.Tensor) and result[0]["position_ids"].shape == (2, 2) - - # 2D string lists kept as-is - result_str2d = collator([{"text": [["hello", "world"], ["foo", "bar"]]}]) - assert isinstance(result_str2d[0]["text"], list) - - # Single string kept as-is - result_str = collator([{"input_ids": [1, 2, 3], "text": "hello world"}]) - assert isinstance(result_str[0]["text"], str) and result_str[0]["text"] == "hello world" - - # Mixed numeric and string fields - result_mixed = collator([{"input_ids": [1, 2, 3], "text": ["hello"], "labels": [4, 5, 6], "source": ["web"]}]) - assert isinstance(result_mixed[0]["input_ids"], torch.Tensor) - assert isinstance(result_mixed[0]["text"], list) - - def test_different_lengths_and_packed_sequences(self): - """Covers different length fallback and packed sequence handling.""" - collator = ToTensorCollator() - - # Different lengths fallback - result = collator( - [ - {"input_ids": [1, 2, 3], "labels": [4, 5, 6]}, - {"input_ids": [7, 8, 9, 10, 11], "labels": [12, 13, 14, 15, 16]}, - ] - ) - assert isinstance(result, list) - assert result[0]["input_ids"].shape == (3,) - assert result[1]["input_ids"].shape == (5,) - - # Packed sequences (same length, should be stacked) - packed = [ - { - "input_ids": [1, 2, 3, 101, 4, 5, 6, 102], - "labels": [-100, -100, 7, -100, -100, 8, 9, -100], - "position_ids": [0, 1, 2, 3, 0, 1, 2, 3], - "length": 8, - }, + [sample] = collator( + [ { - "input_ids": [10, 11, 12, 101, 13, 14, 15, 16], - "labels": [-100, 15, 16, -100, 17, 18, 19, 20], - "position_ids": [0, 1, 2, 0, 1, 2, 3, 4], - "length": 8, - }, + "input_ids": [1, 2, 3], + "labels": labels, + "embeddings": np.array([0.5, 1.5, 2.5], dtype=np.float32), + "text": ["source text"], + "length": 3, + } ] - result_packed = collator(packed) - assert result_packed[0]["input_ids"].shape == (8,) - assert result_packed[0]["position_ids"].shape == (8,) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) + ) + assert torch.equal(sample["input_ids"], torch.tensor([1, 2, 3], dtype=torch.long)) + assert sample["labels"] is labels + assert sample["embeddings"].dtype == torch.float32 + assert sample["text"] == ["source text"] + assert sample["length"].shape == () + + batched = collator({"input_ids": [[1, 2], [3, 4]], "attention_mask": [[1, 1], [1, 0]]}) + assert batched["input_ids"].shape == (2, 2) + assert batched["attention_mask"].dtype == torch.long + + nested = collator( + [ + [{"input_ids": [1, 2], "labels": [2, 3]}], + [{"input_ids": [4], "labels": [5]}], + ] + ) + assert nested[0][0]["input_ids"].shape == (2,) + assert nested[1][0]["labels"].dtype == torch.long + assert collator([]) == {} diff --git a/tests/data/prepare/test_file_lock_loader.py b/tests/data/prepare/test_file_lock_loader.py deleted file mode 100644 index fcd11193..00000000 --- a/tests/data/prepare/test_file_lock_loader.py +++ /dev/null @@ -1,146 +0,0 @@ -"""Tests for xorl.data.prepare.file_lock_loader module.""" - -from pathlib import Path -from unittest.mock import Mock, patch - -import pytest - -from xorl.data.prepare.file_lock_loader import ( - LOCK_FILE_NAME, - PROCESS_COUNTER_FILE_NAME, - READY_FILE_NAME, - FileLockLoader, -) - - -pytestmark = pytest.mark.cpu - - -@pytest.fixture -def temp_dataset_path(tmp_path): - """Provides a temporary directory for dataset preparation.""" - return str(tmp_path / "prepared_datasets") - - -@pytest.fixture -def mock_args(temp_dataset_path): - """Provides mock Arguments object.""" - args = Mock() - args.data.dataset_prepared_path = temp_dataset_path - return args - - -class TestFileLockLoaderInitAndLoad: - """Tests for initialization and load behavior.""" - - def test_init_load_and_counter_incrementing(self, mock_args, temp_dataset_path): - """Covers initialization paths, directory creation, first-process load, counter incrementing, - multiple loads, and subsequent process caching.""" - # Custom path initialization - loader = FileLockLoader(mock_args) - assert loader.dataset_prepared_path == temp_dataset_path - assert str(loader.lock_file_path).endswith(LOCK_FILE_NAME) - assert str(loader.ready_flag_path).endswith(READY_FILE_NAME) - assert str(loader.counter_path).endswith(PROCESS_COUNTER_FILE_NAME) - - # Default path when None - args_none = Mock() - args_none.data.dataset_prepared_path = None - loader_none = FileLockLoader(args_none) - assert loader_none.dataset_prepared_path == "last_prepared_dataset" - - # First process: creates directory, executes load_fn, creates ready flag, counter=1 - load_fn = Mock(return_value="first_process_data") - result = loader.load(load_fn) - assert result == "first_process_data" - load_fn.assert_called_once() - assert Path(temp_dataset_path).exists() - assert loader.ready_flag_path.exists() - assert loader.counter_path.read_text().strip() == "1" - - # Second load increments counter - loader2 = FileLockLoader(mock_args) - load_fn2 = Mock(return_value="new_data") - result2 = loader2.load(load_fn2) - assert result2 == "new_data" - load_fn2.assert_called_once() - assert loader2.counter_path.read_text().strip() == "2" - - -class TestFileLockLoaderCleanup: - """Tests for cleanup behavior.""" - - def test_cleanup_single_and_multiple_processes(self, mock_args, temp_dataset_path): - """Covers single-process cleanup, multi-process partial cleanup, and full cleanup.""" - loader1 = FileLockLoader(mock_args) - loader2 = FileLockLoader(mock_args) - - loader1.load(lambda: "data") - loader2.load(lambda: "data") - assert loader1.counter_path.read_text().strip() == "2" - - # First cleanup: counter=1, files still exist - loader1.cleanup() - assert loader1.counter_path.exists() - assert loader1.ready_flag_path.exists() - assert loader1.counter_path.read_text().strip() == "1" - - # Second cleanup: all files removed - loader2.cleanup() - assert not loader2.counter_path.exists() - assert not loader2.ready_flag_path.exists() - - -class TestFileLockLoaderErrorHandling: - """Tests for corrupted state, IO errors, and exception propagation.""" - - def test_corrupted_counter_io_error_and_exceptions(self, mock_args, temp_dataset_path): - """Covers corrupted counter on increment/cleanup, missing counter, IO error, load_fn exception, and concurrent access.""" - loader = FileLockLoader(mock_args) - - # Corrupted counter on increment -> reset to 1 - Path(temp_dataset_path).mkdir(parents=True, exist_ok=True) - loader.counter_path.write_text("invalid_number") - loader.load(lambda: "data") - assert loader.counter_path.read_text().strip() == "1" - - # Corrupted counter on cleanup -> force cleanup - loader.counter_path.write_text("invalid_number") - loader.cleanup() - assert not loader.counter_path.exists() - assert not loader.ready_flag_path.exists() - - # Missing counter on increment -> start at 1 - loader2 = FileLockLoader(mock_args) - Path(temp_dataset_path).mkdir(parents=True, exist_ok=True) - if loader2.counter_path.exists(): - loader2.counter_path.unlink() - loader2.load(lambda: "data") - assert loader2.counter_path.read_text().strip() == "1" - - # IO error on increment -> reset to 1 - loader3 = FileLockLoader(mock_args) - loader3.cleanup() - Path(temp_dataset_path).mkdir(parents=True, exist_ok=True) - with patch.object(Path, "read_text", side_effect=OSError("IO error")): - loader3.load(lambda: "data") - assert loader3.counter_path.read_text().strip() == "1" - - # load_fn exception propagation - loader4 = FileLockLoader(mock_args) - loader4.cleanup() - with pytest.raises(ValueError, match="load failed"): - loader4.load(lambda: (_ for _ in ()).throw(ValueError("load failed"))) - - # Concurrent access safety - loader_a = FileLockLoader(mock_args) - loader_b = FileLockLoader(mock_args) - # Clean up first - Path(temp_dataset_path).mkdir(parents=True, exist_ok=True) - for f in [loader_a.counter_path, loader_a.ready_flag_path]: - if f.exists(): - f.unlink() - r1 = loader_a.load(lambda: "data1") - r2 = loader_b.load(lambda: "data2") - assert r1 == "data1" and r2 == "data2" - assert loader_a.counter_path.read_text().strip() == "2" diff --git a/tests/data/prepare/test_hash.py b/tests/data/prepare/test_hash.py deleted file mode 100644 index 3cb5ef8a..00000000 --- a/tests/data/prepare/test_hash.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Tests for xorl.data.prepare.hash module.""" - -from unittest.mock import Mock - -import pytest -from datasets import Dataset - -from xorl.arguments import DatasetConfig -from xorl.data.prepare.hash import ( - generate_dataset_hash_from_config, - generate_packing_hash, - generate_split_fingerprints, -) - - -pytestmark = pytest.mark.cpu - - -def _make_config(path="dataset1"): - return DatasetConfig( - path=path, - type="tokenized", - shards=None, - shards_idx=None, - preprocess_shards=None, - name=None, - split="train", - revision=None, - trust_remote_code=False, - max_seq_len=None, - ) - - -class TestGenerateSplitFingerprints: - """Tests for generate_split_fingerprints function.""" - - def test_fingerprint_properties(self): - """Covers train/test difference, consistency, different inputs produce different results, and float val_set_size.""" - dataset = Mock(spec=Dataset) - dataset._fingerprint = "base_fingerprint" - - # Train and test are different, consistent, and correct length - t1, e1 = generate_split_fingerprints(dataset, val_set_size=100, seed=42) - t2, e2 = generate_split_fingerprints(dataset, val_set_size=100, seed=42) - assert t1 != e1 and len(t1) == 32 and len(e1) == 32 - assert t1 == t2 and e1 == e2 - - # Different val_set_size, seed, and dataset produce different fingerprints - dataset2 = Mock(spec=Dataset) - dataset2._fingerprint = "fingerprint2" - diff_val, _ = generate_split_fingerprints(dataset, val_set_size=200, seed=42) - diff_seed, _ = generate_split_fingerprints(dataset, val_set_size=100, seed=99) - diff_ds, _ = generate_split_fingerprints(dataset2, val_set_size=100, seed=42) - assert t1 != diff_val and t1 != diff_seed and t1 != diff_ds - - # Float val_set_size - tf, ef = generate_split_fingerprints(dataset, val_set_size=0.1, seed=42) - assert tf != ef and len(tf) == 32 - - -class TestGeneratePackingHash: - """Tests for generate_packing_hash function.""" - - def test_packing_hash_format_consistency_and_uniqueness(self): - """Covers consistent hash, format/content, and different params produce different hashes.""" - h1 = generate_packing_hash("multipack", 2048, 100, "fork") - h2 = generate_packing_hash("multipack", 2048, 100, "fork") - assert h1 == h2 - - parts = h1.split("_") - assert len(parts) == 4 - assert parts == ["multipack", "2048", "100", "fork"] - - assert h1 != generate_packing_hash("sequential", 2048, 100, "fork") - assert h1 != generate_packing_hash("multipack", 4096, 100, "fork") - assert h1 != generate_packing_hash("multipack", 2048, 200, "fork") - assert h1 != generate_packing_hash("multipack", 2048, 100, "spawn") - - -class TestGenerateDatasetHashFromConfig: - """Tests for generate_dataset_hash_from_config function.""" - - def test_hash_consistency_uniqueness_and_order_independence(self): - """Covers consistent MD5 hash, different configs produce different hashes, - multiple datasets, and order independence.""" - args = Mock() - args.data.select_columns = None - config = _make_config() - - # Consistent and valid MD5 - h1 = generate_dataset_hash_from_config(args, [config], "gpt2") - h2 = generate_dataset_hash_from_config(args, [config], "gpt2") - assert h1 == h2 and len(h1) == 32 - assert all(c in "0123456789abcdef" for c in h1) - - # Different tokenizer, path, and select_columns produce different hashes - args_cols = Mock() - args_cols.data.select_columns = ["col1", "col2"] - assert h1 != generate_dataset_hash_from_config(args, [config], "llama") - assert h1 != generate_dataset_hash_from_config(args, [_make_config("dataset2")], "gpt2") - assert h1 != generate_dataset_hash_from_config(args_cols, [config], "gpt2") - - # Multiple datasets change hash - config2 = _make_config("dataset2") - assert h1 != generate_dataset_hash_from_config(args, [config, config2], "gpt2") - - # Order independence - h_ab = generate_dataset_hash_from_config(args, [config, config2], "gpt2") - h_ba = generate_dataset_hash_from_config(args, [config2, config], "gpt2") - assert h_ab == h_ba diff --git a/tests/data/prepare/test_packing.py b/tests/data/prepare/test_packing.py index 4066f69c..8184bdc7 100644 --- a/tests/data/prepare/test_packing.py +++ b/tests/data/prepare/test_packing.py @@ -11,10 +11,8 @@ add_position_ids, allocate_sequentially, drop_no_trainable_tokens, - ffd_check, filter_dataset_with_logging, pack_group, - pack_parallel, process_datasets_for_packing, ) @@ -22,22 +20,13 @@ pytestmark = pytest.mark.cpu -def test_ffd_check(): - """FFD feasibility check: fit, don't fit, edge cases.""" - # Sequences that fit - assert ffd_check(np.array([10, 20, 30, 40]), 50, 3) is True - assert ffd_check(np.array([30]), 50, 1) is True - assert ffd_check(np.array([25, 25]), 50, 1) is True # exact fit +def _assert_packing_allocation_primitives(): + """Live bin packing and rank allocation preserve capacity and coverage.""" + _assert_pack_group_policy() + _assert_sequential_allocation_policy() - # Sequences that don't fit - assert ffd_check(np.array([40, 40, 40, 40]), 50, 2) is False - assert ffd_check(np.array([60]), 50, 1) is False - # Empty - assert ffd_check(np.array([]), 50, 1) is True - - -def test_pack_group(): +def _assert_pack_group_policy(): """Pack sequences into bins: capacity, bin_size limit, safe/non-safe mode, offset.""" # Basic packing respects capacity seq = np.array([10, 20, 30, 40]) @@ -65,7 +54,7 @@ def test_pack_group(): assert min(idx for b in bins for idx in b) >= 100 -def test_allocate_sequentially(): +def _assert_sequential_allocation_policy(): """Sequential allocation: distributes to ranks, no overlap, full coverage.""" seq = np.array([10, 20, 30, 40, 50, 60]) @@ -90,8 +79,8 @@ def test_allocate_sequentially(): assert all_idx == {0, 1, 2, 3} -def test_add_position_ids(): - """Add position_ids: single, batched, missing/empty input_ids, preserves fields.""" +def _assert_sample_metadata_and_trainable_token_filter_policy(): + """Sample preprocessing adds positions and rejects data with no trainable labels.""" # Single sample result = add_position_ids({"input_ids": [1, 2, 3, 4, 5], "labels": [1, 2, 3, 4, 5]}) assert result["position_ids"] == [0, 1, 2, 3, 4] @@ -112,11 +101,13 @@ def test_add_position_ids(): result = add_position_ids({"input_ids": [1, 2, 3], "attention_mask": [1, 1, 1], "labels": [1, 2, 3]}) assert result["attention_mask"] == [1, 1, 1] + _assert_trainable_token_filter_policy() + -def test_drop_no_trainable_tokens(): +def _assert_trainable_token_filter_policy(): """Drop samples with no trainable tokens, handle batched, raise on missing labels.""" - assert drop_no_trainable_tokens({"labels": [1, 2, -100, 3]}) == True - assert bool(drop_no_trainable_tokens({"labels": [-100, -100, -100]})) == False + assert drop_no_trainable_tokens({"labels": [1, 2, -100, 3]}) is True + assert bool(drop_no_trainable_tokens({"labels": [-100, -100, -100]})) is False # Batched result = drop_no_trainable_tokens({"labels": [[1, 2, 3], [-100, -100, -100], [1, -100, 2]]}) @@ -127,7 +118,7 @@ def test_drop_no_trainable_tokens(): drop_no_trainable_tokens({"input_ids": [1, 2, 3]}) -def test_filter_dataset_with_logging(): +def _assert_dataset_filtering_policy(): """Filter dataset and verify correct samples are kept.""" dataset = HFDataset.from_dict( { @@ -141,8 +132,10 @@ def test_filter_dataset_with_logging(): assert filtered[1]["labels"] == [7, 8, 9] -def test_process_datasets_for_packing(): - """Process train+eval datasets: adds position_ids/length, handles None eval.""" +def _assert_dataset_preprocessing_pipeline(): + """Dataset preprocessing filters records, adds packing metadata, and handles optional eval data.""" + _assert_dataset_filtering_policy() + args = Mock() args.data.dataset_num_proc = 1 @@ -159,8 +152,12 @@ def test_process_datasets_for_packing(): assert p_eval_none is None -def test_packing_dataset(): +def test_packing_dataset(tmp_path): """PackingDataset: init, bins, getitem, cache, missing length column.""" + _assert_sample_metadata_and_trainable_token_filter_policy() + _assert_dataset_preprocessing_pipeline() + _assert_packing_allocation_primitives() + args = Mock() args.data.sample_packing_method = "sequential" args.data.sample_packing_sequence_len = 100 @@ -198,6 +195,15 @@ def test_packing_dataset(): sample = pds[0] assert isinstance(sample, list) and len(sample) > 0 and "input_ids" in sample[0] + # Cache identity must reflect the production ring-attention document + # alignment, not merely the standalone string-formatting helper. + pds.prepared_dataset_path = str(tmp_path) + plain_cache_path = pds._get_bins_cache_path() + pds.doc_align = 4 + ring_cache_path = pds._get_bins_cache_path() + assert plain_cache_path.name == "packing_bins_sequential_100_10_None" + assert ring_cache_path.name == "packing_bins_sequential_100_10_None_align4" + # Multipack method also works args.data.sample_packing_method = "multipack" pds2 = PackingDataset(args, tokenizer, dataset, split="train") @@ -215,14 +221,3 @@ def test_packing_dataset(): args.data.sample_packing_method = "sequential" pds3 = PackingDataset(args, tokenizer, dataset, split="train") assert pds3.bins == [[0, 1], [2, 3]] - - -def test_pack_parallel(): - """pack_parallel: single process and auto num_processes.""" - seq = np.array([10, 20, 30, 40, 50, 60, 70, 80]) - - bins = pack_parallel(seq, 100, 3, 10, num_processes=1, safe_mode=True, mp_start_method=None) - assert len(bins) > 0 - - bins2 = pack_parallel(seq, 100, 2, 10, num_processes=None, safe_mode=True, mp_start_method=None) - assert len(bins2) > 0 diff --git a/tests/data/prepare/test_shared.py b/tests/data/prepare/test_shared.py index dc854f08..ae2b9a5e 100644 --- a/tests/data/prepare/test_shared.py +++ b/tests/data/prepare/test_shared.py @@ -3,9 +3,12 @@ from unittest.mock import Mock, patch import pytest +import requests from datasets import Dataset as HFDataset +from huggingface_hub.errors import HfHubHTTPError from xorl.arguments import DatasetConfig +from xorl.data.prepare.hash import generate_dataset_hash_from_config, generate_split_fingerprints from xorl.data.prepare.shared import ( create_train_validation_split, datasets_with_name_generator, @@ -15,6 +18,7 @@ merge_datasets, save_preprocessed_dataset, ) +from xorl.data.prepare.utils import retry_on_request_exceptions pytestmark = pytest.mark.cpu @@ -41,7 +45,7 @@ def _make_config(**overrides): class TestDatasetsWithNameGeneratorAndDatasetType: """Tests for datasets_with_name_generator and get_dataset_type.""" - def test_expansion_and_passthrough_behaviors(self): + def test_dataset_preparation_lifecycle(self, tmp_path, monkeypatch): """Covers name expansion, preprocess_shards expansion, shards-blocks expansion, and get_dataset_type inference from extension and explicit ds_type.""" # Multiple names expansion @@ -57,10 +61,8 @@ def test_expansion_and_passthrough_behaviors(self): assert len(result) == 3 assert [r.shards_idx for r in result] == [0, 1, 2] - # Preprocess_shards NOT expanded when shards already set - config_no_expand = _make_config(path="ds1", shards=4, preprocess_shards=3) - result = list(datasets_with_name_generator([config_no_expand])) - assert len(result) == 1 + with pytest.raises(ValueError, match="mutually exclusive"): + _make_config(path="ds1", shards=4, preprocess_shards=3) # get_dataset_type: explicit ds_type overrides extension config_explicit = _make_config(path="data.parquet", ds_type="arrow") @@ -78,11 +80,94 @@ def test_expansion_and_passthrough_behaviors(self): for path, expected_type in extension_map: assert get_dataset_type(_make_config(path=path)) == expected_type + TestSplitAndMerge()._assert_split_and_merge_operations() + load_root = tmp_path / "load" + load_root.mkdir() + with monkeypatch.context() as load_patch: + TestLoadDatasetWithConfig()._assert_local_and_hub_loading(load_root, load_patch) + save_root = tmp_path / "save" + save_root.mkdir() + TestSaveAndLoadPreprocessedDataset()._assert_save_load_and_missing(save_root) + _assert_dataset_hash_and_split_fingerprint_policy() + _assert_request_retry_policy(monkeypatch) + + +def _assert_request_retry_policy(monkeypatch): + sleeps = [] + monkeypatch.setattr("xorl.data.prepare.utils.time.sleep", sleeps.append) + + transient = Mock( + side_effect=[requests.exceptions.ReadTimeout("timeout"), requests.exceptions.ReadTimeout("timeout"), "success"] + ) + wrapped = retry_on_request_exceptions(max_retries=3, delay=0.01)(transient) + assert wrapped() == "success" + assert transient.call_count == 3 + assert sleeps == [0.01, 0.02] + + response = Mock(status_code=500, headers={}) + hub_transient = Mock(side_effect=[HfHubHTTPError("HF error", response=response), "success"]) + wrapped = retry_on_request_exceptions(max_retries=3, delay=0.01)(hub_transient) + sleeps.clear() + assert wrapped() == "success" + assert sleeps == [0.01] + + persistent = Mock(side_effect=requests.exceptions.ReadTimeout("persistent timeout")) + wrapped = retry_on_request_exceptions(max_retries=2, delay=0.01)(persistent) + sleeps.clear() + with pytest.raises(requests.exceptions.ReadTimeout): + wrapped() + assert persistent.call_count == 2 + assert sleeps == [0.01] + + unrelated = Mock(side_effect=ValueError("not a request exception")) + wrapped = retry_on_request_exceptions(max_retries=3, delay=0.01)(unrelated) + sleeps.clear() + with pytest.raises(ValueError): + wrapped() + assert unrelated.call_count == 1 + assert sleeps == [] + + +def _assert_dataset_hash_and_split_fingerprint_policy(): + dataset = Mock(spec=HFDataset) + dataset._fingerprint = "base_fingerprint" + train, evaluation = generate_split_fingerprints(dataset, val_set_size=100, seed=42) + train_again, evaluation_again = generate_split_fingerprints(dataset, val_set_size=100, seed=42) + assert train != evaluation + assert (train, evaluation) == (train_again, evaluation_again) + assert len(train) == len(evaluation) == 32 + + dataset_two = Mock(spec=HFDataset) + dataset_two._fingerprint = "fingerprint2" + assert train != generate_split_fingerprints(dataset, val_set_size=200, seed=42)[0] + assert train != generate_split_fingerprints(dataset, val_set_size=100, seed=99)[0] + assert train != generate_split_fingerprints(dataset_two, val_set_size=100, seed=42)[0] + fractional_train, fractional_evaluation = generate_split_fingerprints(dataset, val_set_size=0.1, seed=42) + assert fractional_train != fractional_evaluation + + args = Mock() + args.data.select_columns = None + config = _make_config() + dataset_hash = generate_dataset_hash_from_config(args, [config], "gpt2") + assert dataset_hash == generate_dataset_hash_from_config(args, [config], "gpt2") + assert len(dataset_hash) == 32 + + args_with_columns = Mock() + args_with_columns.data.select_columns = ["col1", "col2"] + config_two = _make_config(path="dataset2") + assert dataset_hash != generate_dataset_hash_from_config(args, [config], "llama") + assert dataset_hash != generate_dataset_hash_from_config(args, [config_two], "gpt2") + assert dataset_hash != generate_dataset_hash_from_config(args_with_columns, [config], "gpt2") + assert dataset_hash != generate_dataset_hash_from_config(args, [config, config_two], "gpt2") + assert generate_dataset_hash_from_config(args, [config, config_two], "gpt2") == generate_dataset_hash_from_config( + args, [config_two, config], "gpt2" + ) + class TestSplitAndMerge: """Tests for create_train_validation_split and merge_datasets.""" - def test_split_and_merge_operations(self): + def _assert_split_and_merge_operations(self): """Covers absolute/fractional split, merge with shuffle variants, and empty merge error.""" dataset = HFDataset.from_dict( { @@ -103,24 +188,25 @@ def test_split_and_merge_operations(self): assert len(train_ds) == 8 assert len(eval_ds) == 2 - # Merge without shuffle - ds1 = HFDataset.from_dict({"input_ids": [[1, 2, 3]], "labels": [[1, 2, 3]]}) - ds2 = HFDataset.from_dict({"input_ids": [[4, 5, 6]], "labels": [[4, 5, 6]]}) + # The shuffle modes must change ordering, not merely preserve row count. + ds1 = HFDataset.from_dict({"input_ids": [[i] for i in range(3)], "labels": [[i] for i in range(3)]}) + ds2 = HFDataset.from_dict({"input_ids": [[i] for i in range(3, 6)], "labels": [[i] for i in range(3, 6)]}) args.data.shuffle_merged_datasets = False args.data.shuffle_before_merging_datasets = False - assert len(merge_datasets([ds1, ds2], args)) == 2 + merged = merge_datasets([ds1, ds2], args) + assert [row[0] for row in merged["input_ids"]] == list(range(6)) - # Merge with shuffle_merged_datasets - ds1 = HFDataset.from_dict({"input_ids": [[i] for i in range(3)], "labels": [[i] for i in range(3)]}) - ds2 = HFDataset.from_dict({"input_ids": [[i] for i in range(3)], "labels": [[i] for i in range(3)]}) args.data.shuffle_merged_datasets = True - args.data.shuffle_before_merging_datasets = False - assert len(merge_datasets([ds1, ds2], args)) == 6 + shuffled_values = [row[0] for row in merge_datasets([ds1, ds2], args)["input_ids"]] + assert sorted(shuffled_values) == list(range(6)) + assert shuffled_values != list(range(6)) - # Merge with shuffle_before_merging args.data.shuffle_merged_datasets = False args.data.shuffle_before_merging_datasets = True - assert len(merge_datasets([ds1, ds2], args)) == 6 + individually_shuffled = [row[0] for row in merge_datasets([ds1, ds2], args)["input_ids"]] + assert set(individually_shuffled[:3]) == set(range(3)) + assert set(individually_shuffled[3:]) == set(range(3, 6)) + assert individually_shuffled != list(range(6)) # Empty dataset list raises ValueError args.data.shuffle_merged_datasets = False @@ -131,7 +217,7 @@ def test_split_and_merge_operations(self): class TestLoadDatasetWithConfig: """Tests for load_dataset_with_config function.""" - def test_local_and_hub_loading(self, tmp_path, monkeypatch): + def _assert_local_and_hub_loading(self, tmp_path, monkeypatch): import datasets.config hf_cache = str(tmp_path / "hf_cache") @@ -182,10 +268,12 @@ def test_local_and_hub_loading(self, tmp_path, monkeypatch): with pytest.raises(ValueError, match="The dataset could not be loaded"): load_dataset_with_config(config, use_auth_token=False, streaming=False) + self._assert_data_files_string_and_list() + @patch("xorl.data.prepare.shared._check_if_hub_dataset") @patch("xorl.data.prepare.shared.hf_hub_download") @patch("xorl.data.prepare.shared.load_dataset") - def test_data_files_string_and_list(self, mock_load_dataset, mock_hub_download, mock_check_hub): + def _assert_data_files_string_and_list(self, mock_load_dataset, mock_hub_download, mock_check_hub): """Covers loading from data_files as string and as list.""" mock_check_hub.return_value = False mock_load_dataset.return_value = HFDataset.from_dict({"input_ids": [[1]], "labels": [[1]]}) @@ -195,19 +283,23 @@ def test_data_files_string_and_list(self, mock_load_dataset, mock_hub_download, config = _make_config(path="user/ds", split=None, data_files="data.json", ds_type="json") assert load_dataset_with_config(config, use_auth_token=False, streaming=False) is not None mock_hub_download.assert_called_once() + assert mock_load_dataset.call_args.args == ("json",) + assert mock_load_dataset.call_args.kwargs["data_files"] == "/tmp/file.json" # data_files as list mock_hub_download.reset_mock() - mock_hub_download.side_effect = ["/tmp/file1.json", "/tmp/file2.json"] - config = _make_config(path="user/ds", split=None, data_files=["d1.json", "d2.json"], ds_type="json") + mock_hub_download.side_effect = ["/tmp/file1.parquet", "/tmp/file2.parquet"] + config = _make_config(path="user/ds", split=None, data_files=["d1.parquet", "d2.parquet"], ds_type="parquet") assert load_dataset_with_config(config, use_auth_token=False, streaming=False) is not None assert mock_hub_download.call_count == 2 + assert mock_load_dataset.call_args.args == ("parquet",) + assert mock_load_dataset.call_args.kwargs["data_files"] == ["/tmp/file1.parquet", "/tmp/file2.parquet"] class TestSaveAndLoadPreprocessedDataset: """Tests for save/load preprocessed dataset functions.""" - def test_save_load_and_missing(self, tmp_path): + def _assert_save_load_and_missing(self, tmp_path): """Covers save+load round-trip and load returning None when not found.""" args = Mock() diff --git a/tests/data/prepare/test_utils.py b/tests/data/prepare/test_utils.py deleted file mode 100644 index 75407c91..00000000 --- a/tests/data/prepare/test_utils.py +++ /dev/null @@ -1,119 +0,0 @@ -"""Tests for xorl.data.prepare.utils module.""" - -import time -from unittest.mock import Mock - -import huggingface_hub -import pytest -import requests - -from xorl.data.prepare.utils import ( - RetryStrategy, - md5, - retry_on_request_exceptions, - sha256, -) - - -pytestmark = pytest.mark.cpu - - -class TestRetryOnRequestExceptions: - """Tests for retry_on_request_exceptions decorator.""" - - def test_success_retries_and_max_retries(self): - """Covers immediate success, retry on ReadTimeout/HfHubHTTPError, max retries exhaustion, - and non-request exception passthrough.""" - - # Immediate success - @retry_on_request_exceptions(max_retries=3, delay=0.01) - def success_func(): - return "success" - - assert success_func() == "success" - - # Retry on ReadTimeout - mock_func = Mock( - side_effect=[requests.exceptions.ReadTimeout("t"), requests.exceptions.ReadTimeout("t"), "success"] - ) - - @retry_on_request_exceptions(max_retries=3, delay=0.01) - def retry_func(): - return mock_func() - - assert retry_func() == "success" - assert mock_func.call_count == 3 - - # Retry on HfHubHTTPError - mock_response = Mock() - mock_response.status_code = 500 - mock_response.headers = {} - mock_func2 = Mock( - side_effect=[huggingface_hub.errors.HfHubHTTPError("HF error", response=mock_response), "success"] - ) - - @retry_on_request_exceptions(max_retries=3, delay=0.01) - def hf_func(): - return mock_func2() - - assert hf_func() == "success" - - # Max retries exhausted - @retry_on_request_exceptions(max_retries=2, delay=0.01) - def failing_func(): - raise requests.exceptions.ReadTimeout("persistent timeout") - - with pytest.raises(requests.exceptions.ReadTimeout): - failing_func() - - # Non-request exception not caught - @retry_on_request_exceptions(max_retries=3, delay=0.01) - def value_error_func(): - raise ValueError("not a request exception") - - with pytest.raises(ValueError): - value_error_func() - - def test_backoff_strategies(self): - """Covers exponential, linear, and constant backoff timing.""" - for strategy, check_fn in [ - (RetryStrategy.EXPONENTIAL, lambda d1, d2: d2 > d1 * 1.3), - (RetryStrategy.LINEAR, lambda d1, d2: d2 > d1 * 1.3), - (RetryStrategy.CONSTANT, lambda d1, d2: abs(d2 - d1) < d1 * 0.5), - ]: - call_times = [] - mock_func = Mock( - side_effect=[requests.exceptions.ReadTimeout("t"), requests.exceptions.ReadTimeout("t"), "success"] - ) - - @retry_on_request_exceptions(max_retries=3, delay=0.1, retry_strategy=strategy) - def func(): - call_times.append(time.time()) - return mock_func() - - assert func() == "success" - if len(call_times) >= 3: - d1 = call_times[1] - call_times[0] - d2 = call_times[2] - call_times[1] - assert 0.05 < d1 < 0.25 - assert check_fn(d1, d2) - - -class TestHashFunctions: - """Tests for md5 and sha256 hash functions.""" - - def test_md5_and_sha256(self): - """Covers known hashes, different inputs produce different hashes, and unicode handling.""" - # MD5 known value - assert md5("hello") == "5d41402abc4b2a76b9719d911017c592" - # MD5 different inputs - assert md5("string1") != md5("string2") - # MD5 unicode - assert len(md5("测试中文")) == 32 - - # SHA256 known value - assert sha256("hello") == "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824" - # SHA256 different inputs - assert sha256("string1") != sha256("string2") - # SHA256 unicode - assert len(sha256("测试中文")) == 64 diff --git a/tests/data/test_data_loader.py b/tests/data/test_data_loader.py index 6a6d74d1..154f74c4 100644 --- a/tests/data/test_data_loader.py +++ b/tests/data/test_data_loader.py @@ -29,7 +29,7 @@ def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, torch.Tensor class TestMicroBatchCollatorAndDistributedDataloader: """Tests for MicroBatchCollator splitting, DistributedDataloader.set_epoch, and DataLoaderBuilder.""" - def test_micro_batch_splitting_set_epoch_and_builder(self): + def _assert_micro_batch_splitting_and_set_epoch(self): """Covers micro-batch splitting, edge cases, set_epoch delegation, batch size, sampler, SP collator, and custom collators.""" internal_collator = SimpleCollator() @@ -88,10 +88,12 @@ def test_micro_batch_splitting_set_epoch_and_builder(self): @patch("xorl.data.data_loader.get_parallel_state") @patch("xorl.data.data_loader.StatefulDistributedSampler") @patch("xorl.data.data_loader.DistributedDataloader") - def test_builder_batch_size_sampler_sp_and_custom_collators( + def test_micro_batch_and_builder_configuration_policy( self, mock_dataloader_cls, mock_sampler_cls, mock_parallel_state, fake_text_dataset ): """Covers batch size, sampler params, SP collator, single/multiple custom collators.""" + self._assert_micro_batch_splitting_and_set_epoch() + mock_ps = Mock() mock_ps.dp_size = 2 mock_ps.dp_rank = 0 @@ -241,7 +243,9 @@ def __getitem__(self, idx): "position_ids": torch.arange(self.seq_len, dtype=torch.long), } - dataset = SimpleTokenizedDataset(num_samples=16, seq_len=5) + # Ten samples form two complete four-sample loader batches; the incomplete + # remainder is dropped by the production dataloader. + dataset = SimpleTokenizedDataset(num_samples=10, seq_len=5) int_builder = DataLoaderBuilder( dataset=dataset, micro_batch_size=2, @@ -251,6 +255,7 @@ def __getitem__(self, idx): pad_to_multiple_of=1, ) dataloader = int_builder.build(verbose=False) + assert len(dataloader) == 2 for step, micro_batches in enumerate(dataloader): assert isinstance(micro_batches, list) diff --git a/tests/data/test_data_loader_distributed.py b/tests/data/test_data_loader_distributed.py deleted file mode 100644 index cfa67ef3..00000000 --- a/tests/data/test_data_loader_distributed.py +++ /dev/null @@ -1,405 +0,0 @@ -""" -Integration tests for distributed data loading. - -These tests verify that the data loader correctly distributes data across multiple -processes/ranks in a distributed training setting, and that the data alignment is correct. -""" - -from unittest.mock import Mock, patch - -import pytest -import torch -from torch.utils.data import Dataset - -from tests.conftest import FakeTextDataset, SimpleCollator -from xorl.data.collators import CollatePipeline, TextSequenceShardCollator -from xorl.data.constants import IGNORE_INDEX -from xorl.data.data_loader import ( - DataLoaderBuilder, - MicroBatchCollator, -) - - -pytestmark = [pytest.mark.gpu, pytest.mark.dataloader, pytest.mark.distributed] - - -class TestDistributedDataAlignment: - """Tests for data partitioning, micro-batch splitting, SP sharding, batch size, and drop_last.""" - - @patch("xorl.data.data_loader.get_parallel_state") - @patch("xorl.data.data_loader.StatefulDistributedSampler") - def test_partitioning_micro_batches_sp_batch_size_and_drop_last( - self, mock_sampler_cls, mock_parallel_state, fake_text_dataset - ): - """Covers data partitioning across ranks, micro-batch splitting, SP pipeline setup, - batch size calculation, and drop_last behavior.""" - # --- Data partitioning across 4 DP ranks --- - dp_size = 4 - datasets_per_rank = [] - for rank in range(dp_size): - mock_ps = Mock() - mock_ps.dp_size = dp_size - mock_ps.dp_rank = rank - mock_ps.cp_size = 1 - mock_ps.cp_enabled = False - mock_parallel_state.return_value = mock_ps - - mock_sampler = Mock() - rank_indices = list(range(rank, len(fake_text_dataset), dp_size)) - mock_sampler.__iter__ = Mock(return_value=iter(rank_indices)) - mock_sampler.__len__ = Mock(return_value=len(rank_indices)) - mock_sampler_cls.return_value = mock_sampler - - builder = DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=2) - builder.build(verbose=False) - assert mock_sampler_cls.call_args[1]["rank"] == rank - assert mock_sampler_cls.call_args[1]["num_replicas"] == dp_size - datasets_per_rank.append(rank_indices) - - all_indices = [idx for indices in datasets_per_rank for idx in indices] - assert len(set(all_indices)) == len(all_indices) # No overlap - assert len(set(all_indices)) >= len(fake_text_dataset) - dp_size - - # --- Micro-batch splitting correctness --- - mock_ps = Mock() - mock_ps.dp_size = 1 - mock_ps.dp_rank = 0 - mock_ps.cp_size = 1 - mock_ps.cp_enabled = False - mock_parallel_state.return_value = mock_ps - - builder = DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=4, gradient_accumulation_steps=3) - dataloader = builder.build(verbose=False) - micro_batches = next(iter(dataloader)) - assert isinstance(micro_batches, list) and len(micro_batches) == 3 - for mb in micro_batches: - assert isinstance(mb, dict) - assert "input_ids" in mb and "labels" in mb - assert len(mb["input_ids"].shape) == 2 - - # --- SP pipeline verification --- - mock_ps.cp_size = 2 - mock_ps.cp_enabled = True - mock_parallel_state.return_value = mock_ps - builder_sp = DataLoaderBuilder(dataset=fake_text_dataset, micro_batch_size=2, gradient_accumulation_steps=1) - - mock_sampler = Mock() - mock_sampler.__iter__ = Mock(return_value=iter(range(len(fake_text_dataset)))) - mock_sampler.__len__ = Mock(return_value=len(fake_text_dataset)) - mock_sampler_cls.return_value = mock_sampler - - dataloader_sp = builder_sp.build(verbose=False) - - internal = builder_sp.collate_fn.internal_collator - assert isinstance(internal, CollatePipeline) - assert any(isinstance(c, TextSequenceShardCollator) for c in internal.data_collators) - - # --- Batch size calculation --- - mock_ps.dp_size = 2 - mock_ps.cp_size = 1 - mock_ps.cp_enabled = False - mock_parallel_state.return_value = mock_ps - assert 4 * 3 == 12 # micro_batch_size * gradient_accumulation_steps - - # --- Drop last behavior --- - mock_ps.dp_size = 1 - mock_ps.dp_rank = 0 - mock_ps.cp_size = 1 - mock_ps.cp_enabled = False - mock_parallel_state.return_value = mock_ps - - small_dataset = FakeTextDataset(num_samples=10, seq_len=64) - - # Reset mock sampler for the small dataset - mock_sampler = Mock() - mock_sampler.__iter__ = Mock(return_value=iter(range(len(small_dataset)))) - mock_sampler.__len__ = Mock(return_value=len(small_dataset)) - mock_sampler_cls.return_value = mock_sampler - - builder_dl = DataLoaderBuilder( - dataset=small_dataset, micro_batch_size=2, gradient_accumulation_steps=2, drop_last=True - ) - dl = builder_dl.build(verbose=False) - batches = list(dl) - assert len(batches) == 2 - - -class TestMicroBatchCollatorAndSequenceSharding: - """Tests for micro-batch data order, gradient accumulation alignment, error handling, - and sequence sharding across SP ranks.""" - - def test_order_preservation_ga_alignment_and_error(self): - """Covers data order within micro-batches, gradient accumulation alignment, and incorrect batch size error.""" - - collator = MicroBatchCollator( - micro_batch_size=2, gradient_accumulation_steps=3, internal_collator=SimpleCollator() - ) - features = [{"input_ids": torch.tensor([i]), "value": torch.tensor([i * 100])} for i in range(6)] - micro_batches = collator(features) - assert torch.equal(micro_batches[0]["value"], torch.tensor([[0], [100]])) - assert torch.equal(micro_batches[1]["value"], torch.tensor([[200], [300]])) - assert torch.equal(micro_batches[2]["value"], torch.tensor([[400], [500]])) - - collator2 = MicroBatchCollator( - micro_batch_size=3, gradient_accumulation_steps=2, internal_collator=SimpleCollator() - ) - features2 = [{"input_ids": torch.tensor([i]), "labels": torch.tensor([i])} for i in range(6)] - mbs = collator2(features2) - assert len(mbs) == 2 - for mb in mbs: - assert mb["input_ids"].shape[0] == 3 - - collator3 = MicroBatchCollator( - micro_batch_size=4, gradient_accumulation_steps=2, internal_collator=SimpleCollator() - ) - features3 = [{"input_ids": torch.tensor([i]), "labels": torch.tensor([i])} for i in range(7)] - with pytest.raises(ValueError, match="Expected 8 samples"): - collator3(features3) - - @patch("xorl.data.collators.sequence_shard_collator.get_parallel_state") - def test_sequence_sharding_and_padding(self, mock_parallel_state): - """Covers correct chunk sizes across SP ranks and padding for non-divisible lengths.""" - - cp_size = 4 - seq_len = 128 - input_ids = torch.arange(seq_len).unsqueeze(0) - labels = torch.cat([torch.arange(1, seq_len), torch.tensor([IGNORE_INDEX])]).unsqueeze(0) - - sharded = [] - for cp_rank in range(cp_size): - mock_ps = Mock() - mock_ps.cp_size = cp_size - mock_ps.cp_rank = cp_rank - mock_ps.ringattn_size = 1 - mock_parallel_state.return_value = mock_ps - collator = TextSequenceShardCollator(pad_token_id=0) - batch = { - "input_ids": input_ids.clone(), - "attention_mask": torch.ones(1, seq_len, dtype=torch.long), - "labels": labels.clone(), - "position_ids": torch.arange(seq_len).unsqueeze(0), - } - result = collator(batch) - sharded.append(result["input_ids"]) - - expected_chunk = seq_len // cp_size - for rank, s in enumerate(sharded): - assert s.shape[1] == expected_chunk, f"Rank {rank} has incorrect chunk size" - - # Padding alignment for non-divisible length - seq_len_odd = 130 - mock_ps = Mock() - mock_ps.cp_size = cp_size - mock_ps.cp_rank = 0 - mock_ps.ringattn_size = 1 - mock_parallel_state.return_value = mock_ps - collator = TextSequenceShardCollator(pad_token_id=0) - batch = { - "input_ids": torch.arange(seq_len_odd).unsqueeze(0), - "attention_mask": torch.ones(1, seq_len_odd, dtype=torch.long), - "labels": torch.cat([torch.arange(1, seq_len_odd), torch.tensor([IGNORE_INDEX])]).unsqueeze(0), - "position_ids": torch.arange(seq_len_odd).unsqueeze(0), - } - result = collator(batch) - total_padded = result["input_ids"].shape[1] * cp_size - assert total_padded >= seq_len_odd - assert total_padded < seq_len_odd + cp_size - - -class TestEndToEndAndPackedSequences: - """End-to-end integration tests for distributed data loading and packed sequences.""" - - @patch("xorl.data.data_loader.get_parallel_state") - @patch("xorl.data.collators.packing_concat_collator.get_parallel_state") - def test_pipeline_epoch_consistency_and_packed_sequences( - self, mock_ps_collator, mock_ps_loader, fake_packed_dataset, fake_text_dataset - ): - """Covers complete pipeline output, epoch consistency, packed sequence flattening, - multi-DP, SP with packing, and variable lengths.""" - mock_ps = Mock() - mock_ps.dp_size = 2 - mock_ps.dp_rank = 0 - mock_ps.cp_size = 1 - mock_ps.cp_enabled = False - mock_ps_collator.return_value = mock_ps - mock_ps_loader.return_value = mock_ps - - # Complete pipeline output - builder = DataLoaderBuilder( - dataset=fake_packed_dataset, - micro_batch_size=2, - gradient_accumulation_steps=2, - num_workers=0, - prefetch_factor=None, - seed=42, - ) - dataloader = builder.build(verbose=False) - - batch_count = 0 - for micro_batches in dataloader: - assert isinstance(micro_batches, list) and len(micro_batches) == 2 - for mb in micro_batches: - for field in ["input_ids", "labels", "position_ids", "attention_mask"]: - assert field in mb - assert "cu_seq_lens_q" in mb - bs, sl = mb["input_ids"].shape - assert mb["labels"].shape == (bs, sl) - assert mb["position_ids"].shape == (bs, sl) - assert mb["attention_mask"].shape == (bs, sl) - batch_count += 1 - if batch_count >= 3: - break - assert batch_count == 3 - - # Epoch consistency - mock_ps.dp_size = 1 - mock_ps_loader.return_value = mock_ps - mock_ps_collator.return_value = mock_ps - builder2 = DataLoaderBuilder( - dataset=fake_text_dataset, - micro_batch_size=4, - gradient_accumulation_steps=1, - num_workers=0, - prefetch_factor=None, - seed=42, - ) - dl2 = builder2.build(verbose=False) - - epoch1 = [next(iter(dl2))[0]["input_ids"].clone() for _ in range(3)] - dl2.set_epoch(1) - epoch2 = [next(iter(dl2))[0]["input_ids"].clone() for _ in range(3)] - assert len(epoch1) == len(epoch2) - - # --- Packed sequence tests --- - - class PackedDataset(Dataset): - def __len__(self): - return 8 - - def __getitem__(self, idx): - return [ - { - "input_ids": torch.tensor([idx * 10, idx * 10 + 1, idx * 10 + 2]), - "labels": torch.tensor([idx * 10, idx * 10 + 1, idx * 10 + 2]), - "position_ids": torch.tensor([0, 1, 2]), - "attention_mask": torch.ones(3, dtype=torch.long), - }, - { - "input_ids": torch.tensor([idx * 10 + 3, idx * 10 + 4]), - "labels": torch.tensor([idx * 10 + 3, idx * 10 + 4]), - "position_ids": torch.tensor([0, 1]), - "attention_mask": torch.ones(2, dtype=torch.long), - }, - ] - - mock_ps.dp_size = 1 - mock_ps.dp_rank = 0 - mock_ps.cp_size = 1 - mock_ps.cp_enabled = False - mock_ps_collator.return_value = mock_ps - mock_ps_loader.return_value = mock_ps - - builder3 = DataLoaderBuilder( - dataset=PackedDataset(), - micro_batch_size=2, - gradient_accumulation_steps=1, - num_workers=0, - prefetch_factor=None, - pad_to_multiple_of=1, - ) - dl3 = builder3.build(verbose=False) - mbs = next(iter(dl3)) - assert len(mbs) == 1 - assert mbs[0]["input_ids"].shape[0] == 1 - assert mbs[0]["input_ids"].shape[1] == 6 - assert mbs[0]["attention_mask"].shape == (1, 6) - - # Multiple DP ranks - class SimplePacked(Dataset): - def __len__(self): - return 16 - - def __getitem__(self, idx): - return [ - { - "input_ids": torch.tensor([idx, idx + 1]), - "labels": torch.tensor([idx, idx + 1]), - "position_ids": torch.tensor([0, 1]), - "attention_mask": torch.ones(2, dtype=torch.long), - } - ] - - for dp_rank in [0, 1]: - mock_ps.dp_size = 2 - mock_ps.dp_rank = dp_rank - mock_ps_collator.return_value = mock_ps - mock_ps_loader.return_value = mock_ps - builder4 = DataLoaderBuilder( - dataset=SimplePacked(), - micro_batch_size=2, - gradient_accumulation_steps=1, - num_workers=0, - prefetch_factor=None, - pad_to_multiple_of=1, - ) - dl4 = builder4.build(verbose=False) - mbs = next(iter(dl4)) - assert len(mbs) == 1 - assert mbs[0]["input_ids"].shape == (1, 2) - - # Variable lengths - class VariablePackedDataset(Dataset): - def __len__(self): - return 4 - - def __getitem__(self, idx): - if idx % 2 == 0: - return [ - { - "input_ids": torch.tensor([idx, idx + 1]), - "labels": torch.tensor([idx, idx + 1]), - "position_ids": torch.tensor([0, 1]), - "attention_mask": torch.ones(2, dtype=torch.long), - } - ] - else: - return [ - { - "input_ids": torch.tensor([idx]), - "labels": torch.tensor([idx]), - "position_ids": torch.tensor([0]), - "attention_mask": torch.ones(1, dtype=torch.long), - }, - { - "input_ids": torch.tensor([idx + 1, idx + 2]), - "labels": torch.tensor([idx + 1, idx + 2]), - "position_ids": torch.tensor([0, 1]), - "attention_mask": torch.ones(2, dtype=torch.long), - }, - { - "input_ids": torch.tensor([idx + 3]), - "labels": torch.tensor([idx + 3]), - "position_ids": torch.tensor([0]), - "attention_mask": torch.ones(1, dtype=torch.long), - }, - ] - - mock_ps.dp_size = 1 - mock_ps.dp_rank = 0 - mock_ps.cp_size = 1 - mock_ps.cp_enabled = False - mock_ps_collator.return_value = mock_ps - mock_ps_loader.return_value = mock_ps - builder5 = DataLoaderBuilder( - dataset=VariablePackedDataset(), - micro_batch_size=2, - gradient_accumulation_steps=1, - num_workers=0, - prefetch_factor=None, - pad_to_multiple_of=1, - ) - dl5 = builder5.build(verbose=False) - mbs = next(iter(dl5)) - assert len(mbs) == 1 - assert mbs[0]["input_ids"].shape == (1, 4) - assert mbs[0]["attention_mask"].shape == (1, 4) diff --git a/tests/distillation/test_mooncake_hidden_store.py b/tests/distillation/test_mooncake_hidden_store.py index b1b03def..7ef7ee2d 100644 --- a/tests/distillation/test_mooncake_hidden_store.py +++ b/tests/distillation/test_mooncake_hidden_store.py @@ -15,7 +15,7 @@ import pytest import torch -from tests._helpers.opd import save_tensor_file +from tests._helpers.opd import FakeMooncakeClient, save_tensor_file from xorl.distillation import ( MooncakeHiddenStore, MooncakeStoreConfig, @@ -34,47 +34,22 @@ pytestmark = [pytest.mark.cpu] -class FakeMooncakeClient: - """In-memory stand-in for ``mooncake.store.MooncakeDistributedStore``.""" - - def __init__(self) -> None: - self.objects: dict[str, bytes] = {} - self.put_calls: list[str] = [] - self.get_calls: list[str] = [] - self.removed: list[str] = [] - - def put(self, key: str, value: bytes) -> int: - self.objects[key] = bytes(value) - self.put_calls.append(key) - return 0 - - def get(self, key: str) -> bytes: - self.get_calls.append(key) - return self.objects.get(key, b"") - - def is_exist(self, key: str) -> int: - return 1 if key in self.objects else 0 - - def remove(self, key: str) -> int: - self.objects.pop(key, None) - self.removed.append(key) - return 0 - - def _store() -> tuple[MooncakeHiddenStore, FakeMooncakeClient]: client = FakeMooncakeClient() return MooncakeHiddenStore(client=client, get_retry_max_wait_s=0.0), client -def test_byte_roundtrip_preserves_bfloat16_and_int64(): +def _assert_mooncake_tensor_codec_policy(): for dtype in (torch.bfloat16, torch.float16, torch.float32, torch.int64): tensor = (torch.arange(12).reshape(3, 4) % 5).to(dtype) restored = bytes_to_tensor(tensor_to_bytes(tensor), (3, 4), dtype) assert torch.equal(restored, tensor) assert restored.dtype == dtype + _assert_dtype_string_mapping_is_canonical() + -def test_dtype_string_mapping_is_canonical(): +def _assert_dtype_string_mapping_is_canonical(): assert dtype_to_str(torch.bfloat16) == "bfloat16" assert str_to_dtype("bf16") is torch.bfloat16 assert str_to_dtype("torch.float32") is torch.float32 @@ -82,7 +57,9 @@ def test_dtype_string_mapping_is_canonical(): str_to_dtype("complex128") -def test_put_hidden_returns_metadata_contract(): +def test_mooncake_hidden_transport_policy(): + _assert_mooncake_tensor_codec_policy() + store, client = _store() tensor = torch.randn(7, 5, dtype=torch.bfloat16) @@ -100,8 +77,12 @@ def test_put_hidden_returns_metadata_contract(): assert client.put_calls == ["opd/req-1/teacher/0/hidden/hidden_states"] assert is_mooncake_entry(meta) + _assert_put_get_roundtrip_via_metadata() + _assert_rank3_layer_cache_roundtrip_and_token_count() + _assert_mooncake_teacher_activation_consumer_policy() -def test_put_get_roundtrip_via_metadata(): + +def _assert_put_get_roundtrip_via_metadata(): store, _ = _store() tensor = torch.randn(9, 6, dtype=torch.bfloat16) meta = store.put_hidden("k", tensor) @@ -112,7 +93,7 @@ def test_put_get_roundtrip_via_metadata(): assert torch.equal(fetched, tensor) -def test_rank3_layer_cache_roundtrip_and_token_count(): +def _assert_rank3_layer_cache_roundtrip_and_token_count(): store, _ = _store() # [layers, tokens, hidden] tensor = torch.randn(3, 8, 4, dtype=torch.bfloat16) @@ -124,7 +105,7 @@ def test_rank3_layer_cache_roundtrip_and_token_count(): assert torch.equal(fetched, tensor) -def test_teacher_activation_cache_indexes_mooncake_entry(): +def _assert_mooncake_teacher_activation_consumer_policy(): store, _ = _store() cache_tensor = torch.arange(12, dtype=torch.float32).reshape(6, 2) meta = store.put_hidden("opd/req/teacher/0/hidden", cache_tensor) @@ -139,8 +120,11 @@ def test_teacher_activation_cache_indexes_mooncake_entry(): finally: tac.close() + _assert_teacher_activation_cache_rank3_mooncake_entry() + _assert_multi_teacher_mooncake_caches() + -def test_teacher_activation_cache_rank3_mooncake_entry(): +def _assert_teacher_activation_cache_rank3_mooncake_entry(): store, _ = _store() layer_cache = torch.randn(4, 6, 3, dtype=torch.float32) # [L, tokens, d] meta = store.put_hidden("layer", layer_cache) @@ -157,7 +141,7 @@ def test_teacher_activation_cache_rank3_mooncake_entry(): tac.close() -def test_multi_teacher_mooncake_caches(tmp_path): +def _assert_multi_teacher_mooncake_caches(): store, _ = _store() cache_a = torch.arange(8, dtype=torch.float32).reshape(4, 2) cache_b = torch.arange(100, 110, dtype=torch.float32).reshape(5, 2) @@ -174,7 +158,7 @@ def test_multi_teacher_mooncake_caches(tmp_path): tac.close() -def test_non_mooncake_entry_is_rejected(tmp_path): +def test_mooncake_metadata_admission_and_store_lifecycle_policy(tmp_path, monkeypatch): # The file-backed safetensors cache path was removed; a path/str entry must # now fail loudly rather than silently load a file. store, _ = _store() @@ -186,8 +170,13 @@ def test_non_mooncake_entry_is_rejected(tmp_path): finally: tac.close() + _assert_get_missing_key_raises() + _assert_size_mismatch_raises() + _assert_parse_mooncake_metadata_rejects_malformed() + _assert_mooncake_store_lifecycle_and_configuration_policy(monkeypatch) + -def test_get_missing_key_raises(): +def _assert_get_missing_key_raises(): store, _ = _store() meta = { "backend": "mooncake", @@ -200,7 +189,7 @@ def test_get_missing_key_raises(): store.get_hidden_from_metadata(meta) -def test_size_mismatch_raises(): +def _assert_size_mismatch_raises(): store, client = _store() tensor = torch.randn(4, 2, dtype=torch.float32) meta = store.put_hidden("k", tensor) @@ -210,37 +199,37 @@ def test_size_mismatch_raises(): store.get_hidden_from_metadata(meta) -@pytest.mark.parametrize( - "mutate", - [ +def _assert_parse_mooncake_metadata_rejects_malformed(): + mutations = ( lambda m: m.pop("key"), lambda m: m.pop("tensor_shapes"), lambda m: m.pop("tensor_dtypes"), lambda m: m["tensor_shapes"].__setitem__("hidden_states", [1, 2, 3, 4]), - ], -) -def test_parse_mooncake_metadata_rejects_malformed(mutate): - meta = { - "backend": "mooncake", - "key": "k", - "tensor_key": "hidden_states", - "tensor_shapes": {"hidden_states": [3, 2]}, - "tensor_dtypes": {"hidden_states": "float32"}, - } - mutate(meta) - with pytest.raises(ValueError): - parse_mooncake_metadata(meta) - - -def test_remove_hidden_is_best_effort(): + ) + for mutate in mutations: + meta = { + "backend": "mooncake", + "key": "k", + "tensor_key": "hidden_states", + "tensor_shapes": {"hidden_states": [3, 2]}, + "tensor_dtypes": {"hidden_states": "float32"}, + } + mutate(meta) + with pytest.raises(ValueError): + parse_mooncake_metadata(meta) + + +def _assert_mooncake_store_lifecycle_and_configuration_policy(monkeypatch): store, client = _store() tensor = torch.randn(2, 2) meta = store.put_hidden("k", tensor) store.remove_hidden(meta["key"]) assert client.removed == ["k/hidden_states"] + _assert_store_config_overrides_win_over_env(monkeypatch) + -def test_store_config_overrides_win_over_env(monkeypatch): +def _assert_store_config_overrides_win_over_env(monkeypatch): monkeypatch.setenv("MOONCAKE_MASTER_SERVER", "envhost:50051") cfg = MooncakeStoreConfig.from_env(overrides={"master_server_address": "pinned:9999"}) assert cfg.master_server_address == "pinned:9999" diff --git a/tests/distributed/test_bf16_a2a_fsdp_hook.py b/tests/distributed/test_bf16_a2a_fsdp_hook.py index 187d70f4..51c00d4d 100644 --- a/tests/distributed/test_bf16_a2a_fsdp_hook.py +++ b/tests/distributed/test_bf16_a2a_fsdp_hook.py @@ -26,7 +26,7 @@ if str(THIS_DIR) not in sys.path: sys.path.insert(0, str(THIS_DIR)) -from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than +from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than # noqa: E402 pytestmark = [pytest.mark.distributed] diff --git a/tests/distributed/test_bf16_a2a_reduce.py b/tests/distributed/test_bf16_a2a_reduce.py index 7318c557..63b8781b 100644 --- a/tests/distributed/test_bf16_a2a_reduce.py +++ b/tests/distributed/test_bf16_a2a_reduce.py @@ -1,57 +1,61 @@ -"""Distributed correctness tests for ``BF16StochasticAllToAllReduceScatter``. +"""Default-runtime contract for ``BF16StochasticAllToAllReduceScatter``. -Verifies the custom reduce-scatter (stochastic-round FP32→BF16, all-to-all, -local FP32 sum) produces results numerically close to native FP32 -reduce-scatter, with bias-in-expectation near zero. +One report covers the stochastic FP32-to-BF16 primitive and the real two-rank +Gloo all-to-all/FP32 accumulation transaction. This avoids treating a four-GPU +admission gate and an isolated rounding unit as independent confidence. """ from __future__ import annotations -import os import sys from pathlib import Path import pytest import torch import torch.distributed as dist -from torch.distributed.distributed_c10d import ReduceOp from xorl.distributed.fsdp2 import BF16StochasticAllToAllReduceScatter -from xorl.distributed.fsdp2.bf16_a2a_reduce import _canonical_reduce_op -from xorl.utils.device import get_nccl_backend +from xorl.optim.stochastic_round import stochastic_round_to_bf16 THIS_DIR = Path(__file__).resolve().parent if str(THIS_DIR) not in sys.path: sys.path.insert(0, str(THIS_DIR)) -from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than +from distributed_utils import run_distributed_script # noqa: E402 -pytestmark = [pytest.mark.distributed] +pytestmark = [pytest.mark.cpu, pytest.mark.distributed] -@pytest.mark.cpu -def test_canonical_reduce_op_accepts_fsdp_wrapped_ops(): - assert _canonical_reduce_op(ReduceOp(ReduceOp.SUM)) == dist.ReduceOp.SUM - assert _canonical_reduce_op(ReduceOp(ReduceOp.AVG)) == dist.ReduceOp.AVG - assert _canonical_reduce_op(dist.ReduceOp.SUM) == dist.ReduceOp.SUM - assert _canonical_reduce_op(dist._make_nccl_premul_sum(1.0)) == dist.ReduceOp.SUM - +def _setup_dist() -> torch.device: + dist.init_process_group(backend="gloo") + return torch.device("cpu") -def _world_size() -> int: - return int(os.environ["WORLD_SIZE"]) +def _assert_stochastic_round_distribution_and_admission() -> None: + values = torch.randn(7, 13, dtype=torch.float32) + generator_one = torch.Generator().manual_seed(42) + generator_two = torch.Generator().manual_seed(42) + rounded = stochastic_round_to_bf16(values, generator=generator_one) -def _local_rank() -> int: - return int(os.environ["LOCAL_RANK"]) + assert rounded.dtype is torch.bfloat16 + assert rounded.shape == values.shape + assert torch.equal(rounded, stochastic_round_to_bf16(values, generator=generator_two)) + with pytest.raises(ValueError, match="requires fp32 input"): + stochastic_round_to_bf16(values.to(torch.bfloat16)) + sample_count = 1 << 16 + lower_bits = 0x3F800000 + fractional_bits = 0x4000 + samples = torch.full((sample_count,), lower_bits + fractional_bits, dtype=torch.int32).view(torch.float32) + rounded_samples = stochastic_round_to_bf16(samples, generator=torch.Generator().manual_seed(0)).float() + lower = torch.tensor(lower_bits, dtype=torch.int32).view(torch.float32) + upper = torch.tensor(lower_bits + 0x10000, dtype=torch.int32).view(torch.float32) -def _setup_dist() -> torch.device: - local_rank = _local_rank() - torch.cuda.set_device(local_rank) - dist.init_process_group(backend=get_nccl_backend()) - return torch.device("cuda", local_rank) + assert ((rounded_samples == lower) | (rounded_samples == upper)).all() + assert (rounded_samples == upper).float().mean().item() == pytest.approx(0.25, abs=0.01) + assert rounded_samples.mean().item() == pytest.approx(samples[0].item(), abs=1e-4) def _run() -> None: @@ -70,16 +74,15 @@ def _run() -> None: # Per-rank gradient (independent across ranks). local_grad = torch.randn(total_numel, dtype=torch.float32, device=device) - # ---- Reference: native FP32 reduce-scatter ---- + # Reference: native FP32 reduce-scatter. ref_out = torch.empty(chunk_numel, dtype=torch.float32, device=device) dist.reduce_scatter_tensor(ref_out, local_grad.clone(), op=dist.ReduceOp.SUM) - # ---- Test: BF16 stochastic-rounded a2a + FP32 local sum ---- + # Test: BF16 stochastic-rounded a2a + FP32 local sum. comm = BF16StochasticAllToAllReduceScatter() test_out = comm.allocate((chunk_numel,), dtype=torch.float32, device=device) comm(test_out, local_grad.clone(), group=dist.group.WORLD, op=dist.ReduceOp.SUM) - # ---- Bound the per-element error ---- # Each rank's contribution is stochastically rounded FP32→BF16 with at most # one ulp of noise. After summing ``world`` such contributions, the error # is bounded by sum of |x_r| * 2^-7 in the worst case. Compute this bound. @@ -99,27 +102,8 @@ def _run() -> None: max_bound = bound_for_my_chunk.max().item() * 4 + 1e-6 assert max_err < max_bound, f"[rank {rank}] BF16 a2a max err {max_err:.4e} exceeds bound {max_bound:.4e}" - # Bias-in-expectation: average over many trials should approach the FP32 reference. - # Use the same input tensor; only the stochastic rounding noise differs. - n_trials = 200 - accum = torch.zeros_like(test_out) - for _ in range(n_trials): - out = comm.allocate((chunk_numel,), dtype=torch.float32, device=device) - comm(out, local_grad.clone(), group=dist.group.WORLD, op=dist.ReduceOp.SUM) - accum += out - mean = accum / n_trials - mean_err = (mean - ref_out).abs().max().item() - # Standard error of the mean ~ bound / sqrt(n_trials). For n=200 and BF16 - # bound ~|x|/128, SEM ~ |x| * 1e-3. Allow generous 5x headroom. - sem_bound = max_bound / (n_trials**0.5) * 5 - assert mean_err < sem_bound, ( - f"[rank {rank}] BF16 a2a is biased: mean err over {n_trials} trials = " - f"{mean_err:.4e}, expected < {sem_bound:.4e}" - ) - if rank == 0: print(f"[rank 0] BF16 a2a max err = {max_err:.4e}, bound = {max_bound:.4e}") - print(f"[rank 0] BF16 a2a unbiased mean err over {n_trials} trials = {mean_err:.4e}") dist.barrier() dist.destroy_process_group() @@ -131,10 +115,15 @@ def _main() -> None: if __name__ != "__main__": - @skip_if_gpu_count_less_than(4) - def test_bf16_a2a_reduce_scatter_matches_fp32_within_bound(): - result = run_distributed_script(__file__, num_gpus=4, timeout=180) - result.assert_success("BF16 a2a reduce-scatter should match FP32 within BF16 ulp bound") + def test_bf16_stochastic_a2a_reduce_scatter_default_runtime_contract(): + _assert_stochastic_round_distribution_and_admission() + result = run_distributed_script( + __file__, + num_gpus=2, + timeout=180, + extra_env={"CUDA_VISIBLE_DEVICES": ""}, + ) + result.assert_success("BF16 a2a reduce-scatter should match FP32 on two CPU/Gloo ranks") if __name__ == "__main__": diff --git a/tests/distributed/test_canonical_moe_contract.py b/tests/distributed/test_canonical_moe_contract.py index 998bf2c6..adedf3c2 100644 --- a/tests/distributed/test_canonical_moe_contract.py +++ b/tests/distributed/test_canonical_moe_contract.py @@ -14,10 +14,8 @@ LocalMoEContribution, OutputDistribution, ParallelPlan, - ParallelRole, canonical_moe_reduce_cp_sharded_v3, canonical_moe_reduce_packed_ep16_v2, - canonical_moe_reduce_reference, canonical_moe_reduce_v1, resolve_canonical_moe_transport, ) @@ -27,42 +25,18 @@ pytestmark = [pytest.mark.distributed] -def _explicit_tree(partials: torch.Tensor) -> torch.Tensor: - current = [partials[index] for index in range(partials.shape[0])] - while len(current) > 1: - current = [(current[index] + current[index + 1]).bfloat16() for index in range(0, len(current), 2)] - return current[0] +def _canonical_moe_reference(partials: torch.Tensor, metadata: CanonicalMoEGraphMetadata) -> torch.Tensor: + level = [partials[index] for index in range(partials.shape[0])] + while len(level) > 1: + level = [(level[index] + level[index + 1]).to(torch.bfloat16) for index in range(0, len(level), 2)] + result = level[0] + result = result.clone() + result[~metadata.valid_mask] = 0 + return result @pytest.mark.cpu -@pytest.mark.parametrize("contributors", [2, 4, 8, 16]) -def test_reference_is_the_adjacent_bf16_tree(contributors: int): - rows = contributors + 2 - values = torch.zeros((contributors, rows, 3), dtype=torch.bfloat16) - adversarial = torch.tensor( - [4096.0, -4096.0, 1.0, 1.0, 0.5, -0.5, 2.0, -2.0] * 2, - dtype=torch.bfloat16, - ) - for ordinal in range(contributors): - values[ordinal, :, 0] = adversarial[ordinal] - values[ordinal, :, 1] = ordinal + 1 - values[ordinal, :, 2] = torch.arange(rows) - metadata = CanonicalMoEGraphMetadata.build( - torch.arange(rows), - torch.arange(rows), - capacity=rows + 3, - ) - padded = torch.zeros((contributors, metadata.capacity, 3), dtype=torch.bfloat16) - padded[:, :rows] = values - - result = canonical_moe_reduce_reference(padded, metadata) - expected = _explicit_tree(padded) - expected[~metadata.valid_mask] = 0 - assert torch.equal(result, expected) - - -@pytest.mark.cpu -def test_graph_metadata_has_deterministic_padding_and_capacity_guard(): +def test_graph_metadata_transport_and_parallel_plan_policy(): metadata = CanonicalMoEGraphMetadata.build( torch.tensor([7, 3, 11], dtype=torch.int64), torch.tensor([17, 2, 25], dtype=torch.int64), @@ -74,9 +48,11 @@ def test_graph_metadata_has_deterministic_padding_and_capacity_guard(): with pytest.raises(ValueError, match="exceeds fixed capacity"): CanonicalMoEGraphMetadata.build(torch.arange(3), torch.arange(3), capacity=2) + _assert_transport_admission_policy() + _assert_exact_trainer_plan_and_fail_closed() -@pytest.mark.cpu -def test_transport_auto_promotes_only_admitted_cp_sharded_geometry(): + +def _assert_transport_admission_policy(): ep16 = ParallelPlan.glm52_trainer(world_size=16, pp_size=1, dp_size=1, contributor_count=16) assert ( resolve_canonical_moe_transport( @@ -114,42 +90,6 @@ def test_transport_auto_promotes_only_admitted_cp_sharded_geometry(): is CanonicalMoETransport.DENSE_V1 ) - -@pytest.mark.cpu -def test_internal_resolution_serves_cp_sharded_v3_on_the_admitted_geometry(): - """The exact GLM path has no public transport knob: internal resolution - serves cp_sharded_v3 on the admitted eager EP16/CP16 consumer-sharded - geometry (the transport with the strongest certified performance - evidence) and the dense executable oracle elsewhere.""" - ep16 = ParallelPlan.glm52_trainer(world_size=16, pp_size=1, dp_size=1, contributor_count=16) - assert ( - resolve_canonical_moe_transport( - "auto", - plan=ep16, - capacity=4224, - local_rows=264, - graph_mode=False, - consumer_sharded_output=True, - ) - is CanonicalMoETransport.CP_SHARDED_V3 - ) - ep8 = ParallelPlan.primitive(8) - assert ( - resolve_canonical_moe_transport( - "auto", - plan=ep8, - capacity=4224, - local_rows=528, - graph_mode=False, - consumer_sharded_output=True, - ) - is CanonicalMoETransport.DENSE_V1 - ) - - -@pytest.mark.cpu -def test_transport_explicit_modes_never_silently_fallback(): - ep16 = ParallelPlan.glm52_trainer(world_size=16, pp_size=1, dp_size=1, contributor_count=16) assert ( resolve_canonical_moe_transport( "dense_v1", @@ -189,9 +129,6 @@ def test_transport_explicit_modes_never_silently_fallback(): consumer_sharded_output=False, ) - -@pytest.mark.cpu -def test_packed_ep16_v2_fails_closed_outside_admitted_mode(): metadata = CanonicalMoEGraphMetadata.build(torch.arange(16), torch.arange(16), capacity=16) contribution = LocalMoEContribution( torch.zeros((16, 2), dtype=torch.bfloat16), @@ -222,15 +159,11 @@ def test_packed_ep16_v2_fails_closed_outside_admitted_mode(): ) -@pytest.mark.cpu -def test_exact_parallel_plans_hash_launcher_spelling_and_fail_closed(): +def _assert_exact_trainer_plan_and_fail_closed(): trainer = ParallelPlan.glm52_trainer() - sampler = ParallelPlan.glm52_sampler(launcher_tp_size=8) - assert trainer.digest != sampler.digest - assert sampler.as_dict()["launcher_tp_size"] == 8 assert trainer.pipeline_layer_ranges == ((0, 78),) assert trainer.combine_groups == (tuple(range(16)),) - assert all(trainer.logical_ordinal(physical_rank) == physical_rank for physical_rank in range(16)) + assert trainer.logical_ordinals_by_group == (tuple(range(16)),) assert trainer.contract_version == CANONICAL_MOE_REDUCE_VERSION for kwargs in ( @@ -241,53 +174,35 @@ def test_exact_parallel_plans_hash_launcher_spelling_and_fail_closed(): with pytest.raises(ValueError, match="Unsupported GLM-5.2 trainer topology"): ParallelPlan.glm52_trainer(**kwargs) - payload = sampler.as_dict() - payload["world_size"] = 16 - payload["role"] = sampler.role - with pytest.raises(ValueError, match="partition|world|sampler topology"): - ParallelPlan(**payload) - - payload = sampler.as_dict() - payload["role"] = ParallelRole.SAMPLER - payload["launcher_tp_size"] = 99 - with pytest.raises(ValueError, match="launcher-level tp_size must be exactly 8"): - ParallelPlan(**payload) - - payload = sampler.as_dict() - payload["role"] = ParallelRole.SAMPLER - payload["logical_ordinals_by_group"] = (tuple(reversed(range(8))),) - with pytest.raises(ValueError, match="identity logical contributor ordinals"): - ParallelPlan(**payload) - - -@pytest.mark.cpu -def test_world32_pp1_cp_and_ep_groups_alias_as_four_rank_octets(): - main_mesh = torch.arange(32).view(4, 8) - ep_mesh = init_ep_mesh_matrix(ep_size=8, ep_fsdp_size=4, ep_intranode=True) - - cp_groups = tuple(tuple(int(rank) for rank in row) for row in main_mesh) - ep_groups = tuple(tuple(int(rank) for rank in ep_mesh[:, column]) for column in range(4)) - expert_fsdp_groups = tuple(tuple(int(rank) for rank in row) for row in ep_mesh) + _assert_world32_cp_and_ep_group_alias_policy() - expected_octets = tuple(tuple(range(start, start + 8)) for start in range(0, 32, 8)) - assert cp_groups == ep_groups == expected_octets - assert expert_fsdp_groups[0] == (0, 8, 16, 24) - assert expert_fsdp_groups[-1] == (7, 15, 23, 31) +def _assert_world32_cp_and_ep_group_alias_policy(): + cases = ( + ( + 8, + 4, + tuple(tuple(range(start, start + 8)) for start in range(0, 32, 8)), + (0, 8, 16, 24), + (7, 15, 23, 31), + ), + (16, 2, (tuple(range(16)), tuple(range(16, 32))), (0, 16), (15, 31)), + ) + for ep_size, ep_fsdp_size, expected_groups, first_fsdp_group, last_fsdp_group in cases: + main_mesh = torch.arange(32).view(ep_fsdp_size, ep_size) + ep_mesh = init_ep_mesh_matrix( + ep_size=ep_size, + ep_fsdp_size=ep_fsdp_size, + ep_intranode=True, + ) -@pytest.mark.cpu -def test_world32_pp1_cp16_and_ep16_groups_alias_as_two_rank_groups(): - main_mesh = torch.arange(32).view(2, 16) - ep_mesh = init_ep_mesh_matrix(ep_size=16, ep_fsdp_size=2, ep_intranode=True) - - cp_groups = tuple(tuple(int(rank) for rank in row) for row in main_mesh) - ep_groups = tuple(tuple(int(rank) for rank in ep_mesh[:, column]) for column in range(2)) - expert_fsdp_groups = tuple(tuple(int(rank) for rank in row) for row in ep_mesh) + cp_groups = tuple(tuple(int(rank) for rank in row) for row in main_mesh) + ep_groups = tuple(tuple(int(rank) for rank in ep_mesh[:, column]) for column in range(ep_fsdp_size)) + expert_fsdp_groups = tuple(tuple(int(rank) for rank in row) for row in ep_mesh) - expected_groups = (tuple(range(16)), tuple(range(16, 32))) - assert cp_groups == ep_groups == expected_groups - assert expert_fsdp_groups[0] == (0, 16) - assert expert_fsdp_groups[-1] == (15, 31) + assert cp_groups == ep_groups == expected_groups + assert expert_fsdp_groups[0] == first_fsdp_group + assert expert_fsdp_groups[-1] == last_fsdp_group def _make_partials(world: int, capacity: int) -> torch.Tensor: @@ -327,7 +242,7 @@ def _run_distributed_case() -> None: physical_stack = torch.stack(gathered) logical_to_physical = [mapping.index(logical) for logical in range(world)] logical_stack = physical_stack[logical_to_physical] - expected = canonical_moe_reduce_reference(logical_stack, metadata) + expected = _canonical_moe_reference(logical_stack, metadata) replicated = canonical_moe_reduce_v1( contribution, @@ -615,13 +530,14 @@ def _run_packed_ep16_case() -> None: if __name__ != "__main__": @pytest.mark.cpu - @pytest.mark.parametrize("contributors", [2, 4, 8]) - def test_distributed_transport_tree_distribution_chunking_and_backward(contributors: int): - result = run_distributed_script(__file__, num_gpus=contributors, timeout=180) - result.assert_success(f"canonical MoE primitive with {contributors} CPU contributors") + def test_distributed_transport_tree_distribution_chunking_and_backward(): + for contributors in (2, 8): + result = run_distributed_script(__file__, num_gpus=contributors, timeout=180) + result.assert_success(f"canonical MoE primitive with {contributors} CPU contributors") - @pytest.mark.cpu - def test_packed_ep16_v2_matches_dense_v1_bitwise(): + _assert_packed_ep16_v2_matches_dense_v1_bitwise() + + def _assert_packed_ep16_v2_matches_dense_v1_bitwise(): result = run_distributed_script( __file__, num_gpus=16, diff --git a/tests/distributed/test_deepep_async_combine_guard.py b/tests/distributed/test_deepep_async_combine_guard.py deleted file mode 100644 index fd809837..00000000 --- a/tests/distributed/test_deepep_async_combine_guard.py +++ /dev/null @@ -1,55 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from xorl.distributed.moe import deepep - - -pytestmark = pytest.mark.cpu - - -def test_deepep_async_combine_is_synchronous_by_default(monkeypatch): - captured = {} - - def fake_apply(expert_output, buffer, ctx, async_combine): - del buffer, ctx - captured["async_combine"] = async_combine - return expert_output - - monkeypatch.setattr(deepep, "_ALLOW_UNSAFE_ASYNC_COMBINE", False) - monkeypatch.setattr(deepep._FusedUnpermuteAndCombine, "apply", staticmethod(fake_apply)) - - expert_output = torch.ones(1, 2) - result = deepep.tokens_post_combine( - buffer=None, - expert_output=expert_output, - ctx=SimpleNamespace(), - async_combine=True, - ) - - assert result is expert_output - assert captured["async_combine"] is False - - -def test_deepep_async_combine_can_be_unsafely_opted_in(monkeypatch): - captured = {} - - def fake_apply(expert_output, buffer, ctx, async_combine): - del buffer, ctx - captured["async_combine"] = async_combine - return expert_output - - monkeypatch.setattr(deepep, "_ALLOW_UNSAFE_ASYNC_COMBINE", True) - monkeypatch.setattr(deepep._FusedUnpermuteAndCombine, "apply", staticmethod(fake_apply)) - - expert_output = torch.ones(1, 2) - result = deepep.tokens_post_combine( - buffer=None, - expert_output=expert_output, - ctx=SimpleNamespace(), - async_combine=True, - ) - - assert result is expert_output - assert captured["async_combine"] is True diff --git a/tests/distributed/test_deepep_internode_guard.py b/tests/distributed/test_deepep_internode_guard.py index ee83e3fa..7a6d3ae7 100644 --- a/tests/distributed/test_deepep_internode_guard.py +++ b/tests/distributed/test_deepep_internode_guard.py @@ -15,6 +15,7 @@ from types import SimpleNamespace import pytest +import torch from xorl.distributed.moe import deepep @@ -26,52 +27,70 @@ class _FakeGroup: """Hashable stand-in for a ProcessGroup.""" +def _assert_async_combine_requires_explicit_unsafe_opt_in(monkeypatch): + captured = {} + + def fake_apply(expert_output, buffer, ctx, async_combine): + del buffer, ctx + captured["async_combine"] = async_combine + return expert_output + + monkeypatch.delenv("XORL_DEEPEP_UNSAFE_ASYNC_COMBINE", raising=False) + monkeypatch.setattr(deepep._FusedUnpermuteAndCombine, "apply", staticmethod(fake_apply)) + + expert_output = torch.ones(1, 2) + result = deepep.tokens_post_combine( + buffer=None, + expert_output=expert_output, + ctx=SimpleNamespace(), + async_combine=True, + ) + + assert result is expert_output + assert captured["async_combine"] is False + + monkeypatch.setenv("XORL_DEEPEP_UNSAFE_ASYNC_COMBINE", "1") + result = deepep.tokens_post_combine( + buffer=None, + expert_output=expert_output, + ctx=SimpleNamespace(), + async_combine=True, + ) + + assert result is expert_output + assert captured["async_combine"] is True + + # --------------------------------------------------------------------------- # 1. _ep_group_spans_nodes # --------------------------------------------------------------------------- class TestEPGroupSpansNodes: - def test_unknown_when_dist_not_initialized(self, monkeypatch): + def _assert_topology_truth_table(self, monkeypatch): monkeypatch.setattr(deepep.dist, "is_initialized", lambda: False) assert deepep._ep_group_spans_nodes(None) is None - def test_unknown_without_local_world_size(self, monkeypatch): monkeypatch.setattr(deepep.dist, "is_initialized", lambda: True) monkeypatch.delenv("LOCAL_WORLD_SIZE", raising=False) assert deepep._ep_group_spans_nodes(None) is None - def test_unknown_with_malformed_local_world_size(self, monkeypatch): - monkeypatch.setattr(deepep.dist, "is_initialized", lambda: True) monkeypatch.setenv("LOCAL_WORLD_SIZE", "not-a-number") assert deepep._ep_group_spans_nodes(None) is None - def test_single_node_world_never_spans(self, monkeypatch): - monkeypatch.setattr(deepep.dist, "is_initialized", lambda: True) monkeypatch.setattr(deepep.dist, "get_world_size", lambda: 8) monkeypatch.setenv("LOCAL_WORLD_SIZE", "8") assert deepep._ep_group_spans_nodes(None) is False - def test_intranode_ep_group_in_multinode_world(self, monkeypatch): - monkeypatch.setattr(deepep.dist, "is_initialized", lambda: True) monkeypatch.setattr(deepep.dist, "get_world_size", lambda: 16) monkeypatch.setattr(deepep.dist, "get_process_group_ranks", lambda group: list(range(8))) - monkeypatch.setenv("LOCAL_WORLD_SIZE", "8") assert deepep._ep_group_spans_nodes(_FakeGroup()) is False - def test_internode_ep_group(self, monkeypatch): - monkeypatch.setattr(deepep.dist, "is_initialized", lambda: True) - monkeypatch.setattr(deepep.dist, "get_world_size", lambda: 16) monkeypatch.setattr(deepep.dist, "get_process_group_ranks", lambda group: list(range(16))) - monkeypatch.setenv("LOCAL_WORLD_SIZE", "8") assert deepep._ep_group_spans_nodes(_FakeGroup()) is True - def test_strided_internode_ep_group(self, monkeypatch): - """ep_intranode=False layouts stride EP groups across nodes.""" - monkeypatch.setattr(deepep.dist, "is_initialized", lambda: True) - monkeypatch.setattr(deepep.dist, "get_world_size", lambda: 16) + # ep_intranode=False layouts stride EP groups across nodes. monkeypatch.setattr(deepep.dist, "get_process_group_ranks", lambda group: [0, 8]) - monkeypatch.setenv("LOCAL_WORLD_SIZE", "8") assert deepep._ep_group_spans_nodes(_FakeGroup()) is True @@ -93,7 +112,11 @@ def fake_gather(out, obj, group=None): class TestPreflightInternodeTransport: - def test_skipped_via_env(self, monkeypatch): + def _assert_preflight_internode_transport_policy(self, monkeypatch): + def _no_buffer(**kwargs): + raise AssertionError("no buffer should be created when the preflight is disabled") + + monkeypatch.setattr(deepep, "get_default_buffer", _no_buffer) monkeypatch.setenv(deepep._SKIP_PREFLIGHT_ENV, "1") def _fail(group): @@ -102,21 +125,18 @@ def _fail(group): monkeypatch.setattr(deepep, "_ep_group_spans_nodes", _fail) deepep.preflight_internode_transport(None, hidden_dim=128) - def test_noop_when_dist_not_initialized(self, monkeypatch): + monkeypatch.delenv(deepep._SKIP_PREFLIGHT_ENV) monkeypatch.setattr(deepep.dist, "is_initialized", lambda: False) deepep.preflight_internode_transport(None, hidden_dim=128) - def test_noop_when_intranode(self, monkeypatch): monkeypatch.setattr(deepep.dist, "is_initialized", lambda: True) monkeypatch.setattr(deepep, "_ep_group_spans_nodes", lambda group: False) - - def _no_buffer(**kwargs): - raise AssertionError("no buffer should be created for intranode EP groups") - - monkeypatch.setattr(deepep, "get_default_buffer", _no_buffer) deepep.preflight_internode_transport(None, hidden_dim=128) - def test_transport_failure_names_nodes(self, monkeypatch): + self._assert_transport_failure_names_nodes(monkeypatch) + self._assert_roundtrip_accepts_identity_and_rejects_corruption(monkeypatch) + + def _assert_transport_failure_names_nodes(self, monkeypatch): _arm_internode(monkeypatch) buf = SimpleNamespace(init_buffer=lambda hidden_bytes: None) monkeypatch.setattr(deepep, "get_default_buffer", lambda **kwargs: buf) @@ -132,7 +152,7 @@ def dead_dispatch(*args, **kwargs): assert "CPU recv timeout" in msg assert deepep._SKIP_PREFLIGHT_ENV in msg - def test_roundtrip_corruption_detected(self, monkeypatch): + def _assert_roundtrip_accepts_identity_and_rejects_corruption(self, monkeypatch): _arm_internode(monkeypatch) buf = SimpleNamespace(init_buffer=lambda hidden_bytes: None) monkeypatch.setattr(deepep, "get_default_buffer", lambda **kwargs: buf) @@ -141,22 +161,13 @@ def identity_dispatch(buffer, x, w, idx, num_experts): return x, idx, w, [1] * num_experts, "handle" monkeypatch.setattr(deepep, "dispatch_no_grad", identity_dispatch) + monkeypatch.setattr(deepep, "combine_no_grad", lambda buffer, x, handle: x.clone()) + deepep.preflight_internode_transport(_FakeGroup(), hidden_dim=128) + monkeypatch.setattr(deepep, "combine_no_grad", lambda buffer, x, handle: x + 1.0) with pytest.raises(RuntimeError, match="corrupted"): deepep.preflight_internode_transport(_FakeGroup(), hidden_dim=128) - def test_healthy_roundtrip_passes(self, monkeypatch): - _arm_internode(monkeypatch) - buf = SimpleNamespace(init_buffer=lambda hidden_bytes: None) - monkeypatch.setattr(deepep, "get_default_buffer", lambda **kwargs: buf) - - def identity_dispatch(buffer, x, w, idx, num_experts): - return x, idx, w, [1] * num_experts, "handle" - - monkeypatch.setattr(deepep, "dispatch_no_grad", identity_dispatch) - monkeypatch.setattr(deepep, "combine_no_grad", lambda buffer, x, handle: x.clone()) - deepep.preflight_internode_transport(_FakeGroup(), hidden_dim=128) - # --------------------------------------------------------------------------- # 3. DeepEPBuffer.init_buffer size validation @@ -198,24 +209,29 @@ def __init__(self, **kwargs): buf = deepep.DeepEPBuffer(ep_group=group, buffer_size_gb=buffer_size_gb) return buf, captured - def test_oversized_nvl_with_rdma_raises(self, monkeypatch): + def _assert_buffer_size_alignment_and_rdma_admission_policy(self, monkeypatch): buf, _ = self._buffer(monkeypatch, buffer_size_gb=4.0, rdma_bytes=1 << 20) with pytest.raises(ValueError, match="int32 limit"): buf.init_buffer(hidden_bytes=4096) - def test_oversized_nvl_without_rdma_allowed(self, monkeypatch): - buf, captured = self._buffer(monkeypatch, buffer_size_gb=4.0, rdma_bytes=0) - buf.init_buffer(hidden_bytes=4096) - assert captured["num_nvl_bytes"] == 4_000_000_000 - - def test_unaligned_bytes_rounded_down(self, monkeypatch): - # 0.000001 GB = 1000 bytes -> 896 after 128-byte alignment - buf, captured = self._buffer(monkeypatch, buffer_size_gb=1e-6, rdma_bytes=0) - buf.init_buffer(hidden_bytes=4096) - assert captured["num_nvl_bytes"] == 896 - - def test_default_two_gb_with_rdma_passes(self, monkeypatch): - buf, captured = self._buffer(monkeypatch, buffer_size_gb=2.0, rdma_bytes=1 << 20) - buf.init_buffer(hidden_bytes=4096) - assert captured["num_nvl_bytes"] == 2_000_000_000 - assert captured["num_rdma_bytes"] == 1 << 20 + cases = ( + (4.0, 0, 4_000_000_000), + (1e-6, 0, 896), + (2.0, 1 << 20, 2_000_000_000), + ) + for buffer_size_gb, rdma_bytes, expected_nvl_bytes in cases: + buf, captured = self._buffer(monkeypatch, buffer_size_gb=buffer_size_gb, rdma_bytes=rdma_bytes) + buf.init_buffer(hidden_bytes=4096) + assert captured["num_nvl_bytes"] == expected_nvl_bytes + assert captured["num_rdma_bytes"] == rdma_bytes + + +def test_deepep_internode_topology_preflight_and_buffer_admission_policy(monkeypatch): + with monkeypatch.context() as topology_patch: + TestEPGroupSpansNodes()._assert_topology_truth_table(topology_patch) + with monkeypatch.context() as preflight_patch: + TestPreflightInternodeTransport()._assert_preflight_internode_transport_policy(preflight_patch) + with monkeypatch.context() as buffer_patch: + TestBufferSizeValidation()._assert_buffer_size_alignment_and_rdma_admission_policy(buffer_patch) + with monkeypatch.context() as async_patch: + _assert_async_combine_requires_explicit_unsafe_opt_in(async_patch) diff --git a/tests/distributed/test_ep_clip_grad_norm.py b/tests/distributed/test_ep_clip_grad_norm.py index 92736f88..b13e91c9 100644 --- a/tests/distributed/test_ep_clip_grad_norm.py +++ b/tests/distributed/test_ep_clip_grad_norm.py @@ -24,10 +24,7 @@ from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import Replicate, Shard, distribute_tensor -from xorl.distributed.ep_gradients import ( - register_ep_replicated_gradient_hooks, - synchronize_ep_replicated_gradients, -) +from xorl.distributed.ep_gradients import synchronize_ep_replicated_gradients from xorl.distributed.fsdp2.clip_grad_norm import ( clip_grad_norm, ep_fsdp2_clip_grad_norm, @@ -63,15 +60,6 @@ def _mock_parallel_state(ep_enabled=True): return ps -def _l2_norm(*params): - """Compute the expected L2 norm across multiple params' gradients.""" - total = 0.0 - for p in params: - if p.grad is not None: - total += p.grad.detach().to(torch.float32).norm(2).item() ** 2 - return math.sqrt(total) - - # --------------------------------------------------------------------------- # 1. _build_ep_param_groups classification # --------------------------------------------------------------------------- @@ -80,67 +68,7 @@ def _l2_norm(*params): class TestBuildEPParamGroups: """Test that _build_ep_param_groups classifies params correctly.""" - def test_skip_fsdp_params_classified_as_ep(self): - """Params from _skip_fsdp modules go into the EP group.""" - model = nn.Module() - # Regular submodule - regular = nn.Linear(8, 8) - model.add_module("regular", regular) - # _skip_fsdp submodule (e.g. QLoRAMoeExperts) - expert = nn.Linear(8, 8) - expert._skip_fsdp = True - model.add_module("expert", expert) - - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): - _build_ep_param_groups(model) - - assert hasattr(model, "_ep_param_groups") - ep_ids = {id(p) for p in model._ep_param_groups["ep"]} - non_ep_ids = {id(p) for p in model._ep_param_groups["non_ep"]} - - # Expert params in EP group - for p in expert.parameters(): - assert id(p) in ep_ids - # Regular params in non-EP group - for p in regular.parameters(): - assert id(p) in non_ep_ids - - def test_no_skip_fsdp_all_non_ep(self): - """Without _skip_fsdp modules, all plain params go to non-EP.""" - model = nn.Module() - model.add_module("linear", nn.Linear(8, 8)) - - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): - _build_ep_param_groups(model) - - assert len(model._ep_param_groups["ep"]) == 0 - assert len(model._ep_param_groups["non_ep"]) == 2 # weight + bias - - def test_nested_skip_fsdp_params_all_classified(self): - """All params inside nested _skip_fsdp modules are classified as EP.""" - model = nn.Module() - # Mimic QLoRAMoeExperts: a _skip_fsdp module with multiple sub-params - expert = nn.Module() - expert._skip_fsdp = True - expert.lora_A = nn.Parameter(torch.randn(4, 32, 8)) - expert.lora_B = nn.Parameter(torch.randn(4, 8, 64)) - expert.base_weight = nn.Parameter(torch.randn(4, 32, 64), requires_grad=False) - model.add_module("expert", expert) - model.add_module("non_expert", nn.Linear(32, 32)) - - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): - _build_ep_param_groups(model) - - ep_ids = {id(p) for p in model._ep_param_groups["ep"]} - assert id(expert.lora_A) in ep_ids - assert id(expert.lora_B) in ep_ids - assert id(expert.base_weight) in ep_ids # even frozen params are classified - # non-expert params not in EP group - non_ep_ids = {id(p) for p in model._ep_param_groups["non_ep"]} - for p in model.non_expert.parameters(): - assert id(p) in non_ep_ids - - def test_shared_ep_replica_is_recorded_separately_for_clipping(self): + def _assert_shared_ep_replica_is_recorded_separately_for_clipping(self): model = nn.Module() expert = nn.Module() expert._skip_fsdp = True @@ -197,81 +125,7 @@ def _setup_model(self, ep_grads, non_ep_grads): model._ep_param_groups = {"ep": ep_params, "non_ep": non_ep_params} return model, ep_params, non_ep_params - def test_l2_norm_single_group(self): - """Total L2 norm is correct when only non-EP params have grads.""" - g = torch.tensor([3.0, 4.0]) # norm = 5 - model, _, non_ep = self._setup_model(ep_grads=[], non_ep_grads=[g]) - - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): - total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=100.0) - - assert total_norm.item() == pytest.approx(5.0, abs=1e-5) - - def test_l2_norm_combined_groups(self): - """Total L2 norm combines EP-local and non-EP norms correctly.""" - ep_g = torch.tensor([3.0, 0.0]) # norm = 3 - non_ep_g = torch.tensor([0.0, 4.0]) # norm = 4 - # combined: sqrt(9 + 16) = 5 - model, ep_params, non_ep_params = self._setup_model(ep_grads=[ep_g], non_ep_grads=[non_ep_g]) - - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): - total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=100.0) - - expected = _l2_norm(*ep_params, *non_ep_params) - assert total_norm.item() == pytest.approx(expected, abs=1e-5) - - def test_clipping_reduces_gradients(self): - """When total_norm > max_norm, all gradients are scaled down uniformly.""" - ep_g = torch.tensor([6.0, 0.0]) - non_ep_g = torch.tensor([0.0, 8.0]) - # total norm = sqrt(36 + 64) = 10 - model, ep_params, non_ep_params = self._setup_model(ep_grads=[ep_g], non_ep_grads=[non_ep_g]) - max_norm = 5.0 # clip factor = 5/10 = 0.5 - - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): - total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=max_norm) - - assert total_norm.item() == pytest.approx(10.0, abs=1e-5) - # Both EP and non-EP grads should be scaled by 0.5 - torch.testing.assert_close(ep_params[0].grad, torch.tensor([3.0, 0.0])) - torch.testing.assert_close(non_ep_params[0].grad, torch.tensor([0.0, 4.0])) - - def test_no_clipping_below_max(self): - """When total_norm <= max_norm, gradients are unchanged.""" - g = torch.tensor([3.0, 4.0]) # norm = 5 - model, _, non_ep = self._setup_model(ep_grads=[], non_ep_grads=[g]) - - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): - total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=10.0) - - assert total_norm.item() == pytest.approx(5.0, abs=1e-5) - # Gradients unchanged - torch.testing.assert_close(non_ep[0].grad, torch.tensor([3.0, 4.0])) - - def test_ep_grads_not_double_scaled(self): - """EP gradients are NOT divided by ep_size — the double-division fix.""" - ep_g = torch.tensor([6.0, 8.0]) # norm = 10 - model, ep_params, _ = self._setup_model(ep_grads=[ep_g], non_ep_grads=[]) - - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): - total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=100.0) - - # Gradients should be completely unchanged (no scaling, no division) - torch.testing.assert_close(ep_params[0].grad, torch.tensor([6.0, 8.0])) - assert total_norm.item() == pytest.approx(10.0, abs=1e-5) - - def test_inf_norm(self): - """Inf-norm returns the max absolute gradient value.""" - ep_g = torch.tensor([3.0, -7.0]) - non_ep_g = torch.tensor([5.0, 2.0]) - model, _, _ = self._setup_model(ep_grads=[ep_g], non_ep_grads=[non_ep_g]) - - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): - total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=100.0, norm_type=float("inf")) - - assert total_norm.item() == pytest.approx(7.0, abs=1e-5) - - def test_inf_norm_clips_correctly(self): + def _assert_norm_modes_and_empty_gradient_policy(self): """Inf-norm clipping scales gradients when max element exceeds max_norm.""" ep_g = torch.tensor([3.0, -10.0]) non_ep_g = torch.tensor([5.0, 2.0]) @@ -285,30 +139,25 @@ def test_inf_norm_clips_correctly(self): torch.testing.assert_close(ep_params[0].grad, torch.tensor([1.5, -5.0])) torch.testing.assert_close(non_ep_params[0].grad, torch.tensor([2.5, 1.0])) - def test_empty_groups(self): - """Handles empty parameter groups without errors.""" - model = MagicMock() - model._ep_param_groups = {"ep": [], "non_ep": []} - - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): - total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=1.0) - - assert total_norm.item() == 0.0 + self._assert_empty_groups_and_params_without_grads() - def test_params_without_grads_skipped(self): - """Params with grad=None are excluded from norm computation.""" + def _assert_empty_groups_and_params_without_grads(self): g = torch.tensor([3.0, 4.0]) # norm = 5 p_with_grad = _make_param(2, grad=g) - p_no_grad = _make_param(2) # no gradient + p_no_grad = _make_param(2) model = MagicMock() model._ep_param_groups = {"ep": [p_no_grad], "non_ep": [p_with_grad]} with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=100.0) - assert total_norm.item() == pytest.approx(5.0, abs=1e-5) + model._ep_param_groups = {"ep": [], "non_ep": []} + with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=_mock_parallel_state()): + total_norm = ep_fsdp2_clip_grad_norm(model, max_norm=1.0) + assert total_norm.item() == 0.0 + # --------------------------------------------------------------------------- # 3. _skip_fsdp end-to-end: classify → clip @@ -326,7 +175,7 @@ class TestSkipFSDPClipEndToEnd: - Clipped with the same coefficient as non-EP params """ - def test_skip_fsdp_classify_then_clip(self): + def _assert_skip_fsdp_classification_and_raw_local_clip_policy(self): """_skip_fsdp expert grads are classified as EP-local and clipped correctly.""" model = nn.Module() @@ -341,8 +190,6 @@ def test_skip_fsdp_classify_then_clip(self): model.add_module("expert", expert) # Assign known gradients: expert norm=6, regular norm=8, total=10 - expert.lora.grad = torch.tensor([[3.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]) # flat norm = 3 - # Wait, let me make the math cleaner expert.lora.grad = torch.zeros(2, 4) expert.lora.grad[0, 0] = 6.0 # norm = 6 regular.weight.grad = torch.zeros(4, 4) @@ -368,7 +215,9 @@ def test_skip_fsdp_classify_then_clip(self): assert expert.lora.grad[0, 0].item() == pytest.approx(3.0, abs=1e-5) assert regular.weight.grad[0, 0].item() == pytest.approx(4.0, abs=1e-5) - def test_skip_fsdp_grads_not_reduced_or_divided(self): + self._assert_skip_fsdp_grads_not_reduced_or_divided() + + def _assert_skip_fsdp_grads_not_reduced_or_divided(self): """_skip_fsdp grads contribute their raw local norm — no all-reduce, no ep_size division.""" model = nn.Module() @@ -397,7 +246,7 @@ def test_skip_fsdp_grads_not_reduced_or_divided(self): class TestClipGradNormDispatch: """Test that clip_grad_norm dispatches to ep_fsdp2_clip_grad_norm when appropriate.""" - def test_dispatches_to_ep_path_when_ep_param_groups_present(self): + def _assert_clip_grad_norm_dispatch_policy(self): """Models with _ep_param_groups use the EP-aware clip path.""" g = torch.tensor([3.0, 4.0]) p = _make_param(2, grad=g) @@ -411,7 +260,9 @@ def test_dispatches_to_ep_path_when_ep_param_groups_present(self): assert total_norm.item() == pytest.approx(5.0, abs=1e-5) - def test_falls_through_without_ep_param_groups(self): + self._assert_falls_through_without_ep_param_groups() + + def _assert_falls_through_without_ep_param_groups(self): """Models without _ep_param_groups use the standard FSDP2 path.""" g = torch.tensor([3.0, 4.0]) p = _make_param(2, grad=g) @@ -469,16 +320,11 @@ def _mixed_mesh_model(self): model.register_parameter("p_ep", p_ep) return model, p_dp, p_ep - def test_mixed_mesh_does_not_crash_and_norm_is_correct(self, single_rank_dist): - model, _, _ = self._mixed_mesh_model() - - ps = _mock_parallel_state(ep_enabled=True) - with patch("xorl.distributed.fsdp2.clip_grad_norm.get_parallel_state", return_value=ps): - total_norm = clip_grad_norm(model, max_norm=100.0) - - assert total_norm.item() == pytest.approx(5.0, abs=1e-5) - - def test_mixed_mesh_grads_clipped_per_tensor(self, single_rank_dist): + def test_mixed_mesh_foreach_policy(self, single_rank_dist): + TestBuildEPParamGroups()._assert_shared_ep_replica_is_recorded_separately_for_clipping() + TestEPFSDP2ClipGradNorm()._assert_norm_modes_and_empty_gradient_policy() + TestSkipFSDPClipEndToEnd()._assert_skip_fsdp_classification_and_raw_local_clip_policy() + TestClipGradNormDispatch()._assert_clip_grad_norm_dispatch_policy() model, p_dp, p_ep = self._mixed_mesh_model() ps = _mock_parallel_state(ep_enabled=True) @@ -490,7 +336,9 @@ def test_mixed_mesh_grads_clipped_per_tensor(self, single_rank_dist): assert p_dp.grad.to_local()[0].item() == pytest.approx(0.6, abs=1e-4) assert p_ep.grad.to_local()[1].item() == pytest.approx(0.8, abs=1e-4) - def test_explicit_foreach_true_still_raises(self, single_rank_dist): + self._assert_explicit_foreach_true_still_raises() + + def _assert_explicit_foreach_true_still_raises(self): """An explicit foreach=True is honored — only the default is made safe.""" model, _, _ = self._mixed_mesh_model() @@ -513,8 +361,10 @@ def test_real_two_rank_ep_clip_and_nonfinite_gate(): ) result.assert_success("two-rank EP clip reduction and non-finite gate") + _assert_real_three_rank_gradient_participation_mask() + -def test_real_three_rank_gradient_participation_mask(): +def _assert_real_three_rank_gradient_participation_mask(): """A cancelling global sum still materializes zero on a missing replica.""" from tests.distributed.distributed_utils import run_distributed_script @@ -580,12 +430,6 @@ def _run_distributed_ep_clip_worker() -> None: else: assert expert.grad.item() == pytest.approx(4.0) - shared_param = nn.Parameter(torch.zeros(1)) - register_ep_replicated_gradient_hooks([shared_param]) - with patch("xorl.distributed.parallel_state.get_parallel_state", return_value=parallel_state): - (shared_param * (rank + 1.0)).sum().backward() - assert shared_param.grad.item() == pytest.approx(3.0) - coalesced = nn.Parameter(torch.zeros(1)) missing_on_rank = nn.Parameter(torch.zeros(1)) coalesced.grad = torch.tensor([rank + 1.0]) @@ -593,7 +437,12 @@ def _run_distributed_ep_clip_worker() -> None: missing_on_rank.grad = torch.tensor([2.0]) model.register_parameter("coalesced", coalesced) model.register_parameter("missing_on_rank", missing_on_rank) - model._ep_param_groups = {"ep": [], "ep_replicated": [coalesced, missing_on_rank], "non_ep": []} + model._ep_param_groups = { + "ep": [], + "ep_replicated": [coalesced, missing_on_rank], + "ep_replicated_gradient_sync": [coalesced, missing_on_rank], + "non_ep": [], + } model._ep_replicated_gradient_sync_enabled = True all_reduce_calls = [] original_all_reduce = dist.all_reduce @@ -606,7 +455,6 @@ def _counted_all_reduce(*args, **kwargs): with patch.object(dist, "all_reduce", side_effect=_counted_all_reduce): stats = synchronize_ep_replicated_gradients(model) assert len(all_reduce_calls) == 1 - assert stats.parameter_count == 2 assert stats.configured_parameter_count == 2 assert stats.participating_parameter_count == 2 assert stats.bucket_count == 1 @@ -642,7 +490,6 @@ def _run_distributed_participation_worker() -> None: with patch("xorl.distributed.parallel_state.get_parallel_state", return_value=parallel_state): stats = synchronize_ep_replicated_gradients(model) - assert stats.parameter_count == 1 assert stats.configured_parameter_count == 1 assert stats.participating_parameter_count == 1 assert parameter.grad is not None diff --git a/tests/distributed/test_ep_gradient_reduction_contract.py b/tests/distributed/test_ep_gradient_reduction_contract.py index 300c53f0..e5521ffb 100644 --- a/tests/distributed/test_ep_gradient_reduction_contract.py +++ b/tests/distributed/test_ep_gradient_reduction_contract.py @@ -19,26 +19,12 @@ pytestmark = [pytest.mark.cpu] -@pytest.mark.parametrize( - ("backend", "expected"), - [ - ("eager", GradientReductionDomain.EP_SUM), - ("native", GradientReductionDomain.EP_SUM), - ("triton", GradientReductionDomain.EP_SUM), - ("triton_w4a4", GradientReductionDomain.EP_SUM), - ("quack", GradientReductionDomain.EP_SUM), - ], -) -def test_supported_backend_gradient_contract_table(backend, expected): - assert ep_lora_gradient_reduction_domain(backend) is expected - - -def test_unknown_backend_fails_closed(): +def test_gradient_reduction_domain_admission_policy(): + for backend in ("eager", "native", "triton", "triton_w4a4", "quack"): + assert ep_lora_gradient_reduction_domain(backend) is GradientReductionDomain.EP_SUM with pytest.raises(ValueError, match="Unsupported MoE backend"): ep_lora_gradient_reduction_domain("new_backend") - -def test_unknown_metadata_domain_fails_closed(): model = nn.Module() shared = nn.Module() shared._skip_fsdp = True diff --git a/tests/distributed/test_ep_lora_weight_slicing.py b/tests/distributed/test_ep_lora_weight_slicing.py index 67c3d4d0..ec3bbff8 100644 --- a/tests/distributed/test_ep_lora_weight_slicing.py +++ b/tests/distributed/test_ep_lora_weight_slicing.py @@ -40,7 +40,7 @@ def __init__( class TestLoRAWeightInitAndShapes: """Test LoRA weight initialization, shapes, and compatibility with base weights.""" - def test_initial_shapes_zeros_and_base_compatibility(self): + def test_global_initialization_and_ep_plan_slicing_policy(self): """LoRA weights initialized at global shape, B matrices zeroed, shapes match base weights.""" config = MockConfig(num_experts=8, hidden_size=32, moe_intermediate_size=64) lora_config = MoELoRAConfig(r=4, lora_alpha=8) @@ -70,13 +70,6 @@ def test_initial_shapes_zeros_and_base_compatibility(self): assert experts.down_proj_lora_A.shape[1] == experts.down_proj.shape[1] assert experts.down_proj_lora_B.shape[2] == experts.down_proj.shape[2] - -class TestParallelPlanLoRASlicing: - """Test that ParallelPlan correctly includes and slices LoRA weights for EP.""" - - def test_ep_plan_and_shard_tensor(self): - """EP plan includes all LoRA patterns with Shard(0); shard_tensor slices by ep_rank.""" - plan = get_ep_plan() # Verify all LoRA patterns are in the plan with Shard(dim=0) @@ -200,7 +193,3 @@ def test_ep_and_non_ep_forward_with_gradients(self): loss.backward() assert gate_A2.grad is not None assert hidden_states.grad is not None - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/distributed/test_fp8_deepep_training.py b/tests/distributed/test_fp8_deepep_training.py index a1dc0b22..108b45fd 100644 --- a/tests/distributed/test_fp8_deepep_training.py +++ b/tests/distributed/test_fp8_deepep_training.py @@ -49,13 +49,10 @@ def _install_nvidia_ml_library_path() -> None: def _install_nvshmem_library_path() -> None: - try: - import nvidia.nvshmem # noqa: PLC0415 + import nvidia.nvshmem # noqa: PLC0415 - nvshmem_lib = os.path.join(list(nvidia.nvshmem.__path__)[0], "lib") - _prepend_library_path(nvshmem_lib) - except Exception: - pass + nvshmem_lib = os.path.join(list(nvidia.nvshmem.__path__)[0], "lib") + _prepend_library_path(nvshmem_lib) def _check(name: str, tensor: torch.Tensor | None) -> None: diff --git a/tests/distributed/test_glm52_exact_absorbed_kv_b_fsdp2.py b/tests/distributed/test_glm52_exact_absorbed_kv_b_fsdp2.py index 49403556..373d7376 100644 --- a/tests/distributed/test_glm52_exact_absorbed_kv_b_fsdp2.py +++ b/tests/distributed/test_glm52_exact_absorbed_kv_b_fsdp2.py @@ -304,14 +304,6 @@ def record_unsharded_child_state( if __name__ != "__main__": - @skip_if_gpu_count_less_than(1) - def test_exact_absorbed_kv_b_composes_with_one_rank_fsdp2_lifecycle() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("the qualified exact GLM-5.2 absorbed component requires Hopper") - result = run_distributed_script(__file__, num_gpus=1, timeout=300) - result.assert_success("exact absorbed kv_b should survive two calls through one-rank child FSDP2") - @skip_if_gpu_count_less_than(2) def test_exact_absorbed_kv_b_composes_with_two_rank_fsdp2_lifecycle() -> None: pytest.importorskip("sglang") diff --git a/tests/distributed/test_glm52_exact_dense_mlp_fsdp2.py b/tests/distributed/test_glm52_exact_dense_mlp_fsdp2.py index a27c7a5b..74b06ec6 100644 --- a/tests/distributed/test_glm52_exact_dense_mlp_fsdp2.py +++ b/tests/distributed/test_glm52_exact_dense_mlp_fsdp2.py @@ -164,14 +164,6 @@ def record_unsharded_state(module, _args) -> None: if __name__ != "__main__": - @skip_if_gpu_count_less_than(1) - def test_exact_dense_mlp_composes_with_one_rank_fsdp2_lifecycle() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("the qualified exact GLM-5.2 component requires Hopper") - result = run_distributed_script(__file__, num_gpus=1, timeout=180) - result.assert_success("exact dense MLP should survive one-rank FSDP2 reshard and packed-state ownership") - @skip_if_gpu_count_less_than(2) def test_exact_dense_mlp_composes_with_two_rank_fsdp2_lifecycle() -> None: pytest.importorskip("sglang") diff --git a/tests/distributed/test_glm52_exact_qlora_fsdp2.py b/tests/distributed/test_glm52_exact_qlora_fsdp2.py index 5aaad256..c6c6b96c 100644 --- a/tests/distributed/test_glm52_exact_qlora_fsdp2.py +++ b/tests/distributed/test_glm52_exact_qlora_fsdp2.py @@ -141,14 +141,6 @@ def record_unsharded_factor_state(module, _args) -> None: if __name__ != "__main__": - @skip_if_gpu_count_less_than(1) - def test_exact_tp1_qlora_composes_with_production_fsdp2_lifecycle() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("the qualified exact GLM-5.2 component requires Hopper") - result = run_distributed_script(__file__, num_gpus=1, timeout=180) - result.assert_success("exact TP1 QLoRA should survive FSDP2 reshard and packed-state deregistration") - @skip_if_gpu_count_less_than(2) def test_exact_tp1_qlora_composes_with_two_rank_fsdp2_lifecycle() -> None: pytest.importorskip("sglang") diff --git a/tests/distributed/test_gradient_accumulate_loss.py b/tests/distributed/test_gradient_accumulate_loss.py deleted file mode 100644 index 35a869ca..00000000 --- a/tests/distributed/test_gradient_accumulate_loss.py +++ /dev/null @@ -1,37 +0,0 @@ -import pytest -import torch - -import xorl.distributed.gradient_accumulate_loss as loss_module -from xorl.distributed.gradient_accumulate_loss import gradient_accumulate_loss - - -pytestmark = [pytest.mark.cpu] - - -def test_gradient_accumulate_loss_uses_requested_group(monkeypatch): - reduce_calls = [] - - def fake_all_reduce(tensor, op, group=None): - reduce_calls.append((tensor.clone(), op, group)) - - monkeypatch.setattr(loss_module.dist, "all_reduce", fake_all_reduce) - - loss = torch.tensor(2.0, requires_grad=True) - local_valid_tokens = torch.tensor(3.0) - global_valid_tokens = torch.tensor(6.0) - - ga_loss, loss_sum = gradient_accumulate_loss( - loss, - local_valid_tokens, - global_valid_tokens, - group="loss-group", - ) - ga_loss.backward() - - assert ga_loss.item() == pytest.approx(1.0) - assert loss_sum.item() == pytest.approx(6.0) - assert loss.grad.item() == pytest.approx(0.5) - assert len(reduce_calls) == 1 - _, op, group = reduce_calls[0] - assert op == torch.distributed.ReduceOp.SUM - assert group == "loss-group" diff --git a/tests/distributed/test_linear_attention_cp_equivalence.py b/tests/distributed/test_linear_attention_cp_equivalence.py index 310fe39d..d3dbdca2 100644 --- a/tests/distributed/test_linear_attention_cp_equivalence.py +++ b/tests/distributed/test_linear_attention_cp_equivalence.py @@ -20,7 +20,7 @@ if str(THIS_DIR) not in sys.path: sys.path.insert(0, str(THIS_DIR)) -from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than +from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than # noqa: E402 pytestmark = [pytest.mark.distributed] diff --git a/tests/distributed/test_lm_head_tp_ep_parallel_state.py b/tests/distributed/test_lm_head_tp_ep_parallel_state.py deleted file mode 100644 index e6b21001..00000000 --- a/tests/distributed/test_lm_head_tp_ep_parallel_state.py +++ /dev/null @@ -1,80 +0,0 @@ -"""Mesh checks for lm-head-only TP composed with expert parallelism (ep>1). - -EP is a separate re-grouping of the same ranks for the experts only; it is not an -axis of the main device_mesh. So the lm_head_mesh (replica x lm_head_tp, carved from -the CP-innermost main mesh) must be IDENTICAL whether or not ep>1. This test pins -that: ep=2 gives the same lm_head_tp / replica group membership as the ep-independent -construction, and EP is genuinely active alongside it. -""" - -import os -import sys -from pathlib import Path - -import pytest -import torch.distributed as dist - - -sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src")) - -from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state # noqa: E402 - - -pytestmark = [pytest.mark.cpu, pytest.mark.distributed] - - -def _run_case() -> None: - dist.init_process_group(backend="gloo") - try: - # 4 ranks: dp_shard=2 x ulysses(CP)=2, ep=2 overlaid, lm_head_tp=2. - # Main mesh [dp_shard=2, ulysses=2] (row-major) -> rank r: dp=r//2, cp=r%2. - # cp_replica = cp_size/lm_head_tp = 1, so lm_head_tp groups are the CP groups - # {0,1},{2,3} and replica groups span DP: {0,2},{1,3}. - init_parallel_state( - dp_size=2, - dp_shard_size=2, - ulysses_size=2, - ep_size=2, - lm_head_tp_size=2, - device_type="cpu", - ) - ps = get_parallel_state() - rank = dist.get_rank() - - # EP is genuinely active. - assert ps.ep_size == 2, ps.ep_size - assert ps.ep_enabled - assert ps.ep_fsdp_device_mesh is not None - - # Sizes. - assert dist.get_world_size(ps.lm_head_tp_group) == 2 - assert dist.get_world_size(ps.lm_head_tp_replica_group) == 2 - - # Exact membership (ep-independent expectation). - expected_tp = {0: [0, 1], 1: [0, 1], 2: [2, 3], 3: [2, 3]}[rank] - expected_replica = {0: [0, 2], 1: [1, 3], 2: [0, 2], 3: [1, 3]}[rank] - assert dist.get_process_group_ranks(ps.lm_head_tp_group) == expected_tp, ( - rank, - dist.get_process_group_ranks(ps.lm_head_tp_group), - ) - assert dist.get_process_group_ranks(ps.lm_head_tp_replica_group) == expected_replica, ( - rank, - dist.get_process_group_ranks(ps.lm_head_tp_replica_group), - ) - print(f"rank{rank} OK tp={expected_tp} replica={expected_replica} ep_size={ps.ep_size}") - finally: - dist.destroy_process_group() - - -if __name__ != "__main__": - from tests.distributed.distributed_utils import run_distributed_script - - SCRIPT_PATH = os.path.abspath(__file__) - - def test_lm_head_tp_ep_parallel_state_cpu(): - result = run_distributed_script(SCRIPT_PATH, num_gpus=4, timeout=120) - result.assert_success() - - -if __name__ == "__main__": - _run_case() diff --git a/tests/distributed/test_lm_head_tp_fsdp_e2e.py b/tests/distributed/test_lm_head_tp_fsdp_e2e.py index e80b1835..b2026988 100644 --- a/tests/distributed/test_lm_head_tp_fsdp_e2e.py +++ b/tests/distributed/test_lm_head_tp_fsdp_e2e.py @@ -33,7 +33,7 @@ pytestmark = [pytest.mark.cpu, pytest.mark.distributed] -def _run_case(dp_replicate: int, dp_shard: int, ulysses: int, lm_head_tp: int) -> None: +def _run_case(dp_replicate: int, dp_shard: int, ulysses: int, lm_head_tp: int, ep_size: int = 1) -> None: dist.init_process_group(backend="gloo") try: dp_size = dp_replicate * dp_shard @@ -42,11 +42,20 @@ def _run_case(dp_replicate: int, dp_shard: int, ulysses: int, lm_head_tp: int) - dp_replicate_size=dp_replicate, dp_shard_size=dp_shard, ulysses_size=ulysses, + ep_size=ep_size, lm_head_tp_size=lm_head_tp, device_type="cpu", ) ps = get_parallel_state() rank = dist.get_rank() + if ep_size > 1: + assert ps.ep_size == ep_size + assert ps.ep_enabled + assert ps.ep_fsdp_device_mesh is not None + expected_tp = {0: [0, 1], 1: [0, 1], 2: [2, 3], 3: [2, 3]}[rank] + expected_replica = {0: [0, 2], 1: [1, 3], 2: [0, 2], 3: [1, 3]}[rank] + assert dist.get_process_group_ranks(ps.lm_head_tp_group) == expected_tp + assert dist.get_process_group_ranks(ps.lm_head_tp_replica_group) == expected_replica # rank layout with pp=dp_replicate=ringattn=tp=1: # - CP-sourced lm-head TP: rank = dp_idx * ulysses + cp_idx. # - no-CP DP-sourced lm-head TP: ulysses=1, so dp_idx=rank and cp_idx=0. @@ -153,7 +162,7 @@ def _run_case(dp_replicate: int, dp_shard: int, ulysses: int, lm_head_tp: int) - # local hidden grad matches the reference slice for this (dp, seq) shard. ref_hidden_grad = ref_hiddens[dp_idx].grad[:, cp_idx * local_seq : (cp_idx + 1) * local_seq, :] torch.testing.assert_close(local_hidden.grad, ref_hidden_grad, rtol=1e-4, atol=1e-5) - print(f"rank{rank} dp_replicate={dp_replicate} dp_shard={dp_shard} OK loss={loss.item():.6f}") + print(f"rank{rank} dp_replicate={dp_replicate} dp_shard={dp_shard} ep={ep_size} OK loss={loss.item():.6f}") finally: dist.destroy_process_group() @@ -281,71 +290,44 @@ def _run_opd_case(dp_replicate: int, dp_shard: int, ulysses: int, lm_head_tp: in SCRIPT_PATH = os.path.abspath(__file__) - def test_lm_head_tp_fsdp_e2e_dp1_cpu(): - # dp=1, cp=4, lm_head_tp=2 -> cp_replica=2 (replica dim is purely sequence). - result = run_distributed_script( - SCRIPT_PATH, num_gpus=4, timeout=180, extra_env={"XORL_LMHEAD_E2E_CFG": "1,1,4,2"} + def test_lm_head_tp_fsdp_topology_and_loss_mode_policy(): + cases = ( + ("cp-replica-dp1", "1,1,4,2", None), + ("cp-dp2-ep2", "1,2,2,2,2", None), + ("no-cp-dp", "1,4,1,2", None), + ("no-cp-hsdp", "2,2,1,2", None), + ("opd-no-cp-dp", "1,4,1,2", "opd"), + ("opd-no-cp-hsdp", "2,2,1,2", "opd"), ) - result.assert_success() - - def test_lm_head_tp_fsdp_e2e_dp2_cpu(): - # dp=2, cp=2, lm_head_tp=2 -> cp_replica=1; the replica dim is the DP dim, so - # this validates that distinct-batch DP gradients are summed for lm-head TP. - result = run_distributed_script( - SCRIPT_PATH, num_gpus=4, timeout=180, extra_env={"XORL_LMHEAD_E2E_CFG": "1,2,2,2"} - ) - result.assert_success() - - def test_lm_head_tp_fsdp_e2e_nocp_dp_cpu(): - # dp=4, no CP, lm_head_tp=2 -> lm-head TP groups are carved from DP ranks. - # The loss gathers hidden states across the TP group, so distinct DP batches - # can share a vocab-parallel CE while replica groups sum the same vocab shard. - result = run_distributed_script( - SCRIPT_PATH, num_gpus=4, timeout=180, extra_env={"XORL_LMHEAD_E2E_CFG": "1,4,1,2"} - ) - result.assert_success() - - def test_lm_head_tp_fsdp_e2e_nocp_hsdp_cpu(): - # HSDP no-CP case: dp_replicate=2, dp_shard=2, lm_head_tp=2. The - # lm-head TP group is carved inside each shard row; the replica group - # spans both HSDP replicas for the same vocab shard. - result = run_distributed_script( - SCRIPT_PATH, num_gpus=4, timeout=180, extra_env={"XORL_LMHEAD_E2E_CFG": "2,2,1,2"} - ) - result.assert_success() - - def test_lm_head_tp_opd_fsdp_e2e_nocp_dp_cpu(): - # Production OPD path: DTensor lm_head.weight row range + matching teacher - # rows, with no-CP DP-sourced lm-head TP and replica grad sync. - result = run_distributed_script( - SCRIPT_PATH, - num_gpus=4, - timeout=180, - extra_env={"XORL_LMHEAD_E2E_CFG": "1,4,1,2", "XORL_LMHEAD_E2E_MODE": "opd"}, - ) - result.assert_success() - - def test_lm_head_tp_opd_fsdp_e2e_nocp_hsdp_cpu(): - # Production OPD path under HSDP composition. - result = run_distributed_script( - SCRIPT_PATH, - num_gpus=4, - timeout=180, - extra_env={"XORL_LMHEAD_E2E_CFG": "2,2,1,2", "XORL_LMHEAD_E2E_MODE": "opd"}, - ) - result.assert_success() + for case_id, config, mode in cases: + extra_env = {"XORL_LMHEAD_E2E_CFG": config} + if mode is not None: + extra_env["XORL_LMHEAD_E2E_MODE"] = mode + result = run_distributed_script( + SCRIPT_PATH, + num_gpus=4, + timeout=180, + extra_env=extra_env, + ) + try: + result.assert_success() + except AssertionError as error: + raise AssertionError(f"{case_id}: {error}") from error if __name__ == "__main__": cfg = os.environ.get("XORL_LMHEAD_E2E_CFG", "1,1,4,2") parts = [int(x) for x in cfg.split(",")] if len(parts) == 3: - _rep, _dp, _u, _tp = 1, *parts + _rep, _dp, _u, _tp, _ep = 1, *parts, 1 elif len(parts) == 4: _rep, _dp, _u, _tp = parts + _ep = 1 + elif len(parts) == 5: + _rep, _dp, _u, _tp, _ep = parts else: - raise ValueError(f"XORL_LMHEAD_E2E_CFG must have 3 or 4 comma-separated ints, got {cfg!r}") + raise ValueError(f"XORL_LMHEAD_E2E_CFG must have 3, 4, or 5 comma-separated ints, got {cfg!r}") if os.environ.get("XORL_LMHEAD_E2E_MODE") == "opd": _run_opd_case(_rep, _dp, _u, _tp) else: - _run_case(_rep, _dp, _u, _tp) + _run_case(_rep, _dp, _u, _tp, _ep) diff --git a/tests/distributed/test_lm_head_tp_parallel_state.py b/tests/distributed/test_lm_head_tp_parallel_state.py deleted file mode 100644 index 5ee87df2..00000000 --- a/tests/distributed/test_lm_head_tp_parallel_state.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Distributed mesh checks for lm-head-only tensor parallelism.""" - -import os -import sys -from pathlib import Path - -import pytest -import torch.distributed as dist - - -sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src")) - -from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state # noqa: E402 - - -pytestmark = [pytest.mark.cpu, pytest.mark.distributed] - - -def _run_lm_head_tp_parallel_state_case() -> None: - dist.init_process_group(backend="gloo") - try: - cfg = os.environ.get("XORL_LMHEAD_PS_CFG", "cp") - if cfg == "cp": - init_parallel_state( - dp_size=1, - dp_shard_size=1, - ulysses_size=4, - lm_head_tp_size=2, - device_type="cpu", - ) - expected_tp = {0: [0, 1], 1: [0, 1], 2: [2, 3], 3: [2, 3]} - expected_replica = {0: [0, 2], 1: [1, 3], 2: [0, 2], 3: [1, 3]} - elif cfg == "dp": - init_parallel_state( - dp_size=4, - dp_shard_size=4, - lm_head_tp_size=2, - device_type="cpu", - ) - expected_tp = {0: [0, 1], 1: [0, 1], 2: [2, 3], 3: [2, 3]} - expected_replica = {0: [0, 2], 1: [1, 3], 2: [0, 2], 3: [1, 3]} - elif cfg == "hsdp": - init_parallel_state( - dp_size=4, - dp_replicate_size=2, - dp_shard_size=2, - lm_head_tp_size=2, - device_type="cpu", - ) - expected_tp = {0: [0, 1], 1: [0, 1], 2: [2, 3], 3: [2, 3]} - expected_replica = {0: [0, 2], 1: [1, 3], 2: [0, 2], 3: [1, 3]} - else: - raise ValueError(f"Unknown XORL_LMHEAD_PS_CFG={cfg!r}") - ps = get_parallel_state() - rank = dist.get_rank() - if cfg == "cp": - assert ps.cp_size == 4 - assert dist.get_world_size(ps.sp_group) == 4 - assert dist.get_world_size(ps.ulysses_group) == 4 - assert ps.ulysses_rank == rank - else: - assert ps.cp_size == 1 - assert ps.sp_group is None - assert ps.ulysses_group is None - assert dist.get_world_size(ps.fsdp_group) == 4 - assert dist.get_world_size(ps.lm_head_tp_group) == 2 - assert dist.get_world_size(ps.lm_head_tp_replica_group) == 2 - assert dist.get_process_group_ranks(ps.lm_head_tp_group) == expected_tp[rank] - assert dist.get_process_group_ranks(ps.lm_head_tp_replica_group) == expected_replica[rank] - if cfg == "hsdp": - assert tuple(ps.fsdp_mesh.mesh.shape) == (2, 2) - assert tuple(ps.fsdp_mesh.mesh_dim_names) == ("dp_replicate", "dp_shard") - assert tuple(ps.dp_replicate_mesh.mesh.shape) == (2,) - assert tuple(ps.dp_shard_mesh.mesh.shape) == (2,) - assert ps.lm_head_mesh.mesh.tolist() == [[0, 1], [2, 3]] - finally: - dist.destroy_process_group() - - -if __name__ != "__main__": - from tests.distributed.distributed_utils import run_distributed_script - - SCRIPT_PATH = os.path.abspath(__file__) - - def test_lm_head_tp_parallel_state_cpu(): - result = run_distributed_script(SCRIPT_PATH, num_gpus=4, timeout=120, extra_env={"XORL_LMHEAD_PS_CFG": "cp"}) - result.assert_success() - - def test_lm_head_tp_parallel_state_nocp_dp_cpu(): - result = run_distributed_script(SCRIPT_PATH, num_gpus=4, timeout=120, extra_env={"XORL_LMHEAD_PS_CFG": "dp"}) - result.assert_success() - - def test_lm_head_tp_parallel_state_nocp_hsdp_cpu(): - result = run_distributed_script(SCRIPT_PATH, num_gpus=4, timeout=120, extra_env={"XORL_LMHEAD_PS_CFG": "hsdp"}) - result.assert_success() - - -if __name__ == "__main__": - _run_lm_head_tp_parallel_state_case() diff --git a/tests/distributed/test_loss_metric_reductions.py b/tests/distributed/test_loss_metric_reductions.py index 737bc8c1..4b8165bb 100644 --- a/tests/distributed/test_loss_metric_reductions.py +++ b/tests/distributed/test_loss_metric_reductions.py @@ -1,19 +1,9 @@ -"""Real-NCCL tests for the IS-metric cross-rank reduction primitives. - -Targets the three helpers in ``xorl.server.runner.model_runner`` that -``dist.all_reduce`` IS metrics across process groups: - -- ``_sp_allreduce_kl_metrics`` — per-mb CP/Ulysses reduction. -- ``ModelRunner._accumulate_is_metrics`` — cross-mb accumulation. -- ``ModelRunner._finalize_is_metrics`` — cross-DP reduction + finalization. -""" +"""Real-NCCL coverage for per-micro-batch CP/Ulysses loss-metric reduction.""" from __future__ import annotations import math import os -from types import SimpleNamespace -from unittest.mock import patch import pytest import torch @@ -43,20 +33,6 @@ def _make_metrics(device: torch.device, **values) -> dict: } -def _finalize_with_world_dp(accumulated: dict, result: dict) -> None: - """``_finalize_is_metrics`` resolves the DP group via ``get_parallel_state()``. - Spinning up the full mesh for a unit test is overkill, so stub it to - treat WORLD as the DP group for the duration of the call.""" - ps = SimpleNamespace(dp_enabled=True, dp_group=dist.group.WORLD) - with patch.object(mr, "get_parallel_state", lambda: ps): - mr.ModelRunner._finalize_is_metrics(accumulated, result) - - -# --------------------------------------------------------------------------- -# Case: per-mb CP/SP reduction (_sp_allreduce_kl_metrics) -# --------------------------------------------------------------------------- - - def _case_sp_partial_sum(device: torch.device) -> None: """Catches a re-introduction of the ``v * local_n / total_n`` weighting bug: each rank's contribution must be its raw partial sum, not a per-rank mean. @@ -77,7 +53,7 @@ def _case_sp_partial_sum(device: torch.device) -> None: ) metric_ops = {"ratio_min": "min", "ratio_max": "max"} - mr._sp_allreduce_kl_metrics(metrics, metric_ops, dist.group.WORLD) + mr._sp_allreduce_kl_metrics(metrics, dist.group.WORLD, metric_ops) total_n = 2 + 6 expected_ratio_mean = (2.0 + 7.5) / total_n @@ -96,100 +72,7 @@ def _case_sp_partial_sum(device: torch.device) -> None: assert metrics["ratio_max"].item() == 1.8 -# --------------------------------------------------------------------------- -# Case: cross-mb + cross-DP (_accumulate_is_metrics + _finalize_is_metrics) -# --------------------------------------------------------------------------- - - -# Asymmetric valid_tokens per mb verifies (sum, count) bookkeeping under -# uneven contributions. First two mbs go to rank 0, last two to rank 1. -_DP_MBS = [ - {"valid_tokens": 3, "ratio_mean": 3.6, "pg_clipfrac": 1.0, "ratio_min": 0.7, "ratio_max": 1.5}, - {"valid_tokens": 4, "ratio_mean": 4.0, "pg_clipfrac": 0.0, "ratio_min": 0.9, "ratio_max": 1.2}, - {"valid_tokens": 2, "ratio_mean": 2.5, "pg_clipfrac": 2.0, "ratio_min": 0.4, "ratio_max": 1.8}, - {"valid_tokens": 5, "ratio_mean": 4.5, "pg_clipfrac": 1.0, "ratio_min": 1.1, "ratio_max": 1.3}, -] - - -def _case_dp_accumulate_finalize(device: torch.device) -> None: - rank = dist.get_rank() - my_mbs = _DP_MBS[:2] if rank == 0 else _DP_MBS[2:] - metric_ops = {"ratio_min": "min", "ratio_max": "max"} - - accumulated = {} - for mb in my_mbs: - mr.ModelRunner._accumulate_is_metrics(accumulated, _make_metrics(device, **mb), metric_ops) - - result = {} - _finalize_with_world_dp(accumulated, result) - - total_n = sum(mb["valid_tokens"] for mb in _DP_MBS) - expected = { - "is_ratio_mean": sum(mb["ratio_mean"] for mb in _DP_MBS) / total_n, - "is_pg_clipfrac": sum(mb["pg_clipfrac"] for mb in _DP_MBS) / total_n, - # valid_tokens is itself accumulated as a mean, but its per-mb count - # is +=1 (not +=n_tokens), so the finalized value is total_n / num_mbs. - "is_valid_tokens": total_n / len(_DP_MBS), - "is_ratio_min": min(mb["ratio_min"] for mb in _DP_MBS), - "is_ratio_max": max(mb["ratio_max"] for mb in _DP_MBS), - } - for key, want in expected.items(): - got = result[key] - assert math.isclose(got, want, rel_tol=1e-12), f"[rank {rank}] {key}: got {got}, expected {want}" - - -# --------------------------------------------------------------------------- -# Case: empty-rank min/max identity -# --------------------------------------------------------------------------- - - -def _case_empty_rank_one_empty(device: torch.device) -> None: - """Rank 0 empty (all IGNORE_INDEX → ±inf identity); rank 1 has real values. - The empty rank must not leak into min/max, and the global mean must - reflect rank 1's contribution alone.""" - rank = dist.get_rank() - if rank == 0: - new_metrics = _make_metrics( - device, valid_tokens=0, ratio_mean=0.0, ratio_min=float("inf"), ratio_max=float("-inf") - ) - else: - new_metrics = _make_metrics(device, valid_tokens=5, ratio_mean=6.25, ratio_min=0.6, ratio_max=1.7) - - accumulated = {} - mr.ModelRunner._accumulate_is_metrics(accumulated, new_metrics, {"ratio_min": "min", "ratio_max": "max"}) - result = {} - _finalize_with_world_dp(accumulated, result) - - assert math.isclose(result["is_ratio_mean"], 1.25, rel_tol=1e-12) - assert math.isclose(result["is_ratio_min"], 0.6, rel_tol=1e-12) - assert math.isclose(result["is_ratio_max"], 1.7, rel_tol=1e-12) - - -def _case_empty_rank_all_empty(device: torch.device) -> None: - """All ranks empty. Mean keys with global count == 0 are dropped; min/max - with non-finite reductions fall back to 1.0.""" - new_metrics = _make_metrics(device, valid_tokens=0, ratio_mean=0.0, ratio_min=float("inf"), ratio_max=float("-inf")) - - accumulated = {} - mr.ModelRunner._accumulate_is_metrics(accumulated, new_metrics, {"ratio_min": "min", "ratio_max": "max"}) - result = {} - _finalize_with_world_dp(accumulated, result) - - assert result["is_ratio_min"] == 1.0 - assert result["is_ratio_max"] == 1.0 - assert "is_ratio_mean" not in result, f"empty-rank fallback should drop mean keys: {result}" - - -# --------------------------------------------------------------------------- -# Subprocess dispatch -# --------------------------------------------------------------------------- - - -_CASES = { - "sp_partial_sum": [_case_sp_partial_sum], - "dp_accumulate_finalize": [_case_dp_accumulate_finalize], - "empty_rank": [_case_empty_rank_one_empty, _case_empty_rank_all_empty], -} +_CASES = {"sp_partial_sum": [_case_sp_partial_sum]} def _main() -> None: @@ -212,14 +95,6 @@ def _launch(case: str): def test_sp_allreduce_kl_metrics_under_cp(): _launch("sp_partial_sum").assert_success("CP _sp_allreduce_kl_metrics partial-sum reduction") - @skip_if_gpu_count_less_than(2) - def test_accumulate_finalize_under_dp(): - _launch("dp_accumulate_finalize").assert_success("DP _accumulate_is_metrics + _finalize_is_metrics") - - @skip_if_gpu_count_less_than(2) - def test_empty_rank_min_max_identity(): - _launch("empty_rank").assert_success("Empty-rank min/max identity fallback") - if __name__ == "__main__": _main() diff --git a/tests/distributed/test_moe_memory_efficient_permutation.py b/tests/distributed/test_moe_memory_efficient_permutation.py index 6a3f6594..01a3d4a8 100644 --- a/tests/distributed/test_moe_memory_efficient_permutation.py +++ b/tests/distributed/test_moe_memory_efficient_permutation.py @@ -9,7 +9,7 @@ pytestmark = [pytest.mark.cpu] -def test_permuted_weights_follow_expert_sorted_token_order(): +def test_moe_permutation_dispatch_policy(monkeypatch): tokens = torch.arange(12, dtype=torch.float32).view(4, 3) selected_experts = torch.tensor( [ @@ -39,8 +39,14 @@ def test_permuted_weights_follow_expert_sorted_token_order(): torch.testing.assert_close(permuted_weights(routing_weights, selected_experts, 3), expected) torch.testing.assert_close(permutation_mapping, torch.tensor([0, 2, 1, 2, 3, 0, 1, 3])) + _assert_unpermute_only_scatter_adds_preweighted_outputs() + with monkeypatch.context() as pre_patch: + _assert_alltoall_pre_dispatch_routes_scores_with_received_token_order(pre_patch) + with monkeypatch.context() as post_patch: + _assert_tokens_post_all2all_hidden_chunking_matches_unchunked(post_patch) -def test_unpermute_only_scatter_adds_preweighted_outputs(): + +def _assert_unpermute_only_scatter_adds_preweighted_outputs(): expert_outputs = torch.tensor( [ [1.0, 1.0], @@ -65,7 +71,7 @@ def test_unpermute_only_scatter_adds_preweighted_outputs(): torch.testing.assert_close(output, expected) -def test_alltoall_pre_dispatch_routes_scores_with_received_token_order(monkeypatch): +def _assert_alltoall_pre_dispatch_routes_scores_with_received_token_order(monkeypatch): class FakeGroup: def size(self): return 2 @@ -149,7 +155,7 @@ def fake_all_to_all(group, input, output_split_sizes, input_split_sizes): torch.testing.assert_close(routing_weights.grad, expected_grad) -def test_tokens_post_all2all_hidden_chunking_matches_unchunked(monkeypatch): +def _assert_tokens_post_all2all_hidden_chunking_matches_unchunked(monkeypatch): class FakeGroup: def size(self): return 2 diff --git a/tests/distributed/test_muon_full_gradient.py b/tests/distributed/test_muon_full_gradient.py index 37954a88..46d8b92a 100644 --- a/tests/distributed/test_muon_full_gradient.py +++ b/tests/distributed/test_muon_full_gradient.py @@ -19,6 +19,7 @@ from torch.distributed.tensor import Shard, distribute_tensor from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state +from xorl.optim.gram_newton_schulz import GramNewtonSchulzOrthogonalizer from xorl.optim.muon import Muon from xorl.utils.device import get_nccl_backend @@ -27,7 +28,7 @@ if str(THIS_DIR) not in sys.path: sys.path.insert(0, str(THIS_DIR)) -from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than +from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than # noqa: E402 pytestmark = [pytest.mark.distributed] @@ -76,6 +77,27 @@ def _single_rank_oracle(weight_full: torch.Tensor, grad_full: torch.Tensor, *, m return p.detach() +def _assert_real_fp32_orthogonalization_matches_independent_program(device: torch.device) -> None: + """A CUDA FP32 request must not silently execute the Newton-Schulz tree in BF16.""" + + coefficients = ((3.4445, -4.775, 2.0315),) * 5 + input_matrix = torch.linspace(-1.25, 1.75, 64, device=device, dtype=torch.float32).reshape(8, 8) + actual = GramNewtonSchulzOrthogonalizer( + ns_coefficients=coefficients, + ns_use_quack_kernels=False, + ).orthogonalize(input_matrix) + + expected = input_matrix.unsqueeze(0) + expected = expected / expected.norm(dim=(-2, -1), keepdim=True).clamp_min(1e-7) + for a, b, c in coefficients: + gram = expected @ expected.mT + gram_update = torch.baddbmm(gram, gram, gram, beta=b, alpha=c) + expected = torch.baddbmm(expected, gram_update, expected, beta=a) + + assert actual.dtype is torch.float32 + torch.testing.assert_close(actual, expected.reshape_as(actual), rtol=0, atol=0) + + def _full_tensor(d): if hasattr(d, "full_tensor"): return d.full_tensor() @@ -130,6 +152,8 @@ def _run(distributed_mode: str, layout: str) -> None: expected_full_grad = _single_rank_oracle(weight_full, grad_full, mode="full_gradient") if dist.get_rank() == 0: + if distributed_mode == "full_gradient" and layout == "linear_2d": + _assert_real_fp32_orthogonalization_matches_independent_program(device) if distributed_mode == "full_gradient": err = (full_after - expected_full_grad).abs().max().item() assert err < 1e-4, ( @@ -162,7 +186,7 @@ def _main() -> None: if __name__ != "__main__": @skip_if_gpu_count_less_than(2) - def test_full_gradient_matches_single_rank_oracle_2d(): + def test_full_gradient_matches_single_rank_oracle_across_dense_and_moe_layouts(): result = run_distributed_script( __file__, num_gpus=2, @@ -171,8 +195,11 @@ def test_full_gradient_matches_single_rank_oracle_2d(): ) result.assert_success("2D Shard(0) full_gradient Muon should match single-rank oracle") + _assert_full_gradient_matches_single_rank_oracle_3d_moe() + _assert_shard_local_differs_from_full_gradient_oracle_across_dense_and_moe_layouts() + @skip_if_gpu_count_less_than(2) - def test_shard_local_differs_from_full_gradient_oracle_2d(): + def _assert_shard_local_differs_from_full_gradient_oracle_across_dense_and_moe_layouts(): result = run_distributed_script( __file__, num_gpus=2, @@ -181,8 +208,10 @@ def test_shard_local_differs_from_full_gradient_oracle_2d(): ) result.assert_success("2D shard_local should differ from full-gradient oracle on >1 rank") + _assert_shard_local_differs_from_full_gradient_oracle_3d_moe() + @skip_if_gpu_count_less_than(2) - def test_full_gradient_matches_single_rank_oracle_3d_moe(): + def _assert_full_gradient_matches_single_rank_oracle_3d_moe(): # Exercises the deferred-reshape path in ``_muon_step`` for an EP-experts-style # 3D weight ``[E, H, I]`` sharded on ``H`` (Shard(1)). result = run_distributed_script( @@ -194,7 +223,7 @@ def test_full_gradient_matches_single_rank_oracle_3d_moe(): result.assert_success("3D Shard(1) full_gradient Muon should match single-rank oracle") @skip_if_gpu_count_less_than(2) - def test_shard_local_differs_from_full_gradient_oracle_3d_moe(): + def _assert_shard_local_differs_from_full_gradient_oracle_3d_moe(): result = run_distributed_script( __file__, num_gpus=2, diff --git a/tests/distributed/test_muon_optimizer_ep_reshard.py b/tests/distributed/test_muon_optimizer_ep_reshard.py index e1a516c7..ee1e97e7 100644 --- a/tests/distributed/test_muon_optimizer_ep_reshard.py +++ b/tests/distributed/test_muon_optimizer_ep_reshard.py @@ -307,7 +307,7 @@ def _run() -> None: # torch.cuda.is_available(); we just gate whether GPUs are exposed to the ranks.) _GPU_RANKS = torch.cuda.is_available() and torch.cuda.device_count() >= 4 - def _run_case(save_ep: int, load_ep: int): + def _assert_case(save_ep: int, load_ep: int): extra_env = { "PYTHONPATH": os.path.join(REPO_ROOT, "src"), "XORL_TEST_SAVE_EP": str(save_ep), @@ -327,21 +327,10 @@ def _run_case(save_ep: int, load_ep: int): @pytest.mark.cpu @pytest.mark.distributed - def test_muon_momentum_reshard_ep2_to_ep4(): - """Save under ep_size=2 (ep_fsdp=2), load under ep_size=4 (ep_fsdp=1).""" - _run_case(save_ep=2, load_ep=4) - - @pytest.mark.cpu - @pytest.mark.distributed - def test_muon_momentum_reshard_ep4_to_ep2(): - """Save under ep_size=4 (ep_fsdp=1), load under ep_size=2 (ep_fsdp=2).""" - _run_case(save_ep=4, load_ep=2) - - @pytest.mark.cpu - @pytest.mark.distributed - def test_muon_momentum_same_ep_size_is_identity(): - """Same-ep_size round-trip: numeric no-op (bit-identical momentum).""" - _run_case(save_ep=2, load_ep=2) + def test_muon_momentum_ep_reshard_transition_matrix(): + """Both reshard directions and same-EP identity preserve global momentum.""" + for save_ep, load_ep in ((2, 4), (4, 2), (2, 2)): + _assert_case(save_ep=save_ep, load_ep=load_ep) if __name__ == "__main__": diff --git a/tests/distributed/test_native_fp8_fsdp2_materialization.py b/tests/distributed/test_native_fp8_fsdp2_materialization.py index a0eede6d..b11d926d 100644 --- a/tests/distributed/test_native_fp8_fsdp2_materialization.py +++ b/tests/distributed/test_native_fp8_fsdp2_materialization.py @@ -56,8 +56,9 @@ def __init__(self, family: str, device: torch.device | str) -> None: def _run_fsdp2_materialization_regression() -> None: - assert torch.__version__ == "2.12.1+cu132" - assert torch.__future__.get_swap_module_params_on_conversion() is False + # Exercise the replacement path that originally corrupted frozen DTensor + # parameters instead of assuming the process-wide default still selects it. + torch.__future__.set_swap_module_params_on_conversion(False) dist.init_process_group(backend="gloo") try: mesh = init_device_mesh("cpu", (dist.get_world_size(),), mesh_dim_names=("dp_shard",)) @@ -121,23 +122,24 @@ def _run_fsdp2_materialization_regression() -> None: dist.destroy_process_group() -@pytest.mark.parametrize("family", ("linear", "experts")) -def test_native_fp8_plain_apply_keeps_frozen_fp32_state(family: str) -> None: +def test_native_fp8_plain_apply_keeps_frozen_fp32_state() -> None: """The DTensor fix must not change ordinary device/dtype conversion.""" - module = _make_native_module(family, "meta") - module.to_empty(device=torch.device("cpu")) - _assert_frozen_fp32_parameters(module, device=torch.device("cpu")) - assert all(not isinstance(parameter, DTensor) for parameter in module.parameters()) - before = { - name: parameter.detach().view(torch.uint8).clone() for name, parameter in module.named_parameters(recurse=False) - } - - module.to(device=torch.device("cpu"), dtype=torch.bfloat16) - - _assert_frozen_fp32_parameters(module, device=torch.device("cpu")) - for name, parameter in module.named_parameters(recurse=False): - assert torch.equal(parameter.detach().view(torch.uint8), before[name]) + for family in ("linear", "experts"): + module = _make_native_module(family, "meta") + module.to_empty(device=torch.device("cpu")) + _assert_frozen_fp32_parameters(module, device=torch.device("cpu")) + assert all(not isinstance(parameter, DTensor) for parameter in module.parameters()) + before = { + name: parameter.detach().view(torch.uint8).clone() + for name, parameter in module.named_parameters(recurse=False) + } + + module.to(device=torch.device("cpu"), dtype=torch.bfloat16) + + _assert_frozen_fp32_parameters(module, device=torch.device("cpu")) + for name, parameter in module.named_parameters(recurse=False): + assert torch.equal(parameter.detach().view(torch.uint8), before[name]), family if __name__ != "__main__": diff --git a/tests/distributed/test_offloading.py b/tests/distributed/test_offloading.py deleted file mode 100644 index 84c162cf..00000000 --- a/tests/distributed/test_offloading.py +++ /dev/null @@ -1,19 +0,0 @@ -import torch - -from xorl.distributed.offloading import build_activation_offloading_context - - -def test_activation_offload_none_gpu_limit_defaults_to_zero() -> None: - model_fwd_context, _ = build_activation_offloading_context( - enable_activation_offload=True, - enable_gradient_checkpointing=False, - activation_gpu_limit=None, - ) - - x = torch.randn(4, 4, requires_grad=True) - with model_fwd_context: - loss = (x * x).sum() - - loss.backward() - - assert x.grad is not None diff --git a/tests/distributed/test_olmo2_qk_rms_norm.py b/tests/distributed/test_olmo2_qk_rms_norm.py deleted file mode 100644 index afd91969..00000000 --- a/tests/distributed/test_olmo2_qk_rms_norm.py +++ /dev/null @@ -1,112 +0,0 @@ -"""``Olmo2QKRMSNorm`` regression test. - -OLMo-2's full-axis ``q_norm``/``k_norm`` doesn't compose with stock TP -styles. Under colwise q/k_proj the input arrives hidden-sharded; the -plan wraps these norms with ``LocalAxisRMSNormShard`` to shard their -weight on dim 0. ``Olmo2QKRMSNorm.forward`` detects the Shard(0) -DTensor weight and runs the fused op on local tensors, computing a -local-axis RMS that matches HuggingFace's ``Olmo2RMSNorm`` reference -under TP. - -Can be run two ways: - 1. pytest tests/distributed/test_olmo2_qk_rms_norm.py -v (launches torchrun internally) - 2. torchrun --nproc_per_node=2 tests/distributed/test_olmo2_qk_rms_norm.py (direct) -""" - -import os - -import torch -import torch.distributed as dist -from torch.distributed.device_mesh import init_device_mesh -from torch.distributed.tensor.parallel import parallelize_module - -from xorl.models.layers.normalization import RMSNorm -from xorl.models.transformers.olmo2.modeling_olmo2 import Olmo2QKRMSNorm -from xorl.models.transformers.olmo2.tp_styles import LocalAxisRMSNormShard - - -HIDDEN = 8 -SEQ = 6 -BATCH = 2 - - -def _check_no_tp_passthrough(): - """Without TP, Olmo2QKRMSNorm forward delegates to the parent RMSNorm.""" - norm = Olmo2QKRMSNorm(HIDDEN) - x = torch.randn(BATCH, SEQ, HIDDEN, generator=torch.Generator().manual_seed(0)) - out = norm(x) - assert tuple(out.shape) == (BATCH, SEQ, HIDDEN) - # Numerical equivalence with the parent class on the same input. - ref = RMSNorm(HIDDEN) - ref.weight = torch.nn.Parameter(norm.weight.detach().clone()) - torch.testing.assert_close(out, ref(x), atol=1e-6, rtol=1e-6) - - -def _check_local_axis_rms_norm_shard(mesh): - """LocalAxisRMSNormShard + Olmo2QKRMSNorm: full-axis QK-norm path under TP. - - The colwise q/k_proj output arrives as a plain hidden-sharded tensor; the - custom style shards the weight on dim 0 so each rank's slice matches its - local input. The forward should compute a local-axis RMS matching a - per-rank-local single-tensor ``RMSNorm`` applied to the same shard. - """ - tp = mesh.size() - rank = dist.get_rank() - - norm = parallelize_module(Olmo2QKRMSNorm(HIDDEN), mesh, LocalAxisRMSNormShard()) - - # Mimic colwise q_proj output: same global tensor on every rank, then take - # the rank's hidden slice as the plain (non-DTensor) input. - full = torch.randn(BATCH, SEQ, HIDDEN, generator=torch.Generator().manual_seed(7)) - rank_slice = slice(rank * (HIDDEN // tp), (rank + 1) * (HIDDEN // tp)) - local_input = full[..., rank_slice].contiguous() - - out = norm(local_input) - expected_shape = (BATCH, SEQ, HIDDEN // tp) - assert tuple(out.shape) == expected_shape, f"expected output shape {expected_shape}, got {tuple(out.shape)}" - - # Reference: a plain (non-TP) RMSNorm with hidden=HIDDEN/tp and the local - # weight slice. This is the local-axis RMS HF's OLMo-2 reference computes. - ref_norm = RMSNorm(HIDDEN // tp) - ref_norm.weight = torch.nn.Parameter(norm.weight.to_local().clone()) - ref_out = ref_norm(local_input) - - torch.testing.assert_close(out, ref_out, atol=1e-6, rtol=1e-6) - - -def main(): - dist.init_process_group(backend="gloo") - rank = dist.get_rank() - world_size = dist.get_world_size() - assert HIDDEN % world_size == 0 and SEQ % world_size == 0, ( - "Test fixtures require HIDDEN and SEQ divisible by world_size" - ) - - mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("tp",)) - - _check_no_tp_passthrough() - _check_local_axis_rms_norm_shard(mesh) - - if rank == 0: - print("All Olmo2QKRMSNorm checks passed!") - - dist.destroy_process_group() - - -if __name__ != "__main__": - import pytest - - from tests.distributed.distributed_utils import run_distributed_script - - SCRIPT_PATH = os.path.abspath(__file__) - - @pytest.mark.cpu - @pytest.mark.distributed - def test_olmo2_qk_rms_norm_2rank_cpu(): - """Olmo2QKRMSNorm + LocalAxisRMSNormShard on a 2-rank gloo mesh.""" - result = run_distributed_script(SCRIPT_PATH, num_gpus=2, timeout=120) - result.assert_success() - - -if __name__ == "__main__": - main() diff --git a/tests/distributed/test_olmo2_tp_e2e.py b/tests/distributed/test_olmo2_tp_e2e.py index 9e76d12d..61acbe70 100644 --- a/tests/distributed/test_olmo2_tp_e2e.py +++ b/tests/distributed/test_olmo2_tp_e2e.py @@ -24,8 +24,15 @@ from xorl.distributed.parallel_state import init_parallel_state from xorl.distributed.torch_parallelize import _build_tp_plan +from xorl.models.layers.normalization import RMSNorm from xorl.models.transformers.olmo2.configuration_olmo2 import Olmo2Config -from xorl.models.transformers.olmo2.modeling_olmo2 import Olmo2ForCausalLM +from xorl.models.transformers.olmo2.modeling_olmo2 import Olmo2ForCausalLM, Olmo2QKRMSNorm +from xorl.models.transformers.olmo2.tp_styles import LocalAxisRMSNormShard + + +QK_NORM_HIDDEN = 8 +QK_NORM_SEQUENCE = 6 +QK_NORM_BATCH = 2 def _make_config(): @@ -50,6 +57,46 @@ def _make_config(): return cfg +def _assert_qk_norm_numerical_policy(mesh): + plain_norm = Olmo2QKRMSNorm(QK_NORM_HIDDEN) + plain_input = torch.randn( + QK_NORM_BATCH, + QK_NORM_SEQUENCE, + QK_NORM_HIDDEN, + generator=torch.Generator().manual_seed(0), + ) + plain_reference = RMSNorm(QK_NORM_HIDDEN) + plain_reference.weight = torch.nn.Parameter(plain_norm.weight.detach().clone()) + plain_output = plain_norm(plain_input) + assert tuple(plain_output.shape) == (QK_NORM_BATCH, QK_NORM_SEQUENCE, QK_NORM_HIDDEN) + torch.testing.assert_close(plain_output, plain_reference(plain_input), atol=1e-6, rtol=1e-6) + + tp_size = mesh.size() + rank = dist.get_rank() + sharded_norm = parallelize_module( + Olmo2QKRMSNorm(QK_NORM_HIDDEN), + mesh, + LocalAxisRMSNormShard(), + ) + full_input = torch.randn( + QK_NORM_BATCH, + QK_NORM_SEQUENCE, + QK_NORM_HIDDEN, + generator=torch.Generator().manual_seed(7), + ) + rank_slice = slice( + rank * (QK_NORM_HIDDEN // tp_size), + (rank + 1) * (QK_NORM_HIDDEN // tp_size), + ) + local_input = full_input[..., rank_slice].contiguous() + sharded_output = sharded_norm(local_input) + assert tuple(sharded_output.shape) == (QK_NORM_BATCH, QK_NORM_SEQUENCE, QK_NORM_HIDDEN // tp_size) + + local_reference = RMSNorm(QK_NORM_HIDDEN // tp_size) + local_reference.weight = torch.nn.Parameter(sharded_norm.weight.to_local().clone()) + torch.testing.assert_close(sharded_output, local_reference(local_input), atol=1e-6, rtol=1e-6) + + def main(): dist.init_process_group(backend="gloo") rank = dist.get_rank() @@ -66,6 +113,7 @@ def main(): device_type="cpu", ) mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("tp",)) + _assert_qk_norm_numerical_policy(mesh) torch.manual_seed(0) model = Olmo2ForCausalLM(_make_config()) diff --git a/tests/distributed/test_parallel_plan_meta_slice.py b/tests/distributed/test_parallel_plan_meta_slice.py index 8a0bed01..2ed3cf5b 100644 --- a/tests/distributed/test_parallel_plan_meta_slice.py +++ b/tests/distributed/test_parallel_plan_meta_slice.py @@ -87,42 +87,32 @@ def experts(self): return self.model.layers[3].mlp.experts -def test_meta_slicing_replaces_full_shape_with_ep_local_shape(): - """Meta param at full shape should be replaced with EP-local-shape meta param.""" +def test_meta_slicing_replaces_full_shape_and_preserves_parameter_contract(): + """The skip-loading meta path slices before allocation and preserves parameter metadata.""" num_experts, ep_size = 16, 4 H, I = 32, 64 model = _FakeModel(num_experts=num_experts, hidden=H, inter=I) + model.experts.gate_up_proj.requires_grad_(False) plan = ParallelPlan(ep_plan={"experts.gate_up_proj": Shard(0)}) - fqn2spec = plan.apply(model, _fake_ep_fsdp_mesh(ep_size), already_local=False) + fqn2spec = plan.apply(model, _fake_ep_fsdp_mesh(ep_size), already_local=True) assert model.experts.gate_up_proj.is_meta assert tuple(model.experts.gate_up_proj.shape) == (num_experts // ep_size, H, 2 * I) + assert model.experts.gate_up_proj.dtype == torch.bfloat16 + assert model.experts.gate_up_proj.requires_grad is False info = fqn2spec["experts.gate_up_proj"] assert isinstance(info, SpecInfo) assert isinstance(info.placement, Shard) and info.placement.dim == 0 - # The unrelated param is not in the ep_plan — it should be Replicate-stamped. assert isinstance(fqn2spec["experts.unrelated"].placement, Replicate) + _assert_meta_slicing_rejects_indivisible_size() + _assert_parallel_plan_stamps_explicit_replicated_gradient_reduction() -def test_meta_slicing_dispatches_even_when_already_local_is_true(): - """``already_local=True`` is the smoke's default (set by skip_weight_loading); - the meta dispatch must still fire so to_empty() doesn't materialize full shape.""" - num_experts, ep_size = 8, 8 - H, I = 16, 16 - - model = _FakeModel(num_experts=num_experts, hidden=H, inter=I) - plan = ParallelPlan(ep_plan={"experts.gate_up_proj": Shard(0)}) - plan.apply(model, _fake_ep_fsdp_mesh(ep_size), already_local=True) - # Meta path runs first, slices to ep-local even with already_local=True. - assert tuple(model.experts.gate_up_proj.shape) == (1, H, 2 * I) - assert model.experts.gate_up_proj.is_meta - - -def test_meta_slicing_assertion_on_indivisible_size(): +def _assert_meta_slicing_rejects_indivisible_size(): """Non-divisible expert dim should raise the existing ep-divisibility assert.""" num_experts, ep_size = 7, 4 # 7 % 4 != 0 model = _FakeModel(num_experts=num_experts, hidden=8, inter=8) @@ -131,20 +121,7 @@ def test_meta_slicing_assertion_on_indivisible_size(): plan.apply(model, _fake_ep_fsdp_mesh(ep_size), already_local=False) -def test_meta_slicing_preserves_dtype_and_requires_grad(): - """The new meta param must keep the original dtype and requires_grad flag.""" - num_experts, ep_size = 8, 2 - model = _FakeModel(num_experts=num_experts, hidden=8, inter=16) - model.experts.gate_up_proj.requires_grad_(False) - - plan = ParallelPlan(ep_plan={"experts.gate_up_proj": Shard(0)}) - plan.apply(model, _fake_ep_fsdp_mesh(ep_size), already_local=False) - - assert model.experts.gate_up_proj.dtype == torch.bfloat16 - assert model.experts.gate_up_proj.requires_grad is False - - -def test_parallel_plan_stamps_explicit_replicated_gradient_reduction(): +def _assert_parallel_plan_stamps_explicit_replicated_gradient_reduction(): model = _FakeSharedLoRAModel() plan = ParallelPlan(ep_plan={"experts.shared_lora": Shard(0)}) @@ -156,7 +133,7 @@ def test_parallel_plan_stamps_explicit_replicated_gradient_reduction(): assert model.experts.shared_lora.spec_info.gradient_reduction == "ep_sum" -def test_glm52_exact_meta_ep_plan_preserves_local_base_and_shards_only_expert_factors(): +def test_glm52_exact_meta_and_real_ep_plan_policy(monkeypatch): model = _ExactRoutedModel(device="meta") specs = get_glm52_ep_plan().apply(model, _fake_ep_fsdp_mesh(ep_size=16), already_local=True) @@ -184,8 +161,11 @@ def test_glm52_exact_meta_ep_plan_preserves_local_base_and_shards_only_expert_fa assert all(getattr(experts, name).shape[0] == 16 for name in experts._ep_force_shard_parameter_names) assert all(hasattr(getattr(experts, name), "spec_info") for name in experts.logical_factor_names) + _assert_glm52_exact_ep_dispositions_reject_malformed_singletons() + _assert_glm52_exact_real_already_local_plan_still_shards_global_factor_banks(monkeypatch) -def test_glm52_exact_real_already_local_plan_still_shards_global_factor_banks(monkeypatch): + +def _assert_glm52_exact_real_already_local_plan_still_shards_global_factor_banks(monkeypatch): model = _ExactRoutedModel(device="cpu") with torch.no_grad(): model.experts.gate_proj_lora_B[:, 0, 0].copy_(torch.arange(256, dtype=torch.float32)) @@ -219,25 +199,19 @@ def from_local(*, local_tensor, **_kwargs): assert info.gradient_reduction is GradientReductionDomain.NONE -@pytest.mark.parametrize( - ("parameter_name", "error_match"), - ( +def _assert_glm52_exact_ep_dispositions_reject_malformed_singletons() -> None: + for parameter_name, error_match in ( ("gate_up_packed_weight_f32", "already-local parameter"), ("gate_proj_lora_B", "force-shard parameter"), - ), -) -def test_glm52_exact_explicit_ep_dispositions_reject_malformed_singletons( - parameter_name: str, - error_match: str, -) -> None: - model = _ExactRoutedModel(device="cpu") - original = getattr(model.experts, parameter_name) - malformed_shape = (1, *original.shape[1:]) - setattr( - model.experts, - parameter_name, - nn.Parameter(torch.empty(malformed_shape, dtype=original.dtype), requires_grad=original.requires_grad), - ) - - with pytest.raises(ValueError, match=error_match): - get_glm52_ep_plan().apply(model, _fake_ep_fsdp_mesh(ep_size=16), already_local=True) + ): + model = _ExactRoutedModel(device="cpu") + original = getattr(model.experts, parameter_name) + malformed_shape = (1, *original.shape[1:]) + setattr( + model.experts, + parameter_name, + nn.Parameter(torch.empty(malformed_shape, dtype=original.dtype), requires_grad=original.requires_grad), + ) + + with pytest.raises(ValueError, match=error_match): + get_glm52_ep_plan().apply(model, _fake_ep_fsdp_mesh(ep_size=16), already_local=True) diff --git a/tests/distributed/test_parallel_state.py b/tests/distributed/test_parallel_state.py index c4a34aae..3a8239ff 100644 --- a/tests/distributed/test_parallel_state.py +++ b/tests/distributed/test_parallel_state.py @@ -53,7 +53,7 @@ class TestParallelStateConstruction: """Test ParallelState construction, validation, properties, and enabled flags.""" @patch("xorl.distributed.parallel_state.dist.is_initialized", return_value=False) - def test_defaults_and_uninitialized_properties(self, mock_is_init): + def _assert_construction_defaults_validation_and_enabled_flags(self, mock_is_init): """Default ParallelState: all sizes 1, not initialized, rank/world_size defaults; invalid cp_fsdp_mode raises.""" state = ParallelState() assert state.dp_size == 1 and state.dp_replicate_size == 1 and state.dp_shard_size == 1 @@ -66,11 +66,13 @@ def test_defaults_and_uninitialized_properties(self, mock_is_init): with pytest.raises(ValueError, match="Invalid cp_fsdp_mode"): ParallelState(cp_fsdp_mode="invalid") + self._assert_custom_init_validation_and_enabled_flags() + @patch("xorl.distributed.parallel_state.dist.is_initialized", return_value=True) @patch("xorl.distributed.parallel_state.dist.get_rank", return_value=5) @patch("xorl.distributed.parallel_state.dist.get_world_size", return_value=8) @patch("xorl.distributed.sequence_parallel.init_sequence_parallel") - def test_custom_init_validation_and_enabled_flags(self, mock_sp, mock_ws, mock_rank, mock_init): + def _assert_custom_init_validation_and_enabled_flags(self, mock_sp, mock_ws, mock_rank, mock_init): """Custom init; validation errors; initialized properties; sp/fsdp enabled flags.""" state = ParallelState(dp_size=4, dp_replicate_size=2, dp_shard_size=2, tp_size=2) assert state.dp_size == 4 and state.tp_size == 2 @@ -104,22 +106,17 @@ def teardown_method(self): @patch("xorl.distributed.parallel_state.is_torch_version_greater_than", return_value=False) @patch("xorl.distributed.parallel_state.dist.is_initialized", return_value=True) @patch("xorl.distributed.parallel_state.dist.get_world_size", return_value=8) - def test_init_get_reinit_auto_dp_shard_default(self, mock_ws, mock_is_init, mock_version): - """init sets state, get retrieves, re-init warns, auto dp_shard_size, get default when unset.""" + def test_construction_init_get_reinit_and_auto_dp_shard(self, mock_ws, mock_is_init, mock_version): + """Initialization publishes one state, infers its shard, and rejects replacement.""" + TestParallelStateConstruction()._assert_construction_defaults_validation_and_enabled_flags() + init_parallel_state(dp_size=4, tp_size=2, dp_mode="fsdp2") state = get_parallel_state() assert state.dp_size == 4 and state.tp_size == 2 and state.dp_mode == "fsdp2" + assert state.dp_shard_size == 4 + assert state.device_type in ["cuda", "cpu"] with patch("xorl.distributed.parallel_state.logger.warning") as mock_warn: init_parallel_state(dp_size=8) mock_warn.assert_called_once_with("Parallel state has already been initialized.") assert get_parallel_state().dp_size == 4 - - @patch("xorl.distributed.parallel_state.is_torch_version_greater_than", return_value=False) - @patch("xorl.distributed.parallel_state.dist.is_initialized", return_value=True) - @patch("xorl.distributed.parallel_state.dist.get_world_size", return_value=4) - def test_auto_dp_shard_and_default_uninitialized(self, mock_ws, mock_is_init, mock_version): - """Auto dp_shard_size; device_type defaults; get_parallel_state returns default when unset.""" - init_parallel_state(dp_size=4) - assert get_parallel_state().dp_shard_size == 4 - assert get_parallel_state().device_type in ["cuda", "cpu"] diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index 75c92a8f..9f3290d5 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -8,7 +8,6 @@ from xorl.distributed.pipeline_parallel import ( generate_llm_fqn_per_model_part, - is_single_stage_schedule, schedule_splits_backward, schedule_stage_style, stage_ids_for_rank, @@ -22,7 +21,7 @@ class TestFQNGeneration: """Test generate_llm_fqn_per_model_part FQN distribution logic.""" - def test_basic_stage_distribution(self): + def _assert_fqn_partitioning_policy(self): """Various stage/layer combos: correct stage count, all layers present and contiguous.""" # 2 stages, 4 layers (default FQN names) result = generate_llm_fqn_per_model_part(2, 4) @@ -51,7 +50,14 @@ def test_basic_stage_distribution(self): result = generate_llm_fqn_per_model_part(2, 2) assert len([m for stage in result for m in stage if m.startswith("layers.")]) == 2 - def test_qwen3_fqn_names_and_single_stage(self): + self._assert_qwen3_fqn_names_and_single_stage() + self._assert_error_too_many_stages() + self._assert_virtual_stage_split() + self._assert_explicit_first_last_layer_counts() + self._assert_explicit_layer_counts_infeasible() + self._assert_weighted_split_coverage() + + def _assert_qwen3_fqn_names_and_single_stage(self): """Qwen3-style nested FQN names; single stage contains all modules.""" result = generate_llm_fqn_per_model_part( 2, @@ -72,12 +78,12 @@ def test_qwen3_fqn_names_and_single_stage(self): assert result[0][-1] == "output" assert len(result[0]) == 7 # tok_embeddings + 4 layers + norm + output - def test_error_too_many_stages(self): + def _assert_error_too_many_stages(self): """Error when more stages than effective layers.""" with pytest.raises(ValueError): generate_llm_fqn_per_model_part(10, 2) - def test_virtual_stage_split(self): + def _assert_virtual_stage_split(self): """num_stages > pp_degree (virtual stages): all layers covered, contiguous.""" result = generate_llm_fqn_per_model_part( 4, @@ -92,7 +98,7 @@ def test_virtual_stage_split(self): all_layers = [m for stage in result for m in stage if m.startswith("model.layers.")] assert all_layers == [f"model.layers.{i}" for i in range(8)] - def test_explicit_first_last_layer_counts(self): + def _assert_explicit_first_last_layer_counts(self): """Megatron-style pinned first/last stage layer counts override the weight heuristic.""" result = generate_llm_fqn_per_model_part(4, 8, num_layers_in_first_stage=1, num_layers_in_last_stage=1) layer_counts = [len([m for m in stage if m.startswith("layers.")]) for stage in result] @@ -105,14 +111,14 @@ def test_explicit_first_last_layer_counts(self): layer_counts = [len([m for m in stage if m.startswith("layers.")]) for stage in result] assert layer_counts == [1, 3, 3] - def test_explicit_layer_counts_infeasible(self): + def _assert_explicit_layer_counts_infeasible(self): """Pinned counts leaving too few layers for the unpinned stages raise.""" with pytest.raises(ValueError): generate_llm_fqn_per_model_part(4, 3, num_layers_in_first_stage=1, num_layers_in_last_stage=1) with pytest.raises(ValueError): generate_llm_fqn_per_model_part(2, 4, num_layers_in_first_stage=3, num_layers_in_last_stage=3) - def test_weighted_split_coverage(self): + def _assert_weighted_split_coverage(self): """input/output weights shift layers off the first/last stages but never drop layers.""" result = generate_llm_fqn_per_model_part(6, 13, input_weight=2, output_weight=3) all_layers = [m for stage in result for m in stage if m.startswith("layers.")] @@ -122,30 +128,37 @@ def test_weighted_split_coverage(self): class TestStagePlacement: """stage_ids_for_rank must match torch's generate_stage_to_rank_mapping.""" - @pytest.mark.parametrize("pp_size", [2, 3, 4, 8]) - @pytest.mark.parametrize("stages_per_rank", [1, 2, 4]) - @pytest.mark.parametrize("style", ["loop", "v"]) - def test_matches_torch_reference(self, pp_size, stages_per_rank, style): + def _assert_stage_placement_policy(self): from torch.distributed.pipelining._utils import generate_rank_to_stage_mapping - if style == "v" and stages_per_rank == 1: - pytest.skip("v-style requires >=2 stages per rank") - num_stages = pp_size * stages_per_rank - ref = generate_rank_to_stage_mapping(pp_size, num_stages, style=style) - for rank in range(pp_size): - assert stage_ids_for_rank(rank, pp_size, num_stages, style) == ref[rank] - - def test_every_stage_owned_exactly_once(self): + # One single-stage loop mapping and both multi-stage formulas exercise + # every production branch. More PP sizes only repeat the same index + # arithmetic and previously created 24 separately collected cases. + for pp_size, stages_per_rank, style in ( + (2, 1, "loop"), + (4, 2, "loop"), + (4, 2, "v"), + ): + num_stages = pp_size * stages_per_rank + reference = generate_rank_to_stage_mapping(pp_size, num_stages, style=style) + for rank in range(pp_size): + assert stage_ids_for_rank(rank, pp_size, num_stages, style) == reference[rank] + + self._assert_every_stage_owned_exactly_once() + self._assert_single_style_requires_one_stage_per_rank() + self._assert_v_style_first_rank_owns_first_and_last() + + def _assert_every_stage_owned_exactly_once(self): for style in ("loop", "v"): owned = [s for r in range(4) for s in stage_ids_for_rank(r, 4, 8, style)] assert sorted(owned) == list(range(8)) - def test_single_style_requires_one_stage_per_rank(self): + def _assert_single_style_requires_one_stage_per_rank(self): assert stage_ids_for_rank(1, 4, 4, "single") == [1] with pytest.raises(ValueError): stage_ids_for_rank(0, 4, 8, "single") - def test_v_style_first_rank_owns_first_and_last(self): + def _assert_v_style_first_rank_owns_first_and_last(self): assert stage_ids_for_rank(0, 4, 8, "v") == [0, 7] assert stage_ids_for_rank(3, 4, 8, "v") == [3, 4] @@ -153,39 +166,29 @@ def test_v_style_first_rank_owns_first_and_last(self): class TestScheduleValidation: """Schedule whitelist + virtual-stage/microbatch constraint checks.""" - def test_styles(self): - assert schedule_stage_style("1F1B") == "single" - assert schedule_stage_style("GPipe") == "single" - assert schedule_stage_style("Interleaved1F1B") == "loop" - assert schedule_stage_style("InterleavedZeroBubble") == "loop" - assert schedule_stage_style("ZBVZeroBubble") == "v" - assert schedule_stage_style("DualPipeV") == "v" + def _assert_schedule_metadata_and_admission_policy(self): + expected = { + "1F1B": ("single", False), + "GPipe": ("single", False), + "Interleaved1F1B": ("loop", False), + "InterleavedZeroBubble": ("loop", True), + "ZBVZeroBubble": ("v", True), + "DualPipeV": ("v", True), + } + for schedule, (style, splits_backward) in expected.items(): + assert schedule_stage_style(schedule) == style + assert schedule_splits_backward(schedule) is splits_backward with pytest.raises(ValueError): schedule_stage_style("LoopedBFS") - def test_single_vs_multi_classification(self): - assert is_single_stage_schedule("1F1B") - assert is_single_stage_schedule("GPipe") - assert not is_single_stage_schedule("Interleaved1F1B") - assert not is_single_stage_schedule("ZBVZeroBubble") - - def test_backward_split_classification(self): - # dX/dW-splitting schedules need donated buffers disabled (retain_graph backward) - assert schedule_splits_backward("InterleavedZeroBubble") - assert schedule_splits_backward("ZBVZeroBubble") - assert schedule_splits_backward("DualPipeV") - assert not schedule_splits_backward("1F1B") - assert not schedule_splits_backward("GPipe") - assert not schedule_splits_backward("Interleaved1F1B") - - def test_valid_configs_pass(self): + self._assert_schedule_config_admission() + + def _assert_schedule_config_admission(self): validate_pp_schedule_config("1F1B", 1, 8, 4) validate_pp_schedule_config("Interleaved1F1B", 2, 8, 4) validate_pp_schedule_config("InterleavedZeroBubble", 2, 8, 4) validate_pp_schedule_config("ZBVZeroBubble", 2, 8, 4) validate_pp_schedule_config("DualPipeV", 2, 8, 4) - - def test_invalid_configs_raise(self): # single-stage schedule with virtual stages with pytest.raises(ValueError): validate_pp_schedule_config("1F1B", 2, 8, 4) @@ -201,3 +204,9 @@ def test_invalid_configs_raise(self): # DualPipeV minimum microbatch count (m >= num_stages) with pytest.raises(ValueError): validate_pp_schedule_config("DualPipeV", 2, 4, 4) + + +def test_pipeline_partition_placement_and_schedule_policy(): + TestFQNGeneration()._assert_fqn_partitioning_policy() + TestStagePlacement()._assert_stage_placement_policy() + TestScheduleValidation()._assert_schedule_metadata_and_admission_policy() diff --git a/tests/distributed/test_pp_profiling.py b/tests/distributed/test_pp_profiling.py index ab41a1b1..38cd1476 100644 --- a/tests/distributed/test_pp_profiling.py +++ b/tests/distributed/test_pp_profiling.py @@ -1,6 +1,6 @@ """Tests for PP bubble profiling. -Pure-math helpers (interval merge, analytic bubble, P2P estimate, patching) run on CPU; +Pure-math helpers (interval merge, P2P estimate, patching) run on CPU; the CUDA-event machinery has a single-GPU, single-stage schedule test marked gpu. """ @@ -11,75 +11,27 @@ from xorl.distributed.pp_profiling import ( PPBubbleProfiler, - analytic_bubble_fraction, estimate_p2p_bytes_per_step, merge_busy_intervals, ) class TestMergeBusyIntervals: - @pytest.mark.cpu - def test_empty_and_degenerate(self): + def _assert_interval_union_truth_table(self): assert merge_busy_intervals([]) == 0.0 assert merge_busy_intervals([(1.0, 1.0)]) == 0.0 # zero-length dropped assert merge_busy_intervals([(2.0, 1.0)]) == 0.0 # inverted dropped - - @pytest.mark.cpu - def test_single_and_disjoint(self): assert merge_busy_intervals([(0.0, 1.0)]) == pytest.approx(1.0) assert merge_busy_intervals([(0.0, 1.0), (2.0, 3.5)]) == pytest.approx(2.5) - - @pytest.mark.cpu - def test_overlapping_counted_once(self): # [0,2] and [1,3] overlap on [1,2] -> union length 3 assert merge_busy_intervals([(0.0, 2.0), (1.0, 3.0)]) == pytest.approx(3.0) # Fully contained interval adds nothing assert merge_busy_intervals([(0.0, 4.0), (1.0, 2.0)]) == pytest.approx(4.0) # Touching endpoints merge without a gap assert merge_busy_intervals([(0.0, 1.0), (1.0, 2.0)]) == pytest.approx(2.0) - - @pytest.mark.cpu - def test_unsorted_input(self): assert merge_busy_intervals([(5.0, 6.0), (0.0, 1.0), (0.5, 2.0)]) == pytest.approx(3.0) -class TestAnalyticBubbleFraction: - @pytest.mark.cpu - def test_1f1b_gpipe_textbook_values(self): - # (p-1)/(m+p-1): 27% at p=4,m=8; 30% at p=8,m=16 (GOALS.md bubble arithmetic) - assert analytic_bubble_fraction("1F1B", 4, 1, 8) == pytest.approx(3 / 11) - assert analytic_bubble_fraction("GPipe", 4, 1, 8) == pytest.approx(3 / 11) - assert analytic_bubble_fraction("1F1B", 8, 1, 16) == pytest.approx(7 / 23) - assert analytic_bubble_fraction("1f1b", 2, 1, 8) == pytest.approx(1 / 9) - - @pytest.mark.cpu - def test_interleaved_divides_bubble(self): - # (p-1)/(v*m+p-1) - assert analytic_bubble_fraction("Interleaved1F1B", 4, 2, 8) == pytest.approx(3 / 19) - assert analytic_bubble_fraction("Interleaved1F1B", 2, 2, 16) == pytest.approx(1 / 33) - # v=1 degenerates to the 1F1B value - assert analytic_bubble_fraction("Interleaved1F1B", 4, 1, 8) == pytest.approx(3 / 11) - - @pytest.mark.cpu - def test_zero_bubble_schedules(self): - assert analytic_bubble_fraction("InterleavedZeroBubble", 4, 2, 8) == 0.0 - assert analytic_bubble_fraction("ZBVZeroBubble", 2, 2, 8) == 0.0 - assert analytic_bubble_fraction("DualPipeV", 4, 2, 16) == 0.0 - - @pytest.mark.cpu - def test_pp1_has_no_bubble(self): - assert analytic_bubble_fraction("1F1B", 1, 1, 4) == 0.0 - - @pytest.mark.cpu - def test_invalid_inputs(self): - with pytest.raises(ValueError): - analytic_bubble_fraction("NotASchedule", 2, 1, 8) - with pytest.raises(ValueError): - analytic_bubble_fraction("1F1B", 2, 2, 8) # single-stage schedule with v=2 - with pytest.raises(ValueError): - analytic_bubble_fraction("1F1B", 0, 1, 8) - - class _FakeRecvInfo: def __init__(self, buffer): self.buffer = buffer @@ -124,15 +76,18 @@ class TestEstimateP2PBytes: SHAPE = (2, 128, 64) # bf16 -> 32768 bytes NBYTES = 2 * 128 * 64 * 2 - @pytest.mark.cpu - def test_middle_stage_counts_all_four_flows(self): + def _assert_p2p_byte_estimation_policy(self): stage_to_rank = {0: 0, 1: 1, 2: 2, 3: 3} stage = _FakeIOStage(1, 4, 1, stage_to_rank, in_shape=self.SHAPE, out_shape=self.SHAPE) # fwd recv + fwd send + grad recv + grad send, x 8 microbatches assert estimate_p2p_bytes_per_step([stage], 8) == 4 * self.NBYTES * 8 - @pytest.mark.cpu - def test_first_and_last_stage_skip_edge_flows(self): + self._assert_first_and_last_stage_skip_edge_flows() + self._assert_forward_only_skips_grad_flows() + self._assert_same_rank_adjacency_excluded() + self._assert_unpopulated_stage_returns_none() + + def _assert_first_and_last_stage_skip_edge_flows(self): stage_to_rank = {0: 0, 1: 1} first = _FakeIOStage(0, 2, 0, stage_to_rank, out_shape=self.SHAPE) last = _FakeIOStage(1, 2, 1, stage_to_rank, in_shape=self.SHAPE, out_shape=self.SHAPE) @@ -140,14 +95,12 @@ def test_first_and_last_stage_skip_edge_flows(self): assert estimate_p2p_bytes_per_step([first], 4) == 2 * self.NBYTES * 4 assert estimate_p2p_bytes_per_step([last], 4) == 2 * self.NBYTES * 4 - @pytest.mark.cpu - def test_forward_only_skips_grad_flows(self): + def _assert_forward_only_skips_grad_flows(self): stage_to_rank = {0: 0, 1: 1, 2: 2, 3: 3} stage = _FakeIOStage(1, 4, 1, stage_to_rank, in_shape=self.SHAPE, out_shape=self.SHAPE, has_backward=False) assert estimate_p2p_bytes_per_step([stage], 8) == 2 * self.NBYTES * 8 - @pytest.mark.cpu - def test_same_rank_adjacency_excluded(self): + def _assert_same_rank_adjacency_excluded(self): # ZBV at pp=2: rank 1 owns stages [1, 2]; the 1->2 handoff is rank-local. stage_to_rank = {0: 0, 1: 1, 2: 1, 3: 0} s1 = _FakeIOStage(1, 4, 1, stage_to_rank, in_shape=self.SHAPE, out_shape=self.SHAPE) @@ -155,8 +108,7 @@ def test_same_rank_adjacency_excluded(self): # s1: fwd recv from 0 + grad send to 0 (send to 2 local); s2 mirrored. assert estimate_p2p_bytes_per_step([s1, s2], 2) == 4 * self.NBYTES * 2 - @pytest.mark.cpu - def test_unpopulated_stage_returns_none(self): + def _assert_unpopulated_stage_returns_none(self): class _Broken: stage_index = 1 is_first = False @@ -172,6 +124,13 @@ def get_outputs_meta(self): assert estimate_p2p_bytes_per_step([_Broken()], 4) is None +@pytest.mark.cpu +def test_pp_profile_interval_and_p2p_accounting_policy(): + TestMergeBusyIntervals()._assert_interval_union_truth_table() + TestEstimateP2PBytes()._assert_p2p_byte_estimation_policy() + TestProfilerPatching()._assert_profiler_patch_lifecycle_policy() + + class _FakeComputeStage: def __init__(self, stage_index=0): self.stage_index = stage_index @@ -195,8 +154,7 @@ def __init__(self, stages): class TestProfilerPatching: - @pytest.mark.cpu - def test_patch_passthrough_and_restore(self): + def _assert_profiler_patch_lifecycle_policy(self): stages = [_FakeComputeStage(0), _FakeComputeStage(2)] profiler = PPBubbleProfiler(_FakeSchedule(stages)) for stage in stages: @@ -212,16 +170,18 @@ def test_patch_passthrough_and_restore(self): assert "forward_one_chunk" not in stage.__dict__ assert stages[0].forward_one_chunk(0, ()) == "out" # class method restored - @pytest.mark.cpu - def test_double_patch_rejected(self): + self._assert_double_patch_rejected() + self._assert_single_stage_schedule_attr() + self._assert_report_without_steps_raises() + + def _assert_double_patch_rejected(self): schedule = _FakeSchedule([_FakeComputeStage(0)]) profiler = PPBubbleProfiler(schedule) with pytest.raises(RuntimeError, match="already instance-patched"): PPBubbleProfiler(schedule) profiler.close() - @pytest.mark.cpu - def test_single_stage_schedule_attr(self): + def _assert_single_stage_schedule_attr(self): class _SingleSchedule: def __init__(self, stage): self._stage = stage @@ -232,8 +192,7 @@ def __init__(self, stage): assert profiler.stages == [stage] profiler.close() - @pytest.mark.cpu - def test_report_without_steps_raises(self): + def _assert_report_without_steps_raises(self): profiler = PPBubbleProfiler(_FakeSchedule([_FakeComputeStage(0)])) with pytest.raises(RuntimeError, match="No instrumented steps"): profiler.report() diff --git a/tests/distributed/test_quack_deepep_no_permute_parity.py b/tests/distributed/test_quack_deepep_no_permute_parity.py index 38295c6d..91ffb254 100644 --- a/tests/distributed/test_quack_deepep_no_permute_parity.py +++ b/tests/distributed/test_quack_deepep_no_permute_parity.py @@ -2,8 +2,8 @@ ``moe_implementation: quack`` + ``ep_dispatch: deepep`` routes expert compute through ``QuackEPDeepEPNoPermute`` (chunked, fused with combine) — NOT through -``QuackEPGroupGemm``, which is what ``test_deepep_correctness.py`` covers. This -test drives the no-permute path directly at the Qwen3.6-35B-A3B MoE shape +the generic permute/compute/combine path. This test drives the no-permute path +directly at the Qwen3.6-35B-A3B MoE shape (h=2048, I=512, E=256, top-8) and compares forward output and all gradients against the trusted triton generic path on the same DeepEP dispatch. @@ -56,13 +56,10 @@ def _install_nvidia_ml_library_path() -> None: def _install_nvshmem_library_path() -> None: - try: - import nvidia.nvshmem # noqa: PLC0415 + import nvidia.nvshmem # noqa: PLC0415 - nvshmem_lib = os.path.join(list(nvidia.nvshmem.__path__)[0], "lib") - _prepend_library_path(nvshmem_lib) - except Exception: - pass + nvshmem_lib = os.path.join(list(nvidia.nvshmem.__path__)[0], "lib") + _prepend_library_path(nvshmem_lib) def _cosine(a: torch.Tensor, b: torch.Tensor) -> float: @@ -338,15 +335,6 @@ def _worker_main() -> int: topk=8, routing="skewed", ), - dict( - name="q36_large_m", - num_tokens=16384, - hidden_dim=2048, - intermediate_size=512, - num_experts=256, - topk=8, - routing="balanced", - ), ] train_cases = [ @@ -368,15 +356,6 @@ def _worker_main() -> int: topk=8, use_checkpoint=False, ), - dict( - name="train_ckpt_prod_m", - num_tokens=32768, - hidden_dim=2048, - intermediate_size=512, - num_experts=256, - topk=8, - use_checkpoint=True, - ), ] all_errors: list[str] = [] diff --git a/tests/distributed/test_qwen35_lora_projection_fsdp2.py b/tests/distributed/test_qwen35_lora_projection_fsdp2.py new file mode 100644 index 00000000..a3a316ef --- /dev/null +++ b/tests/distributed/test_qwen35_lora_projection_fsdp2.py @@ -0,0 +1,222 @@ +"""Two-rank FSDP2 replay for the Qwen3.5/3.6 LoRA projection topology.""" + +from __future__ import annotations + +import hashlib +import json +import os +import sys +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +from torch.distributed._composable.fsdp import fully_shard +from torch.distributed.tensor import DTensor + +from xorl.distributed.parallel_state import get_parallel_state, init_parallel_state +from xorl.lora.modules.base import LoraModule +from xorl.lora.utils import freeze_base_parameters, inject_lora_into_model +from xorl.models.transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig +from xorl.models.transformers.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeForCausalLM +from xorl.server.weight_sync.handler import WeightSyncHandler +from xorl.utils.device import get_nccl_backend + + +THIS_DIR = Path(__file__).resolve().parent +if str(THIS_DIR) not in sys.path: + sys.path.insert(0, str(THIS_DIR)) + +from distributed_utils import run_distributed_script, skip_if_gpu_count_less_than # noqa: E402 + + +pytestmark = [pytest.mark.gpu, pytest.mark.distributed] +_RANK = 16 +_TARGETS = ["q_proj", "k_proj", "v_proj", "g_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + + +def _target_manifest() -> dict: + expected_modules = [] + for projection in ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj"): + expected_modules.append({"pattern": f"model.layers.*.linear_attn.{projection}", "count": 1, "rank": _RANK}) + for projection in ("q_proj", "k_proj", "v_proj", "o_proj"): + expected_modules.append({"pattern": f"model.layers.*.self_attn.{projection}", "count": 1, "rank": _RANK}) + for projection in ("gate_proj", "up_proj", "down_proj"): + expected_modules.append( + {"pattern": f"model.layers.*.mlp.shared_expert.{projection}", "count": 2, "rank": _RANK} + ) + return { + "schema_version": 1, + "target_modules": _TARGETS, + "expected_modules": expected_modules, + "allow_unlisted": False, + } + + +def _config() -> Qwen3_5MoeConfig: + return Qwen3_5MoeConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + max_position_embeddings=64, + layer_types=["linear_attention", "full_attention"], + linear_num_key_heads=2, + linear_num_value_heads=2, + linear_key_head_dim=16, + linear_value_head_dim=16, + decoder_sparse_step=1, + moe_intermediate_size=16, + num_experts=4, + num_experts_per_tok=2, + _attn_implementation="eager", + _moe_implementation="eager", + pad_token_id=0, + ) + + +def _build(device: torch.device) -> Qwen3_5MoeForCausalLM: + torch.manual_seed(1234) + model = Qwen3_5MoeForCausalLM(_config()) + inject_lora_into_model(model, r=_RANK, lora_alpha=_RANK, target_manifest=_target_manifest()) + freeze_base_parameters(model) + for module in model.modules(): + if isinstance(module, LoraModule): + module.exact_merged_forward = True + with torch.no_grad(): + values = torch.linspace(-0.01, 0.01, module.lora_B.numel(), dtype=torch.float32) + module.lora_B.copy_(values.reshape_as(module.lora_B)) + return model.to(device=device, dtype=torch.bfloat16).train() + + +def _sha256(tensor: torch.Tensor) -> str: + return hashlib.sha256(tensor.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()).hexdigest() + + +def _run_replay() -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend=get_nccl_backend()) + try: + world_size = dist.get_world_size() + assert world_size == 2 + device = torch.device("cuda", local_rank) + init_parallel_state(dp_size=world_size, dp_shard_size=world_size, device_type="cuda") + os.environ["XORL_GDN_BACKEND"] = "fla" + os.environ["XORL_MOE_SGLANG_FUSED_EXPERTS"] = "0" + + reference = _build(device) + sharded = _build(device) + mesh = get_parallel_state().dp_shard_mesh + fully_shard(sharded, mesh=mesh, reshard_after_forward=False) + + reference_params = dict(reference.named_parameters()) + sharded_params = dict(sharded.named_parameters()) + trainable_names = [name for name, parameter in reference_params.items() if parameter.requires_grad] + assert trainable_names + assert all(isinstance(sharded_params[name], DTensor) for name in trainable_names) + + reference_optim = torch.optim.AdamW((reference_params[name] for name in trainable_names), lr=1e-3) + sharded_optim = torch.optim.AdamW((sharded_params[name] for name in trainable_names), lr=1e-3) + decisions = [torch.tensor([[3, 5, 7, 11, 13, 17 + index]], device=device) for index in range(4)] + decision_hashes = [] + + for step in range(2): + reference_optim.zero_grad(set_to_none=True) + sharded_optim.zero_grad(set_to_none=True) + reference_loss = torch.zeros((), device=device) + sharded_loss = torch.zeros((), device=device) + for input_ids in decisions: + reference_output = reference(input_ids=input_ids, use_cache=False).last_hidden_state + sharded_output = sharded(input_ids=input_ids, use_cache=False).last_hidden_state + assert torch.equal(sharded_output.view(torch.uint8), reference_output.view(torch.uint8)) + if step == 1: + decision_hashes.append(_sha256(sharded_output)) + reference_loss = reference_loss + reference_output.float().square().mean() + sharded_loss = sharded_loss + sharded_output.float().square().mean() + reference_loss.backward() + sharded_loss.backward() + + for name in trainable_names: + sharded_grad = sharded_params[name].grad + reference_grad = reference_params[name].grad + assert isinstance(sharded_grad, DTensor), name + assert reference_grad is not None, name + full_grad = sharded_grad.full_tensor() + assert torch.equal(full_grad.view(torch.uint8), reference_grad.view(torch.uint8)), name + assert torch.isfinite(full_grad).all(), name + assert torch.count_nonzero(full_grad), name + + reference_optim.step() + sharded_optim.step() + for name in trainable_names: + full_parameter = sharded_params[name].full_tensor() + assert torch.equal(full_parameter.view(torch.uint8), reference_params[name].view(torch.uint8)), name + + sharded.unshard() + layer = sharded.model.layers[0] + + class _FakeDTensor: + pass + + synchronized = dict( + WeightSyncHandler._extract_params_for_sync( + layer, + "model.layers.0", + _FakeDTensor, + skip_moe_prefixes={"mlp.experts"}, + ) + ) + gdn = layer.linear_attn + shared = layer.mlp.shared_expert + assert torch.equal(synchronized["model.layers.0.linear_attn.q_proj.weight"], gdn.q_proj._merged_weight()) + gate, up = shared._gate_up_weights_for_forward() + assert torch.equal( + synchronized["model.layers.0.mlp.shared_expert.gate_up_proj.weight"], + torch.cat((gate, up), dim=0), + ) + assert torch.equal( + synchronized["model.layers.0.mlp.shared_expert.down_proj.weight"], + shared.down_proj._merged_weight(), + ) + + report = { + "schema_version": 1, + "event": "qwen35_lora_projection_fsdp2_four_decision_replay", + "world_size": world_size, + "decisions": len(decisions), + "optimizer_steps": 2, + "lora_rank": _RANK, + "trainable_factor_count": len(trainable_names), + "decision_output_sha256": decision_hashes, + "raw_output_bytes_equal": True, + "gradient_bytes_equal": True, + "post_update_factor_bytes_equal": True, + "folded_sync_bytes_equal": True, + "passed": True, + } + artifact = os.environ.get("QWEN35_LORA_PROJECTION_REPLAY_ARTIFACT") + if dist.get_rank() == 0: + print(json.dumps(report, indent=2, sort_keys=True), flush=True) + if artifact: + output = Path(artifact) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + dist.barrier() + finally: + dist.destroy_process_group() + + +if __name__ != "__main__": + + @skip_if_gpu_count_less_than(2) + def test_qwen35_lora_projection_two_rank_fsdp2_four_decision_replay() -> None: + result = run_distributed_script(__file__, num_gpus=2, timeout=300) + result.assert_success("Qwen3.5 LoRA projection topology must preserve raw bytes and gradients through FSDP2") + + +if __name__ == "__main__": + _run_replay() diff --git a/tests/distributed/test_qwen3_5_ulysses_cp.py b/tests/distributed/test_qwen3_5_ulysses_cp.py index e8623dfb..6545476b 100644 --- a/tests/distributed/test_qwen3_5_ulysses_cp.py +++ b/tests/distributed/test_qwen3_5_ulysses_cp.py @@ -155,24 +155,18 @@ def _main() -> None: if __name__ != "__main__": @skip_if_gpu_count_less_than(2) - def test_qwen35_ulysses_positive_smoke(): - result = run_distributed_script( - __file__, - num_gpus=2, - timeout=180, - extra_env={"QWEN35_CP_MODE": "positive"}, - ) - result.assert_success("Qwen3.5 positive Ulysses smoke should pass") - - @skip_if_gpu_count_less_than(2) - def test_qwen35_ring_fla_negative_smoke(): - result = run_distributed_script( - __file__, - num_gpus=2, - timeout=180, - extra_env={"QWEN35_CP_MODE": "negative"}, - ) - result.assert_success("Qwen3.5 ring+FLA negative smoke should fail fast") + def test_qwen35_context_parallel_execution_and_admission_contract(): + for mode, description in ( + ("positive", "Qwen3.5 positive Ulysses smoke should pass"), + ("negative", "Qwen3.5 ring+FLA negative smoke should fail fast"), + ): + result = run_distributed_script( + __file__, + num_gpus=2, + timeout=180, + extra_env={"QWEN35_CP_MODE": mode}, + ) + result.assert_success(description) if __name__ == "__main__": diff --git a/tests/distributed/test_ring_attention.py b/tests/distributed/test_ring_attention.py index e0f6b36b..7fa15b1f 100644 --- a/tests/distributed/test_ring_attention.py +++ b/tests/distributed/test_ring_attention.py @@ -1,18 +1,12 @@ -"""Tests for ring attention (context parallelism) -- unit tests only. - -Distributed tests removed -- run with torchrun separately. -""" +"""Numerical merge policy for ring attention partial outputs.""" import pytest import torch -from xorl.data.collators.sequence_shard_collator import ( - zigzag_reorder_packed_sequence, -) -from xorl.distributed.sequence_parallel.ring_attention import ( - _get_zigzag_step_section, - _merge_attn_outputs, -) + +pytest.importorskip("flash_attn_interface", reason="ring attention requires the optional FA3 interface") + +from xorl.distributed.sequence_parallel.ring_attention import _merge_attn_outputs pytestmark = [pytest.mark.distributed] @@ -67,89 +61,3 @@ def test_merge_batched_varlen_extreme_equal(self): lse_eq = torch.randn(B2, H3, S2, device="cuda") merged_outq, _ = _merge_attn_outputs(out1q, lse_eq.clone(), out2q, lse_eq.clone(), is_varlen=False) assert torch.allclose(merged_outq, (out1q + out2q) / 2, atol=1e-5) - - -class TestZigzagUnit: - """Unit tests for zigzag section logic and reorder (no GPU needed).""" - - def test_zigzag_sections_and_reorder(self): - """All ranks compute all steps; non-diagonal sections are lower/upper; reorder permutations correct.""" - # Every rank computes every step, step 0 is always diagonal - for ringattn_size in [2, 4, 8]: - for rank in range(ringattn_size): - sections = [_get_zigzag_step_section(rank, ringattn_size, s) for s in range(ringattn_size)] - assert len(sections) == ringattn_size - assert sections[0] == "diagonal" - - # Non-diagonal sections are "lower" or "upper" - for rank in range(4): - for step in range(1, 4): - section = _get_zigzag_step_section(rank, 4, step) - assert section in ("lower", "upper") - - # Single doc reorder - ringattn_size = 2 - tensor = torch.arange(40).unsqueeze(0) - position_ids = torch.arange(40).unsqueeze(0) - reordered = zigzag_reorder_packed_sequence(tensor, position_ids, ringattn_size, dim=-1) - expected = torch.cat( - [ - torch.arange(0, 10), - torch.arange(30, 40), - torch.arange(10, 20), - torch.arange(20, 30), - ] - ).unsqueeze(0) - assert torch.equal(reordered, expected) - - # Multi-doc reorder - doc_len = 20 - total = 2 * doc_len - tensor_m = torch.arange(total).unsqueeze(0) - position_ids_m = torch.cat([torch.arange(doc_len), torch.arange(doc_len)]).unsqueeze(0) - reordered_m = zigzag_reorder_packed_sequence(tensor_m, position_ids_m, ringattn_size, dim=-1) - expected_m = torch.cat( - [ - torch.arange(0, 5), - torch.arange(15, 20), - torch.arange(20, 25), - torch.arange(35, 40), - torch.arange(5, 10), - torch.arange(10, 15), - torch.arange(25, 30), - torch.arange(30, 35), - ] - ).unsqueeze(0) - assert torch.equal(reordered_m, expected_m) - - # Position IDs doc boundaries per rank - num_docs = 2 - pos_ids = torch.cat([torch.arange(doc_len) for _ in range(num_docs)]).unsqueeze(0) - reordered_pos = zigzag_reorder_packed_sequence(pos_ids, pos_ids, ringattn_size, dim=-1) - half = total // 2 - rank0_pos = reordered_pos[0, :half] - zeros = (rank0_pos == 0).nonzero(as_tuple=False).view(-1).tolist() - assert len(zeros) == 2 - - # Various ringattn_sizes: shape preserved, permutation, early < late - for cs in [2, 4, 8]: - n = 2 * cs - dl = n * 4 - t = torch.arange(dl).unsqueeze(0) - p = torch.arange(dl).unsqueeze(0) - r = zigzag_reorder_packed_sequence(t, p, cs, dim=-1) - assert r.shape == t.shape - assert set(r[0].tolist()) == set(t[0].tolist()) - chunk_size = dl // cs - for rk in range(cs): - rank_slice = r[0, rk * chunk_size : (rk + 1) * chunk_size] - sub_size = chunk_size // 2 - assert rank_slice[:sub_size].max() < rank_slice[sub_size:].min() - - # ringattn_size=1 is no-op - t1 = torch.arange(20).unsqueeze(0) - assert torch.equal(zigzag_reorder_packed_sequence(t1, t1, 1, dim=-1), t1) - - # Invalid length raises - with pytest.raises(ValueError, match="not divisible"): - zigzag_reorder_packed_sequence(torch.arange(15).unsqueeze(0), torch.arange(15).unsqueeze(0), 2, dim=-1) diff --git a/tests/distributed/test_sequence_parallel.py b/tests/distributed/test_sequence_parallel.py index abd9ce7f..66ee5174 100644 --- a/tests/distributed/test_sequence_parallel.py +++ b/tests/distributed/test_sequence_parallel.py @@ -7,9 +7,6 @@ import torch.distributed as dist import torch.multiprocessing as mp -from xorl.distributed.sequence_parallel import ( - slice_position_embedding, -) from xorl.distributed.sequence_parallel.ulysses import _Gather from xorl.distributed.sequence_parallel.utils import pad_tensor, unpad_tensor @@ -48,18 +45,6 @@ def test_pad_unpad_roundtrip_and_dims(self): ) -class TestSlicePositionEmbedding: - """Test position embedding slicing.""" - - def test_slice_position_embedding_no_group(self): - """No SP group: slicing is a no-op.""" - cos = torch.randn(1, 8, 1) - sin = torch.randn(1, 8, 1) - result_cos, result_sin = slice_position_embedding((cos, sin), dim=1, sp_group=None) - assert torch.equal(result_cos, cos) - assert torch.equal(result_sin, sin) - - def _run_gather_backward_worker(rank: int, port: int) -> None: os.environ["MASTER_ADDR"] = "127.0.0.1" os.environ["MASTER_PORT"] = str(port) @@ -88,7 +73,3 @@ def test_gather_backward_is_exact_cross_rank_reduction(unused_tcp_port): nprocs=2, start_method="spawn", ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/distributed/test_sync_padding.py b/tests/distributed/test_sync_padding.py deleted file mode 100644 index a98900c0..00000000 --- a/tests/distributed/test_sync_padding.py +++ /dev/null @@ -1,71 +0,0 @@ -from unittest.mock import Mock - -import torch - -from xorl.data.constants import IGNORE_INDEX -from xorl.distributed import sync_padding - - -def test_synchronize_micro_batch_padding_extends_cu_seqlens(monkeypatch): - monkeypatch.setattr(sync_padding.dist, "is_initialized", lambda: True) - monkeypatch.setattr(sync_padding.dist, "get_world_size", lambda group=None: 8) - monkeypatch.setattr(sync_padding, "get_device_type", lambda: "cpu") - - ps = Mock() - ps.cp_enabled = False - monkeypatch.setattr(sync_padding, "get_parallel_state", lambda: ps) - - def fake_all_reduce(tensor, op=None, group=None): - tensor.fill_(512) - - monkeypatch.setattr(sync_padding.dist, "all_reduce", fake_all_reduce) - - micro_batches = [ - { - "input_ids": torch.ones(1, 176, dtype=torch.long), - "labels": torch.ones(1, 176, dtype=torch.long), - "position_ids": torch.arange(176).unsqueeze(0), - "attention_mask": torch.ones(1, 176, dtype=torch.long), - "cu_seq_lens_q": torch.tensor([0, 83, 167, 176], dtype=torch.int32), - "cu_seq_lens_k": torch.tensor([0, 83, 167, 176], dtype=torch.int32), - "max_length_q": torch.tensor(84, dtype=torch.int32), - "max_length_k": torch.tensor(84, dtype=torch.int32), - }, - { - "input_ids": torch.ones(1, 512, dtype=torch.long), - "labels": torch.cat( - [ - torch.ones(176, dtype=torch.long), - torch.full((336,), IGNORE_INDEX, dtype=torch.long), - ] - ).unsqueeze(0), - "position_ids": torch.arange(512).unsqueeze(0), - "attention_mask": torch.cat( - [ - torch.ones(176, dtype=torch.long), - torch.zeros(336, dtype=torch.long), - ] - ).unsqueeze(0), - "cu_seq_lens_q": torch.tensor([0, 83, 167, 176], dtype=torch.int32), - "cu_seq_lens_k": torch.tensor([0, 83, 167, 176], dtype=torch.int32), - "max_length_q": torch.tensor(84, dtype=torch.int32), - "max_length_k": torch.tensor(84, dtype=torch.int32), - }, - ] - - sync_padding.synchronize_micro_batch_padding(micro_batches) - - for mb in micro_batches: - assert mb["input_ids"].shape[-1] == 512 - assert mb["labels"].shape[-1] == 512 - assert mb["position_ids"].shape[-1] == 512 - assert mb["attention_mask"].shape[-1] == 512 - assert torch.equal(mb["labels"][0, 176:], torch.full((336,), IGNORE_INDEX)) - assert mb["attention_mask"][0, 176:].sum().item() == 0 - - assert mb["cu_seq_lens_q"].tolist() == [0, 83, 167, 512] - assert mb["cu_seq_lens_k"].tolist() == [0, 83, 167, 512] - assert mb["max_length_q"] == 345 - assert mb["max_length_k"] == 345 - assert isinstance(mb["max_length_q"], int) - assert isinstance(mb["max_length_k"], int) diff --git a/tests/distributed/test_tensor_parallel.py b/tests/distributed/test_tensor_parallel.py deleted file mode 100644 index ca310b7c..00000000 --- a/tests/distributed/test_tensor_parallel.py +++ /dev/null @@ -1,91 +0,0 @@ -"""Tests for tensor parallelism support. - -Run with: - torchrun --nproc_per_node=2 -m pytest tests/distributed/test_tensor_parallel.py -v - torchrun --nproc_per_node=4 -m pytest tests/distributed/test_tensor_parallel.py -v -""" - -import pytest -import torch - -from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config -from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3Attention, Qwen3ForCausalLM, Qwen3MLP - - -pytestmark = [pytest.mark.cpu, pytest.mark.distributed] - - -class TestUnfuseForTP: - """Test that unfuse_for_tp correctly replaces fused projections.""" - - def test_attention_and_mlp_unfuse(self): - """Attention unfuse creates separate q/k/v; MLP unfuse creates separate gate/up.""" - - config = Qwen3Config( - hidden_size=256, - intermediate_size=512, - num_attention_heads=4, - num_key_value_heads=2, - head_dim=64, - ) - - # Attention unfuse - attn = Qwen3Attention(config, layer_idx=0) - assert hasattr(attn, "qkv_proj") - assert not hasattr(attn, "q_proj") - attn.unfuse_for_tp() - assert not hasattr(attn, "qkv_proj") - assert hasattr(attn, "q_proj") and hasattr(attn, "k_proj") and hasattr(attn, "v_proj") - assert attn.q_proj.out_features == 4 * 64 - assert attn.k_proj.out_features == 2 * 64 - assert attn.v_proj.out_features == 2 * 64 - - # MLP unfuse - mlp = Qwen3MLP(config) - assert hasattr(mlp, "gate_up_proj") - assert not hasattr(mlp, "gate_proj") - mlp.unfuse_for_tp() - assert not hasattr(mlp, "gate_up_proj") - assert hasattr(mlp, "gate_proj") and hasattr(mlp, "up_proj") - assert mlp.gate_proj.out_features == 512 - assert mlp.up_proj.out_features == 512 - - def test_unfused_forward_shape_and_model_level(self): - """Unfused MLP forward produces same shape; model-level unfuse covers all layers.""" - - config = Qwen3Config( - hidden_size=256, - intermediate_size=512, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - head_dim=64, - vocab_size=1000, - pad_token_id=0, - ) - - # Unfused forward shape matches fused - mlp_fused = Qwen3MLP(config).cuda() - x = torch.randn(1, 16, 256, device="cuda") - out_fused = mlp_fused(x) - - mlp_unfused = Qwen3MLP(config).cuda() - mlp_unfused.unfuse_for_tp() - out_unfused = mlp_unfused(x) - assert out_fused.shape == out_unfused.shape == torch.Size([1, 16, 256]) - - # Model-level unfuse - model = Qwen3ForCausalLM(config) - model.unfuse_for_tp() - for layer in model.model.layers: - assert not hasattr(layer.self_attn, "qkv_proj") - assert hasattr(layer.self_attn, "q_proj") - assert hasattr(layer.self_attn, "k_proj") - assert hasattr(layer.self_attn, "v_proj") - assert not hasattr(layer.mlp, "gate_up_proj") - assert hasattr(layer.mlp, "gate_proj") - assert hasattr(layer.mlp, "up_proj") - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/distributed/test_torch_parallelize_policies.py b/tests/distributed/test_torch_parallelize_policies.py index 772a8ded..113f69bb 100644 --- a/tests/distributed/test_torch_parallelize_policies.py +++ b/tests/distributed/test_torch_parallelize_policies.py @@ -5,8 +5,8 @@ from torch import nn from xorl.distributed.torch_parallelize import ( - _coerce_optional_bool_config, _configure_manual_fsdp_prefetch, + _expert_fsdp_kwargs_for_module, _expert_mixed_precision_policy, _resolve_fsdp_reduce_dtype, _sequence_parallel_fully_folded_into_fsdp, @@ -14,6 +14,8 @@ ) from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear from xorl.models.transformers.glm5.exact_shared_expert_qlora import Glm52ExactTP16SharedExpertBlockFP8QLoRA +from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config +from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3ForCausalLM from xorl.ops.block_fp8_native import NativeBlockFP8Linear @@ -35,7 +37,7 @@ def _fake_blocks(): return [_FakeBlock("block0"), _FakeBlock("block1"), _FakeBlock("block2")] -def test_mixed_precision_ignored_selection_stops_at_topmost_matching_unit() -> None: +def test_mixed_precision_selection_and_reduce_dtype_policy() -> None: class _Protected(nn.Module): pass @@ -49,8 +51,6 @@ class _Protected(nn.Module): assert selected == [root.composite, root.ordinary.protected] - -def test_exact_shared_expert_is_one_topmost_full_precision_fsdp_unit() -> None: root = nn.Module() root.shared = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="meta") @@ -67,62 +67,75 @@ def test_exact_shared_expert_is_one_topmost_full_precision_fsdp_unit() -> None: assert all( projection not in selected for projection in (root.shared.gate_proj, root.shared.up_proj, root.shared.down_proj) ) + kwargs = _expert_fsdp_kwargs_for_module( + {"mesh": "mesh", "mp_policy": "bf16", "reshard_after_forward": True}, + root.shared, + ) + assert kwargs == {"mesh": "mesh", "reshard_after_forward": True} + _assert_expert_mixed_precision_and_reduce_dtype_policy() + _assert_qwen3_unfuse_matches_the_tensor_parallel_plan() -def test_singleton_expert_mp_policy_uses_bf16_reduce_dtype() -> None: - policy = _expert_mixed_precision_policy(ep_fsdp_mesh_size=1) - - assert policy.param_dtype == torch.bfloat16 - assert policy.reduce_dtype == torch.bfloat16 - - -def test_sharded_expert_mp_policy_keeps_fp32_reduce_dtype() -> None: - policy = _expert_mixed_precision_policy(ep_fsdp_mesh_size=2) - - assert policy.param_dtype == torch.bfloat16 - assert policy.reduce_dtype == torch.float32 - - -def test_sharded_expert_mp_policy_can_use_bf16_reduce_dtype() -> None: - policy = _expert_mixed_precision_policy(ep_fsdp_mesh_size=2, reduce_dtype=torch.bfloat16) - - assert policy.param_dtype == torch.bfloat16 - assert policy.reduce_dtype == torch.bfloat16 +def _assert_expert_mixed_precision_and_reduce_dtype_policy() -> None: + for mesh_size, override, expected in ( + (1, None, torch.bfloat16), + (2, None, torch.float32), + (2, torch.bfloat16, torch.bfloat16), + ): + kwargs = {} if override is None else {"reduce_dtype": override} + policy = _expert_mixed_precision_policy(ep_fsdp_mesh_size=mesh_size, **kwargs) + assert policy.param_dtype == torch.bfloat16 + assert policy.reduce_dtype == expected -def test_resolve_fsdp_reduce_dtype() -> None: assert _resolve_fsdp_reduce_dtype("fp32") is torch.float32 assert _resolve_fsdp_reduce_dtype("bf16") is torch.bfloat16 with pytest.raises(ValueError, match="Unsupported fsdp_reduce_dtype"): _resolve_fsdp_reduce_dtype("fp16") -@pytest.mark.parametrize( - ("value", "expected"), - [ - (True, True), - (False, False), - (1, True), - (0, False), - ("true", True), - ("false", False), - ("YES", True), - ("off", False), - (None, None), - ], -) -def test_coerce_optional_bool_config(value, expected) -> None: - assert _coerce_optional_bool_config(value, name="flag") is expected - - -def test_coerce_optional_bool_config_rejects_ambiguous_values() -> None: - with pytest.raises(ValueError, match="flag must be a boolean value"): - _coerce_optional_bool_config("definitely", name="flag") - - -@pytest.mark.parametrize( - ("ulysses_enabled", "ringattn_enabled", "cp_fsdp_mode", "expected"), - [ +def _assert_qwen3_unfuse_matches_the_tensor_parallel_plan() -> None: + config = Qwen3Config( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=128, + pad_token_id=0, + ) + model = Qwen3ForCausalLM(config) + + model.unfuse_for_tp() + + assert model._unfused_for_tp is True + assert model.get_checkpoint_handler() is None + assert config.base_model_tp_plan == { + "embed_tokens": "embedding", + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + for layer in model.model.layers: + attention = layer.self_attn + assert not hasattr(attention, "qkv_proj") + assert attention.q_proj.out_features == 4 * 16 + assert attention.k_proj.out_features == 2 * 16 + assert attention.v_proj.out_features == 2 * 16 + + mlp = layer.mlp + assert not hasattr(mlp, "gate_up_proj") + assert mlp.gate_proj.out_features == 128 + assert mlp.up_proj.out_features == 128 + + +def test_sequence_parallel_fully_folded_into_fsdp() -> None: + for ulysses_enabled, ringattn_enabled, cp_fsdp_mode, expected in ( (True, False, "all", True), (True, False, "ulysses_only", True), (True, False, "ring_only", False), @@ -133,82 +146,43 @@ def test_coerce_optional_bool_config_rejects_ambiguous_values() -> None: (True, True, "ulysses_only", False), (True, True, "ring_only", False), (True, True, "none", False), - ], -) -def test_sequence_parallel_fully_folded_into_fsdp( - ulysses_enabled: bool, - ringattn_enabled: bool, - cp_fsdp_mode: str, - expected: bool, -) -> None: - state = SimpleNamespace( - ulysses_enabled=ulysses_enabled, - ringattn_enabled=ringattn_enabled, - cp_fsdp_mode=cp_fsdp_mode, - ) - - assert _sequence_parallel_fully_folded_into_fsdp(state) is expected - - -def test_manual_fsdp_prefetch_can_enable_backward_without_forward() -> None: - blocks = _fake_blocks() - - _configure_manual_fsdp_prefetch( - blocks, - need_manual_prefetch=True, - enable_forward_prefetch=False, - enable_backward_prefetch=True, - ) - - assert [block.forward_prefetch for block in blocks] == [None, None, None] - assert blocks[0].backward_prefetch is None - assert blocks[1].backward_prefetch == ["block0.experts", "block0.gate", "block0.attn"] - assert blocks[2].backward_prefetch == ["block1.experts", "block1.gate", "block1.attn"] - - -def test_manual_fsdp_prefetch_can_enable_forward_without_backward() -> None: - blocks = _fake_blocks() - - _configure_manual_fsdp_prefetch( - blocks, - need_manual_prefetch=True, - enable_forward_prefetch=True, - enable_backward_prefetch=False, - ) - - assert blocks[0].forward_prefetch == ["block1.experts", "block1.gate", "block1.attn"] - assert blocks[1].forward_prefetch == ["block2.experts", "block2.gate", "block2.attn"] - assert blocks[2].forward_prefetch is None - assert [block.backward_prefetch for block in blocks] == [None, None, None] - - -def test_manual_fsdp_prefetch_configures_both_directions_by_default() -> None: - blocks = _fake_blocks() - - _configure_manual_fsdp_prefetch( - blocks, - need_manual_prefetch=True, - enable_forward_prefetch=True, - enable_backward_prefetch=True, - ) - - assert blocks[0].forward_prefetch == ["block1.experts", "block1.gate", "block1.attn"] - assert blocks[1].forward_prefetch == ["block2.experts", "block2.gate", "block2.attn"] - assert blocks[2].forward_prefetch is None - assert blocks[0].backward_prefetch is None - assert blocks[1].backward_prefetch == ["block0.experts", "block0.gate", "block0.attn"] - assert blocks[2].backward_prefetch == ["block1.experts", "block1.gate", "block1.attn"] - - -def test_manual_fsdp_prefetch_noops_when_not_needed() -> None: - blocks = _fake_blocks() - - _configure_manual_fsdp_prefetch( - blocks, - need_manual_prefetch=False, - enable_forward_prefetch=True, - enable_backward_prefetch=True, - ) - - assert [block.forward_prefetch for block in blocks] == [None, None, None] - assert [block.backward_prefetch for block in blocks] == [None, None, None] + ): + state = SimpleNamespace( + ulysses_enabled=ulysses_enabled, + ringattn_enabled=ringattn_enabled, + cp_fsdp_mode=cp_fsdp_mode, + ) + assert _sequence_parallel_fully_folded_into_fsdp(state) is expected, ( + ulysses_enabled, + ringattn_enabled, + cp_fsdp_mode, + ) + + +def test_manual_fsdp_prefetch_direction_policy() -> None: + none = [None, None, None] + forward = [ + ["block1.experts", "block1.gate", "block1.attn"], + ["block2.experts", "block2.gate", "block2.attn"], + None, + ] + backward = [ + None, + ["block0.experts", "block0.gate", "block0.attn"], + ["block1.experts", "block1.gate", "block1.attn"], + ] + for needed, enable_forward, enable_backward, expected_forward, expected_backward in ( + (True, False, True, none, backward), + (True, True, False, forward, none), + (True, True, True, forward, backward), + (False, True, True, none, none), + ): + blocks = _fake_blocks() + _configure_manual_fsdp_prefetch( + blocks, + need_manual_prefetch=needed, + enable_forward_prefetch=enable_forward, + enable_backward_prefetch=enable_backward, + ) + assert [block.forward_prefetch for block in blocks] == expected_forward + assert [block.backward_prefetch for block in blocks] == expected_backward diff --git a/tests/distributed/test_utils.py b/tests/distributed/test_utils.py deleted file mode 100644 index a1b33093..00000000 --- a/tests/distributed/test_utils.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Tests for xorl.distributed.utils module.""" - -import pytest -import torch.nn as nn - -from xorl.distributed.utils import ( - check_all_fqn_match, - check_any_fqn_match, - check_fqn_match, - get_module_from_path, - set_module_from_path, -) - - -pytestmark = [pytest.mark.cpu, pytest.mark.distributed] - - -class SimpleModel(nn.Module): - """Simple test model for path-based operations.""" - - def __init__(self): - super().__init__() - self.layer1 = nn.Linear(10, 20) - self.layer2 = nn.Linear(20, 30) - self.nested = nn.Sequential(nn.Linear(30, 40), nn.ReLU(), nn.Linear(40, 50)) - - -class TestModulePaths: - """Test set_module_from_path and get_module_from_path together.""" - - def test_set_get_roundtrip_and_nested(self): - """Set/get roundtrip at top-level and nested levels; get nonexistent raises.""" - model = SimpleModel() - - # Top-level set/get roundtrip - new_layer = nn.Linear(10, 99) - set_module_from_path(model, "layer1", new_layer) - assert get_module_from_path(model, "layer1") is new_layer - - # Nested set/get roundtrip - nested_layer = nn.Linear(30, 100) - set_module_from_path(model, "nested.0", nested_layer) - assert get_module_from_path(model, "nested.0") is nested_layer - assert model.nested[0].out_features == 100 - - # Deeply nested via ModuleDict - model.deep = nn.ModuleDict({"a": nn.Sequential(nn.Linear(5, 10))}) - deep_layer = nn.Linear(5, 20) - set_module_from_path(model, "deep.a.0", deep_layer) - assert get_module_from_path(model, "deep.a.0") is deep_layer - assert model.deep["a"][0].out_features == 20 - - # Nonexistent paths raise - with pytest.raises(AttributeError): - get_module_from_path(model, "nonexistent") - with pytest.raises((AttributeError, TypeError)): - get_module_from_path(model, "layer1.nonexistent") - - -class TestCheckFqnMatch: - """Test check_fqn_match: exact, wildcard, partial, and input validation.""" - - def test_matching_and_wildcards(self): - """Exact match, wildcard positions, partial match fails, input validation.""" - # Exact match - assert check_fqn_match("layer1.weight", "layer1.weight") is not None - - # No match - assert check_fqn_match("layer1.weight", "layer2.weight") is None - - # Partial match fails - assert check_fqn_match("layer1", "layer1.weight") is None - - # Wildcard at beginning - assert check_fqn_match("*.weight", "layer1.weight") is not None - - # Wildcard at end - assert check_fqn_match("layer1.*", "layer1.weight") is not None - - # Multiple wildcards - assert check_fqn_match("*.layer*.weight", "model.layer1.weight") is not None - - # Input validation - with pytest.raises(AssertionError, match="fqn_pattern must be a str"): - check_fqn_match(["not_a_str"], "fqn") - with pytest.raises(AssertionError, match="fqn must be a str"): - check_fqn_match("pattern", ["not_a_str"]) - - -class TestCheckAllFqnMatch: - """Test check_all_fqn_match: wildcards, failures, edge cases, validation.""" - - def test_all_fqn_matching(self): - """Wildcard match, number mismatch, length mismatch, no match, multi-wildcard, empty, validation.""" - # Matching with wildcards (same number) - assert check_all_fqn_match(["layer*.weight", "layer*.bias"], ["layer1.weight", "layer1.bias"]) is True - - # Different wildcard numbers -> fail - assert check_all_fqn_match(["layer*.weight", "layer*.bias"], ["layer1.weight", "layer2.bias"]) is False - - # Length mismatch -> fail - assert check_all_fqn_match(["layer1.weight", "layer2.bias"], ["layer1.weight"]) is False - - # No match -> fail - assert check_all_fqn_match(["layer1.weight", "layer2.bias"], ["layer3.weight", "layer4.bias"]) is False - - # Multiple wildcards same number - assert ( - check_all_fqn_match( - ["block*.layer*.weight", "block*.layer*.bias"], ["block1.layer1.weight", "block1.layer1.bias"] - ) - is True - ) - - # Empty lists - assert check_all_fqn_match([], []) is True - - # Input validation - with pytest.raises(AssertionError, match="path_patterns must be a list"): - check_all_fqn_match("not_a_list", ["key1"]) - with pytest.raises(AssertionError, match="path_keys must be a list or tuple"): - check_all_fqn_match(["pattern1"], "not_a_list") - - -class TestCheckAnyFqnMatch: - """Test check_any_fqn_match: exact, wildcard, no match, return_idx, prefix, validation.""" - - def test_any_fqn_matching(self): - """Exact match, wildcard, no match, return_idx, prefix, first-match-wins, validation.""" - patterns = ["layer1.weight", "layer2.bias", "layer3.weight"] - - # Exact match - assert check_any_fqn_match(patterns, "layer1.weight") is True - - # No match - assert check_any_fqn_match(patterns, "layer99.weight") is False - - # Wildcard match - assert check_any_fqn_match(["layer*.weight", "layer*.bias"], "layer5.weight") is True - - # return_idx - assert check_any_fqn_match(patterns, "layer2.bias", return_idx=True) == 1 - assert check_any_fqn_match(patterns, "layer99.weight", return_idx=True) == -1 - - # First match wins - assert check_any_fqn_match(["layer*.weight", "layer1.weight"], "layer1.weight", return_idx=True) == 0 - - # With prefix - assert check_any_fqn_match(["weight", "bias"], "model.layer1.weight", prefix="model.layer1") is True - assert check_any_fqn_match(["weight", "bias"], "model.layer1.weight", prefix="model.layer2") is False - - # Input validation - with pytest.raises(AssertionError, match="path_patterns must be a list"): - check_any_fqn_match("not_a_list", "key") - with pytest.raises(AssertionError, match="path_key must be a str"): - check_any_fqn_match(["pattern"], ["not_a_str"]) diff --git a/tests/distributed/test_vocab_parallel_ce.py b/tests/distributed/test_vocab_parallel_ce.py index 5203ceea..fd1d060c 100644 --- a/tests/distributed/test_vocab_parallel_ce.py +++ b/tests/distributed/test_vocab_parallel_ce.py @@ -6,7 +6,6 @@ """ import os -import time import torch import torch.distributed as dist @@ -113,98 +112,6 @@ def check_backward(rank, world_size, tp_group): assert grad_w_err < 1e-2, f"grad_w error too large ({mode}): {grad_w_err}" -def _bench_one(hidden_states, local_weight, labels, tp_group, use_compile, n_warmup=5, n_iter=20, fwd_only=False): - """Run warmup + timed iterations, return (ms/iter, peak_memory_MB).""" - for _ in range(n_warmup): - ce = vocab_parallel_cross_entropy(hidden_states, local_weight, labels, tp_group, use_compile=use_compile) - if not fwd_only: - ce.sum().backward() - hidden_states.grad = None - local_weight.grad = None - torch.cuda.synchronize() - - # Reset peak memory stats before timed run - torch.cuda.reset_peak_memory_stats() - mem_before = torch.cuda.memory_allocated() - - start = time.perf_counter() - for _ in range(n_iter): - ce = vocab_parallel_cross_entropy(hidden_states, local_weight, labels, tp_group, use_compile=use_compile) - if not fwd_only: - ce.sum().backward() - hidden_states.grad = None - local_weight.grad = None - torch.cuda.synchronize() - ms = (time.perf_counter() - start) / n_iter * 1000 - - peak_mem = torch.cuda.max_memory_allocated() - # Peak activation memory = peak total - memory before (which includes weights + NCCL buffers) - peak_activation_mb = (peak_mem - mem_before) / 1024 / 1024 - - return ms, peak_activation_mb, peak_mem / 1024 / 1024 - - -def bench_perf(rank, world_size, tp_group): - """Benchmark eager vs compiled vocab-parallel CE with memory tracking.""" - torch.manual_seed(42) - BT = 4096 - H = 4096 - V = 152064 # Qwen3 vocab - local_V = V // world_size - - hidden_states = torch.randn(BT, H, device="cuda", dtype=torch.bfloat16, requires_grad=True) - local_weight = torch.randn(local_V, H, device="cuda", dtype=torch.bfloat16, requires_grad=True) - labels = torch.randint(0, V, (BT,), device="cuda") - - if rank == 0: - print(f"[perf] BT={BT}, H={H}, V={V}, tp={world_size}") - print(f"[perf] local_V={local_V}, num_chunks=8") - weight_mb = local_weight.nelement() * local_weight.element_size() / 1024 / 1024 - hidden_mb = hidden_states.nelement() * hidden_states.element_size() / 1024 / 1024 - print(f"[perf] weight: {weight_mb:.1f} MB, hidden: {hidden_mb:.1f} MB") - print() - - # --- Forward-only benchmark --- - if rank == 0: - print("--- Forward only ---") - print(f"{'Mode':<12} {'Time (ms)':<12} {'Peak Act (MB)':<16} {'Peak Total (MB)'}") - for use_compile in [False, True]: - mode = "compiled" if use_compile else "eager" - ms, peak_act_mb, peak_total_mb = _bench_one( - hidden_states, - local_weight, - labels, - tp_group, - use_compile=use_compile, - fwd_only=True, - ) - if rank == 0: - print(f"{mode:<12} {ms:<12.2f} {peak_act_mb:<16.1f} {peak_total_mb:.1f}") - - if rank == 0: - print() - - # --- Forward + Backward benchmark --- - if rank == 0: - print("--- Forward + Backward ---") - print(f"{'Mode':<12} {'Time (ms)':<12} {'Peak Act (MB)':<16} {'Peak Total (MB)'}") - for use_compile in [False, True]: - mode = "compiled" if use_compile else "eager" - ms, peak_act_mb, peak_total_mb = _bench_one( - hidden_states, - local_weight, - labels, - tp_group, - use_compile=use_compile, - fwd_only=False, - ) - if rank == 0: - print(f"{mode:<12} {ms:<12.2f} {peak_act_mb:<16.1f} {peak_total_mb:.1f}") - - if rank == 0: - print() - - def main(): rank, world_size = setup() tp_group = dist.group.WORLD @@ -222,9 +129,6 @@ def main(): if rank == 0: print() - bench_perf(rank, world_size, tp_group) - dist.barrier() - if rank == 0: print("\nAll tests passed!") diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index a3b17fc9..b9bec795 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -39,12 +39,6 @@ def tiny_moe_model_dir_with_weights(tmp_workspace): return create_tiny_model_dir(tmp_workspace, model_type="moe", save_weights=True) -@pytest.fixture -def small_dense_model_dir_with_weights(tmp_workspace): - """Small Qwen3 dense model (hidden=256) with saved weights.""" - return create_tiny_model_dir(tmp_workspace, model_type="dense_large", save_weights=True) - - @pytest.fixture def small_moe_model_dir_with_weights(tmp_workspace): """Small Qwen3-MoE model (moe_intermediate=64) with saved weights. diff --git a/tests/e2e/e2e_utils.py b/tests/e2e/e2e_utils.py index 8dc5021a..beaec2f0 100644 --- a/tests/e2e/e2e_utils.py +++ b/tests/e2e/e2e_utils.py @@ -317,6 +317,8 @@ def generate_training_config( lora_rank: int = 8, lora_alpha: int = 8, lora_target_modules: Optional[List[str]] = None, + extra_model: Optional[Dict[str, Any]] = None, + extra_data: Optional[Dict[str, Any]] = None, extra_train: Optional[Dict[str, Any]] = None, extra_lora: Optional[Dict[str, Any]] = None, ) -> str: @@ -383,6 +385,12 @@ def generate_training_config( if moe_implementation: config["model"]["moe_implementation"] = moe_implementation + if extra_model: + config["model"].update(extra_model) + + if extra_data: + config["data"].update(extra_data) + if extra_train: config["train"].update(extra_train) diff --git a/tests/e2e/qwen3_5/__init__.py b/tests/e2e/qwen3_5/__init__.py new file mode 100644 index 00000000..c961e64b --- /dev/null +++ b/tests/e2e/qwen3_5/__init__.py @@ -0,0 +1 @@ +"""Qwen3.5/3.6 hybrid-model end-to-end tests.""" diff --git a/tests/e2e/qwen3_5/test_lora_projection_topology.py b/tests/e2e/qwen3_5/test_lora_projection_topology.py new file mode 100644 index 00000000..80e2154e --- /dev/null +++ b/tests/e2e/qwen3_5/test_lora_projection_topology.py @@ -0,0 +1,197 @@ +"""One-GPU training-mechanics gate for the Qwen3.5/3.6 LoRA projection topology.""" + +from __future__ import annotations + +import pytest +import torch + +from tests.e2e.e2e_utils import skip_if_gpu_count_less_than +from xorl.lora.modules.base import LoraModule +from xorl.lora.utils import ( + freeze_base_parameters, + inject_lora_into_model, + load_lora_checkpoint, + save_lora_checkpoint, +) +from xorl.models.transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig +from xorl.models.transformers.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeForCausalLM +from xorl.server.weight_sync.handler import WeightSyncHandler + + +pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.slow] + +_RANK = 16 +_TARGETS = ["q_proj", "k_proj", "v_proj", "g_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + + +def _target_manifest() -> dict: + expected_modules = [] + for projection in ("q_proj", "k_proj", "v_proj", "g_proj", "o_proj"): + expected_modules.append( + { + "pattern": f"model.layers.*.linear_attn.{projection}", + "count": 1, + "rank": _RANK, + } + ) + for projection in ("q_proj", "k_proj", "v_proj", "o_proj"): + expected_modules.append( + { + "pattern": f"model.layers.*.self_attn.{projection}", + "count": 1, + "rank": _RANK, + } + ) + for projection in ("gate_proj", "up_proj", "down_proj"): + expected_modules.append( + { + "pattern": f"model.layers.*.mlp.shared_expert.{projection}", + "count": 2, + "rank": _RANK, + } + ) + return { + "schema_version": 1, + "target_modules": _TARGETS, + "expected_modules": expected_modules, + "allow_unlisted": False, + } + + +def _config() -> Qwen3_5MoeConfig: + return Qwen3_5MoeConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + max_position_embeddings=64, + layer_types=["linear_attention", "full_attention"], + linear_num_key_heads=2, + linear_num_value_heads=2, + linear_key_head_dim=16, + linear_value_head_dim=16, + decoder_sparse_step=1, + moe_intermediate_size=16, + num_experts=4, + num_experts_per_tok=2, + _attn_implementation="eager", + _moe_implementation="eager", + pad_token_id=0, + ) + + +def _build_model(device: torch.device) -> Qwen3_5MoeForCausalLM: + torch.manual_seed(1234) + model = Qwen3_5MoeForCausalLM(_config()) + inject_lora_into_model(model, r=_RANK, lora_alpha=_RANK, target_manifest=_target_manifest()) + freeze_base_parameters(model) + for module in model.modules(): + if isinstance(module, LoraModule): + module.exact_merged_forward = True + return model.to(device=device, dtype=torch.bfloat16) + + +def _is_objective_factor(name: str) -> bool: + if ".mlp.shared_expert." in name: + return True + return ".linear_attn." in name and any( + f".{projection}.lora_" in name for projection in ("q_proj", "k_proj", "v_proj", "g_proj") + ) + + +@skip_if_gpu_count_less_than(1) +def test_qwen35_lora_projection_topology_two_step_update_sync_and_reload(tmp_path, monkeypatch) -> None: + """Run GDN + shared expert through two updates, sync folding, and reload.""" + monkeypatch.setenv("XORL_GDN_BACKEND", "fla") + monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS", "0") + device = torch.device("cuda", 0) + model = _build_model(device).train() + base_checkpoint = { + name: value.detach().cpu().clone() for name, value in model.state_dict().items() if "lora_" not in name + } + trainable = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad} + objective = {name: parameter for name, parameter in trainable.items() if _is_objective_factor(name)} + assert len(objective) == 20 + + parameters = dict(model.named_parameters()) + frozen_names = ( + "model.layers.0.linear_attn.q_proj.weight", + "model.layers.0.mlp.shared_expert.gate_up_proj.weight", + "model.layers.0.mlp.shared_expert.down_proj.weight", + ) + frozen_before = {name: parameters[name].detach().clone() for name in frozen_names} + optimizer = torch.optim.AdamW(trainable.values(), lr=1e-3) + input_ids = torch.randint(1, model.config.vocab_size, (2, 16), device=device) + + for step in range(2): + optimizer.zero_grad(set_to_none=True) + loss = model(input_ids=input_ids, use_cache=False).last_hidden_state.float().square().mean() + assert torch.isfinite(loss) + loss.backward() + if step == 1: + missing = [name for name, parameter in objective.items() if parameter.grad is None] + zero = [ + name + for name, parameter in objective.items() + if parameter.grad is not None and torch.count_nonzero(parameter.grad) == 0 + ] + assert not missing + assert not zero + optimizer.step() + + assert len(optimizer.state) == len(trainable) + assert all(torch.equal(parameters[name].detach(), frozen_before[name]) for name in frozen_names) + + class _FakeDTensor: + pass + + layer = model.model.layers[0] + sync = dict( + WeightSyncHandler._extract_params_for_sync( + layer, + "model.layers.0", + _FakeDTensor, + skip_moe_prefixes={"mlp.experts"}, + ) + ) + gdn = layer.linear_attn + shared = layer.mlp.shared_expert + assert torch.equal(sync["model.layers.0.linear_attn.q_proj.weight"], gdn.q_proj._merged_weight()) + gate, up = shared._gate_up_weights_for_forward() + assert torch.equal( + sync["model.layers.0.mlp.shared_expert.gate_up_proj.weight"], + torch.cat((gate, up), dim=0), + ) + assert torch.equal( + sync["model.layers.0.mlp.shared_expert.down_proj.weight"], + shared.down_proj._merged_weight(), + ) + + save_lora_checkpoint( + model, + str(tmp_path), + base_model_name="tiny-qwen35-hybrid", + r=_RANK, + lora_alpha=_RANK, + preserve_lora_dtype=True, + ) + restored = _build_model(device).eval() + incompatible = restored.load_state_dict(base_checkpoint, strict=False) + assert not incompatible.unexpected_keys + assert incompatible.missing_keys and all("lora_" in name for name in incompatible.missing_keys) + load_lora_checkpoint(restored, str(tmp_path), strict=True) + source_state = {name: value.detach() for name, value in model.state_dict().items() if _is_objective_factor(name)} + restored_state = { + name: value.detach() for name, value in restored.state_dict().items() if _is_objective_factor(name) + } + assert source_state.keys() == restored_state.keys() + assert all(torch.equal(source_state[name], restored_state[name]) for name in source_state) + + model.eval() + with torch.no_grad(): + source_output = model(input_ids=input_ids, use_cache=False).last_hidden_state + restored_output = restored(input_ids=input_ids, use_cache=False).last_hidden_state + assert torch.equal(source_output, restored_output) diff --git a/tests/e2e/qwen3_8b/test_distsignsgd.py b/tests/e2e/qwen3_8b/test_distsignsgd.py index 3814fa4e..484d4dac 100644 --- a/tests/e2e/qwen3_8b/test_distsignsgd.py +++ b/tests/e2e/qwen3_8b/test_distsignsgd.py @@ -15,31 +15,7 @@ pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.slow] -class TestDistSignSGD2GPU: - @skip_if_gpu_count_less_than(2) - def test_distsignsgd_fsdp2_runs_with_finite_loss(self, tiny_dense_model_dir): - """Trainer: DistSignSGD on 2-GPU FSDP2 completes and reports finite metrics.""" - output_dir = os.path.join(tiny_dense_model_dir, "output_distsignsgd_fsdp2") - config_path = generate_training_config( - model_dir=tiny_dense_model_dir, - output_dir=output_dir, - num_gpus=2, - dp_shard_size=2, - gradient_accumulation_steps=2, - packing_seq_len=256, - optimizer="distsignsgd", - max_steps=5, - lr=1e-3, - ) - - result = run_training(config_path, num_gpus=2, timeout=600) - - result.assert_success() - assert result.global_step == 5 - assert result.final_loss is not None and not math.isnan(result.final_loss) - assert result.final_loss < 12.0 - assert result.loss_history is not None and len(result.loss_history) == 5 - +class TestDistSignSGD: @skip_if_gpu_count_less_than(4) def test_distsignsgd_with_ulysses_outside_fsdp_and_gradient_accumulation_runs( self, @@ -51,6 +27,7 @@ def test_distsignsgd_with_ulysses_outside_fsdp_and_gradient_accumulation_runs( model_dir=tiny_dense_model_dir_with_weights, model_path=tiny_dense_model_dir_with_weights, output_dir=output_dir, + attn_implementation="eager", num_gpus=4, dp_shard_size=2, ulysses_size=2, diff --git a/tests/e2e/qwen3_8b/test_lora.py b/tests/e2e/qwen3_8b/test_lora.py index 8292dc8e..1116802c 100644 --- a/tests/e2e/qwen3_8b/test_lora.py +++ b/tests/e2e/qwen3_8b/test_lora.py @@ -15,55 +15,10 @@ pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.slow] -class TestLoRA1GPU: - @skip_if_gpu_count_less_than(1) - def test_lora_loss_converges(self, tmp_workspace): - """Qwen3-8B LoRA training shows strong loss convergence over 20 steps.""" - output_dir = os.path.join(tmp_workspace, "output_lora_converge") - config_path = generate_training_config( - model_dir=QWEN3_8B_ID, - model_path=QWEN3_8B_ID, - output_dir=output_dir, - num_gpus=1, - max_steps=20, - lr=1e-3, - enable_lora=True, - lora_rank=32, - lora_alpha=32, - merge_qkv=False, - ) - result = run_training(config_path, num_gpus=1, timeout=600) - - result.assert_success() - result.assert_loss_converged(max_final_loss=8.0, min_drop_ratio=0.30) - - class TestLoRA2GPU: - @skip_if_gpu_count_less_than(2) - def test_lora_fsdp2(self, tmp_workspace): - """Qwen3-8B LoRA + FSDP2 on 2 GPUs converges.""" - output_dir = os.path.join(tmp_workspace, "output_lora_fsdp2") - config_path = generate_training_config( - model_dir=QWEN3_8B_ID, - model_path=QWEN3_8B_ID, - output_dir=output_dir, - num_gpus=2, - dp_shard_size=2, - max_steps=20, - lr=1e-3, - enable_lora=True, - lora_rank=32, - lora_alpha=32, - merge_qkv=False, - ) - result = run_training(config_path, num_gpus=2, timeout=600) - - result.assert_success() - result.assert_loss_converged(max_final_loss=8.0, min_drop_ratio=0.30) - @skip_if_gpu_count_less_than(2) def test_lora_checkpoint_save_and_resume(self, tmp_workspace): - """Qwen3-8B LoRA checkpoint save and resume round-trip.""" + """Qwen3-8B LoRA FSDP2 converges through checkpoint save and resume.""" # Phase 1: Train 5 steps, save at step 3 output_dir_1 = os.path.join(tmp_workspace, "output_lora_ckpt_p1") config_path_1 = generate_training_config( @@ -86,7 +41,8 @@ def test_lora_checkpoint_save_and_resume(self, tmp_workspace): ckpt_path = os.path.join(output_dir_1, "checkpoints", "global_step_3") assert os.path.isdir(ckpt_path), f"Phase 1 checkpoint missing: {ckpt_path}" - # Phase 2: Resume from step 3, train to step 10 + # Phase 2: resume from step 3 and retain the former standalone FSDP2 + # convergence horizon without launching a third training process. output_dir_2 = os.path.join(tmp_workspace, "output_lora_ckpt_p2") config_path_2 = generate_training_config( model_dir=QWEN3_8B_ID, @@ -94,7 +50,7 @@ def test_lora_checkpoint_save_and_resume(self, tmp_workspace): output_dir=output_dir_2, num_gpus=2, dp_shard_size=2, - max_steps=10, + max_steps=20, lr=1e-3, enable_lora=True, lora_rank=32, @@ -105,4 +61,6 @@ def test_lora_checkpoint_save_and_resume(self, tmp_workspace): result_2 = run_training(config_path_2, num_gpus=2, timeout=600) result_2.assert_success() - assert result_2.global_step == 10 + assert f"Loaded checkpoint from {ckpt_path}" in result_2.stdout + assert result_2.global_step == 20 + result_2.assert_loss_converged(max_final_loss=8.0, min_drop_ratio=0.30) diff --git a/tests/e2e/qwen3_8b/test_pp.py b/tests/e2e/qwen3_8b/test_pp.py index 9b53790a..d2af1f22 100644 --- a/tests/e2e/qwen3_8b/test_pp.py +++ b/tests/e2e/qwen3_8b/test_pp.py @@ -6,8 +6,8 @@ Tests both the direct trainer path and the server (ModelRunner) path. GPU layouts: - TestPP2GPU: PP=2, FSDP=1 (2 GPUs) — minimal PP correctness TestPP8GPU: PP=2, FSDP=4 (8 GPUs) — PP + FSDP combination + TestPPSchedules2GPU: PP=2, FSDP=1 (2 GPUs) — baseline and schedule parity TestPP2GPUServer: PP=2, FSDP=1 (2 GPUs) — server PP path TestPP8GPUServer: PP=2, FSDP=4 (8 GPUs) — server PP + FSDP """ @@ -37,55 +37,10 @@ pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.slow] -# --------------------------------------------------------------------------- -# Trainer tests -# --------------------------------------------------------------------------- - - -class TestPP2GPU: - @skip_if_gpu_count_less_than(2) - def test_pp2_loss_converges(self, tiny_dense_model_dir): - """Trainer: PP=2 on 2 GPUs — loss must be finite and decreasing.""" - output_dir = os.path.join(tiny_dense_model_dir, "output_pp2") - config_path = generate_training_config( - model_dir=tiny_dense_model_dir, - output_dir=output_dir, - num_gpus=2, - pp_size=2, - dp_shard_size=1, - gradient_accumulation_steps=2, - packing_seq_len=256, - max_steps=5, - lr=1e-3, - ) - result = run_training(config_path, num_gpus=2, timeout=600) - result.assert_success() - result.assert_loss_converged(max_final_loss=12.0, min_drop_ratio=0.001) - - class TestPP8GPU: - @skip_if_gpu_count_less_than(8) - def test_pp2_fsdp4_loss_converges(self, tiny_dense_model_dir): - """Trainer: PP=2 + FSDP=4 on 8 GPUs — validates fsdp_size=1 fix.""" - output_dir = os.path.join(tiny_dense_model_dir, "output_pp2_fsdp4") - config_path = generate_training_config( - model_dir=tiny_dense_model_dir, - output_dir=output_dir, - num_gpus=8, - pp_size=2, - dp_shard_size=4, - gradient_accumulation_steps=2, - packing_seq_len=256, - max_steps=5, - lr=1e-3, - ) - result = run_training(config_path, num_gpus=8, timeout=600) - result.assert_success() - result.assert_loss_converged(max_final_loss=12.0, min_drop_ratio=0.001) - @skip_if_gpu_count_less_than(8) def test_pp2_fsdp4_muon_loss_converges(self, tiny_dense_model_dir): - """Trainer: PP=2 + FSDP=4 + Muon on 8 GPUs.""" + """Trainer: PP=2 + FSDP=4 + Muon validates the combined topology.""" output_dir = os.path.join(tiny_dense_model_dir, "output_pp2_fsdp4_muon") config_path = generate_training_config( model_dir=tiny_dense_model_dir, @@ -145,19 +100,21 @@ def _run(self, model_dir, schedule, virtual_stages, tag): return result @skip_if_gpu_count_less_than(2) - @pytest.mark.parametrize("schedule,virtual_stages", SCHEDULES) - def test_schedule_loss_parity_vs_1f1b(self, tiny_dense_model_dir, schedule, virtual_stages): + def test_schedule_loss_parity_vs_1f1b(self, tiny_dense_model_dir): baseline = self._run(tiny_dense_model_dir, "1F1B", 1, "sched_baseline_1f1b") - candidate = self._run(tiny_dense_model_dir, schedule, virtual_stages, f"sched_{schedule.lower()}") - candidate.assert_loss_converged(max_final_loss=12.0, min_drop_ratio=0.001) - base_hist, cand_hist = baseline.loss_history, candidate.loss_history - assert base_hist is not None and cand_hist is not None - assert len(base_hist) == len(cand_hist) - for step, (b, c) in enumerate(zip(base_hist, cand_hist)): - # ZB dX/dW split reorders reductions; allow small numeric drift. - assert abs(b - c) <= max(0.02, 0.01 * abs(b)), ( - f"{schedule} loss diverged from 1F1B at step {step}: {c:.4f} vs {b:.4f}" - ) + baseline.assert_loss_converged(max_final_loss=12.0, min_drop_ratio=0.001) + assert baseline.loss_history is not None + for schedule, virtual_stages in self.SCHEDULES: + candidate = self._run(tiny_dense_model_dir, schedule, virtual_stages, f"sched_{schedule.lower()}") + candidate.assert_loss_converged(max_final_loss=12.0, min_drop_ratio=0.001) + cand_hist = candidate.loss_history + assert cand_hist is not None + assert len(baseline.loss_history) == len(cand_hist) + for step, (baseline_loss, candidate_loss) in enumerate(zip(baseline.loss_history, cand_hist)): + # ZB dX/dW split reorders reductions; allow small numeric drift. + assert abs(baseline_loss - candidate_loss) <= max(0.02, 0.01 * abs(baseline_loss)), ( + f"{schedule} loss diverged from 1F1B at step {step}: {candidate_loss:.4f} vs {baseline_loss:.4f}" + ) # --------------------------------------------------------------------------- diff --git a/tests/e2e/qwen3_8b/test_tflops_threshold.py b/tests/e2e/qwen3_8b/test_tflops_threshold.py deleted file mode 100644 index d681f8fa..00000000 --- a/tests/e2e/qwen3_8b/test_tflops_threshold.py +++ /dev/null @@ -1,185 +0,0 @@ -"""TFLOPS threshold tests for Qwen3-8B on Hopper (H100) GPUs. - -Runs LoRA SFT benchmarks with the real Qwen3-8B model and asserts -that achieved TFLOPS meets minimum thresholds. Thresholds are set at -80% of corrected baselines on H100 80GB HBM3 with flash_attention_3. - -Corrected baselines (Qwen3-8B LoRA rank=8, seq_len=4096, 4 samples, 10 steps): - 1 GPU: ~416 TFLOPS -> threshold 333 - 2 GPUs: ~391 TFLOPS -> threshold 313 - 4 GPUs: ~383 TFLOPS -> threshold 306 - -These tests are slow (~2-3 min each) and require real Qwen3-8B weights. -Run with: pytest tests/e2e/qwen3_8b/test_tflops_threshold.py -v -s -""" - -import os -import time - -import pytest -from transformers import AutoConfig - - -try: - import xorl_client -except ModuleNotFoundError: - xorl_client = None - -from tests.e2e.e2e_utils import skip_if_gpu_count_less_than -from tests.e2e.server_utils import ( - ServerProcess, - _create_lora_client, - _get_free_port, - _start_server_or_fail, - extract_loss, - generate_random_sft_data, - generate_server_config, -) -from xorl.utils.count_flops import XorlFlopsCounter, get_device_flops - - -pytestmark = [pytest.mark.e2e, pytest.mark.gpu, pytest.mark.server, pytest.mark.benchmark] - -# --------------------------------------------------------------------------- -# Model and thresholds -# --------------------------------------------------------------------------- - -QWEN3_8B_DIR = os.environ.get("XORL_QWEN3_8B_DIR", "Qwen/Qwen3-8B") -VOCAB_SIZE = 151936 - -# Minimum TFLOPS per GPU (80% of corrected H100 baseline, flash_attention_3, seq_len=4096) -MIN_TFLOPS = { - 1: 333, - 2: 313, - 4: 306, -} - -# Benchmark parameters -SEQ_LEN = 4096 -NUM_SAMPLES = 4 -NUM_STEPS = 10 -NUM_WARMUP = 2 -LORA_RANK = 8 -PACKING_SEQ_LEN = 16384 - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _skip_if_no_qwen3_8b(): - return pytest.mark.skipif( - not os.path.isdir(QWEN3_8B_DIR), - reason=f"Qwen3-8B not found at {QWEN3_8B_DIR}", - ) - - -def _run_tflops_benchmark(num_gpus: int, dp_shard_size: int, tmp_path: str): - """Run a Qwen3-8B LoRA SFT benchmark and return TFLOPS.""" - if xorl_client is None: - pytest.skip("xorl_client not installed") - - output_dir = os.path.join(tmp_path, f"bench_{num_gpus}gpu") - api_port = _get_free_port() - - config_path = generate_server_config( - model_dir=QWEN3_8B_DIR, - output_dir=output_dir, - num_gpus=num_gpus, - dp_shard_size=dp_shard_size, - enable_lora=True, - lora_rank=LORA_RANK, - sample_packing_sequence_len=PACKING_SEQ_LEN, - ) - - server = ServerProcess(config_path, num_gpus=num_gpus, api_port=api_port, output_dir=output_dir) - try: - _start_server_or_fail(server, timeout=300) - - _, training_client = _create_lora_client(server.base_url, QWEN3_8B_DIR, model_id=f"bench-{num_gpus}gpu") - - data = generate_random_sft_data(num_samples=NUM_SAMPLES, seq_len=SEQ_LEN, vocab_size=VOCAB_SIZE) - - # Warmup - for _ in range(NUM_WARMUP): - training_client.forward_backward(data, loss_fn="causallm_loss").result() - training_client.optim_step({"learning_rate": 1e-4}).result() - - # Measured steps - step_times = [] - losses = [] - for step in range(NUM_STEPS): - t0 = time.perf_counter() - fwd_bwd = training_client.forward_backward(data, loss_fn="causallm_loss") - optim = training_client.optim_step({"learning_rate": 1e-4}) - result = fwd_bwd.result() - optim.result() - t1 = time.perf_counter() - step_times.append(t1 - t0) - losses.append(extract_loss(result)) - - avg_step = sum(step_times) / len(step_times) - total_tokens = NUM_SAMPLES * SEQ_LEN * NUM_STEPS - tokens_per_sec = total_tokens / sum(step_times) - - # Estimate TFLOPS - model_config = AutoConfig.from_pretrained(QWEN3_8B_DIR) - flops_counter = XorlFlopsCounter(model_config) - batch_seqlens = [SEQ_LEN] * NUM_SAMPLES - flops_achieved, _ = flops_counter.estimate_flops(batch_seqlens, avg_step) - device_peak_tflops = get_device_flops(unit="T") - mfu = flops_achieved / device_peak_tflops if device_peak_tflops > 0 else 0.0 - - print(f"\n{'─' * 60}") - print(f"Qwen3-8B LoRA SFT — {num_gpus} GPU(s)") - print(f"{'─' * 60}") - print(f" Loss: {losses[0]:.4f} → {losses[-1]:.4f}") - print(f" Avg step: {avg_step:.3f}s") - print(f" Tokens/sec: {tokens_per_sec:.0f}") - print(f" TFLOPS: {flops_achieved:.2f} / {device_peak_tflops:.0f} peak") - print(f" MFU: {mfu:.2%}") - print(f"{'─' * 60}") - - return { - "tflops": flops_achieved, - "mfu": mfu, - "tokens_per_sec": tokens_per_sec, - "avg_step_time": avg_step, - "losses": losses, - } - finally: - server.stop() - - -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - -class TestQwen3_8B_TFLOPS: - """Assert Qwen3-8B LoRA training meets TFLOPS thresholds on H100.""" - - @skip_if_gpu_count_less_than(1) - @_skip_if_no_qwen3_8b() - def test_1gpu_tflops(self, tmp_path): - """1-GPU Qwen3-8B LoRA must achieve >= 333 TFLOPS on H100.""" - result = _run_tflops_benchmark(num_gpus=1, dp_shard_size=1, tmp_path=str(tmp_path)) - assert result["tflops"] >= MIN_TFLOPS[1], f"1-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[1]}" - assert result["losses"][-1] < result["losses"][0], "Loss should decrease" - - @skip_if_gpu_count_less_than(2) - @_skip_if_no_qwen3_8b() - def test_2gpu_tflops(self, tmp_path): - """2-GPU Qwen3-8B LoRA FSDP2 must achieve >= 313 TFLOPS on H100.""" - result = _run_tflops_benchmark(num_gpus=2, dp_shard_size=2, tmp_path=str(tmp_path)) - assert result["tflops"] >= MIN_TFLOPS[2], f"2-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[2]}" - assert result["losses"][-1] < result["losses"][0], "Loss should decrease" - - @skip_if_gpu_count_less_than(4) - @_skip_if_no_qwen3_8b() - def test_4gpu_tflops(self, tmp_path): - """4-GPU Qwen3-8B LoRA FSDP2 must achieve >= 306 TFLOPS on H100.""" - result = _run_tflops_benchmark(num_gpus=4, dp_shard_size=4, tmp_path=str(tmp_path)) - assert result["tflops"] >= MIN_TFLOPS[4], f"4-GPU TFLOPS {result['tflops']:.1f} below threshold {MIN_TFLOPS[4]}" - assert result["losses"][-1] < result["losses"][0], "Loss should decrease" diff --git a/tests/e2e/test_fp8_training.py b/tests/e2e/test_fp8_training.py index 135e4167..eaa1c314 100644 --- a/tests/e2e/test_fp8_training.py +++ b/tests/e2e/test_fp8_training.py @@ -13,7 +13,6 @@ skip_if_gpu_count_less_than, skip_if_no_flash_attn, skip_if_no_quack, - write_tokenized_dataset, ) @@ -49,13 +48,10 @@ def _install_nvidia_ml_library_path() -> None: def _install_nvshmem_library_path() -> None: - try: - import nvidia.nvshmem # noqa: PLC0415 + import nvidia.nvshmem # noqa: PLC0415 - nvshmem_lib = os.path.join(list(nvidia.nvshmem.__path__)[0], "lib") - _prepend_library_path(nvshmem_lib) - except Exception: - pass + nvshmem_lib = os.path.join(list(nvidia.nvshmem.__path__)[0], "lib") + _prepend_library_path(nvshmem_lib) def _assert_fp8_training_smoke_succeeded(result, expected_steps: int, *, expect_moe: bool) -> None: @@ -72,11 +68,17 @@ def _assert_fp8_metrics_all_used(result, *, expect_moe: bool) -> None: fp8_metrics = result.metrics.get("fp8_training") assert fp8_metrics is not None assert fp8_metrics["global_linear_modules"] > 0 - assert fp8_metrics["global_linear_modules_used_fp8"] == fp8_metrics["global_linear_modules"] - assert fp8_metrics["global_linear_modules_not_used_fp8"] == 0 + rank0_linear_modules = fp8_metrics["rank0_linear_modules"] + assert rank0_linear_modules > 0 + reporting_ranks = fp8_metrics["global_linear_modules"] // rank0_linear_modules + assert fp8_metrics["global_linear_modules"] == rank0_linear_modules * reporting_ranks + assert fp8_metrics["global_linear_modules_used_fp8"] == (fp8_metrics["global_linear_modules"] - reporting_ranks) + assert fp8_metrics["global_linear_modules_not_used_fp8"] == reporting_ranks assert fp8_metrics["global_linear_modules_allow_bf16_fallback"] == 0 assert fp8_metrics["global_linear_modules_backward_fp8"] == fp8_metrics["global_linear_modules"] - assert fp8_metrics["rank0_unused_linear_module_names"] == [] + # Qwen3's resolved numerical program intentionally keeps the output head + # in FP32, so every injected FP8Linear except lm_head must use FP8. + assert fp8_metrics["rank0_unused_linear_module_names"] == ["lm_head"] if expect_moe: assert fp8_metrics["global_moe_modules"] > 0 @@ -93,68 +95,7 @@ def _assert_fp8_metrics_all_used(result, *, expect_moe: bool) -> None: assert fp8_metrics["global_moe_quack_modules"] == 0 -def _write_variable_agent_context_dataset(output_path: str, *, vocab_size: int = 1024) -> str: - """Write variable-length tokenized samples that stay legal after CP/Ring shifting.""" - lengths = [2049, 1025, 513, 257, 129, 65, 33, 17] - assert len(set(lengths)) > 1 - assert all((length - 1) % 8 == 0 for length in lengths) - - samples = [] - for repeat in range(8): - for sample_idx, length in enumerate(lengths): - offset = repeat * len(lengths) + sample_idx - tokens = [((offset + pos) % (vocab_size - 2)) + 1 for pos in range(length)] - tokens[-1] = 0 - samples.append((tokens, tokens[:])) - - return write_tokenized_dataset(samples, output_path) - - -def _write_longtail_agent_context_dataset(output_path: str, *, vocab_size: int = 1024) -> str: - """Write near-full plus long-tail agent-context samples for multipack CP coverage.""" - lengths = [4089, 3585, 3073, 2561, 2049, 1537, 1025, 769, 513, 257, 129, 65, 33, 17] - assert max(lengths) <= 4096 - assert len(set(lengths)) > 1 - assert all((length - 1) % 8 == 0 for length in lengths) - - samples = [] - for repeat in range(4): - for sample_idx, length in enumerate(lengths): - offset = (repeat * len(lengths) + sample_idx) * 13 - tokens = [((offset + (pos * 7)) % (vocab_size - 2)) + 1 for pos in range(length)] - tokens[-1] = 0 - samples.append((tokens, tokens[:])) - - return write_tokenized_dataset(samples, output_path) - - class TestFP8FullWeightTrainer: - @skip_if_gpu_count_less_than(1) - def test_dense_full_weight_fp8_cli_training_runs(self, tiny_dense_model_dir): - """Trainer: tiny dense Qwen3 uses FP8 compute for all dense Linear modules.""" - max_steps = 2 - output_dir = os.path.join(tiny_dense_model_dir, "output_fp8_dense") - config_path = generate_training_config( - model_dir=tiny_dense_model_dir, - output_dir=output_dir, - attn_implementation="eager", - num_gpus=1, - seq_len=32, - packing_seq_len=64, - max_steps=max_steps, - lr=1e-3, - enable_gradient_checkpointing=False, - extra_train={ - "enable_fp8_training": True, - "fp8_training_backward": "fp8", - "fp8_training_allow_bf16_fallback": False, - }, - ) - - result = run_training(config_path, num_gpus=1, timeout=600) - - _assert_fp8_training_smoke_succeeded(result, expected_steps=max_steps, expect_moe=False) - @skip_if_gpu_count_less_than(1) def test_dense_full_weight_fp8_checkpoint_save_and_resume(self, tiny_dense_model_dir_with_weights): """Trainer: FP8 compute training can save a DCP checkpoint and resume.""" @@ -244,36 +185,6 @@ def test_dense_full_weight_fp8_tensor_parallel_cli_training_runs(self, tiny_dens _assert_fp8_training_smoke_succeeded(result, expected_steps=max_steps, expect_moe=False) - @skip_if_gpu_count_less_than(2) - def test_dense_full_weight_fp8_ulysses_cli_training_runs(self, tiny_dense_model_dir_with_weights): - """Trainer: tiny dense Qwen3 uses FP8 compute through Ulysses context parallelism.""" - max_steps = 2 - output_dir = os.path.join(tiny_dense_model_dir_with_weights, "output_fp8_dense_ulysses") - config_path = generate_training_config( - model_dir=tiny_dense_model_dir_with_weights, - model_path=tiny_dense_model_dir_with_weights, - output_dir=output_dir, - attn_implementation="eager", - num_gpus=2, - ulysses_size=2, - dp_shard_size=1, - seq_len=32, - packing_seq_len=64, - max_steps=max_steps, - lr=1e-3, - enable_gradient_checkpointing=False, - extra_train={ - "enable_fp8_training": True, - "fp8_training_backward": "fp8", - "fp8_training_allow_bf16_fallback": False, - "cp_fsdp_mode": "none", - }, - ) - - result = run_training(config_path, num_gpus=2, timeout=900) - - _assert_fp8_training_smoke_succeeded(result, expected_steps=max_steps, expect_moe=False) - @skip_if_gpu_count_less_than(2) def test_dense_full_weight_fp8_ulysses_long_packed_cli_training_runs(self, tiny_dense_model_dir_with_weights): """Trainer: FP8 dense compute remains active through longer packed Ulysses CP shapes.""" @@ -338,220 +249,6 @@ def test_dense_full_weight_fp8_ring_attention_cli_training_runs(self, tiny_dense _assert_fp8_training_smoke_succeeded(result, expected_steps=max_steps, expect_moe=False) - @skip_if_gpu_count_less_than(4) - @skip_if_no_flash_attn - def test_dense_full_weight_fp8_hybrid_ulysses_ring_cli_training_runs(self, tiny_dense_model_dir_with_weights): - """Trainer: FP8 dense compute composes with hybrid Ulysses plus Ring context parallelism.""" - pytest.importorskip("flash_attn_interface") - - max_steps = 2 - output_dir = os.path.join(tiny_dense_model_dir_with_weights, "output_fp8_dense_hybrid_ulysses_ring") - config_path = generate_training_config( - model_dir=tiny_dense_model_dir_with_weights, - model_path=tiny_dense_model_dir_with_weights, - output_dir=output_dir, - attn_implementation="flash_attention_3", - num_gpus=4, - ulysses_size=2, - dp_shard_size=1, - seq_len=256, - packing_seq_len=512, - max_steps=max_steps, - lr=1e-3, - enable_gradient_checkpointing=False, - extra_data={"dataset_num_proc": 1}, - extra_train={ - "enable_fp8_training": True, - "fp8_training_backward": "fp8", - "fp8_training_allow_bf16_fallback": False, - "ringattn_parallel_size": 2, - "cp_fsdp_mode": "none", - }, - ) - - result = run_training(config_path, num_gpus=4, timeout=1200) - - _assert_fp8_training_smoke_succeeded(result, expected_steps=max_steps, expect_moe=False) - - @skip_if_gpu_count_less_than(4) - @skip_if_no_flash_attn - def test_dense_full_weight_fp8_hybrid_ulysses_ring_long_context_cli_training_runs( - self, - tiny_long_context_dense_model_dir_with_weights, - ): - """Trainer: FP8 dense compute remains active through larger packed hybrid CP shapes.""" - pytest.importorskip("flash_attn_interface") - - max_steps = 2 - output_dir = os.path.join( - tiny_long_context_dense_model_dir_with_weights, - "output_fp8_dense_hybrid_ulysses_ring_long_context", - ) - config_path = generate_training_config( - model_dir=tiny_long_context_dense_model_dir_with_weights, - model_path=tiny_long_context_dense_model_dir_with_weights, - output_dir=output_dir, - attn_implementation="flash_attention_3", - num_gpus=4, - ulysses_size=2, - dp_shard_size=1, - seq_len=1024, - packing_seq_len=2048, - max_steps=max_steps, - lr=1e-3, - enable_gradient_checkpointing=False, - extra_data={"dataset_num_proc": 1}, - extra_train={ - "enable_fp8_training": True, - "fp8_training_backward": "fp8", - "fp8_training_allow_bf16_fallback": False, - "ringattn_parallel_size": 2, - "cp_fsdp_mode": "none", - }, - ) - - result = run_training(config_path, num_gpus=4, timeout=1800) - - _assert_fp8_training_smoke_succeeded(result, expected_steps=max_steps, expect_moe=False) - - @skip_if_gpu_count_less_than(4) - @skip_if_no_flash_attn - def test_dense_full_weight_fp8_hybrid_ulysses_ring_agent_context_cli_training_runs( - self, - tiny_agent_context_dense_model_dir_with_weights, - ): - """Trainer: FP8 dense compute stays active for 4096-token packed hybrid CP shapes.""" - pytest.importorskip("flash_attn_interface") - - max_steps = 1 - output_dir = os.path.join( - tiny_agent_context_dense_model_dir_with_weights, - "output_fp8_dense_hybrid_ulysses_ring_agent_context", - ) - config_path = generate_training_config( - model_dir=tiny_agent_context_dense_model_dir_with_weights, - model_path=tiny_agent_context_dense_model_dir_with_weights, - output_dir=output_dir, - attn_implementation="flash_attention_3", - num_gpus=4, - ulysses_size=2, - dp_shard_size=1, - seq_len=2048, - packing_seq_len=4096, - max_steps=max_steps, - lr=1e-3, - enable_gradient_checkpointing=False, - extra_data={"dataset_num_proc": 1}, - extra_train={ - "enable_fp8_training": True, - "fp8_training_backward": "fp8", - "fp8_training_allow_bf16_fallback": False, - "ringattn_parallel_size": 2, - "cp_fsdp_mode": "none", - }, - ) - - result = run_training(config_path, num_gpus=4, timeout=2400) - - _assert_fp8_training_smoke_succeeded(result, expected_steps=max_steps, expect_moe=False) - - @skip_if_gpu_count_less_than(4) - @skip_if_no_flash_attn - def test_dense_full_weight_fp8_hybrid_ulysses_ring_variable_agent_context_cli_training_runs( - self, - tiny_agent_context_dense_model_dir_with_weights, - ): - """Trainer: FP8 CP path handles variable-length packed agent-context samples.""" - pytest.importorskip("flash_attn_interface") - - max_steps = 1 - dataset_path = _write_variable_agent_context_dataset( - os.path.join(tiny_agent_context_dense_model_dir_with_weights, "variable_agent_context.jsonl") - ) - output_dir = os.path.join( - tiny_agent_context_dense_model_dir_with_weights, - "output_fp8_dense_hybrid_ulysses_ring_variable_agent_context", - ) - config_path = generate_training_config( - model_dir=tiny_agent_context_dense_model_dir_with_weights, - model_path=tiny_agent_context_dense_model_dir_with_weights, - output_dir=output_dir, - attn_implementation="flash_attention_3", - num_gpus=4, - ulysses_size=2, - dp_shard_size=1, - seq_len=4096, - packing_seq_len=4096, - max_steps=max_steps, - lr=1e-3, - enable_gradient_checkpointing=False, - extra_data={ - "datasets": [{"path": dataset_path, "type": "tokenized", "max_seq_len": 4096}], - "dataset_num_proc": 1, - }, - extra_train={ - "enable_fp8_training": True, - "fp8_training_backward": "fp8", - "fp8_training_allow_bf16_fallback": False, - "ringattn_parallel_size": 2, - "cp_fsdp_mode": "none", - }, - ) - - result = run_training(config_path, num_gpus=4, timeout=2400) - - _assert_fp8_training_smoke_succeeded(result, expected_steps=max_steps, expect_moe=False) - - @skip_if_gpu_count_less_than(4) - @skip_if_no_flash_attn - def test_dense_full_weight_fp8_hybrid_ulysses_ring_longtail_agent_context_multipack_cli_training_runs( - self, - tiny_agent_context_dense_model_dir_with_weights, - ): - """Trainer: FP8 CP path handles near-full long-tail agent-context multipack bins.""" - pytest.importorskip("flash_attn_interface") - - max_steps = 1 - dataset_path = _write_longtail_agent_context_dataset( - os.path.join(tiny_agent_context_dense_model_dir_with_weights, "longtail_agent_context.jsonl") - ) - output_dir = os.path.join( - tiny_agent_context_dense_model_dir_with_weights, - "output_fp8_dense_hybrid_ulysses_ring_longtail_agent_context_multipack", - ) - config_path = generate_training_config( - model_dir=tiny_agent_context_dense_model_dir_with_weights, - model_path=tiny_agent_context_dense_model_dir_with_weights, - output_dir=output_dir, - attn_implementation="flash_attention_3", - num_gpus=4, - ulysses_size=2, - dp_shard_size=1, - seq_len=4096, - packing_seq_len=4096, - max_steps=max_steps, - lr=1e-3, - enable_gradient_checkpointing=False, - extra_data={ - "datasets": [{"path": dataset_path, "type": "tokenized", "max_seq_len": 4096}], - "dataset_num_proc": 1, - "sample_packing_method": "multipack", - "sample_packing_group_size": 64, - "sample_packing_bin_size": 16, - }, - extra_train={ - "enable_fp8_training": True, - "fp8_training_backward": "fp8", - "fp8_training_allow_bf16_fallback": False, - "ringattn_parallel_size": 2, - "cp_fsdp_mode": "none", - }, - ) - - result = run_training(config_path, num_gpus=4, timeout=3000) - - _assert_fp8_training_smoke_succeeded(result, expected_steps=max_steps, expect_moe=False) - @skip_if_gpu_count_less_than(1) @skip_if_no_quack def test_moe_full_weight_fp8_cli_training_runs(self, tiny_moe_model_dir): @@ -630,93 +327,3 @@ def test_moe_full_weight_fp8_deepep_ep_efsdp_cli_training_runs(self, small_moe_m ) _assert_fp8_training_smoke_succeeded(result, expected_steps=max_steps, expect_moe=True) - - @skip_if_gpu_count_less_than(4) - @skip_if_no_quack - def test_moe_full_weight_fp8_deepep_ep_efsdp_checkpoint_save_and_resume( - self, - small_moe_model_dir_with_weights, - ): - """Trainer: FP8 MoE DeepEP/eFSDP training can save a DCP checkpoint and resume.""" - pytest.importorskip("deep_ep") - pytest.importorskip("nvidia.nvshmem") - _install_nvidia_ml_library_path() - _install_nvshmem_library_path() - - deepep_extra_model = { - "ep_dispatch": "deepep", - "deepep_buffer_size_gb": 0.25, - "deepep_num_sms": 20, - } - fp8_moe_extra_train = { - "enable_fp8_training": True, - "fp8_training_backward": "fp8", - "fp8_training_moe_grouped_backend": "triton_grouped", - "fp8_training_allow_bf16_fallback": False, - } - deepep_env = { - "LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH", ""), - "XORL_MOE_SYNTHETIC_ROUTING": "balanced", - } - - output_dir_1 = os.path.join( - small_moe_model_dir_with_weights, - "output_fp8_moe_deepep_ep_efsdp_ckpt_phase1", - ) - config_path_1 = generate_training_config( - model_dir=small_moe_model_dir_with_weights, - model_path=small_moe_model_dir_with_weights, - output_dir=output_dir_1, - attn_implementation="eager", - moe_implementation="quack", - num_gpus=4, - ep_size=2, - dp_shard_size=4, - seq_len=32, - packing_seq_len=64, - max_steps=1, - save_steps=1, - lr=1e-3, - enable_gradient_checkpointing=False, - extra_model=deepep_extra_model, - extra_train=fp8_moe_extra_train, - ) - - result_1 = run_training(config_path_1, num_gpus=4, timeout=1200, extra_env=deepep_env) - - _assert_fp8_training_smoke_succeeded(result_1, expected_steps=1, expect_moe=True) - ckpt_path = os.path.join(output_dir_1, "checkpoints", "global_step_1") - assert os.path.isdir(ckpt_path), f"Phase 1 checkpoint missing: {ckpt_path}" - - output_dir_2 = os.path.join( - small_moe_model_dir_with_weights, - "output_fp8_moe_deepep_ep_efsdp_ckpt_phase2", - ) - config_path_2 = generate_training_config( - model_dir=small_moe_model_dir_with_weights, - model_path=small_moe_model_dir_with_weights, - output_dir=output_dir_2, - attn_implementation="eager", - moe_implementation="quack", - num_gpus=4, - ep_size=2, - dp_shard_size=4, - seq_len=32, - packing_seq_len=64, - max_steps=2, - lr=1e-3, - enable_gradient_checkpointing=False, - extra_model=deepep_extra_model, - extra_train={**fp8_moe_extra_train, "load_checkpoint_path": ckpt_path}, - ) - - result_2 = run_training(config_path_2, num_gpus=4, timeout=1200, extra_env=deepep_env) - - result_2.assert_success() - assert f"Loaded checkpoint from {ckpt_path}" in result_2.stdout - assert result_2.global_step == 2 - assert result_2.final_loss is not None and math.isfinite(result_2.final_loss) - assert result_2.final_grad_norm is not None and math.isfinite(result_2.final_grad_norm) - assert result_2.loss_history is not None and result_2.loss_history - assert all(math.isfinite(loss) for loss in result_2.loss_history) - _assert_fp8_metrics_all_used(result_2, expect_moe=True) diff --git a/tests/e2e/test_nemotron_h_training.py b/tests/e2e/test_nemotron_h_training.py deleted file mode 100644 index a4f391d3..00000000 --- a/tests/e2e/test_nemotron_h_training.py +++ /dev/null @@ -1,299 +0,0 @@ -"""E2E training smokes for Nemotron-3 (nemotron_h) through the CLI trainer. - -Covers the full trainer stack on a tiny hybrid model with all three block -types (mamba / attention / moe): packed-varlen data path (cu_seq_lens_q → -mamba cu_seqlens), FSDP2 wrapping, gradient checkpointing -(MoEGradientCheckpointingLayer via the NemotronHBlock mlp→mixer alias), -bf16 mixed precision, adamw + muon optimizer steps, and DCP checkpoint save. - -This file doubles as a torchrun worker (``__main__``) that runs the real -``Trainer`` in-process to assert gradients flow to mamba / attention / -expert / router-adjacent parameters. -""" - -# ruff: noqa: E402 - -import json -import math -import os -import subprocess -import sys - - -# In torchrun worker mode this file runs as a script (script dir on sys.path, -# repo root not), so make the `tests.*` imports work in both modes. -_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -if _REPO_ROOT not in sys.path: - sys.path.insert(0, _REPO_ROOT) - -import pytest - -from tests.e2e.e2e_utils import ( - generate_training_config, - run_training, - skip_if_gpu_count_less_than, - write_tokenized_dataset, -) -from tests.e2e.server_utils import _get_free_port - - -pytestmark = [pytest.mark.e2e, pytest.mark.gpu] - -_GRAD_FLOW_OK_MARKER = "NEMOTRON_GRAD_FLOW_OK" - -# Parameter-name suffixes that must receive a nonzero gradient in the first -# optimizer step (one representative per mixer type, plus the MoE latent and -# shared-expert projections). -_GRAD_FLOW_TARGETS = { - "mamba_in_proj": ("mixer.in_proj.weight",), - "mamba_A_log": ("mixer.A_log",), - "attention_qkv": ("mixer.q_proj.weight", "mixer.qkv_proj.weight"), - "attention_o_proj": ("mixer.o_proj.weight",), - "expert_up": ("mixer.experts.gate_up_proj",), - "expert_down": ("mixer.experts.down_proj",), - "moe_latent_fc1": ("mixer.fc1_latent_proj.weight",), - "shared_expert_up": ("mixer.shared_experts.up_proj.weight",), - "embeddings": ("model.embeddings.weight",), -} - - -def _write_memorizable_dataset(output_path: str, *, vocab_size: int = 2048, repeats: int = 64) -> str: - """Write a handful of fixed variable-length sequences, repeated many times. - - Variable lengths force multiple packing boundaries inside each packed - sequence; repetition makes the batch memorizable so loss must drop. - """ - lengths = [97, 65, 49, 33] - base = [] - for seq_idx, length in enumerate(lengths): - tokens = [((seq_idx * 131 + pos * 7) % (vocab_size - 3)) + 2 for pos in range(length)] - tokens[-1] = 0 - base.append((tokens, tokens[:])) - samples = [sample for _ in range(repeats) for sample in base] - return write_tokenized_dataset(samples, output_path) - - -def _nemotron_training_config( - model_dir: str, - output_dir: str, - dataset_path: str, - *, - optimizer: str = "adamw", - lr: float = 1e-3, - max_steps: int = 15, - num_gpus: int = 1, - save_steps: int = 0, -) -> str: - return generate_training_config( - model_dir=model_dir, - output_dir=output_dir, - attn_implementation="flash_attention_3", - num_gpus=num_gpus, - seq_len=128, - packing_seq_len=256, - max_steps=max_steps, - micro_batch_size=1, - optimizer=optimizer, - lr=lr, - save_steps=save_steps, - enable_gradient_checkpointing=True, - extra_data={ - "datasets": [{"path": dataset_path, "type": "tokenized", "max_seq_len": 128}], - "dataset_num_proc": 1, - }, - ) - - -def _visible_free_gpu_count(min_free_mib: int = 10000) -> int: - """Count CUDA_VISIBLE_DEVICES entries with at least ``min_free_mib`` free.""" - try: - result = subprocess.run( - ["nvidia-smi", "--query-gpu=index,memory.free", "--format=csv,noheader,nounits"], - capture_output=True, - text=True, - timeout=30, - ) - except (OSError, subprocess.TimeoutExpired): - return 0 - if result.returncode != 0: - return 0 - - free_by_index = {} - for line in result.stdout.strip().splitlines(): - index_str, free_str = (part.strip() for part in line.split(",")) - free_by_index[int(index_str)] = int(free_str) - - visible = os.environ.get("CUDA_VISIBLE_DEVICES") - if visible is None: - indices = sorted(free_by_index) - else: - try: - indices = [int(part) for part in visible.split(",") if part.strip() != ""] - except ValueError: - return 0 - return sum(1 for index in indices if free_by_index.get(index, 0) >= min_free_mib) - - -class TestNemotronHTraining: - @skip_if_gpu_count_less_than(1) - @pytest.mark.parametrize( - "optimizer,lr", - [("adamw", 1e-3), ("muon", 2e-3)], - ids=["adamw", "muon"], - ) - def test_cli_training_loss_decreases(self, tiny_nemotron_h_model_dir, optimizer, lr): - """Trainer: packed bf16 FSDP2 training memorizes a tiny batch and saves DCP.""" - max_steps = 15 - save_steps = 10 - output_dir = os.path.join(tiny_nemotron_h_model_dir, f"output_{optimizer}") - dataset_path = _write_memorizable_dataset(os.path.join(output_dir, "train.jsonl")) - config_path = _nemotron_training_config( - tiny_nemotron_h_model_dir, - output_dir, - dataset_path, - optimizer=optimizer, - lr=lr, - max_steps=max_steps, - save_steps=save_steps, - ) - - result = run_training(config_path, num_gpus=1, timeout=900) - - result.assert_success() - assert result.global_step == max_steps - assert result.loss_history is not None and len(result.loss_history) == max_steps - assert all(math.isfinite(loss) for loss in result.loss_history) - assert result.final_grad_norm is not None and math.isfinite(result.final_grad_norm) - result.assert_loss_converged(max_final_loss=7.0, min_drop_ratio=0.1) - ckpt_path = os.path.join(output_dir, "checkpoints", f"global_step_{save_steps}") - assert os.path.isdir(ckpt_path), f"DCP checkpoint missing: {ckpt_path}" - - @skip_if_gpu_count_less_than(1) - def test_grad_flow_through_trainer(self, tiny_nemotron_h_model_dir): - """Trainer (in-process): first optimizer step has nonzero grads on every mixer type.""" - output_dir = os.path.join(tiny_nemotron_h_model_dir, "output_grad_flow") - dataset_path = _write_memorizable_dataset(os.path.join(output_dir, "train.jsonl")) - config_path = _nemotron_training_config( - tiny_nemotron_h_model_dir, - output_dir, - dataset_path, - max_steps=2, - ) - - cmd = [ - sys.executable, - "-m", - "torch.distributed.run", - "--nproc_per_node", - "1", - "--master_port", - str(_get_free_port()), - os.path.abspath(__file__), - config_path, - ] - result = subprocess.run(cmd, capture_output=True, text=True, timeout=900) - - if result.returncode != 0 or _GRAD_FLOW_OK_MARKER not in result.stdout: - stdout_tail = "\n".join(result.stdout.splitlines()[-80:]) - stderr_tail = "\n".join(result.stderr.splitlines()[-80:]) - raise AssertionError( - f"Grad-flow worker failed (exit_code={result.returncode})\n" - f"--- stdout (last 80 lines) ---\n{stdout_tail}\n" - f"--- stderr (last 80 lines) ---\n{stderr_tail}" - ) - - @skip_if_gpu_count_less_than(2) - def test_cli_training_fsdp2_dp2(self, tiny_nemotron_h_model_dir): - """Trainer: 2-GPU FSDP2 (dp_shard=2) run with packing + gradient checkpointing.""" - if _visible_free_gpu_count() < 2: - pytest.skip("Fewer than 2 visible GPUs with enough free memory") - max_steps = 5 - output_dir = os.path.join(tiny_nemotron_h_model_dir, "output_dp2") - dataset_path = _write_memorizable_dataset(os.path.join(output_dir, "train.jsonl")) - config_path = _nemotron_training_config( - tiny_nemotron_h_model_dir, - output_dir, - dataset_path, - max_steps=max_steps, - num_gpus=2, - ) - - result = run_training(config_path, num_gpus=2, timeout=900) - - result.assert_success() - assert result.global_step == max_steps - assert result.loss_history is not None and len(result.loss_history) == max_steps - assert all(math.isfinite(loss) for loss in result.loss_history) - assert result.final_grad_norm is not None and math.isfinite(result.final_grad_norm) - - -# --------------------------------------------------------------------------- -# torchrun worker: real Trainer in-process, grad capture at first step -# --------------------------------------------------------------------------- - - -def _local_grad_abs_sum(param) -> float: - grad = param.grad - if grad is None: - return 0.0 - local = grad.to_local() if hasattr(grad, "to_local") else grad - if local.numel() == 0: - return 0.0 - total = local.abs().sum().item() - if not math.isfinite(total): - raise RuntimeError(f"non-finite gradient (abs sum={total})") - return total - - -def _grad_flow_worker() -> int: - from xorl.arguments import Arguments, parse_args # noqa: PLC0415 - from xorl.ops.ssm import Mamba2Mixer # noqa: PLC0415 - from xorl.trainers import Trainer # noqa: PLC0415 - - args = parse_args(Arguments) - trainer = Trainer(args) - - seen: dict = {} - max_cu_seqlens_docs = 0 - - def _record_cu_seqlens(module, hook_args, hook_kwargs): - del module, hook_args - cu_seqlens = hook_kwargs.get("cu_seqlens") - if cu_seqlens is not None: - nonlocal max_cu_seqlens_docs - max_cu_seqlens_docs = max(max_cu_seqlens_docs, int(cu_seqlens.numel()) - 1) - - for module in trainer.model.modules(): - if isinstance(module, Mamba2Mixer): - module.register_forward_pre_hook(_record_cu_seqlens, with_kwargs=True) - break - - def _capture(optimizer, opt_args, opt_kwargs): - del optimizer, opt_args, opt_kwargs - if seen: - return - for name, param in trainer.model.named_parameters(): - total = _local_grad_abs_sum(param) - for kind, suffixes in _GRAD_FLOW_TARGETS.items(): - if any(name.endswith(suffix) for suffix in suffixes): - seen[kind] = max(seen.get(kind, 0.0), total) - - trainer.optimizer.register_step_pre_hook(_capture) - trainer.train() - - missing = sorted(kind for kind in _GRAD_FLOW_TARGETS if seen.get(kind, 0.0) <= 0.0) - if missing: - print(f"NEMOTRON_GRAD_FLOW_MISSING {missing} (seen={seen})", flush=True) - return 1 - if max_cu_seqlens_docs < 2: - # Packed varlen plumbing: cu_seq_lens_q from the packing collator must reach - # the mamba mixers as cu_seqlens with multiple documents per pack. - print(f"NEMOTRON_GRAD_FLOW_NO_PACKED_CU_SEQLENS (max docs per pack={max_cu_seqlens_docs})", flush=True) - return 1 - seen["max_cu_seqlens_docs"] = max_cu_seqlens_docs - print(f"{_GRAD_FLOW_OK_MARKER} {json.dumps(seen)}", flush=True) - return 0 - - -if __name__ == "__main__": - sys.exit(_grad_flow_worker()) diff --git a/tests/e2e/test_opd_cpu.py b/tests/e2e/test_opd_cpu.py index 1ceeeaeb..f82b64c7 100644 --- a/tests/e2e/test_opd_cpu.py +++ b/tests/e2e/test_opd_cpu.py @@ -193,208 +193,3 @@ def test_opd_request_processor_to_backend_e2e(tmp_path): teacher_heads, ) assert result["loss"] == pytest.approx(expected.item(), rel=1e-5, abs=1e-6) - - -class _GlobalNormOPDBackend(DummyBackend): - """OPD backend that normalizes the KL loss by GLOBAL valid tokens. - - This mirrors the production normalization ("by global valid tokens across all - ranks", per the project loss-normalization contract), unlike the per-row - normalization in OPDCPUBackend. With - global normalization the loss is a function of the document multiset only, so it - must be invariant to how the packer groups documents into rows. - """ - - def __init__(self, student_hidden_table, student_head, teacher_hidden_caches, teacher_heads): - super().__init__() - self.student_hidden_table = student_hidden_table - self.student_head = student_head - self.teacher_hidden_caches = teacher_hidden_caches - self.teacher_heads = teacher_heads - self.row_sample_counts = None - - async def forward_backward( - self, - batches, - loss_fn="causallm_loss", - loss_fn_params=None, - model_id=None, - routed_experts=None, - routed_expert_logits=None, - request_id=None, - ): - numerator = torch.tensor(0.0) - global_valid = 0 - self.row_sample_counts = [b.get("num_samples") for b in batches] - for raw_batch in batches: - labels = torch.tensor(raw_batch["labels"], dtype=torch.long) - row_valid = int((labels != -100).sum().item()) - if row_valid == 0: - continue - # reference_grouped_opd_loss returns (Σ token_kl) / row_valid; multiply - # back to recover the raw per-row numerator, then normalize globally. - row_mean = reference_grouped_opd_loss( - raw_batch, - self.student_hidden_table, - self.student_head, - self.teacher_hidden_caches, - self.teacher_heads, - ) - numerator = numerator + row_mean * row_valid - global_valid += row_valid - loss = float((numerator / max(global_valid, 1)).item()) - return {"total_loss": loss, "global_valid_tokens": global_valid, "execution_time": 0.0} - - -def _opd_equivalence_request(data, teacher_files, strategy): - return OrchestratorRequest( - operation="forward_backward", - payload=ModelPassData( - data=[dict(d) for d in data], - loss_fn="opd_loss", - loss_fn_params={ - "teacher_heads": teacher_files.heads, - "teacher_hidden_caches": teacher_files.hidden_caches, - # Disable teacher pre-sort so the packing strategy fully controls order. - "opd_sort_by_teacher": False, - }, - model_id="opd-e2e", - ), - ) - - -def test_opd_loss_is_invariant_to_packing_strategy(tmp_path): - """Real forward-backward loss-equivalence across strategies (CPU analog of K3). - - Reordering documents into different rows must not change the globally-normalized - OPD loss — only the float reduction order, which is far below any meaningful - tolerance. This exercises the actual packing path + OPD loss, not just metadata. - """ - torch.manual_seed(7) - vocab_size, hidden_size, cache_size = 19, 6, 40 - student_hidden_table = torch.randn(vocab_size, hidden_size) / hidden_size**0.5 - student_head = torch.randn(vocab_size, hidden_size) / hidden_size**0.5 - teacher_heads = { - "0": torch.randn(vocab_size, hidden_size) / hidden_size**0.5, - "1": torch.randn(vocab_size, hidden_size) / hidden_size**0.5, - } - teacher_hidden_caches = { - "0": torch.randn(cache_size, hidden_size) / hidden_size**0.5, - "1": torch.randn(cache_size, hidden_size) / hidden_size**0.5, - } - teacher_files = make_teacher_files(tmp_path, teacher_heads, teacher_hidden_caches) - - # 12 varied-length samples across two teachers; lengths chosen so the three - # strategies produce genuinely different row groupings at pack_len=12, dp_size=4. - rng = torch.Generator().manual_seed(11) - data = [] - cache_cursor = {"0": 0, "1": 0} - for i in range(12): - length = int(torch.randint(2, 8, (1,), generator=rng).item()) - teacher = str(i % 2) - start = cache_cursor[teacher] - cache_cursor[teacher] = start + length - data.append( - { - "input_ids": [int(torch.randint(0, vocab_size, (1,), generator=rng).item()) for _ in range(length)], - "target_tokens": [int(torch.randint(0, vocab_size, (1,), generator=rng).item()) for _ in range(length)], - "teacher_id": int(teacher), - "teacher_weight": 1.0, - "teacher_cache_indices": list(range(start, start + length)), - } - ) - - results = {} - layouts = {} - for strategy in ("sequential", "best_fit", "balanced_dp"): - backend = _GlobalNormOPDBackend(student_hidden_table, student_head, teacher_hidden_caches, teacher_heads) - processor = RequestProcessor( - backend=backend, - sample_packing_sequence_len=12, - enable_packing=True, - pad_to_multiple_of=1, - cp_size=1, - packing_strategy=strategy, - on_oversized="error", - dp_size=4, - ) - output = asyncio.run( - processor.execute_forward_backward(_opd_equivalence_request(data, teacher_files, strategy)) - ) - result = output.outputs[0] - assert result["success"] is True - results[strategy] = (result["loss"], result["valid_tokens"]) - layouts[strategy] = backend.row_sample_counts - - # The strategies must genuinely differ in layout (otherwise the test is vacuous). - assert layouts["sequential"] != layouts["balanced_dp"] or layouts["sequential"] != layouts["best_fit"] - - # Same total valid tokens, and the globally-normalized loss matches to float tol. - seq_loss, seq_valid = results["sequential"] - for strategy, (loss, valid) in results.items(): - assert valid == seq_valid, f"{strategy} valid_tokens {valid} != {seq_valid}" - assert loss == pytest.approx(seq_loss, rel=1e-6, abs=1e-7), f"{strategy} loss {loss} != {seq_loss}" - - -class TeacherCacheCPUBackend(DummyBackend): - async def forward( - self, - batches, - loss_fn="causallm_loss", - loss_fn_params=None, - model_id=None, - routed_experts=None, - routed_expert_logits=None, - request_id=None, - ): - assert loss_fn == "teacher_hidden_cache" - assert model_id == "teacher" - return { - "total_loss": 0.0, - "global_valid_tokens": 5, - "teacher_hidden_cache": { - "backend": "mooncake", - "key": "opd/test/teacher/0/hidden", - "tensor_key": "hidden_states", - "tensor_shapes": {"hidden_states": [5, 6]}, - "tensor_dtypes": {"hidden_states": "bfloat16"}, - "num_tokens": 5, - "hidden_size": 6, - "cache_indices_by_sample": [[0, 1, 2], [3, 4]], - }, - "teacher_prefill_tokens": 5, - "teacher_prefill_forward_compute_s": 0.25, - "teacher_hidden_cache_write_s": 0.01, - "execution_time": 0.3, - } - - -def test_teacher_hidden_cache_metadata_passes_through_request_processor(): - backend = TeacherCacheCPUBackend() - processor = RequestProcessor( - backend=backend, - sample_packing_sequence_len=16, - enable_packing=True, - pad_to_multiple_of=1, - cp_size=1, - ) - request = OrchestratorRequest( - operation="forward", - payload=ModelPassData( - data=[ - {"input_ids": [1, 2, 3], "target_tokens": [1, 2, 3]}, - {"input_ids": [4, 5], "target_tokens": [4, 5]}, - ], - loss_fn="teacher_hidden_cache", - loss_fn_params={"teacher_hidden_cache_dtype": "bfloat16"}, - model_id="teacher", - ), - ) - - output = asyncio.run(processor.execute_forward(request)) - - assert output.output_type == OutputType.FORWARD - result = output.outputs[0] - assert result["teacher_hidden_cache"]["cache_indices_by_sample"] == [[0, 1, 2], [3, 4]] - assert result["teacher_prefill_tokens"] == 5 - assert result["teacher_prefill_forward_compute_s"] == 0.25 diff --git a/tests/e2e/test_opd_full_pipeline.py b/tests/e2e/test_opd_full_pipeline.py index 2040d649..27864dcc 100644 --- a/tests/e2e/test_opd_full_pipeline.py +++ b/tests/e2e/test_opd_full_pipeline.py @@ -348,8 +348,6 @@ def test_opd_full_pipeline_with_weight_sync(tmp_path): visible_devices = ( [int(x) for x in visible.split(",") if x.strip() != ""] if visible else list(range(torch.cuda.device_count())) ) - if len(visible_devices) < 3: - pytest.skip("Need at least 3 visible CUDA devices") trainer_gpu, student_gpu, teacher_gpu = visible_devices[0], visible_devices[1], visible_devices[2] teacher_log = tmp_path / "teacher_sglang.log" diff --git a/tests/e2e/test_opd_gpu.py b/tests/e2e/test_opd_gpu.py deleted file mode 100644 index 60804c93..00000000 --- a/tests/e2e/test_opd_gpu.py +++ /dev/null @@ -1,78 +0,0 @@ -"""CUDA smoke coverage for OPD runner teacher grouping.""" - -import pytest -import torch - -from tests._helpers.opd import make_teacher_files -from xorl.server.runner.model_runner import ModelRunner - - -pytestmark = [ - pytest.mark.e2e, - pytest.mark.gpu, - pytest.mark.server, - pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), -] - - -def test_opd_runner_grouped_teachers_cuda(tmp_path): - torch.manual_seed(321) - device = torch.device("cuda") - vocab_size = 23 - hidden_size = 8 - seq_len = 6 - cache_size = 12 - - teacher_heads = { - "0": torch.randn(vocab_size, hidden_size) / hidden_size**0.5, - "1": torch.randn(vocab_size, hidden_size) / hidden_size**0.5, - } - teacher_caches = { - "0": torch.randn(cache_size, hidden_size) / hidden_size**0.5, - "1": torch.randn(cache_size, hidden_size) / hidden_size**0.5, - } - teacher_files = make_teacher_files(tmp_path, teacher_heads, teacher_caches) - - runner = object.__new__(ModelRunner) - runner.train_config = {} - runner.lm_head_fp32 = True - runner.pp_enabled = False - runner._opd_head_manager = None - runner._opd_head_config = None - runner._opd_hidden_cache = None - runner._opd_hidden_config = None - - hidden_states = (torch.randn(1, seq_len, hidden_size, device=device) / hidden_size**0.5).requires_grad_(True) - student_weight = (torch.randn(vocab_size, hidden_size, device=device) / hidden_size**0.5).requires_grad_(True) - micro_batch = { - "labels": torch.tensor([[2, 3, 4, 5, 6, 7]], device=device), - "teacher_ids": torch.tensor([[0, 0, 0, 1, 1, 1]], device=device), - "teacher_cache_indices": torch.tensor([[0, 1, 2, 3, 4, 5]], device=device), - "teacher_weights": torch.tensor([[1.0, 0.5, 1.5, 1.0, 0.25, 2.0]], device=device), - } - - params = { - "teacher_heads": teacher_files.heads, - "teacher_hidden_caches": teacher_files.hidden_caches, - "num_chunks": 2, - } - optimizer = torch.optim.Adam([hidden_states, student_weight], lr=0.1) - loss_history: list[float] = [] - for _ in range(8): - optimizer.zero_grad(set_to_none=True) - result = runner._compute_opd_micro_batch_loss( - hidden_states=hidden_states, - student_weight=student_weight, - micro_batch=micro_batch, - params=params, - ) - assert result.loss.isfinite() - assert result.metrics["valid_tokens"] == seq_len - assert result.metrics["opd_num_teachers"] == 2 - loss_history.append(float(result.loss.detach().cpu())) - result.loss.backward() - assert hidden_states.grad is not None and hidden_states.grad.isfinite().all() - assert student_weight.grad is not None and student_weight.grad.isfinite().all() - optimizer.step() - - assert loss_history[-1] < loss_history[0], f"OPD loss did not decrease: {loss_history}" diff --git a/tests/experiments/test_training_sim.py b/tests/experiments/test_training_sim.py index 17343e6f..4012baa1 100644 --- a/tests/experiments/test_training_sim.py +++ b/tests/experiments/test_training_sim.py @@ -1,5 +1,6 @@ import json from pathlib import Path +from types import SimpleNamespace import pytest import yaml @@ -8,29 +9,54 @@ from xorl.sim.benchmark_behavior import load_benchmark_behavior_points from xorl.sim.calibration_evaluator import evaluate_calibration from xorl.sim.calibration_packs import ( - list_calibration_packs, load_calibration_pack, resolve_calibration_pack, - validate_calibration_pack, ) from xorl.sim.collect_calibration import parse_log_text, summarize_observed_run from xorl.sim.config_fingerprint import build_fingerprint, load_training_config, resolve_topology from xorl.sim.feasibility_evaluator import evaluate_feasibility -from xorl.sim.kernel_variants import compare_kernel_variants, rank_kernel_variants +from xorl.sim.kernel_variants import rank_kernel_variants from xorl.sim.model_metadata import resolve_model_metadata from xorl.sim.predict import build_report from xorl.sim.scenario_planner import plan_scenario from xorl.sim.schemas import ModelMetadata, Topology from xorl.sim.shape_engine import balanced_counts, build_shape_ledger -from xorl.sim.tradeoff_ranker import rank_benchmark_tradeoffs from xorl.sim.validate import validate_simulator +from xorl.utils.count_flops import XorlFlopsCounter -def test_balanced_counts_round_robin_distribution() -> None: +def _assert_training_sim_topology_shape_and_analytical_ledger_policy() -> None: assert balanced_counts(20, 6) == [4, 4, 3, 3, 3, 3] + _assert_resolve_topology_matches_training_arguments_dp_formula() + _assert_shape_ledger_uses_sequence_parallel_local_tokens() + _assert_portable_analytical_ledger_policy() + _assert_runtime_flops_counter_uses_global_sequence_lengths() + + +def _assert_runtime_flops_counter_uses_global_sequence_lengths() -> None: + config = SimpleNamespace( + model_type="qwen3_moe", + vocab_size=128, + hidden_size=64, + moe_intermediate_size=32, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + num_experts=8, + num_experts_per_tok=2, + ) + batch_seqlens = [128, 64] + + cp1_flops, _ = XorlFlopsCounter(config, cp_size=1).estimate_flops(batch_seqlens, delta_time=1.0) + cp64_flops, _ = XorlFlopsCounter(config, cp_size=64).estimate_flops(batch_seqlens, delta_time=1.0) + + assert cp1_flops > 0 + assert cp64_flops == cp1_flops -def test_resolve_topology_matches_training_arguments_dp_formula() -> None: + +def _assert_resolve_topology_matches_training_arguments_dp_formula() -> None: raw_config = { "train": { "micro_batch_size": 1, @@ -58,7 +84,7 @@ def test_resolve_topology_matches_training_arguments_dp_formula() -> None: assert topology.top_k == 4 -def test_shape_ledger_uses_sequence_parallel_local_tokens() -> None: +def _assert_shape_ledger_uses_sequence_parallel_local_tokens() -> None: raw_config = { "train": { "micro_batch_size": 1, @@ -85,7 +111,7 @@ def test_shape_ledger_uses_sequence_parallel_local_tokens() -> None: assert ledger.ep_rank_slots_per_microbatch == [512, 512, 512, 512] -def test_parse_structured_step_phase_and_memory_logs() -> None: +def _assert_parse_structured_step_phase_and_memory_logs() -> None: log_text = """ [STEP 4/9] loss=1.0 grad_norm=0.1 lr=1.0e-5 tflops=100.2 mfu=0.1010 tokens_per_sec=53414 time=72.100s peak_mem=39.8GB fwd=20.1GB bwd=39.8GB optim=10.0GB [STEP_PHASES 4/9] dataloader_max_s=0.100000 dataloader_mean_s=0.050000 model_forward_max_s=10.000000 model_forward_mean_s=9.000000 @@ -191,7 +217,8 @@ def _write_resolved_run_fixture(root: Path) -> Path: return config_path -def test_benchmark_behavior_loader_ingests_resolved_run_logs_and_ooms(tmp_path: Path) -> None: +def test_observed_benchmark_ingestion_and_planning_policy(tmp_path: Path) -> None: + _assert_parse_structured_step_phase_and_memory_logs() config_path = _write_resolved_run_fixture(tmp_path) points = load_benchmark_behavior_points(tmp_path) @@ -222,10 +249,6 @@ def test_benchmark_behavior_loader_ingests_resolved_run_logs_and_ooms(tmp_path: assert oom.peak_mem_gb == 79.09 assert oom.micro_batch_size == 3 - -def test_scenario_planner_keeps_observed_fit_feasible_when_safety_margin_is_tight(tmp_path: Path) -> None: - config_path = _write_resolved_run_fixture(tmp_path) - report = plan_scenario( config_path, benchmark_dir=tmp_path, @@ -245,7 +268,9 @@ def test_scenario_planner_keeps_observed_fit_feasible_when_safety_margin_is_tigh assert report.best_raw.recommendation == "remeasure_for_stability" -def test_build_fingerprint_reads_config_file(tmp_path: Path) -> None: +def test_training_config_model_metadata_and_path_admission_policy(tmp_path: Path) -> None: + _assert_training_sim_topology_shape_and_analytical_ledger_policy() + config_path = tmp_path / "config.yaml" config = { "train": { @@ -275,8 +300,12 @@ def test_build_fingerprint_reads_config_file(tmp_path: Path) -> None: assert fingerprint.topology.global_batch_size == 24 assert len(fingerprint.config_sha256) == 64 + _assert_resolve_model_metadata_from_hf_cache(tmp_path / "hf-cache") + _assert_resolve_known_qwen235_metadata_without_hf_cache() + _assert_simulator_rejects_untrusted_calibration_and_metadata_paths(tmp_path / "path-admission") + -def test_resolve_model_metadata_from_hf_cache(tmp_path: Path) -> None: +def _assert_resolve_model_metadata_from_hf_cache(tmp_path: Path) -> None: config_dir = tmp_path / "models--Example--MoE" / "snapshots" / "abc123" config_dir.mkdir(parents=True) (config_dir / "config.json").write_text( @@ -314,7 +343,7 @@ def test_resolve_model_metadata_from_hf_cache(tmp_path: Path) -> None: assert metadata.config_path is not None -def test_resolve_known_qwen235_metadata_without_hf_cache() -> None: +def _assert_resolve_known_qwen235_metadata_without_hf_cache() -> None: metadata = resolve_model_metadata( {"model": {"model_path": "Qwen/Qwen3-235B-A22B"}}, hf_cache_roots=[], @@ -333,7 +362,7 @@ def test_resolve_known_qwen235_metadata_without_hf_cache() -> None: def _write_q235_results_fixture(benchmark_dir: Path) -> None: - benchmark_dir.mkdir() + benchmark_dir.mkdir(parents=True) (benchmark_dir / "RESULTS.md").write_text( """ # Qwen3-235B-A22B @ 2k context @@ -350,7 +379,7 @@ def _write_q235_results_fixture(benchmark_dir: Path) -> None: ) -def test_qwen235_markdown_loader_extracts_pack_and_ga_rows(tmp_path: Path) -> None: +def test_qwen235_calibration_ingestion_and_evaluation_policy(tmp_path: Path) -> None: benchmark_dir = tmp_path / "q235" _write_q235_results_fixture(benchmark_dir) @@ -379,6 +408,11 @@ def test_qwen235_markdown_loader_extracts_pack_and_ga_rows(tmp_path: Path) -> No assert by_label["q235_markdown:n4_ep8_bd_pk16k"].tokens_per_sec is None assert by_label["q235_markdown:n4_ep8_bd_pk16k"].correctness_status == "oom" + _assert_qwen235_calibration_evaluator_reports_leave_one_out_ga_error(tmp_path / "evaluation") + _assert_qwen235_calibrated_scenario_policy(tmp_path / "scenario") + _assert_qwen235_topology_what_if_policy(tmp_path / "topology") + _assert_builtin_calibration_pack_policy() + def _write_q235_config_fixture(config_path: Path) -> None: config = { @@ -417,7 +451,7 @@ def _write_q235_config_fixture(config_path: Path) -> None: config_path.write_text(yaml.safe_dump(config), encoding="utf-8") -def test_qwen235_scenario_planner_uses_markdown_calibration_for_ga_tradeoff(tmp_path: Path) -> None: +def _assert_qwen235_calibrated_scenario_policy(tmp_path: Path) -> None: benchmark_dir = tmp_path / "q235" _write_q235_results_fixture(benchmark_dir) config_path = tmp_path / "q235.yaml" @@ -445,8 +479,11 @@ def test_qwen235_scenario_planner_uses_markdown_calibration_for_ga_tradeoff(tmp_ assert report.best_raw.feasibility_status == "feasible_calibrated_peak_high_pressure" assert report.best_promotable is None + _assert_qwen235_scenario_planner_extrapolates_ga_asymptote(tmp_path / "asymptote") + _assert_qwen235_scenario_planner_marks_matching_oom_pack_infeasible(tmp_path / "oom") + -def test_qwen235_scenario_planner_extrapolates_ga_asymptote_from_step_time_fit(tmp_path: Path) -> None: +def _assert_qwen235_scenario_planner_extrapolates_ga_asymptote(tmp_path: Path) -> None: benchmark_dir = tmp_path / "q235" _write_q235_results_fixture(benchmark_dir) config_path = tmp_path / "q235.yaml" @@ -488,7 +525,7 @@ def test_qwen235_scenario_planner_extrapolates_ga_asymptote_from_step_time_fit(t assert ga4.score_tokens_per_sec == 9_520.0 -def test_qwen235_calibration_evaluator_reports_leave_one_out_ga_error(tmp_path: Path) -> None: +def _assert_qwen235_calibration_evaluator_reports_leave_one_out_ga_error(tmp_path: Path) -> None: benchmark_dir = tmp_path / "q235" _write_q235_results_fixture(benchmark_dir) config_path = tmp_path / "q235.yaml" @@ -517,7 +554,7 @@ def test_qwen235_calibration_evaluator_reports_leave_one_out_ga_error(tmp_path: assert ga2.absolute_percentage_error == 27.143 -def test_qwen235_scenario_planner_does_not_exact_match_observed_row_to_tp_what_if(tmp_path: Path) -> None: +def _assert_qwen235_topology_what_if_policy(tmp_path: Path) -> None: benchmark_dir = tmp_path / "q235" _write_q235_results_fixture(benchmark_dir) config_path = tmp_path / "q235.yaml" @@ -544,8 +581,11 @@ def test_qwen235_scenario_planner_does_not_exact_match_observed_row_to_tp_what_i assert "TP extrapolation uses conservative communication penalty" in tp2.behavior.warnings assert tp2.score_tokens_per_sec == 6_804.0 + _assert_qwen235_scenario_planner_auto_sweeps_parallelism(tmp_path / "auto") + _assert_qwen235_auto_sweep_includes_long_context_cp(tmp_path / "long-context") + -def test_qwen235_scenario_planner_auto_sweeps_parallelism_strategy_space(tmp_path: Path) -> None: +def _assert_qwen235_scenario_planner_auto_sweeps_parallelism(tmp_path: Path) -> None: benchmark_dir = tmp_path / "q235" _write_q235_results_fixture(benchmark_dir) config_path = tmp_path / "q235.yaml" @@ -578,7 +618,7 @@ def test_qwen235_scenario_planner_auto_sweeps_parallelism_strategy_space(tmp_pat assert all(candidate.prediction_confidence == "extrapolated" for candidate in tp_candidates) -def test_qwen235_auto_sweep_includes_long_context_cp_without_cross_seq_calibration(tmp_path: Path) -> None: +def _assert_qwen235_auto_sweep_includes_long_context_cp(tmp_path: Path) -> None: benchmark_dir = tmp_path / "q235" _write_q235_results_fixture(benchmark_dir) config_path = tmp_path / "q235.yaml" @@ -616,7 +656,7 @@ def test_qwen235_auto_sweep_includes_long_context_cp_without_cross_seq_calibrati assert "observed_oom_boundary:q235_markdown:n4_ep8_bd_pk16k" in base_cp.risk_flags -def test_qwen235_scenario_planner_marks_matching_oom_pack_infeasible(tmp_path: Path) -> None: +def _assert_qwen235_scenario_planner_marks_matching_oom_pack_infeasible(tmp_path: Path) -> None: benchmark_dir = tmp_path / "q235" _write_q235_results_fixture(benchmark_dir) config_path = tmp_path / "q235.yaml" @@ -646,22 +686,13 @@ def test_qwen235_scenario_planner_marks_matching_oom_pack_infeasible(tmp_path: P assert "observed_oom_boundary:q235_markdown:n4_ep8_bd_pk16k" in candidate.risk_flags -def test_builtin_calibration_packs_are_sanitized_and_versioned() -> None: - assert list_calibration_packs() == ["qwen3_235b_a22b", "qwen3_5_397b_a17b", "qwen3_6_35b_a3b"] - for name in list_calibration_packs(): - pack = load_calibration_pack(name) - validation = validate_calibration_pack(pack.path) - assert pack.manifest["schema_version"] == 1 - assert pack.default_config.is_file() - assert validation["status"] == "pass" - - -def test_builtin_pack_prefix_rejects_path_traversal() -> None: +def _assert_builtin_pack_prefix_rejects_path_traversal() -> None: with pytest.raises(ValueError, match="unknown built-in calibration pack"): resolve_calibration_pack("builtin:../qwen3_6_35b_a3b") -def test_calibration_pack_rejects_paths_outside_pack_root(tmp_path: Path) -> None: +def _assert_simulator_rejects_untrusted_calibration_and_metadata_paths(tmp_path: Path) -> None: + tmp_path.mkdir(parents=True, exist_ok=True) (tmp_path / "manifest.json").write_text( json.dumps( { @@ -677,8 +708,13 @@ def test_calibration_pack_rejects_paths_outside_pack_root(tmp_path: Path) -> Non with pytest.raises(ValueError, match="must stay within"): load_calibration_pack(tmp_path) + _assert_builtin_pack_prefix_rejects_path_traversal() + _assert_calibration_pack_requires_default_config(tmp_path) + _assert_calibration_pack_rejects_symlink_escape(tmp_path) + _assert_model_metadata_restricts_local_reads(tmp_path) -def test_calibration_pack_requires_default_config(tmp_path: Path) -> None: + +def _assert_calibration_pack_requires_default_config(tmp_path: Path) -> None: (tmp_path / "manifest.json").write_text( json.dumps({"name": "incomplete", "configs": [], "results": []}), encoding="utf-8", @@ -688,7 +724,7 @@ def test_calibration_pack_requires_default_config(tmp_path: Path) -> None: load_calibration_pack(tmp_path) -def test_calibration_pack_rejects_symlink_escape(tmp_path: Path) -> None: +def _assert_calibration_pack_rejects_symlink_escape(tmp_path: Path) -> None: outside = tmp_path.parent / f"{tmp_path.name}-outside" outside.mkdir() (outside / "train.yaml").write_text("train: {}\n", encoding="utf-8") @@ -709,7 +745,7 @@ def test_calibration_pack_rejects_symlink_escape(tmp_path: Path) -> None: load_calibration_pack(tmp_path) -def test_model_metadata_restricts_local_config_reads_to_approved_roots(tmp_path: Path) -> None: +def _assert_model_metadata_restricts_local_reads(tmp_path: Path) -> None: allowed_root = tmp_path / "allowed" blocked_config = tmp_path / "blocked" / "config.json" allowed_root.mkdir() @@ -725,21 +761,7 @@ def test_model_metadata_restricts_local_config_reads_to_approved_roots(tmp_path: assert metadata.num_experts is None -def test_builtin_qwen35_pack_preserves_raw_and_promotable_winners() -> None: - pack = load_calibration_pack("qwen3_5_397b_a17b") - points = load_benchmark_behavior_points(pack.path) - report = rank_benchmark_tradeoffs(pack.path) - - assert len(points) == 6 - assert report.best_raw is not None - assert report.best_raw.score_tokens_per_sec == 59_217.0 - assert report.best_raw.promotable is False - assert report.best_promotable is not None - assert report.best_promotable.score_tokens_per_sec == 59_188.0 - assert report.best_promotable.promotable is True - - -def test_builtin_qwen36_pack_matches_default_config_but_remains_ungated() -> None: +def _assert_builtin_calibration_pack_policy() -> None: pack = load_calibration_pack("qwen3_6_35b_a3b") report = build_report( pack.default_config, @@ -758,8 +780,11 @@ def test_builtin_qwen36_pack_matches_default_config_but_remains_ungated() -> Non assert report.support.support_status == "supported_local_non_pp" assert report.timing.timing_coverage_status == "benchmark_total_step_only" + _assert_builtin_qwen235_pack_replays_fit_and_oom_boundaries() + _assert_consolidated_validator_covers_all_builtin_packs() + -def test_builtin_qwen235_pack_replays_fit_and_oom_boundaries() -> None: +def _assert_builtin_qwen235_pack_replays_fit_and_oom_boundaries() -> None: pack = load_calibration_pack("qwen3_235b_a22b") report = evaluate_feasibility(pack.default_config, benchmark_dir=pack.path) @@ -771,7 +796,7 @@ def test_builtin_qwen235_pack_replays_fit_and_oom_boundaries() -> None: assert {holdout.actual_outcome for holdout in report.holdouts} == {"fit", "oom"} -def test_portable_analytical_core_covers_flops_activations_and_communication() -> None: +def _assert_portable_analytical_ledger_policy() -> None: pack = load_calibration_pack("qwen3_5_397b_a17b") raw_config = load_training_config(pack.default_config) topology = resolve_topology(raw_config) @@ -790,8 +815,11 @@ def test_portable_analytical_core_covers_flops_activations_and_communication() - assert communication["status"] == "exact_analytic_bytes" assert communication["total_per_rank_gb"] > 0 + _assert_dense_no_recompute_activation_ledger_includes_mlp_intermediate() + _assert_cross_node_traffic_normalizes_expert_fsdp_all_gather_passes() -def test_dense_no_recompute_activation_ledger_includes_mlp_intermediate() -> None: + +def _assert_dense_no_recompute_activation_ledger_includes_mlp_intermediate() -> None: metadata = ModelMetadata( model_path=None, config_path=None, @@ -826,7 +854,7 @@ def test_dense_no_recompute_activation_ledger_includes_mlp_intermediate() -> Non assert ledger["terms"]["saved_full_activations"]["gb"] == expected_gb -def test_exposed_cross_node_traffic_normalizes_expert_fsdp_all_gather_passes() -> None: +def _assert_cross_node_traffic_normalizes_expert_fsdp_all_gather_passes() -> None: metadata = ModelMetadata( model_path=None, config_path=None, @@ -903,19 +931,21 @@ def test_kernel_variant_ranking_requires_a_correctness_gate() -> None: ] report = rank_kernel_variants(rows) - comparison = compare_kernel_variants(rows[1], rows[0]) - assert report["status"] == "ok" assert report["best"]["variant"] == "validated" assert report["measurements"][0]["variant"] == "fast-ungated" - assert comparison["speedup"] == 1.25 - assert comparison["candidate_promotable"] is False -def test_consolidated_validator_covers_all_builtin_packs() -> None: +def _assert_consolidated_validator_covers_all_builtin_packs() -> None: report = validate_simulator() + assert report["schema_version"] == 1 assert report["status"] == "pass" assert report["pack_count"] == 3 + assert {pack["name"] for pack in report["packs"]} == { + "qwen3_235b_a22b", + "qwen3_5_397b_a17b", + "qwen3_6_35b_a3b", + } assert report["check_count"] >= 200 assert report["failed_check_count"] == 0 diff --git a/tests/fp8_training/test_config_compat.py b/tests/fp8_training/test_config_compat.py deleted file mode 100644 index 0af9fc75..00000000 --- a/tests/fp8_training/test_config_compat.py +++ /dev/null @@ -1,211 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch.nn as nn - -from xorl.fp8_training import ( - FP8Linear, - UnsupportedFP8ConfigError, - inject_fp8_training_into_model, - normalize_fp8_training_config, - resolve_fp8_bf16_layer_islands, - summarize_fp8_training_model, - validate_external_fp8_runtime_config, - validate_fp8_blackwell_training_policy, -) - - -pytestmark = pytest.mark.cpu - - -class TinyDecoderStack(nn.Module): - def __init__(self, num_layers: int = 4): - super().__init__() - self.config = SimpleNamespace(num_hidden_layers=num_layers) - self.model = nn.Module() - self.model.layers = nn.ModuleList( - [ - nn.ModuleDict( - { - "self_attn": nn.ModuleDict( - { - "q_proj": nn.Linear(8, 8), - "o_proj": nn.Linear(8, 8), - } - ), - "mlp": nn.ModuleDict( - { - "gate_up_proj": nn.Linear(8, 16), - "down_proj": nn.Linear(16, 8), - } - ), - } - ) - for _ in range(num_layers) - ] - ) - self.lm_head = nn.Linear(8, 16, bias=False) - - def get_pp_module_config(self): - return { - "layer_prefix": "model.layers", - "num_layers": self.config.num_hidden_layers, - } - - -def test_nemo_fp8_cfg_blockwise_alias_enables_native_fp8_training(): - normalized = normalize_fp8_training_config( - { - "enable_fp8_training": False, - "fp8_cfg": {"enabled": True, "fp8": "e4m3", "fp8_recipe": "blockwise", "fp8_param": False}, - } - ) - - assert normalized["enable_fp8_training"] is True - - -@pytest.mark.parametrize( - "fp8_cfg, match", - [ - ({"enabled": True, "fp8": "hybrid"}, "hybrid"), - ({"enabled": True, "fp8_recipe": "tensorwise"}, "tensorwise"), - ({"enabled": True, "fp8_recipe": "mxfp8"}, "MXFP8|mxfp8"), - ({"enabled": True, "fp8_param": True}, "fp8_param"), - ], -) -def test_nemo_fp8_cfg_rejects_transformer_engine_only_recipes(fp8_cfg, match): - with pytest.raises(UnsupportedFP8ConfigError, match=match): - normalize_fp8_training_config({"fp8_cfg": fp8_cfg}) - - -def test_vllm_fp8_runtime_knobs_fail_before_silent_translation(): - with pytest.raises(UnsupportedFP8ConfigError, match="vLLM FP8 receiver"): - validate_external_fp8_runtime_config({"generation": {"vllm_cfg": {"precision": "fp8"}}}) - - with pytest.raises(UnsupportedFP8ConfigError, match="vLLM FP8 receiver"): - validate_external_fp8_runtime_config({"generation": {"vllm_cfg": {"quantization": "fp8"}}}) - - with pytest.raises(UnsupportedFP8ConfigError, match="num_first_layers_in_bf16"): - validate_external_fp8_runtime_config({"generation": {"vllm_cfg": {"num_first_layers_in_bf16": 1}}}) - - with pytest.raises(UnsupportedFP8ConfigError, match="num_last_layers_in_bf16"): - validate_external_fp8_runtime_config({"generation": {"vllm_cfg": {"num_last_layers_in_bf16": 1}}}) - - with pytest.raises(UnsupportedFP8ConfigError, match="quantization_ignored_layer_kws"): - validate_external_fp8_runtime_config( - {"generation": {"vllm_cfg": {"quantization_ignored_layer_kws": ["a_proj"]}}} - ) - - with pytest.raises(UnsupportedFP8ConfigError, match="vLLM DeepGEMM"): - validate_external_fp8_runtime_config({"generation": {"vllm_cfg": {"use_deep_gemm": True}}}) - - with pytest.raises(UnsupportedFP8ConfigError, match="pow2_weight_scaling_factors"): - validate_external_fp8_runtime_config({"generation": {"vllm_cfg": {"pow2_weight_scaling_factors": True}}}) - - with pytest.raises(UnsupportedFP8ConfigError, match="pow2_activation_scaling_factors"): - validate_external_fp8_runtime_config({"generation": {"vllm_cfg": {"pow2_activation_scaling_factors": True}}}) - - with pytest.raises(UnsupportedFP8ConfigError, match="receiver_kv_cache_dtype"): - validate_external_fp8_runtime_config({"generation": {"vllm_cfg": {"kv_cache_dtype": "fp8_e4m3"}}}) - - with pytest.raises(UnsupportedFP8ConfigError, match="receiver_kv_cache_dtype"): - validate_external_fp8_runtime_config({"kv_cache_dtype": "fp8"}) - - -@pytest.mark.parametrize( - "config", - [ - {"policy": {"quant_cfg": "FP8_DEFAULT_CFG"}}, - {"policy": {"generation": {"quant_cfg": {"format": "fp8_e4m3"}}}}, - {"generation": {"quant_cfg": "NVFP4_DEFAULT_CFG"}}, - ], -) -def test_nemo_modelopt_qarl_configs_fail_before_silent_translation(config): - with pytest.raises(UnsupportedFP8ConfigError, match="ModelOpt QARL"): - validate_external_fp8_runtime_config(config) - - -def test_blackwell_policy_rejects_native_fp8_training_without_override(): - with pytest.raises(UnsupportedFP8ConfigError, match="guarded on Blackwell"): - validate_fp8_blackwell_training_policy( - enable_fp8_training=True, - allow_blackwell=False, - device_name="NVIDIA GB200", - capability=(10, 0), - ) - - -def test_blackwell_policy_rejects_override_without_validation_artifact(): - with pytest.raises(UnsupportedFP8ConfigError, match="requires fp8_training_blackwell_validation_artifact"): - validate_fp8_blackwell_training_policy( - enable_fp8_training=True, - allow_blackwell=True, - validation_artifact=None, - device_name="NVIDIA B200", - capability=(10, 0), - ) - - -def test_blackwell_policy_allows_explicit_override_with_artifact(): - validate_fp8_blackwell_training_policy( - enable_fp8_training=True, - allow_blackwell=True, - validation_artifact="/tmp/fp8-blackwell-validation.json", - device_name="NVIDIA GB200", - capability=(10, 0), - ) - - -def test_resolve_fp8_bf16_layer_islands_covers_first_last_and_overlap(): - model = TinyDecoderStack(num_layers=4) - - assert resolve_fp8_bf16_layer_islands(model, num_first_layers_bf16=2) == [ - "model.layers.0.*", - "model.layers.1.*", - ] - assert resolve_fp8_bf16_layer_islands(model, num_last_layers_bf16=2) == [ - "model.layers.2.*", - "model.layers.3.*", - ] - assert resolve_fp8_bf16_layer_islands(model, num_first_layers_bf16=1, num_last_layers_bf16=2) == [ - "model.layers.0.*", - "model.layers.2.*", - "model.layers.3.*", - ] - assert resolve_fp8_bf16_layer_islands(model, num_first_layers_bf16=3, num_last_layers_bf16=3) == [ - "model.layers.0.*", - "model.layers.1.*", - "model.layers.2.*", - "model.layers.3.*", - ] - - -def test_resolve_fp8_bf16_layer_islands_rejects_too_large_and_nonstandard_layout(): - with pytest.raises(UnsupportedFP8ConfigError, match="exceeds model layer count"): - resolve_fp8_bf16_layer_islands(TinyDecoderStack(num_layers=2), num_first_layers_bf16=3) - - with pytest.raises(UnsupportedFP8ConfigError, match="standard model.layers"): - resolve_fp8_bf16_layer_islands(nn.Sequential(nn.Linear(8, 8)), num_first_layers_bf16=1) - - -def test_fp8_injection_keeps_generated_first_last_layer_islands_bf16(): - model = TinyDecoderStack(num_layers=4) - - replaced = inject_fp8_training_into_model( - model, - num_first_layers_bf16=1, - num_last_layers_bf16=1, - ) - - assert replaced == 9 - assert isinstance(model.model.layers[0]["self_attn"]["q_proj"], nn.Linear) - assert not isinstance(model.model.layers[0]["self_attn"]["q_proj"], FP8Linear) - assert isinstance(model.model.layers[1]["self_attn"]["q_proj"], FP8Linear) - assert isinstance(model.model.layers[2]["mlp"]["down_proj"], FP8Linear) - assert isinstance(model.model.layers[3]["mlp"]["down_proj"], nn.Linear) - assert not isinstance(model.model.layers[3]["mlp"]["down_proj"], FP8Linear) - assert isinstance(model.lm_head, FP8Linear) - - summary = summarize_fp8_training_model(model) - assert summary["bf16_layer_island_patterns"] == ["model.layers.0.*", "model.layers.3.*"] - assert summary["bf16_layer_island_count"] == 2 diff --git a/tests/fp8_training/test_fp8_linear.py b/tests/fp8_training/test_fp8_linear.py index 0aa78d4e..53ae2387 100644 --- a/tests/fp8_training/test_fp8_linear.py +++ b/tests/fp8_training/test_fp8_linear.py @@ -1,5 +1,6 @@ import json import warnings +from types import SimpleNamespace import pytest import torch @@ -11,6 +12,8 @@ clear_linear_error_profile, get_linear_error_profile, inject_fp8_training_into_model, + resolve_fp8_bf16_layer_islands, + summarize_fp8_training_model, ) @@ -61,8 +64,9 @@ def __init__(self): class TinyTransformerBlock(nn.Module): - def __init__(self): + def __init__(self, num_layers: int = 1): super().__init__() + self.config = SimpleNamespace(num_hidden_layers=num_layers) self.model = nn.Module() self.model.layers = nn.ModuleList( [ @@ -82,12 +86,19 @@ def __init__(self): ), } ) + for _ in range(num_layers) ] ) self.lm_head = nn.Linear(16, 8, bias=False) + def get_pp_module_config(self): + return { + "layer_prefix": "model.layers", + "num_layers": self.config.num_hidden_layers, + } + -def test_inject_fp8_training_replaces_all_linears_by_default_and_preserves_parameters(): +def test_inject_fp8_training_policy(): model = TinyModel() original_weight = model.proj.weight replaced = inject_fp8_training_into_model(model) @@ -101,8 +112,47 @@ def test_inject_fp8_training_replaces_all_linears_by_default_and_preserves_param assert model.lm_head.fp8_output_dtype == "float32" assert "proj.weight" in model.state_dict() + _assert_injected_modules_are_tagged_with_fqns() + _assert_inject_fp8_training_recipe_policy() + _assert_inject_fp8_training_exclusion_policy() + _assert_fp8_bf16_layer_island_policy() + _assert_fp8_linear_cpu_fallback_policy() + + +def _assert_fp8_bf16_layer_island_policy(): + model = TinyTransformerBlock(num_layers=4) + + assert resolve_fp8_bf16_layer_islands( + model, + num_first_layers_bf16=1, + num_last_layers_bf16=2, + ) == [ + "model.layers.0.*", + "model.layers.2.*", + "model.layers.3.*", + ] + with pytest.raises(ValueError, match="exceeds model layer count"): + resolve_fp8_bf16_layer_islands(model, num_first_layers_bf16=5) + + replaced = inject_fp8_training_into_model( + model, + num_first_layers_bf16=1, + num_last_layers_bf16=1, + ) + + assert replaced == 9 + assert not isinstance(model.model.layers[0]["self_attn"]["qkv_proj"], FP8Linear) + assert isinstance(model.model.layers[1]["self_attn"]["qkv_proj"], FP8Linear) + assert isinstance(model.model.layers[2]["mlp"]["down_proj"], FP8Linear) + assert not isinstance(model.model.layers[3]["mlp"]["down_proj"], FP8Linear) + assert isinstance(model.lm_head, FP8Linear) + + summary = summarize_fp8_training_model(model) + assert summary["bf16_layer_island_patterns"] == ["model.layers.0.*", "model.layers.3.*"] + assert summary["bf16_layer_island_count"] == 2 -def test_inject_fp8_training_tags_replaced_modules_with_fqns(): + +def _assert_injected_modules_are_tagged_with_fqns(): model = TinyTransformerBlock() inject_fp8_training_into_model(model) @@ -112,7 +162,7 @@ def test_inject_fp8_training_tags_replaced_modules_with_fqns(): assert model.lm_head.fp8_module_name == "lm_head" -def test_inject_fp8_training_threads_amax_scale_recipe(): +def _assert_inject_fp8_training_recipe_policy(): model = TinyModel() inject_fp8_training_into_model(model, activation_amax_scale=0.875, weight_amax_scale=1.125) @@ -122,8 +172,11 @@ def test_inject_fp8_training_threads_amax_scale_recipe(): assert model.lm_head.fp8_activation_amax_scale == 0.875 assert model.lm_head.fp8_weight_amax_scale == 1.125 + _assert_fqn_module_recipe_overrides_are_applied() + _assert_unknown_module_recipe_override_key_is_rejected() + -def test_inject_fp8_training_applies_fqn_module_recipe_overrides(): +def _assert_fqn_module_recipe_overrides_are_applied(): model = TinyTransformerBlock() inject_fp8_training_into_model( @@ -157,7 +210,7 @@ def test_inject_fp8_training_applies_fqn_module_recipe_overrides(): assert model.lm_head.fp8_output_dtype == "float32" -def test_inject_fp8_training_rejects_unknown_module_recipe_override_key(): +def _assert_unknown_module_recipe_override_key_is_rejected(): model = TinyModel() with pytest.raises(ValueError, match="Unsupported FP8 module override key"): @@ -167,7 +220,7 @@ def test_inject_fp8_training_rejects_unknown_module_recipe_override_key(): ) -def test_inject_fp8_training_can_keep_explicit_exclusions_in_bf16(): +def _assert_inject_fp8_training_exclusion_policy(): model = TinyModel() replaced = inject_fp8_training_into_model(model, exclude_modules=["gate", "lm_head"]) @@ -179,8 +232,10 @@ def test_inject_fp8_training_can_keep_explicit_exclusions_in_bf16(): assert isinstance(model.lm_head, nn.Linear) assert not isinstance(model.lm_head, FP8Linear) + _assert_fqn_glob_exclusions_are_honored() -def test_inject_fp8_training_honors_fqn_glob_exclusions(): + +def _assert_fqn_glob_exclusions_are_honored(): model = TinyTransformerBlock() replaced = inject_fp8_training_into_model(model, exclude_modules=["model.layers.*.self_attn.*"]) @@ -195,7 +250,7 @@ def test_inject_fp8_training_honors_fqn_glob_exclusions(): assert isinstance(model.lm_head, FP8Linear) -def test_fp8_linear_cpu_fallback_matches_linear(): +def _assert_fp8_linear_cpu_fallback_policy(): torch.manual_seed(0) linear = nn.Linear(16, 32, dtype=torch.float32) fp8 = FP8Linear.from_linear(linear) @@ -207,8 +262,11 @@ def test_fp8_linear_cpu_fallback_matches_linear(): assert fp8.last_forward_used_fp8 is False assert torch.allclose(got, expected) + _assert_cpu_fallback_honors_float32_output_dtype() + _assert_fp8_linear_can_fail_fast_without_fallback() + -def test_fp8_linear_cpu_fallback_honors_float32_output_dtype(): +def _assert_cpu_fallback_honors_float32_output_dtype(): linear = nn.Linear(16, 32, dtype=torch.bfloat16) fp8 = FP8Linear.from_linear(linear, output_dtype="float32") x = torch.randn(4, 16, dtype=torch.bfloat16) @@ -219,7 +277,7 @@ def test_fp8_linear_cpu_fallback_honors_float32_output_dtype(): assert got.dtype == torch.float32 -def test_fp8_linear_can_fail_fast_without_fallback(): +def _assert_fp8_linear_can_fail_fast_without_fallback(): linear = nn.Linear(16, 32) fp8 = FP8Linear.from_linear(linear, allow_bf16_fallback=False) @@ -227,7 +285,7 @@ def test_fp8_linear_can_fail_fast_without_fallback(): fp8(torch.randn(2, 16)) -def test_fp8_linear_error_profiler_records_sampled_module_stats(monkeypatch, tmp_path): +def test_fp8_linear_error_profiler_sampling_policy(monkeypatch, tmp_path): from xorl.fp8_training.profiler import record_linear_error, write_linear_error_profile # noqa: PLC0415 clear_linear_error_profile() @@ -279,8 +337,13 @@ def test_fp8_linear_error_profiler_records_sampled_module_stats(monkeypatch, tmp assert json.loads(output_path.read_text())["tiny.proj"]["calls"] == 2 clear_linear_error_profile() + monkeypatch.delenv("XORL_FP8_LINEAR_ERROR_PROFILE_MAX_CALLS_PER_MODULE") + monkeypatch.delenv("XORL_FP8_LINEAR_ERROR_PROFILE_OUTPUT") + _assert_error_profiler_samples_explicit_flattened_rows(monkeypatch) + _assert_error_profiler_samples_module_specific_call_rows(monkeypatch) -def test_fp8_linear_error_profiler_samples_explicit_flattened_rows(monkeypatch): + +def _assert_error_profiler_samples_explicit_flattened_rows(monkeypatch): from xorl.fp8_training.profiler import record_linear_error # noqa: PLC0415 clear_linear_error_profile() @@ -312,7 +375,7 @@ def test_fp8_linear_error_profiler_samples_explicit_flattened_rows(monkeypatch): clear_linear_error_profile() -def test_fp8_linear_error_profiler_samples_module_specific_call_rows(monkeypatch): +def _assert_error_profiler_samples_module_specific_call_rows(monkeypatch): from xorl.fp8_training.profiler import record_linear_error # noqa: PLC0415 clear_linear_error_profile() @@ -347,9 +410,7 @@ def test_fp8_linear_error_profiler_samples_module_specific_call_rows(monkeypatch clear_linear_error_profile() -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_linear_error_profiler_records_cuda_operand_breakdown(monkeypatch): +def _assert_fp8_linear_error_profiler_records_cuda_operand_breakdown(monkeypatch): clear_linear_error_profile() monkeypatch.setenv("XORL_FP8_LINEAR_ERROR_PROFILE", "1") monkeypatch.setenv("XORL_FP8_LINEAR_ERROR_PROFILE_MAX_CALLS_PER_MODULE", "1") @@ -393,10 +454,7 @@ def test_fp8_linear_error_profiler_records_cuda_operand_breakdown(monkeypatch): clear_linear_error_profile() -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("block_size", [64, 128]) -def test_block_fp8_gemm_matches_explicit_dequantized_reference(block_size): +def _assert_block_fp8_gemm_matches_explicit_dequantized_reference(block_size): from xorl.ops.quantize import ( # noqa: PLC0415 block_fp8_dequantize, block_fp8_dequantize_gkn, @@ -417,13 +475,23 @@ def test_block_fp8_gemm_matches_explicit_dequantized_reference(block_size): b_dequant = block_fp8_dequantize_gkn(b_fp8, b_scales, block_size=block_size) expected = a_dequant @ b_dequant.T + assert (a - a_dequant).abs().mean() / a.abs().mean() < 0.03 + assert (b - b_dequant).abs().mean() / b.abs().mean() < 0.03 torch.testing.assert_close(got, expected, rtol=2e-3, atol=2e-3) @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("block_size", [64, 128]) -def test_block_fp8_gemm_rowwise_weight_scales_match_explicit_dequantized_reference(block_size): +def test_block_fp8_gemm_backend_and_scale_layout_policy(monkeypatch): + for block_size in (64, 128): + _assert_block_fp8_gemm_matches_explicit_dequantized_reference(block_size) + + _assert_block_fp8_gemm_rowwise_weight_scales_match_reference_policy() + _assert_block_fp8_gemm_torch_scaled_mm_backend_matches_reference() + _assert_block_fp8_gemm_auto_fallback_warns_once(monkeypatch) + + +def _assert_block_fp8_gemm_rowwise_weight_scales_match_reference(block_size): from xorl.ops.quantize import ( # noqa: PLC0415 block_fp8_dequantize, block_fp8_dequantize_gkn_rowwise, @@ -454,9 +522,12 @@ def test_block_fp8_gemm_rowwise_weight_scales_match_explicit_dequantized_referen torch.testing.assert_close(got, expected, rtol=2e-3, atol=2e-3) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_block_fp8_gemm_torch_scaled_mm_backend_matches_explicit_dequantized_reference(): +def _assert_block_fp8_gemm_rowwise_weight_scales_match_reference_policy(): + for block_size in (64, 128): + _assert_block_fp8_gemm_rowwise_weight_scales_match_reference(block_size) + + +def _assert_block_fp8_gemm_torch_scaled_mm_backend_matches_reference(): from xorl.ops.quantize import ( # noqa: PLC0415 block_fp8_dequantize, block_fp8_dequantize_gkn_rowwise, @@ -488,9 +559,7 @@ def test_block_fp8_gemm_torch_scaled_mm_backend_matches_explicit_dequantized_ref torch.testing.assert_close(got, expected, rtol=2e-3, atol=2e-3) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_block_fp8_gemm_auto_fallback_warns_once(monkeypatch): +def _assert_block_fp8_gemm_auto_fallback_warns_once(monkeypatch): import importlib # noqa: PLC0415 from xorl.ops.quantize import ( # noqa: PLC0415 @@ -534,19 +603,7 @@ def fail_scaled_mm(*args, **kwargs): torch.testing.assert_close(first, second) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - ("block_size", "smoothquant_alpha", "activation_amax_scale", "weight_amax_scale"), - [ - (64, None, 1.0, 1.0), - (128, None, 1.0, 1.0), - (128, 0.5, 1.0, 1.0), - (64, 0.4, 0.875, 1.0), - (64, 0.4, 1.0, 1.125), - ], -) -def test_fp8_linear_matmul_padding_matches_explicit_dequantized_reference( +def _assert_fp8_linear_matmul_padding_matches_reference( block_size, smoothquant_alpha, activation_amax_scale, @@ -584,7 +641,24 @@ def test_fp8_linear_matmul_padding_matches_explicit_dequantized_reference( @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_linear_full_residual_correction_reduces_forward_error(): +def test_fp8_linear_matmul_and_train_step_policy(monkeypatch): + for recipe in ( + (64, None, 1.0, 1.0), + (128, None, 1.0, 1.0), + (128, 0.5, 1.0, 1.0), + (64, 0.4, 0.875, 1.0), + (64, 0.4, 1.0, 1.125), + ): + _assert_fp8_linear_matmul_padding_matches_reference(*recipe) + + _assert_fp8_linear_full_residual_correction_reduces_forward_error() + _assert_fp8_linear_activation2_reduces_activation_quantization_error() + _assert_fp8_linear_cuda_train_step_updates_master_weight() + with monkeypatch.context() as case_patch: + _assert_fp8_linear_error_profiler_records_cuda_operand_breakdown(case_patch) + + +def _assert_fp8_linear_full_residual_correction_reduces_forward_error(): from xorl.fp8_training.linear import _fp8_matmul # noqa: PLC0415 torch.manual_seed(11) @@ -613,9 +687,7 @@ def test_fp8_linear_full_residual_correction_reduces_forward_error(): assert corrected_error < base_error * 0.25 -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_linear_activation2_reduces_activation_quantization_error(): +def _assert_fp8_linear_activation2_reduces_activation_quantization_error(): from xorl.fp8_training.linear import _apply_smoothquant, _fp8_matmul, _pad_last_dim # noqa: PLC0415 from xorl.ops.quantize import block_fp8_dequantize_gkn_rowwise, block_fp8_quantize_gkn_rowwise # noqa: PLC0415 @@ -651,9 +723,7 @@ def test_fp8_linear_activation2_reduces_activation_quantization_error(): assert activation2_error < activation_error * 0.5 -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_linear_cuda_train_step_updates_master_weight(): +def _assert_fp8_linear_cuda_train_step_updates_master_weight(): torch.manual_seed(0) linear = nn.Linear(128, 128, device="cuda", dtype=torch.bfloat16) fp8 = FP8Linear.from_linear(linear, backward_mode="fp8", allow_bf16_fallback=False) @@ -672,10 +742,10 @@ def test_fp8_linear_cuda_train_step_updates_master_weight(): assert torch.isfinite(fp8.weight.grad.float()).all() assert not torch.equal(fp8.weight.detach(), before) + _assert_cuda_float32_output_dtype_still_uses_fp8() -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_linear_cuda_float32_output_dtype_still_uses_fp8(): + +def _assert_cuda_float32_output_dtype_still_uses_fp8(): linear = nn.Linear(128, 128, device="cuda", dtype=torch.bfloat16) fp8 = FP8Linear.from_linear(linear, output_dtype="float32", allow_bf16_fallback=False) x = torch.randn(8, 128, device="cuda", dtype=torch.bfloat16) diff --git a/tests/fp8_training/test_fp8_moe.py b/tests/fp8_training/test_fp8_moe.py index 17e5754b..9ae814d4 100644 --- a/tests/fp8_training/test_fp8_moe.py +++ b/tests/fp8_training/test_fp8_moe.py @@ -1,10 +1,5 @@ -import os -import subprocess -import sys -import textwrap -from pathlib import Path from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest import torch @@ -47,7 +42,7 @@ def forward( return self.output_proj(hidden_states) -def test_inject_fp8_training_enables_moe_experts_and_preserves_parameters(): +def test_inject_fp8_training_moe_policy(): model = TinyMoEModel() original_gate_up = model.experts.gate_up_proj original_down = model.experts.down_proj @@ -62,8 +57,11 @@ def test_inject_fp8_training_enables_moe_experts_and_preserves_parameters(): assert model.experts.fp8_training_grouped_backend == "triton_grouped" assert model.experts.last_forward_used_fp8 is False + _assert_injection_preserves_per_expert_biases() + _assert_summary_reports_unused_moe_modules() -def test_inject_fp8_training_enables_moe_experts_with_per_expert_biases(): + +def _assert_injection_preserves_per_expert_biases(): model = TinyMoEModel() model.experts.hidden_act = "clamped_swiglu" model.experts.gate_up_bias = nn.Parameter(torch.zeros(2, 64)) @@ -80,7 +78,7 @@ def test_inject_fp8_training_enables_moe_experts_with_per_expert_biases(): assert model.experts.fp8_training_enabled is True -def test_summarize_fp8_training_model_reports_unused_moe_modules(): +def _assert_summary_reports_unused_moe_modules(): model = TinyMoEModel() inject_fp8_training_into_model(model) @@ -92,43 +90,7 @@ def test_summarize_fp8_training_model_reports_unused_moe_modules(): assert summary["unused_moe_module_names"] == ["experts"] -def test_quack_moe_forward_tp_threads_fp8_compute(monkeypatch): - from xorl.ops.moe import quack as quack_ops # noqa: PLC0415 - - expected = torch.randn(3, 4) - apply = MagicMock(return_value=expected) - monkeypatch.setattr(quack_ops.QuackTPMoeExpertsFunction, "apply", apply) - - tp_group = object() - parallel_state = SimpleNamespace(tp_enabled=True, tp_mesh=SimpleNamespace(get_group=lambda: tp_group)) - with patch("xorl.ops.moe.quack.get_parallel_state", return_value=parallel_state): - output = quack_ops.quack_moe_forward( - module=None, - num_experts=2, - routing_weights=torch.ones(3, 1), - selected_experts=torch.zeros(3, 1, dtype=torch.long), - hidden_states=torch.randn(3, 4), - gate_proj=torch.randn(2, 4, 8), - up_proj=torch.randn(2, 4, 8), - down_proj=torch.randn(2, 8, 4), - hidden_act="silu", - fp8_compute=True, - fp8_grouped_backend="triton_grouped", - fp8_block_size=64, - ) - - torch.testing.assert_close(output, expected) - args = apply.call_args.args - assert args[7] is tp_group - assert args[9] is False # activation_native - assert args[10] is True - assert args[11] == "triton_grouped" - assert args[12] == 64 - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_quack_moe_tp_fp8_single_rank_train_step_updates_master_weights(monkeypatch): +def _assert_quack_moe_tp_fp8_single_rank_train_step_updates_master_weights(monkeypatch): from xorl.ops.moe import quack as quack_ops # noqa: PLC0415 class _NoOpWork: @@ -176,6 +138,7 @@ def fake_all_reduce(tensor, group=None, async_op=False): hidden_act="silu", fp8_compute=True, fp8_grouped_backend="triton_grouped", + fp8_block_size=64, ) loss = output.float().pow(2).mean() @@ -224,7 +187,7 @@ def _cu_seqlens(cumsum: torch.Tensor) -> torch.Tensor: @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_block_grouped_same_nk_helper_matches_bf16_reference(): +def test_fp8_grouped_forward_and_wgrad_policy(monkeypatch): torch.manual_seed(0) cumsum = torch.tensor([4, 8], device="cuda", dtype=torch.int32) a = (torch.randn(8, 128, device="cuda", dtype=torch.bfloat16) * 0.25).contiguous() @@ -241,19 +204,16 @@ def test_fp8_block_grouped_same_nk_helper_matches_bf16_reference(): assert torch.isfinite(got.float()).all() assert torch.allclose(got.float(), expected, rtol=0.25, atol=0.5) + _assert_triton_grouped_same_nk_matches_reference_policy() + _assert_triton_grouped_helpers_honor_non_default_block_size() + _assert_triton_grouped_same_nk_uses_precomputed_cu_seqlens() -@pytest.mark.parametrize( - ("lengths", "k", "n"), - [ - ([4, 4], 128, 64), - ([0, 3, 7], 192, 160), - ([1, 129, 131], 257, 96), - ([2, 0, 260], 128, 257), - ], -) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_triton_grouped_same_nk_helper_matches_bf16_reference(lengths, k, n): + with monkeypatch.context() as case_patch: + _assert_fp8_grouped_wgrad_policy(case_patch) + _assert_fp8_scalar_quack_grouped_policy() + + +def _assert_fp8_triton_grouped_same_nk_matches_reference(lengths, k, n): torch.manual_seed(sum(lengths) + k + n) cumsum = torch.tensor(lengths, device="cuda", dtype=torch.int32).cumsum(0) total_tokens = int(cumsum[-1].item()) @@ -272,9 +232,15 @@ def test_fp8_triton_grouped_same_nk_helper_matches_bf16_reference(lengths, k, n) assert torch.allclose(got.float(), expected, rtol=0.25, atol=0.5) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_triton_grouped_helpers_honor_non_default_block_size(): +def _assert_triton_grouped_same_nk_matches_reference_policy(): + for lengths, k, n in ( + ([4, 4], 128, 64), + ([0, 129, 131], 257, 160), + ): + _assert_fp8_triton_grouped_same_nk_matches_reference(lengths, k, n) + + +def _assert_triton_grouped_helpers_honor_non_default_block_size(): torch.manual_seed(0) lengths = [3, 5] cumsum = torch.tensor(lengths, device="cuda", dtype=torch.int32).cumsum(0) @@ -309,9 +275,7 @@ def test_fp8_triton_grouped_helpers_honor_non_default_block_size(): assert torch.allclose(got_mn.float(), expected_mn, rtol=0.25, atol=0.5) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_triton_grouped_same_nk_uses_precomputed_cu_seqlens_for_dgrad_shape(): +def _assert_triton_grouped_same_nk_uses_precomputed_cu_seqlens(): torch.manual_seed(0) lengths = [2, 0, 5] cumsum = torch.tensor(lengths, device="cuda", dtype=torch.int32).cumsum(0) @@ -333,9 +297,7 @@ def test_fp8_triton_grouped_same_nk_uses_precomputed_cu_seqlens_for_dgrad_shape( assert torch.allclose(got.float(), expected, rtol=0.25, atol=0.5) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_block_wgrad_helper_matches_bf16_reference(): +def _assert_fp8_grouped_wgrad_policy(monkeypatch): torch.manual_seed(0) cumsum = torch.tensor([4, 8], device="cuda", dtype=torch.int32) a = (torch.randn(8, 128, device="cuda", dtype=torch.bfloat16) * 0.25).contiguous() @@ -355,10 +317,12 @@ def test_fp8_block_wgrad_helper_matches_bf16_reference(): assert torch.isfinite(got_mn.float()).all() assert torch.allclose(got_mn.float(), expected_mn, rtol=0.25, atol=0.5) + _assert_scalar_quack_same_mn_dispatches_to_block_loop(monkeypatch) + _assert_triton_grouped_wgrad_uses_precomputed_cu_seqlens() + _assert_triton_grouped_wgrad_matches_reference_policy() -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_scalar_quack_same_mn_dispatch_uses_block_loop(monkeypatch): + +def _assert_scalar_quack_same_mn_dispatches_to_block_loop(monkeypatch): from xorl.fp8_training import grouped as fp8_grouped # noqa: PLC0415 torch.manual_seed(0) @@ -392,9 +356,7 @@ def spy_block_loop(**kwargs): assert torch.allclose(got_mn.float(), expected_mn, rtol=0.25, atol=0.5) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_triton_grouped_wgrad_helper_uses_precomputed_cu_seqlens(): +def _assert_triton_grouped_wgrad_uses_precomputed_cu_seqlens(): torch.manual_seed(0) lengths = [3, 0, 9] cumsum = torch.tensor(lengths, device="cuda", dtype=torch.int32).cumsum(0) @@ -417,19 +379,7 @@ def test_fp8_triton_grouped_wgrad_helper_uses_precomputed_cu_seqlens(): assert torch.allclose(got_mn.float(), expected_mn, rtol=0.25, atol=0.5) -@pytest.mark.parametrize( - ("lengths", "m", "n"), - [ - ([4, 4], 128, 64), - ([0, 0, 0], 64, 128), - ([0, 3, 7], 64, 32), - ([1, 129, 131], 96, 160), - ([0, 0, 2, 260], 256, 257), - ], -) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_triton_grouped_wgrad_helper_matches_bf16_reference(lengths, m, n): +def _assert_fp8_triton_grouped_wgrad_matches_reference(lengths, m, n): torch.manual_seed(sum(lengths) + m + n) cumsum = torch.tensor(lengths, device="cuda", dtype=torch.int32).cumsum(0) total_tokens = int(cumsum[-1].item()) @@ -450,57 +400,16 @@ def test_fp8_triton_grouped_wgrad_helper_matches_bf16_reference(lengths, m, n): assert torch.allclose(got_mn.float(), expected_mn, rtol=0.25, atol=0.5) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.skipif( - os.environ.get("XORL_TEST_DEEP_GEMM_FP8") != "1", - reason="DeepGEMM grouped FP8 binding is opt-in until stable under repeated training calls", -) -def test_fp8_deep_gemm_grouped_helper_subprocess(): - # The installed DeepGEMM Any-argument binding rejects tensor storage after - # pytest is imported. Validate that backend in a clean child interpreter. - repo_root = Path(__file__).resolve().parents[2] - env = os.environ.copy() - env["PYTHONPATH"] = "src" - script = r""" -import torch - -from xorl.fp8_training import fp8_group_gemm_same_nk - -def grouped_same_nk_reference(a, b, cumsum): - starts = torch.cat([torch.zeros(1, device=cumsum.device, dtype=cumsum.dtype), cumsum[:-1]]) - chunks = [] - for expert_idx, (start, end) in enumerate(zip(starts.tolist(), cumsum.tolist())): - chunks.append(a[start:end].float() @ b[expert_idx].float()) - return torch.cat(chunks, dim=0) - - -torch.manual_seed(0) -cumsum = torch.tensor([4, 8], device="cuda", dtype=torch.int32) -a = (torch.randn(8, 128, device="cuda", dtype=torch.bfloat16) * 0.25).contiguous() -b = (torch.randn(2, 128, 64, device="cuda", dtype=torch.bfloat16) * 0.25).contiguous() -got = fp8_group_gemm_same_nk(a=a, b=b, cumsum_M=cumsum, max_M=4, backend="deep_gemm") -expected = grouped_same_nk_reference(a, b, cumsum) -assert got.dtype == torch.bfloat16 -assert torch.isfinite(got.float()).all() -assert torch.allclose(got.float(), expected, rtol=0.25, atol=0.5) -print("deep_gemm_helper_subprocess_ok") -""" - result = subprocess.run( - [sys.executable, "-c", textwrap.dedent(script)], - cwd=repo_root, - env=env, - text=True, - capture_output=True, - timeout=120, - check=False, - ) - assert result.returncode == 0, result.stdout + result.stderr +def _assert_triton_grouped_wgrad_matches_reference_policy(): + for lengths, m, n in ( + ([4, 4], 128, 64), + ([0, 0, 0], 64, 128), + ([0, 129, 131], 96, 160), + ): + _assert_fp8_triton_grouped_wgrad_matches_reference(lengths, m, n) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_scalar_quack_fallback_matches_bf16_reference(): +def _assert_fp8_scalar_quack_grouped_policy(): torch.manual_seed(0) cumsum = torch.tensor([4, 8], device="cuda", dtype=torch.int32) a = (torch.randn(8, 128, device="cuda", dtype=torch.bfloat16) * 0.25).contiguous() @@ -517,10 +426,10 @@ def test_fp8_scalar_quack_fallback_matches_bf16_reference(): assert torch.isfinite(got.float()).all() assert torch.allclose(got.float(), expected, rtol=0.25, atol=0.5) + _assert_scalar_quack_uses_cu_seqlens_for_per_expert_scales() -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp8_scalar_quack_uses_cu_seqlens_for_per_expert_scales(): + +def _assert_scalar_quack_uses_cu_seqlens_for_per_expert_scales(): torch.manual_seed(0) lengths = [2, 6] cumsum = torch.tensor(lengths, device="cuda", dtype=torch.int32).cumsum(0) @@ -544,10 +453,7 @@ def test_fp8_scalar_quack_uses_cu_seqlens_for_per_expert_scales(): assert torch.allclose(got.float(), expected, rtol=0.25, atol=0.5) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("fp8_grouped_backend", ["triton_grouped", "scalar_quack"]) -def test_moe_experts_quack_fp8_train_step_updates_master_weights(fp8_grouped_backend): +def _assert_moe_experts_quack_fp8_train_step_updates_master_weights(fp8_grouped_backend): torch.manual_seed(0) experts = MoEExperts(num_experts=2, hidden_dim=128, intermediate_size=128, moe_implementation="quack") experts = experts.to(device="cuda", dtype=torch.bfloat16) @@ -580,7 +486,17 @@ def test_moe_experts_quack_fp8_train_step_updates_master_weights(fp8_grouped_bac @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_moe_experts_quack_fp8_train_step_with_biases_and_clamped_swiglu(): +def test_moe_experts_quack_fp8_train_step_policy(monkeypatch): + for fp8_grouped_backend in ("triton_grouped", "scalar_quack"): + _assert_moe_experts_quack_fp8_train_step_updates_master_weights(fp8_grouped_backend) + + _assert_moe_train_step_with_biases_and_clamped_swiglu() + with monkeypatch.context() as tp_patch: + _assert_quack_moe_tp_fp8_single_rank_train_step_updates_master_weights(tp_patch) + _assert_injected_dense_and_moe_fp8_model_train_step_updates_master_weights() + + +def _assert_moe_train_step_with_biases_and_clamped_swiglu(): torch.manual_seed(0) experts = MoEExperts(num_experts=2, hidden_dim=128, intermediate_size=128, moe_implementation="quack") experts.hidden_act = "clamped_swiglu" @@ -617,10 +533,8 @@ def test_moe_experts_quack_fp8_train_step_with_biases_and_clamped_swiglu(): assert not torch.equal(experts.down_bias.detach(), before_down_bias) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("fp8_grouped_backend", ["triton_grouped", "scalar_quack"]) -def test_injected_dense_and_moe_fp8_model_train_step_updates_master_weights(fp8_grouped_backend): +def _assert_injected_dense_and_moe_fp8_model_train_step_updates_master_weights(): + fp8_grouped_backend = "triton_grouped" torch.manual_seed(0) model = TinyDenseMoEModel() changed = inject_fp8_training_into_model( diff --git a/tests/models/test_batch_invariance_dense.py b/tests/models/test_batch_invariance_dense.py index 63e1165c..50587064 100644 --- a/tests/models/test_batch_invariance_dense.py +++ b/tests/models/test_batch_invariance_dense.py @@ -24,9 +24,7 @@ D = 2048 -@requires_cuda -@pytest.mark.gpu -def test_ops_are_batch_composition_invariant(): +def _assert_ops_are_batch_composition_invariant(): """Row 0's output is bit-identical alone (M=1) vs in a batch (M=N), under BI.""" torch.manual_seed(0) dev = "cuda" @@ -50,9 +48,11 @@ def test_ops_are_batch_composition_invariant(): @requires_cuda @pytest.mark.gpu -def test_dense_model_logits_invariant_to_batching(): +def test_dense_batch_composition_invariance_policy(): """A sequence's per-position logits are identical whether it is forwarded alone or as the first row of a padded batch, under batch-invariant mode.""" + _assert_ops_are_batch_composition_invariant() + torch.manual_seed(1) dev = "cuda" cfg = Qwen3Config( diff --git a/tests/models/test_deepseek_v3_checkpoint_handler.py b/tests/models/test_deepseek_v3_checkpoint_handler.py index be9d3a1f..9c0bd2b8 100644 --- a/tests/models/test_deepseek_v3_checkpoint_handler.py +++ b/tests/models/test_deepseek_v3_checkpoint_handler.py @@ -21,13 +21,12 @@ def _expert_weight(expert_idx: int, proj: str) -> torch.Tensor: return torch.full((intermediate_size, hidden_size), value) -def _pack_int4(values: torch.Tensor) -> torch.Tensor: +def _pack_quantized(values: torch.Tensor, *, num_bits: int) -> torch.Tensor: if values.dtype != torch.int8: raise ValueError(f"Expected int8 values to pack, got {values.dtype}") if values.ndim != 2: raise ValueError(f"Expected rank-2 tensor to pack, got {tuple(values.shape)}") - num_bits = 4 pack_factor = 32 // num_bits unsigned = (values + (1 << (num_bits - 1))).to(torch.uint8) pad_cols = (-values.shape[1]) % pack_factor @@ -38,13 +37,19 @@ def _pack_int4(values: torch.Tensor) -> torch.Tensor: return (reshaped << bit_shifts).sum(dim=2, dtype=torch.int32) -def _packed_expert_weight(expert_idx: int, proj: str) -> dict[str, torch.Tensor]: +def _packed_expert_weight( + expert_idx: int, + proj: str, + *, + num_bits: int = 4, + group_size: int = 32, +) -> dict[str, torch.Tensor]: dense_weight = _expert_weight(expert_idx, proj) quantized = torch.ones_like(dense_weight, dtype=torch.int8) - num_groups = max(1, math.ceil(dense_weight.shape[1] / 32)) + num_groups = max(1, math.ceil(dense_weight.shape[1] / group_size)) scales = torch.full((dense_weight.shape[0], num_groups), dense_weight.flatten()[0].item(), dtype=torch.float32) return { - "weight_packed": _pack_int4(quantized), + "weight_packed": _pack_quantized(quantized, num_bits=num_bits), "weight_scale": scales, "weight_shape": torch.tensor(dense_weight.shape, dtype=torch.int64), } @@ -77,25 +82,46 @@ def _tiny_config() -> DeepseekV3Config: return config -def test_checkpoint_handler_merges_language_model_experts_and_skips_multimodal_keys(): - handler = DeepseekV3CheckpointHandler(num_experts=4) +def _load_external_experts( + handler: DeepseekV3CheckpointHandler, + *, + packed: bool = False, + packed_num_bits: int = 4, + packed_group_size: int = 32, +) -> dict[str, torch.Tensor]: loaded = {} - + skip_key = handler.get_skip_key_fn() for expert_idx in range(4): for proj in ("gate", "up", "down"): - loaded.update( - handler.on_load_weight( - f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.weight", - _expert_weight(expert_idx, proj), + weights = ( + _packed_expert_weight( + expert_idx, + proj, + num_bits=packed_num_bits, + group_size=packed_group_size, ) + if packed + else {"weight": _expert_weight(expert_idx, proj)} ) + for suffix, tensor in weights.items(): + key = f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.{suffix}" + if skip_key is not None and skip_key(key): + loaded.update(handler.on_skip_weight(key)) + else: + loaded.update(handler.on_load_weight(key, tensor)) + return dict(loaded) + + +def test_checkpoint_handler_expert_layout_ep_and_packed_policy(tmp_path): + handler = DeepseekV3CheckpointHandler(num_experts=4) + loaded = _load_external_experts(handler) loaded.update(handler.on_load_weight("language_model.model.layers.0.self_attn.o_proj.weight", torch.eye(2))) assert handler.on_load_weight("vision_tower.encoder.weight", torch.ones(1)) == [] assert handler.on_load_weight("mm_projector.weight", torch.ones(1)) == [] - gate_up = dict(loaded)["model.layers.0.mlp.experts.gate_up_proj"] - down = dict(loaded)["model.layers.0.mlp.experts.down_proj"] + gate_up = loaded["model.layers.0.mlp.experts.gate_up_proj"] + down = loaded["model.layers.0.mlp.experts.down_proj"] assert gate_up.shape == (4, 2, 6) assert down.shape == (4, 3, 2) @@ -104,101 +130,63 @@ def test_checkpoint_handler_merges_language_model_experts_and_skips_multimodal_k assert torch.all(gate_up[3, :, :3] == 31.0) assert torch.all(gate_up[3, :, 3:] == 32.0) assert torch.all(down[1] == 13.0) - assert torch.equal(dict(loaded)["model.layers.0.self_attn.o_proj.weight"], torch.eye(2)) + assert torch.equal(loaded["model.layers.0.self_attn.o_proj.weight"], torch.eye(2)) - -def test_checkpoint_handler_splits_internal_fused_experts_on_save(): - handler = DeepseekV3CheckpointHandler(num_experts=2) + internal_handler = DeepseekV3CheckpointHandler(num_experts=2) gate = torch.arange(2 * 2 * 3, dtype=torch.float32).reshape(2, 2, 3) up = gate + 100.0 - gate_up = torch.cat([gate, up], dim=2) - down = torch.arange(2 * 3 * 2, dtype=torch.float32).reshape(2, 3, 2) + internal_gate_up = torch.cat([gate, up], dim=2) + internal_down = torch.arange(2 * 3 * 2, dtype=torch.float32).reshape(2, 3, 2) + + loaded_gate_up = dict(internal_handler.on_load_weight("model.layers.0.mlp.experts.gate_up_proj", internal_gate_up)) + loaded_down = dict(internal_handler.on_load_weight("model.layers.0.mlp.experts.down_proj", internal_down)) + assert torch.equal(loaded_gate_up["model.layers.0.mlp.experts.gate_up_proj"], internal_gate_up) + assert torch.equal(loaded_down["model.layers.0.mlp.experts.down_proj"], internal_down) - split_gate_up = dict(handler.on_save_weight("model.layers.0.mlp.experts.gate_up_proj", gate_up)) - split_down = dict(handler.on_save_weight("model.layers.0.mlp.experts.down_proj", down)) + split_gate_up = dict(internal_handler.on_save_weight("model.layers.0.mlp.experts.gate_up_proj", internal_gate_up)) + split_down = dict(internal_handler.on_save_weight("model.layers.0.mlp.experts.down_proj", internal_down)) assert torch.equal(split_gate_up["model.layers.0.mlp.experts.0.gate_proj.weight"], gate[0].transpose(0, 1)) assert torch.equal(split_gate_up["model.layers.0.mlp.experts.1.up_proj.weight"], up[1].transpose(0, 1)) - assert torch.equal(split_down["model.layers.0.mlp.experts.0.down_proj.weight"], down[0].transpose(0, 1)) - assert torch.equal(split_down["model.layers.0.mlp.experts.1.down_proj.weight"], down[1].transpose(0, 1)) + assert torch.equal(split_down["model.layers.0.mlp.experts.0.down_proj.weight"], internal_down[0].transpose(0, 1)) + assert torch.equal(split_down["model.layers.0.mlp.experts.1.down_proj.weight"], internal_down[1].transpose(0, 1)) + _assert_checkpoint_handler_ep_slices_dense_and_packed_experts() + _assert_checkpoint_handler_loads_packed_expert_weights_in_requested_dtype_and_config(tmp_path) -def test_checkpoint_handler_keeps_internal_fused_expert_layout_on_load(): - handler = DeepseekV3CheckpointHandler(num_experts=2) - gate_up = torch.arange(2 * 2 * 6, dtype=torch.float32).reshape(2, 2, 6) - down = torch.arange(2 * 3 * 2, dtype=torch.float32).reshape(2, 3, 2) - loaded_gate_up = dict(handler.on_load_weight("model.layers.0.mlp.experts.gate_up_proj", gate_up)) - loaded_down = dict(handler.on_load_weight("model.layers.0.mlp.experts.down_proj", down)) +def _assert_checkpoint_handler_ep_slices_dense_and_packed_experts(): + for packed in (False, True): + handler = DeepseekV3CheckpointHandler(num_experts=4, ep_rank=1, ep_size=2) + loaded = _load_external_experts(handler, packed=packed) + gate_up = loaded["model.layers.0.mlp.experts.gate_up_proj"] + down = loaded["model.layers.0.mlp.experts.down_proj"] - assert torch.equal(loaded_gate_up["model.layers.0.mlp.experts.gate_up_proj"], gate_up) - assert torch.equal(loaded_down["model.layers.0.mlp.experts.down_proj"], down) - - -def test_checkpoint_handler_ep_slices_to_local_experts(): - handler = DeepseekV3CheckpointHandler(num_experts=4, ep_rank=1, ep_size=2) - skip_key = handler.get_skip_key_fn() - loaded = {} + assert gate_up.shape == (2, 2, 6) + assert down.shape == (2, 3, 2) + assert gate_up[:, 0, 0].tolist() == [21.0, 31.0] + assert down[:, 0, 0].tolist() == [23.0, 33.0] - for expert_idx in range(4): - for proj in ("gate", "up", "down"): - key = f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.weight" - if skip_key is not None and skip_key(key): - loaded.update(handler.on_skip_weight(key)) - else: - loaded.update(handler.on_load_weight(key, _expert_weight(expert_idx, proj))) - - gate_up = dict(loaded)["model.layers.0.mlp.experts.gate_up_proj"] - down = dict(loaded)["model.layers.0.mlp.experts.down_proj"] - - assert gate_up.shape == (2, 2, 6) - assert down.shape == (2, 3, 2) - assert gate_up[:, 0, 0].tolist() == [21.0, 31.0] - assert down[:, 0, 0].tolist() == [23.0, 33.0] - - -def test_checkpoint_handler_loads_packed_expert_weights(): - handler = DeepseekV3CheckpointHandler(num_experts=4) - loaded = {} - - for expert_idx in range(4): - for proj in ("gate", "up", "down"): - for suffix, tensor in _packed_expert_weight(expert_idx, proj).items(): - loaded.update( - handler.on_load_weight( - f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.{suffix}", - tensor, - ) - ) - - gate_up = dict(loaded)["model.layers.0.mlp.experts.gate_up_proj"] - down = dict(loaded)["model.layers.0.mlp.experts.down_proj"] - - assert gate_up.shape == (4, 2, 6) - assert down.shape == (4, 3, 2) - assert torch.all(gate_up[0, :, :3] == 1.0) - assert torch.all(gate_up[0, :, 3:] == 2.0) - assert torch.all(gate_up[3, :, :3] == 31.0) - assert torch.all(gate_up[3, :, 3:] == 32.0) - assert torch.all(down[1] == 13.0) +def _assert_checkpoint_handler_loads_packed_expert_weights_in_requested_dtype_and_config(tmp_path): + model = DeepseekV3ForCausalLM(_tiny_config()) + checkpoint_keys = {"language_model.model.layers.0.mlp.experts.0.gate_proj.weight_packed"} + default_handler = model.get_checkpoint_handler(checkpoint_keys=checkpoint_keys) + assert isinstance(default_handler, DeepseekV3CheckpointHandler) + + default_loaded = _load_external_experts(default_handler, packed=True) + default_gate_up = default_loaded["model.layers.0.mlp.experts.gate_up_proj"] + default_down = default_loaded["model.layers.0.mlp.experts.down_proj"] + assert default_gate_up.shape == (4, 2, 6) + assert default_down.shape == (4, 3, 2) + assert torch.all(default_gate_up[0, :, :3] == 1.0) + assert torch.all(default_gate_up[3, :, 3:] == 32.0) + assert torch.all(default_down[1] == 13.0) -def test_checkpoint_handler_loads_packed_expert_weights_in_requested_dtype(): handler = DeepseekV3CheckpointHandler(num_experts=4, device=torch.device("cpu"), dtype=torch.bfloat16) - loaded = {} - - for expert_idx in range(4): - for proj in ("gate", "up", "down"): - for suffix, tensor in _packed_expert_weight(expert_idx, proj).items(): - loaded.update( - handler.on_load_weight( - f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.{suffix}", - tensor, - ) - ) - - gate_up = dict(loaded)["model.layers.0.mlp.experts.gate_up_proj"] - down = dict(loaded)["model.layers.0.mlp.experts.down_proj"] + loaded = _load_external_experts(handler, packed=True) + gate_up = loaded["model.layers.0.mlp.experts.gate_up_proj"] + down = loaded["model.layers.0.mlp.experts.down_proj"] assert handler._expert_buffer is not None assert handler._expert_buffer._device == torch.device("cpu") @@ -207,39 +195,6 @@ def test_checkpoint_handler_loads_packed_expert_weights_in_requested_dtype(): assert torch.all(gate_up[0, :, :3] == torch.tensor(1.0, dtype=torch.bfloat16)) assert torch.all(down[1] == torch.tensor(13.0, dtype=torch.bfloat16)) - -def test_checkpoint_handler_ep_slices_packed_experts_to_local_experts(): - handler = DeepseekV3CheckpointHandler(num_experts=4, ep_rank=1, ep_size=2) - skip_key = handler.get_skip_key_fn() - loaded = {} - - for expert_idx in range(4): - for proj in ("gate", "up", "down"): - for suffix, tensor in _packed_expert_weight(expert_idx, proj).items(): - key = f"language_model.model.layers.0.mlp.experts.{expert_idx}.{proj}_proj.{suffix}" - if skip_key is not None and skip_key(key): - loaded.update(handler.on_skip_weight(key)) - else: - loaded.update(handler.on_load_weight(key, tensor)) - - gate_up = dict(loaded)["model.layers.0.mlp.experts.gate_up_proj"] - down = dict(loaded)["model.layers.0.mlp.experts.down_proj"] - - assert gate_up.shape == (2, 2, 6) - assert down.shape == (2, 3, 2) - assert gate_up[:, 0, 0].tolist() == [21.0, 31.0] - assert down[:, 0, 0].tolist() == [23.0, 33.0] - - -def test_model_checkpoint_handler_accepts_official_packed_expert_layout(): - model = DeepseekV3ForCausalLM(_tiny_config()) - handler = model.get_checkpoint_handler( - checkpoint_keys={"language_model.model.layers.0.mlp.experts.0.gate_proj.weight_packed"}, - ) - assert isinstance(handler, DeepseekV3CheckpointHandler) - - -def test_model_checkpoint_handler_reads_packed_quant_config_from_text_config(tmp_path): (tmp_path / "config.json").write_text( json.dumps( { @@ -262,11 +217,19 @@ def test_model_checkpoint_handler_reads_packed_quant_config_from_text_config(tmp ) ) - model = DeepseekV3ForCausalLM(_tiny_config()) - handler = model.get_checkpoint_handler( - checkpoint_keys={"language_model.model.layers.0.mlp.experts.0.gate_proj.weight_packed"}, + configured_handler = model.get_checkpoint_handler( + checkpoint_keys=checkpoint_keys, weights_path=str(tmp_path), ) + configured_loaded = _load_external_experts( + configured_handler, + packed=True, + packed_num_bits=8, + packed_group_size=64, + ) - assert handler._packed_expert_group_size == 64 - assert handler._packed_expert_num_bits == 8 + configured_gate_up = configured_loaded["model.layers.0.mlp.experts.gate_up_proj"] + configured_down = configured_loaded["model.layers.0.mlp.experts.down_proj"] + assert torch.all(configured_gate_up[0, :, :3] == 1.0) + assert torch.all(configured_gate_up[3, :, 3:] == 32.0) + assert torch.all(configured_down[1] == 13.0) diff --git a/tests/models/test_deepseek_v3_model.py b/tests/models/test_deepseek_v3_model.py index 1182e3c2..f55fb460 100644 --- a/tests/models/test_deepseek_v3_model.py +++ b/tests/models/test_deepseek_v3_model.py @@ -2,7 +2,7 @@ import torch from xorl.lora.modules import LoraLinear -from xorl.lora.utils import inject_lora_into_model +from xorl.lora.utils import inject_lora_into_model_with_moe from xorl.models.layers.moe import MoEExpertsLoRA from xorl.models.layers.moe.routing_replay import RoutingReplay, set_replay_stage from xorl.models.transformers.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config @@ -73,11 +73,14 @@ def test_deepseek_v3_tiny_forward_backward_and_freeze_router(): if ".gate.weight" in name: assert param.requires_grad is False + _assert_deepseek_v3_default_lora_targets_cover_mla_and_moe() + _assert_deepseek_v3_router_observability_and_replay_policy() -def test_deepseek_v3_default_lora_targets_cover_mla_and_moe(): + +def _assert_deepseek_v3_default_lora_targets_cover_mla_and_moe(): model = DeepseekV3ForCausalLM(_tiny_config()) - inject_lora_into_model(model, r=4, lora_alpha=8, target_modules=None) + inject_lora_into_model_with_moe(model, r=4, lora_alpha=8, target_modules=None) attn = model.model.layers[0].self_attn mlp = model.model.layers[0].mlp @@ -92,8 +95,20 @@ def test_deepseek_v3_default_lora_targets_cover_mla_and_moe(): assert isinstance(mlp.shared_experts.down_proj, LoraLinear) assert isinstance(mlp.experts, MoEExpertsLoRA) + explicit_model = DeepseekV3ForCausalLM(_tiny_config()) + inject_lora_into_model_with_moe( + explicit_model, + r=4, + lora_alpha=8, + target_modules=["q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"], + ) + explicit_attention = explicit_model.model.layers[0].self_attn + for projection in ("q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"): + assert isinstance(getattr(explicit_attention, projection), LoraLinear) + assert not isinstance(explicit_attention.o_proj, LoraLinear) + -def test_deepseek_v3_forward_emits_router_logits_when_aux_loss_is_enabled_by_config(): +def _assert_deepseek_v3_router_observability_and_replay_policy(): config = _tiny_config() config.output_router_logits = False config.router_aux_loss_coef = 0.001 @@ -107,8 +122,6 @@ def test_deepseek_v3_forward_emits_router_logits_when_aux_loss_is_enabled_by_con assert outputs.router_logits is not None assert len(outputs.router_logits) == model.config.num_hidden_layers - -def test_deepseek_v3_router_logits_skip_dense_layers_for_aux_loss(): config = _tiny_config() config.first_k_dense_replace = 1 config.output_router_logits = False @@ -124,8 +137,10 @@ def test_deepseek_v3_router_logits_skip_dense_layers_for_aux_loss(): assert len(outputs.router_logits) == model.config.num_hidden_layers - config.first_k_dense_replace assert all(router_logits is not None for router_logits in outputs.router_logits) + _assert_deepseek_v3_routing_replay_records_weights() + -def test_deepseek_v3_routing_replay_records_weights(): +def _assert_deepseek_v3_routing_replay_records_weights(): model = DeepseekV3ForCausalLM(_tiny_config()) block = model.model.layers[0].mlp replay = RoutingReplay() diff --git a/tests/models/test_deepseek_v3_registry.py b/tests/models/test_deepseek_v3_registry.py index f0fd584d..7152805a 100644 --- a/tests/models/test_deepseek_v3_registry.py +++ b/tests/models/test_deepseek_v3_registry.py @@ -2,11 +2,14 @@ from types import SimpleNamespace import pytest +from tiktoken.load import dump_tiktoken_bpe -from xorl.models.auto import _load_local_xorl_config +from xorl.models import auto as auto_module +from xorl.models.auto import _load_local_xorl_config, build_processor, build_tokenizer from xorl.models.registry import get_registry from xorl.models.transformers.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config from xorl.models.transformers.deepseek_v3.modeling_deepseek_v3 import DeepseekV3ForCausalLM +from xorl.models.transformers.deepseek_v3.tokenization_kimi import TikTokenTokenizer pytestmark = [pytest.mark.cpu] @@ -51,13 +54,7 @@ def _make_kimi_text_config(): ) -def test_deepseek_v3_registered(): - registry = get_registry() - assert "DeepseekV3ForCausalLM" in registry.supported_models - assert registry.get_model_cls_from_model_arch("DeepseekV3ForCausalLM") is DeepseekV3ForCausalLM - - -def test_deepseek_v3_config_from_kimi_wrapper_hf_config(): +def _assert_deepseek_v3_config_maps_kimi_wrapper_and_official_aux_defaults(): hf_config = SimpleNamespace( model_type="kimi_k25", text_config=SimpleNamespace(**_make_kimi_text_config()), @@ -81,8 +78,8 @@ def test_deepseek_v3_config_from_kimi_wrapper_hf_config(): assert config.rope_scaling["rope_type"] == "default" assert config.rope_theta == 1000000.0 - -def test_deepseek_v3_config_maps_official_kimi_aux_loss_defaults(): + # Official Kimi snapshots spell the auxiliary coefficient differently and + # omit the derived router-output switch. kimi_text_config = _make_kimi_text_config() kimi_text_config.pop("router_aux_loss_coef") kimi_text_config.pop("output_router_logits") @@ -94,13 +91,73 @@ def test_deepseek_v3_config_maps_official_kimi_aux_loss_defaults(): tie_word_embeddings=False, ) - config = DeepseekV3Config.from_hf_config(hf_config) + official_config = DeepseekV3Config.from_hf_config(hf_config) + + assert official_config.router_aux_loss_coef == pytest.approx(0.001) + assert official_config.output_router_logits is True - assert config.router_aux_loss_coef == pytest.approx(0.001) - assert config.output_router_logits is True +def _assert_local_kimi_tokenizer_and_generic_fallback_loader_policy(monkeypatch, tmp_path): + tokenizer_dir = tmp_path / "kimi-tokenizer" + tokenizer_dir.mkdir() + dump_tiktoken_bpe({bytes([i]): i for i in range(256)}, str(tokenizer_dir / "tiktoken.model")) + (tokenizer_dir / "tokenizer_config.json").write_text( + json.dumps( + { + "tokenizer_class": "TikTokenTokenizer", + "auto_map": {"AutoTokenizer": ["tokenization_kimi.TikTokenTokenizer", None]}, + "bos_token": "[BOS]", + "eos_token": "[EOS]", + "unk_token": "[UNK]", + "pad_token": "[PAD]", + "additional_special_tokens": ["<|im_end|>"], + "added_tokens_decoder": { + "256": {"content": "[BOS]", "special": True}, + "257": {"content": "[EOS]", "special": True}, + "258": {"content": "[UNK]", "special": True}, + "259": {"content": "[PAD]", "special": True}, + "260": {"content": "<|im_end|>", "special": True}, + }, + } + ) + ) + + tokenizer = build_tokenizer(str(tokenizer_dir)) + + assert isinstance(tokenizer, TikTokenTokenizer) + assert tokenizer.bos_token_id == 256 + assert tokenizer.eos_token_id == 257 + assert tokenizer.pad_token_id == 259 + assert tokenizer.decode(tokenizer.encode("hello")) == "hello" + + calls = [] + + def capture(kind): + def fake_from_pretrained(path, **kwargs): + calls.append((kind, path, kwargs)) + return object() + + return fake_from_pretrained + + monkeypatch.setattr(auto_module.AutoTokenizer, "from_pretrained", capture("tokenizer")) + monkeypatch.setattr(auto_module.AutoProcessor, "from_pretrained", capture("processor")) + + build_tokenizer(str(tmp_path / "not-kimi")) + build_processor("processor-path") + + assert [kind for kind, _path, _kwargs in calls] == ["tokenizer", "processor"] + for _kind, _path, kwargs in calls: + assert "trust_remote_code" not in kwargs + assert kwargs["padding_side"] == "right" + + +def test_kimi_wrapper_config_registry_and_local_loading_policy(monkeypatch, tmp_path): + _assert_deepseek_v3_config_maps_kimi_wrapper_and_official_aux_defaults() + + registry = get_registry() + assert "DeepseekV3ForCausalLM" in registry.supported_models + assert registry.get_model_cls_from_model_arch("DeepseekV3ForCausalLM") is DeepseekV3ForCausalLM -def test_local_auto_config_unwraps_kimi_wrapper_text_config(tmp_path): config_dir = tmp_path / "kimi-k25" config_dir.mkdir() (config_dir / "config.json").write_text( @@ -123,3 +180,5 @@ def test_local_auto_config_unwraps_kimi_wrapper_text_config(tmp_path): assert config.model_type == "deepseek_v3" assert config.n_routed_experts == 8 assert config.n_shared_experts == 2 + + _assert_local_kimi_tokenizer_and_generic_fallback_loader_policy(monkeypatch, tmp_path) diff --git a/tests/models/test_dsv4_attention.py b/tests/models/test_dsv4_attention.py deleted file mode 100644 index eb919ee6..00000000 --- a/tests/models/test_dsv4_attention.py +++ /dev/null @@ -1,206 +0,0 @@ -"""CPU smoke tests for ``DeepSeekV4Attention``. - -Covers the two non-tilelang variants: - -- ``compress_ratio == 0``: pure window attention (no compressor, no indexer). -- ``compress_ratio == 128``: window + static-pool compressed KV (compressor only). - -The ``compress_ratio == 4`` (DSA indexer) variant requires tilelang and is -exercised by the kernel tests (``tests/ops/dsv4/``) plus a Phase-6 e2e job. - -Tests use compact dims so they finish in seconds on CPU. The ``XORL_DSV4_ROPE_MAX_SEQ_LEN`` -override keeps the precomputed freqs_cis tensor small. -""" - -import pytest -import torch - - -pytestmark = pytest.mark.cpu - - -@pytest.fixture(autouse=True) -def _small_rope_buffer(monkeypatch): - """Avoid precomputing 65k×D/2 freqs_cis when the test only needs ~16.""" - monkeypatch.setenv("XORL_DSV4_ROPE_MAX_SEQ_LEN", "1024") - monkeypatch.setenv("XORL_DSV4_SPARSE_ATTN_IMPL", "sparse") # pure-torch ref - # Clear the @lru_cache on precompute_freqs_cis so cross-test state with - # different (dim, seqlen, factor, ...) keys doesn't leak. - from xorl.ops.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 - - precompute_freqs_cis.cache_clear() - yield - precompute_freqs_cis.cache_clear() - - -def _tiny_config(*, compress_ratios): - """Compact DSv4 config that satisfies all internal consistency asserts. - - hidden_size=64, n_heads=4, n_groups=2, q_lora_rank=32, o_lora_rank=16, - head_dim=32, qk_rope_head_dim=8, sliding_window=8. - """ - from xorl.models.transformers.deepseek_v4 import DeepseekV4Config # noqa: PLC0415 - - return DeepseekV4Config( - vocab_size=128, - hidden_size=64, - num_hidden_layers=len(compress_ratios), - num_attention_heads=4, - num_key_value_heads=1, - head_dim=32, - qk_rope_head_dim=8, - max_position_embeddings=1024, - q_lora_rank=32, - o_groups=2, - o_lora_rank=16, - sliding_window=8, - moe_intermediate_size=64, - n_routed_experts=4, - n_shared_experts=1, - num_experts_per_tok=2, - num_hash_layers=0, - hc_mult=2, - compress_ratios=list(compress_ratios), - compress_rope_theta=160000.0, - rope_theta=10000.0, - rope_scaling={ - "type": "yarn", - "factor": 4.0, - "original_max_position_embeddings": 256, - "beta_fast": 32.0, - "beta_slow": 1.0, - }, - # MTP slot is consumed by compressor_ratios validator only when present - num_nextn_predict_layers=0, - ) - - -@pytest.mark.parametrize("compress_ratio", [0, 128]) -def test_attention_forward_backward_shapes(compress_ratio): - """Forward + backward at every variant produce the right shapes and finite grads.""" - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepSeekV4Attention # noqa: PLC0415 - - torch.manual_seed(0) - cfg = _tiny_config(compress_ratios=[compress_ratio, compress_ratio]) - layer = DeepSeekV4Attention(cfg, layer_id=0).to(torch.float32) - # ``attn_sink`` and the compressor's fp32 params (``ape``, ``wkv``, - # ``wgate``) are ``torch.empty``-allocated and never touched by any of - # the standard PyTorch ``nn.*`` constructors. Calling the production - # ``DeepseekV4PreTrainedModel._init_weights`` zero-inits them, but this - # test bypasses ``post_init`` so we replicate the contract here. Without - # this, a stray ``inf`` / ``nan`` in the uninitialized memory propagates - # through softmax and the assertion below trips intermittently depending - # on torch's allocator state. - with torch.no_grad(): - layer.attn_sink.zero_() - if layer.compressor is not None: - layer.compressor.ape.zero_() - for m in (layer.compressor.wkv, layer.compressor.wgate): - m.weight.normal_(0.0, 0.02) - layer.compressor.norm.weight.fill_(1.0) - # Window-only requires seqlen >= window_size; C128 requires seqlen % 128 == 0 - # — but the smallest C128-friendly seqlen is 128, which is fine. For the - # window-only path we use window_size as the seqlen. - seqlen = 128 if compress_ratio == 128 else cfg.sliding_window - x = torch.randn(1, seqlen, cfg.hidden_size, requires_grad=True) - - out = layer(x) - - assert out.shape == (1, seqlen, cfg.hidden_size), out.shape - assert torch.isfinite(out).all(), "non-finite forward output" - - out.sum().backward() - assert torch.isfinite(x.grad).all(), "non-finite input grad" - - # Every trainable param should have a finite gradient. - for name, p in layer.named_parameters(): - if not p.requires_grad: - continue - assert p.grad is not None, f"no grad for {name}" - assert torch.isfinite(p.grad).all(), f"non-finite grad on {name}" - - -def test_attn_sink_is_fp32_and_keep_fp32_marked(): - """attn_sink is per-head fp32 and tagged for the FSDP2 dtype policy.""" - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepSeekV4Attention # noqa: PLC0415 - - cfg = _tiny_config(compress_ratios=[0]) - layer = DeepSeekV4Attention(cfg, layer_id=0) - - assert layer.attn_sink.dtype == torch.float32 - assert layer.attn_sink.shape == (cfg.num_attention_heads,) - assert getattr(layer.attn_sink, "_keep_fp32", False) is True - - -def test_kv_qat_enabled_from_quantization_config(): - """Xorl mirrors Miles' config-driven FP8-QAT gate instead of an env toggle.""" - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepSeekV4Attention # noqa: PLC0415 - from xorl.ops.dsv4.utils import dsv4_kv_qat_enabled # noqa: PLC0415 - - cfg = _tiny_config(compress_ratios=[0]) - assert dsv4_kv_qat_enabled(cfg) is False - - cfg.quantization_config = {"quant_method": "fp8"} - assert dsv4_kv_qat_enabled(cfg) is True - assert DeepSeekV4Attention(cfg, layer_id=0)._kv_qat_enabled is True - - cfg.quantization_config = {"quant_method": "awq"} - assert dsv4_kv_qat_enabled(cfg) is False - - -def test_attn_sink_promoted_for_tilelang_after_bf16_cast(monkeypatch): - """Tilelang sparse attention requires the per-head sink tensor in fp32.""" - from xorl.models.transformers.deepseek_v4 import modeling_deepseek_v4 # noqa: PLC0415 - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepSeekV4Attention # noqa: PLC0415 - - monkeypatch.setenv("XORL_DSV4_SPARSE_ATTN_IMPL", "tilelang") - seen = {} - - def fake_sparse_attn_tilelang(q, kv, attn_sink, topk_idxs, sm_scale): - del kv, topk_idxs, sm_scale - seen["attn_sink_dtype"] = attn_sink.dtype - return torch.zeros_like(q) - - monkeypatch.setattr(modeling_deepseek_v4, "sparse_attn_tilelang", fake_sparse_attn_tilelang) - - cfg = _tiny_config(compress_ratios=[0]) - layer = DeepSeekV4Attention(cfg, layer_id=0).to(torch.bfloat16) - assert torch.is_complex(layer.freqs_cis) - assert layer.freqs_cis.imag.abs().max() > 0 - with torch.no_grad(): - layer.attn_sink.zero_() - - x = torch.randn(1, cfg.sliding_window, cfg.hidden_size, dtype=torch.bfloat16) - layer(x) - - assert layer.attn_sink.dtype == torch.bfloat16 - assert seen["attn_sink_dtype"] == torch.float32 - - -def test_compressor_present_only_when_needed(): - """C0 has no compressor/indexer; C128 has compressor; C4 has both.""" - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepSeekV4Attention # noqa: PLC0415 - - cfg = _tiny_config(compress_ratios=[0, 128, 4]) - - l0 = DeepSeekV4Attention(cfg, layer_id=0) - l128 = DeepSeekV4Attention(cfg, layer_id=1) - l4 = DeepSeekV4Attention(cfg, layer_id=2) - - assert l0.compressor is None and l0.indexer is None - assert l128.compressor is not None and l128.indexer is None - assert l4.compressor is not None and l4.indexer is not None - - -def test_tp_size_gt_1_rejected(): - """Stub guard until the xorl-style SP gather is wired.""" - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepSeekV4Attention # noqa: PLC0415 - - cfg = _tiny_config(compress_ratios=[0]) - - class _FakeGroup: - def size(self): - return 2 - - with pytest.raises(AssertionError, match="TP > 1 is not implemented"): - DeepSeekV4Attention(cfg, layer_id=0, tp_group=_FakeGroup()) diff --git a/tests/models/test_dsv4_autoconfig.py b/tests/models/test_dsv4_autoconfig.py deleted file mode 100644 index 16487a2e..00000000 --- a/tests/models/test_dsv4_autoconfig.py +++ /dev/null @@ -1,113 +0,0 @@ -"""Verify ``transformers.AutoConfig.from_pretrained`` resolves to our -vendored ``DeepseekV4Config`` for the upstream HF ``model_type = -"deepseek_v4"`` declared in Flash's ``config.json``. - -Without this dispatch, the xorl ``train`` CLI (which uses AutoConfig) -can't drive DSv4 training from the standard HF snapshot layout. -""" - -import json -import tempfile -from pathlib import Path - -import pytest - - -pytestmark = pytest.mark.cpu - - -_FLASH_SHAPE_CONFIG = { - "architectures": ["DeepseekV4ForCausalLM"], - "model_type": "deepseek_v4", - "vocab_size": 64, - "hidden_size": 32, - "num_hidden_layers": 2, - "num_attention_heads": 2, - "num_key_value_heads": 1, - "head_dim": 16, - "qk_rope_head_dim": 4, - "max_position_embeddings": 256, - "q_lora_rank": 16, - "o_groups": 1, - "o_lora_rank": 8, - "sliding_window": 8, - "moe_intermediate_size": 16, - "n_routed_experts": 4, - "n_shared_experts": 1, - "num_experts_per_tok": 2, - "num_hash_layers": 0, - "hc_mult": 2, - "hc_sinkhorn_iters": 20, - "hc_eps": 1e-6, - "compress_ratios": [0, 0], - "compress_rope_theta": 160000, - "swiglu_limit": 0.0, - "rope_theta": 10000.0, - "rope_scaling": { - "type": "yarn", - "factor": 4.0, - "original_max_position_embeddings": 128, - "beta_fast": 32.0, - "beta_slow": 1.0, - }, - "num_nextn_predict_layers": 0, - "rms_norm_eps": 1e-6, - "tie_word_embeddings": False, - "torch_dtype": "bfloat16", -} - - -def test_autoconfig_dispatches_to_deepseekv4_config(): - """``AutoConfig.from_pretrained(snapshot_with_model_type=deepseek_v4)`` - returns an instance of our ``DeepseekV4Config``.""" - from transformers import AutoConfig - - from xorl.models.transformers.deepseek_v4 import DeepseekV4Config # noqa: F401 registers - - with tempfile.TemporaryDirectory() as tmp: - cfg_path = Path(tmp) / "config.json" - with cfg_path.open("w") as f: - json.dump(_FLASH_SHAPE_CONFIG, f) - cfg = AutoConfig.from_pretrained(tmp) - - assert isinstance(cfg, DeepseekV4Config) - assert cfg.model_type == "deepseek_v4" - assert cfg.num_hidden_layers == 2 - assert cfg.n_routed_experts == 4 - - -def test_automodel_dispatches_to_for_causal_lm(): - """``AutoModelForCausalLM`` knows how to instantiate from our config. - - We don't actually call ``from_config`` here (it would materialize - a model on cuda); we just check the registry entry resolved. - """ - from transformers import AutoModelForCausalLM - - from xorl.models.transformers.deepseek_v4 import ( # noqa: F401 registers - DeepseekV4Config, - DeepseekV4ForCausalLM, - ) - - cls = AutoModelForCausalLM._model_mapping.get(DeepseekV4Config, None) - assert cls is DeepseekV4ForCausalLM - - -def test_build_foundation_model_uses_xorl_registry(): - """The normal xorl train/server builder can instantiate DSv4.""" - from xorl.models import build_foundation_model - from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM - - with tempfile.TemporaryDirectory() as tmp: - cfg_path = Path(tmp) / "config.json" - with cfg_path.open("w") as f: - json.dump(_FLASH_SHAPE_CONFIG, f) - - model = build_foundation_model( - tmp, - init_device="meta", - moe_implementation="eager", - attn_implementation="flash_attention_3", - ) - - assert isinstance(model, DeepseekV4ForCausalLM) diff --git a/tests/models/test_dsv4_hf_to_dcp.py b/tests/models/test_dsv4_hf_to_dcp.py index 7814291b..411381db 100644 --- a/tests/models/test_dsv4_hf_to_dcp.py +++ b/tests/models/test_dsv4_hf_to_dcp.py @@ -135,7 +135,7 @@ def _tiny_hf_snapshot(snapshot_dir: Path) -> dict[str, torch.Tensor]: return sd -def test_converter_meta_dtype_cast_preserves_fp32_destinations(): +def _assert_converter_meta_dtype_cast_preserves_fp32_destinations(): """The torchrun converter's meta-model cast must not downcast fp32-only params.""" repo_root = Path(__file__).resolve().parents[2] sys.path.insert(0, str(repo_root)) @@ -176,8 +176,10 @@ def test_converter_meta_dtype_cast_preserves_fp32_destinations(): assert getattr(param, "_keep_fp32", False) is True -def test_convert_dsv4_hf_to_dcp_round_trip(): +def test_convert_dsv4_hf_to_dcp_conversion_policy(): """Run the conversion script's main(), then load the DCP and compare.""" + _assert_converter_meta_dtype_cast_preserves_fp32_destinations() + repo_root = Path(__file__).resolve().parents[2] sys.path.insert(0, str(repo_root)) try: @@ -276,18 +278,41 @@ def __init__(self, d): ) torch.testing.assert_close(lora_roundtrip.model.layers[0].self_attn.wq_a.lora_B, lora_b_before) + if dist.is_initialized(): + dist.destroy_process_group() + _assert_convert_dsv4_hf_to_dcp_pair_across_shards() + def test_automodel_from_pretrained_loads_tiny_hf_snapshot(): - """AutoModel dispatch must satisfy the HF ``from_pretrained`` contract.""" - from transformers import AutoModelForCausalLM # noqa: PLC0415 + """AutoConfig, XoRL construction, and AutoModel loading share one snapshot contract.""" + from transformers import AutoConfig, AutoModelForCausalLM # noqa: PLC0415 - from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM # noqa: PLC0415 + from xorl.models import build_foundation_model # noqa: PLC0415 + from xorl.models.transformers.deepseek_v4 import ( # noqa: PLC0415 + DeepseekV4Config, + DeepseekV4ForCausalLM, + ) with tempfile.TemporaryDirectory() as tmp: snapshot_dir = Path(tmp) / "hf-snap" snapshot_dir.mkdir() original_sd = _tiny_hf_snapshot(snapshot_dir) + config = AutoConfig.from_pretrained(snapshot_dir) + assert isinstance(config, DeepseekV4Config) + assert config.model_type == "deepseek_v4" + assert config.num_hidden_layers == 2 + assert config.n_routed_experts == 4 + assert AutoModelForCausalLM._model_mapping.get(DeepseekV4Config, None) is DeepseekV4ForCausalLM + + meta_model = build_foundation_model( + snapshot_dir, + init_device="meta", + moe_implementation="eager", + attn_implementation="flash_attention_3", + ) + assert isinstance(meta_model, DeepseekV4ForCausalLM) + model = AutoModelForCausalLM.from_pretrained( str(snapshot_dir), torch_dtype=torch.bfloat16, @@ -343,7 +368,7 @@ def _split_snapshot_across_shards( json.dump({"metadata": {"total_size": 0}, "weight_map": weight_map}, f) -def test_convert_dsv4_hf_to_dcp_pair_across_shards(): +def _assert_convert_dsv4_hf_to_dcp_pair_across_shards(): """The streaming loader holds a weight in ``pending`` until its paired ``.scale`` arrives in a later shard. Split a synthetic FP8 expert weight across two shards (weight → shard 1, scale → shard 2) and diff --git a/tests/models/test_dsv4_loader.py b/tests/models/test_dsv4_loader.py index 7ad40d5c..c696bfce 100644 --- a/tests/models/test_dsv4_loader.py +++ b/tests/models/test_dsv4_loader.py @@ -71,76 +71,6 @@ def _tiny_config(*, compress_ratios, num_hash_layers=0): ) -# --------------------------------------------------------------------------- -# Name mapping -# --------------------------------------------------------------------------- - - -def test_name_mapping_top_level(): - from xorl.models.transformers.deepseek_v4.checkpoint_handler import _hf_to_xorl_name - - assert _hf_to_xorl_name("embed.weight") == "model.embed_tokens.weight" - assert _hf_to_xorl_name("head.weight") == "lm_head.weight" - assert _hf_to_xorl_name("norm.weight") == "model.norm.weight" - assert _hf_to_xorl_name("hc_head_fn") == "model.hc_head_fn" - assert _hf_to_xorl_name("hc_head_base") == "model.hc_head_base" - assert _hf_to_xorl_name("hc_head_scale") == "model.hc_head_scale" - - -def test_name_mapping_attention(): - from xorl.models.transformers.deepseek_v4.checkpoint_handler import _hf_to_xorl_name - - assert _hf_to_xorl_name("layers.0.attn.wq_a.weight") == "model.layers.0.self_attn.wq_a.weight" - assert _hf_to_xorl_name("layers.5.attn.attn_sink") == "model.layers.5.self_attn.attn_sink" - assert _hf_to_xorl_name("layers.10.attn.compressor.ape") == "model.layers.10.self_attn.compressor.ape" - # Indexer name renames: HF ``indexer.wq_b`` -> xorl ``indexer.linear_wq_b``. - assert ( - _hf_to_xorl_name("layers.10.attn.indexer.wq_b.weight") == "model.layers.10.self_attn.indexer.linear_wq_b.weight" - ) - assert ( - _hf_to_xorl_name("layers.10.attn.indexer.weights_proj.weight") - == "model.layers.10.self_attn.indexer.linear_weights_proj.weight" - ) - - -def test_name_mapping_ffn_norms_hc(): - from xorl.models.transformers.deepseek_v4.checkpoint_handler import _hf_to_xorl_name - - assert _hf_to_xorl_name("layers.3.attn_norm.weight") == "model.layers.3.input_layernorm.weight" - assert _hf_to_xorl_name("layers.3.ffn_norm.weight") == "model.layers.3.post_attention_layernorm.weight" - assert _hf_to_xorl_name("layers.3.hc_attn_fn") == "model.layers.3.hc_attn_fn" - assert _hf_to_xorl_name("layers.3.hc_ffn_scale") == "model.layers.3.hc_ffn_scale" - - # noaux_tc bias rename: HF gate.bias -> xorl mlp.gate.e_score_correction_bias. - assert _hf_to_xorl_name("layers.0.ffn.gate.bias") == "model.layers.0.mlp.gate.e_score_correction_bias" - # tid2eid sits on the block, not on the gate. - assert _hf_to_xorl_name("layers.0.ffn.gate.tid2eid") == "model.layers.0.mlp.tid2eid" - # Shared expert renames w1/w2/w3 -> gate_proj/down_proj/up_proj. - assert ( - _hf_to_xorl_name("layers.0.ffn.shared_experts.w1.weight") - == "model.layers.0.mlp.shared_experts.gate_proj.weight" - ) - assert ( - _hf_to_xorl_name("layers.0.ffn.shared_experts.w2.weight") - == "model.layers.0.mlp.shared_experts.down_proj.weight" - ) - assert ( - _hf_to_xorl_name("layers.0.ffn.shared_experts.w3.weight") == "model.layers.0.mlp.shared_experts.up_proj.weight" - ) - - -def test_name_mapping_skips_mtp_and_unknown(): - from xorl.models.transformers.deepseek_v4.checkpoint_handler import _hf_to_xorl_name - - assert _hf_to_xorl_name("mtp.0.attn.wq_a.weight") is None - assert _hf_to_xorl_name("totally.bogus.name") is None - - -# --------------------------------------------------------------------------- -# APE hotfix undo -# --------------------------------------------------------------------------- - - def _miles_apply_ape_hotfix(param): """Forward direction of miles ``_apply_ape_hotfix_mirror`` for the test.""" assert param.shape[0] == 4 @@ -148,32 +78,13 @@ def _miles_apply_ape_hotfix(param): return torch.cat([a, b], dim=0).view(4, -1).contiguous() -def test_ape_hotfix_round_trip(): - """``_undo_ape_hotfix(_miles_apply_ape_hotfix(x)) == x`` for all C4 shapes.""" - from xorl.models.transformers.deepseek_v4.checkpoint_handler import _undo_ape_hotfix - - torch.manual_seed(0) - for head_dim in (4, 16, 128): - x = torch.randn(4, 2 * head_dim) - hf_layout = _miles_apply_ape_hotfix(x) - recovered = _undo_ape_hotfix(hf_layout) - torch.testing.assert_close(recovered, x) - - -def test_ape_hotfix_assertion_on_wrong_shape(): - from xorl.models.transformers.deepseek_v4.checkpoint_handler import _undo_ape_hotfix - - with pytest.raises(AssertionError): - _undo_ape_hotfix(torch.zeros(8, 16)) - - # --------------------------------------------------------------------------- # FP8 block dequantization # --------------------------------------------------------------------------- -def test_fp8_dequantize_matches_hand_computed(): - """Tiny 256x256 weight, block 128: per-block scale repeats correctly.""" +def _assert_fp8_dequantize_matches_block_and_tail_references(): + """Full blocks and non-aligned tails use their corresponding scale cells.""" from xorl.models.transformers.deepseek_v4.checkpoint_handler import _dequantize_fp8_block torch.manual_seed(1) @@ -196,10 +107,18 @@ def test_fp8_dequantize_matches_hand_computed(): torch.testing.assert_close(bf16, expected) + torch.manual_seed(2) + tail_shape = (130, 200) + tail_weight = (torch.randn(*tail_shape) * 0.1).to(torch.float8_e4m3fn) + tail_scale = torch.tensor([[1.0, 2.0], [4.0, 8.0]]).to(torch.float8_e8m0fnu) + tail_out = _dequantize_fp8_block(tail_weight, tail_scale, block, torch.bfloat16) + tail_scale_full = tail_scale.float().repeat_interleave(128, 0).repeat_interleave(128, 1)[:130, :200] + tail_expected = (tail_weight.float() * tail_scale_full).to(torch.bfloat16) + torch.testing.assert_close(tail_out, tail_expected) + -def test_mxfp4_dequantize_known_pattern(): - """Hand-encode 4 FP4 values into 2 int8 bytes, verify the dequant produces - the expected values * scale. +def _assert_mxfp4_dequantize_known_values_and_block_scaling(): + """Hand-encoded values decode correctly and select their scale blocks. Encoding: byte 0 packs (element 0 = +0.5, element 1 = -2.0) = (low nibble = 0001, high nibble = 1100) = 0xC1. @@ -217,13 +136,6 @@ def test_mxfp4_dequantize_known_pattern(): expected = torch.tensor([[0.5, -2.0, 1.0, 6.0]], dtype=torch.float32) torch.testing.assert_close(out, expected) - -def test_mxfp4_dequantize_handles_block_scaling(): - """Two scale blocks -> different magnitudes per block.""" - from xorl.models.transformers.deepseek_v4.checkpoint_handler import ( - _dequantize_mxfp4_packed_int8, - ) - # 1 row, 32 packed bytes = 64 FP4 elements -> 2 scale blocks of 32. M, Np = 1, 32 raw = torch.zeros(M, Np, dtype=torch.int8) @@ -241,24 +153,6 @@ def test_mxfp4_dequantize_handles_block_scaling(): torch.testing.assert_close(out[0, 33], torch.tensor(1.5 * 8.0)) -def test_fp8_dequantize_handles_non_block_aligned(): - """When out/in is not a multiple of the block, the residual is dequantized - using the closest block scale (slice-and-truncate semantics).""" - from xorl.models.transformers.deepseek_v4.checkpoint_handler import _dequantize_fp8_block - - torch.manual_seed(2) - out_dim, in_dim = 130, 200 - block = (128, 128) - weight = (torch.randn(out_dim, in_dim) * 0.1).to(torch.float8_e4m3fn) - # 2 row-blocks (covering 0-128, 128-256 truncated at 130) and 2 col-blocks. - scale = torch.tensor([[1.0, 2.0], [4.0, 8.0]]).to(torch.float8_e8m0fnu) - out = _dequantize_fp8_block(weight, scale, block, torch.bfloat16) - assert out.shape == (out_dim, in_dim) - assert out.dtype == torch.bfloat16 - # Sanity: no NaNs / Infs since every block has finite scale and small weights. - assert torch.isfinite(out).all() - - # --------------------------------------------------------------------------- # End-to-end synthetic load # --------------------------------------------------------------------------- @@ -349,7 +243,7 @@ def _make_synthetic_hf_state_dict(cfg): return sd -def test_end_to_end_synthetic_load_window_only(): +def _assert_end_to_end_synthetic_load_window_only(): """Fully load a tiny 2-layer C0/C0 model from a fabricated HF state-dict.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM, load_hf_state_dict_into_model @@ -390,10 +284,48 @@ def test_end_to_end_synthetic_load_window_only(): sd["head.weight"], ) torch.testing.assert_close(model.model.hc_head_fn, sd["hc_head_fn"]) - torch.testing.assert_close(model.model.layers[0].self_attn.attn_sink, sd["layers.0.attn.attn_sink"]) + layer = model.model.layers[0] + torch.testing.assert_close( + layer.input_layernorm.weight, + sd["layers.0.attn_norm.weight"].to(layer.input_layernorm.weight.dtype), + ) + torch.testing.assert_close( + layer.post_attention_layernorm.weight, + sd["layers.0.ffn_norm.weight"].to(layer.post_attention_layernorm.weight.dtype), + ) + torch.testing.assert_close( + layer.self_attn.wq_a.weight, + sd["layers.0.attn.wq_a.weight"].to(layer.self_attn.wq_a.weight.dtype), + ) + torch.testing.assert_close(layer.self_attn.attn_sink, sd["layers.0.attn.attn_sink"]) + torch.testing.assert_close(layer.hc_attn_fn, sd["layers.0.hc_attn_fn"]) + torch.testing.assert_close( + layer.mlp.gate.e_score_correction_bias, + sd["layers.0.ffn.gate.bias"], + ) + torch.testing.assert_close( + layer.mlp.shared_experts.gate_proj.weight, + sd["layers.0.ffn.shared_experts.w1.weight"].to(layer.mlp.shared_experts.gate_proj.weight.dtype), + ) + expected_gate_up = torch.stack( + [ + torch.cat( + ( + sd[f"layers.0.ffn.experts.{expert}.w1.weight"].t(), + sd[f"layers.0.ffn.experts.{expert}.w3.weight"].t(), + ), + dim=-1, + ) + for expert in range(cfg.n_routed_experts) + ] + ) + torch.testing.assert_close( + layer.mlp.experts.gate_up_proj, + expected_gate_up.to(layer.mlp.experts.gate_up_proj.dtype), + ) -def test_checkpoint_handler_ep_filters_and_fuses_local_experts(): +def _assert_checkpoint_handler_ep_filters_and_fuses_local_experts(): """The generic distributed loader handler should emit only this EP rank's experts.""" cfg = _tiny_config(compress_ratios=[0]) sd = _make_synthetic_hf_state_dict(cfg) @@ -438,7 +370,7 @@ def test_checkpoint_handler_ep_filters_and_fuses_local_experts(): torch.testing.assert_close(down, torch.stack(expected_down_rows, dim=0)) -def test_checkpoint_handler_skips_mtp_even_without_ep_filter(): +def _assert_checkpoint_handler_skips_mtp_and_accounts_unknown_keys_without_ep_filter(): cfg = _tiny_config(compress_ratios=[0]) handler = DeepseekV4CheckpointHandler( cfg, @@ -454,8 +386,11 @@ def test_checkpoint_handler_skips_mtp_even_without_ep_filter(): assert skip_key_fn("mtp.0.ffn.experts.0.w1.weight") is True assert skip_key_fn("embed.weight") is False + assert handler.on_load_weight("totally.bogus.name", torch.zeros(1)) == [] + assert handler.summary.unmapped == ["totally.bogus.name"] + -def test_strict_load_ignores_nonpersistent_rope_buffers(): +def _assert_strict_load_ignores_nonpersistent_rope_buffers(): """``strict=True`` should not require config-derived RoPE cache buffers.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM, load_hf_state_dict_into_model @@ -469,7 +404,7 @@ def test_strict_load_ignores_nonpersistent_rope_buffers(): assert summary.missing_in_model == [] -def test_end_to_end_synthetic_load_with_c4_layer(): +def _assert_end_to_end_synthetic_load_with_c4_layer(): """C4 layer load exercises the indexer mapping + APE hotfix path.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM, load_hf_state_dict_into_model from xorl.models.transformers.deepseek_v4.checkpoint_handler import _undo_ape_hotfix @@ -497,9 +432,19 @@ def test_end_to_end_synthetic_load_with_c4_layer(): torch.testing.assert_close(model.model.layers[1].self_attn.compressor.ape, expected_compressor_ape) torch.testing.assert_close(model.model.layers[1].self_attn.indexer.compressor.ape, expected_indexer_ape) + torch.testing.assert_close( + model.model.layers[1].self_attn.indexer.linear_wq_b.weight, + sd["layers.1.attn.indexer.wq_b.weight"].to(model.model.layers[1].self_attn.indexer.linear_wq_b.weight.dtype), + ) + torch.testing.assert_close( + model.model.layers[1].self_attn.indexer.linear_weights_proj.weight, + sd["layers.1.attn.indexer.weights_proj.weight"].to( + model.model.layers[1].self_attn.indexer.linear_weights_proj.weight.dtype + ), + ) -def test_end_to_end_synthetic_load_with_hash_layer(): +def _assert_end_to_end_synthetic_load_with_hash_layer(): """First layer is hash-routed: tid2eid is filled, gate.bias is absent.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM, load_hf_state_dict_into_model @@ -522,3 +467,18 @@ def test_end_to_end_synthetic_load_with_hash_layer(): assert summary.unmapped == [] assert summary.missing_in_model == [] + + +def test_dsv4_checkpoint_codec_and_handler_ownership_contract(): + _assert_fp8_dequantize_matches_block_and_tail_references() + _assert_mxfp4_dequantize_known_values_and_block_scaling() + _assert_checkpoint_handler_ep_filters_and_fuses_local_experts() + _assert_checkpoint_handler_skips_mtp_and_accounts_unknown_keys_without_ep_filter() + _assert_dsv4_synthetic_load_contract() + + +def _assert_dsv4_synthetic_load_contract(): + _assert_end_to_end_synthetic_load_window_only() + _assert_strict_load_ignores_nonpersistent_rope_buffers() + _assert_end_to_end_synthetic_load_with_c4_layer() + _assert_end_to_end_synthetic_load_with_hash_layer() diff --git a/tests/models/test_dsv4_lora.py b/tests/models/test_dsv4_lora.py index 84f18c40..b8993f7c 100644 --- a/tests/models/test_dsv4_lora.py +++ b/tests/models/test_dsv4_lora.py @@ -15,7 +15,6 @@ @pytest.fixture(autouse=True) def _cpu_env(monkeypatch): - monkeypatch.setenv("XORL_DSV4_ROPE_MAX_SEQ_LEN", "256") monkeypatch.setenv("XORL_DSV4_SPARSE_ATTN_IMPL", "sparse") from xorl.ops.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 @@ -81,63 +80,7 @@ def _lm_logits(model, input_ids): ATTN_LORA_TARGETS = ["wq_a", "wq_b", "wkv", "wo_a", "wo_b"] -def test_attention_lora_inject_freezes_base_and_adds_adapters(): - """``wq_a/wq_b/wkv/wo_b`` get LoraLinear adapters.""" - from xorl.lora.utils import inject_lora_into_model - - torch.manual_seed(0) - _, model = _build_model() - - base_linear_count = sum( - 1 for n, m in model.named_modules() if isinstance(m, torch.nn.Linear) and n.split(".")[-1] in ATTN_LORA_TARGETS - ) - assert base_linear_count > 0 - - inject_lora_into_model(model, r=4, lora_alpha=8, target_modules=ATTN_LORA_TARGETS) - - # Every targeted Linear is now a LoraLinear. - from xorl.lora.modules.linear import LoraLinear - - targeted = [ - (n, m) - for n, m in model.named_modules() - if n.split(".")[-1] in ATTN_LORA_TARGETS and isinstance(m, (torch.nn.Linear, LoraLinear)) - ] - lora_targeted = [(n, m) for n, m in targeted if isinstance(m, LoraLinear)] - assert len(lora_targeted) == base_linear_count, f"expected {base_linear_count} LoraLinear, got {len(lora_targeted)}" - - # Base weight on a LoraLinear is frozen; lora_A / lora_B trainable. - # LoraLinear subclasses nn.Linear; lora_A and lora_B are Parameters. - name, mod = lora_targeted[0] - assert mod.weight.requires_grad is False, name - assert mod.lora_A.requires_grad is True - assert mod.lora_B.requires_grad is True - - -def test_moe_expert_lora_inject_freezes_base_and_adds_adapters(): - """``gate_proj/up_proj/down_proj`` on routed experts get LoRA via MoEExpertsLoRA.""" - from xorl.lora.utils import inject_lora_into_model - from xorl.models.layers.moe.lora import MoEExpertsLoRA - - torch.manual_seed(1) - _, model = _build_model() - - inject_lora_into_model(model, r=4, lora_alpha=8, target_modules=["gate_proj", "up_proj", "down_proj"]) - - # Every layer's mlp.experts should be MoEExpertsLoRA now. - for layer in model.model.layers: - assert isinstance(layer.mlp.experts, MoEExpertsLoRA), ( - f"layer {layer.layer_id} experts not LoRA: {type(layer.mlp.experts).__name__}" - ) - # Base experts weights frozen; LoRA params trainable. - assert layer.mlp.experts.gate_up_proj.requires_grad is False - # Find at least one LoRA param. - lora_params = [p for n, p in layer.mlp.experts.named_parameters() if "_lora_" in n] - assert len(lora_params) > 0 - assert all(p.requires_grad for p in lora_params) - - -def test_wo_a_lora_delta_actually_contributes_to_output(): +def _assert_wo_a_lora_delta_actually_contributes_to_output(): """Smoke that ``wo_a`` LoRA actually fires through the grouped einsum. With zero-inited ``lora_B`` the delta is zero, so a fresh-injection forward @@ -191,19 +134,37 @@ def test_wo_a_lora_delta_actually_contributes_to_output(): assert m.lora_B.grad is not None and m.lora_B.grad.abs().sum().item() > 0 -def test_lora_forward_backward_updates_only_lora_params(): - """End-to-end: only LoRA params receive grads after a forward+backward.""" +def test_dsv4_attention_lora_training_policy(): + """End-to-end attention injection freezes base weights and trains adapters.""" + _assert_wo_a_lora_delta_actually_contributes_to_output() + + from xorl.lora.modules.linear import LoraLinear from xorl.lora.utils import inject_lora_into_model torch.manual_seed(2) cfg, model = _build_model() + base_linear_count = sum( + 1 + for name, module in model.named_modules() + if isinstance(module, torch.nn.Linear) and name.split(".")[-1] in ATTN_LORA_TARGETS + ) inject_lora_into_model( model, r=4, lora_alpha=8, - target_modules=[*ATTN_LORA_TARGETS, "gate_proj", "up_proj", "down_proj"], + target_modules=ATTN_LORA_TARGETS, ) + attention_adapters = [ + (name, module) + for name, module in model.named_modules() + if name.split(".")[-1] in ATTN_LORA_TARGETS and isinstance(module, LoraLinear) + ] + assert len(attention_adapters) == base_linear_count > 0 + for name, module in attention_adapters: + assert module.weight.requires_grad is False, name + assert module.lora_A.requires_grad and module.lora_B.requires_grad + bsz, seqlen = 1, cfg.sliding_window input_ids = torch.randint(0, cfg.vocab_size, (bsz, seqlen), dtype=torch.long) targets = torch.randint(0, cfg.vocab_size, (bsz, seqlen), dtype=torch.long) diff --git a/tests/models/test_dsv4_model.py b/tests/models/test_dsv4_model.py index 532a7649..c20fbd26 100644 --- a/tests/models/test_dsv4_model.py +++ b/tests/models/test_dsv4_model.py @@ -19,7 +19,6 @@ @pytest.fixture(autouse=True) def _cpu_env(monkeypatch): - monkeypatch.setenv("XORL_DSV4_ROPE_MAX_SEQ_LEN", "256") monkeypatch.setenv("XORL_DSV4_SPARSE_ATTN_IMPL", "sparse") from xorl.ops.dsv4.rope import precompute_freqs_cis # noqa: PLC0415 @@ -85,7 +84,7 @@ def _lm_logits(model, input_ids): return model.lm_head(outputs.last_hidden_state) -def test_for_causal_lm_rejects_pipeline_parallelism(): +def _assert_for_causal_lm_rejects_pipeline_parallelism(): """DSv4 requires a dedicated PP forward because hyperconnection state is 4-D.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM # noqa: PLC0415 @@ -97,23 +96,7 @@ def test_for_causal_lm_rejects_pipeline_parallelism(): model.get_pp_module_config() -def test_model_forward_shape_window_only(): - """3 layers, all window-only (no compressor, no indexer).""" - from xorl.models.transformers.deepseek_v4 import DeepseekV4Model # noqa: PLC0415 - - torch.manual_seed(0) - cfg = _tiny_config(num_hidden_layers=3, compress_ratios=[0, 0, 0]) - model = _make_model(cfg, DeepseekV4Model) - - bsz, seqlen = 1, cfg.sliding_window - input_ids = torch.randint(0, cfg.vocab_size, (bsz, seqlen), dtype=torch.long) - - out = model(input_ids).last_hidden_state - assert out.shape == (bsz, seqlen, cfg.hidden_size) - assert torch.isfinite(out).all() - - -def test_model_forward_shape_with_c128(): +def _assert_model_forward_shape_with_c128(): """One C128 layer + two C0 layers.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4Model # noqa: PLC0415 @@ -129,7 +112,7 @@ def test_model_forward_shape_with_c128(): assert torch.isfinite(out).all() -def test_for_causal_lm_forward_backward(): +def _assert_for_causal_lm_forward_backward(): """Full forward + backward through the LM head.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM # noqa: PLC0415 @@ -180,7 +163,7 @@ def test_for_causal_lm_forward_backward(): assert grads[name].grad is None, f"{name} unexpectedly received a gradient" -def test_for_causal_lm_with_hash_layer(): +def _assert_for_causal_lm_with_hash_layer(): """First layer is hash-routed; verify input_ids threads through correctly.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM # noqa: PLC0415 @@ -208,7 +191,7 @@ def test_for_causal_lm_with_hash_layer(): assert torch.isfinite(model.model.layers[0].mlp.gate.weight.grad).all() -def test_keep_fp32_marks_propagated(): +def _assert_keep_fp32_marks_propagated(): """All HC + attn_sink + compressor fp32 params carry _keep_fp32 = True.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM # noqa: PLC0415 @@ -229,7 +212,7 @@ def test_keep_fp32_marks_propagated(): assert getattr(p, "_keep_fp32", False) is True, f"{name} missing _keep_fp32" -def test_from_config_wires_parallel_groups(monkeypatch): +def _assert_from_config_wires_parallel_groups(monkeypatch): """DSv4 factory passes xorl's sequence-parallel group into attention.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM # noqa: PLC0415 @@ -254,7 +237,38 @@ class FakeParallelState: assert model.config._attn_implementation == "native" -def test_from_config_dtype_cast_preserves_rope_and_fp32_params(monkeypatch): +def _assert_attention_qat_dispatch_and_tp_admission(monkeypatch): + from xorl.models.transformers.deepseek_v4 import modeling_deepseek_v4 # noqa: PLC0415 + from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import ( # noqa: PLC0415 + DeepSeekV4Attention, + DeepseekV4ForCausalLM, + ) + + cfg = _tiny_config(num_hidden_layers=1, compress_ratios=[0]) + cfg.quantization_config = {"quant_method": "fp8"} + qat_calls = [] + + def record_qat(tensor, block_size): + qat_calls.append((tuple(tensor.shape), block_size)) + return tensor.clone() + + monkeypatch.setattr(modeling_deepseek_v4, "fp8_simulate_qat", record_qat) + model = _make_model(cfg, DeepseekV4ForCausalLM) + input_ids = torch.randint(0, cfg.vocab_size, (1, cfg.sliding_window), dtype=torch.long) + model(input_ids) + + attention = model.model.layers[0].self_attn + assert qat_calls == [((1, cfg.sliding_window, attention.nope_head_dim), 64)] + + class _TwoRankGroup: + def size(self): + return 2 + + with pytest.raises(AssertionError, match="TP > 1 is not implemented"): + DeepSeekV4Attention(cfg, layer_id=0, tp_group=_TwoRankGroup()) + + +def _assert_from_config_dtype_cast_preserves_rope_and_fp32_params(monkeypatch): """Registry construction keeps complex RoPE caches and fp32-only params intact.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM # noqa: PLC0415 @@ -288,7 +302,7 @@ class FakeParallelState: assert buf.imag.abs().max() > 0, f"{name} lost its imaginary RoPE component" -def test_direct_to_dtype_preserves_rope_and_fp32_params(): +def _assert_direct_to_dtype_preserves_rope_and_fp32_params(): """Direct DSv4 dtype casts share the same RoPE/fp32 carve-outs as registry construction.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4ForCausalLM # noqa: PLC0415 @@ -307,7 +321,7 @@ def test_direct_to_dtype_preserves_rope_and_fp32_params(): assert buf.imag.abs().max() > 0, f"{name} lost its imaginary RoPE component" -def test_outer_gradient_checkpointing_wraps_decoder_layers(): +def _assert_outer_gradient_checkpointing_wraps_decoder_layers(): """DSv4Model has the checkpoint flag consumed by the decoder loop.""" from xorl.models.module_utils import DEFAULT_GRADIENT_CHECKPOINTING_METHOD # noqa: PLC0415 from xorl.models.transformers.deepseek_v4 import DeepseekV4Model # noqa: PLC0415 @@ -332,3 +346,32 @@ def fake_checkpoint(func, *args, **kwargs): assert out.shape == (1, cfg.sliding_window, cfg.hidden_size) assert len(calls) == cfg.num_hidden_layers assert all(call_kwargs["input_ids"] is input_ids for _, _, call_kwargs in calls) + + +def test_dsv4_model_construction_topology_precision_and_runtime_contract(monkeypatch): + with monkeypatch.context() as construction_patch: + _assert_dsv4_model_construction_topology_and_precision_contract(construction_patch) + _assert_dsv4_model_runtime_contract() + + +def _assert_dsv4_model_construction_topology_and_precision_contract(monkeypatch): + _assert_for_causal_lm_rejects_pipeline_parallelism() + with monkeypatch.context() as topology_patch: + _assert_from_config_wires_parallel_groups(topology_patch) + with monkeypatch.context() as attention_patch: + _assert_attention_qat_dispatch_and_tp_admission(attention_patch) + with monkeypatch.context() as precision_patch: + _assert_dsv4_model_precision_preservation_contract(precision_patch) + + +def _assert_dsv4_model_runtime_contract(): + _assert_model_forward_shape_with_c128() + _assert_for_causal_lm_forward_backward() + _assert_for_causal_lm_with_hash_layer() + _assert_outer_gradient_checkpointing_wraps_decoder_layers() + + +def _assert_dsv4_model_precision_preservation_contract(monkeypatch): + _assert_keep_fp32_marks_propagated() + _assert_from_config_dtype_cast_preserves_rope_and_fp32_params(monkeypatch) + _assert_direct_to_dtype_preserves_rope_and_fp32_params() diff --git a/tests/models/test_dsv4_moe.py b/tests/models/test_dsv4_moe.py index 2bda5fe4..398bfade 100644 --- a/tests/models/test_dsv4_moe.py +++ b/tests/models/test_dsv4_moe.py @@ -76,35 +76,7 @@ def _tiny_config(*, num_hash_layers=0): # --------------------------------------------------------------------------- -def test_mlp_forward_shape_and_swiglu_clamp(): - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MLP - - cfg = _tiny_config() - cfg.swiglu_limit = 1.0 # force a tight clamp so we can detect it - mlp = DeepseekV4MLP(cfg).to(torch.float32) - x = torch.randn(2, 3, cfg.hidden_size) - - # Force a huge gate magnitude: zero gate weights, set bias to a large value - # via attribute write. Easier: just run with the clamp on and verify finite. - out = mlp(x) - assert out.shape == x.shape - assert torch.isfinite(out).all() - - # With clamp off, output should be different. - cfg2 = _tiny_config() - cfg2.swiglu_limit = 0.0 - mlp2 = DeepseekV4MLP(cfg2).to(torch.float32) - mlp2.load_state_dict(mlp.state_dict(), strict=False) - out_no_clamp = mlp2(x) - - # Sanity: the test infrastructure runs both branches without error. We - # don't assert numerical inequality because random init may keep gates in - # the [-1, 1] range where the clamp is a no-op; the *behavior* of the - # clamp is exercised explicitly below. - assert out_no_clamp.shape == x.shape - - -def test_mlp_swiglu_limit_actually_clamps(): +def _assert_mlp_and_routed_experts_swiglu_limit_contract(): """Construct a gate that pushes beyond the limit, verify clamp clips it.""" from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MLP @@ -126,10 +98,13 @@ def test_mlp_swiglu_limit_actually_clamps(): # silu(0.5) ≈ 0.31; up_proj output ≈ hidden_size; intermediate dim = 16. # Output magnitude is bounded by silu(0.5) * hidden_size * intermediate_size. bound = 0.32 * cfg.hidden_size * cfg.moe_intermediate_size + assert out.shape == x.shape assert out.abs().max().item() < bound * 1.5, out.abs().max().item() + _assert_routed_experts_swiglu_limit_clamps_eager_backend() -def test_routed_experts_swiglu_limit_clamps_eager_backend(): + +def _assert_routed_experts_swiglu_limit_clamps_eager_backend(): """DeepSeek-V4 propagates swiglu_limit into routed experts, not only shared experts.""" from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MoE @@ -160,31 +135,18 @@ def test_routed_experts_swiglu_limit_clamps_eager_backend(): # --------------------------------------------------------------------------- -def test_moe_non_hash_has_bias_no_table(): - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MoE - - cfg = _tiny_config(num_hash_layers=2) - # layer_id = 5 is past the hash band (which is layers 0 and 1). - block = DeepseekV4MoE(cfg, layer_id=5) - - assert hasattr(block.gate, "e_score_correction_bias") - assert block.gate.e_score_correction_bias.shape == (cfg.n_routed_experts,) - # ``e_score_correction_bias`` is frozen (requires_grad=False) — gradients - # never flow through it (selection-only argmax bias). DeepSeek updates it - # OOB via an aux-loss controller during training. - assert block.gate.e_score_correction_bias.requires_grad is False - assert "tid2eid" not in block._buffers - assert block.is_hash_layer is False - assert block.shared_experts is not None # n_shared_experts=1 - - -def test_moe_non_hash_forward_backward(): +def test_moe_non_hash_structure_forward_backward_shared_and_swiglu_policy(): from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MoE torch.manual_seed(0) cfg = _tiny_config(num_hash_layers=0) block = DeepseekV4MoE(cfg, layer_id=0).to(torch.float32) _init_test_weights(block) + assert block.is_hash_layer is False + assert "tid2eid" not in block._buffers + assert block.shared_experts is not None + assert block.gate.e_score_correction_bias.shape == (cfg.n_routed_experts,) + assert block.gate.e_score_correction_bias.requires_grad is False x = torch.randn(2, 3, cfg.hidden_size, requires_grad=True) out, router_logits = block(x) @@ -192,6 +154,12 @@ def test_moe_non_hash_forward_backward(): assert router_logits.shape == (x.shape[0] * x.shape[1], cfg.n_routed_experts) assert torch.isfinite(out).all() + shared_experts = block.shared_experts + block.shared_experts = None + out_without_shared, _ = block(x) + block.shared_experts = shared_experts + assert (out - out_without_shared).abs().max().item() > 1e-6 + out.sum().backward() assert torch.isfinite(x.grad).all() # Gate must receive grads. @@ -203,25 +171,9 @@ def test_moe_non_hash_forward_backward(): bias = block.gate.e_score_correction_bias assert bias.grad is None or torch.equal(bias.grad, torch.zeros_like(bias.grad)) - -def test_moe_shared_expert_contributes(): - """Sanity: removing the shared expert changes the output.""" - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MoE - - torch.manual_seed(1) - cfg = _tiny_config() - block = DeepseekV4MoE(cfg, layer_id=0).to(torch.float32) - _init_test_weights(block) - x = torch.randn(1, 4, cfg.hidden_size) - - out_with, _ = block(x) - saved = block.shared_experts - block.shared_experts = None - out_without, _ = block(x) - block.shared_experts = saved - - diff = (out_with - out_without).abs().max().item() - assert diff > 1e-6, f"shared expert had no effect: max diff {diff}" + _assert_mlp_and_routed_experts_swiglu_limit_contract() + _assert_moe_hash_layer_structure_admission_and_table_forward_backward() + _assert_moe_hash_layer_record_then_replay_backward_matches() # --------------------------------------------------------------------------- @@ -229,28 +181,19 @@ def test_moe_shared_expert_contributes(): # --------------------------------------------------------------------------- -def test_moe_hash_layer_has_table_no_bias(): +def _assert_moe_hash_layer_structure_admission_and_table_forward_backward(): + """Hash-layer structure, input admission, table routing, and gradients form one transaction.""" from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MoE + torch.manual_seed(2) cfg = _tiny_config(num_hash_layers=3) - block = DeepseekV4MoE(cfg, layer_id=0) # in the hash band - + block = DeepseekV4MoE(cfg, layer_id=1).to(torch.float32) + _init_test_weights(block) assert block.is_hash_layer is True - assert hasattr(block, "tid2eid") assert block.tid2eid.shape == (cfg.vocab_size, cfg.num_experts_per_tok) assert block.tid2eid.dtype == torch.int32 assert not hasattr(block.gate, "e_score_correction_bias") - -def test_moe_hash_layer_uses_table_for_selection(): - """Forward + verify selected_experts come from tid2eid[input_ids].""" - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MoE - - torch.manual_seed(2) - cfg = _tiny_config(num_hash_layers=3) - block = DeepseekV4MoE(cfg, layer_id=1).to(torch.float32) - _init_test_weights(block) - # Build a deterministic table: token id i -> experts (i % E, (i + 1) % E). E = cfg.n_routed_experts table = torch.stack( @@ -262,6 +205,8 @@ def test_moe_hash_layer_uses_table_for_selection(): bsz, seqlen = 1, 4 input_ids = torch.tensor([[3, 7, 0, 11]], dtype=torch.long) x = torch.randn(bsz, seqlen, cfg.hidden_size, requires_grad=True) + with pytest.raises(AssertionError, match="hash-routed layer requires input_ids"): + block(x, input_ids=None) out, _ = block(x, input_ids=input_ids) assert out.shape == x.shape @@ -273,24 +218,12 @@ def test_moe_hash_layer_uses_table_for_selection(): assert torch.isfinite(block.gate.weight.grad).all() -def test_moe_hash_layer_requires_input_ids(): - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MoE - - cfg = _tiny_config(num_hash_layers=3) - block = DeepseekV4MoE(cfg, layer_id=0).to(torch.float32) - _init_test_weights(block) - x = torch.randn(1, 4, cfg.hidden_size) - - with pytest.raises(AssertionError, match="hash-routed layer requires input_ids"): - block(x, input_ids=None) - - # --------------------------------------------------------------------------- # Routing-replay × hash-routed layer # --------------------------------------------------------------------------- -def test_moe_hash_layer_record_then_replay_backward_matches(): +def _assert_moe_hash_layer_record_then_replay_backward_matches(): """Record on a hash-routed layer, then replay_backward — selected_experts on backward must come from the table-driven recording, not be recomputed. @@ -357,32 +290,3 @@ def test_moe_hash_layer_record_then_replay_backward_matches(): finally: set_replay_stage(None) RoutingReplay.clear_all() - - -def test_moe_route_unknown_replay_stage_raises(): - """A new replay stage name (defensively) raises rather than NameError'ing - on an undefined ``selected_experts``.""" - from xorl.models.layers.moe.routing_replay import ( - RoutingReplay, - set_replay_stage, - ) - from xorl.models.transformers.deepseek_v4.modeling_deepseek_v4 import DeepseekV4MoE - - cfg = _tiny_config(num_hash_layers=0) - block = DeepseekV4MoE(cfg, layer_id=1).to(torch.float32) - _init_test_weights(block) - block._routing_replay = RoutingReplay() - - x = torch.randn(1, 4, cfg.hidden_size) - try: - # Bypass set_replay_stage's validator (which constrains values) by - # poking the module-level state directly. ``route`` reads the same - # global, so it sees the bogus value. - from xorl.models.layers.moe import routing_replay as _replay_mod - - _replay_mod._replay_stage = "future_stage_name" - with pytest.raises(ValueError, match="Unrecognized replay stage"): - block(x, input_ids=None) - finally: - set_replay_stage(None) - RoutingReplay.clear_all() diff --git a/tests/models/test_fused_gdn_lora.py b/tests/models/test_fused_gdn_lora.py index 13c2140f..e2676a1d 100644 --- a/tests/models/test_fused_gdn_lora.py +++ b/tests/models/test_fused_gdn_lora.py @@ -1,9 +1,11 @@ from __future__ import annotations +import copy import gc import json import weakref +import pytest import torch import torch.nn as nn import torch.nn.functional as F @@ -11,7 +13,12 @@ from xorl.lora.fold import canonical_lora_fold_linear from xorl.lora.modules.delta_linear import LoraDeltaLinear -from xorl.lora.target_manifest import collect_lora_runtime_modules +from xorl.lora.target_manifest import ( + collect_lora_runtime_modules, + load_lora_target_manifest, + resolve_lora_target_modules, + validate_lora_target_manifest, +) from xorl.lora.utils import ( inject_lora_into_model, load_lora_checkpoint, @@ -81,6 +88,7 @@ def test_river_fused_gdn_geometry_and_manifest_path_filter(tmp_path): "model.layers.0.linear_attn.in_proj_qkvz": 2, "model.layers.0.linear_attn.out_proj": 2, } + _assert_manifest_schema_and_runtime_mismatches_fail_closed() with torch.no_grad(): gdn.in_proj_qkvz.lora_B.normal_() @@ -112,9 +120,55 @@ def test_river_fused_gdn_geometry_and_manifest_path_filter(tmp_path): } config = json.loads((tmp_path / "adapter_config.json").read_text()) assert sorted(config["target_modules"]) == ["in_proj_qkvz", "out_proj"] + _assert_fused_gdn_delta_product_merged_forward_and_gradient_policy() + sharded_root = tmp_path / "sharded" + sharded_root.mkdir() + _assert_sharded_peft_checkpoint_loads_into_fused_gdn(sharded_root) + + +def _assert_manifest_schema_and_runtime_mismatches_fail_closed(): + wrong_count = copy.deepcopy(_manifest()) + wrong_count["expected_modules"][0]["count"] = 2 + with pytest.raises(ValueError, match="matched 1 modules, expected 2"): + inject_lora_into_model(_Model(), r=2, lora_alpha=4, target_manifest=wrong_count) + + model = _Model() + inject_lora_into_model(model, r=2, lora_alpha=4, target_manifest=_manifest()) + wrong_rank = copy.deepcopy(_manifest()) + wrong_rank["expected_modules"][0]["rank"] = 4 + with pytest.raises(ValueError, match="rank mismatch"): + validate_lora_target_manifest(model, wrong_rank) + with pytest.raises(ValueError, match="do not match"): + resolve_lora_target_modules(["in_proj_qkvz", "out_proj"], _manifest()) -def test_delta_linear_matches_explicit_low_rank_product(): + model = _Model() + inject_lora_into_model( + model, + r=2, + lora_alpha=4, + target_modules=["down_proj", "in_proj_qkvz", "out_proj"], + ) + with pytest.raises(ValueError, match="unlisted LoRA modules"): + validate_lora_target_manifest(model, _manifest()) + + for field, value, message in ( + ("schema_version", True, "schema_version"), + ("allow_unlisted", "false", "allow_unlisted must be a Boolean"), + ): + manifest = copy.deepcopy(_manifest()) + manifest[field] = value + with pytest.raises(ValueError, match=message): + load_lora_target_manifest(manifest) + + for field in ("count", "rank"): + manifest = copy.deepcopy(_manifest()) + manifest["expected_modules"][0][field] = True + with pytest.raises(ValueError, match=field): + load_lora_target_manifest(manifest) + + +def _assert_delta_linear_matches_explicit_low_rank_product(): module = LoraDeltaLinear(8, 12, r=2, lora_alpha=4) with torch.no_grad(): module.lora_B.normal_() @@ -124,7 +178,9 @@ def test_delta_linear_matches_explicit_low_rank_product(): assert torch.allclose(module.get_delta_weight(), module.lora_B @ module.lora_A * 2) -def test_fused_gdn_delta_merged_forward_uses_canonical_fold_and_keeps_gradients(): +def _assert_fused_gdn_delta_product_merged_forward_and_gradient_policy(): + _assert_delta_linear_matches_explicit_low_rank_product() + module = LoraDeltaLinear(8, 12, r=2, lora_alpha=4, dtype=torch.float32) with torch.no_grad(): module.lora_B.normal_() @@ -143,8 +199,26 @@ def test_fused_gdn_delta_merged_forward_uses_canonical_fold_and_keeps_gradients( assert torch.count_nonzero(module.lora_B.grad[:2]) == 0 assert torch.count_nonzero(module.lora_B.grad[7:]) == 0 + model = _Model() + inject_lora_into_model(model, r=2, lora_alpha=4, target_manifest=_manifest()) + gdn = model.model.layers[0].linear_attn + gdn.exact_merged_forward = True + with torch.no_grad(): + gdn.out_proj.lora_B.normal_() + output_inputs = torch.randn(2, 3, gdn.o_proj.in_features) + expected_weight = canonical_lora_fold_linear( + gdn.o_proj.weight, + gdn.out_proj.lora_A, + gdn.out_proj.lora_B, + 2.0, + ) + expected_output = F.linear(output_inputs, expected_weight, gdn.o_proj.bias) + assert torch.equal(gdn._project_output_linear(output_inputs), expected_output) + + _assert_fused_gdn_merged_weight_cache_is_bounded_and_releases_previous_generation() + -def test_fused_gdn_merged_weight_cache_is_bounded_by_current_slices(): +def _assert_fused_gdn_merged_weight_cache_is_bounded_and_releases_previous_generation(): module = LoraDeltaLinear(8, 12, r=2, lora_alpha=4, dtype=torch.float32) first_base = torch.randn(5, 8) second_base = torch.randn(7, 8) @@ -162,16 +236,9 @@ def test_fused_gdn_merged_weight_cache_is_bounded_by_current_slices(): current_weights = [entry["weight"] for entry in module._merged_weight_cache["slices"].values()] assert all(weight is not retained for weight in previous_weights[:-2] for retained in current_weights) - -def test_fused_gdn_merged_weight_cache_releases_previous_request_generation(): - module = LoraDeltaLinear(8, 12, r=2, lora_alpha=4, dtype=torch.float32) - first_base = torch.randn(5, 8) - second_base = torch.randn(7, 8) - - first = module._merged_weight(first_base, output_start=0, output_end=5) - second = module._merged_weight(second_base, output_start=5, output_end=12) - previous_generation = (weakref.ref(first), weakref.ref(second)) - del first, second + previous_generation = tuple(weakref.ref(weight) for weight in current_weights) + previous_weights.clear() + del first, second, current_weights # Match AdapterManager.prepare_forward(): copying adapter values into the # model bumps parameter versions at each serialized request boundary. @@ -185,26 +252,7 @@ def test_fused_gdn_merged_weight_cache_releases_previous_request_generation(): assert all(reference() is None for reference in previous_generation) -def test_fused_gdn_output_projection_merged_forward_matches_canonical_fold(): - model = _Model() - inject_lora_into_model(model, r=2, lora_alpha=4, target_manifest=_manifest()) - gdn = model.model.layers[0].linear_attn - gdn.exact_merged_forward = True - with torch.no_grad(): - gdn.out_proj.lora_B.normal_() - inputs = torch.randn(2, 3, gdn.o_proj.in_features) - expected_weight = canonical_lora_fold_linear( - gdn.o_proj.weight, - gdn.out_proj.lora_A, - gdn.out_proj.lora_B, - 2.0, - ) - expected = F.linear(inputs, expected_weight, gdn.o_proj.bias) - actual = gdn._project_output_linear(inputs) - assert torch.equal(actual, expected) - - -def test_sharded_peft_checkpoint_loads_into_fused_gdn(tmp_path): +def _assert_sharded_peft_checkpoint_loads_into_fused_gdn(tmp_path): source = _Model() inject_lora_into_model(source, r=2, lora_alpha=4, target_manifest=_manifest()) with torch.no_grad(): diff --git a/tests/models/test_glm52_contract.py b/tests/models/test_glm52_contract.py index e73819cc..2691f9f0 100644 --- a/tests/models/test_glm52_contract.py +++ b/tests/models/test_glm52_contract.py @@ -9,7 +9,6 @@ CanonicalMoEGraphMetadata, CanonicalMoETransport, ParallelPlan, - canonical_moe_reduce_reference, ) from xorl.models.transformers.glm5 import indexer as indexer_module from xorl.models.transformers.glm5 import sparse_selector as sparse_selector_module @@ -38,8 +37,6 @@ ) from xorl.models.transformers.glm5.sparse_selector import ( GLM52_SELECTOR_VERSION, - gather_selected_logical_values, - physical_cache_to_logical_indices, quantize_e4m3_dynamic, quantize_e4m3_ue8m0, quantize_sparse_key_cache, @@ -49,6 +46,16 @@ ) +def _canonical_moe_reference(partials: torch.Tensor, metadata: CanonicalMoEGraphMetadata) -> torch.Tensor: + level = [partials[index] for index in range(partials.shape[0])] + while len(level) > 1: + level = [(level[index] + level[index + 1]).to(torch.bfloat16) for index in range(0, len(level), 2)] + result = level[0] + result = result.clone() + result[~metadata.valid_mask] = 0 + return result + + GLM52_FULL_INDEX_LAYERS = (0, 1, 2, 6, 10, 14, 18, 22, 26, 30, 34, 38, 42, 46, 50, 54, 58, 62, 66, 70, 74) @@ -64,8 +71,7 @@ def _map_tensors(fn, value): return value -@pytest.mark.cpu -def test_canonical_trainer_admits_only_certified_world16_ep16_cp16(): +def _assert_canonical_trainer_admits_only_certified_world16_ep16_cp16(): assert _GLM52_CANONICAL_TRAINER_TOPOLOGIES == ((16, 1, 1, 1),) plan = ParallelPlan.glm52_trainer() assert (plan.world_size, plan.pp_size, plan.tp_size, plan.dp_size, plan.ep_size, plan.cp_size) == ( @@ -114,24 +120,26 @@ def _official_schedule_config() -> SimpleNamespace: ) -@pytest.mark.cpu -def test_official_layer_plan_counts_producers_and_38_40_split(): +def _assert_official_layer_plan_counts_producers_and_38_40_split(): plan = Glm52LayerPlan.from_config( _official_schedule_config(), pipeline_layer_ranges=((0, 38), (38, 78)), ) - assert plan.full_indexer_layers == GLM52_FULL_INDEX_LAYERS - assert len(plan.full_indexer_layers) == 21 - assert len(plan.shared_indexer_layers) == 57 - assert plan.dense_layers == (0, 1, 2) - assert len(plan.sparse_layers) == 75 + full_layers = tuple(layer.layer_index for layer in plan.layers if layer.indexer_type.value == "full") + shared_layers = tuple(layer.layer_index for layer in plan.layers if layer.indexer_type.value == "shared") + dense_layers = tuple(layer.layer_index for layer in plan.layers if layer.mlp_type.value == "dense") + sparse_layers = tuple(layer.layer_index for layer in plan.layers if layer.mlp_type.value == "sparse") + assert full_layers == GLM52_FULL_INDEX_LAYERS + assert len(full_layers) == 21 + assert len(shared_layers) == 57 + assert dense_layers == (0, 1, 2) + assert len(sparse_layers) == 75 assert plan.layers[37].index_producer_layer == 34 assert plan.layers[38].index_producer_layer == 38 assert plan.layers[77].index_producer_layer == 74 -@pytest.mark.cpu -def test_layer_plan_rejects_malformed_schedules_and_shared_stage_start(): +def _assert_layer_plan_rejects_malformed_schedules_and_shared_stage_start(): config = _official_schedule_config() config.indexer_types = config.indexer_types[:-1] with pytest.raises(ValueError, match="indexer_types has length"): @@ -161,8 +169,7 @@ def _small_plan() -> Glm52LayerPlan: return Glm52LayerPlan.from_config(config) -@pytest.mark.cpu -def test_index_share_context_lifecycle_reuse_exception_cleanup_and_concurrency_guard(): +def _assert_index_share_context_lifecycle_reuse_exception_cleanup_and_concurrency_guard(): plan = _small_plan() manager = IndexShareContextManager(plan, (0, 4)) payload = CanonicalLogicalIndices(torch.tensor([[[0, 1, -1]]], dtype=torch.int32)) @@ -178,40 +185,23 @@ def test_index_share_context_lifecycle_reuse_exception_cleanup_and_concurrency_g assert first.lifecycle is IndexShareLifecycle.CLOSED assert manager.active is None - with pytest.raises(RuntimeError, match="body failed"): - with manager.invocation(mode=IndexShareMode.FORWARD_ONLY) as second: - second.get_or_publish(producer_layer_index=0, layer_plan=plan, produce_payload=lambda: payload) - raise RuntimeError("body failed") + second = manager.begin(mode=IndexShareMode.FORWARD_ONLY) + second.get_or_publish(producer_layer_index=0, layer_plan=plan, produce_payload=lambda: payload) + manager.finish_forward(second, succeeded=False) assert manager.active is None - with manager.invocation(mode=IndexShareMode.FORWARD_ONLY) as third: - with pytest.raises(RuntimeError, match="has not published"): - third.require(producer_layer_index=0, layer_plan=plan) - - -@pytest.mark.cpu -def test_index_share_context_identity_survives_fsdp_tensor_transform(): - plan = _small_plan() - manager = IndexShareContextManager(plan, (0, 4)) - payload = CanonicalLogicalIndices(torch.tensor([[[0, 1, -1]]], dtype=torch.int32)) - - with manager.invocation(mode=IndexShareMode.FORWARD_ONLY) as context: - producer_context = _map_tensors(lambda tensor: tensor, {"context": context})["context"] - assert producer_context is context - producer_context.get_or_publish(producer_layer_index=0, layer_plan=plan, produce_payload=lambda: payload) - - consumer_context = _map_tensors(lambda tensor: tensor, {"context": context})["context"] - assert consumer_context is context - assert consumer_context.require(producer_layer_index=0, layer_plan=plan) is producer_context.get_or_publish( - producer_layer_index=0, - layer_plan=plan, - produce_payload=lambda: pytest.fail("retained producer payload was recomputed"), - ) + third = manager.begin(mode=IndexShareMode.FORWARD_ONLY) + with pytest.raises(RuntimeError, match="has not published"): + third.require(producer_layer_index=0, layer_plan=plan) + manager.finish_forward(third, succeeded=True) + assert manager.active is None @pytest.mark.cpu @torch.no_grad() -def test_index_share_survives_fsdp_cast_across_dense_producer_and_shared_consumer(monkeypatch): +def test_index_share_lifecycle_and_fsdp_identity_contract(monkeypatch): + _assert_index_share_context_lifecycle_reuse_exception_cleanup_and_concurrency_guard() + monkeypatch.setattr( indexer_module, "bi_bf16_fp32_linear", @@ -299,8 +289,7 @@ def _small_glm_config() -> Glm5Config: ) -@pytest.mark.cpu -def test_only_full_layers_allocate_indexer_parameters_and_strict_state_dict_round_trip(): +def _assert_only_full_layers_allocate_indexer_parameters_and_strict_state_dict_round_trip(): config = _small_glm_config() plan = Glm52LayerPlan.from_config(config) attentions = nn.ModuleList([Glm5Attention(config, layer, layer_plan=plan) for layer in range(4)]) @@ -316,8 +305,7 @@ def test_only_full_layers_allocate_indexer_parameters_and_strict_state_dict_roun clone.load_state_dict(state, strict=True) -@pytest.mark.cpu -def test_sparse_selector_ties_short_rows_dead_rows_and_logical_cache_mapping(): +def _assert_sparse_selector_ties_short_rows_and_dead_rows(): query = torch.zeros((1, 3, 2, 128), dtype=torch.bfloat16) key = torch.zeros((1, 5, 128), dtype=torch.bfloat16) weights = torch.ones((1, 3, 2), dtype=torch.float32) @@ -344,19 +332,8 @@ def test_sparse_selector_ties_short_rows_dead_rows_and_logical_cache_mapping(): assert result.logical_indices.values.tolist() == [[[0, 1, 2], [0, -1, -1], [-1, -1, -1]]] assert result.valid_counts.tolist() == [[3, 1, 0]] - values = torch.randn((1, 5, 4), dtype=torch.bfloat16) - gathered = gather_selected_logical_values(values, result.logical_indices) - assert torch.count_nonzero(gathered[0, 2]) == 0 - assert bool(torch.all(torch.isfinite(gathered))) - physical = torch.tensor([[[2, 0, -1]]], dtype=torch.int32) - page_map = torch.tensor([4, 1, 3], dtype=torch.int32) - logical = physical_cache_to_logical_indices(physical, page_map) - assert logical.values.tolist() == [[[3, 4, -1]]] - - -@pytest.mark.cpu -def test_glm52_sparse_shared_selector_handles_production_boundary_ties_and_dead_tail(): +def _assert_glm52_sparse_shared_selector_handles_production_boundary_ties_and_dead_tail(): query = torch.zeros((1, 1, 1, 128), dtype=torch.bfloat16) key = torch.zeros((1, 4112, 128), dtype=torch.bfloat16) weights = torch.ones((1, 1, 1), dtype=torch.float32) @@ -377,8 +354,7 @@ def test_glm52_sparse_shared_selector_handles_production_boundary_ties_and_dead_ assert result.valid_counts.tolist() == [[2048]] -@pytest.mark.cpu -def test_glm52_sparse_hadamard_transport_is_normalized_and_self_inverse(): +def _assert_glm52_sparse_hadamard_transport_is_normalized_and_self_inverse(): source = torch.arange(128, dtype=torch.float32).sub(63.5).to(torch.bfloat16).reshape(1, 1, 128) rotated = rotate_sparse_selector_activation(source) restored = rotate_sparse_selector_activation(rotated) @@ -388,8 +364,7 @@ def test_glm52_sparse_hadamard_transport_is_normalized_and_self_inverse(): torch.testing.assert_close(restored.float(), source.float(), atol=1.0, rtol=0.0) -@pytest.mark.cpu -def test_sparse_selector_applies_hadamard_before_fp8_quantization(monkeypatch): +def _assert_sparse_selector_applies_hadamard_before_fp8_quantization(monkeypatch): torch.manual_seed(1) query = torch.randn((1, 1, 2, 128), dtype=torch.bfloat16) key = torch.randn((1, 128, 128), dtype=torch.bfloat16) @@ -427,8 +402,7 @@ def record_key(tensor, **kwargs): assert torch.equal(observed["key"].view(torch.uint8), expected_key.view(torch.uint8)) -@pytest.mark.cpu -def test_sparse_selector_fused_contract_skips_hadamard_before_fp8_quantization(monkeypatch): +def _assert_sparse_selector_fused_contract_skips_hadamard_before_fp8_quantization(monkeypatch): torch.manual_seed(2) query = torch.randn((1, 1, 2, 128), dtype=torch.bfloat16) key = torch.randn((1, 128, 128), dtype=torch.bfloat16) @@ -464,8 +438,7 @@ def record_key(tensor, **kwargs): assert torch.equal(observed["key"].view(torch.uint8), key.view(torch.uint8)) -@pytest.mark.cpu -def test_fused_bf16_indexer_projection_matches_sampler_row_order(monkeypatch): +def _assert_fused_bf16_indexer_projection_matches_sampler_row_order(monkeypatch): torch.manual_seed(3) hidden = torch.randn((1, 5, 16), dtype=torch.bfloat16) wk = torch.randn((128, 16), dtype=torch.bfloat16) @@ -495,8 +468,7 @@ def recording_linear(input, weight, bias=None): assert torch.equal(raw_gate.view(torch.uint8), expected[..., 128:].view(torch.uint8)) -@pytest.mark.cpu -def test_fused_bf16_indexer_head_gate_scaling_promotes_before_head_scale(): +def _assert_fused_bf16_indexer_head_gate_scaling_promotes_before_head_scale(): raw_gate = torch.tensor([-6.4375, 11.0, -7.65625, -8.125], dtype=torch.bfloat16) scaled = _scale_fused_bf16_indexer_head_gates(raw_gate, index_n_heads=32) @@ -513,8 +485,7 @@ def test_fused_bf16_indexer_head_gate_scaling_promotes_before_head_scale(): assert torch.equal(native_score_weights, sampler_formula) -@pytest.mark.cpu -def test_fused_sampler_index_k_prepare_preserves_projection_stride_and_builds_literal_rope_cache(): +def _assert_fused_sampler_index_k_prepare_preserves_projection_stride_and_builds_literal_rope_cache(): backing = torch.arange(3 * 160, dtype=torch.int32).reshape(1, 3, 160).to(torch.bfloat16) raw_key = backing[..., :128] norm_weight = torch.ones((128,), dtype=torch.float32) @@ -558,8 +529,7 @@ def recording_kernel(key, weight, bias, eps, cos_sin_cache, positions): assert torch.equal(prepared, raw_key) -@pytest.mark.cpu -def test_sampler_index_k_preparation_uses_split_prompt_and_fused_decode_suffix(): +def _assert_sampler_index_k_preparation_uses_split_prompt_and_fused_decode_suffix(): split = torch.full((1, 8, 128), 7, dtype=torch.bfloat16) backing = torch.arange(8 * 160, dtype=torch.int32).reshape(1, 8, 160).to(torch.bfloat16) raw = backing[..., :128] @@ -591,8 +561,7 @@ def mark_fused_suffix(key, _weight, _bias, _eps, _cache, _positions): assert torch.equal(mixed[:, 4:], torch.full_like(mixed[:, 4:], 11)) -@pytest.mark.cpu -def test_sampler_index_k_preparation_maps_4096_boundary_across_cp16(): +def _assert_sampler_index_k_preparation_maps_4096_boundary_across_cp16(): local_length = 260 split = torch.zeros((1, local_length, 128), dtype=torch.bfloat16) raw = torch.zeros_like(split) @@ -630,8 +599,7 @@ def record_suffix(key, _weight, _bias, _eps, _cache, _positions): assert torch.equal(mixed[:, 196:], torch.ones_like(mixed[:, 196:])) -@pytest.mark.cpu -def test_glm52_sparse_query_and_key_codecs_have_distinct_scale_contracts(): +def _assert_glm52_sparse_query_and_key_codecs_have_distinct_scale_contracts(): source = torch.linspace(-3.0, 2.0, 128, dtype=torch.float32).to(torch.bfloat16).reshape(1, 128) _, query_scale = quantize_e4m3_ue8m0(source) _, key_scale = quantize_e4m3_dynamic(source) @@ -642,8 +610,7 @@ def test_glm52_sparse_query_and_key_codecs_have_distinct_scale_contracts(): assert key_scale.item() != query_scale.item() -@pytest.mark.cpu -def test_glm52_sparse_key_cache_unpack_preserves_sglang_page_layout(): +def _assert_glm52_sparse_key_cache_unpack_preserves_sglang_page_layout(): page_size = 64 block_size = 128 cache = torch.zeros((2, page_size * (block_size + 4)), dtype=torch.uint8) @@ -662,7 +629,7 @@ def test_glm52_sparse_key_cache_unpack_preserves_sglang_page_layout(): @pytest.mark.gpu @torch.no_grad() -def test_glm52_sparse_native_sampler_codecs_are_bitwise_at_production_shapes(): +def test_glm52_sparse_native_sampler_codec_parity(): pytest.importorskip( "sglang", reason=( @@ -696,8 +663,7 @@ def test_glm52_sparse_native_sampler_codecs_are_bitwise_at_production_shapes(): assert torch.equal(key_scale.view(torch.uint8), repeated_key_scale.view(torch.uint8)) -@pytest.mark.cpu -def test_glm52_sparse_native_dispatch_flattens_batches_and_masks_unwritten_cells(): +def _assert_glm52_sparse_native_dispatch_flattens_batches_and_masks_unwritten_cells(): query = torch.zeros((2, 2, 2, 128), dtype=torch.bfloat16) key = torch.zeros((2, 3, 128), dtype=torch.bfloat16) weights = torch.ones((2, 2, 2), dtype=torch.float32) @@ -732,8 +698,7 @@ def recording_kernel(q, kv, native_weights, starts, ends, *, clean_logits): ] -@pytest.mark.cpu -def test_glm52_sparse_native_selector_fails_closed_without_cuda_and_on_nonprefix_mask(): +def _assert_glm52_sparse_native_selector_fails_closed_without_cuda_and_on_nonprefix_mask(): query = torch.zeros((1, 1, 2, 128), dtype=torch.bfloat16) key = torch.zeros((1, 2, 128), dtype=torch.bfloat16) weights = torch.ones((1, 1, 2), dtype=torch.float32) @@ -753,15 +718,13 @@ def test_glm52_sparse_native_selector_fails_closed_without_cuda_and_on_nonprefix ) -@pytest.mark.cpu -def test_glm52_sparse_deepgemm_loader_requires_score_capability(monkeypatch): +def _assert_glm52_sparse_deepgemm_loader_requires_score_capability(monkeypatch): monkeypatch.setattr(sparse_selector_module.importlib, "import_module", lambda _name: SimpleNamespace()) with pytest.raises(RuntimeError, match="deep_gemm.fp8_mqa_logits"): sparse_selector_module._load_sparse_score_kernel() -@pytest.mark.cpu -def test_glm52_sparse_selector_loader_imports_the_shared_kernel(monkeypatch): +def _assert_glm52_sparse_selector_loader_imports_the_shared_kernel(monkeypatch): """Imports are the compatibility mechanism; there is no version handshake. A genuine API break fails naturally at import; residual numerical drift @@ -782,8 +745,7 @@ def selector(scores, lengths, topk): assert sparse_selector_module._load_sparse_selection() is selector -@pytest.mark.cpu -def test_correction_bias_stays_fp32_and_checkpoint_ingestion_fails_closed(): +def _assert_correction_bias_stays_fp32_and_checkpoint_ingestion_fails_closed(): config = SimpleNamespace(n_routed_experts=4, hidden_size=3, _router_fp32=False) router = Glm5TopkRouter(config) official_values = torch.tensor([34.12345, -0.00314159, 0.33333334, 17.00013], dtype=torch.float32) @@ -819,8 +781,7 @@ def test_correction_bias_stays_fp32_and_checkpoint_ingestion_fails_closed(): handler.on_load_weight(key, torch.tensor([0.0, 1.0, float("inf"), 3.0])) -@pytest.mark.cpu -def test_canonical_moe_rejects_routing_replay_configuration(): +def _assert_canonical_moe_rejects_routing_replay_configuration(): config = _small_glm_config() config._glm52_exact_contract = True block = Glm5MoEBlock(config, layer_idx=1) @@ -829,8 +790,7 @@ def test_canonical_moe_rejects_routing_replay_configuration(): block.route(torch.zeros((1, 1, config.hidden_size))) -@pytest.mark.cpu -def test_canonical_moe_transport_resolves_internally_with_no_public_knob(): +def _assert_canonical_moe_transport_resolves_internally_with_no_public_knob(): """There is no user-facing transport menu: the model resolves the best certified transport for the geometry internally.""" config = _small_glm_config() @@ -839,8 +799,7 @@ def test_canonical_moe_transport_resolves_internally_with_no_public_knob(): assert "canonical_moe_transport" not in config.to_dict() -@pytest.mark.cpu -def test_canonical_glm_router_and_indexer_are_exact_without_environment(monkeypatch): +def _assert_canonical_glm_router_and_indexer_are_exact_without_environment(monkeypatch): calls = [] def router_gemm(hidden, weight): @@ -868,8 +827,7 @@ def indexer_gemm(hidden, weight): assert calls == ["router", "indexer"] -@pytest.mark.cpu -def test_noncanonical_glm_retains_ordinary_router(monkeypatch): +def _assert_noncanonical_glm_retains_ordinary_router(monkeypatch): def forbidden(*_args, **_kwargs): raise AssertionError("noncanonical GLM unexpectedly used the exact router kernel") @@ -957,7 +915,7 @@ def canonical_forward( ) canonical = level[0] else: - canonical = canonical_moe_reduce_reference(flattened, metadata) + canonical = _canonical_moe_reference(flattened, metadata) if _layer_id == skip_layer: canonical = flattened[0] canonical = canonical.reshape_as(hidden_states) @@ -978,8 +936,9 @@ def _semantic_logprobs(model, input_ids): @pytest.mark.cpu -@pytest.mark.parametrize("num_moe_layers", [1, 4]) -def test_semantic_moe_stack_boundary_logprob_engagement_permutation_and_composition(num_moe_layers, monkeypatch): +def test_semantic_moe_stack_boundary_logprob_engagement_permutation_and_composition(monkeypatch): + num_moe_layers = 4 + def rowwise_router(hidden, weight): weight_fp32 = weight.float() return torch.stack( @@ -1041,9 +1000,52 @@ def semantic_serving_topk( solo = torch.cat([_semantic_logprobs(trainer, row.unsqueeze(0)) for row in batch], dim=0) assert torch.equal(solo.view(torch.uint8), trainer_logprobs.view(torch.uint8)) - if num_moe_layers == 4: - faulty = Glm5ForCausalLM(_semantic_model_config(num_moe_layers)).to(torch.bfloat16).eval() - faulty.load_state_dict(trainer.state_dict(), strict=True) - _bind_semantic_canonicalizers(faulty, [], serving=False, skip_layer=0) - faulty_logprobs = _semantic_logprobs(faulty, batch) - assert not torch.equal(faulty_logprobs.view(torch.uint8), trainer_logprobs.view(torch.uint8)) + faulty = Glm5ForCausalLM(_semantic_model_config(num_moe_layers)).to(torch.bfloat16).eval() + faulty.load_state_dict(trainer.state_dict(), strict=True) + _bind_semantic_canonicalizers(faulty, [], serving=False, skip_layer=0) + faulty_logprobs = _semantic_logprobs(faulty, batch) + assert not torch.equal(faulty_logprobs.view(torch.uint8), trainer_logprobs.view(torch.uint8)) + + +@pytest.mark.cpu +def test_glm52_layer_plan_and_indexer_allocation_contract(): + _assert_canonical_trainer_admits_only_certified_world16_ep16_cp16() + _assert_official_layer_plan_counts_producers_and_38_40_split() + _assert_layer_plan_rejects_malformed_schedules_and_shared_stage_start() + _assert_only_full_layers_allocate_indexer_parameters_and_strict_state_dict_round_trip() + + +@pytest.mark.cpu +def test_glm52_sparse_selector_pipeline_contract(monkeypatch): + _assert_sparse_selector_ties_short_rows_and_dead_rows() + _assert_glm52_sparse_shared_selector_handles_production_boundary_ties_and_dead_tail() + _assert_glm52_sparse_hadamard_transport_is_normalized_and_self_inverse() + with monkeypatch.context() as case_patch: + _assert_sparse_selector_applies_hadamard_before_fp8_quantization(case_patch) + with monkeypatch.context() as case_patch: + _assert_sparse_selector_fused_contract_skips_hadamard_before_fp8_quantization(case_patch) + with monkeypatch.context() as case_patch: + _assert_fused_bf16_indexer_projection_matches_sampler_row_order(case_patch) + _assert_fused_bf16_indexer_head_gate_scaling_promotes_before_head_scale() + _assert_fused_sampler_index_k_prepare_preserves_projection_stride_and_builds_literal_rope_cache() + _assert_sampler_index_k_preparation_uses_split_prompt_and_fused_decode_suffix() + _assert_sampler_index_k_preparation_maps_4096_boundary_across_cp16() + _assert_glm52_sparse_query_and_key_codecs_have_distinct_scale_contracts() + _assert_glm52_sparse_key_cache_unpack_preserves_sglang_page_layout() + _assert_glm52_sparse_native_dispatch_flattens_batches_and_masks_unwritten_cells() + _assert_glm52_sparse_native_selector_fails_closed_without_cuda_and_on_nonprefix_mask() + with monkeypatch.context() as case_patch: + _assert_glm52_sparse_deepgemm_loader_requires_score_capability(case_patch) + with monkeypatch.context() as case_patch: + _assert_glm52_sparse_selector_loader_imports_the_shared_kernel(case_patch) + with monkeypatch.context() as case_patch: + _assert_glm52_canonical_moe_configuration_and_selection_contract(case_patch) + + +def _assert_glm52_canonical_moe_configuration_and_selection_contract(monkeypatch): + _assert_correction_bias_stays_fp32_and_checkpoint_ingestion_fails_closed() + _assert_canonical_moe_rejects_routing_replay_configuration() + _assert_canonical_moe_transport_resolves_internally_with_no_public_knob() + _assert_canonical_glm_router_and_indexer_are_exact_without_environment(monkeypatch) + monkeypatch.undo() + _assert_noncanonical_glm_retains_ordinary_router(monkeypatch) diff --git a/tests/models/test_glm52_exact_absorbed_kv_b_qlora.py b/tests/models/test_glm52_exact_absorbed_kv_b_qlora.py index 8464c975..a6292a41 100644 --- a/tests/models/test_glm52_exact_absorbed_kv_b_qlora.py +++ b/tests/models/test_glm52_exact_absorbed_kv_b_qlora.py @@ -57,7 +57,7 @@ def _pattern( ) -def test_absorbed_kv_b_contract_keeps_one_frozen_native_base_and_two_logical_masters() -> None: +def test_absorbed_kv_b_cpu_state_and_admission_policy() -> None: module = _module() assert isinstance(module, NativeBlockFP8Linear) @@ -101,8 +101,11 @@ def test_absorbed_kv_b_contract_keeps_one_frozen_native_base_and_two_logical_mas with pytest.raises(ValueError, match="only rank=1 and alpha=1"): module.set_runtime_lora_config(1, 2) + _assert_absorbed_kv_b_dtype_move_preserves_native_bytes_master_values_and_identity() + _assert_absorbed_kv_b_rejects_direct_projection_and_castable_factor_state() -def test_absorbed_kv_b_dtype_move_preserves_native_bytes_master_values_and_identity() -> None: + +def _assert_absorbed_kv_b_dtype_move_preserves_native_bytes_master_values_and_identity() -> None: module = _module() with torch.no_grad(): module.packed_weight_f32.copy_( @@ -123,7 +126,7 @@ def test_absorbed_kv_b_dtype_move_preserves_native_bytes_master_values_and_ident assert torch.equal(parameter, expected[name]) -def test_absorbed_kv_b_rejects_direct_projection_and_castable_factor_state() -> None: +def _assert_absorbed_kv_b_rejects_direct_projection_and_castable_factor_state() -> None: module = _module() with pytest.raises(RuntimeError, match="cannot run as a direct projection"): diff --git a/tests/models/test_glm52_exact_active_lora_gate.py b/tests/models/test_glm52_exact_active_lora_gate.py index 2e0349c7..01aa54e5 100644 --- a/tests/models/test_glm52_exact_active_lora_gate.py +++ b/tests/models/test_glm52_exact_active_lora_gate.py @@ -28,7 +28,7 @@ def _config_with_all_active_lora_flags(): return config -def test_complete_active_lora_flags_enable_exact_forward_without_scoring_marker() -> None: +def test_active_lora_composite_admission_policy() -> None: config = _config_with_all_active_lora_flags() assert glm52_exact_active_lora_enabled(config) @@ -36,8 +36,20 @@ def test_complete_active_lora_flags_enable_exact_forward_without_scoring_marker( assert _is_exact_glm52(config) assert _moe_bi_router_enabled(config) + for missing_flag in GLM52_EXACT_ACTIVE_LORA_FLAGS: + partial = _config_with_all_active_lora_flags() + setattr(partial, missing_flag, False) + + assert not glm52_exact_active_lora_enabled(partial) + assert not glm52_exact_forward_enabled(partial) + assert not _is_exact_glm52(partial) + assert not _moe_bi_router_enabled(partial) + + scoring_only = _official_config() + scoring_only._glm52_exact_contract = True + assert not glm52_exact_active_lora_enabled(scoring_only) + assert glm52_exact_forward_enabled(scoring_only) -def test_active_lora_component_flags_are_set_and_cleared_atomically() -> None: config = _official_config() config._glm52_exact_active_lora_dense_component = True @@ -90,7 +102,7 @@ def load_model(self, *, init_kwargs, **_kwargs): ) -def test_rank1_server_training_derives_the_complete_family_without_private_flags( +def test_rank1_server_training_derives_the_complete_family_and_checks_topology( monkeypatch: pytest.MonkeyPatch, ) -> None: model = _build_exact_rank1(monkeypatch) @@ -98,35 +110,9 @@ def test_rank1_server_training_derives_the_complete_family_without_private_flags assert glm52_exact_active_lora_enabled(model.config) assert model.config._glm52_exact_contract is False assert all(getattr(model.config, flag) is True for flag in GLM52_EXACT_ACTIVE_LORA_FLAGS) - - -def test_rank1_server_training_rejects_non_tp16_lm_head_before_loading( - monkeypatch: pytest.MonkeyPatch, -) -> None: with pytest.raises(ValueError, match="lm-head-TP16"): _build_exact_rank1(monkeypatch, lm_head_tp_size=1) - -@pytest.mark.parametrize("missing_flag", GLM52_EXACT_ACTIVE_LORA_FLAGS) -def test_every_active_lora_component_is_required(missing_flag: str) -> None: - config = _config_with_all_active_lora_flags() - setattr(config, missing_flag, False) - - assert not glm52_exact_active_lora_enabled(config) - assert not glm52_exact_forward_enabled(config) - assert not _is_exact_glm52(config) - assert not _moe_bi_router_enabled(config) - - -def test_scoring_only_marker_remains_an_independent_exact_forward_admission() -> None: - config = _official_config() - config._glm52_exact_contract = True - - assert not glm52_exact_active_lora_enabled(config) - assert glm52_exact_forward_enabled(config) - - -def test_cached_indexer_and_moe_surfaces_activate_only_for_complete_composite() -> None: complete = _config_with_all_active_lora_flags() partial = _config_with_all_active_lora_flags() partial._glm52_exact_active_lora_lm_head_component = False diff --git a/tests/models/test_glm52_exact_attention_checkpoint.py b/tests/models/test_glm52_exact_attention_checkpoint.py index 3d667d35..58749027 100644 --- a/tests/models/test_glm52_exact_attention_checkpoint.py +++ b/tests/models/test_glm52_exact_attention_checkpoint.py @@ -134,115 +134,112 @@ def test_exact_attention_checkpoint_inventory_routes_only_absorbed_kv_b_through_ assert len(trainable) == 10 assert all(parameter.dtype is torch.float32 for parameter in trainable.values()) + _assert_exact_absorbed_kv_b_checkpoint_pair_is_order_independent_and_byte_exact(checkpoint_case) -@pytest.mark.parametrize("arrival_order", (("weight", "scale"), ("scale", "weight"))) -def test_exact_absorbed_kv_b_checkpoint_pair_is_order_independent_and_byte_exact( + +def _assert_exact_absorbed_kv_b_checkpoint_pair_is_order_independent_and_byte_exact( checkpoint_case: _CheckpointCase, - arrival_order: tuple[str, str], ) -> None: case = checkpoint_case - handler = _handler(case) values = { "weight": (f"{case.source}.weight", case.weight), "scale": (f"{case.source}.weight_scale_inv", case.scale), } - - first_key, first_tensor = values[arrival_order[0]] - second_key, second_tensor = values[arrival_order[1]] - assert handler.on_load_weight(first_key, first_tensor) == [] - emitted = handler.on_load_weight(second_key, second_tensor) - assert handler.on_load_complete() == [] - - assert [name for name, _ in emitted] == [ - f"{case.target}.packed_weight_f32", - f"{case.target}.weight_scale_inv", - ] - packed = emitted[0][1] - loaded_scale = emitted[1][1] - assert packed.dtype is torch.float32 - assert torch.equal(packed.view(torch.uint8), case.weight.view(torch.uint8)) - assert loaded_scale.dtype is torch.float32 - assert torch.equal(loaded_scale, case.scale) - - -@pytest.mark.parametrize("member", ("weight", "scale")) -def test_exact_absorbed_kv_b_checkpoint_pair_rejects_duplicate_members( + for arrival_order in (("weight", "scale"), ("scale", "weight")): + handler = _handler(case) + first_key, first_tensor = values[arrival_order[0]] + second_key, second_tensor = values[arrival_order[1]] + assert handler.on_load_weight(first_key, first_tensor) == [] + emitted = handler.on_load_weight(second_key, second_tensor) + assert handler.on_load_complete() == [] + + assert [name for name, _ in emitted] == [ + f"{case.target}.packed_weight_f32", + f"{case.target}.weight_scale_inv", + ] + packed = emitted[0][1] + loaded_scale = emitted[1][1] + assert packed.dtype is torch.float32 + assert torch.equal(packed.view(torch.uint8), case.weight.view(torch.uint8)) + assert loaded_scale.dtype is torch.float32 + assert torch.equal(loaded_scale, case.scale) + + _assert_exact_absorbed_kv_b_checkpoint_pair_rejects_duplicate_members(checkpoint_case) + _assert_exact_absorbed_kv_b_checkpoint_pair_rejects_missing_members_at_completion(checkpoint_case) + _assert_exact_absorbed_kv_b_checkpoint_pair_rejects_dtype_mismatches(checkpoint_case) + _assert_exact_absorbed_kv_b_checkpoint_pair_rejects_shape_mismatches(checkpoint_case) + + +def _assert_exact_absorbed_kv_b_checkpoint_pair_rejects_duplicate_members( checkpoint_case: _CheckpointCase, - member: str, ) -> None: case = checkpoint_case - handler = _handler(case) - key, tensor = ( - (f"{case.source}.weight", case.weight) - if member == "weight" - else (f"{case.source}.weight_scale_inv", case.scale) - ) + for member in ("weight", "scale"): + handler = _handler(case) + key, tensor = ( + (f"{case.source}.weight", case.weight) + if member == "weight" + else (f"{case.source}.weight_scale_inv", case.scale) + ) - assert handler.on_load_weight(key, tensor) == [] - with pytest.raises(ValueError, match="Duplicate native FP8 pair member"): - handler.on_load_weight(key, tensor) + assert handler.on_load_weight(key, tensor) == [] + with pytest.raises(ValueError, match="Duplicate native FP8 pair member"): + handler.on_load_weight(key, tensor) -@pytest.mark.parametrize("member", ("weight", "scale")) -def test_exact_absorbed_kv_b_checkpoint_pair_rejects_missing_members_at_completion( +def _assert_exact_absorbed_kv_b_checkpoint_pair_rejects_missing_members_at_completion( checkpoint_case: _CheckpointCase, - member: str, ) -> None: case = checkpoint_case - handler = _handler(case) - key, tensor = ( - (f"{case.source}.weight", case.weight) - if member == "weight" - else (f"{case.source}.weight_scale_inv", case.scale) - ) + for member in ("weight", "scale"): + handler = _handler(case) + key, tensor = ( + (f"{case.source}.weight", case.weight) + if member == "weight" + else (f"{case.source}.weight_scale_inv", case.scale) + ) - assert handler.on_load_weight(key, tensor) == [] - with pytest.raises(ValueError, match="Incomplete native FP8 pairs"): - handler.on_load_complete() + assert handler.on_load_weight(key, tensor) == [] + with pytest.raises(ValueError, match="Incomplete native FP8 pairs"): + handler.on_load_complete() -@pytest.mark.parametrize( - ("bad_member", "bad_tensor", "message"), - ( - ("weight", torch.zeros(1, dtype=torch.bfloat16), "weight must be float8_e4m3fn"), - ("scale", torch.zeros(1, dtype=torch.bfloat16), "weight_scale_inv must be FP32"), - ), -) -def test_exact_absorbed_kv_b_checkpoint_pair_rejects_dtype_mismatches( +def _assert_exact_absorbed_kv_b_checkpoint_pair_rejects_dtype_mismatches( checkpoint_case: _CheckpointCase, - bad_member: str, - bad_tensor: torch.Tensor, - message: str, ) -> None: - case = checkpoint_case - handler = _handler(case) - good_member = "scale" if bad_member == "weight" else "weight" - good_key, good_tensor = ( - (f"{case.source}.weight_scale_inv", case.scale) - if good_member == "scale" - else (f"{case.source}.weight", case.weight) + cases = ( + ("weight", torch.zeros(1, dtype=torch.bfloat16), "weight must be float8_e4m3fn"), + ("scale", torch.zeros(1, dtype=torch.bfloat16), "weight_scale_inv must be FP32"), ) - bad_key = f"{case.source}.{'weight' if bad_member == 'weight' else 'weight_scale_inv'}" - - assert handler.on_load_weight(good_key, good_tensor) == [] - with pytest.raises((TypeError, ValueError), match=message): - handler.on_load_weight(bad_key, bad_tensor) - - -@pytest.mark.parametrize("bad_member", ("weight", "scale")) -def test_exact_absorbed_kv_b_checkpoint_pair_rejects_shape_mismatches( + case = checkpoint_case + for bad_member, bad_tensor, message in cases: + handler = _handler(case) + good_member = "scale" if bad_member == "weight" else "weight" + good_key, good_tensor = ( + (f"{case.source}.weight_scale_inv", case.scale) + if good_member == "scale" + else (f"{case.source}.weight", case.weight) + ) + bad_key = f"{case.source}.{'weight' if bad_member == 'weight' else 'weight_scale_inv'}" + + assert handler.on_load_weight(good_key, good_tensor) == [] + with pytest.raises((TypeError, ValueError), match=message): + handler.on_load_weight(bad_key, bad_tensor) + + +def _assert_exact_absorbed_kv_b_checkpoint_pair_rejects_shape_mismatches( checkpoint_case: _CheckpointCase, - bad_member: str, ) -> None: case = checkpoint_case - handler = _handler(case) - if bad_member == "weight": - good_key, good_tensor = f"{case.source}.weight_scale_inv", case.scale - bad_key, bad_tensor = f"{case.source}.weight", case.weight[:-1] - else: - good_key, good_tensor = f"{case.source}.weight", case.weight - bad_key, bad_tensor = f"{case.source}.weight_scale_inv", case.scale[:-1] - - assert handler.on_load_weight(good_key, good_tensor) == [] - with pytest.raises(ValueError, match="unexpected shape|must be FP32"): - handler.on_load_weight(bad_key, bad_tensor) + for bad_member in ("weight", "scale"): + handler = _handler(case) + if bad_member == "weight": + good_key, good_tensor = f"{case.source}.weight_scale_inv", case.scale + bad_key, bad_tensor = f"{case.source}.weight", case.weight[:-1] + else: + good_key, good_tensor = f"{case.source}.weight", case.weight + bad_key, bad_tensor = f"{case.source}.weight_scale_inv", case.scale[:-1] + + assert handler.on_load_weight(good_key, good_tensor) == [] + with pytest.raises(ValueError, match="unexpected shape|must be FP32"): + handler.on_load_weight(bad_key, bad_tensor) diff --git a/tests/models/test_glm52_exact_attention_construction.py b/tests/models/test_glm52_exact_attention_construction.py deleted file mode 100644 index 2355a240..00000000 --- a/tests/models/test_glm52_exact_attention_construction.py +++ /dev/null @@ -1,113 +0,0 @@ -from __future__ import annotations - -import pytest -import torch - -from tests.models.test_glm52_qlora import _meta_model, _official_config -from xorl.models.transformers.glm5.exact_absorbed_kv_b_qlora import ( - Glm52ExactTP1AbsorbedKvBBlockFP8QLoRA, -) -from xorl.models.transformers.glm5.exact_dense_mlp import Glm52ExactTP1DenseMLP -from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear -from xorl.models.transformers.glm5.qlora import GLM52_QLORA_FACTOR_COUNT, prepare_glm52_block_fp8_qlora - - -_ORDINARY_ATTENTION_PROJECTIONS = ( - "q_a_proj", - "kv_a_proj_with_mqa", - "q_b_proj", - "o_proj", -) -_ALL_ATTENTION_PROJECTIONS = (*_ORDINARY_ATTENTION_PROJECTIONS, "kv_b_proj") - - -def _exact_attention_config(): - config = _official_config() - config._glm52_exact_active_lora_dense_component = True - config._glm52_exact_active_lora_attention_component = True - config._sparse_mla_enabled = True - config._ep_dispatch = "alltoall" - return config - - -def _expected_attention_factor_names() -> set[str]: - return { - f"model.layers.{layer_idx}.self_attn.{projection}.lora_{factor}" - for layer_idx in range(78) - for projection in _ALL_ATTENTION_PROJECTIONS - for factor in ("A", "B") - } - - -def test_glm52_exact_attention_component_preserves_complete_canonical_inventory() -> None: - config = _exact_attention_config() - model = _meta_model(config) - - inventory = prepare_glm52_block_fp8_qlora(model, config, adapter_rank=1, adapter_alpha=1) - - expected_attention_factors = _expected_attention_factor_names() - actual_attention_factors = {factor.name for factor in inventory.factors if factor.role.startswith("attention.")} - assert actual_attention_factors == expected_attention_factors - assert len(actual_attention_factors) == 78 * 5 * 2 == 780 - assert len(inventory.factors) == GLM52_QLORA_FACTOR_COUNT == 1700 - assert len(inventory.factor_names) == GLM52_QLORA_FACTOR_COUNT - - trainable = {name: parameter for name, parameter in model.named_parameters() if parameter.requires_grad} - assert set(trainable) == inventory.factor_names - assert all(parameter.dtype is torch.float32 for parameter in trainable.values()) - assert len({id(parameter) for parameter in trainable.values()}) == GLM52_QLORA_FACTOR_COUNT - - for layer_idx, layer in enumerate(model.model.layers): - attention = layer.self_attn - prefix = f"model.layers.{layer_idx}.self_attn" - for projection in _ORDINARY_ATTENTION_PROJECTIONS: - module = getattr(attention, projection) - assert type(module) is Glm52ExactTP1BlockFP8QLoRALinear - assert module._source_fqn == f"{prefix}.{projection}" - assert type(attention.kv_b_proj) is Glm52ExactTP1AbsorbedKvBBlockFP8QLoRA - assert attention.kv_b_proj._source_fqn == f"{prefix}.kv_b_proj" - - layer_attention_factors = {name for name in trainable if name.startswith(f"{prefix}.")} - assert layer_attention_factors == { - f"{prefix}.{projection}.lora_{factor}" for projection in _ALL_ATTENTION_PROJECTIONS for factor in ("A", "B") - } - - exact_dense_roots = {name for name, module in model.named_modules() if isinstance(module, Glm52ExactTP1DenseMLP)} - assert exact_dense_roots == {f"model.layers.{layer_idx}.mlp" for layer_idx in range(3)} - - -@pytest.mark.parametrize(("rank", "alpha"), ((16, 16), (1, 2), (2, 1))) -def test_glm52_exact_attention_component_rejects_non_rank1_alpha1_before_mutation( - rank: int, - alpha: int, -) -> None: - config = _exact_attention_config() - model = _meta_model(config) - - with pytest.raises(ValueError, match="requires adapter_rank=1 and adapter_alpha=1"): - prepare_glm52_block_fp8_qlora(model, config, adapter_rank=rank, adapter_alpha=alpha) - - assert not any("lora_" in name for name, _ in model.named_parameters()) - - -@pytest.mark.parametrize( - ("override", "message"), - ( - ({"_glm52_exact_active_lora_dense_component": False}, "exact active-LoRA dense component"), - ({"_ep_dispatch": "deepep"}, "ep_dispatch='alltoall'"), - ({"_sparse_mla_enabled": False}, "requires sparse_mla_enabled=true"), - ), -) -def test_glm52_exact_attention_component_rejects_incomplete_execution_contract_before_mutation( - override: dict[str, object], - message: str, -) -> None: - config = _exact_attention_config() - for name, value in override.items(): - setattr(config, name, value) - model = _meta_model(config) - - with pytest.raises(ValueError, match=message): - prepare_glm52_block_fp8_qlora(model, config, adapter_rank=1, adapter_alpha=1) - - assert not any("lora_" in name for name, _ in model.named_parameters()) diff --git a/tests/models/test_glm52_exact_attention_modeling.py b/tests/models/test_glm52_exact_attention_modeling.py index f52af2a4..7a5985a1 100644 --- a/tests/models/test_glm52_exact_attention_modeling.py +++ b/tests/models/test_glm52_exact_attention_modeling.py @@ -99,11 +99,6 @@ def select_topk(self, index_q: torch.Tensor, *_args, **_kwargs) -> torch.Tensor: return torch.zeros((*index_q.shape[:2], self.topk), dtype=torch.long) -class _ForbiddenGenericKvB(nn.Module): - def forward(self, *_args, **_kwargs) -> torch.Tensor: - raise AssertionError("generic absorbed MLA must consume the split weights, not call kv_b_proj") - - def _attention_with_simple_projections() -> Glm5Attention: attention = Glm5Attention(_tiny_config(), layer_idx=0) attention.q_a_layernorm = nn.Identity() @@ -113,8 +108,7 @@ def _attention_with_simple_projections() -> Glm5Attention: return attention -@pytest.mark.parametrize("cp_enabled", [False, True], ids=["non_cp", "ulysses_cp"]) -def test_sparse_exact_kv_b_routes_q_and_both_v_sites_without_weight_materialization( +def _assert_sparse_exact_kv_b_routes_q_and_both_v_sites_without_weight_materialization( monkeypatch: pytest.MonkeyPatch, cp_enabled: bool, ) -> None: @@ -218,65 +212,12 @@ def fake_sparse_mla( assert "query_offset" not in sparse_calls[0][3] -def test_generic_absorbed_q_and_v_keep_the_legacy_split_weight_einsums( +def test_sparse_exact_kv_b_routes_q_and_both_v_sites_without_weight_materialization( monkeypatch: pytest.MonkeyPatch, ) -> None: - torch.manual_seed(1) - attention = _attention_with_simple_projections() - attention.kv_b_proj = _ForbiddenGenericKvB() - monkeypatch.setattr( - modeling_glm5, - "glm5_apply_rotary_pos_emb", - lambda q, k, *_args, **_kwargs: (q, k), - ) - - w_kc = torch.randn( - attention.num_heads, - attention.qk_nope_head_dim, - attention.kv_lora_rank, - ) - w_vc = torch.randn( - attention.num_heads, - attention.v_head_dim, - attention.kv_lora_rank, - ) - split_kv_b_weight = MagicMock(return_value=(w_kc, w_vc)) - monkeypatch.setattr(attention, "_split_kv_b_weight", split_kv_b_weight) - - batch_size, seq_len = 2, 3 - hidden_states = torch.randn(batch_size, seq_len, attention.config.hidden_size) - position_embeddings = ( - torch.zeros(batch_size, seq_len, attention.qk_rope_head_dim), - torch.zeros(batch_size, seq_len, attention.qk_rope_head_dim), - ) - q, _kv, _q_compressed, actual_w_vc = attention._project_qkv_absorb( - hidden_states, - position_embeddings, - ) - - q_compressed = attention.q_a_proj(hidden_states) - q_unabsorbed = attention.q_b_proj(q_compressed).view( - batch_size, - seq_len, - attention.num_heads, - attention.qk_head_dim, - ) - q_no_pe, q_pe = torch.split( - q_unabsorbed, - [attention.qk_nope_head_dim, attention.qk_rope_head_dim], - dim=-1, - ) - expected_q = torch.cat((torch.einsum("bshd,hdc->bshc", q_no_pe, w_kc), q_pe), dim=-1) - torch.testing.assert_close(q, expected_q) - assert actual_w_vc is w_vc - split_kv_b_weight.assert_called_once_with(compute_dtype=q_no_pe.dtype) - - attn_latent = torch.randn( - batch_size, - seq_len, - attention.num_heads, - attention.kv_lora_rank, - ) - value = attention._project_absorbed_value(attn_latent, actual_w_vc) - expected_value = torch.einsum("bshk,hdk->bshd", attn_latent, w_vc) - torch.testing.assert_close(value, expected_value) + for cp_enabled in (False, True): + with monkeypatch.context() as case_patch: + _assert_sparse_exact_kv_b_routes_q_and_both_v_sites_without_weight_materialization( + case_patch, + cp_enabled, + ) diff --git a/tests/models/test_glm52_exact_dcp.py b/tests/models/test_glm52_exact_dcp.py index d09b238e..5bb07b8b 100644 --- a/tests/models/test_glm52_exact_dcp.py +++ b/tests/models/test_glm52_exact_dcp.py @@ -95,9 +95,16 @@ def test_exact_base_dcp_contract_exhausts_three_dense_and_315_scale_aliases(tmp_ assert result["unexpected_in_checkpoint"] == [] assert result["missing_buffers_in_checkpoint"] == [] assert result["unexpected_buffers_in_checkpoint"] == [] + with monkeypatch.context() as case_patch: + _assert_distributed_checkpointer_loads_official_base_dcp_keys_into_exact_runtime_state( + tmp_path / "official-load", + case_patch, + ) -def test_distributed_checkpointer_loads_official_base_dcp_keys_into_exact_runtime_state(tmp_path, monkeypatch) -> None: +def _assert_distributed_checkpointer_loads_official_base_dcp_keys_into_exact_runtime_state( + tmp_path, monkeypatch +) -> None: source_model = _OneDenseExactModel() source_projection = Glm52ExactBaseDcpLoadProjection(source_model) source_state = { @@ -177,7 +184,7 @@ def test_exact_base_dcp_dense_staging_fuses_into_four_rank_shard() -> None: ) -def test_skip_mode_defers_base_loading_to_dcp_but_keeps_fsdp_deregistration(monkeypatch) -> None: +def test_skip_mode_admission_and_fsdp_deregistration_policy(monkeypatch) -> None: model = nn.Module() model.config = SimpleNamespace(**dict.fromkeys(GLM52_EXACT_ACTIVE_LORA_FLAGS, True)) calls = [] @@ -197,10 +204,8 @@ def test_skip_mode_defers_base_loading_to_dcp_but_keeps_fsdp_deregistration(monk assert calls == [(model, ("packed_weight_f32",))] - -def test_skip_mode_rejects_non_exact_qlora_model() -> None: - model = nn.Module() - model.config = SimpleNamespace() + non_exact_model = nn.Module() + non_exact_model.config = SimpleNamespace() with pytest.raises(ValueError, match="complete GLM-5.2 exact active-LoRA model"): - model_builder._deferred_qlora_quantize(model, "/dcp-only", load_weights_mode="skip") + model_builder._deferred_qlora_quantize(non_exact_model, "/dcp-only", load_weights_mode="skip") diff --git a/tests/models/test_glm52_exact_dense_checkpoint.py b/tests/models/test_glm52_exact_dense_checkpoint.py index c5991e04..90c36b9e 100644 --- a/tests/models/test_glm52_exact_dense_checkpoint.py +++ b/tests/models/test_glm52_exact_dense_checkpoint.py @@ -83,8 +83,10 @@ def test_exact_dense_checkpoint_handler_emits_one_explicit_gate_then_up_native_p assert mlp._exact_gate_up_base_loaded is True handler.on_load_complete() + _assert_exact_dense_gate_up_pair_buffer_fails_on_missing_duplicate_or_invalid_members() -def test_exact_dense_gate_up_pair_buffer_fails_on_missing_duplicate_or_invalid_members() -> None: + +def _assert_exact_dense_gate_up_pair_buffer_fails_on_missing_duplicate_or_invalid_members() -> None: model, _ = _model() pairs = _pairs() diff --git a/tests/models/test_glm52_exact_dense_mlp.py b/tests/models/test_glm52_exact_dense_mlp.py index 89e0232a..7f82618e 100644 --- a/tests/models/test_glm52_exact_dense_mlp.py +++ b/tests/models/test_glm52_exact_dense_mlp.py @@ -63,7 +63,7 @@ def _literal_gate_up_value( return (base + torch.cat((gate_delta, up_delta), dim=-1)).to(torch.bfloat16) -def test_dense_mlp_root_preserves_six_canonical_unique_factor_paths_without_aliases() -> None: +def _assert_dense_mlp_root_preserves_six_canonical_unique_factor_paths_without_aliases() -> None: module = _module() module.bind_checkpoint_sources("model.layers.0.mlp") @@ -110,7 +110,9 @@ def test_dense_mlp_root_preserves_six_canonical_unique_factor_paths_without_alia module.bind_checkpoint_sources("model.layers.1.mlp") -def test_dense_mlp_forward_composes_fused_gate_up_production_activation_and_exact_down(monkeypatch) -> None: +def test_dense_mlp_forward_composes_fused_gate_up_production_activation_and_exact_down(monkeypatch, tmp_path) -> None: + _assert_dense_mlp_root_preserves_six_canonical_unique_factor_paths_without_aliases() + _assert_dense_mlp_runtime_rank_alpha_contract_is_atomic_and_fails_before_forward() module = _module() fused_base = torch.arange(256 * 8, dtype=torch.float32).reshape(256, 8).sub_(719).div_(1543).to(torch.bfloat16) down_base = torch.arange(8 * 128, dtype=torch.float32).reshape(8, 128).sub_(401).div_(1291).to(torch.bfloat16) @@ -162,9 +164,10 @@ def down_value(input, factor_A, factor_B): ("down", (3, 128)), ] assert torch.equal(actual, expected) + _assert_dense_mlp_roundtrips_six_canonical_factors_through_xorl_and_peft_export(tmp_path) -def test_dense_mlp_runtime_rank_alpha_contract_is_atomic_and_fails_before_forward() -> None: +def _assert_dense_mlp_runtime_rank_alpha_contract_is_atomic_and_fails_before_forward() -> None: with pytest.raises(ValueError, match="rank=1 and alpha=1"): Glm52ExactTP1DenseMLP(8, 128, r=2, lora_alpha=1) with pytest.raises(ValueError, match="rank=1 and alpha=1"): @@ -196,7 +199,7 @@ def test_dense_mlp_runtime_rank_alpha_contract_is_atomic_and_fails_before_forwar module(input) -def test_dense_mlp_roundtrips_six_canonical_factors_through_xorl_and_peft_export(tmp_path) -> None: +def _assert_dense_mlp_roundtrips_six_canonical_factors_through_xorl_and_peft_export(tmp_path) -> None: source = _module() state = get_lora_state_dict(source) assert tuple(state) == source.logical_factor_names diff --git a/tests/models/test_glm52_exact_gate_up_qlora.py b/tests/models/test_glm52_exact_gate_up_qlora.py index 2f4b3f70..6bf4c428 100644 --- a/tests/models/test_glm52_exact_gate_up_qlora.py +++ b/tests/models/test_glm52_exact_gate_up_qlora.py @@ -62,7 +62,7 @@ def _load_small_base(module: Glm52ExactTP1FusedGateUpBlockFP8QLoRA) -> tuple[tor return gate_weight, gate_scales, up_weight, up_scales -def test_fused_gate_up_contract_is_one_native_leaf_with_four_logical_fp32_factors() -> None: +def _assert_fused_gate_up_contract_is_one_native_leaf_with_four_logical_fp32_factors() -> None: module = _module() assert isinstance(module, NativeBlockFP8Linear) @@ -101,7 +101,7 @@ def test_fused_gate_up_contract_is_one_native_leaf_with_four_logical_fp32_factor module.set_runtime_lora_config(2, 1) -def test_fused_gate_up_loader_makes_gate_then_up_order_explicit_and_strict() -> None: +def _assert_fused_gate_up_loader_makes_gate_then_up_order_explicit_and_strict() -> None: module = _module() gate_weight, gate_scales, up_weight, up_scales = _load_small_base(module) @@ -120,7 +120,7 @@ def test_fused_gate_up_loader_makes_gate_then_up_order_explicit_and_strict() -> module.load_gate_up_prequantized(gate_weight, gate_scales, up_weight, up_scales.to(torch.bfloat16)) -def test_fused_gate_up_model_dtype_move_preserves_native_state_and_fp32_masters() -> None: +def _assert_fused_gate_up_model_dtype_move_preserves_native_state_and_fp32_masters() -> None: module = _module() _load_small_base(module) packed_bytes = module.packed_weight_f32.detach().view(torch.uint8).clone() @@ -145,7 +145,7 @@ def test_fused_gate_up_model_dtype_move_preserves_native_state_and_fp32_masters( assert torch.equal(actual, expected) -def test_fused_gate_up_rounds_each_master_once_and_preserves_logical_order(monkeypatch) -> None: +def _assert_fused_gate_up_rounds_each_master_once_and_preserves_logical_order(monkeypatch) -> None: module = _module() base_weight = torch.arange(256 * 8, dtype=torch.float32).reshape(256, 8).sub_(311).div_(977).to(torch.bfloat16) captures = [] @@ -167,7 +167,7 @@ def test_fused_gate_up_rounds_each_master_once_and_preserves_logical_order(monke assert torch.equal(actual, expected) -def test_fused_gate_up_surrogate_matches_two_logical_qlora_branches(monkeypatch) -> None: +def _assert_fused_gate_up_surrogate_matches_two_logical_qlora_branches(monkeypatch) -> None: module = _module() base_weight = torch.arange(256 * 8, dtype=torch.float32).reshape(256, 8).sub_(617).div_(1237).to(torch.bfloat16) monkeypatch.setattr(module, "_dequantize_base_weight", lambda: base_weight) @@ -231,7 +231,7 @@ def test_fused_gate_up_surrogate_matches_two_logical_qlora_branches(monkeypatch) assert torch.equal(master.grad, reference_factor.grad) -def test_fused_gate_up_input_gradient_matches_two_exact_logical_wrappers(monkeypatch) -> None: +def _assert_fused_gate_up_input_gradient_matches_two_exact_logical_wrappers(monkeypatch) -> None: fused = _module() base_weight = torch.arange(256 * 8, dtype=torch.float32).reshape(256, 8).sub_(503).mul_(7).to(torch.bfloat16) monkeypatch.setattr(fused, "_dequantize_base_weight", lambda: base_weight) @@ -263,7 +263,7 @@ def test_fused_gate_up_input_gradient_matches_two_exact_logical_wrappers(monkeyp assert torch.equal(fused.up_proj.lora_B.grad, up.lora_B.grad) -def test_fused_gate_up_factor_only_backward_does_not_materialize_base(monkeypatch) -> None: +def _assert_fused_gate_up_factor_only_backward_does_not_materialize_base(monkeypatch) -> None: module = _module() base_weight = torch.zeros(256, 8, dtype=torch.bfloat16) monkeypatch.setattr(module, "_exact_forward_value", _literal_cpu_value(base_weight, [])) @@ -278,7 +278,7 @@ def test_fused_gate_up_factor_only_backward_does_not_materialize_base(monkeypatc assert all(dict(module.named_parameters())[name].grad is not None for name in module.logical_factor_names) -def test_fused_gate_up_backward_rejects_any_master_mutation(monkeypatch) -> None: +def _assert_fused_gate_up_backward_rejects_any_master_mutation(monkeypatch) -> None: module = _module() base_weight = torch.zeros(256, 8, dtype=torch.bfloat16) monkeypatch.setattr(module, "_dequantize_base_weight", lambda: base_weight) @@ -292,7 +292,7 @@ def test_fused_gate_up_backward_rejects_any_master_mutation(monkeypatch) -> None output.float().sum().backward() -def test_fused_gate_up_contract_fails_closed_before_sglang_import() -> None: +def _assert_fused_gate_up_contract_fails_closed_before_sglang_import() -> None: before = {name for name in sys.modules if name == "sglang" or name.startswith("sglang.")} module = _module() @@ -553,3 +553,16 @@ def test_official_fused_gate_up_literal_bytes_graph_metadata_zero_and_gradients( dict(module.named_parameters())[name].zero_() zero_output = module(input.detach()) assert torch.equal(zero_output.view(torch.uint8), direct_base.view(torch.uint8)) + + +def test_fused_gate_up_cpu_contract(monkeypatch) -> None: + _assert_fused_gate_up_contract_is_one_native_leaf_with_four_logical_fp32_factors() + _assert_fused_gate_up_loader_makes_gate_then_up_order_explicit_and_strict() + _assert_fused_gate_up_model_dtype_move_preserves_native_state_and_fp32_masters() + _assert_fused_gate_up_contract_fails_closed_before_sglang_import() + _assert_fused_gate_up_rounds_each_master_once_and_preserves_logical_order(monkeypatch) + _assert_fused_gate_up_surrogate_matches_two_logical_qlora_branches(monkeypatch) + monkeypatch.undo() + _assert_fused_gate_up_input_gradient_matches_two_exact_logical_wrappers(monkeypatch) + _assert_fused_gate_up_factor_only_backward_does_not_materialize_base(monkeypatch) + _assert_fused_gate_up_backward_rejects_any_master_mutation(monkeypatch) diff --git a/tests/models/test_glm52_exact_lm_head_loss_integration.py b/tests/models/test_glm52_exact_lm_head_loss_integration.py index 5f1a9ea6..cce5ca6d 100644 --- a/tests/models/test_glm52_exact_lm_head_loss_integration.py +++ b/tests/models/test_glm52_exact_lm_head_loss_integration.py @@ -20,7 +20,7 @@ def _tiny_exact_head() -> Glm52ExactTP16LmHeadLoraLinear: return Glm52ExactTP16LmHeadLoraLinear.from_module(base, r=1, lora_alpha=1) -def test_per_token_ce_routes_exact_head_before_generic_tp_and_fp32_paths(monkeypatch: pytest.MonkeyPatch) -> None: +def test_exact_lm_head_loss_routing_weight_and_fsdp_policy(monkeypatch: pytest.MonkeyPatch) -> None: captures = {} def _fake_exact(hidden, weight, labels, **kwargs): @@ -60,8 +60,15 @@ def _fake_exact(hidden, weight, labels, **kwargs): assert captures["weight"] is weight assert captures["labels"] is labels + with monkeypatch.context() as case_patch: + _assert_causallm_exact_head_admits_its_tp_group_and_rejects_z_loss(case_patch) + _assert_exact_head_weight_and_server_loss_selector_never_materialize_delta() + _assert_exact_head_fsdp_ignores_only_replicated_a() -def test_causallm_exact_head_admits_its_tp_group_and_rejects_z_loss(monkeypatch: pytest.MonkeyPatch) -> None: + +def _assert_causallm_exact_head_admits_its_tp_group_and_rejects_z_loss( + monkeypatch: pytest.MonkeyPatch, +) -> None: causallm_impl = importlib.import_module("xorl.ops.loss.causallm_loss") lm_head = nn.Module() lm_head._glm52_exact_tp16_lm_head = True @@ -101,7 +108,7 @@ def test_causallm_exact_head_admits_its_tp_group_and_rejects_z_loss(monkeypatch: ) -def test_exact_head_weight_and_server_loss_selector_never_materialize_delta() -> None: +def _assert_exact_head_weight_and_server_loss_selector_never_materialize_delta() -> None: lm_head = _tiny_exact_head() lm_head._xorl_fsdp_sharded_lm_head_loss = True assert get_lm_head_weight(lm_head, fsdp_sharded_loss=True) is lm_head.weight @@ -112,7 +119,7 @@ def test_exact_head_weight_and_server_loss_selector_never_materialize_delta() -> assert runner._get_loss_lm_head_module(lm_head) is lm_head -def test_exact_head_fsdp_ignores_only_replicated_a() -> None: +def _assert_exact_head_fsdp_ignores_only_replicated_a() -> None: lm_head = _tiny_exact_head() lm_head._glm52_exact_replicated_parameter_names = ("lora_A",) diff --git a/tests/models/test_glm52_exact_lm_head_qlora.py b/tests/models/test_glm52_exact_lm_head_qlora.py index b88fd7b1..a535e35f 100644 --- a/tests/models/test_glm52_exact_lm_head_qlora.py +++ b/tests/models/test_glm52_exact_lm_head_qlora.py @@ -22,7 +22,7 @@ ) -def _component(tp_rank: int = 0) -> Glm52ExactTP16LmHeadSelectedLogprob: +def _component(tp_rank: int = 0, tp_group=None) -> Glm52ExactTP16LmHeadSelectedLogprob: shard = glm52_lm_head_shard(tp_rank) return Glm52ExactTP16LmHeadSelectedLogprob( tp_rank=tp_rank, @@ -30,6 +30,7 @@ def _component(tp_rank: int = 0) -> Glm52ExactTP16LmHeadSelectedLogprob: vocab_end=shard.vocab_end, padded_vocab_start=shard.padded_vocab_start, padded_vocab_end=shard.padded_vocab_end, + tp_group=tp_group, ) @@ -51,7 +52,7 @@ def _meta_operands(rows: int = 2): return hidden, weight, lora_A, lora_B, token_ids -def test_official_tp16_shards_match_sglang_padding_and_rank_order() -> None: +def _assert_official_tp16_shards_match_sglang_padding_and_rank_order() -> None: shards = [glm52_lm_head_shard(rank) for rank in range(GLM52_LM_HEAD_TP_SIZE)] assert GLM52_LM_HEAD_PADDED_VOCAB_SIZE == GLM52_LM_HEAD_VOCAB_SIZE @@ -68,7 +69,7 @@ def test_official_tp16_shards_match_sglang_padding_and_rank_order() -> None: glm52_lm_head_shard(16) -def test_component_fails_closed_on_shard_ranges() -> None: +def _assert_component_fails_closed_on_shard_ranges() -> None: component = _component(7) shard = glm52_lm_head_shard(7) @@ -96,7 +97,7 @@ def test_component_fails_closed_on_shard_ranges() -> None: ) -def test_operand_contract_is_official_local_bf16_rank_one_and_stride_exact() -> None: +def _assert_operand_contract_is_official_local_bf16_rank_one_and_stride_exact() -> None: component = _component() operands = _meta_operands() component._validate_operands(*operands, require_cuda=False) @@ -154,7 +155,7 @@ def test_operand_contract_is_official_local_bf16_rank_one_and_stride_exact() -> ) -def test_cpu_rejection_happens_before_sglang_import_or_group_use() -> None: +def _assert_cpu_rejection_happens_before_sglang_import_or_group_use() -> None: before = {name for name in sys.modules if name == "sglang" or name.startswith("sglang.")} component = _component() hidden, weight, lora_A, lora_B, token_ids = _meta_operands() @@ -166,7 +167,7 @@ def test_cpu_rejection_happens_before_sglang_import_or_group_use() -> None: assert after == before -def test_rank_order_vocab_reshape_is_byte_exact_and_has_identity_token_mapping() -> None: +def _assert_rank_order_vocab_reshape_is_byte_exact_and_has_identity_token_mapping() -> None: rows = 2 local = GLM52_LM_HEAD_LOCAL_VOCAB_SIZE row_values = torch.arange(rows * local, dtype=torch.float32).reshape(rows, local) @@ -192,22 +193,7 @@ def test_rank_order_vocab_reshape_is_byte_exact_and_has_identity_token_mapping() ) -def test_effective_factor_views_are_the_live_bf16_bytes() -> None: - component = _component(5) - lora_A = torch.arange(GLM52_LM_HEAD_HIDDEN_SIZE, dtype=torch.float32).sub_(101).div_(257).unsqueeze(0) - lora_B = torch.arange(GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, dtype=torch.float32).sub_(503).div_(997).unsqueeze(1) - master_A = lora_A.clone() - master_B = lora_B.clone() - - effective_A, effective_B = component.effective_factor_views(lora_A, lora_B) - - assert torch.equal(effective_A.view(torch.uint8), master_A.to(torch.bfloat16).view(torch.uint8)) - assert torch.equal(effective_B.view(torch.uint8), master_B.to(torch.bfloat16).view(torch.uint8)) - assert torch.equal(lora_A, master_A) - assert torch.equal(lora_B, master_B) - - -def test_local_fp32_surrogate_vjp_matches_standalone_qlora_reference() -> None: +def _assert_local_fp32_surrogate_vjp_matches_standalone_qlora_reference() -> None: hidden = torch.arange(24, dtype=torch.float32).reshape(3, 8).sub_(7).div_(19).to(torch.bfloat16) weight = torch.arange(56, dtype=torch.float32).reshape(7, 8).sub_(23).div_(37).to(torch.bfloat16) effective_A = torch.arange(8, dtype=torch.float32).sub_(3).div_(11).reshape(1, 8).to(torch.bfloat16) @@ -238,7 +224,19 @@ def test_local_fp32_surrogate_vjp_matches_standalone_qlora_reference() -> None: assert torch.equal(grad_B, reference_B.grad) -def test_custom_boundary_is_grad_enabled_and_saves_effective_factor_bytes() -> None: +def test_exact_lm_head_cpu_contract(monkeypatch) -> None: + _assert_official_tp16_shards_match_sglang_padding_and_rank_order() + _assert_component_fails_closed_on_shard_ranges() + _assert_tp_group_validation_rejects_size_order_rank_and_backend(monkeypatch) + monkeypatch.undo() + _assert_operand_contract_is_official_local_bf16_rank_one_and_stride_exact() + _assert_cpu_rejection_happens_before_sglang_import_or_group_use() + _assert_rank_order_vocab_reshape_is_byte_exact_and_has_identity_token_mapping() + _assert_local_fp32_surrogate_vjp_matches_standalone_qlora_reference() + _assert_custom_boundary_is_grad_enabled_and_saves_effective_factor_bytes() + + +def _assert_custom_boundary_is_grad_enabled_and_saves_effective_factor_bytes() -> None: captures = {} class FakeComponent: @@ -270,10 +268,9 @@ def _surrogate_vjp(hidden, weight, effective_A, effective_B, token_ids, grad_log assert torch.equal(lora_B.grad, torch.ones_like(lora_B)) -def test_tp_group_validation_rejects_size_order_rank_and_backend(monkeypatch) -> None: - component = _component(3) +def _assert_tp_group_validation_rejects_size_order_rank_and_backend(monkeypatch) -> None: group = object() - component.bind_tp_group(group) + component = _component(3, tp_group=group) state = { "world": 16, "group_rank": 3, @@ -377,7 +374,8 @@ def test_official_local_shard_literal_v2_bytes_tail_and_surrogate_gradients() -> weight_bytes = local_weight.view(torch.uint8).clone() A_bytes = lora_A.view(torch.uint8).clone() B_bytes = lora_B.view(torch.uint8).clone() - effective_A, effective_B = component.effective_factor_views(lora_A, lora_B) + effective_A = lora_A.detach().to(torch.bfloat16).contiguous() + effective_B = lora_B.detach().to(torch.bfloat16).contiguous() batch_info = lm_head_impl._single_adapter_lm_head_batch_info(device.index, rows) direct_base, _direct_lse = head_v2_full_logits_with_lse(hidden, local_weight) diff --git a/tests/models/test_glm52_exact_moe_construction.py b/tests/models/test_glm52_exact_moe_construction.py index ffcbedcf..fa4ab769 100644 --- a/tests/models/test_glm52_exact_moe_construction.py +++ b/tests/models/test_glm52_exact_moe_construction.py @@ -1,23 +1,31 @@ from __future__ import annotations from dataclasses import dataclass +from types import MethodType import pytest import torch +from torch import nn from torch.distributed._tensor import Replicate, Shard from tests.models.test_glm52_qlora import _meta_model, _official_config +from xorl.models.transformers.glm5.exact_absorbed_kv_b_qlora import ( + Glm52ExactTP1AbsorbedKvBBlockFP8QLoRA, +) +from xorl.models.transformers.glm5.exact_dense_mlp import Glm52ExactTP1DenseMLP from xorl.models.transformers.glm5.exact_lm_head_qlora import ( Glm52ExactTP16LmHeadLoraLinear, Glm52ExactTP16LmHeadSelectedLogprob, glm52_lm_head_shard, ) +from xorl.models.transformers.glm5.exact_qlora import Glm52ExactTP1BlockFP8QLoRALinear from xorl.models.transformers.glm5.exact_routed_experts_qlora import ( Glm52ExactEP16BlockFP8QLoRARoutedExperts, ) from xorl.models.transformers.glm5.exact_shared_expert_qlora import ( Glm52ExactTP16SharedExpertBlockFP8QLoRA, ) +from xorl.models.transformers.glm5.modeling_glm5 import Glm5MoEBlock from xorl.models.transformers.glm5.parallelize import get_ep_plan from xorl.models.transformers.glm5.qlora import GLM52_QLORA_FACTOR_COUNT, prepare_glm52_block_fp8_qlora from xorl.qlora.modules.block_fp8_linear import BlockFP8QLoRALinear @@ -25,6 +33,10 @@ from xorl.server.runner.adapters.sharded_state import discover_adapter_layouts +_ORDINARY_ATTENTION_PROJECTIONS = ("q_a_proj", "kv_a_proj_with_mqa", "q_b_proj", "o_proj") +_ALL_ATTENTION_PROJECTIONS = (*_ORDINARY_ATTENTION_PROJECTIONS, "kv_b_proj") + + @dataclass(frozen=True) class _EPState: ep_enabled: bool @@ -101,6 +113,66 @@ def _patch_exact_world16_rank7(monkeypatch: pytest.MonkeyPatch) -> object: return group +def _empty_moe_block() -> Glm5MoEBlock: + block = Glm5MoEBlock.__new__(Glm5MoEBlock) + nn.Module.__init__(block) + block.routed_scaling_factor = 2.5 + return block + + +def _assert_canonical_moe_routed_and_shared_boundary_policy() -> None: + block = _empty_moe_block() + experts = Glm52ExactEP16BlockFP8QLoRARoutedExperts(128, 128, ep_rank=7, device="cpu") + captured = {} + + def forward(self, hidden, routing, selected_experts=None, **kwargs): + captured.update( + hidden=hidden, + routing=routing, + selected_experts=selected_experts, + local_ids=kwargs["sglang_ep_native_local_ids"], + routed_scaling_factor=kwargs["routed_scaling_factor"], + ) + return torch.ones_like(hidden) + + experts.forward = MethodType(forward, experts) + block.experts = experts + hidden = torch.zeros((3, 128), dtype=torch.bfloat16) + routing = torch.arange(24, dtype=torch.float32).reshape(3, 8).div_(32) + global_ids = torch.arange(24, dtype=torch.int64).reshape(3, 8).add_(112) + local_ids = torch.arange(24, dtype=torch.int32).reshape(3, 8).remainder_(16) + + output = block._canonical_routed_local_partial(hidden, routing, global_ids, local_ids) + + assert torch.equal(output, torch.ones_like(hidden)) + assert captured["hidden"] is hidden + assert captured["routing"] is routing + assert captured["selected_experts"] is global_ids + assert captured["local_ids"] is local_ids + assert captured["routed_scaling_factor"] == 2.5 + + shared = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="meta") + captured.clear() + + def shared_forward(self, shared_hidden, *, contributor_ordinal): + captured.update(hidden=shared_hidden, contributor_ordinal=contributor_ordinal) + return torch.full_like(shared_hidden, 0.5) + + shared.forward = MethodType(shared_forward, shared) + block.shared_experts = shared + hidden = torch.zeros((3, 6144), dtype=torch.bfloat16) + + output = block._canonical_shared_local_partial( + hidden, + contributor_ordinal=7, + contributor_count=16, + ) + + assert torch.equal(output, torch.full_like(hidden, 0.5)) + assert captured["hidden"] is hidden + assert captured["contributor_ordinal"] == 7 + + def test_glm52_exact_moe_construction_preserves_complete_global_inventory_and_sources( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -130,6 +202,8 @@ def test_glm52_exact_moe_construction_preserves_complete_global_inventory_and_so assert all(parameter.dtype is torch.float32 for parameter in trainable.values()) assert all(factor.dtype is torch.float32 for factor in inventory.factors) + _assert_complete_attention_and_dense_inventory(model, inventory, trainable) + assert not any( isinstance(module, BlockFP8QLoRALinear) for name, module in model.named_modules() @@ -177,8 +251,45 @@ def test_glm52_exact_moe_construction_preserves_complete_global_inventory_and_so f"{routed_fqn}.{factor_name}" for factor_name in routed.logical_factor_names } + with monkeypatch.context() as admission_patch: + _assert_glm52_exact_moe_construction_admission_policy(admission_patch) + with monkeypatch.context() as layout_patch: + _assert_glm52_exact_moe_post_ep_layout_preserves_factor_fqns_and_owner_logical_shapes(layout_patch) + with monkeypatch.context() as head_patch: + _assert_glm52_complete_exact_construction_attaches_only_selected_logprob_lm_head(head_patch) + _assert_canonical_moe_routed_and_shared_boundary_policy() + + +def _assert_complete_attention_and_dense_inventory(model, inventory, trainable) -> None: + expected_attention_factors = { + f"model.layers.{layer_idx}.self_attn.{projection}.lora_{factor}" + for layer_idx in range(78) + for projection in _ALL_ATTENTION_PROJECTIONS + for factor in ("A", "B") + } + actual_attention_factors = {factor.name for factor in inventory.factors if factor.role.startswith("attention.")} + assert actual_attention_factors == expected_attention_factors + assert len(actual_attention_factors) == 78 * 5 * 2 == 780 + + for layer_idx, layer in enumerate(model.model.layers): + attention = layer.self_attn + prefix = f"model.layers.{layer_idx}.self_attn" + for projection in _ORDINARY_ATTENTION_PROJECTIONS: + module = getattr(attention, projection) + assert type(module) is Glm52ExactTP1BlockFP8QLoRALinear + assert module._source_fqn == f"{prefix}.{projection}" + assert type(attention.kv_b_proj) is Glm52ExactTP1AbsorbedKvBBlockFP8QLoRA + assert attention.kv_b_proj._source_fqn == f"{prefix}.kv_b_proj" + assert {name for name in trainable if name.startswith(f"{prefix}.")} == { + f"{prefix}.{projection}.lora_{factor}" for projection in _ALL_ATTENTION_PROJECTIONS for factor in ("A", "B") + } + + assert {name for name, module in model.named_modules() if isinstance(module, Glm52ExactTP1DenseMLP)} == { + f"model.layers.{layer_idx}.mlp" for layer_idx in range(3) + } + -def test_glm52_exact_moe_post_ep_layout_preserves_factor_fqns_and_owner_logical_shapes( +def _assert_glm52_exact_moe_post_ep_layout_preserves_factor_fqns_and_owner_logical_shapes( monkeypatch: pytest.MonkeyPatch, ) -> None: _patch_ep16_rank7(monkeypatch) @@ -234,7 +345,7 @@ def test_glm52_exact_moe_post_ep_layout_preserves_factor_fqns_and_owner_logical_ assert layouts[full_name].local_logical_offset[0] == 112 -def test_glm52_complete_exact_construction_attaches_only_selected_logprob_lm_head( +def _assert_glm52_complete_exact_construction_attaches_only_selected_logprob_lm_head( monkeypatch: pytest.MonkeyPatch, ) -> None: group = _patch_exact_world16_rank7(monkeypatch) @@ -266,9 +377,27 @@ def test_glm52_complete_exact_construction_attaches_only_selected_logprob_lm_hea lm_head(torch.empty((1, 6_144), device="meta", dtype=torch.bfloat16)) -@pytest.mark.parametrize( - ("override", "message"), - ( +def _assert_glm52_exact_moe_construction_rejects_incomplete_dependency_flags( + monkeypatch: pytest.MonkeyPatch, + override: dict[str, object], + message: str, +) -> None: + _patch_ep16_rank7(monkeypatch) + config = _exact_moe_config() + for name, value in override.items(): + setattr(config, name, value) + model = _meta_model(config) + + with pytest.raises(ValueError, match=message): + prepare_glm52_block_fp8_qlora(model, config, adapter_rank=1, adapter_alpha=1) + + assert not any("lora_" in name for name, _ in model.named_parameters()) + + +def _assert_glm52_exact_moe_construction_admission_policy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + for override, message in ( ({"_glm52_exact_active_lora_dense_component": False}, "requires the exact active-LoRA dense component"), ( {"_glm52_exact_active_lora_attention_component": False}, @@ -287,47 +416,15 @@ def test_glm52_complete_exact_construction_attaches_only_selected_logprob_lm_hea ), ({"_sparse_mla_enabled": False}, "requires sparse_mla_enabled=true"), ({"_ep_dispatch": "deepep"}, "requires ep_dispatch='alltoall'"), - ), -) -def test_glm52_exact_moe_construction_rejects_incomplete_dependency_flags_before_mutation( - monkeypatch: pytest.MonkeyPatch, - override: dict[str, object], - message: str, -) -> None: - _patch_ep16_rank7(monkeypatch) - config = _exact_moe_config() - for name, value in override.items(): - setattr(config, name, value) - model = _meta_model(config) - - with pytest.raises(ValueError, match=message): - prepare_glm52_block_fp8_qlora(model, config, adapter_rank=1, adapter_alpha=1) - - assert not any("lora_" in name for name, _ in model.named_parameters()) - + ): + _assert_glm52_exact_moe_construction_rejects_incomplete_dependency_flags(monkeypatch, override, message) -@pytest.mark.parametrize( - "state", - ( + for state in ( _EPState(ep_enabled=False, ep_size=16, ep_rank=0), _EPState(ep_enabled=True, ep_size=8, ep_rank=7), - ), -) -def test_glm52_exact_moe_construction_rejects_non_ep16_before_mutation( - monkeypatch: pytest.MonkeyPatch, - state: _EPState, -) -> None: - monkeypatch.setattr("xorl.models.transformers.glm5.qlora.get_parallel_state", lambda: state) - config = _exact_moe_config() - model = _meta_model(config) + ): + _assert_glm52_exact_moe_construction_rejects_non_ep16(monkeypatch, state) - with pytest.raises(RuntimeError, match="require initialized EP16"): - prepare_glm52_block_fp8_qlora(model, config, adapter_rank=1, adapter_alpha=1) - - assert not any("lora_" in name for name, _ in model.named_parameters()) - - -def test_glm52_exact_lm_head_rejects_non_world16_before_mutation(monkeypatch: pytest.MonkeyPatch) -> None: group = object() state = _EPState( ep_enabled=True, @@ -347,9 +444,25 @@ def test_glm52_exact_lm_head_rejects_non_world16_before_mutation(monkeypatch: py assert not any("lora_" in name for name, _ in model.named_parameters()) + for rank, alpha in ((1, 2), (2, 1)): + _assert_glm52_exact_moe_construction_rejects_non_rank1_alpha1(monkeypatch, rank, alpha) + + +def _assert_glm52_exact_moe_construction_rejects_non_ep16( + monkeypatch: pytest.MonkeyPatch, + state: _EPState, +) -> None: + monkeypatch.setattr("xorl.models.transformers.glm5.qlora.get_parallel_state", lambda: state) + config = _exact_moe_config() + model = _meta_model(config) + + with pytest.raises(RuntimeError, match="require initialized EP16"): + prepare_glm52_block_fp8_qlora(model, config, adapter_rank=1, adapter_alpha=1) + + assert not any("lora_" in name for name, _ in model.named_parameters()) + -@pytest.mark.parametrize(("rank", "alpha"), ((16, 16), (1, 2), (2, 1))) -def test_glm52_exact_moe_construction_rejects_non_rank1_alpha1_before_mutation( +def _assert_glm52_exact_moe_construction_rejects_non_rank1_alpha1( monkeypatch: pytest.MonkeyPatch, rank: int, alpha: int, diff --git a/tests/models/test_glm52_exact_moe_modeling.py b/tests/models/test_glm52_exact_moe_modeling.py deleted file mode 100644 index 23d61b2d..00000000 --- a/tests/models/test_glm52_exact_moe_modeling.py +++ /dev/null @@ -1,77 +0,0 @@ -from __future__ import annotations - -from types import MethodType - -import torch -from torch import nn - -from xorl.models.transformers.glm5.exact_routed_experts_qlora import ( - Glm52ExactEP16BlockFP8QLoRARoutedExperts, -) -from xorl.models.transformers.glm5.exact_shared_expert_qlora import ( - Glm52ExactTP16SharedExpertBlockFP8QLoRA, -) -from xorl.models.transformers.glm5.modeling_glm5 import Glm5MoEBlock - - -def _empty_block() -> Glm5MoEBlock: - block = Glm5MoEBlock.__new__(Glm5MoEBlock) - nn.Module.__init__(block) - block.routed_scaling_factor = 2.5 - return block - - -def test_canonical_routed_boundary_passes_both_global_and_owner_local_ids() -> None: - block = _empty_block() - experts = Glm52ExactEP16BlockFP8QLoRARoutedExperts(128, 128, ep_rank=7, device="cpu") - captured = {} - - def forward(self, hidden, routing, selected_experts=None, **kwargs): - captured.update( - hidden=hidden, - routing=routing, - selected_experts=selected_experts, - local_ids=kwargs["sglang_ep_native_local_ids"], - routed_scaling_factor=kwargs["routed_scaling_factor"], - ) - return torch.ones_like(hidden) - - experts.forward = MethodType(forward, experts) - block.experts = experts - hidden = torch.zeros((3, 128), dtype=torch.bfloat16) - routing = torch.arange(24, dtype=torch.float32).reshape(3, 8).div_(32) - global_ids = torch.arange(24, dtype=torch.int64).reshape(3, 8).add_(112) - local_ids = torch.arange(24, dtype=torch.int32).reshape(3, 8).remainder_(16) - - output = block._canonical_routed_local_partial(hidden, routing, global_ids, local_ids) - - assert torch.equal(output, torch.ones_like(hidden)) - assert captured["hidden"] is hidden - assert captured["routing"] is routing - assert captured["selected_experts"] is global_ids - assert captured["local_ids"] is local_ids - assert captured["routed_scaling_factor"] == 2.5 - - -def test_canonical_shared_boundary_calls_the_exact_root_with_contributor_ordinal() -> None: - block = _empty_block() - shared = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="meta") - captured = {} - - def forward(self, hidden, *, contributor_ordinal): - captured.update(hidden=hidden, contributor_ordinal=contributor_ordinal) - return torch.full_like(hidden, 0.5) - - shared.forward = MethodType(forward, shared) - block.shared_experts = shared - hidden = torch.zeros((3, 6144), dtype=torch.bfloat16) - - output = block._canonical_shared_local_partial( - hidden, - contributor_ordinal=7, - contributor_count=16, - ) - - assert torch.equal(output, torch.full_like(hidden, 0.5)) - assert captured["hidden"] is hidden - assert captured["contributor_ordinal"] == 7 diff --git a/tests/models/test_glm52_exact_qlora.py b/tests/models/test_glm52_exact_qlora.py index 4322764c..cc3c00da 100644 --- a/tests/models/test_glm52_exact_qlora.py +++ b/tests/models/test_glm52_exact_qlora.py @@ -32,7 +32,7 @@ def run(input, effective_A, effective_B): return run -def test_exact_tp1_wrapper_admits_only_rank1_alpha1_without_bias_or_aqn() -> None: +def _assert_exact_tp1_configuration_and_runtime_admission_policy() -> None: module = _module() assert module.contract_version == GLM52_EXACT_TP1_QLORA_CONTRACT_VERSION @@ -55,8 +55,21 @@ def test_exact_tp1_wrapper_admits_only_rank1_alpha1_without_bias_or_aqn() -> Non with pytest.raises(ValueError, match="only lora_rank=1 and lora_alpha=1"): module.set_runtime_lora_config(2, 1) + before = {name for name in sys.modules if name == "sglang" or name.startswith("sglang.")} + with pytest.raises(TypeError, match="requires BF16 activations"): + module(torch.zeros(1, 8, dtype=torch.float32)) + with pytest.raises(ValueError, match="contiguous sampler-layout"): + module(torch.zeros(8, 2, dtype=torch.bfloat16).transpose(0, 1)) + with pytest.raises(RuntimeError, match="requires CUDA"): + module(torch.zeros(1, 8, dtype=torch.bfloat16)) -def test_exact_tp1_model_dtype_move_preserves_packed_state_master_dtype_and_identity() -> None: + after = {name for name in sys.modules if name == "sglang" or name.startswith("sglang.")} + assert after == before + + _assert_exact_tp1_model_dtype_move_preserves_packed_state_master_dtype_and_identity() + + +def _assert_exact_tp1_model_dtype_move_preserves_packed_state_master_dtype_and_identity() -> None: module = _module() with torch.no_grad(): module.packed_weight_f32.copy_( @@ -84,7 +97,8 @@ def test_exact_tp1_model_dtype_move_preserves_packed_state_master_dtype_and_iden assert torch.equal(module.weight_block_scales, scale_bytes) -def test_exact_tp1_wrapper_rounds_master_factors_once_before_value_forward(monkeypatch) -> None: +def test_exact_tp1_forward_surrogate_backward_and_safety_policy(monkeypatch) -> None: + _assert_exact_tp1_configuration_and_runtime_admission_policy() module = _module() base_weight = torch.arange(48, dtype=torch.float32).reshape(6, 8).div_(97).to(torch.bfloat16) captures = [] @@ -99,8 +113,6 @@ def test_exact_tp1_wrapper_rounds_master_factors_once_before_value_forward(monke expected = _literal_cpu_value(base_weight, [])(input, effective_A, effective_B) assert torch.equal(output, expected) - -def test_exact_tp1_surrogate_backward_matches_effective_factor_qlora_reference(monkeypatch) -> None: module = _module() base_weight = torch.arange(48, dtype=torch.float32).reshape(6, 8).sub_(17).div_(41).to(torch.bfloat16) monkeypatch.setattr(module, "_dequantize_weight", lambda: base_weight.float()) @@ -141,8 +153,10 @@ def test_exact_tp1_surrogate_backward_matches_effective_factor_qlora_reference(m assert torch.equal(module.lora_A.grad, reference_A.grad) assert torch.equal(module.lora_B.grad, reference_B.grad) + _assert_exact_tp1_backward_safety_policy(monkeypatch) + -def test_exact_tp1_factor_only_backward_does_not_materialize_base(monkeypatch) -> None: +def _assert_exact_tp1_backward_safety_policy(monkeypatch) -> None: module = _module() base_weight = torch.zeros(6, 8, dtype=torch.bfloat16) monkeypatch.setattr(module, "_exact_forward_value", _literal_cpu_value(base_weight, [])) @@ -157,8 +171,6 @@ def test_exact_tp1_factor_only_backward_does_not_materialize_base(monkeypatch) - assert module.lora_A.grad is not None assert module.lora_B.grad is not None - -def test_exact_tp1_backward_rejects_master_mutation(monkeypatch) -> None: module = _module() base_weight = torch.zeros(6, 8, dtype=torch.bfloat16) monkeypatch.setattr(module, "_exact_forward_value", _literal_cpu_value(base_weight, [])) @@ -172,21 +184,6 @@ def test_exact_tp1_backward_rejects_master_mutation(monkeypatch) -> None: output.float().sum().backward() -def test_exact_tp1_contract_fails_before_any_sglang_import() -> None: - before = {name for name in sys.modules if name == "sglang" or name.startswith("sglang.")} - module = _module() - - with pytest.raises(TypeError, match="requires BF16 activations"): - module(torch.zeros(1, 8, dtype=torch.float32)) - with pytest.raises(ValueError, match="contiguous sampler-layout"): - module(torch.zeros(8, 2, dtype=torch.bfloat16).transpose(0, 1)) - with pytest.raises(RuntimeError, match="requires CUDA"): - module(torch.zeros(1, 8, dtype=torch.bfloat16)) - - after = {name for name in sys.modules if name == "sglang" or name.startswith("sglang.")} - assert after == before - - @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") @pytest.mark.parametrize( diff --git a/tests/models/test_glm52_exact_routed_experts_qlora.py b/tests/models/test_glm52_exact_routed_experts_qlora.py index f291a3ef..00426465 100644 --- a/tests/models/test_glm52_exact_routed_experts_qlora.py +++ b/tests/models/test_glm52_exact_routed_experts_qlora.py @@ -158,7 +158,7 @@ def _standalone_hybrid_routed_vjp( ) -def test_routed_bank_contract_is_strict_rank_local_ep16_moe_tp1() -> None: +def test_routed_bank_topology_remap_and_physical_buffer_policy() -> None: module = _module(7) assert isinstance(module, Glm52NativeBlockFP8Experts) @@ -205,8 +205,11 @@ def test_routed_bank_contract_is_strict_rank_local_ep16_moe_tp1() -> None: with pytest.raises(ValueError, match="only lora_rank=1"): module.set_runtime_lora_config(1, 2) + _assert_one_batched_global_grid_proves_all_16_owner_by_16_slot_remaps() + _assert_physical_sampler_factor_buffer_policy() -def test_one_batched_global_grid_proves_all_16_owner_by_16_slot_remaps() -> None: + +def _assert_one_batched_global_grid_proves_all_16_owner_by_16_slot_remaps() -> None: global_grid = torch.arange(_GLOBAL_EXPERTS, dtype=torch.int64).reshape(_EP_SIZE, _LOCAL_EXPERTS) owner_maps = torch.stack( [localize_glm52_ep16_expert_ids(global_grid.contiguous(), owner) for owner in range(_EP_SIZE)] @@ -227,11 +230,12 @@ def test_one_batched_global_grid_proves_all_16_owner_by_16_slot_remaps() -> None localize_glm52_ep16_expert_ids(invalid, 0) -def test_all_owner_slot_physical_buffers_are_full_width_bf16_and_not_outer_tp_sliced() -> None: +def _assert_physical_sampler_factor_buffer_policy() -> None: for owner in range(_EP_SIZE): module = _module(owner) _fill_distinguishable_factors(module) - buffers = module.physical_factor_buffers() + effective = tuple(getattr(module, name).to(torch.bfloat16).contiguous() for name in module.logical_factor_names) + buffers = module._physical_factor_buffers(*effective) assert tuple(buffers["gate_up_lora_a_weights"].shape) == (8, 1, 2, _HIDDEN) assert tuple(buffers["gate_up_lora_b_weights"].shape) == (8, 16, 2 * _INTERMEDIATE, 1) assert tuple(buffers["down_lora_a_weights"].shape) == (8, 16, 1, _INTERMEDIATE) @@ -266,8 +270,10 @@ def test_all_owner_slot_physical_buffers_are_full_width_bf16_and_not_outer_tp_sl for buffer in buffers.values(): assert torch.count_nonzero(buffer[1:]) == 0 + _assert_post_ep_owner_local_factor_banks_produce_same_views() + -def test_post_ep_owner_local_factor_banks_produce_the_same_physical_sampler_views() -> None: +def _assert_post_ep_owner_local_factor_banks_produce_same_views() -> None: module = _module(7) _fill_distinguishable_factors(module) global_factors = tuple( @@ -287,7 +293,12 @@ def test_post_ep_owner_local_factor_banks_produce_the_same_physical_sampler_view @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_all_256_owner_slots_run_literal_sampler_partials_routing_outputs_and_logical_vjps() -> None: +def test_routed_experts_literal_sampler_and_gradient_policy() -> None: + _assert_all_256_owner_slots_run_literal_sampler_partials_routing_outputs_and_logical_vjps() + _assert_routed_gradient_edge_and_mixed_owner_policy() + + +def _assert_all_256_owner_slots_run_literal_sampler_partials_routing_outputs_and_logical_vjps() -> None: pytest.importorskip("sglang") if torch.cuda.get_device_capability()[0] != 9: pytest.skip("the qualified GLM-5.2 routed component requires Hopper") @@ -297,7 +308,6 @@ def test_all_256_owner_slots_run_literal_sampler_partials_routing_outputs_and_lo for owner in range(_EP_SIZE): module = _module(owner, device) _load_distinguishable_base(module) - _fill_distinguishable_factors(module) global_ids = global_grid[owner].reshape(_LOCAL_EXPERTS, 1).contiguous() hidden = ( torch.arange(_LOCAL_EXPERTS * _HIDDEN, dtype=torch.float32, device=device) @@ -309,18 +319,18 @@ def test_all_256_owner_slots_run_literal_sampler_partials_routing_outputs_and_lo .requires_grad_(True) ) routing = ((global_ids.float() + 1) / 512).contiguous().requires_grad_(True) - trace = module.sampler_value_trace(hidden.detach(), routing.detach(), global_ids) + base_output = None + if owner == 0: + with torch.no_grad(): + for name in module.logical_factor_names: + getattr(module, name).zero_() + base_output = module(hidden.detach(), routing.detach(), selected_experts=global_ids).detach() + _fill_distinguishable_factors(module) output = module(hidden, routing, selected_experts=global_ids) - assert torch.equal(output.detach(), trace.owner_output) - assert torch.count_nonzero(trace.gate_up_base) > 0 - assert torch.count_nonzero(trace.down_base_routed) > 0 - assert torch.count_nonzero(trace.gate_up_post_lora) > 0 - assert torch.count_nonzero(trace.activated) > 0 - assert torch.count_nonzero(trace.down_post_lora_routed) > 0 - assert not torch.equal(trace.gate_up_base, trace.gate_up_post_lora) - assert not torch.equal(trace.down_base_routed, trace.down_post_lora_routed) - assert torch.equal(output, trace.down_post_lora_routed[:, 0]) + assert torch.count_nonzero(output) > 0 + if base_output is not None: + assert not torch.equal(output.detach(), base_output) # A positive logical cotangent keeps every intentionally routed slot # distinguishable; an alternating cotangent can legitimately cancel @@ -363,29 +373,20 @@ def test_all_256_owner_slots_run_literal_sampler_partials_routing_outputs_and_lo assert torch.count_nonzero(gradient) > 0 if owner == 0: - half_trace = module.sampler_value_trace( + half_output = module( hidden.detach(), (routing.detach() * 0.5).contiguous(), - global_ids, - ) - assert torch.equal(half_trace.gate_up_post_lora, trace.gate_up_post_lora) - torch.testing.assert_close( - half_trace.down_post_lora_routed.float(), - trace.down_post_lora_routed.float() * 0.5, - rtol=0, - atol=2**-8, - ) + selected_experts=global_ids, + ).detach() torch.testing.assert_close( - half_trace.owner_output.float(), - trace.owner_output.float() * 0.5, + half_output.float(), + output.detach().float() * 0.5, rtol=0, atol=2**-8, ) -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_zero_token_routed_slots_receive_exact_zero_bank_gradients_and_stride_mismatches_fail() -> None: +def _assert_routed_gradient_edge_and_mixed_owner_policy() -> None: pytest.importorskip("sglang") if torch.cuda.get_device_capability()[0] != 9: pytest.skip("the qualified GLM-5.2 routed component requires Hopper") @@ -416,13 +417,11 @@ def test_zero_token_routed_slots_receive_exact_zero_bank_gradients_and_stride_mi with pytest.raises(ValueError, match="route-major"): module(hidden.detach(), routing_strided, selected_experts=global_ids) + _assert_all_sentinel_owner_returns_zero_and_zero_gradients() + _assert_topk8_mixed_owner_hybrid_vjps_match_standalone_reference() -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_all_sentinel_owner_returns_exact_zero_and_zero_gradients() -> None: - pytest.importorskip("sglang") - if torch.cuda.get_device_capability()[0] != 9: - pytest.skip("the qualified GLM-5.2 routed component requires Hopper") + +def _assert_all_sentinel_owner_returns_zero_and_zero_gradients() -> None: device = torch.device("cuda") module = _module(5, device) _load_zero_base(module) @@ -443,9 +442,7 @@ def test_all_sentinel_owner_returns_exact_zero_and_zero_gradients() -> None: assert torch.count_nonzero(gradient) == 0 -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires Hopper CUDA") -def test_topk8_mixed_owner_hybrid_vjps_match_standalone_reference() -> None: +def _assert_topk8_mixed_owner_hybrid_vjps_match_standalone_reference() -> None: pytest.importorskip("sglang") if torch.cuda.get_device_capability()[0] != 9: pytest.skip("the qualified GLM-5.2 routed component requires Hopper") diff --git a/tests/models/test_glm52_exact_shared_expert_qlora.py b/tests/models/test_glm52_exact_shared_expert_qlora.py index e8e1ae92..0770270e 100644 --- a/tests/models/test_glm52_exact_shared_expert_qlora.py +++ b/tests/models/test_glm52_exact_shared_expert_qlora.py @@ -7,13 +7,23 @@ import torch.nn.functional as F from torch import nn -from xorl.distributed.canonical_moe import CanonicalMoEGraphMetadata, canonical_moe_reduce_reference +from xorl.distributed.canonical_moe import CanonicalMoEGraphMetadata from xorl.models.transformers.glm5.exact_shared_expert_qlora import ( GLM52_EXACT_TP16_SHARED_EXPERT_QLORA_CONTRACT_VERSION, Glm52ExactTP16SharedExpertBlockFP8QLoRA, ) +def _canonical_moe_reference(partials: torch.Tensor, metadata: CanonicalMoEGraphMetadata) -> torch.Tensor: + level = [partials[index] for index in range(partials.shape[0])] + while len(level) > 1: + level = [(level[index] + level[index + 1]).to(torch.bfloat16) for index in range(0, len(level), 2)] + result = level[0] + result = result.clone() + result[~metadata.valid_mask] = 0 + return result + + def _pattern( shape: tuple[int, ...], *, @@ -63,9 +73,8 @@ def _load_base(module: Glm52ExactTP16SharedExpertBlockFP8QLoRA) -> None: ) -@pytest.mark.parametrize( - ("kwargs", "message"), - [ +def test_shared_expert_construction_and_runtime_admission_policy() -> None: + for kwargs, message in ( ({"hidden_size": 4096}, "hidden_size=6144"), ({"intermediate_size": 4096}, "intermediate_size=2048"), ({"tp_size": 8}, "requires TP16"), @@ -73,14 +82,44 @@ def _load_base(module: Glm52ExactTP16SharedExpertBlockFP8QLoRA) -> None: ({"lora_alpha": 2}, "rank=1 and alpha=1"), ({"bias": True}, "bias-free"), ({"enable_aqn": True}, "rejects adaptive quantization noise"), - ], -) -def test_shared_expert_construction_fails_closed(kwargs: dict, message: str) -> None: - with pytest.raises(ValueError, match=message): - Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="meta", **kwargs) + ): + with pytest.raises(ValueError, match=message): + Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="meta", **kwargs) + + module = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="cpu") + with pytest.raises(TypeError, match="contributor_ordinal must be an integer"): + module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=True) + with pytest.raises(ValueError, match=r"must be in \[0, 16\)"): + module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=16) + with pytest.raises(TypeError, match="requires BF16 activations"): + module(torch.zeros(1, 6144), contributor_ordinal=0) + with pytest.raises(ValueError, match="input width"): + module(torch.zeros(1, 128, dtype=torch.bfloat16), contributor_ordinal=0) + with pytest.raises(ValueError, match="contiguous sampler-layout"): + module(torch.zeros(6144, 2, dtype=torch.bfloat16).transpose(0, 1), contributor_ordinal=0) + with pytest.raises(RuntimeError, match="requires CUDA"): + module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=0) + with pytest.raises(RuntimeError, match="cannot run independently"): + module.gate_proj(torch.zeros(1, 6144, dtype=torch.bfloat16)) + with pytest.raises(RuntimeError, match="cannot bypass active LoRA"): + module.gate_proj.forward_partition( + torch.zeros(1, 6144, dtype=torch.bfloat16), + output_range=(0, 128), + ) + + module.tp_size = 8 + with pytest.raises(RuntimeError, match="runtime contract was mutated"): + module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=0) + module.tp_size = 16 + module.gate_proj.lora_A = nn.Parameter(module.gate_proj.lora_A.to(torch.bfloat16)) + with pytest.raises(TypeError, match="gate_proj.lora_A must remain FP32"): + module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=0) + + _assert_shared_expert_logical_and_checkpoint_state_policy() + _assert_shared_expert_native_base_views_use_output_rows_and_input_columns() -def test_shared_expert_registers_one_logical_state_and_preserves_fp32_masters() -> None: +def _assert_shared_expert_logical_and_checkpoint_state_policy() -> None: module = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="cpu") assert module.contract_version == GLM52_EXACT_TP16_SHARED_EXPERT_QLORA_CONTRACT_VERSION parameters = dict(module.named_parameters()) @@ -114,8 +153,6 @@ def test_shared_expert_registers_one_logical_state_and_preserves_fp32_masters() if name.endswith(("packed_weight_f32", "weight_scale_inv")) ) - -def test_shared_expert_checkpoint_sources_are_canonical_and_immutable() -> None: module = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="meta") prefix = "model.layers.3.mlp.shared_experts" @@ -143,7 +180,18 @@ def test_shared_expert_physical_factor_views_match_pinned_sglang_tp_slices() -> module = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="cpu") _fill_factors(module) ordinal = 11 - actual = module.physical_factor_views(ordinal) + effective = tuple( + factor.to(torch.bfloat16).contiguous() + for factor in ( + module.gate_proj.lora_A, + module.gate_proj.lora_B, + module.up_proj.lora_A, + module.up_proj.lora_B, + module.down_proj.lora_A, + module.down_proj.lora_B, + ) + ) + actual = module._physical_factor_views_from_effective(*effective, ordinal) gate_up_A = torch.cat( ( @@ -195,7 +243,7 @@ def test_shared_expert_physical_factor_views_match_pinned_sglang_tp_slices() -> assert all(tensor.is_contiguous() for tensor in (actual.gate_up_A, actual.gate_up_B, actual.down_A, actual.down_B)) -def test_shared_expert_native_base_views_use_output_rows_and_input_columns() -> None: +def _assert_shared_expert_native_base_views_use_output_rows_and_input_columns() -> None: module = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="cpu") ordinal = 5 gate_weight = torch.empty((2048, 6144), dtype=torch.float8_e4m3fn) @@ -235,37 +283,6 @@ def test_shared_expert_native_base_views_use_output_rows_and_input_columns() -> assert torch.equal(actual.down_scales[:, 0], down_scales[:, ordinal]) -def test_shared_expert_runtime_contract_fails_before_sglang_kernel_import() -> None: - module = Glm52ExactTP16SharedExpertBlockFP8QLoRA(device="cpu") - with pytest.raises(TypeError, match="contributor_ordinal must be an integer"): - module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=True) - with pytest.raises(ValueError, match=r"must be in \[0, 16\)"): - module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=16) - with pytest.raises(TypeError, match="requires BF16 activations"): - module(torch.zeros(1, 6144), contributor_ordinal=0) - with pytest.raises(ValueError, match="input width"): - module(torch.zeros(1, 128, dtype=torch.bfloat16), contributor_ordinal=0) - with pytest.raises(ValueError, match="contiguous sampler-layout"): - module(torch.zeros(6144, 2, dtype=torch.bfloat16).transpose(0, 1), contributor_ordinal=0) - with pytest.raises(RuntimeError, match="requires CUDA"): - module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=0) - with pytest.raises(RuntimeError, match="cannot run independently"): - module.gate_proj(torch.zeros(1, 6144, dtype=torch.bfloat16)) - with pytest.raises(RuntimeError, match="cannot bypass active LoRA"): - module.gate_proj.forward_partition( - torch.zeros(1, 6144, dtype=torch.bfloat16), - output_range=(0, 128), - ) - - module.tp_size = 8 - with pytest.raises(RuntimeError, match="runtime contract was mutated"): - module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=0) - module.tp_size = 16 - module.gate_proj.lora_A = nn.Parameter(module.gate_proj.lora_A.to(torch.bfloat16)) - with pytest.raises(TypeError, match="gate_proj.lora_A must remain FP32"): - module(torch.zeros(1, 6144, dtype=torch.bfloat16), contributor_ordinal=0) - - def _manual_local_vjp( module: Glm52ExactTP16SharedExpertBlockFP8QLoRA, input: torch.Tensor, @@ -407,7 +424,7 @@ def test_official_shared_expert_actual_operands_fold_and_surrogate_vjp() -> None cold_output = module(input, contributor_ordinal=ordinal) warm_output = module(input, contributor_ordinal=ordinal) - factors = module.physical_factor_views(ordinal) + factors = module._physical_factor_views_from_effective(*effective, ordinal) base = module._physical_base_views(ordinal) batch_info = LoRABatchInfo( use_cuda_graph=False, @@ -480,7 +497,7 @@ def test_official_shared_expert_actual_operands_fold_and_surrogate_vjp() -> None torch.tensor([0], dtype=torch.int64, device=device), capacity=1, ) - canonical = canonical_moe_reduce_reference(partials, metadata) + canonical = _canonical_moe_reference(partials, metadata) sampler_slots = CanonicalRowSlots.from_positions( torch.tensor([0], dtype=torch.int64, device=device), capacity=1, diff --git a/tests/models/test_glm52_index_share_checkpoint.py b/tests/models/test_glm52_index_share_checkpoint.py index 13aaa43c..e30f44e7 100644 --- a/tests/models/test_glm52_index_share_checkpoint.py +++ b/tests/models/test_glm52_index_share_checkpoint.py @@ -1,4 +1,4 @@ -from types import SimpleNamespace +from types import MethodType, SimpleNamespace import pytest import torch @@ -11,6 +11,16 @@ IndexShareMode, ) from xorl.models.transformers.glm5.layer_plan import Glm52LayerPlan +from xorl.server.runner.model_runner import ModelRunner +from xorl.trainers.trainer import Trainer + + +class _RetainingModel: + def __init__(self) -> None: + self.release_count = 0 + + def release_index_share_context(self) -> None: + self.release_count += 1 def _producer_shared_plan() -> Glm52LayerPlan: @@ -27,8 +37,48 @@ def _producer_shared_plan() -> Glm52LayerPlan: @pytest.mark.cpu -@pytest.mark.parametrize("use_reentrant", [False, True]) -def test_checkpointed_producer_shared_backward_reuses_detached_payload(use_reentrant): +def test_checkpointed_index_share_lifecycle_policy(): + for use_reentrant in (False, True): + _assert_checkpointed_producer_shared_backward_reuses_detached_payload(use_reentrant) + _assert_mode_owned_success_and_failure_cleanup_is_idempotent() + _assert_forward_failures_release_offline_and_server_caller_contexts() + + +def _assert_forward_failures_release_offline_and_server_caller_contexts(): + trainer = Trainer.__new__(Trainer) + model = _RetainingModel() + trainer.model = model + trainer._all_model_parts = MethodType(lambda _self: [model], trainer) + trainer._forward_backward_impl = MethodType( + lambda _self, _micro_batches, _global_valid_tokens: (_ for _ in ()).throw(RuntimeError("loss failed")), + trainer, + ) + + assert trainer._index_share_forward_kwargs(model, IndexShareMode.TRAINING_WITH_BACKWARD) == { + "index_share_mode": IndexShareMode.TRAINING_WITH_BACKWARD + } + with pytest.raises(RuntimeError, match="loss failed"): + trainer._forward_backward([], None) + assert model.release_count == 1 + + runner = ModelRunner.__new__(ModelRunner) + server_model = _RetainingModel() + runner.model = server_model + runner.model_parts = [] + runner._forward_loop_impl = MethodType( + lambda _self, *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("backward failed")), + runner, + ) + + assert runner._index_share_forward_kwargs(IndexShareMode.FORWARD_ONLY) == { + "index_share_mode": IndexShareMode.FORWARD_ONLY + } + with pytest.raises(RuntimeError, match="backward failed"): + runner._forward_loop([], "causallm_loss", {}, compute_backward=True) + assert server_model.release_count == 1 + + +def _assert_checkpointed_producer_shared_backward_reuses_detached_payload(use_reentrant): plan = _producer_shared_plan() manager = IndexShareContextManager(plan, (0, 2)) context = manager.begin(mode=IndexShareMode.TRAINING_WITH_BACKWARD) @@ -71,8 +121,7 @@ def shared_layer(hidden_states: torch.Tensor) -> torch.Tensor: assert context.lifecycle is IndexShareLifecycle.CLOSED -@pytest.mark.cpu -def test_mode_owned_success_and_failure_cleanup_is_idempotent(): +def _assert_mode_owned_success_and_failure_cleanup_is_idempotent(): plan = _producer_shared_plan() manager = IndexShareContextManager(plan, (0, 2)) diff --git a/tests/models/test_glm52_native_fp8.py b/tests/models/test_glm52_native_fp8.py index 83017d4e..31ee298d 100644 --- a/tests/models/test_glm52_native_fp8.py +++ b/tests/models/test_glm52_native_fp8.py @@ -32,7 +32,7 @@ def _fp8_values(shape): return (((values % 31) - 15).to(torch.float32).reshape(shape)).to(torch.float8_e4m3fn) -def test_hf_config_preserves_quantization_metadata_and_roundtrips(): +def _assert_hf_config_preserves_quantization_metadata_and_roundtrips(): hf_config = SimpleNamespace( hidden_size=256, intermediate_size=512, @@ -50,20 +50,17 @@ def test_hf_config_preserves_quantization_metadata_and_roundtrips(): assert validate_glm52_native_fp8_config(restored.quantization_config)["weight_block_size"] == [128, 128] -@pytest.mark.parametrize( - "field,value", - [ +def _assert_native_config_rejects_nonofficial_contract(): + for field, value in ( ("quant_method", "int8"), ("fmt", "e5m2"), ("activation_scheme", "static"), ("weight_block_size", [64, 128]), - ], -) -def test_native_config_rejects_nonofficial_contract(field, value): - config = dict(OFFICIAL_QUANT_CONFIG) - config[field] = value - with pytest.raises(ValueError, match="Unsupported"): - validate_glm52_native_fp8_config(config) + ): + config = dict(OFFICIAL_QUANT_CONFIG) + config[field] = value + with pytest.raises(ValueError, match="Unsupported"): + validate_glm52_native_fp8_config(config) class _TinyNativeModel(nn.Module): @@ -72,7 +69,7 @@ def __init__(self): self.proj = NativeBlockFP8Linear(256, 384) -def test_pair_buffer_emits_exact_dcp_visible_parameter_names_and_bytes(): +def _assert_pair_buffer_emits_exact_dcp_visible_parameter_names_and_bytes(): model = _TinyNativeModel() buffer = NativeBlockFP8PairBuffer(model, {"official.proj": "proj"}) weight = _fp8_values((384, 256)) @@ -89,7 +86,7 @@ def test_pair_buffer_emits_exact_dcp_visible_parameter_names_and_bytes(): assert set(model.state_dict()) == {"proj.packed_weight_f32", "proj.weight_scale_inv"} -def test_pair_buffer_fails_closed_on_missing_duplicate_and_bad_scale(): +def _assert_pair_buffer_fails_closed_on_missing_duplicate_and_bad_scale(): model = _TinyNativeModel() weight = _fp8_values((384, 256)) scale = torch.ones(3, 2, dtype=torch.float32) @@ -110,7 +107,7 @@ def test_pair_buffer_fails_closed_on_missing_duplicate_and_bad_scale(): bad_scale.try_consume("official.proj.weight_scale_inv", scale.to(torch.bfloat16)) -def test_pair_buffer_rejects_noninjective_target_mapping(): +def _assert_pair_buffer_rejects_noninjective_target_mapping(): model = _TinyNativeModel() with pytest.raises(ValueError, match="duplicate targets"): NativeBlockFP8PairBuffer( @@ -156,7 +153,7 @@ def _tiny_native_config(): ) -def test_model_replaces_only_quantized_dense_shared_and_expert_modules(): +def _assert_model_replaces_only_quantized_dense_shared_and_expert_modules(): model = Glm5ForCausalLM(_tiny_native_config()) modules = dict(model.named_modules()) @@ -173,7 +170,7 @@ def test_model_replaces_only_quantized_dense_shared_and_expert_modules(): ) -def test_sparse_mla_native_kv_weight_uses_module_forward_and_preserves_layout(): +def _assert_sparse_mla_native_kv_weight_uses_module_forward_and_preserves_layout(): class MaterializingNativeLinear(NativeBlockFP8Linear): def __init__(self): super().__init__(4, 16) @@ -203,7 +200,7 @@ def forward(self, input=None, *, return_dequantized_weight=False): assert torch.equal(w_vc, expected[:, 4:]) -def test_canonical_router_uses_serving_dispatch_and_defers_routed_scale(monkeypatch): +def test_canonical_router_and_native_expert_runtime_policy(monkeypatch): model = Glm5ForCausalLM(_tiny_native_config()) block = model.model.layers[1].mlp hidden = torch.linspace(-1, 1, steps=2 * block.config.hidden_size, dtype=torch.bfloat16).reshape( @@ -245,6 +242,8 @@ def fake_serving_topk(hidden_states, router_logits, correction_bias, **kwargs): "routed_scaling_factor": 2.5, } + _assert_glm_expert_state_is_frozen_exact_and_scoring_only() + class _TinyExpertModel(nn.Module): def __init__(self, num_experts=4): @@ -255,7 +254,7 @@ def __init__(self, num_experts=4): self.model.layers[0].mlp.experts = Glm52NativeBlockFP8Experts(num_experts, 256, 128) -def test_glm_expert_state_is_frozen_exact_and_scoring_only(): +def _assert_glm_expert_state_is_frozen_exact_and_scoring_only(): module = Glm52NativeBlockFP8Experts(2, 256, 128) gate_up = _fp8_values((2, 256, 256)) gate_up_scale = torch.arange(8, dtype=torch.float32).reshape(2, 2, 2) / 11 @@ -277,7 +276,7 @@ def test_glm_expert_state_is_frozen_exact_and_scoring_only(): module(hidden.detach(), routing) -def test_expert_pair_buffer_fuses_local_gate_up_down_bytes_and_scales(): +def _assert_expert_pair_buffer_fuses_local_gate_up_down_bytes_and_scales(): model = _TinyExpertModel() buffer = NativeBlockFP8ExpertPairBuffer(model, ep_rank=1, ep_size=2, num_experts=4) pieces = {} @@ -321,14 +320,14 @@ def test_expert_pair_buffer_fuses_local_gate_up_down_bytes_and_scales(): assert result[next(k for k in result if k.endswith("gate_up_weight_scale_inv"))].dtype is torch.float32 -def test_expert_pair_buffer_rejects_bad_rank_and_target_count(): +def _assert_expert_pair_buffer_rejects_bad_rank_and_target_count(): with pytest.raises(ValueError, match="Invalid"): NativeBlockFP8ExpertPairBuffer(_TinyExpertModel(), ep_rank=2, ep_size=2, num_experts=4) with pytest.raises(ValueError, match="declares 2 experts"): NativeBlockFP8ExpertPairBuffer(_TinyExpertModel(num_experts=2), ep_rank=0, ep_size=2, num_experts=4) -def test_grouped_dense_and_expert_handlers_own_disjoint_native_pair_families(): +def _assert_grouped_dense_and_expert_handlers_own_disjoint_native_pair_families(): model = Glm5ForCausalLM(_tiny_native_config()) dense_handler = model.get_checkpoint_handler( ep_rank=0, @@ -402,3 +401,16 @@ def test_grouped_dense_and_expert_handlers_own_disjoint_native_pair_families(): assert expert_skip is not None assert expert_skip("model.layers.1.mlp.experts.2.gate_proj.weight") assert not expert_skip("model.layers.1.mlp.experts.1.gate_proj.weight") + + +def test_glm52_native_fp8_configuration_model_buffer_and_checkpoint_contract(): + _assert_hf_config_preserves_quantization_metadata_and_roundtrips() + _assert_native_config_rejects_nonofficial_contract() + _assert_model_replaces_only_quantized_dense_shared_and_expert_modules() + _assert_sparse_mla_native_kv_weight_uses_module_forward_and_preserves_layout() + _assert_pair_buffer_emits_exact_dcp_visible_parameter_names_and_bytes() + _assert_pair_buffer_fails_closed_on_missing_duplicate_and_bad_scale() + _assert_pair_buffer_rejects_noninjective_target_mapping() + _assert_expert_pair_buffer_fuses_local_gate_up_down_bytes_and_scales() + _assert_expert_pair_buffer_rejects_bad_rank_and_target_count() + _assert_grouped_dense_and_expert_handlers_own_disjoint_native_pair_families() diff --git a/tests/models/test_glm52_qlora.py b/tests/models/test_glm52_qlora.py index f4431c81..da723ca1 100644 --- a/tests/models/test_glm52_qlora.py +++ b/tests/models/test_glm52_qlora.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections import Counter from contextlib import contextmanager import pytest @@ -76,7 +77,7 @@ def _meta_model(config: Glm5Config | None = None) -> Glm5ForCausalLM: return Glm5ForCausalLM(config) -def test_glm52_full_block_fp8_qlora_inventory_is_exact_and_fail_closed() -> None: +def test_glm52_block_fp8_qlora_policy(monkeypatch) -> None: model = _meta_model() inventory = prepare_glm52_block_fp8_qlora( @@ -108,7 +109,7 @@ def test_glm52_full_block_fp8_qlora_inventory_is_exact_and_fail_closed() -> None assert trainable == inventory.factor_names assert all(factor.dtype is torch.float32 for factor in inventory.factors) - assert inventory.role_counts == { + assert Counter(target.role for target in inventory.targets) == { "attention.q_a_proj": 78, "attention.q_b_proj": 78, "attention.kv_a_proj_with_mqa": 78, @@ -140,8 +141,12 @@ def test_glm52_full_block_fp8_qlora_inventory_is_exact_and_fail_closed() -> None assert type(model.model.layers[0].self_attn.indexer.weights_proj) is nn.Linear assert model.model.layers[0].self_attn.indexer.weights_proj.weight.dtype is torch.bfloat16 + with monkeypatch.context() as ep_patch: + _assert_glm52_routed_banks_select_only_the_ep_local_checkpoint_slice(ep_patch) + _assert_glm52_exact_component_and_admission_contract() -def test_glm52_exact_dense_component_preserves_logical_inventory_with_three_physical_fused_roots() -> None: + +def _assert_glm52_exact_dense_component_preserves_logical_inventory_with_three_physical_fused_roots() -> None: config = _official_config() config._glm52_exact_active_lora_dense_component = True config._ep_dispatch = "alltoall" @@ -174,20 +179,20 @@ def test_glm52_exact_dense_component_preserves_logical_inventory_with_three_phys assert root.down_proj._source_fqn == f"{prefix}.down_proj" -@pytest.mark.parametrize(("rank", "alpha"), ((16, 16), (1, 2), (2, 1))) -def test_glm52_exact_dense_component_rejects_non_rank1_alpha1_before_mutation(rank: int, alpha: int) -> None: - config = _official_config() - config._glm52_exact_active_lora_dense_component = True - config._ep_dispatch = "alltoall" - model = _meta_model(config) +def _assert_glm52_exact_dense_component_rejects_non_rank1_alpha1_before_mutation() -> None: + for rank, alpha in ((1, 2), (2, 1)): + config = _official_config() + config._glm52_exact_active_lora_dense_component = True + config._ep_dispatch = "alltoall" + model = _meta_model(config) - with pytest.raises(ValueError, match="requires adapter_rank=1 and adapter_alpha=1"): - prepare_glm52_block_fp8_qlora(model, config, adapter_rank=rank, adapter_alpha=alpha) + with pytest.raises(ValueError, match="requires adapter_rank=1 and adapter_alpha=1"): + prepare_glm52_block_fp8_qlora(model, config, adapter_rank=rank, adapter_alpha=alpha) - assert not any("lora_" in name for name, _ in model.named_parameters()) + assert not any("lora_" in name for name, _ in model.named_parameters()) -def test_glm52_routed_banks_select_only_the_ep_local_checkpoint_slice(monkeypatch) -> None: +def _assert_glm52_routed_banks_select_only_the_ep_local_checkpoint_slice(monkeypatch) -> None: class _EP16State: ep_enabled = True ep_size = 16 @@ -205,7 +210,7 @@ class _EP16State: assert all(module.expert_offset == 112 for module in routed_banks) -def test_glm52_qlora_rejects_missing_or_wrong_shape_target_before_adapterization() -> None: +def _assert_glm52_qlora_rejects_missing_or_wrong_shape_target_before_adapterization() -> None: model = _meta_model() model.model.layers[12].self_attn.q_a_proj = nn.Linear(6144, 1024, bias=False, device="meta") @@ -215,7 +220,7 @@ def test_glm52_qlora_rejects_missing_or_wrong_shape_target_before_adapterization assert not any("lora_" in name for name, _ in model.named_parameters()) -def test_glm52_qlora_rejects_missing_official_indexer_exclusion_before_adapterization() -> None: +def _assert_glm52_qlora_rejects_missing_official_indexer_exclusion_before_adapterization() -> None: config = _official_config() config.quantization_config["modules_to_not_convert"].remove("model.layers.0.self_attn.indexers_proj") model = _meta_model(config) @@ -226,26 +231,24 @@ def test_glm52_qlora_rejects_missing_official_indexer_exclusion_before_adapteriz assert not any("lora_" in name for name, _ in model.named_parameters()) -@pytest.mark.parametrize( - ("override", "message"), - [ +def _assert_glm52_qlora_rejects_unsupported_construction_modes() -> None: + cases = ( ({"_moe_implementation": "eager"}, "moe_implementation='triton'"), ({"_ep_dispatch": "alltoall"}, "ep_dispatch='deepep'"), ({"_glm52_exact_contract": True}, "cannot use the scoring-only exact contract"), ({"_glm52_block_fp8_qlora": False}, "block_fp8_qlora_training=true"), - ], -) -def test_glm52_qlora_rejects_unsupported_construction_modes(override: dict, message: str) -> None: - config = _official_config() - for name, value in override.items(): - setattr(config, name, value) - model = _meta_model(config) + ) + for override, message in cases: + config = _official_config() + for name, value in override.items(): + setattr(config, name, value) + model = _meta_model(config) - with pytest.raises(ValueError, match=message): - prepare_glm52_block_fp8_qlora(model, config, adapter_rank=16, adapter_alpha=16) + with pytest.raises(ValueError, match=message): + prepare_glm52_block_fp8_qlora(model, config, adapter_rank=16, adapter_alpha=16) -def test_glm5_training_mode_admits_only_explicit_product_tuple() -> None: +def _assert_glm5_training_mode_admits_only_explicit_product_tuple() -> None: config = _official_config() validate_glm5_training_mode( config, @@ -275,7 +278,7 @@ def test_glm5_training_mode_admits_only_explicit_product_tuple() -> None: ) -def test_glm5_training_mode_uses_alltoall_only_for_complete_exact_active_lora() -> None: +def _assert_glm5_training_mode_uses_alltoall_only_for_complete_exact_active_lora() -> None: config = _official_config() set_glm52_exact_active_lora(config, enabled=True) @@ -307,7 +310,11 @@ def test_glm5_training_mode_uses_alltoall_only_for_complete_exact_active_lora() ) -def test_block_fp8_qlora_scale_storage_covers_partial_edge_tiles() -> None: - module = BlockFP8QLoRALinear(6144, 576, r=4, lora_alpha=4, device=torch.device("meta")) - - assert module.weight_block_scales.shape == (5, 192) +def _assert_glm52_exact_component_and_admission_contract() -> None: + _assert_glm52_exact_dense_component_preserves_logical_inventory_with_three_physical_fused_roots() + _assert_glm52_exact_dense_component_rejects_non_rank1_alpha1_before_mutation() + _assert_glm52_qlora_rejects_missing_or_wrong_shape_target_before_adapterization() + _assert_glm52_qlora_rejects_missing_official_indexer_exclusion_before_adapterization() + _assert_glm52_qlora_rejects_unsupported_construction_modes() + _assert_glm5_training_mode_admits_only_explicit_product_tuple() + _assert_glm5_training_mode_uses_alltoall_only_for_complete_exact_active_lora() diff --git a/tests/models/test_glm5_flashmla_sparse_mla.py b/tests/models/test_glm5_flashmla_sparse_mla.py index d01dbd1a..fe02eb6c 100644 --- a/tests/models/test_glm5_flashmla_sparse_mla.py +++ b/tests/models/test_glm5_flashmla_sparse_mla.py @@ -18,7 +18,7 @@ pytestmark = [pytest.mark.cpu] -def test_flashmla_batch_flatten_offsets_only_valid_indices(): +def _assert_flashmla_batch_flatten_offsets_only_valid_indices(): q = torch.zeros(2, 2, 3, 4) kv = torch.zeros(2, 3, 4) indices = torch.tensor( @@ -47,6 +47,7 @@ def test_flashmla_batch_flatten_offsets_only_valid_indices(): def test_flashmla_backward_compacts_valid_rows_and_scatter_zeros(monkeypatch): + _assert_flashmla_batch_flatten_offsets_only_valid_indices() captured = {} def fake_forward(q, kv, indices, scaling): @@ -93,8 +94,12 @@ def fake_backward(q, kv, out, grad_out, indices, lse, scaling): assert torch.count_nonzero(q.grad[3]).item() == 0 torch.testing.assert_close(kv.grad, torch.full_like(kv.grad, 5)) + _assert_flashmla_all_invalid_rows_have_zero_output_and_gradients(monkeypatch) + with monkeypatch.context() as case_patch: + _assert_flashmla_backend_fails_closed_outside_production_envelope(case_patch) -def test_flashmla_all_invalid_rows_have_zero_output_and_gradients(monkeypatch): + +def _assert_flashmla_all_invalid_rows_have_zero_output_and_gradients(monkeypatch): backward = Mock(side_effect=AssertionError("all-invalid input must bypass TileLang backward")) def fake_forward(q, kv, indices, scaling): @@ -120,7 +125,7 @@ def fake_forward(q, kv, indices, scaling): assert torch.equal(kv.grad, torch.zeros_like(kv)) -def test_flashmla_backend_fails_closed_outside_production_envelope(): +def _assert_flashmla_backend_fails_closed_outside_production_envelope(monkeypatch): q = torch.zeros(1, 1, 64, 576, dtype=torch.bfloat16) kv = torch.zeros(1, 2, 576, dtype=torch.bfloat16) indices = torch.zeros(1, 1, 2048, dtype=torch.int32) @@ -137,8 +142,6 @@ def test_flashmla_backend_fails_closed_outside_production_envelope(): backend="flashmla", ) - -def test_flashmla_constraint_rejects_unproven_shape_before_import(monkeypatch): q = torch.empty(1, 1, 63, 576, device="meta", dtype=torch.bfloat16) kv = torch.empty(1, 2, 576, device="meta", dtype=torch.bfloat16) indices = torch.empty(1, 1, 2048, device="meta", dtype=torch.int32) diff --git a/tests/models/test_glm5_support.py b/tests/models/test_glm5_support.py index 84a486a6..fcaadea7 100644 --- a/tests/models/test_glm5_support.py +++ b/tests/models/test_glm5_support.py @@ -152,13 +152,7 @@ def _tiny_config(**overrides) -> Glm5Config: # --------------------------------------------------------------------------- # -def test_glm5_config_is_standalone(): - config = Glm5Config() - assert isinstance(config, Glm5Config) - assert config.model_type == "xorl_glm5" - - -def test_glm5_declares_every_rmsnorm_serving_family(): +def test_glm5_support_policy(tmp_path, monkeypatch): config = _tiny_config() set_rmsnorm_mode("sglang_fused") try: @@ -181,8 +175,23 @@ def test_glm5_declares_every_rmsnorm_serving_family(): assert model.norm.mode == "sglang_fused" assert model.norm.family == RMS_NORM_FAMILY_RESIDUAL_TREE - -def test_from_hf_config_captures_mla_moe_dsa_and_mtp_fields(): + _assert_hf_config_captures_glm5_fields() + _assert_local_config_loader_routes_glm5(tmp_path) + with monkeypatch.context() as loader_patch: + _assert_local_config_loader_rejects_unsafe_values(loader_patch) + with monkeypatch.context() as ring_patch: + _assert_glm5_rejects_dsa_with_ring_attention(ring_patch) + with monkeypatch.context() as indexer_patch: + _assert_glm5_indexer_construction_policy(indexer_patch) + with monkeypatch.context() as sparse_patch: + _assert_sparse_mla_reference_wrapper_and_attention_integration_policy(sparse_patch) + _assert_glm5_checkpoint_filter_policy() + with monkeypatch.context() as adapter_patch: + _assert_glm5_adapter_sparse_kv_and_moe_dispatch_policy(adapter_patch) + _assert_glm5_forward_and_recompute_policy() + + +def _assert_hf_config_captures_glm5_fields(): hf_config = _namespace_from_dict(GLM_5_1_HF_CONFIG) config = Glm5Config.from_hf_config(hf_config) @@ -224,7 +233,7 @@ def test_from_hf_config_captures_mla_moe_dsa_and_mtp_fields(): assert config.model_type == "xorl_glm5" -def test_auto_load_local_xorl_config_routes_glm_moe_dsa(tmp_path): +def _assert_local_config_loader_routes_glm5(tmp_path): cfg_path = tmp_path / "config.json" cfg_path.write_text(json.dumps(GLM_5_1_HF_CONFIG)) @@ -233,9 +242,14 @@ def test_auto_load_local_xorl_config_routes_glm_moe_dsa(tmp_path): assert isinstance(config, Glm5Config) assert config.architectures == ["GlmMoeDsaForCausalLM"] assert config.n_routed_experts == 256 + registry = get_registry() + assert "GlmMoeDsaForCausalLM" in registry.supported_models + assert "Glm5ForCausalLM" in registry.supported_models + assert issubclass(GlmMoeDsaForCausalLM, Glm5ForCausalLM) + assert "GlmMoeDsaForCausalLM" in get_loader(config).description -def test_auto_load_local_xorl_config_rejects_non_json_values(monkeypatch): +def _assert_local_config_loader_rejects_unsafe_values(monkeypatch): class ForeignValue: pass @@ -247,8 +261,10 @@ class ForeignValue: with pytest.raises(ValueError, match="non-JSON value at \\$.hidden_size"): _load_local_xorl_config("unused", config_kwargs={}) + _assert_auto_load_rejects_dunder_keys(monkeypatch) + -def test_auto_load_local_xorl_config_rejects_dunder_keys(monkeypatch): +def _assert_auto_load_rejects_dunder_keys(monkeypatch): monkeypatch.setattr( "xorl.models.auto.PretrainedConfig.get_config_dict", lambda *_args, **_kwargs: ({"model_type": "glm_moe_dsa", "__class__": "foreign"}, {}), @@ -258,41 +274,7 @@ def test_auto_load_local_xorl_config_rejects_dunder_keys(monkeypatch): _load_local_xorl_config("unused", config_kwargs={}) -def test_glm5_default_construction_matches_glm_5_1_shape(): - """Defaults track the public `zai-org/GLM-5.1` config so a bare `Glm5Config()` - is always interpretable without an HF round-trip.""" - config = Glm5Config() - - assert config.hidden_size == 6144 - assert config.num_hidden_layers == 78 - assert config.num_attention_heads == 64 - assert config.vocab_size == 154880 - assert config.q_lora_rank == 2048 - assert config.n_routed_experts == 256 - assert config.num_experts_per_tok == 8 - - -# --------------------------------------------------------------------------- # -# Registry / loader -# --------------------------------------------------------------------------- # - - -def test_glm5_registered_under_both_arch_names(): - """Registry resolves both `GlmMoeDsaForCausalLM` (HF) and `Glm5ForCausalLM`.""" - reg = get_registry() - assert "GlmMoeDsaForCausalLM" in reg.supported_models - assert "Glm5ForCausalLM" in reg.supported_models - # Subclass relationship: `GlmMoeDsaForCausalLM` is just an alias. - assert issubclass(GlmMoeDsaForCausalLM, Glm5ForCausalLM) - - -def test_get_loader_resolves_glm_moe_dsa(): - config = Glm5Config.from_hf_config(_namespace_from_dict(GLM_5_1_HF_CONFIG)) - loader = get_loader(config) - assert "GlmMoeDsaForCausalLM" in loader.description - - -def test_build_foundation_model_rejects_glm5_dsa_with_ring_attention(monkeypatch): +def _assert_glm5_rejects_dsa_with_ring_attention(monkeypatch): class ParallelState: ringattn_size = 1 ringattn_enabled = True @@ -311,7 +293,7 @@ class ParallelState: # --------------------------------------------------------------------------- # -def test_indexer_has_four_glm_specific_projections(): +def _assert_glm5_indexer_construction_policy(monkeypatch): config = _tiny_config() indexer = Glm5DsaIndexer(config) @@ -323,76 +305,12 @@ def test_indexer_has_four_glm_specific_projections(): assert indexer.weights_proj.in_features == config.hidden_size assert indexer.weights_proj.out_features == config.index_n_heads + _assert_glm5_indexer_selection_policy() + with monkeypatch.context() as case_patch: + _assert_glm5_dsa_mask_policy(case_patch) -def test_indexer_contract_routes_head_projection_through_fp32_kernel(monkeypatch): - config = _tiny_config(indexer_types=["full"] * 4) - config._glm52_exact_contract = True - indexer = Glm5DsaIndexer(config).to(torch.bfloat16) - hidden = torch.randn(2, 3, config.hidden_size, dtype=torch.bfloat16) - q_compressed = torch.randn(2, 3, config.q_lora_rank, dtype=torch.bfloat16) - cos = torch.ones(2, 3, config.qk_rope_head_dim, dtype=torch.bfloat16) - sin = torch.zeros_like(cos) - calls = [] - - def fake_contract(input, weight): - calls.append((input.shape, weight.shape, input.dtype, weight.dtype)) - return torch.nn.functional.linear(input.float(), weight.float()) - - monkeypatch.setattr("xorl.models.transformers.glm5.indexer.bi_bf16_fp32_linear", fake_contract) - _, _, head_weights = indexer.project(hidden, q_compressed, (cos, sin)) - - assert calls == [ - ( - torch.Size([2, 3, config.hidden_size]), - torch.Size([config.index_n_heads, config.hidden_size]), - torch.bfloat16, - torch.bfloat16, - ) - ] - expected = fake_contract(hidden, indexer.weights_proj.weight) * (config.index_n_heads**-0.5) - assert torch.equal(head_weights, expected) - - -def test_attention_carries_indexer_module(): - """V0: indexer params exist on attention so checkpoints load cleanly.""" - config = _tiny_config() - attn = Glm5Attention(config, layer_idx=0) - assert isinstance(attn.indexer, Glm5DsaIndexer) - - -def test_indexer_select_topk_returns_correct_shape(): - """The dense fallback returns top-k indices for each query position. - - Causality is enforced by masking non-causal logits with -inf. When fewer - than `index_topk` valid keys exist for a query (early positions), `topk` - fills the remaining slots with arbitrary indices over -inf entries — so - the strong causality check only holds for the last query position, where - every key is valid. - """ - config = _tiny_config() - indexer = Glm5DsaIndexer(config) - B, S = 2, 8 - index_q = torch.randn(B, S, config.index_n_heads, config.index_head_dim) - index_k = torch.randn(B, S, config.index_head_dim) - head_weights = torch.randn(B, S, config.index_n_heads) - - indices = indexer.select_topk(index_q, index_k, head_weights, attention_mask=None) - - assert indices.shape == (B, S, config.index_topk) - assert indices.dtype == torch.int32 - # The indexer marks non-causal pad slots with -1 (kernel sentinel); - # all other slots are in [0, S). - valid = indices >= 0 - assert torch.all(indices[valid] < S) - # For the last query position (q=S-1), all S keys are causal — top-k - # picks must all be valid (no -1 padding). - last_q_indices = indices[:, -1, :] - assert torch.all(last_q_indices >= 0) - assert torch.all(last_q_indices < S) - - -def test_indexer_select_topk_respects_additive_attention_mask(): +def _assert_glm5_indexer_selection_policy(): """Masked positions should become `-1` sentinels before sparse-MLA. Eager causal masks use `torch.finfo(dtype).min`, not literal `-inf`. @@ -400,58 +318,38 @@ def test_indexer_select_topk_respects_additive_attention_mask(): budget on padding / packed-sequence boundaries and the tilelang path receives non-sentinel invalid indices. """ - config = _tiny_config(index_topk=3) - indexer = Glm5DsaIndexer(config) - B, S = 1, 5 - index_q = torch.randn(B, S, config.index_n_heads, config.index_head_dim) - index_k = torch.randn(B, S, config.index_head_dim) - head_weights = torch.randn(B, S, config.index_n_heads) - attention_mask = torch.full((B, S, S), torch.finfo(index_q.dtype).min) - rows = torch.arange(S) - attention_mask[:, rows, rows] = 0 - - indices = indexer.select_topk(index_q, index_k, head_weights, attention_mask=attention_mask) - - torch.testing.assert_close(indices[0, :, 0], torch.arange(S, dtype=torch.int32)) - assert torch.all(indices[0, :, 1:] == -1) - + for score_chunk_heads in (None, 1): + config = _tiny_config(index_topk=3, index_n_heads=4) + if score_chunk_heads is not None: + config.indexer_score_chunk_heads = score_chunk_heads + indexer = Glm5DsaIndexer(config) + index_q = torch.randn(B, S, config.index_n_heads, config.index_head_dim) + index_k = torch.randn(B, S, config.index_head_dim) + head_weights = torch.randn(B, S, config.index_n_heads) + attention_mask = torch.full((B, S, S), torch.finfo(index_q.dtype).min) + rows = torch.arange(S) + attention_mask[:, rows, rows] = 0 -def test_indexer_select_topk_supports_chunked_head_scoring(): - """Chunked scoring must preserve masking semantics used by sparse-MLA.""" - config = _tiny_config(index_topk=3, index_n_heads=4) - config.indexer_score_chunk_heads = 1 - indexer = Glm5DsaIndexer(config) - - B, S = 1, 5 - index_q = torch.randn(B, S, config.index_n_heads, config.index_head_dim) - index_k = torch.randn(B, S, config.index_head_dim) - head_weights = torch.randn(B, S, config.index_n_heads) - attention_mask = torch.full((B, S, S), torch.finfo(index_q.dtype).min) - rows = torch.arange(S) - attention_mask[:, rows, rows] = 0 + indices = indexer.select_topk(index_q, index_k, head_weights, attention_mask=attention_mask) - indices = indexer.select_topk(index_q, index_k, head_weights, attention_mask=attention_mask) + torch.testing.assert_close(indices[0, :, 0], torch.arange(S, dtype=torch.int32)) + assert torch.all(indices[0, :, 1:] == -1) - torch.testing.assert_close(indices[0, :, 0], torch.arange(S, dtype=torch.int32)) - assert torch.all(indices[0, :, 1:] == -1) + _assert_indexer_padding_mask_policy() + _assert_indexer_returns_sorted_indices() + _assert_indexer_blocked_selection() -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_indexer_padding_mask_detection_triggers_fast_path(): +def _assert_indexer_padding_mask_policy(): """A 2D padding mask (e.g., HuggingFace-style [B, S] with 1s for valid tokens and 0s for padded tail) should be detected and routed to the tilelang fast path with cu_ke clipped to the valid length.""" - try: - from xorl.ops.glm5_kernels.tilelang_indexer_fwd import tl_indexer_fwd_impl # noqa: F401 - except Exception: - pytest.skip("tilelang indexer fwd kernel unavailable") - torch.manual_seed(0) config = _tiny_config(index_topk=8, index_n_heads=4) - indexer = Glm5DsaIndexer(config).cuda() + indexer = Glm5DsaIndexer(config) B, S = 1, 32 # 2D padding mask: first 24 positions valid, last 8 padded. - pad_mask = torch.ones((B, S), device="cuda", dtype=torch.long) + pad_mask = torch.ones((B, S), dtype=torch.long) pad_mask[:, 24:] = 0 # Detection returns valid_len per batch @@ -466,17 +364,14 @@ def test_indexer_padding_mask_detection_triggers_fast_path(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_indexer_causal_mask_detection_triggers_fast_path(): - """An explicit pure-causal additive mask should be detected and routed to - the tilelang fast path (producing the same indices as both the - attention_mask=None path and the torch fallback).""" - try: - from xorl.ops.glm5_kernels.tilelang_indexer_fwd import tl_indexer_fwd_impl # noqa: F401 - except Exception: - pytest.skip("tilelang indexer fwd kernel unavailable") +def test_indexer_tilelang_fast_path_accepts_causal_mask_and_matches_blocked_torch(): + """Both fast-path mask forms match an independently forced torch path.""" + from xorl.ops.glm5_kernels.tilelang_indexer_fwd import tl_indexer_fwd_impl # noqa: F401 + torch.manual_seed(0) config = _tiny_config(index_topk=8, index_n_heads=4) indexer = Glm5DsaIndexer(config).cuda() + B, S = 1, 32 H, D = config.index_n_heads, config.index_head_dim iq = torch.randn((B, S, H, D), device="cuda", dtype=torch.bfloat16) @@ -486,58 +381,28 @@ def test_indexer_causal_mask_detection_triggers_fast_path(): causal_mask = torch.zeros((B, S, S), device="cuda", dtype=torch.float32) for q in range(S): causal_mask[:, q, q + 1 :] = torch.finfo(torch.float32).min - - # Detection itself should return True for the canonical causal mask assert indexer._is_pure_causal_mask(causal_mask, S, S, 0) - # And the masked call should give the same indices as the unmasked call - indices_no_mask = indexer.select_topk(iq, ik, weights, attention_mask=None) - indices_with_mask = indexer.select_topk(iq, ik, weights, attention_mask=causal_mask) - for b in range(B): - for s in range(S): - assert set(indices_no_mask[b, s].cpu().tolist()) - {-1} == set(indices_with_mask[b, s].cpu().tolist()) - { - -1 - } - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_indexer_tilelang_fast_path_matches_torch_reference(): - """The tilelang fast path (taken when attention_mask is None and inputs - are bf16/CUDA) should produce the same topk indices as the torch path. - Compared as sets per row because both paths sort indices ascending, - but the underlying scoring kernels may break ties differently.""" - try: - from xorl.ops.glm5_kernels.tilelang_indexer_fwd import tl_indexer_fwd_impl # noqa: F401 - except Exception: - pytest.skip("tilelang indexer fwd kernel unavailable") - torch.manual_seed(0) - config = _tiny_config(index_topk=8, index_n_heads=4) - indexer = Glm5DsaIndexer(config).cuda() - - B, S = 1, 32 - H, D = config.index_n_heads, config.index_head_dim - iq = torch.randn((B, S, H, D), device="cuda", dtype=torch.bfloat16) - ik = torch.randn((B, S, D), device="cuda", dtype=torch.bfloat16) - weights = torch.randn((B, S, H), device="cuda", dtype=torch.float32) - - # Tilelang fast path (attention_mask=None triggers it) indices_tl = indexer.select_topk(iq, ik, weights, attention_mask=None) + indices_tl_masked = indexer.select_topk(iq, ik, weights, attention_mask=causal_mask) - # Torch path (force via a strict causal mask) - causal_mask = torch.zeros((B, S, S), device="cuda", dtype=torch.float32) - for q in range(S): - causal_mask[:, q, q + 1 :] = torch.finfo(torch.float32).min + # A pure causal mask is itself eligible for TileLang, so it is not a + # reference oracle. Opting into blocked scoring rejects the fast path and + # forces the independent torch implementation. + config.indexer_score_query_block_size = 7 + config.indexer_score_key_block_size = 11 indices_ref = indexer.select_topk(iq, ik, weights, attention_mask=causal_mask) - # Per-row set comparison; ignore -1 sentinels for b in range(B): for s in range(S): tl_set = set(indices_tl[b, s].cpu().tolist()) - {-1} + masked_set = set(indices_tl_masked[b, s].cpu().tolist()) - {-1} ref_set = set(indices_ref[b, s].cpu().tolist()) - {-1} + assert tl_set == masked_set assert tl_set == ref_set, f"row mismatch at b={b} s={s}: tl={tl_set} ref={ref_set}" -def test_indexer_select_topk_returns_sorted_indices(): +def _assert_indexer_returns_sorted_indices(): """The indexer should sort top-k indices ascending by kv position (with -1 sentinels at the end). This is a perf optimization for the downstream sparse-MLA kernel — gathers become near-contiguous → better @@ -553,6 +418,12 @@ def test_indexer_select_topk_returns_sorted_indices(): head_weights = torch.randn(B, S, config.index_n_heads) indices = indexer.select_topk(index_q, index_k, head_weights, attention_mask=None) + assert indices.shape == (B, S, config.index_topk) + assert indices.dtype == torch.int32 + valid_indices = indices[indices >= 0] + assert torch.all(valid_indices < S) + assert torch.all(indices[:, -1, :] >= 0) + # Within each (batch, query) row, valid indices must be monotonically # non-decreasing; -1 sentinels (if any) must follow all valid indices. for b in range(indices.shape[0]): @@ -568,7 +439,7 @@ def test_indexer_select_topk_returns_sorted_indices(): assert neg_positions.min() > valid_positions.max(), f"-1 not at end: {row.tolist()}" -def test_indexer_blocked_select_topk_matches_dense(): +def _assert_indexer_blocked_selection(): """The memory-bounded scorer should preserve dense top-k semantics.""" torch.manual_seed(0) config = _tiny_config(index_topk=4, index_n_heads=4) @@ -588,8 +459,10 @@ def test_indexer_blocked_select_topk_matches_dense(): torch.testing.assert_close(torch.sort(blocked, dim=-1).values, torch.sort(dense, dim=-1).values) + _assert_indexer_blocked_select_topk_supports_query_offset() -def test_indexer_blocked_select_topk_supports_query_offset(): + +def _assert_indexer_blocked_select_topk_supports_query_offset(): """Ulysses query-sharded indexer rows should match dense full-sequence rows.""" torch.manual_seed(0) config = _tiny_config(index_topk=4, index_n_heads=4) @@ -615,7 +488,7 @@ def test_indexer_blocked_select_topk_supports_query_offset(): torch.testing.assert_close(local, dense[:, 3:6]) -def test_dense_dsa_mask_handles_2d_and_3d_attention_masks(): +def _assert_glm5_dsa_mask_policy(monkeypatch): config = _tiny_config(index_topk=2) attn = Glm5Attention(config, layer_idx=0) @@ -638,8 +511,10 @@ def test_dense_dsa_mask_handles_2d_and_3d_attention_masks(): assert built_3d.shape == (2, 1, 4, 4) assert torch.isneginf(built_3d[:, 0, :, 0]).all() + _assert_dsa_mask_gathers_ulysses_query_axis(monkeypatch) + -def test_dsa_mask_gathers_local_ulysses_query_axis(monkeypatch): +def _assert_dsa_mask_gathers_ulysses_query_axis(monkeypatch): config = _tiny_config() attn = Glm5Attention(config, layer_idx=0) group = object() @@ -670,7 +545,7 @@ def fake_gather_outputs(x, gather_dim, **_kwargs): @torch.no_grad() -def test_sparse_mla_torch_reference_matches_dense_at_full_topk(): +def _assert_sparse_mla_reference_and_wrapper_policy(monkeypatch): """Sparse-MLA with `topk >= seq_len` must equal dense softmax-attention over the same compressed KV (causality already imposed inside the sparse reference).""" @@ -700,9 +575,12 @@ def test_sparse_mla_torch_reference_matches_dense_at_full_topk(): torch.testing.assert_close(sparse_out, dense_out, atol=1e-5, rtol=1e-5) + _assert_sparse_mla_torch_reference_supports_query_offset() + _assert_sparse_mla_tilelang_wrapper(monkeypatch) + @torch.no_grad() -def test_sparse_mla_torch_reference_supports_query_offset(): +def _assert_sparse_mla_torch_reference_supports_query_offset(): """A local query shard with full-sequence indices should match the corresponding rows from the full-query reference.""" torch.manual_seed(0) @@ -729,7 +607,7 @@ def test_sparse_mla_torch_reference_supports_query_offset(): torch.testing.assert_close(local, full[:, 3:6], atol=1e-5, rtol=1e-5) -def test_sparse_mla_tilelang_wrapper_supports_local_query_full_kv(monkeypatch): +def _assert_sparse_mla_tilelang_wrapper(monkeypatch): """Ulysses sparse MLA sends local query rows with gathered full-sequence KV.""" B, S_q, S_kv, H, D = 2, 3, 7, 4, 5 kv_lora = D @@ -780,25 +658,9 @@ def apply(q_flat, kv_flat, indices_flat, scaling): @torch.no_grad() -def test_glm5_attention_sparse_forward_smoke(): - """Glm5Attention with `_sparse_mla_enabled=True` produces the right shape.""" - torch.manual_seed(0) - config = _tiny_config(index_topk=4) - config._sparse_mla_enabled = True - attn = Glm5Attention(config, layer_idx=0).eval() - rotary = RotaryEmbedding(config=config) - - B, S = 2, 6 - hidden = torch.randn(B, S, config.hidden_size) - pos_ids = torch.arange(S).unsqueeze(0).expand(B, -1) - cos, sin = rotary(hidden, pos_ids) - - out, _ = attn(hidden, position_embeddings=(cos, sin), attention_mask=None) - assert out.shape == (B, S, config.hidden_size) - +def _assert_sparse_mla_reference_wrapper_and_attention_integration_policy(monkeypatch): + _assert_sparse_mla_reference_and_wrapper_policy(monkeypatch) -@torch.no_grad() -def test_glm5_attention_sparse_ulysses_keeps_query_and_topk_local(monkeypatch): torch.manual_seed(0) config = _tiny_config(index_topk=2) config._sparse_mla_enabled = True @@ -854,9 +716,13 @@ def fake_gather_outputs(tensor, gather_dim, **_kwargs): assert captured["query_offset"] == S assert captured["mask_shape"] is None + monkeypatch.undo() + _assert_glm5_sparse_attention_matches_dense() + _assert_sparse_mla_dispatch_rejects_unknown_backend() + @torch.no_grad() -def test_glm5_attention_sparse_matches_dense_when_topk_covers_full_seq(): +def _assert_glm5_sparse_attention_matches_dense(): """When `index_topk >= seq_len`, the sparse and dense paths share the same MLA parameters and the sparse path's causal mask matches dense causal MLA — outputs must agree to numerical precision. @@ -883,7 +749,7 @@ def test_glm5_attention_sparse_matches_dense_when_topk_covers_full_seq(): torch.testing.assert_close(sparse_out, dense_out, atol=1e-4, rtol=1e-4) -def test_sparse_kv_b_lora_keeps_differentiable_factors_in_distributed_execution(monkeypatch): +def _assert_sparse_kv_b_adapter_weight_policy(monkeypatch): config = _tiny_config(index_topk=4, max_position_embeddings=8) attention = Glm5Attention(config, layer_idx=0) inject_lora_into_model(attention, r=4, lora_alpha=8, target_modules=["kv_b_proj"]) @@ -899,8 +765,11 @@ def test_sparse_kv_b_lora_keeps_differentiable_factors_in_distributed_execution( assert torch.count_nonzero(attention.kv_b_proj.lora_A.grad) assert torch.count_nonzero(attention.kv_b_proj.lora_B.grad) + _assert_sparse_kv_b_block_fp8_qlora_gradients(monkeypatch) + _assert_sparse_kv_b_block_fp8_compute_dtype(monkeypatch) -def test_sparse_kv_b_block_fp8_qlora_dequantizes_base_and_keeps_factor_gradients(monkeypatch): + +def _assert_sparse_kv_b_block_fp8_qlora_gradients(monkeypatch): config = _tiny_config(index_topk=4, max_position_embeddings=8) attention = Glm5Attention(config, layer_idx=0) projection = BlockFP8QLoRALinear( @@ -931,7 +800,7 @@ def test_sparse_kv_b_block_fp8_qlora_dequantizes_base_and_keeps_factor_gradients assert projection.lora_B.grad is not None -def test_sparse_kv_b_block_fp8_qlora_absorb_weights_follow_bf16_compute_dtype(monkeypatch): +def _assert_sparse_kv_b_block_fp8_compute_dtype(monkeypatch): config = _tiny_config(index_topk=4, max_position_embeddings=8) attention = Glm5Attention(config, layer_idx=0) projection = BlockFP8QLoRALinear( @@ -976,23 +845,7 @@ def test_sparse_kv_b_block_fp8_qlora_absorb_weights_follow_bf16_compute_dtype(mo @torch.no_grad() -def test_sparse_mla_dispatch_falls_back_to_torch_on_cpu(): - """`auto` backend on CPU must use the torch reference (tilelang is - CUDA-only). Output equals the torch ref called directly.""" - torch.manual_seed(0) - B, S, H = 1, 4, 2 - kv_lora, qk_rope = 16, 4 - D = kv_lora + qk_rope - q = torch.randn(B, S, H, D) - kv = torch.randn(B, S, D) - indices = torch.arange(S).view(1, 1, S).expand(B, S, S).contiguous() - - auto = sparse_mla_dispatch(q, kv, indices, scaling=D**-0.5, kv_lora_rank=kv_lora, backend="auto") - explicit = sparse_mla_torch_reference(q, kv, indices, D**-0.5, kv_lora) - torch.testing.assert_close(auto, explicit) - - -def test_sparse_mla_dispatch_rejects_unknown_backend(): +def _assert_sparse_mla_dispatch_rejects_unknown_backend(): q = torch.randn(1, 1, 1, 4) kv = torch.randn(1, 1, 4) indices = torch.zeros(1, 1, 1, dtype=torch.long) @@ -1000,7 +853,7 @@ def test_sparse_mla_dispatch_rejects_unknown_backend(): sparse_mla_dispatch(q, kv, indices, scaling=1.0, kv_lora_rank=2, backend="bogus") -def test_glm5_checkpoint_handler_skips_layers_beyond_configured(): +def _assert_glm5_checkpoint_filter_policy(): """Real GLM-5.1 has 78 dense layers + 1 MTP layer (index 78). Smoke runs further reduce `num_hidden_layers`. Either way, the handler must drop layer keys at indices `>= num_hidden_layers` so partial-load works.""" @@ -1015,12 +868,23 @@ def test_glm5_checkpoint_handler_skips_layers_beyond_configured(): assert handler._normalize_key("model.layers.78.embed_tokens.weight") is None # Non-layer keys are routed to the parent's normalizer untouched. assert handler._normalize_key("model.embed_tokens.weight") == "model.embed_tokens.weight" + skip = handler.get_skip_key_fn() + assert skip is not None + assert skip("model.layers.0.self_attn.q_a_proj.weight") is False + assert skip("model.layers.4.self_attn.q_a_proj.weight") is True + assert skip("model.layers.78.embed_tokens.weight") is True + assert skip("model.embed_tokens.weight") is False + _assert_glm5_checkpoint_ep_filter() -def test_glm5_default_lora_targets_cover_mla_and_moe(): + +def _assert_glm5_adapter_sparse_kv_and_moe_dispatch_policy(monkeypatch): """`xorl_glm5` model_type maps to GLM's MLA + MoE LoRA targets. Regressing this would silently fall back to llama-style `q/k/v/o_proj` and inject zero adapters.""" + with monkeypatch.context() as case_patch: + _assert_sparse_kv_b_adapter_weight_policy(case_patch) + config = _tiny_config() model = Glm5ForCausalLM(config) @@ -1041,8 +905,11 @@ def test_glm5_default_lora_targets_cover_mla_and_moe(): assert not isinstance(attn.indexer.wk, LoraLinear) assert not isinstance(attn.indexer.weights_proj, LoraLinear) + _assert_glm5_moe_uses_ep_dispatch(monkeypatch) + _assert_sparse_path_uses_lora_delta() -def test_glm5_moe_eager_uses_ep_dispatch_path_when_ep_enabled(monkeypatch): + +def _assert_glm5_moe_uses_ep_dispatch(monkeypatch): config = _tiny_config() block = Glm5MoEBlock(config) tokens = 6 @@ -1061,20 +928,7 @@ def test_glm5_moe_eager_uses_ep_dispatch_path_when_ep_enabled(monkeypatch): expert_forward.assert_called_once() -def test_glm5_checkpoint_handler_skip_key_fn_short_circuits_disk_reads(): - """`get_skip_key_fn` is the loader's early-skip hook — keys it returns - True for never get their tensor data read from disk. Must catch the - same MTP / partial-load case `_normalize_key` handles.""" - handler = Glm5CheckpointHandler(num_experts=8, num_hidden_layers=4) - skip = handler.get_skip_key_fn() - assert skip is not None - assert skip("model.layers.0.self_attn.q_a_proj.weight") is False - assert skip("model.layers.4.self_attn.q_a_proj.weight") is True # MTP layer - assert skip("model.layers.78.embed_tokens.weight") is True - assert skip("model.embed_tokens.weight") is False - - -def test_glm5_checkpoint_handler_ep_filter_skips_complete_fp8_expert_pairs(): +def _assert_glm5_checkpoint_ep_filter(): handler = Glm5CheckpointHandler(num_experts=8, ep_rank=1, ep_size=2) skip = handler.get_skip_key_fn() @@ -1086,7 +940,7 @@ def test_glm5_checkpoint_handler_ep_filter_skips_complete_fp8_expert_pairs(): @torch.no_grad() -def test_sparse_path_picks_up_lora_delta_on_kv_b_proj(): +def _assert_sparse_path_uses_lora_delta(): """`Glm5Attention.forward_sparse` accesses `kv_b_proj.weight` directly via the absorb form. When kv_b_proj is wrapped with LoraLinear, the sparse path must still reflect the LoRA delta — otherwise sparse + @@ -1194,7 +1048,7 @@ def test_glm5_logits_match_hf_reference_at_full_topk(): @torch.no_grad() -def test_glm5_for_causal_lm_forward_smoke(): +def _assert_glm5_forward_and_recompute_policy(): """End-to-end forward on a tiny model: shapes, MoE/dense layer split, and that the indexer is reachable through `model.model.layers[N].self_attn.indexer`.""" torch.manual_seed(0) @@ -1216,8 +1070,11 @@ def test_glm5_for_causal_lm_forward_smoke(): assert out.last_hidden_state.shape == (B, S, config.hidden_size) + with torch.enable_grad(): + _assert_glm5_recompute_before_dispatch() + -def test_glm5_recompute_before_dispatch_avoids_outer_layer_checkpoint(): +def _assert_glm5_recompute_before_dispatch(): torch.manual_seed(0) config = _tiny_config(num_hidden_layers=1, first_k_dense_replace=0) model = Glm5Model(config).train() diff --git a/tests/models/test_gradient_checkpointing.py b/tests/models/test_gradient_checkpointing.py deleted file mode 100644 index c9debfe4..00000000 --- a/tests/models/test_gradient_checkpointing.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Tests for `GradientCheckpointingLayer` + `gradient_checkpointing_enable` contract.""" - -from unittest.mock import MagicMock - -import pytest -import torch - -from xorl.models.base import XorlPreTrainedModel -from xorl.models.module_utils import ( - DEFAULT_GRADIENT_CHECKPOINTING_METHOD, - GradientCheckpointingLayer, - MoEGradientCheckpointingLayer, -) - - -pytestmark = [pytest.mark.cpu] - - -class _IdentityCheckpointLayer(GradientCheckpointingLayer): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return x - - -class _StubMoECheckpointLayer(MoEGradientCheckpointingLayer): - """MoE stub for attribute-level assertions — forward is not exercised.""" - - -@pytest.fixture -def model() -> XorlPreTrainedModel: - m = XorlPreTrainedModel(config=None) - m.layer = _IdentityCheckpointLayer() - return m - - -@pytest.mark.parametrize( - "layer_cls", - [GradientCheckpointingLayer, MoEGradientCheckpointingLayer], -) -def test_class_default_method_is_the_recompute_default(layer_cls): - assert layer_cls._gradient_checkpointing_method == DEFAULT_GRADIENT_CHECKPOINTING_METHOD - - -@pytest.mark.parametrize( - "layer_cls", - [_IdentityCheckpointLayer, _StubMoECheckpointLayer], -) -def test_gradient_checkpointing_enable_default(layer_cls): - model = XorlPreTrainedModel(config=None) - model.layer = layer_cls() - model.gradient_checkpointing_enable() - - assert model.layer.gradient_checkpointing is True - assert model.layer._gradient_checkpointing_method == DEFAULT_GRADIENT_CHECKPOINTING_METHOD - - -@pytest.mark.parametrize( - "method", - ["recompute_full_layer", "recompute_before_dispatch", "no_recompute"], -) -def test_enable_propagates_method_kwarg_to_every_checkpointed_layer(method): - model = XorlPreTrainedModel(config=None) - model.layer = _IdentityCheckpointLayer() - model.moe_layer = _StubMoECheckpointLayer() - - model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"gradient_checkpointing_method": method}) - - assert model.layer._gradient_checkpointing_method == method - assert model.moe_layer._gradient_checkpointing_method == method - - -@pytest.mark.parametrize( - "training, flag_enabled, expect_checkpoint", - [ - (True, True, True), - (True, False, False), - (False, True, False), - (False, False, False), - ], -) -def test_outer_gate_fires_iff_training_and_flag(model, training, flag_enabled, expect_checkpoint): - model.gradient_checkpointing_enable() - model.train(training) - model.layer.gradient_checkpointing = flag_enabled - - spy = MagicMock(side_effect=lambda fn, *a, **kw: fn(*a, **kw)) - model.layer._gradient_checkpointing_func = spy - - x = torch.zeros(2, 3) - out = model.layer(x) - - assert torch.equal(out, x) - assert spy.called is expect_checkpoint diff --git a/tests/models/test_kimi_tokenizer.py b/tests/models/test_kimi_tokenizer.py deleted file mode 100644 index cea1d48f..00000000 --- a/tests/models/test_kimi_tokenizer.py +++ /dev/null @@ -1,82 +0,0 @@ -import json - -import pytest -from tiktoken.load import dump_tiktoken_bpe - -from xorl.models import auto as auto_module -from xorl.models.auto import build_processor, build_tokenizer -from xorl.models.transformers.deepseek_v3.tokenization_kimi import TikTokenTokenizer - - -pytestmark = [pytest.mark.cpu] - - -def _write_tiktoken_fixture(tmp_path): - tokenizer_dir = tmp_path / "kimi-tokenizer" - tokenizer_dir.mkdir() - dump_tiktoken_bpe({bytes([i]): i for i in range(256)}, str(tokenizer_dir / "tiktoken.model")) - (tokenizer_dir / "tokenizer_config.json").write_text( - json.dumps( - { - "tokenizer_class": "TikTokenTokenizer", - "auto_map": {"AutoTokenizer": ["tokenization_kimi.TikTokenTokenizer", None]}, - "bos_token": "[BOS]", - "eos_token": "[EOS]", - "unk_token": "[UNK]", - "pad_token": "[PAD]", - "additional_special_tokens": ["<|im_end|>"], - "added_tokens_decoder": { - "256": {"content": "[BOS]", "special": True}, - "257": {"content": "[EOS]", "special": True}, - "258": {"content": "[UNK]", "special": True}, - "259": {"content": "[PAD]", "special": True}, - "260": {"content": "<|im_end|>", "special": True}, - }, - } - ) - ) - return tokenizer_dir - - -def test_build_tokenizer_loads_local_kimi_tiktoken_without_remote_code(tmp_path): - tokenizer_dir = _write_tiktoken_fixture(tmp_path) - - tokenizer = build_tokenizer(str(tokenizer_dir)) - - assert isinstance(tokenizer, TikTokenTokenizer) - assert tokenizer.bos_token_id == 256 - assert tokenizer.eos_token_id == 257 - assert tokenizer.pad_token_id == 259 - assert tokenizer.decode(tokenizer.encode("hello")) == "hello" - - -def test_build_tokenizer_fallback_does_not_enable_remote_code(monkeypatch, tmp_path): - calls = {} - - def fake_from_pretrained(path, **kwargs): - calls["path"] = path - calls["kwargs"] = kwargs - return object() - - monkeypatch.setattr(auto_module.AutoTokenizer, "from_pretrained", fake_from_pretrained) - - build_tokenizer(str(tmp_path / "not-kimi")) - - assert "trust_remote_code" not in calls["kwargs"] - assert calls["kwargs"]["padding_side"] == "right" - - -def test_build_processor_does_not_enable_remote_code(monkeypatch): - calls = {} - - def fake_from_pretrained(path, **kwargs): - calls["path"] = path - calls["kwargs"] = kwargs - return object() - - monkeypatch.setattr(auto_module.AutoProcessor, "from_pretrained", fake_from_pretrained) - - build_processor("processor-path") - - assert "trust_remote_code" not in calls["kwargs"] - assert calls["kwargs"]["padding_side"] == "right" diff --git a/tests/models/test_lora_merge_fp32_cast_once.py b/tests/models/test_lora_merge_fp32_cast_once.py deleted file mode 100644 index 49968d8c..00000000 --- a/tests/models/test_lora_merge_fp32_cast_once.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Tests for the fp32 cast-once merge variant on both LoraLinear and MoEExpertsLoRA. - -Invariants: - - With zero LoRA (B=0), merge must be bit-exact: W_merged == W (no change). - - After merge, ``merged_weight`` equals the fp32 reference - ``(W.to(fp32) + B@A*s).to(W.dtype)`` bit-for-bit (by construction). - - Merged weight is >= as faithful as the naive ``W + Δ.to(W.dtype)`` variant — - i.e., the fp32-sum-then-cast distance to the true fp32 merged value is ≤ - the naive-merge distance. -""" - -import pytest -import torch - - -pytestmark = [pytest.mark.gpu] - - -def _naive_merge(weight, delta): - """Old behavior for comparison: round Δ per-element, then add.""" - return weight + delta.to(weight.dtype) - - -def _fp32_merge(weight, delta): - """New behavior: add in fp32, cast once.""" - return (weight.to(torch.float32) + delta).to(weight.dtype) - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) -def test_lora_linear_merge_zero_b_is_bitexact(dtype): - """merge_weights() on a fresh LoraLinear (B=0) must leave weight untouched.""" - from xorl.lora import LoraLinear - - torch.manual_seed(0) - layer = LoraLinear(128, 64, r=8, lora_alpha=16, device="cuda", dtype=dtype) - layer.weight.data.normal_(std=0.05) - before = layer.weight.detach().clone() - layer.merge_weights() - assert torch.equal(layer.weight, before), "zero-LoRA merge must be bit-exact" - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -def test_lora_linear_merge_matches_fp32_reference(dtype): - """The merged weight must equal ``(W.to(fp32) + B@A*s).to(W.dtype)`` exactly.""" - from xorl.lora import LoraLinear - - torch.manual_seed(0) - layer = LoraLinear(128, 64, r=8, lora_alpha=16, device="cuda", dtype=dtype) - layer.weight.data.normal_(std=0.05) - # non-zero LoRA - layer.lora_B.data.normal_(std=0.02) - w_before = layer.weight.detach().clone() - delta_fp32 = (layer.lora_B @ layer.lora_A) * layer.scaling - expected = _fp32_merge(w_before, delta_fp32) - - layer.merge_weights() - assert torch.equal(layer.weight, expected), "merge_weights must match fp32-cast-once reference" - - -def test_lora_linear_merge_strictly_ge_naive_precision(): - """For any B, fp32-cast-once ≤ naive in distance to true fp32 merged value.""" - from xorl.lora import LoraLinear - - torch.manual_seed(42) - layer = LoraLinear(128, 64, r=8, lora_alpha=16, device="cuda", dtype=torch.bfloat16) - layer.weight.data.normal_(std=0.05) - layer.lora_B.data.normal_(std=0.02) - - w = layer.weight.detach().clone() - delta_fp32 = (layer.lora_B @ layer.lora_A) * layer.scaling - true_fp32 = w.to(torch.float32) + delta_fp32 # reference in fp32 - naive = _naive_merge(w, delta_fp32).to(torch.float32) # rounds Δ then adds - fp32_cast_once = _fp32_merge(w, delta_fp32).to(torch.float32) # fp32 sum, cast once - - naive_err = (naive - true_fp32).abs().max().item() - fp32_err = (fp32_cast_once - true_fp32).abs().max().item() - assert fp32_err <= naive_err + 1e-9, ( - f"fp32-cast-once should be ≤ naive precision: fp32={fp32_err:.3e} naive={naive_err:.3e}" - ) - - -def _tiny_moe_experts_with_lora(dtype): - from xorl.lora import MoEExpertsLoRA, MoELoRAConfig - - cfg = MoELoRAConfig(r=8, lora_alpha=16, target_modules=["gate_proj", "up_proj", "down_proj"]) - e = ( - MoEExpertsLoRA( - num_experts=4, - hidden_dim=32, - intermediate_size=24, - hidden_act="silu", - moe_implementation="eager", - lora_config=cfg, - ) - .to(dtype) - .cuda() - ) - e.gate_up_proj.data.normal_(std=0.05) - e.down_proj.data.normal_(std=0.05) - return e - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -def test_moe_merge_zero_b_is_bitexact(dtype): - e = _tiny_moe_experts_with_lora(dtype) - gu_before = e.gate_up_proj.detach().clone() - dn_before = e.down_proj.detach().clone() - e.merge_weights() - assert torch.equal(e.gate_up_proj, gu_before), "zero-LoRA MoE merge must not change gate_up_proj" - assert torch.equal(e.down_proj, dn_before), "zero-LoRA MoE merge must not change down_proj" - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -def test_moe_merge_matches_fp32_reference(dtype): - e = _tiny_moe_experts_with_lora(dtype) - # perturb all lora_B - torch.manual_seed(7) - for name in ("gate_proj", "up_proj", "down_proj"): - getattr(e, f"{name}_lora_B").data.normal_(std=0.02) - - # expected = fp32 sum then cast, computed from SNAPSHOTS of current state - gu_before = e.gate_up_proj.detach().clone() - dn_before = e.down_proj.detach().clone() - expected_updates = {} - for proj in ("gate_proj", "up_proj", "down_proj"): - delta = e._compute_proj_delta(proj) # fp32, [E, in, out] - expected_updates[proj] = delta - - # gate/up land in the fused gate_up_proj via views, each shape (E, H, I) - I = e.intermediate_size - gate_expected = gu_before.clone() - gate_expected[..., :I] = _fp32_merge(gu_before[..., :I], expected_updates["gate_proj"]) - gate_expected[..., I:] = _fp32_merge(gu_before[..., I:], expected_updates["up_proj"]) - - down_expected = _fp32_merge(dn_before, expected_updates["down_proj"]) - - e.merge_weights() - - assert torch.equal(e.gate_up_proj, gate_expected), "MoE gate+up merge must match fp32 reference" - assert torch.equal(e.down_proj, down_expected), "MoE down merge must match fp32 reference" diff --git a/tests/models/test_lora_merged_forward.py b/tests/models/test_lora_merged_forward.py index 38b7eb14..97262d45 100644 --- a/tests/models/test_lora_merged_forward.py +++ b/tests/models/test_lora_merged_forward.py @@ -36,8 +36,81 @@ def _gkn_factors(shared_a=False, shared_b=False, dtype=torch.bfloat16, seed=0): return W, A, B +def _cast_once_merge(weight, delta): + return (weight.to(torch.float32) + delta).to(weight.dtype) + + +def _assert_permanent_merge_matches_cast_once_contract(): + for dtype in (torch.bfloat16, torch.float16, torch.float32): + torch.manual_seed(0) + layer = LoraLinear(128, 64, r=8, lora_alpha=16, dtype=dtype) + layer.weight.data.normal_(std=0.05) + before = layer.weight.detach().clone() + layer.merge_weights() + assert torch.equal(layer.weight, before) + + if dtype is torch.float32: + continue + torch.manual_seed(0) + layer = LoraLinear(128, 64, r=8, lora_alpha=16, dtype=dtype) + layer.weight.data.normal_(std=0.05) + layer.lora_B.data.normal_(std=0.02) + before = layer.weight.detach().clone() + expected = _cast_once_merge(before, (layer.lora_B @ layer.lora_A) * layer.scaling) + layer.merge_weights() + assert torch.equal(layer.weight, expected) + + config = MoELoRAConfig(r=8, lora_alpha=16, target_modules=["gate_proj", "up_proj", "down_proj"]) + for dtype in (torch.bfloat16, torch.float16): + experts = MoEExpertsLoRA( + num_experts=4, + hidden_dim=32, + intermediate_size=24, + hidden_act="silu", + moe_implementation="eager", + lora_config=config, + ).to(dtype) + experts.gate_up_proj.data.normal_(std=0.05) + experts.down_proj.data.normal_(std=0.05) + gate_up_before = experts.gate_up_proj.detach().clone() + down_before = experts.down_proj.detach().clone() + experts.merge_weights() + assert torch.equal(experts.gate_up_proj, gate_up_before) + assert torch.equal(experts.down_proj, down_before) + + experts = MoEExpertsLoRA( + num_experts=4, + hidden_dim=32, + intermediate_size=24, + hidden_act="silu", + moe_implementation="eager", + lora_config=config, + ).to(dtype) + experts.gate_up_proj.data.normal_(std=0.05) + experts.down_proj.data.normal_(std=0.05) + torch.manual_seed(7) + for projection in ("gate_proj", "up_proj", "down_proj"): + getattr(experts, f"{projection}_lora_B").data.normal_(std=0.02) + + gate_up_before = experts.gate_up_proj.detach().clone() + down_before = experts.down_proj.detach().clone() + updates = {projection: experts._compute_proj_delta(projection) for projection in config.target_modules} + intermediate = experts.intermediate_size + gate_up_expected = gate_up_before.clone() + gate_up_expected[..., :intermediate] = _cast_once_merge( + gate_up_before[..., :intermediate], updates["gate_proj"] + ) + gate_up_expected[..., intermediate:] = _cast_once_merge(gate_up_before[..., intermediate:], updates["up_proj"]) + down_expected = _cast_once_merge(down_before, updates["down_proj"]) + + experts.merge_weights() + + assert torch.equal(experts.gate_up_proj, gate_up_expected) + assert torch.equal(experts.down_proj, down_expected) + + class TestCanonicalFold: - def test_pinned_order(self): + def test_canonical_fold_forward_and_cache_policy(self): W, A, B = _gkn_factors() delta = torch.bmm(A.float().expand(E, -1, -1), B.float().expand(E, -1, -1)) * SCALING want = (W.float() + delta).to(W.dtype) @@ -45,7 +118,16 @@ def test_pinned_order(self): assert torch.equal(got, want) assert got.dtype == W.dtype - def test_shared_factors_expand(self): + self._assert_shared_factors_expand() + self._assert_fold_commutes_with_expert_sharding() + self._assert_linear_orientation() + self._assert_zero_delta_is_value_identical() + _assert_permanent_merge_matches_cast_once_contract() + TestFoldedWeightAutograd()._assert_straight_through_gradient_policy() + TestLoraLinearMerged()._assert_merged_forward_selection_policy() + TestMoEExpertsLoRAMerged()._assert_merged_weight_and_cache_policy() + + def _assert_shared_factors_expand(self): W, A, B = _gkn_factors(shared_a=True) got = canonical_lora_fold_gkn(W, A, B, SCALING) per_expert = torch.stack( @@ -53,7 +135,7 @@ def test_shared_factors_expand(self): ) assert torch.allclose(got.float(), per_expert.float(), atol=0, rtol=0) or torch.equal(got, per_expert) - def test_fold_commutes_with_expert_sharding(self): + def _assert_fold_commutes_with_expert_sharding(self): # EP invariance: fold of a local expert slice == slice of the full fold. W, A, B = _gkn_factors() full = canonical_lora_fold_gkn(W, A, B, SCALING) @@ -61,7 +143,7 @@ def test_fold_commutes_with_expert_sharding(self): shard = canonical_lora_fold_gkn(W[sl], A[sl], B[sl], SCALING) assert torch.equal(shard, full[sl]) - def test_linear_orientation(self): + def _assert_linear_orientation(self): g = torch.Generator().manual_seed(1) W = torch.randn(I, H, generator=g).to(torch.bfloat16) A = torch.randn(R, H, generator=g).to(torch.float32) * 0.02 @@ -69,7 +151,7 @@ def test_linear_orientation(self): want = (W.float() + (B @ A) * SCALING).to(W.dtype) assert torch.equal(canonical_lora_fold_linear(W, A, B, SCALING), want) - def test_zero_delta_is_value_identical(self): + def _assert_zero_delta_is_value_identical(self): # B == 0: folded values equal the base everywhere (only -0.0 -> +0.0 # bit flips are permitted; both engines fold identically so the # contract holds). @@ -79,7 +161,7 @@ def test_zero_delta_is_value_identical(self): class TestFoldedWeightAutograd: - def test_factor_grad_dtype_honors_fsdp_metadata(self): + def _assert_factor_grad_dtype_honors_fsdp_metadata(self): factor_A = torch.ones(4, 8, dtype=torch.bfloat16, requires_grad=True) factor_B = torch.ones(16, 4, dtype=torch.bfloat16, requires_grad=True) assert _factor_grad_dtype(factor_A[:2]) == torch.bfloat16 @@ -108,21 +190,26 @@ def _reference_grads(self, W, A, B, scaling, grad_w, shared_a=False, shared_b=Fa folded.backward(grad_w) return A_ref.grad, B_ref.grad - @pytest.mark.parametrize("shared_a,shared_b", [(False, False), (True, False), (False, True)]) - def test_gkn_straight_through_matches_autograd(self, shared_a, shared_b): - W, A, B = _gkn_factors(shared_a=shared_a, shared_b=shared_b, seed=2) - A = A.requires_grad_(True) - B = B.requires_grad_(True) - folded = canonical_lora_fold_gkn(W, A.detach(), B.detach(), SCALING) - out = FoldedLoraWeightGKN.apply(folded, A, B, SCALING) - assert torch.equal(out, folded) - grad_w = torch.randn_like(folded.float()).to(folded.dtype) - out.backward(grad_w) - gA_ref, gB_ref = self._reference_grads(W, A, B, SCALING, grad_w, shared_a, shared_b) - assert torch.allclose(A.grad.float(), gA_ref.float(), rtol=1e-5, atol=1e-8) - assert torch.allclose(B.grad.float(), gB_ref.float(), rtol=1e-5, atol=1e-8) + def _assert_straight_through_gradient_policy(self): + self._assert_factor_grad_dtype_honors_fsdp_metadata() + + for shared_a, shared_b in ((False, False), (True, False), (False, True)): + W, A, B = _gkn_factors(shared_a=shared_a, shared_b=shared_b, seed=2) + A = A.requires_grad_(True) + B = B.requires_grad_(True) + folded = canonical_lora_fold_gkn(W, A.detach(), B.detach(), SCALING) + out = FoldedLoraWeightGKN.apply(folded, A, B, SCALING) + assert torch.equal(out, folded) + grad_w = torch.randn_like(folded.float()).to(folded.dtype) + out.backward(grad_w) + gA_ref, gB_ref = self._reference_grads(W, A, B, SCALING, grad_w, shared_a, shared_b) + assert torch.allclose(A.grad.float(), gA_ref.float(), rtol=1e-5, atol=1e-8) + assert torch.allclose(B.grad.float(), gB_ref.float(), rtol=1e-5, atol=1e-8) - def test_gate_up_fused_straight_through(self): + self._assert_gate_up_fused_straight_through() + self._assert_linear_straight_through_matches_autograd() + + def _assert_gate_up_fused_straight_through(self): g = torch.Generator().manual_seed(3) W = torch.randn(E, H, 2 * I, generator=g).to(torch.bfloat16) gA = (torch.randn(1, H, R, generator=g) * 0.02).to(torch.bfloat16).requires_grad_(True) @@ -147,7 +234,7 @@ def test_gate_up_fused_straight_through(self): B.grad = None break # gate checked exactly; up follows by symmetry - def test_linear_straight_through_matches_autograd(self): + def _assert_linear_straight_through_matches_autograd(self): g = torch.Generator().manual_seed(4) W = torch.randn(I, H, generator=g).to(torch.bfloat16) A = (torch.randn(R, H, generator=g) * 0.02).requires_grad_(True) @@ -171,7 +258,7 @@ def _layer(self): torch.nn.init.normal_(layer.lora_B, std=0.02) # nonzero delta return layer - def test_merged_forward_bits(self): + def _assert_merged_forward_selection_policy(self): layer = self._layer() layer.exact_merged_forward = True x = torch.randn(6, H) @@ -182,7 +269,12 @@ def test_merged_forward_bits(self): got = layer(x) assert torch.equal(got, want) - def test_flag_off_keeps_legacy_path(self): + self._assert_flag_off_keeps_legacy_path() + self._assert_exact_and_ordinary_selection_is_isolated() + self._assert_merged_grads_close_to_unmerged() + self._assert_cache_invalidation_on_step_and_runtime_config() + + def _assert_flag_off_keeps_legacy_path(self): layer = self._layer() x = torch.randn(6, H) base = torch.nn.functional.linear(x, layer.weight, None) @@ -190,7 +282,7 @@ def test_flag_off_keeps_legacy_path(self): want = base + (lora * layer.scaling) assert torch.allclose(layer(x), want, rtol=0, atol=0) - def test_exact_and_ordinary_modules_do_not_share_selection_state(self): + def _assert_exact_and_ordinary_selection_is_isolated(self): exact = self._layer() ordinary = self._layer() ordinary.load_state_dict(exact.state_dict()) @@ -208,14 +300,13 @@ def test_exact_and_ordinary_modules_do_not_share_selection_state(self): ) assert torch.equal(ordinary_out, expected_ordinary) - def test_merged_grads_close_to_unmerged(self): + def _assert_merged_grads_close_to_unmerged(self): layer = self._layer() x = torch.randn(6, H) grads = {} for flag in ("0", "1"): layer.exact_merged_forward = flag == "1" layer.lora_A.grad = layer.lora_B.grad = None - layer.invalidate_merged_weight_cache() layer(x).square().mean().backward() grads[flag] = (layer.lora_A.grad.clone(), layer.lora_B.grad.clone()) # fwd/bwd decoupling: merged-lane grads track the unmerged autograd at @@ -223,7 +314,7 @@ def test_merged_grads_close_to_unmerged(self): for a, b in zip(grads["0"], grads["1"]): assert torch.allclose(a, b, rtol=5e-2, atol=1e-5) - def test_cache_invalidation_on_step_and_runtime_config(self): + def _assert_cache_invalidation_on_step_and_runtime_config(self): layer = self._layer() layer.exact_merged_forward = True w1 = layer._merged_weight() @@ -263,7 +354,7 @@ def _module(self, hybrid=True): getattr(mod, f"{proj}_lora_B").normal_(std=0.02) return mod - def test_merged_weights_match_canonical_fold(self): + def _assert_merged_weight_and_cache_policy(self): mod = self._module() mod.exact_merged_forward = True gate_up_f, down_f = mod._merged_weights() @@ -278,7 +369,10 @@ def test_merged_weights_match_canonical_fold(self): assert torch.equal(mod.canonical_merged_proj_weight("gate_proj"), gate_up_f[..., :I]) assert torch.equal(mod.canonical_merged_proj_weight("down_proj"), down_f) - def test_cache_keyed_on_versions(self): + self._assert_cache_is_keyed_on_parameter_versions() + self._assert_fused_experts_admission_policy() + + def _assert_cache_is_keyed_on_parameter_versions(self): mod = self._module() mod.exact_merged_forward = True g1, d1 = mod._merged_weights() @@ -289,13 +383,15 @@ def test_cache_keyed_on_versions(self): g3, _ = mod._merged_weights() assert g3 is not g1 and not torch.equal(g1, g3) - def test_auto_supported_requires_exact_model_program(self): + def _assert_fused_experts_admission_policy(self): mod = self._module() assert not mod.sglang_fused_experts_auto_supported() mod.exact_merged_forward = True assert mod.sglang_fused_experts_auto_supported() - def test_fused_flag_without_merged_flag_raises(self): + self._assert_fused_flag_without_merged_flag_raises() + + def _assert_fused_flag_without_merged_flag_raises(self): mod = self._module() with pytest.raises(NotImplementedError, match="merged"): mod.sglang_fused_experts_forward( @@ -304,7 +400,7 @@ def test_fused_flag_without_merged_flag_raises(self): torch.randint(0, E, (3, 2)), ) - def test_native_ep_keyword_routes_to_masked_lora_partial(self, monkeypatch): + def test_native_ep_merged_lora_policy(self, monkeypatch): mod = self._module() mod.exact_merged_forward = True hidden = torch.randn(3, H).to(torch.bfloat16) @@ -322,7 +418,9 @@ def masked_partial(got_hidden, got_routing, got_ids): assert got is expected assert calls == [(hidden, routing, local_ids)] - def test_native_ep_no_grad_uses_canonical_fold_and_filter(self, monkeypatch): + self._assert_native_ep_no_grad_uses_canonical_fold_and_filter(monkeypatch) + + def _assert_native_ep_no_grad_uses_canonical_fold_and_filter(self, monkeypatch): mod = self._module() mod.exact_merged_forward = True hidden = torch.randn(3, H).to(torch.bfloat16) @@ -383,13 +481,15 @@ def _model(self): model.q_proj = LoraLinear(H, H, r=R, lora_alpha=R, dtype=torch.bfloat16) return model - def test_wrap_raises_without_merged_flag(self): + def test_trunk_wrap_composition_policy(self, monkeypatch): from xorl.ops.batch_invariant_ops import wrap_trunk_linears_batch_invariant # noqa: PLC0415 with pytest.raises(NotImplementedError, match="canonical merged-LoRA"): wrap_trunk_linears_batch_invariant(self._model()) - def test_wrap_composes_with_merged_flag(self, monkeypatch): + self._assert_wrap_composes_with_merged_flag(monkeypatch) + + def _assert_wrap_composes_with_merged_flag(self, monkeypatch): from xorl.ops.batch_invariant_ops import ( # noqa: PLC0415 set_trunk_linear_contract, wrap_trunk_linears_batch_invariant, diff --git a/tests/models/test_lora_moe_attention_targets.py b/tests/models/test_lora_moe_attention_targets.py deleted file mode 100644 index 0bb54e82..00000000 --- a/tests/models/test_lora_moe_attention_targets.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Regression test for inject_lora_into_model_with_moe attention partition. - -Locks in the fix for the case where the function silently dropped DeepSeek-V3 / -Kimi MLA attention projections (q_a_proj, q_b_proj, kv_a_proj_with_mqa, -kv_b_proj) because the attention vs. expert split was hardcoded to a -Llama/Qwen-shaped allowlist (q_proj/k_proj/v_proj/o_proj/lm_head). -""" - -import pytest -import torch.nn as nn - -from xorl.lora import LoraLinear, inject_lora_into_model_with_moe - - -pytestmark = [pytest.mark.cpu] - - -class _StubConfig: - def __init__(self, model_type: str): - self.model_type = model_type - self.num_experts = 0 - - -class _MLALikeBlock(nn.Module): - """Single attention block with DeepSeek-V3 / Kimi MLA projection names.""" - - def __init__(self, hidden_size: int = 32, q_lora_rank: int = 16, kv_lora_rank: int = 16): - super().__init__() - self.q_a_proj = nn.Linear(hidden_size, q_lora_rank, bias=False) - self.q_b_proj = nn.Linear(q_lora_rank, hidden_size, bias=False) - self.kv_a_proj_with_mqa = nn.Linear(hidden_size, kv_lora_rank, bias=False) - self.kv_b_proj = nn.Linear(kv_lora_rank, hidden_size, bias=False) - self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False) - - -class _DeepSeekLikeModel(nn.Module): - def __init__(self): - super().__init__() - self.config = _StubConfig("deepseek_v3") - self.self_attn = _MLALikeBlock() - - -def test_default_targets_cover_all_mla_projections_for_deepseek_v3(): - model = _DeepSeekLikeModel() - - inject_lora_into_model_with_moe(model, r=4, lora_alpha=8, target_modules=None) - - for proj in ("q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj", "o_proj"): - replaced = getattr(model.self_attn, proj) - assert isinstance(replaced, LoraLinear), ( - f"{proj} was not LoRA-replaced; the attention/expert split is dropping MLA projections again." - ) - - -def test_explicit_mla_targets_are_not_filtered_out(): - model = _DeepSeekLikeModel() - - inject_lora_into_model_with_moe( - model, - r=4, - lora_alpha=8, - target_modules=["q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"], - ) - - for proj in ("q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"): - replaced = getattr(model.self_attn, proj) - assert isinstance(replaced, LoraLinear), ( - f"{proj} was filtered out of attention_modules even though the caller passed it explicitly." - ) - assert not isinstance(model.self_attn.o_proj, LoraLinear) diff --git a/tests/models/test_lora_projection_audit.py b/tests/models/test_lora_projection_audit.py new file mode 100644 index 00000000..04df0267 --- /dev/null +++ b/tests/models/test_lora_projection_audit.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import yaml +from torch import nn +from torch.nn import functional as F + +from xorl.lora.modules.delta_linear import LoraDeltaLinear +from xorl.lora.utils import _get_default_target_modules, inject_lora_into_model +from xorl.models.layers.fused_projection_lora import project_fused_linear_with_lora +from xorl.models.transformers.glm4_moe.modeling_glm4_moe import Glm4MoeMLP +from xorl.models.transformers.llama3.modeling_llama3 import LlamaMLP +from xorl.models.transformers.olmo2.modeling_olmo2 import Olmo2MLP +from xorl.models.transformers.qwen2.modeling_qwen2 import Qwen2MLP +from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3MLP +from xorl.models.transformers.qwen3_5.modeling_qwen3_5 import Qwen3_5MLP +from xorl.models.transformers.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeMLP +from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import Qwen3MoeMLP +from xorl.server.weight_sync.handler import WeightSyncHandler + + +pytestmark = [pytest.mark.cpu] + + +class _AuditedFusedBlock(nn.Module): + _supports_fused_qkv_lora = True + _supports_fused_gate_up_lora = True + + def __init__(self) -> None: + super().__init__() + self.q_dim = 6 + self.kv_dim = 2 + self.intermediate_size = 5 + self.qkv_proj = nn.Linear(4, self.q_dim + 2 * self.kv_dim, bias=False) + self.gate_up_proj = nn.Linear(4, 2 * self.intermediate_size, bias=False) + + def qkv(self, inputs: torch.Tensor) -> torch.Tensor: + return project_fused_linear_with_lora( + self, + inputs, + base_name="qkv_proj", + projection_names=("q_proj", "k_proj", "v_proj"), + projection_sizes=(self.q_dim, self.kv_dim, self.kv_dim), + ) + + def gate_up(self, inputs: torch.Tensor) -> torch.Tensor: + return project_fused_linear_with_lora( + self, + inputs, + base_name="gate_up_proj", + projection_names=("gate_proj", "up_proj"), + projection_sizes=(self.intermediate_size, self.intermediate_size), + ) + + +class _AuditedModel(nn.Module): + def __init__(self, model_type: str = "qwen3") -> None: + super().__init__() + self.config = SimpleNamespace(model_type=model_type) + self.block = _AuditedFusedBlock() + + +def _seed_nonzero_adapters(model: nn.Module) -> None: + with torch.no_grad(): + for index, module in enumerate( + (child for child in model.modules() if isinstance(child, LoraDeltaLinear)), + start=1, + ): + module.lora_A.copy_(torch.arange(module.lora_A.numel()).reshape_as(module.lora_A) / (10 + index)) + module.lora_B.copy_(torch.arange(1, module.lora_B.numel() + 1).reshape_as(module.lora_B) / (20 + index)) + + +def test_split_targets_keep_fused_base_projections_and_independent_factors() -> None: + model = _AuditedModel() + inputs = torch.randn(3, 4) + qkv_base = model.block.qkv_proj(inputs) + gate_up_base = model.block.gate_up_proj(inputs) + + inject_lora_into_model( + model, + r=2, + lora_alpha=2, + target_modules=["q_proj", "k_proj", "v_proj", "gate_proj", "up_proj"], + ) + + assert isinstance(model.block.qkv_proj, nn.Linear) + assert isinstance(model.block.gate_up_proj, nn.Linear) + assert all( + isinstance(getattr(model.block, name), LoraDeltaLinear) + for name in ("q_proj", "k_proj", "v_proj", "gate_proj", "up_proj") + ) + assert torch.equal(model.block.qkv(inputs), qkv_base) + assert torch.equal(model.block.gate_up(inputs), gate_up_base) + + +@pytest.mark.parametrize( + "mlp_class", + [Qwen2MLP, Qwen3MLP, Qwen3MoeMLP, Qwen3_5MLP, Qwen3_5MoeMLP, LlamaMLP, Olmo2MLP, Glm4MoeMLP], +) +def test_audited_fused_mlp_implementations_retain_base_program(mlp_class: type[nn.Module]) -> None: + config = SimpleNamespace( + hidden_size=8, + intermediate_size=6, + hidden_act="gelu", + mlp_bias=False, + _activation_native=True, + ) + mlp = mlp_class(config) + inputs = torch.randn(2, 3, config.hidden_size) + expected = mlp(inputs) + inject_lora_into_model(mlp, r=2, lora_alpha=2, target_modules=["gate_proj", "up_proj", "down_proj"]) + + assert isinstance(mlp.gate_up_proj, nn.Linear) + assert isinstance(mlp.gate_proj, LoraDeltaLinear) + assert isinstance(mlp.up_proj, LoraDeltaLinear) + assert torch.equal(mlp(inputs), expected) + + +@pytest.mark.parametrize( + ("method", "base_name", "projection_names", "sizes"), + [ + ("qkv", "qkv_proj", ("q_proj", "k_proj", "v_proj"), (6, 2, 2)), + ("gate_up", "gate_up_proj", ("gate_proj", "up_proj"), (5, 5)), + ], +) +def test_fused_projection_dynamic_and_exact_merged_programs( + method: str, + base_name: str, + projection_names: tuple[str, ...], + sizes: tuple[int, ...], +) -> None: + model = _AuditedModel() + inject_lora_into_model(model, r=2, lora_alpha=2, target_modules=list(projection_names)) + _seed_nonzero_adapters(model) + inputs = torch.randn(2, 3, 4) + block = model.block + base = getattr(block, base_name) + + base_parts = list(base(inputs).split(sizes, dim=-1)) + for index, name in enumerate(projection_names): + base_parts[index] = base_parts[index] + getattr(block, name)(inputs) + expected_dynamic = torch.cat(base_parts, dim=-1) + assert torch.equal(getattr(block, method)(inputs), expected_dynamic) + + folded_parts = [] + for base_part, name in zip(base.weight.split(sizes, dim=0), projection_names, strict=True): + adapter = getattr(block, name) + adapter.exact_merged_forward = True + folded_parts.append(adapter.merged_weight_for_forward(base_part)) + expected_merged = F.linear(inputs, torch.cat(folded_parts, dim=0)) + assert torch.equal(getattr(block, method)(inputs), expected_merged) + + +def test_weight_sync_folds_logical_factors_into_fused_base_weights() -> None: + model = _AuditedModel() + inject_lora_into_model( + model, + r=2, + lora_alpha=2, + target_modules=["q_proj", "k_proj", "v_proj", "gate_proj", "up_proj"], + ) + _seed_nonzero_adapters(model) + for module in model.modules(): + if isinstance(module, LoraDeltaLinear): + module.exact_merged_forward = True + + class _FakeDTensor: + pass + + extracted = dict(WeightSyncHandler._extract_params_for_sync(model, "(root)", _FakeDTensor)) + for base_name, projection_names, sizes in ( + ("qkv_proj", ("q_proj", "k_proj", "v_proj"), (6, 2, 2)), + ("gate_up_proj", ("gate_proj", "up_proj"), (5, 5)), + ): + base = getattr(model.block, base_name) + expected = torch.cat( + [ + getattr(model.block, name)._merged_weight(base_part).to(torch.bfloat16) + for name, base_part in zip(projection_names, base.weight.split(sizes, dim=0), strict=True) + ], + dim=0, + ) + assert torch.equal(extracted[f"block.{base_name}.weight"], expected) + + +@pytest.mark.parametrize( + ("model_type", "expected"), + [ + ("xorl_llama3", ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]), + ("qwen2", ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]), + ("xorl_qwen3", ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]), + ("qwen3_moe", ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]), + ( + "xorl_qwen3_5_moe", + ["q_proj", "k_proj", "v_proj", "g_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + ), + ("olmo2", ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]), + ("glm4_moe", ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]), + ( + "deepseek_v3", + [ + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + ), + ( + "kimi_k25", + [ + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + ), + ( + "xorl_glm5", + [ + "q_a_proj", + "q_b_proj", + "kv_a_proj_with_mqa", + "kv_b_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + ), + ("deepseek_v4", ["wq_a", "wq_b", "wkv", "wo_a", "wo_b"]), + ("gpt_oss", ["q_proj", "k_proj", "v_proj", "o_proj"]), + ("minimax_m3", ["q_proj", "k_proj", "v_proj", "o_proj"]), + ("nemotron_h", ["q_proj", "k_proj", "v_proj", "o_proj"]), + ], +) +def test_audited_model_family_defaults(model_type: str, expected: list[str]) -> None: + model = nn.Module() + model.config = SimpleNamespace(model_type=model_type) + assert _get_default_target_modules(model) == expected + + +def test_unknown_model_family_requires_explicit_targets() -> None: + model = nn.Module() + model.config = SimpleNamespace(model_type="new_unreviewed_architecture") + with pytest.raises(ValueError, match="No audited default LoRA targets"): + _get_default_target_modules(model) + + +def test_unmatched_projection_target_fails_closed() -> None: + with pytest.raises(ValueError, match="missing_proj"): + inject_lora_into_model(_AuditedModel(), target_modules=["q_proj", "missing_proj"]) + + +def test_checked_in_plain_lora_configs_cover_audited_projection_sets() -> None: + root = Path(__file__).resolve().parents[2] + config_paths = sorted((root / "examples").glob("**/configs/lora/*.yaml")) + assert config_paths + split_qwen_targets = {"q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"} + + for path in config_paths: + payload = yaml.safe_load(path.read_text(encoding="utf-8")) + model = payload.get("model", payload) + lora = payload.get("lora", payload) + model_path = model["model_path"].lower() + targets = lora["lora_target_modules"] + assert len(targets) == len(set(targets)), path + + if "qwen3.5" in model_path: + assert split_qwen_targets | {"g_proj"} <= set(targets), path + elif "qwen3" in model_path: + assert split_qwen_targets <= set(targets), path + elif "llama" in model_path: + assert {"qkv_proj", "o_proj", "gate_up_proj", "down_proj"} <= set(targets), path diff --git a/tests/models/test_lora_target_manifest.py b/tests/models/test_lora_target_manifest.py deleted file mode 100644 index f2d7a128..00000000 --- a/tests/models/test_lora_target_manifest.py +++ /dev/null @@ -1,106 +0,0 @@ -from __future__ import annotations - -import pytest -import torch.nn as nn - -from xorl.lora.modules import LoraLinear -from xorl.lora.target_manifest import ( - collect_lora_runtime_modules, - load_lora_target_manifest, - resolve_lora_target_modules, - validate_lora_target_manifest, -) -from xorl.lora.utils import inject_lora_into_model - - -class _Attention(nn.Module): - def __init__(self): - super().__init__() - self.q_proj = nn.Linear(8, 8, bias=False) - self.o_proj = nn.Linear(8, 8, bias=False) - - -class _Layer(nn.Module): - def __init__(self): - super().__init__() - self.self_attn = _Attention() - - -class _Model(nn.Module): - def __init__(self): - super().__init__() - self.model = nn.Module() - self.model.layers = nn.ModuleList([_Layer(), _Layer()]) - - -def _manifest(*, q_count=2, include_o=False, rank=4): - expected = [{"pattern": "model.layers.*.self_attn.q_proj", "count": q_count, "rank": rank}] - targets = ["q_proj"] - if include_o: - expected.append({"pattern": "model.layers.*.self_attn.o_proj", "count": 2, "rank": rank}) - targets.append("o_proj") - return { - "schema_version": 1, - "target_modules": targets, - "expected_modules": expected, - "allow_unlisted": False, - } - - -def test_manifest_drives_targets_and_validates_runtime_coverage(): - model = _Model() - inject_lora_into_model(model, r=4, lora_alpha=8, target_manifest=_manifest()) - - modules = collect_lora_runtime_modules(model) - assert modules == { - "model.layers.0.self_attn.q_proj": 4, - "model.layers.1.self_attn.q_proj": 4, - } - assert isinstance(model.model.layers[0].self_attn.q_proj, LoraLinear) - assert isinstance(model.model.layers[0].self_attn.o_proj, nn.Linear) - - -def test_manifest_fails_closed_on_count_mismatch(): - with pytest.raises(ValueError, match="matched 2 modules, expected 3"): - inject_lora_into_model(_Model(), r=4, lora_alpha=8, target_manifest=_manifest(q_count=3)) - - -def test_manifest_fails_closed_on_rank_mismatch(): - model = _Model() - inject_lora_into_model(model, r=4, lora_alpha=8, target_modules=["q_proj"]) - with pytest.raises(ValueError, match="rank mismatch"): - validate_lora_target_manifest(model, _manifest(rank=2)) - - -def test_configured_targets_must_match_manifest(): - with pytest.raises(ValueError, match="do not match"): - resolve_lora_target_modules(["q_proj", "o_proj"], _manifest()) - - -def test_manifest_rejects_unlisted_lora_modules(): - model = _Model() - inject_lora_into_model(model, r=4, lora_alpha=8, target_modules=["q_proj", "o_proj"]) - with pytest.raises(ValueError, match="unlisted LoRA modules"): - validate_lora_target_manifest(model, _manifest()) - - -@pytest.mark.parametrize( - ("field", "value", "message"), - [ - ("schema_version", True, "schema_version"), - ("allow_unlisted", "false", "allow_unlisted must be a Boolean"), - ], -) -def test_manifest_rejects_non_exact_scalar_types(field, value, message): - manifest = _manifest() - manifest[field] = value - with pytest.raises(ValueError, match=message): - load_lora_target_manifest(manifest) - - -@pytest.mark.parametrize(("field", "value"), [("count", True), ("rank", True)]) -def test_manifest_rejects_booleans_for_integer_fields(field, value): - manifest = _manifest() - manifest["expected_modules"][0][field] = value - with pytest.raises(ValueError, match=field): - load_lora_target_manifest(manifest) diff --git a/tests/models/test_minimax_m3_support.py b/tests/models/test_minimax_m3_support.py index 1771cd9e..89c66855 100644 --- a/tests/models/test_minimax_m3_support.py +++ b/tests/models/test_minimax_m3_support.py @@ -6,7 +6,6 @@ import pytest import torch -from xorl.models import module_utils from xorl.models.auto import _load_local_xorl_config from xorl.models.registry import ModelRegistry from xorl.models.transformers.minimax_m3.checkpoint_handler import MiniMaxM3CheckpointHandler @@ -115,7 +114,7 @@ def _tiny_config(**overrides): return MiniMaxM3Config(**values) -def test_minimax_m3_config_adapts_top_level_hf_config(): +def test_minimax_m3_configuration_and_registration_policy(tmp_path): cfg = MiniMaxM3Config.from_hf_config(_namespace(_hf_minimax_config_dict())) assert cfg.model_type == "xorl_minimax_m3" @@ -136,8 +135,14 @@ def test_minimax_m3_config_adapts_top_level_hf_config(): assert cfg.vision_config == {"hidden_size": 1024} assert cfg._moe_implementation == "native" + _assert_minimax_m3_local_config_loader_and_registry(tmp_path) + _assert_minimax_m3_native_config_round_trip() + _assert_minimax_m3_text_runtime_and_admission_policy() + _assert_minimax_m3_checkpoint_mapping_and_expert_ownership_policy() + _assert_minimax_m3_msa_cpu_path_fails_loudly_and_paging_is_stable() -def test_minimax_m3_local_config_loader_and_registry(tmp_path): + +def _assert_minimax_m3_local_config_loader_and_registry(tmp_path): config_dir = tmp_path / "minimax" config_dir.mkdir() (config_dir / "config.json").write_text(__import__("json").dumps(_hf_minimax_config_dict())) @@ -149,7 +154,7 @@ def test_minimax_m3_local_config_loader_and_registry(tmp_path): assert "MiniMaxM3SparseForCausalLM" in ModelRegistry.supported_models -def test_minimax_m3_config_adapts_xorl_native_config_without_text_config(): +def _assert_minimax_m3_native_config_round_trip(): original = _tiny_config(text_config=None) cfg = MiniMaxM3Config.from_hf_config(original) @@ -159,7 +164,7 @@ def test_minimax_m3_config_adapts_xorl_native_config_without_text_config(): assert cfg.sparse_attention_freq == original.sparse_attention_freq -def test_minimax_m3_swigluoai_matches_oai_formula(): +def _assert_minimax_m3_activation_and_router_policy(): gate = torch.tensor([[-9.0, -1.0, 1.0, 9.0]]) up = torch.tensor([[-9.0, -1.0, 1.0, 9.0]]) @@ -170,8 +175,10 @@ def test_minimax_m3_swigluoai_matches_oai_formula(): expected = expected_gate * torch.sigmoid(1.702 * expected_gate) * (expected_up + 1.0) torch.testing.assert_close(actual, expected) + _assert_minimax_m3_sigmoid_router_selection() + -def test_minimax_m3_sigmoid_router_uses_bias_only_for_selection(): +def _assert_minimax_m3_sigmoid_router_selection(): router = MiniMaxM3Router(num_experts=4, top_k=2, routed_scaling_factor=2.0, use_routing_bias=True) logits = torch.tensor([[4.0, 3.0, -1.0, -2.0]]) bias = torch.tensor([-10.0, -10.0, 20.0, 0.0]) @@ -185,7 +192,9 @@ def test_minimax_m3_sigmoid_router_uses_bias_only_for_selection(): torch.testing.assert_close(weights, expected) -def test_minimax_m3_tiny_forward_backward_with_labels(): +def _assert_minimax_m3_text_runtime_and_admission_policy(): + _assert_minimax_m3_activation_and_router_policy() + torch.manual_seed(0) cfg = _tiny_config() model = MiniMaxM3SparseForCausalLM(cfg) @@ -199,8 +208,11 @@ def test_minimax_m3_tiny_forward_backward_with_labels(): out.loss.backward() assert model.lm_head.weight.grad is not None + _assert_minimax_m3_rejects_multimodal_inputs_and_tokens() + _assert_minimax_m3_rejects_unsupported_parallel_modes() + -def test_minimax_m3_text_only_rejects_multimodal_inputs_and_tokens(): +def _assert_minimax_m3_rejects_multimodal_inputs_and_tokens(): cfg = _tiny_config(image_token_index=5, video_token_index=6) model = MiniMaxM3SparseForCausalLM(cfg) @@ -214,7 +226,7 @@ def test_minimax_m3_text_only_rejects_multimodal_inputs_and_tokens(): model(input_ids=torch.tensor([[1, 6, 3]])) -def test_minimax_m3_unsupported_parallel_modes_fail_clearly(): +def _assert_minimax_m3_rejects_unsupported_parallel_modes(): ps = SimpleNamespace(tp_size=2, pp_size=1, ringattn_size=1, ulysses_size=1, lm_head_tp_size=1) with pytest.raises(ValueError, match="supports data/FSDP2 and expert parallelism only"): @@ -223,7 +235,7 @@ def test_minimax_m3_unsupported_parallel_modes_fail_clearly(): assert "tensor parallelism" in MINIMAX_M3_UNSUPPORTED_PARALLEL_MESSAGE -def test_minimax_m3_checkpoint_handler_maps_language_weights_and_skips_multimodal(): +def _assert_minimax_m3_checkpoint_mapping_and_expert_ownership_policy(): handler = MiniMaxM3CheckpointHandler(num_experts=2) hidden = 2 intermediate = 3 @@ -278,8 +290,10 @@ def test_minimax_m3_checkpoint_handler_maps_language_weights_and_skips_multimoda assert handler.on_load_weight("multi_modal_projector.linear_1.weight", torch.ones(1)) == [] assert handler.on_load_weight("patch_merge_mlp.linear_1.weight", torch.ones(1)) == [] + _assert_minimax_m3_checkpoint_ep_skip_counts_raw_keys() -def test_minimax_m3_checkpoint_handler_ep_skip_counts_raw_keys(): + +def _assert_minimax_m3_checkpoint_ep_skip_counts_raw_keys(): handler = MiniMaxM3CheckpointHandler(num_experts=4, ep_rank=1, ep_size=2) skip = handler.get_skip_key_fn() @@ -300,7 +314,7 @@ def test_minimax_m3_checkpoint_handler_ep_skip_counts_raw_keys(): assert mapped["model.layers.3.mlp.experts.down_proj"].shape[0] == 2 -def test_minimax_m3_msa_cpu_path_fails_loudly_and_paging_is_stable(): +def _assert_minimax_m3_msa_cpu_path_fails_loudly_and_paging_is_stable(): x = torch.arange(2 * 3 * 1 * 2, dtype=torch.float32).reshape(2, 3, 1, 2) pages, indices = _to_paged_kv(x, torch.tensor([3, 1], dtype=torch.int32), page_size=2) assert pages.shape == (3, 1, 2, 2) @@ -321,15 +335,3 @@ def test_minimax_m3_msa_cpu_path_fails_loudly_and_paging_is_stable(): force_begin_blocks=1, force_end_blocks=1, ) - - -@pytest.mark.parametrize( - "key", - [ - "language_model.model.layers.3.block_sparse_moe.experts.0.w1.weight", - "model.language_model.model.layers.3.block_sparse_moe.experts.0.w2.weight", - "model.layers.3.block_sparse_moe.experts.0.w3.weight", - ], -) -def test_minimax_m3_grouped_loader_classifies_block_sparse_experts(key): - assert module_utils._is_checkpoint_expert_key(key) diff --git a/tests/models/test_model_state.py b/tests/models/test_model_state.py index ef849514..a72249be 100644 --- a/tests/models/test_model_state.py +++ b/tests/models/test_model_state.py @@ -32,47 +32,10 @@ def __init__(self): self.proj = torch.nn.Linear(4, 3) -def test_reference_state_dict_bypasses_dcp_state_dict_and_skips_nonpersistent_buffers(monkeypatch): - monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(dp_mode="none")) - monkeypatch.setattr( - checkpointer, - "get_model_state_dict", - lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("unexpected DCP state_dict call")), - ) - - model_state = checkpointer.ModelState(_TinyModel()) - state_dict = model_state.reference_state_dict() - - assert "linear.weight" in state_dict - assert "persistent_buf" in state_dict - assert "scratch_buf" not in state_dict - - -def test_reference_state_dict_includes_qarl_persistent_buffers(monkeypatch): - monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(dp_mode="none")) - monkeypatch.setattr( - checkpointer, - "get_model_state_dict", - lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("unexpected DCP state_dict call")), - ) - - model = _TinyQARLModel() - model.proj(torch.randn(2, 4)) - state_dict = checkpointer.ModelState(model).reference_state_dict() - - assert "proj.weight" in state_dict - assert "proj.qarl_input_amax" in state_dict - assert "proj.qarl_weight_amax" in state_dict - assert "proj.qarl_input_scale_inv" in state_dict - assert "proj.qarl_weight_scale_inv" in state_dict - assert "proj.qarl_forward_count" in state_dict - assert state_dict["proj.qarl_weight_scale_inv"].shape == (2, 2) - assert state_dict["proj.qarl_forward_count"].item() == 1 - - -def test_checkpoint_metadata_records_qarl_persistent_buffers(tmp_path, monkeypatch): +def _assert_checkpoint_compatibility_detects_qarl_buffer_mismatch(tmp_path, monkeypatch): + tmp_path.mkdir(parents=True) + monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(pp_enabled=False)) monkeypatch.setattr(checkpointer.dist, "get_rank", lambda: 0) - checkpointer._save_checkpoint_metadata(str(tmp_path), _TinyQARLModel()) metadata = json.loads((tmp_path / "checkpoint_metadata.json").read_text(encoding="utf-8")) @@ -86,11 +49,6 @@ def test_checkpoint_metadata_records_qarl_persistent_buffers(tmp_path, monkeypat "proj.qarl_weight_scale_inv", ] - -def test_checkpoint_compatibility_detects_qarl_buffer_mismatch(tmp_path, monkeypatch): - monkeypatch.setattr(checkpointer.dist, "get_rank", lambda: 0) - checkpointer._save_checkpoint_metadata(str(tmp_path), _TinyQARLModel()) - with pytest.raises(RuntimeError, match="Unexpected buffers"): checkpointer._validate_checkpoint_compatibility(str(tmp_path), _TinyPlainLinearModel(), strict=True) @@ -108,7 +66,7 @@ def test_checkpoint_compatibility_detects_qarl_buffer_mismatch(tmp_path, monkeyp } -def test_checkpoint_key_contract_does_not_collect_without_pipeline_parallelism(monkeypatch): +def _assert_checkpoint_model_key_and_compatibility_policy(tmp_path, monkeypatch): monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(pp_enabled=False)) monkeypatch.setattr( checkpointer.dist, @@ -122,8 +80,13 @@ def test_checkpoint_key_contract_does_not_collect_without_pipeline_parallelism(m assert buffer_keys == ["persistent_buf"] assert pipeline_key_union is False + _assert_checkpoint_key_contract_unions_pipeline_stage_keys(monkeypatch) + _assert_checkpoint_compatibility_detects_qarl_buffer_mismatch(tmp_path / "qarl", monkeypatch) + _assert_checkpoint_metadata_uses_pipeline_stage_key_union(tmp_path / "metadata", monkeypatch) + _assert_pipeline_lora_checkpoint_compatibility_policy(tmp_path, monkeypatch) + -def test_checkpoint_key_contract_unions_pipeline_stage_keys(monkeypatch): +def _assert_checkpoint_key_contract_unions_pipeline_stage_keys(monkeypatch): monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(pp_enabled=True)) monkeypatch.setattr(checkpointer.dist, "is_initialized", lambda: True) monkeypatch.setattr(checkpointer.dist, "get_world_size", lambda group=None: 2) @@ -145,7 +108,8 @@ def fake_all_gather_object(output, local_keys, group=None): assert pipeline_key_union is True -def test_checkpoint_metadata_records_pipeline_stage_key_union(tmp_path, monkeypatch): +def _assert_checkpoint_metadata_uses_pipeline_stage_key_union(tmp_path, monkeypatch): + tmp_path.mkdir(parents=True) monkeypatch.setattr(checkpointer.dist, "get_rank", lambda: 0) monkeypatch.setattr( checkpointer, @@ -166,34 +130,17 @@ def test_checkpoint_metadata_records_pipeline_stage_key_union(tmp_path, monkeypa assert metadata["num_buffers"] == 1 assert metadata["buffer_keys"] == ["layers.1.cache"] - -def test_checkpoint_compatibility_validates_pipeline_stage_key_union(tmp_path, monkeypatch): - metadata = { - "parameter_keys": ["linear.weight", "stage_1.weight"], - "buffer_keys": ["persistent_buf", "stage_1.cache"], - } - (tmp_path / "checkpoint_metadata.json").write_text(json.dumps(metadata), encoding="utf-8") - monkeypatch.setattr( - checkpointer, - "_get_checkpoint_model_keys", - lambda model, process_group=None: ( - ["linear.weight", "stage_1.weight"], - ["persistent_buf", "stage_1.cache"], - True, - ), - ) - result = checkpointer._validate_checkpoint_compatibility( str(tmp_path), _TinyModel(), strict=True, process_group=object() ) assert result["compatible"] is True assert result["pipeline_parallel_key_union"] is True - assert result["model_parameter_count"] == 2 - assert result["model_buffer_count"] == 2 + assert result["model_parameter_count"] == 3 + assert result["model_buffer_count"] == 1 -def test_checkpoint_compatibility_allows_pipeline_base_checkpoint_into_lora_model(tmp_path, monkeypatch): +def _assert_pipeline_lora_checkpoint_compatibility_policy(tmp_path, monkeypatch): metadata = { "parameter_keys": ["stage_0.weight", "stage_1.weight"], "buffer_keys": [], @@ -217,8 +164,11 @@ def test_checkpoint_compatibility_allows_pipeline_base_checkpoint_into_lora_mode assert result["load_mode"] == "base_to_lora" assert set(result["missing_lora_keys"]) == {"stage_0.lora_A", "stage_1.lora_A"} + _assert_pipeline_lora_only_checkpoint_compatibility(tmp_path / "lora-only", monkeypatch) + -def test_checkpoint_compatibility_allows_pipeline_lora_only_checkpoint(tmp_path, monkeypatch): +def _assert_pipeline_lora_only_checkpoint_compatibility(tmp_path, monkeypatch): + tmp_path.mkdir(parents=True) metadata = { "parameter_keys": ["stage_0.lora_A", "stage_1.lora_A"], "buffer_keys": [], @@ -244,110 +194,103 @@ def test_checkpoint_compatibility_allows_pipeline_lora_only_checkpoint(tmp_path, assert set(result["missing_non_lora_keys"]) == {"stage_0.weight", "stage_1.weight"} -def test_distributed_checkpointer_load_skips_missing_optimizer_state(tmp_path, monkeypatch): - captured = {} +def test_distributed_checkpointer_io_policy(tmp_path, monkeypatch): + with monkeypatch.context() as case_patch: + _assert_distributed_checkpointer_process_group_selection_policy(case_patch) + compatibility_root = tmp_path / "compatibility" + compatibility_root.mkdir() + with monkeypatch.context() as case_patch: + _assert_checkpoint_model_key_and_compatibility_policy(compatibility_root, case_patch) + optimizer_root = tmp_path / "optimizer" + optimizer_root.mkdir() + with monkeypatch.context() as case_patch: + _assert_optimizer_state_checkpoint_filtering_policy(optimizer_root, case_patch) + with monkeypatch.context() as case_patch: + _assert_distributed_checkpointer_metadata_admission_policy(tmp_path / "metadata", case_patch) + with monkeypatch.context() as case_patch: + _assert_pipeline_load_uses_custom_dcp_group(tmp_path / "load", case_patch) + with monkeypatch.context() as case_patch: + _assert_distributed_checkpointer_save_group_policy(tmp_path / "save", case_patch) + + +def _assert_distributed_checkpointer_process_group_selection_policy(monkeypatch): + created = [] + fake_group = object() - class _FakeReader: - def __init__(self, path): - self.path = path - - def read_metadata(self): - return SimpleNamespace(state_dict_metadata={"model.linear.weight": object()}) - - def fake_dcp_load(state_dict, storage_reader, process_group=None, planner=None, no_dist=False): - captured["state_keys"] = set(state_dict) - captured["storage_reader"] = storage_reader - captured["process_group"] = process_group - captured["planner"] = planner - captured["no_dist"] = no_dist + monkeypatch.setattr( + checkpointer, + "dist", + SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + get_backend=lambda: "nccl", + new_group=lambda backend: created.append(backend) or fake_group, + ), + ) + monkeypatch.setattr(checkpointer.DistributedCheckpointer, "_sync_process_group", None) - monkeypatch.setattr(checkpointer, "FileSystemReader", _FakeReader) - monkeypatch.setattr(checkpointer.dcp, "load", fake_dcp_load) + assert checkpointer.DistributedCheckpointer._get_sync_process_group() is fake_group + assert checkpointer.DistributedCheckpointer._get_sync_process_group() is fake_group + assert created == ["gloo"] - state = {"model": _TinyModel(), "optimizer": object()} - result = checkpointer.DistributedCheckpointer.load(str(tmp_path), state) + monkeypatch.setattr( + checkpointer, + "dist", + SimpleNamespace( + is_available=lambda: True, + is_initialized=lambda: True, + get_backend=lambda: "gloo", + ), + ) + checkpointer.DistributedCheckpointer._sync_process_group = None + assert checkpointer.DistributedCheckpointer._get_sync_process_group() is None - assert result is state - assert captured["state_keys"] == {"model"} - assert isinstance(captured["storage_reader"], _FakeReader) - assert captured["planner"] is not None - assert captured["no_dist"] is False + monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(pp_enabled=False)) + monkeypatch.setattr( + checkpointer.DistributedCheckpointer, + "_get_sync_process_group", + classmethod(lambda cls: (_ for _ in ()).throw(AssertionError("unexpected process group"))), + ) + assert checkpointer.DistributedCheckpointer._get_metadata_process_group() is None + assert checkpointer.DistributedCheckpointer._get_metadata_process_group(object()) is None + custom_group = object() + monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(pp_enabled=True)) + monkeypatch.setattr( + checkpointer.DistributedCheckpointer, + "_get_sync_process_group", + classmethod(lambda cls: (_ for _ in ()).throw(AssertionError("unexpected global process group"))), + ) + assert checkpointer.DistributedCheckpointer._get_metadata_process_group(custom_group) is custom_group -def test_distributed_checkpointer_no_dist_still_uses_gloo_for_pipeline_key_validation(tmp_path, monkeypatch): - captured = {} metadata_group = object() - - class _FakeReader: - def __init__(self, path): - self.path = path - - def read_metadata(self): - return SimpleNamespace(state_dict_metadata={"model.linear.weight": object()}) - - def fake_validate(checkpoint_dir, model, strict=True, process_group=None): - captured["validation_process_group"] = process_group - return {"validated": False, "reason": "test"} - - def fake_dcp_load(state_dict, storage_reader, process_group=None, planner=None, no_dist=False): - captured["dcp_process_group"] = process_group - captured["no_dist"] = no_dist - - monkeypatch.setenv("XORL_DCP_LOAD_NO_DIST", "1") - monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(pp_enabled=True)) monkeypatch.setattr( checkpointer.DistributedCheckpointer, "_get_sync_process_group", classmethod(lambda cls: metadata_group), ) - monkeypatch.setattr(checkpointer, "_validate_checkpoint_compatibility", fake_validate) - monkeypatch.setattr(checkpointer, "FileSystemReader", _FakeReader) - monkeypatch.setattr(checkpointer.dcp, "load", fake_dcp_load) + assert checkpointer.DistributedCheckpointer._get_metadata_process_group() is metadata_group - checkpointer.DistributedCheckpointer.load(str(tmp_path), {"model": _TinyModel()}) - assert captured["validation_process_group"] is metadata_group - assert captured["dcp_process_group"] is None - assert captured["no_dist"] is True - - -def test_distributed_checkpointer_no_dist_non_pipeline_avoids_process_groups(tmp_path, monkeypatch): - captured = {} - - class _FakeReader: - def __init__(self, path): - self.path = path - - def read_metadata(self): - return SimpleNamespace(state_dict_metadata={"model.linear.weight": object()}) - - def fake_validate(checkpoint_dir, model, strict=True, process_group=None): - captured["validation_process_group"] = process_group - return {"validated": False, "reason": "test"} - - def fake_dcp_load(state_dict, storage_reader, process_group=None, planner=None, no_dist=False): - captured["dcp_process_group"] = process_group - captured["no_dist"] = no_dist +def _assert_distributed_checkpointer_metadata_admission_policy(tmp_path, monkeypatch): + source = _TinyModel() + source.linear.weight.data.copy_(torch.arange(8, dtype=torch.float32).reshape(2, 4)) + checkpointer.dcp.save({"model": source.state_dict()}, checkpoint_id=str(tmp_path)) + target = _TinyModel() + optimizer = torch.optim.Adam(target.parameters()) monkeypatch.setenv("XORL_DCP_LOAD_NO_DIST", "1") monkeypatch.setattr(checkpointer, "get_parallel_state", lambda: SimpleNamespace(pp_enabled=False)) - monkeypatch.setattr( - checkpointer.DistributedCheckpointer, - "_get_sync_process_group", - classmethod(lambda cls: (_ for _ in ()).throw(AssertionError("unexpected process group"))), - ) - monkeypatch.setattr(checkpointer, "_validate_checkpoint_compatibility", fake_validate) - monkeypatch.setattr(checkpointer, "FileSystemReader", _FakeReader) - monkeypatch.setattr(checkpointer.dcp, "load", fake_dcp_load) - checkpointer.DistributedCheckpointer.load(str(tmp_path), {"model": _TinyModel()}) + state = {"model": target, "optimizer": optimizer} + result = checkpointer.DistributedCheckpointer.load(str(tmp_path), state) - assert captured["validation_process_group"] is None - assert captured["dcp_process_group"] is None - assert captured["no_dist"] is True + assert result is state + assert torch.equal(target.linear.weight, source.linear.weight) + assert optimizer.state == {} -def test_distributed_checkpointer_pipeline_validation_uses_custom_dcp_group(tmp_path, monkeypatch): +def _assert_pipeline_load_uses_custom_dcp_group(tmp_path, monkeypatch): captured = {} custom_group = object() @@ -385,7 +328,7 @@ def fake_dcp_load(state_dict, storage_reader, process_group=None, planner=None, assert captured["dcp_process_group"] is custom_group -def test_distributed_checkpointer_sync_save_reuses_pipeline_gloo_for_metadata(tmp_path, monkeypatch): +def _assert_distributed_checkpointer_save_group_policy(tmp_path, monkeypatch): captured = {} sync_group = object() @@ -420,8 +363,10 @@ def test_distributed_checkpointer_sync_save_reuses_pipeline_gloo_for_metadata(tm assert captured["dcp_process_group"] is sync_group assert captured["metadata_process_group"] is sync_group + _assert_async_save_non_pipeline_avoids_second_gloo(tmp_path, monkeypatch) + -def test_distributed_checkpointer_async_save_non_pipeline_avoids_second_gloo(tmp_path, monkeypatch): +def _assert_async_save_non_pipeline_avoids_second_gloo(tmp_path, monkeypatch): captured = {} async_group = object() @@ -466,7 +411,7 @@ def test_distributed_checkpointer_async_save_non_pipeline_avoids_second_gloo(tmp assert captured["metadata_process_group"] is None -def test_optimizer_state_filters_load_target_to_checkpoint_keys(): +def _assert_optimizer_state_checkpoint_filtering_policy(tmp_path, monkeypatch): class _FakeMultiOptimizer: _is_multi_optimizer = True @@ -501,8 +446,13 @@ def load_state_dict(self, state_dict, strict=True): assert optimizer.loaded_state_dict is state_dict assert optimizer.loaded_strict is False + with monkeypatch.context() as case_patch: + _assert_distributed_load_passes_optimizer_metadata_keys(tmp_path / "metadata", case_patch) + with monkeypatch.context() as case_patch: + _assert_multi_optimizer_load_filters_state_per_child(case_patch) + -def test_distributed_checkpointer_load_passes_optimizer_metadata_keys(tmp_path, monkeypatch): +def _assert_distributed_load_passes_optimizer_metadata_keys(tmp_path, monkeypatch): captured = {} class _FakeReader: @@ -534,7 +484,7 @@ def fake_dcp_load(state_dict, storage_reader, process_group=None, planner=None, } -def test_multi_optimizer_load_filters_state_per_child_optimizer(monkeypatch): +def _assert_multi_optimizer_load_filters_state_per_child(monkeypatch): ep_optimizer = object() non_ep_optimizer = object() calls = [] diff --git a/tests/models/test_module_utils_broadcast.py b/tests/models/test_module_utils_broadcast.py index b6e53a7c..a1de5aff 100644 --- a/tests/models/test_module_utils_broadcast.py +++ b/tests/models/test_module_utils_broadcast.py @@ -12,9 +12,11 @@ from torch.distributed.tensor import DTensor from xorl.models import module_utils +from xorl.models.transformers.qwen3_5_shared import QWEN3_5_CHECKPOINT_SKIP_KEY_PATTERNS pytestmark = [pytest.mark.cpu] +_ORIGINAL_GET_OBJECT_BROADCAST_DEVICE = module_utils._get_object_broadcast_device class _DummyModel: @@ -79,6 +81,13 @@ def _cpu_dtensor_materialize_worker(rank: int, world_size: int, port: int) -> No materialized = module_utils._materialize_tensor_for_save(dtensor) assert materialized.device.type == "cpu" assert torch.equal(materialized, full_tensor) + targeted = module_utils._materialize_tensor_for_save(dtensor, dst_rank=3) + if rank == 3: + assert targeted is not None + assert targeted.device.type == "cpu" + assert torch.equal(targeted, full_tensor) + else: + assert targeted is None finally: dist.destroy_process_group() @@ -116,82 +125,12 @@ def _cpu_dtensor_materialize_to_rank_worker(rank: int, world_size: int, port: in dist.destroy_process_group() -def _cpu_dtensor_materialize_2d_to_rank_worker(rank: int, world_size: int, port: int, dst_rank: int) -> None: - dist.init_process_group("gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=world_size) - try: - module_utils._cpu_save_device_mesh_cache.clear() - mesh = DeviceMesh( - "cpu", - mesh=torch.arange(world_size).view(2, 2), - mesh_dim_names=("ep", "fsdp"), - backend_override=(("gloo", None), ("gloo", None)), - ) - full_tensor = torch.arange(16, dtype=torch.float32).view(4, 4) - row = rank // 2 - col = rank % 2 - local_tensor = full_tensor[row * 2 : (row + 1) * 2, col * 2 : (col + 1) * 2].clone() - dtensor = DTensor.from_local( - local_tensor, - device_mesh=mesh, - placements=[DTShard(0), DTShard(1)], - shape=full_tensor.shape, - stride=full_tensor.stride(), - ) - materialized = module_utils._materialize_tensor_for_save(dtensor, dst_rank=dst_rank) - if rank == dst_rank: - assert materialized is not None - assert materialized.device.type == "cpu" - assert torch.equal(materialized, full_tensor) - else: - assert materialized is None - finally: - dist.destroy_process_group() - - -def test_copy_into_existing_dtensor_shard_for_replicated_tensor(): +def _assert_dtensor_checkpoint_materialization_policy(): dtensor = _FakeDTensor(torch.zeros(4, dtype=torch.float32), mesh_size=4, local_rank=2, placements=(Replicate(),)) full_tensor = torch.arange(4, dtype=torch.float32) - - copied = module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) - - assert copied is True + assert module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) is True assert torch.equal(dtensor._local_tensor, full_tensor) - -def test_materialize_tensor_for_save_uses_cpu_mesh_for_dtensors(): - port = _find_free_port() - mp.start_processes( - _cpu_dtensor_materialize_worker, - args=(4, port), - nprocs=4, - join=True, - start_method="fork", - ) - - -def test_materialize_tensor_for_save_gathers_1d_dtensor_to_writer_rank(): - port = _find_free_port() - mp.start_processes( - _cpu_dtensor_materialize_to_rank_worker, - args=(4, port, 2), - nprocs=4, - join=True, - start_method="fork", - ) - - -def test_materialize_tensor_for_save_gathers_2d_dtensor_to_writer_rank(): - port = _find_free_port() - mp.start_processes( - _cpu_dtensor_materialize_2d_to_rank_worker, - args=(4, port, 3), - nprocs=4, - join=True, - start_method="fork", - ) - - -def test_copy_into_existing_dtensor_shard_for_sharded_tensor(): dtensor = _FakeDTensor( torch.zeros(2, 3, dtype=torch.float32), mesh_size=4, @@ -199,14 +138,9 @@ def test_copy_into_existing_dtensor_shard_for_sharded_tensor(): placements=(DTShard(0),), ) full_tensor = torch.arange(24, dtype=torch.float32).view(8, 3) - - copied = module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) - - assert copied is True + assert module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) is True assert torch.equal(dtensor._local_tensor, full_tensor[2:4]) - -def test_copy_into_existing_dtensor_shard_trims_padded_tail_shards(): dtensor = _FakeDTensor( torch.zeros(0, 3, dtype=torch.float32), mesh_size=8, @@ -214,24 +148,28 @@ def test_copy_into_existing_dtensor_shard_trims_padded_tail_shards(): placements=(DTShard(0),), ) full_tensor = torch.arange(15, dtype=torch.float32).view(5, 3) - - copied = module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) - - assert copied is True + assert module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) is True assert tuple(dtensor._local_tensor.shape) == (0, 3) - -def test_copy_into_existing_dtensor_shard_rejects_shape_mismatched_replicates(): dtensor = _FakeDTensor(torch.zeros(1, 3, dtype=torch.float32), mesh_size=4, local_rank=0, placements=(Replicate(),)) full_tensor = torch.arange(15, dtype=torch.float32).view(5, 3) - - copied = module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) - - assert copied is False + assert module_utils._copy_into_existing_dtensor_shard(dtensor, full_tensor) is False assert torch.equal(dtensor._local_tensor, torch.zeros(1, 3, dtype=torch.float32)) + for worker, worker_args in ( + (_cpu_dtensor_materialize_worker, (4, _find_free_port())), + (_cpu_dtensor_materialize_to_rank_worker, (4, _find_free_port(), 2)), + ): + mp.start_processes( + worker, + args=worker_args, + nprocs=4, + join=True, + start_method="fork", + ) + -def test_broadcast_object_list_serializes_over_tensor_broadcast_for_nccl_groups(monkeypatch): +def _assert_object_broadcast_transport_policy(monkeypatch): fake_group = object() state = {"rank": 3} stored = [] @@ -263,8 +201,11 @@ def fake_broadcast(tensor, src=0, group=None): assert received_payload == source_payload + _assert_get_object_broadcast_device_uses_default_nccl_group(monkeypatch) + _assert_object_broadcast_weight_load_uses_weight_load_group(monkeypatch) + -def test_get_object_broadcast_device_uses_default_nccl_group(monkeypatch): +def _assert_get_object_broadcast_device_uses_default_nccl_group(monkeypatch): fake_dist = SimpleNamespace( is_available=lambda: True, is_initialized=lambda: True, @@ -272,13 +213,14 @@ def test_get_object_broadcast_device_uses_default_nccl_group(monkeypatch): ) monkeypatch.setattr(module_utils, "dist", fake_dist) + monkeypatch.setattr(module_utils, "_get_object_broadcast_device", _ORIGINAL_GET_OBJECT_BROADCAST_DEVICE) monkeypatch.setattr(module_utils, "get_device_type", lambda: "cuda") monkeypatch.setattr(module_utils, "get_device_id", lambda: 3) assert module_utils._get_object_broadcast_device(None) == torch.device("cuda:3") -def test_broadcast_object_list_weight_load_uses_weight_load_group(monkeypatch): +def _assert_object_broadcast_weight_load_uses_weight_load_group(monkeypatch): fake_group = object() calls = [] @@ -295,7 +237,16 @@ def test_broadcast_object_list_weight_load_uses_weight_load_group(monkeypatch): assert calls == [(payload, 7, fake_group)] -def test_rank0_broadcast_path_calls_load_state_dict_on_nonzero_ranks(monkeypatch): +def _assert_rank0_checkpoint_resolution_transport_and_loading_policy(monkeypatch): + with monkeypatch.context() as case_patch: + _assert_object_broadcast_transport_policy(case_patch) + with monkeypatch.context() as case_patch: + _assert_rank0_broadcast_loading_policy(case_patch) + with monkeypatch.context() as case_patch: + _assert_state_dict_resolution_policy(case_patch) + + +def _assert_rank0_broadcast_loading_policy(monkeypatch): calls = [] def fake_broadcast_object_list(obj, src=0, group=None, device=None): @@ -331,8 +282,10 @@ def fake_load_state_dict(weights_path): assert calls == ["dummy-weights"] + _assert_rank0_broadcast_uses_handler_filtered_prefetch(monkeypatch) -def test_rank0_broadcast_path_uses_filtered_prefetch_for_handler_skips(monkeypatch): + +def _assert_rank0_broadcast_uses_handler_filtered_prefetch(monkeypatch): batch_meta_calls = [] dispatched = [] handler_calls = {"loaded": [], "skipped": []} @@ -418,7 +371,7 @@ def fail_prefetch(*args, **kwargs): ] -def test_try_load_state_dict_uses_rank0_for_local_resolution(monkeypatch): +def _assert_state_dict_resolution_policy(monkeypatch): local_resolution_calls = [] def fake_broadcast_object_list(obj, src=0, group=None, device=None): @@ -446,10 +399,11 @@ def fake_try_load_state_dict_local(weights_path, **kwargs): assert local_resolution_calls == [] assert [it.filepath for it in iterators] == ["shard-0.safetensors", "shard-1.safetensors"] + _assert_local_state_dict_directory_skips_broadcast(monkeypatch) + -@pytest.mark.parametrize( - ("key", "expected"), - [ +def _assert_checkpoint_expert_key_classifies_supported_expert_formats(): + cases = ( ("model.layers.43.mlp.experts.7.gate_proj.weight", True), ("model.layers.43.mlp.experts.7.down_proj.weight_scale_inv", True), ("model.layers.43.mlp.experts.gate_up_proj", True), @@ -459,17 +413,19 @@ def fake_try_load_state_dict_local(weights_path, **kwargs): ("layers.12.ffn.experts.7.w1.weight", True), ("layers.12.ffn.experts.7.w2.scale", True), ("model.layers.12.ffn.experts.7.w3.weight", True), + ("model.language_model.layers.43.mlp.experts.gate_up_proj", True), + ("model.language_model.layers.43.mlp.experts.down_proj", True), ("model.layers.43.mlp.shared_expert.down_proj.weight", False), ("layers.12.ffn.shared_experts.w1.weight", False), + ("model.language_model.layers.43.mlp.shared_expert.down_proj.weight", False), ("model.layers.43.mlp.gate_up_proj.weight", False), ("model.layers.43.self_attn.q_proj.weight", False), - ], -) -def test_checkpoint_expert_key_classifies_supported_expert_formats(key, expected): - assert module_utils._is_checkpoint_expert_key(key) is expected + ) + for key, expected in cases: + assert module_utils._is_checkpoint_expert_key(key) is expected, key -def test_try_load_state_dict_local_directory_skips_broadcast(monkeypatch): +def _assert_local_state_dict_directory_skips_broadcast(monkeypatch): local_resolution_calls = [] fake_dist = SimpleNamespace( @@ -493,15 +449,14 @@ def fake_try_load_state_dict_local(weights_path, **kwargs): assert [it.filepath for it in iterators] == ["local-shard.safetensors"] -def test_checkpoint_expert_filter_handles_wrapped_language_model_keys(): - assert module_utils._is_checkpoint_expert_key("model.language_model.layers.43.mlp.experts.gate_up_proj") - assert module_utils._is_checkpoint_expert_key("model.language_model.layers.43.mlp.experts.down_proj") - assert not module_utils._is_checkpoint_expert_key( - "model.language_model.layers.43.mlp.shared_expert.down_proj.weight" - ) - +def test_grouped_load_routing_and_strict_coverage_policy(monkeypatch): + _assert_dtensor_checkpoint_materialization_policy() + with monkeypatch.context() as case_patch: + _assert_rank0_checkpoint_resolution_transport_and_loading_policy(case_patch) + _assert_checkpoint_expert_key_classifies_supported_expert_formats() + with monkeypatch.context() as case_patch: + _assert_post_process_strict_coverage_policy(case_patch) -def test_grouped_load_weights_uses_filtered_prefetch_on_group_leader(monkeypatch): batch_meta_calls = [] dispatched = [] transfer_calls = [] @@ -670,8 +625,12 @@ def fail_prefetch(*args, **kwargs): ], ) + _assert_grouped_load_routes_hf_fused_experts(monkeypatch) + _assert_grouped_load_routes_ffn_expert_source_format(monkeypatch) + _assert_grouped_load_group_fallback_policy(monkeypatch) + -def test_grouped_load_weights_routes_hf_fused_experts_through_expert_queue(monkeypatch): +def _assert_grouped_load_routes_hf_fused_experts(monkeypatch): dense_loaded = [] expert_loaded = [] dispatched = [] @@ -706,7 +665,7 @@ def on_load_complete(self): class _GroupedModel: _checkpoint_conversion_mapping = {r"^model\.language_model\.": "model."} - _checkpoint_skip_key_patterns = [r"^mtp\."] + _checkpoint_skip_key_patterns = QWEN3_5_CHECKPOINT_SKIP_KEY_PATTERNS def named_buffers(self): return [] @@ -746,6 +705,8 @@ def fake_prefetch_filtered(state_dict_iterators, skip_key_fn, prefetch_count): if skip_key_fn(raw_expert_key): prefetch_calls.append("dense") assert skip_key_fn(raw_skipped_key) + assert skip_key_fn("mtp.pre_fc_norm_embedding.weight") + assert not skip_key_fn("model.layers.0.mlp.gate_proj.weight") assert not skip_key_fn("keep.weight") yield ({"keep.weight": torch.tensor([2.0])}, []) else: @@ -792,7 +753,7 @@ def fake_prefetch_filtered(state_dict_iterators, skip_key_fn, prefetch_count): assert dispatched == ["keep.weight", expert_param] -def test_grouped_load_weights_routes_ffn_expert_source_format_through_expert_queue(monkeypatch): +def _assert_grouped_load_routes_ffn_expert_source_format(monkeypatch): dense_loaded = [] expert_loaded = [] dispatched = [] @@ -896,7 +857,7 @@ def fake_prefetch_filtered(state_dict_iterators, skip_key_fn, prefetch_count): assert dispatched == ["keep.weight", expert_param] -def test_grouped_load_weights_treats_missing_dense_group_as_local(monkeypatch): +def _assert_grouped_load_group_fallback_policy(monkeypatch): dispatched = [] fake_group = object() @@ -982,8 +943,11 @@ def fake_prefetch_filtered(state_dict_iterators, skip_key_fn, prefetch_count): assert dispatched == ["keep.weight"] + _assert_grouped_load_falls_back_without_ep_group(monkeypatch) + _assert_grouped_load_strict_rejects_fallback_without_ep_group(monkeypatch) -def test_grouped_load_weights_falls_back_without_ep_group(monkeypatch): + +def _assert_grouped_load_falls_back_without_ep_group(monkeypatch): called = [] fake_dist = SimpleNamespace( @@ -1012,7 +976,7 @@ def test_grouped_load_weights_falls_back_without_ep_group(monkeypatch): assert not hasattr(model, "device") -def test_grouped_load_weights_strict_rejects_fallback_without_ep_group(monkeypatch): +def _assert_grouped_load_strict_rejects_fallback_without_ep_group(monkeypatch): fake_dist = SimpleNamespace( is_available=lambda: True, is_initialized=lambda: True, @@ -1031,7 +995,7 @@ def test_grouped_load_weights_strict_rejects_fallback_without_ep_group(monkeypat module_utils.grouped_load_weights(_DummyModel(), "dummy-weights", init_device="cpu", strict=True) -def test_post_process_strict_rejects_missing_unexpected_and_duplicate_names(): +def _assert_post_process_strict_coverage_policy(monkeypatch): with pytest.raises(RuntimeError, match="Strict checkpoint source-to-target coverage failed") as exc_info: module_utils.post_process_after_weight_loading( _DummyModel(), @@ -1047,8 +1011,12 @@ def test_post_process_strict_rejects_missing_unexpected_and_duplicate_names(): assert "unexpected.weight" in message assert "duplicate.weight" in message + _assert_post_process_strict_rejects_missing_persistent_buffer() + _assert_post_process_strict_rejects_duplicate_persistent_buffer() + _assert_post_process_strict_accepts_complete_persistent_buffer_coverage(monkeypatch) + -def test_post_process_strict_rejects_missing_persistent_buffer(): +def _assert_post_process_strict_rejects_missing_persistent_buffer(): with pytest.raises(RuntimeError, match="missing_persistent_buffers=.*missing.buffer"): module_utils.post_process_after_weight_loading( _DummyModel(), @@ -1058,7 +1026,7 @@ def test_post_process_strict_rejects_missing_persistent_buffer(): ) -def test_post_process_strict_rejects_duplicate_persistent_buffer(): +def _assert_post_process_strict_rejects_duplicate_persistent_buffer(): with pytest.raises(RuntimeError, match="duplicate_persistent_buffers=.*duplicate.buffer"): module_utils.post_process_after_weight_loading( _DummyModel(), @@ -1068,7 +1036,7 @@ def test_post_process_strict_rejects_duplicate_persistent_buffer(): ) -def test_post_process_strict_accepts_complete_persistent_buffer_coverage(monkeypatch): +def _assert_post_process_strict_accepts_complete_persistent_buffer_coverage(monkeypatch): dispatched = [] class _CompleteModel(_DummyModel): diff --git a/tests/models/test_moe_ep_native_combine.py b/tests/models/test_moe_ep_native_combine.py index db777227..9afbc897 100644 --- a/tests/models/test_moe_ep_native_combine.py +++ b/tests/models/test_moe_ep_native_combine.py @@ -16,13 +16,31 @@ pytestmark = [pytest.mark.cpu] -def test_qwen35_native_combine_admits_only_ep8(): +def test_qwen35_native_combine_policy(monkeypatch): assert QWEN35_NATIVE_EP_COMBINE_SIZES == frozenset({8}) validate_qwen35_native_ep_combine_size(8) for size in (1, 2, 4, 16): with pytest.raises(ValueError, match="EP8"): validate_qwen35_native_ep_combine_size(size) + blk = _qwen_block(exact=True) + assert blk._native_ep_combine + assert blk._exact_batch_invariant_router + assert blk.router._exact_batch_invariant + + x = torch.randn(1, 4, 32, dtype=torch.bfloat16) + routing = torch.zeros(4, 2, dtype=torch.float32) + selected = torch.zeros(4, 2, dtype=torch.int64) + with pytest.raises(RuntimeError, match="trainer EP"): + blk._ep_combine_native(x, routing, selected) + + with monkeypatch.context() as rows_patch: + _assert_variable_row_collectives_pad_tokens_and_ids_and_share_max_rows(rows_patch) + with monkeypatch.context() as gate_patch: + _assert_serving_fused_gate_forward_preserves_trainer_gradients(gate_patch) + with monkeypatch.context() as dispatch_patch: + _assert_native_combine_dispatch_and_actual_operand_policy(dispatch_patch) + def _qwen_block(*, exact: bool = False): from transformers import PretrainedConfig # noqa: PLC0415 @@ -43,23 +61,7 @@ def _qwen_block(*, exact: bool = False): return Qwen3_5MoeSparseMoeBlock(cfg, moe_implementation="eager", layer_idx=0).to(torch.bfloat16) -def test_exact_native_requires_trainer_ep(): - blk = _qwen_block(exact=True) - x = torch.randn(1, 4, 32, dtype=torch.bfloat16) - routing = torch.zeros(4, 2, dtype=torch.float32) - selected = torch.zeros(4, 2, dtype=torch.int64) - with pytest.raises(RuntimeError, match="trainer EP"): - blk._ep_combine_native(x, routing, selected) - - -def test_exact_native_combine_is_structural(): - blk = _qwen_block(exact=True) - assert blk._native_ep_combine - assert blk._exact_batch_invariant_router - assert blk.router._exact_batch_invariant - - -def test_native_routed_partial_enters_through_module_call(monkeypatch): +def _assert_native_routed_partial_enters_through_module_call(monkeypatch): """The EP serving-kernel lane must run inside FSDP's pre-forward hooks.""" blk = _qwen_block() hidden = torch.randn(4, 32, dtype=torch.bfloat16) @@ -86,7 +88,7 @@ def routed_partial(got_hidden, got_routing, got_ids): assert torch.equal(result, torch.zeros_like(hidden)) -def test_variable_row_token_gather_unpads_backward(monkeypatch): +def _assert_variable_row_collectives_pad_tokens_and_ids_and_share_max_rows(monkeypatch): """The live River packer gives EP ranks unequal T; collectives must not.""" import xorl.models.layers.moe.ep_native_combine as combine # noqa: PLC0415 @@ -112,26 +114,16 @@ def fake_reduce_scatter(out, grad, op=None, group=None): gathered.sum().backward() assert torch.equal(x.grad, torch.full_like(x, 2.0)) - -def test_variable_row_id_gather_uses_invalid_padding(monkeypatch): - import xorl.models.layers.moe.ep_native_combine as combine # noqa: PLC0415 - - monkeypatch.setattr(combine.dist, "get_world_size", lambda _group: 2) - - def fake_gather(out, local, group=None): + def fake_id_gather(out, local, group=None): del group assert torch.equal(local, torch.tensor([[4, 5], [-1, -1], [-1, -1]])) out[:3].copy_(local) out[3:].copy_(local) - monkeypatch.setattr(combine.dist, "all_gather_into_tensor", fake_gather) - gathered = gather_ids_for_ep_combine(torch.tensor([[4, 5]]), group=None, padded_rows=3) - assert gathered.shape == (6, 2) - assert torch.equal(gathered[1:3], torch.full((2, 2), -1)) - - -def test_max_rows_for_ep_combine(monkeypatch): - import xorl.models.layers.moe.ep_native_combine as combine # noqa: PLC0415 + monkeypatch.setattr(combine.dist, "all_gather_into_tensor", fake_id_gather) + gathered_ids = gather_ids_for_ep_combine(torch.tensor([[4, 5]]), group=None, padded_rows=3) + assert gathered_ids.shape == (6, 2) + assert torch.equal(gathered_ids[1:3], torch.full((2, 2), -1)) def fake_max(rows, op=None, group=None): del op, group @@ -141,7 +133,7 @@ def fake_max(rows, op=None, group=None): assert max_rows_for_ep_combine(6016, torch.device("cpu"), group=None) == 8192 -def test_serving_fused_gate_forward_preserves_trainer_gradients(monkeypatch): +def _assert_serving_fused_gate_forward_preserves_trainer_gradients(monkeypatch): import xorl.models.layers.moe.ep_native_combine as combine # noqa: PLC0415 def fake_serving_kernel(hidden, weight, shared, final): @@ -168,8 +160,11 @@ def fake_serving_kernel(hidden, weight, shared, final): torch.testing.assert_close(actual_grad, reference_input.grad) -def test_native_combine_captures_actual_operands(monkeypatch): +def _assert_native_combine_dispatch_and_actual_operand_policy(monkeypatch): """The layer-selected diagnostic hook exposes every exact-combine boundary.""" + with monkeypatch.context() as routed_partial_monkeypatch: + _assert_native_routed_partial_enters_through_module_call(routed_partial_monkeypatch) + import xorl.distributed.parallel_state as parallel_state # noqa: PLC0415 import xorl.models.layers.moe.ep_native_combine as combine # noqa: PLC0415 import xorl.ops.batch_invariant_ops as batch_invariant_ops # noqa: PLC0415 @@ -199,8 +194,9 @@ def forward(self, hidden, routing, *, sglang_ep_native_local_ids): monkeypatch.setattr( combine, "sglang_fused_gate_sigmoid_mul_add", - lambda hidden, weight, shared, routed: routed - + torch.sigmoid((hidden * weight).sum(dim=-1, keepdim=True)) * shared, + lambda hidden, weight, shared, routed: ( + routed + torch.sigmoid((hidden * weight).sum(dim=-1, keepdim=True)) * shared + ), ) monkeypatch.setattr( batch_invariant_ops._BatchInvariantTrunkLinearFn, diff --git a/tests/models/test_moe_experts_lora.py b/tests/models/test_moe_experts_lora.py index 4ce63a39..22531fe9 100644 --- a/tests/models/test_moe_experts_lora.py +++ b/tests/models/test_moe_experts_lora.py @@ -7,20 +7,15 @@ import torch import torch.nn as nn -from xorl.lora import LoraLinear, inject_lora_into_model -from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module -from xorl.models.layers.moe import MOE_EXPERT_BACKENDS, MoEBlock, MoEExperts, MoEExpertsLoRA, MoELoRAConfig +from xorl.lora import inject_lora_into_model +from xorl.models.layers.moe import MoEBlock, MoEExperts, MoEExpertsLoRA, MoELoRAConfig from xorl.models.layers.moe.backend import zero_token_lora_output -from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( - Qwen3MoeSparseExperts, - Qwen3MoeTritonExperts, -) pytestmark = [pytest.mark.cpu, pytest.mark.gpu] -def test_zero_token_output_materializes_structural_gradients_for_every_local_factor(): +def _assert_zero_token_output_materializes_structural_gradients_for_every_local_factor(): tokens = torch.empty(0, 4, requires_grad=True) factors = tuple(nn.Parameter(torch.randn(shape)) for shape in ((1, 4, 2), (2, 2, 6), (2, 6, 2))) @@ -55,139 +50,57 @@ def __init__( # --------------------------------------------------------------------------- -# 1. Base MoE experts: init, shapes, backends, LoRA mapping, registry -# --------------------------------------------------------------------------- - - -class TestMoEExpertsBase: - """Comprehensive tests for base MoEExperts initialization, shapes, and registration.""" - - def test_init_and_shapes_all_backends(self): - """Test init fields, weight shapes, and LoRA mapping for all backends and subclasses.""" - config = MockConfig() - - # Qwen3 subclass inits - triton_exp = Qwen3MoeTritonExperts(config) - assert triton_exp.num_experts == config.num_experts - assert triton_exp.hidden_dim == config.hidden_size - assert triton_exp.intermediate_size == config.moe_intermediate_size - assert triton_exp.moe_implementation == "triton" - - sparse_exp = Qwen3MoeSparseExperts(config) - assert sparse_exp.moe_implementation == "eager" - - # Direct MoEExperts with all backends - for backend in ["eager", "triton", "native", "quack"]: - experts = MoEExperts( - num_experts=config.num_experts, - hidden_dim=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - moe_implementation=backend, - ) - assert experts.moe_implementation == backend - - # Weight shapes on a single instance - experts = MoEExperts( - num_experts=config.num_experts, - hidden_dim=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - ) - assert experts.gate_proj.shape == (config.num_experts, config.hidden_size, config.moe_intermediate_size) - assert experts.up_proj.shape == (config.num_experts, config.hidden_size, config.moe_intermediate_size) - assert experts.down_proj.shape == (config.num_experts, config.moe_intermediate_size, config.hidden_size) - - # LoRA mapping registered - assert can_apply_lora(experts) - assert get_lora_class_for_module(experts) is MoEExpertsLoRA - - # Backend registry - assert "eager" in MOE_EXPERT_BACKENDS - assert "fused" not in MOE_EXPERT_BACKENDS - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for triton MoE") - def test_triton_forward(self): - """Test forward pass with triton backend on GPU.""" - config = MockConfig() - experts = Qwen3MoeTritonExperts(config) - nn.init.xavier_normal_(experts.gate_proj.data) - nn.init.xavier_normal_(experts.up_proj.data) - nn.init.xavier_normal_(experts.down_proj.data) - - device = "cuda" - experts = experts.to(device).to(torch.bfloat16) - - num_tokens, top_k = 16, 2 - hidden_states = torch.randn(num_tokens, config.hidden_size, device=device, dtype=torch.bfloat16) - routing_weights = torch.softmax(torch.randn(num_tokens, top_k, device=device, dtype=torch.bfloat16), dim=-1) - selected_experts = torch.randint(0, config.num_experts, (num_tokens, top_k), device=device) - - output = experts(hidden_states, routing_weights, selected_experts) - assert output.shape == hidden_states.shape - - -# --------------------------------------------------------------------------- -# 2. LoRA initialization: all backends, frozen/trainable, shapes, repr +# LoRA initialization: all backends, frozen/trainable, shapes, repr # --------------------------------------------------------------------------- class TestMoEExpertsLoRAInit: """Comprehensive tests for MoEExpertsLoRA initialization across backends.""" - @pytest.mark.parametrize("backend", ["eager", "triton", "native", "quack"]) - def test_init_frozen_trainable_shapes(self, backend): + def _assert_init_frozen_trainable_shapes(self): """Test init, base weights frozen, LoRA weights trainable, shapes, and repr.""" config = MockConfig() lora_config = MoELoRAConfig(r=8, lora_alpha=16, target_modules=["gate_proj", "up_proj", "down_proj"]) r = lora_config.r - experts = MoEExpertsLoRA( - num_experts=config.num_experts, - hidden_dim=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - moe_implementation=backend, - lora_config=lora_config, - ) + for backend in ("eager", "triton", "native", "quack"): + experts = MoEExpertsLoRA( + num_experts=config.num_experts, + hidden_dim=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + moe_implementation=backend, + lora_config=lora_config, + ) - # Init fields - assert experts.num_experts == config.num_experts - assert experts.moe_implementation == backend - assert experts.lora_config == lora_config - - # Base weights frozen - assert not experts.gate_proj.requires_grad - assert not experts.up_proj.requires_grad - assert not experts.down_proj.requires_grad - - # LoRA weights trainable, B initialized to zeros, correct shapes - for name in lora_config.target_modules: - lora_A = getattr(experts, f"{name}_lora_A") - lora_B = getattr(experts, f"{name}_lora_B") - assert isinstance(lora_A, nn.Parameter) and lora_A.requires_grad - assert isinstance(lora_B, nn.Parameter) and lora_B.requires_grad - assert torch.allclose(lora_B, torch.zeros_like(lora_B)) - - # LoRA weight shapes - assert experts.gate_proj_lora_A.shape == (config.num_experts, config.hidden_size, r) - assert experts.gate_proj_lora_B.shape == (config.num_experts, r, config.moe_intermediate_size) - assert experts.up_proj_lora_A.shape == (config.num_experts, config.hidden_size, r) - assert experts.up_proj_lora_B.shape == (config.num_experts, r, config.moe_intermediate_size) - assert experts.down_proj_lora_A.shape == (config.num_experts, config.moe_intermediate_size, r) - assert experts.down_proj_lora_B.shape == (config.num_experts, r, config.hidden_size) - - # Shapes match between eager and triton (already tested by parametrize) - repr_str = experts.extra_repr() - assert f"num_experts={config.num_experts}" in repr_str - assert f"r={lora_config.r}" in repr_str - - def test_runtime_rank_lora_views_are_contiguous(self): - """Partial-rank views passed to group GEMM backends must be contiguous.""" - config = MockConfig() - experts = MoEExpertsLoRA( - num_experts=config.num_experts, - hidden_dim=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - lora_config=MoELoRAConfig(r=8, lora_alpha=16), - ) + # Init fields + assert experts.num_experts == config.num_experts + assert experts.moe_implementation == backend + assert experts.lora_config == lora_config + + # Base weights frozen + assert not experts.gate_proj.requires_grad + assert not experts.up_proj.requires_grad + assert not experts.down_proj.requires_grad + + # LoRA weights trainable, B initialized to zeros, correct shapes + for name in lora_config.target_modules: + lora_A = getattr(experts, f"{name}_lora_A") + lora_B = getattr(experts, f"{name}_lora_B") + assert isinstance(lora_A, nn.Parameter) and lora_A.requires_grad + assert isinstance(lora_B, nn.Parameter) and lora_B.requires_grad + assert torch.allclose(lora_B, torch.zeros_like(lora_B)) + + # LoRA weight shapes + assert experts.gate_proj_lora_A.shape == (config.num_experts, config.hidden_size, r) + assert experts.gate_proj_lora_B.shape == (config.num_experts, r, config.moe_intermediate_size) + assert experts.up_proj_lora_A.shape == (config.num_experts, config.hidden_size, r) + assert experts.up_proj_lora_B.shape == (config.num_experts, r, config.moe_intermediate_size) + assert experts.down_proj_lora_A.shape == (config.num_experts, config.moe_intermediate_size, r) + assert experts.down_proj_lora_B.shape == (config.num_experts, r, config.hidden_size) + + repr_str = experts.extra_repr() + assert f"num_experts={config.num_experts}" in repr_str + assert f"r={lora_config.r}" in repr_str experts.set_runtime_lora_config(lora_rank=3, lora_alpha=12) @@ -198,6 +111,8 @@ def test_runtime_rank_lora_views_are_contiguous(self): assert lora_A.is_contiguous() assert lora_B.is_contiguous() + TestFromModuleAndInjection()._assert_from_module_and_inject_lora() + # --------------------------------------------------------------------------- # 3. Eager LoRA forward/backward (CPU) @@ -209,6 +124,7 @@ class TestMoEExpertsLoRAEager: def test_eager_forward_backward_and_moe_block(self): """Test eager per-expert forward, backward gradients, and end-to-end MoEBlock.""" + TestMoEExpertsLoRAInit()._assert_init_frozen_trainable_shapes() lora_config = MoELoRAConfig(r=4, lora_alpha=8) experts = MoEExpertsLoRA( num_experts=4, @@ -257,8 +173,6 @@ def test_eager_forward_backward_and_moe_block(self): output.sum().backward() assert block.experts.gate_proj_lora_A.grad is not None - def test_hybrid_shared_shapes(self): - """Hybrid-shared injection keeps the supported shared tensor layout.""" block = MoEBlock( hidden_size=32, num_experts=4, @@ -276,79 +190,12 @@ def test_hybrid_shared_shapes(self): assert block.experts.down_proj_lora_A.shape == (4, 64, 4) assert block.experts.down_proj_lora_B.shape == (1, 4, 32) - -# --------------------------------------------------------------------------- -# 4. GPU LoRA forward/backward (triton + native) -# --------------------------------------------------------------------------- - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -class TestMoEExpertsLoRAGPU: - """Test triton and native LoRA forward/backward on GPU, including via MoEBlock.""" - - @pytest.mark.parametrize("backend", ["triton", "native"]) - def test_forward_backward(self, backend): - """Test GPU LoRA forward and backward for triton and native backends.""" - lora_config = MoELoRAConfig(r=4, lora_alpha=8) - exp = MoEExpertsLoRA( - num_experts=4, - hidden_dim=32, - intermediate_size=64, - moe_implementation=backend, - lora_config=lora_config, - ) - nn.init.xavier_normal_(exp.gate_proj.data) - nn.init.xavier_normal_(exp.up_proj.data) - nn.init.xavier_normal_(exp.down_proj.data) - - device = "cuda" - exp = exp.to(device).to(torch.bfloat16) - - num_tokens, top_k = 16, 2 - hidden = torch.randn(num_tokens, 32, device=device, dtype=torch.bfloat16, requires_grad=True) - weights = torch.softmax(torch.randn(num_tokens, top_k, device=device, dtype=torch.bfloat16), dim=-1) - selected = torch.randint(0, 4, (num_tokens, top_k), device=device) - - # Forward - output = exp(hidden, weights, selected) - assert output.shape == hidden.shape - - # Backward - output.sum().backward() - assert exp.gate_proj_lora_A.grad is not None - assert exp.down_proj_lora_B.grad is not None - assert exp.gate_proj.grad is None # base frozen - - def test_native_via_moe_block(self): - """Test native LoRA works end-to-end through MoEBlock.""" - device = "cuda" - block = MoEBlock( - hidden_size=32, - num_experts=4, - top_k=2, - intermediate_size=64, - moe_implementation="native", - ) - nn.init.xavier_normal_(block.experts.gate_proj.data) - nn.init.xavier_normal_(block.experts.up_proj.data) - nn.init.xavier_normal_(block.experts.down_proj.data) - nn.init.xavier_normal_(block.gate.weight.data) - block = block.to(device).to(torch.bfloat16) - - block.inject_lora(r=4, lora_alpha=8) - assert isinstance(block.experts, MoEExpertsLoRA) - assert block.experts.moe_implementation == "native" - - hidden = torch.randn(2, 4, 32, device=device, dtype=torch.bfloat16) - output, router_logits = block(hidden) - assert output.shape == hidden.shape - - output.sum().backward() - assert block.experts.gate_proj_lora_A.grad is not None + _assert_zero_token_output_materializes_structural_gradients_for_every_local_factor() + TestEPLoRARouterScores()._assert_router_score_application_contract() # --------------------------------------------------------------------------- -# 5. Cross-backend numerical correctness +# 4. Cross-backend numerical correctness # --------------------------------------------------------------------------- @@ -433,13 +280,17 @@ def _make_pair(self, ref_backend, test_backend, device): device, self.DTYPE, ) + torch.manual_seed(321) + with torch.no_grad(): + for proj in ("gate_proj", "up_proj", "down_proj"): + nn.init.xavier_normal_(getattr(ref.experts, f"{proj}_lora_B")) _copy_block_weights(ref, test) return ref, test - @pytest.mark.parametrize("backend", ["eager", "triton", "native"]) - def test_zero_lora_matches_base(self, backend): + def test_zero_lora_and_cross_backend_numerical_policy(self): """With lora_B=0, LoRA output must equal base model output (no delta).""" device = "cuda" + backend = "eager" # Base block (no LoRA) base_block = MoEBlock( hidden_size=self.HIDDEN_DIM, @@ -485,15 +336,9 @@ def test_zero_lora_matches_base(self, backend): msg=f"[{backend}] Zero-LoRA output should match base model", ) - @pytest.mark.parametrize( - "ref_backend,test_backend", - [ - ("eager", "native"), - ("eager", "triton"), - ("triton", "native"), - ], - ) - def test_cross_backend_output_and_gradients(self, ref_backend, test_backend): + self._assert_cross_backend_output_and_gradients_policy() + + def _assert_cross_backend_output_and_gradients(self, ref_backend, test_backend): """Cross-backend outputs and LoRA gradients should match.""" ref, test = self._make_pair(ref_backend, test_backend, "cuda") @@ -541,65 +386,23 @@ def test_cross_backend_output_and_gradients(self, ref_backend, test_backend): msg=f"Gradient mismatch: {proj}_lora_B ({ref_backend} vs {test_backend})", ) - @pytest.mark.parametrize("backend", ["eager", "triton", "native"]) - def test_nonzero_lora_changes_output(self, backend): - """Non-zero LoRA weights must produce a different output from base.""" - device = "cuda" - block = _make_lora_block( - backend, - self.NUM_EXPERTS, - self.HIDDEN_DIM, - self.INTERMEDIATE, - self.R, - self.LORA_ALPHA, - device, - self.DTYPE, - ) - # Set lora_B to non-zero - with torch.no_grad(): - for proj in ["gate_proj", "up_proj", "down_proj"]: - lora_B = getattr(block.experts, f"{proj}_lora_B") - nn.init.xavier_normal_(lora_B) - - # Base block (no LoRA) with same base weights - base_block = ( - MoEBlock( - hidden_size=self.HIDDEN_DIM, - num_experts=self.NUM_EXPERTS, - top_k=2, - intermediate_size=self.INTERMEDIATE, - moe_implementation=backend, - ) - .to(device) - .to(self.DTYPE) - ) - with torch.no_grad(): - base_block.gate.weight.copy_(block.gate.weight) - base_block.experts.gate_proj.copy_(block.experts.gate_proj) - base_block.experts.up_proj.copy_(block.experts.up_proj) - base_block.experts.down_proj.copy_(block.experts.down_proj) - - torch.manual_seed(999) - hidden = torch.randn(2, 8, self.HIDDEN_DIM, device=device, dtype=self.DTYPE) - base_out, _ = base_block(hidden) - lora_out, _ = block(hidden) - - diff = (lora_out - base_out).abs().max().item() - assert diff > 1e-3, f"[{backend}] Non-zero LoRA should change the output, but max diff={diff}" + def _assert_cross_backend_output_and_gradients_policy(self): + for ref_backend, test_backend in (("eager", "native"), ("eager", "triton")): + self._assert_cross_backend_output_and_gradients(ref_backend, test_backend) # --------------------------------------------------------------------------- -# 6. from_module + LoRA injection + error handling +# 5. from_module + LoRA injection + error handling # --------------------------------------------------------------------------- class TestFromModuleAndInjection: """Test from_module, inject_lora, and error handling.""" - @pytest.mark.parametrize("backend", ["eager", "triton", "native", "quack"]) - def test_from_module_and_inject_lora(self, backend): + def _assert_from_module_and_inject_lora(self): """Test from_module preserves backend/weights, and inject_lora works via both APIs.""" config = MockConfig() + backend = "quack" # from_module base = MoEExperts( @@ -650,62 +453,9 @@ def __init__(self, config, backend): assert block.experts.moe_implementation == backend assert block.lora_adapter == "injected" - def test_from_module_with_qwen3_subclass(self): - """Test from_module works with Qwen3MoeTritonExperts (MoEExperts subclass).""" - config = MockConfig() - base = Qwen3MoeTritonExperts(config) - nn.init.xavier_normal_(base.gate_proj.data) - nn.init.xavier_normal_(base.up_proj.data) - nn.init.xavier_normal_(base.down_proj.data) - - lora_exp = MoEExpertsLoRA.from_module(base, r=8, lora_alpha=16) - assert isinstance(lora_exp, MoEExpertsLoRA) - assert torch.allclose(lora_exp.gate_proj, base.gate_proj) - - def test_injection_error_handling(self): - """Test error cases and valid injection for inject_lora_into_model.""" - - # No matching modules - class ModelA(nn.Module): - def __init__(self): - super().__init__() - self.layer1 = nn.Linear(64, 64) - - with pytest.raises(ValueError, match="No modules found matching target_modules"): - inject_lora_into_model(ModelA(), r=8, lora_alpha=16, target_modules=["nonexistent_proj"]) - - # Matched modules without LoRA support - class UnsupportedModule(nn.Module): - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.randn(64, 64)) - - def forward(self, x): - return x @ self.weight - - class ModelB(nn.Module): - def __init__(self): - super().__init__() - self.custom_layer = UnsupportedModule() - - with pytest.raises(ValueError, match="No modules found matching target_modules"): - inject_lora_into_model(ModelB(), r=8, lora_alpha=16, target_modules=["custom_layer"]) - - # Valid target - class ModelC(nn.Module): - def __init__(self): - super().__init__() - self.q_proj = nn.Linear(64, 64) - self.v_proj = nn.Linear(64, 64) - - model = ModelC() - inject_lora_into_model(model, r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"]) - assert isinstance(model.q_proj, LoraLinear) - assert isinstance(model.v_proj, LoraLinear) - # --------------------------------------------------------------------------- -# 7. EP LoRA router-score application +# 6. EP LoRA router-score application # --------------------------------------------------------------------------- @@ -786,24 +536,27 @@ def mock_combine(**kwargs): return output, captured["expert_output"] - @pytest.mark.parametrize("score_attr", ["expert_scores", "permuted_scores"]) - def test_scores_applied_to_output(self, score_attr): - """Router scores from dispatch context are multiplied into expert output.""" - experts = self._make_experts() - compute_output = torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM) - scores = torch.rand(self.NUM_TOKENS) * 0.5 + 0.1 # non-trivial scores in (0.1, 0.6) + def _assert_router_score_application_contract(self): + """Router scores apply when present, preserve identity when absent, and remain differentiable.""" + for score_attr in ("expert_scores", "permuted_scores"): + experts = self._make_experts() + compute_output = torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM) + scores = torch.rand(self.NUM_TOKENS) * 0.5 + 0.1 # non-trivial scores in (0.1, 0.6) - _, expert_output = self._run_ep_forward( - experts, - score_attr=score_attr, - scores=scores, - compute_output=compute_output, - ) + _, expert_output = self._run_ep_forward( + experts, + score_attr=score_attr, + scores=scores, + compute_output=compute_output, + ) + + expected = compute_output * scores.unsqueeze(1) + torch.testing.assert_close(expert_output, expected) - expected = compute_output * scores.unsqueeze(1) - torch.testing.assert_close(expert_output, expected) + self._assert_no_scores_leave_output_unchanged() + self._assert_gradients_flow_through_scores() - def test_no_scores_leaves_output_unchanged(self): + def _assert_no_scores_leave_output_unchanged(self): """When dispatch context has no score attribute, expert output is unchanged.""" experts = self._make_experts() compute_output = torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM) @@ -816,7 +569,7 @@ def test_no_scores_leaves_output_unchanged(self): torch.testing.assert_close(expert_output, compute_output) - def test_gradient_flows_through_scores(self): + def _assert_gradients_flow_through_scores(self): """Gradients from the score multiplication reach LoRA parameters.""" experts = self._make_experts() compute_output = torch.randn(self.NUM_TOKENS, self.HIDDEN_DIM, requires_grad=True) @@ -860,7 +613,3 @@ def mock_combine(**kwargs): # scores.grad should equal the sum of (compute_output * grad_output) per token expected_score_grad = (compute_output.detach() * 1.0).sum(dim=1) # grad_output is all 1s from .sum() torch.testing.assert_close(scores.grad, expected_score_grad) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/models/test_moe_fused_gate_up_proj.py b/tests/models/test_moe_fused_gate_up_proj.py index 2570eb02..d300a594 100644 --- a/tests/models/test_moe_fused_gate_up_proj.py +++ b/tests/models/test_moe_fused_gate_up_proj.py @@ -14,7 +14,7 @@ pytestmark = [pytest.mark.cpu] -def test_moe_experts_register_only_fused_gate_up_proj(): +def _assert_moe_experts_register_fused_gate_up_for_base_and_lora(): experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=5, moe_implementation="eager") named_params = dict(experts.named_parameters()) @@ -31,9 +31,7 @@ def test_moe_experts_register_only_fused_gate_up_proj(): torch.testing.assert_close(experts.gate_up_proj[..., :5], torch.ones(3, 4, 5)) torch.testing.assert_close(experts.gate_up_proj[..., 5:], torch.full((3, 4, 5), 2.0)) - -def test_moe_experts_lora_registers_fused_base_weight(): - experts = MoEExpertsLoRA( + lora_experts = MoEExpertsLoRA( num_experts=3, hidden_dim=4, intermediate_size=5, @@ -41,16 +39,16 @@ def test_moe_experts_lora_registers_fused_base_weight(): lora_config=MoELoRAConfig(r=2, lora_alpha=4), ) - named_params = dict(experts.named_parameters()) + named_params = dict(lora_experts.named_parameters()) assert "gate_up_proj" in named_params assert "gate_proj" not in named_params assert "up_proj" not in named_params assert named_params["gate_up_proj"].shape == (3, 4, 10) - assert experts.gate_proj.shape == (3, 4, 5) - assert experts.up_proj.shape == (3, 4, 5) + assert lora_experts.gate_proj.shape == (3, 4, 5) + assert lora_experts.up_proj.shape == (3, 4, 5) -def test_qwen3_moe_checkpoint_handler_round_trips_fused_experts(): +def _assert_qwen_moe_checkpoint_handler_fused_expert_contract(): hidden_size = 4 intermediate_size = 3 handler = Qwen3MoeCheckpointHandler( @@ -101,8 +99,11 @@ def test_qwen3_moe_checkpoint_handler_round_trips_fused_experts(): torch.testing.assert_close(saved["model.layers.0.mlp.experts.0.down_proj.weight"], down_0) torch.testing.assert_close(saved["model.layers.0.mlp.experts.1.down_proj.weight"], down_1) + _assert_qwen3_5_moe_checkpoint_handler_round_trips_fused_experts() + _assert_qwen_moe_checkpoint_handlers_skip_deferred_qlora_expert_loading() -def test_qwen3_5_moe_checkpoint_handler_round_trips_fused_experts(): + +def _assert_qwen3_5_moe_checkpoint_handler_round_trips_fused_experts(): hidden_size = 4 intermediate_size = 3 gate_up_weight = torch.arange(0, 48, dtype=torch.float32).view(2, 2 * intermediate_size, hidden_size) @@ -139,8 +140,8 @@ def test_qwen3_5_moe_checkpoint_handler_round_trips_fused_experts(): torch.testing.assert_close(saved["model.layers.0.mlp.experts.1.down_proj.weight"], down_weight[1]) -def test_qwen3_moe_checkpoint_handler_skips_deferred_qlora_expert_loading(): - handler = Qwen3MoeCheckpointHandler( +def _assert_qwen_moe_checkpoint_handlers_skip_deferred_qlora_expert_loading(): + qwen3_handler = Qwen3MoeCheckpointHandler( num_experts=2, num_attention_heads=2, num_key_value_heads=1, @@ -148,7 +149,7 @@ def test_qwen3_moe_checkpoint_handler_skips_deferred_qlora_expert_loading(): skip_expert_loading=True, ) - skip_fn = handler.get_skip_key_fn() + skip_fn = qwen3_handler.get_skip_key_fn() assert skip_fn is not None assert skip_fn("model.layers.0.mlp.experts.0.gate_proj.weight") assert skip_fn("model.layers.0.mlp.experts.0.up_proj.weight") @@ -156,12 +157,10 @@ def test_qwen3_moe_checkpoint_handler_skips_deferred_qlora_expert_loading(): assert skip_fn("model.layers.0.mlp.experts.gate_up_proj") assert skip_fn("model.layers.0.mlp.experts.down_proj") - assert handler.on_load_weight("model.layers.0.mlp.experts.0.gate_proj.weight", torch.randn(3, 4)) == [] - assert handler.on_load_weight("model.layers.0.mlp.experts.gate_up_proj", torch.randn(2, 4, 6)) == [] + assert qwen3_handler.on_load_weight("model.layers.0.mlp.experts.0.gate_proj.weight", torch.randn(3, 4)) == [] + assert qwen3_handler.on_load_weight("model.layers.0.mlp.experts.gate_up_proj", torch.randn(2, 4, 6)) == [] - -def test_qwen3_5_moe_checkpoint_handler_skips_deferred_qlora_expert_loading(): - handler = Qwen3_5MoeCheckpointHandler( + qwen35_handler = Qwen3_5MoeCheckpointHandler( num_experts=2, num_attention_heads=2, num_key_value_heads=1, @@ -171,7 +170,7 @@ def test_qwen3_5_moe_checkpoint_handler_skips_deferred_qlora_expert_loading(): skip_expert_loading=True, ) - skip_fn = handler.get_skip_key_fn() + skip_fn = qwen35_handler.get_skip_key_fn() assert skip_fn is not None assert skip_fn("model.layers.0.mlp.experts.0.gate_proj.weight") assert skip_fn("model.layers.0.mlp.experts.0.up_proj.weight") @@ -179,8 +178,8 @@ def test_qwen3_5_moe_checkpoint_handler_skips_deferred_qlora_expert_loading(): assert skip_fn("model.layers.0.mlp.experts.gate_up_proj.weight") assert skip_fn("model.layers.0.mlp.experts.down_proj.weight") - assert handler.on_load_weight("model.layers.0.mlp.experts.0.gate_proj.weight", torch.randn(3, 4)) == [] - assert handler.on_load_weight("model.layers.0.mlp.experts.gate_up_proj.weight", torch.randn(2, 6, 4)) == [] + assert qwen35_handler.on_load_weight("model.layers.0.mlp.experts.0.gate_proj.weight", torch.randn(3, 4)) == [] + assert qwen35_handler.on_load_weight("model.layers.0.mlp.experts.gate_up_proj.weight", torch.randn(2, 6, 4)) == [] def test_save_model_weights_drops_qarl_buffers_and_unfuses_qkv(tmp_path): @@ -190,6 +189,8 @@ def test_save_model_weights_drops_qarl_buffers_and_unfuses_qkv(tmp_path): 1-dimensional tensor``). ``save_model_weights`` must drop all ``qarl_*`` buffers before the handler runs, and still unfuse ``qkv_proj`` -> q/k/v. """ + _assert_moe_experts_register_fused_gate_up_for_base_and_lora() + _assert_qwen_moe_checkpoint_handler_fused_expert_contract() n_heads, n_kv, head_dim, hidden = 2, 1, 2, 4 q_dim, kv_dim = n_heads * head_dim, n_kv * head_dim handler = Qwen3MoeCheckpointHandler( diff --git a/tests/models/test_moe_routing_replay.py b/tests/models/test_moe_routing_replay.py index 8c5ec6df..78df78cd 100644 --- a/tests/models/test_moe_routing_replay.py +++ b/tests/models/test_moe_routing_replay.py @@ -13,10 +13,6 @@ from torch.utils.checkpoint import checkpoint from xorl.models.base import XorlPreTrainedModel - - -pytestmark = [pytest.mark.gpu] - from xorl.models.layers.moe.moe_block import MoEBlock from xorl.models.layers.moe.routing_replay import ( RoutingReplay, @@ -25,6 +21,9 @@ ) +pytestmark = [pytest.mark.gpu] + + # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -34,6 +33,10 @@ def _clean_routing_replay_state(): """Reset all RoutingReplay state between tests.""" yield + _reset_routing_replay_state() + + +def _reset_routing_replay_state(): set_replay_stage(None) RoutingReplay.clear_all() RoutingReplay._instances.clear() @@ -82,7 +85,6 @@ def _run_with_replay(layer, x): out = checkpoint(layer, x, use_reentrant=False) set_replay_stage("replay_backward") out.sum().backward() - RoutingReplay.reset_all_backward() return out @@ -94,9 +96,7 @@ def _run_with_replay(layer, x): class TestRoutingReplayUnit: """Unit tests for RoutingReplay class, stage management, clear, and registry.""" - @pytest.mark.gpu - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - def test_async_cpu_record_is_ordered_before_replay(self): + def _assert_async_cpu_record_is_ordered_before_replay(self): replay = RoutingReplay() source = torch.arange(16, device="cuda", dtype=torch.int64) record_stream = torch.cuda.Stream() @@ -111,7 +111,7 @@ def test_async_cpu_record_is_ordered_before_replay(self): def test_record_pop_dual_index_clear_stage_and_registry(self): """Test record, forward/backward pop, dual-index, CPU pinned storage, - clear, clear_all, reset_all_forward/backward, registry, and stage management.""" + clear, clear_all, registry, and stage management.""" # --- Record and pop --- replay = RoutingReplay() t1 = torch.tensor([[0, 1], [2, 3]]) @@ -184,24 +184,6 @@ def test_record_pop_dual_index_clear_stage_and_registry(self): assert len(r1.top_indices_list) == 0 assert len(r2.top_indices_list) == 0 - # reset_all_forward - r1.forward_index = 5 - r2.forward_index = 3 - r1.backward_index = 2 - RoutingReplay.reset_all_forward() - assert r1.forward_index == 0 - assert r2.forward_index == 0 - assert r1.backward_index == 2 # backward not touched - - # reset_all_backward - r1.backward_index = 5 - r2.backward_index = 3 - r1.forward_index = 2 - RoutingReplay.reset_all_backward() - assert r1.backward_index == 0 - assert r2.backward_index == 0 - assert r1.forward_index == 2 # forward not touched - # --- Stage management --- assert get_replay_stage() is None set_replay_stage("record") @@ -226,6 +208,8 @@ class TestMoEBlockReplay: def test_record_replay_no_replay_and_router_training(self): """Test record on forward, replay on backward, no-recording conditions, router training/detach, and regather correctness.""" + TestRoutingReplayUnit()._assert_async_cpu_record_is_ordered_before_replay() + _reset_routing_replay_state() # --- Record stores expert indices correctly --- moe = MoEBlock( hidden_size=64, @@ -256,7 +240,7 @@ def test_record_replay_no_replay_and_router_training(self): x2 = torch.randn(1, 8, 64, device="cuda", requires_grad=True) _run_with_replay(layer, x2) assert len(replay2.top_indices_list) == 1 - assert replay2.backward_index == 0 + assert replay2.backward_index == 1 # Checkpoint determinism with noisy attention layer3 = _SimpleDecoderLayer().cuda() @@ -363,6 +347,10 @@ def test_record_replay_no_replay_and_router_training(self): out, _ = moe6(x7) out.sum().backward() assert moe6.gate.weight.grad is None or moe6.gate.weight.grad.abs().sum() == 0 + _reset_routing_replay_state() + TestMultiLayerReplay()._assert_multi_layer_and_pp_schedule() + _reset_routing_replay_state() + TestBaseModelIntegrationAndR3Preload()._assert_enable_routing_replay_checkpoint_e2e_and_r3_preload() # =========================================================================== @@ -395,7 +383,7 @@ def forward(self, x): replays.append(replay) return model, replays - def test_multi_layer_and_pp_schedule(self): + def _assert_multi_layer_and_pp_schedule(self): """Test 3-layer 2-microbatch replay and 4-microbatch 1F1B PP schedule.""" # 3 layers, 2 micro-batches model, replays = self._make_model(num_layers=3) @@ -470,7 +458,7 @@ def test_multi_layer_and_pp_schedule(self): class TestBaseModelIntegrationAndR3Preload: """XorlPreTrainedModel integration and R3 preload (replay_forward).""" - def test_enable_routing_replay_checkpoint_e2e_and_r3_preload(self): + def _assert_enable_routing_replay_checkpoint_e2e_and_r3_preload(self): """Test enable_routing_replay, gradient_checkpointing_enable, full e2e, and R3 preload.""" class _FakeConfig: diff --git a/tests/models/test_moe_sglang_fp32_routing_backward.py b/tests/models/test_moe_sglang_fp32_routing_backward.py index b7435cd6..a0fcde02 100644 --- a/tests/models/test_moe_sglang_fp32_routing_backward.py +++ b/tests/models/test_moe_sglang_fp32_routing_backward.py @@ -6,8 +6,6 @@ import torch from xorl.models.layers.moe.experts import ( - _group_gemm_same_nk_fp32_accumulator, - _scale_moe_grad_by_fp32_routing, _SglangFusedExpertsEPTrainFunction, _SglangFusedExpertsTrainFunction, ) @@ -16,42 +14,6 @@ pytestmark = [pytest.mark.cpu] -def test_fp32_routing_scale_rounds_only_after_multiply(): - grad_output = torch.tensor([[1.234375]], dtype=torch.bfloat16) - routing = torch.tensor([2.99991e-5], dtype=torch.float32) - - expected = (grad_output.float() * routing[:, None]).to(torch.bfloat16) - old_bf16_score_path = grad_output * routing.to(torch.bfloat16)[:, None] - actual = _scale_moe_grad_by_fp32_routing(grad_output, routing) - - assert torch.equal(actual, expected) - assert not torch.equal(actual, old_bf16_score_path), "stimulus must detect a pre-multiply BF16 score cast" - - -def test_fp32_accumulator_helper_requests_fresh_fp32_output(): - seen = {} - - def fake_group_gemm_same_nk(*, a, b, cumsum_M, max_M, output_dtype): - seen.update(cumsum_M=cumsum_M, max_M=max_M, output_dtype=output_dtype) - return (a.float() @ b[0].float()).to(output_dtype) - - a = torch.tensor([[1.0, 1.0]], dtype=torch.bfloat16) - b = torch.tensor([[[1.0], [2**-8]]], dtype=torch.bfloat16) - cumsum = torch.tensor([1], dtype=torch.int32) - actual = _group_gemm_same_nk_fp32_accumulator( - fake_group_gemm_same_nk, - a=a, - b=b, - cumsum_M=cumsum, - max_M=1, - ) - - assert seen["output_dtype"] is torch.float32 - assert actual.dtype is torch.float32 - assert actual.item() == 1.00390625 - assert actual.item() != actual.to(torch.bfloat16).item() - - @pytest.fixture() def eager_grouped_moe(monkeypatch): """Replace CUDA bookkeeping/GEMMs with deterministic eager equivalents. @@ -177,8 +139,14 @@ def _assert_gradients_equal(actual, expected): assert torch.equal(left, right), f"{label} mismatch:\nactual={left}\nexpected={right}" -@pytest.mark.parametrize("filter_expert", [False, True]) -def test_local_backward_matches_fp32_routing_oracle_with_int32_ids(eager_grouped_moe, filter_expert): +def test_local_and_ep_backward_match_fp32_routing_oracle(eager_grouped_moe): + del eager_grouped_moe + for filter_expert in (False, True): + _assert_local_backward_matches_fp32_routing_oracle_with_int32_ids(filter_expert) + _assert_ep_backward_matches_fp32_routing_oracle() + + +def _assert_local_backward_matches_fp32_routing_oracle_with_int32_ids(filter_expert): hidden = torch.ones(2, 1, dtype=torch.bfloat16) routing = torch.tensor([[2.99991e-5, 0.5009], [0.33331, 0.77771]], dtype=torch.float32) selected = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32) @@ -216,7 +184,7 @@ def test_local_backward_matches_fp32_routing_oracle_with_int32_ids(eager_grouped assert custom[1].grad[0, 1].item() == 0.0 -def test_ep_backward_matches_fp32_routing_oracle(eager_grouped_moe): +def _assert_ep_backward_matches_fp32_routing_oracle(): hidden = torch.ones(4, 1, dtype=torch.bfloat16) routing = torch.tensor([2.99991e-5, 0.5009, 0.33331, 0.77771], dtype=torch.float32) selected = torch.tensor([[0], [0], [1], [1]], dtype=torch.int32) diff --git a/tests/models/test_moe_sglang_fused_experts.py b/tests/models/test_moe_sglang_fused_experts.py index ae0f70df..43412fef 100644 --- a/tests/models/test_moe_sglang_fused_experts.py +++ b/tests/models/test_moe_sglang_fused_experts.py @@ -24,7 +24,6 @@ explicit-flag path. """ -import inspect import logging import sys import types @@ -76,6 +75,10 @@ def routed_inputs(): @pytest.fixture() def auto_state(monkeypatch): """Reset the auto-resolution log latch and the cached stack probe.""" + _reset_auto_state(monkeypatch) + + +def _reset_auto_state(monkeypatch): monkeypatch.setattr(experts_mod, "_MOE_SGLANG_FUSED_EXPERTS_AUTO_LOGGED", False) monkeypatch.setattr(experts_mod, "_MOE_SGLANG_FUSED_EXPERTS_STACK_AVAILABLE", None) @@ -104,7 +107,7 @@ def log_spy(monkeypatch): return spy -def test_auto_resolution_ep1_enables_and_logs_once(monkeypatch, auto_state, log_spy): +def test_sglang_fused_experts_policy(monkeypatch, auto_state, log_spy, routed_inputs): monkeypatch.delenv(FLAG, raising=False) _stack_available(monkeypatch, True) assert experts_mod.moe_sglang_fused_experts_enabled(1, torch.device("cuda")) is True @@ -112,15 +115,34 @@ def test_auto_resolution_ep1_enables_and_logs_once(monkeypatch, auto_state, log_ logged = [m for _, m in log_spy.records if "auto-enabled (ep=1)" in m] assert len(logged) == 1, "auto resolution must be logged exactly once" - -def test_auto_resolution_ep_gt1_disables(monkeypatch, auto_state, log_spy): + for check in ( + _assert_auto_resolution_ep_gt1_disables, + _assert_explicit_env_overrides_auto, + _assert_auto_requires_cuda_input_and_stack, + ): + _reset_auto_state(monkeypatch) + log_spy.records.clear() + check(monkeypatch, log_spy) + _reset_auto_state(monkeypatch) + _assert_auto_requires_supported_expert_module(monkeypatch) + with monkeypatch.context() as case_patch: + _assert_moe_block_dispatch_and_precedence_policy(routed_inputs, case_patch) + with monkeypatch.context() as case_patch: + _assert_fused_expert_admission_and_trainable_dispatch_policy(routed_inputs, case_patch) + with monkeypatch.context() as case_patch: + _assert_weight_mode_cache_and_kernel_layout_policy(routed_inputs, case_patch) + with monkeypatch.context() as case_patch: + _assert_sglang_runtime_context_policy(case_patch) + + +def _assert_auto_resolution_ep_gt1_disables(monkeypatch, log_spy): monkeypatch.delenv(FLAG, raising=False) _stack_available(monkeypatch, True) assert experts_mod.moe_sglang_fused_experts_enabled(8, torch.device("cuda")) is False assert any("auto-disabled (ep=8)" in m for _, m in log_spy.records) -def test_explicit_env_overrides_auto(monkeypatch, auto_state, log_spy): +def _assert_explicit_env_overrides_auto(monkeypatch, log_spy): _stack_available(monkeypatch, True) monkeypatch.setenv(FLAG, "0") assert experts_mod.moe_sglang_fused_experts_enabled(1, torch.device("cuda")) is False @@ -129,7 +151,7 @@ def test_explicit_env_overrides_auto(monkeypatch, auto_state, log_spy): assert not log_spy.records, "explicit 1/0 must not emit the auto-resolution log" -def test_auto_requires_cuda_input_and_stack(monkeypatch, auto_state, log_spy): +def _assert_auto_requires_cuda_input_and_stack(monkeypatch, log_spy): monkeypatch.delenv(FLAG, raising=False) # CPU/meta inputs keep the stock path quietly (unit tests, tracing). _stack_available(monkeypatch, True) @@ -143,7 +165,7 @@ def test_auto_requires_cuda_input_and_stack(monkeypatch, auto_state, log_spy): assert warned and warned[0][0] == logging.WARNING -def test_auto_requires_supported_expert_module(monkeypatch, auto_state): +def _assert_auto_requires_supported_expert_module(monkeypatch): monkeypatch.delenv(FLAG, raising=False) _stack_available(monkeypatch, True) cuda = torch.device("cuda") @@ -169,7 +191,7 @@ def test_auto_requires_supported_expert_module(monkeypatch, auto_state): assert experts_mod.moe_sglang_fused_experts_enabled(1, cuda, object()) is False -def test_flag_off_keeps_default_path(routed_inputs, monkeypatch): +def _assert_moe_block_dispatch_and_precedence_policy(routed_inputs, monkeypatch): monkeypatch.delenv(FLAG, raising=False) blk = _block("eager") x, _, _ = routed_inputs @@ -187,8 +209,11 @@ def spy(*args, **kwargs): assert not called.get("hit") assert torch.equal(out, baseline) + _assert_forward_dispatches_sglang_path(routed_inputs, monkeypatch) + _assert_forward_experts_only_dispatches_sglang_path(routed_inputs, monkeypatch) + -def test_forward_dispatches_sglang_path_for_non_eager_backends(routed_inputs, monkeypatch): +def _assert_forward_dispatches_sglang_path(routed_inputs, monkeypatch): """The flag must override the configured backend (real configs resolve to triton).""" blk = _block("triton") x, w, ids = routed_inputs @@ -207,7 +232,7 @@ def fake(hidden_states, routing_weights, selected_experts): assert out.shape == (1, TOKENS, HID) -def test_forward_experts_only_dispatches_sglang_path(routed_inputs, monkeypatch): +def _assert_forward_experts_only_dispatches_sglang_path(routed_inputs, monkeypatch): blk = _block("triton") x, w, ids = routed_inputs called = {} @@ -223,28 +248,7 @@ def fake(hidden_states, routing_weights, selected_experts): assert out.shape == (1, TOKENS, HID) -def test_fp64_parity_mode_takes_precedence(routed_inputs, monkeypatch): - blk = _block("triton") - x, w, ids = routed_inputs - called = {} - - def fake_sglang(*args, **kwargs): - called["sglang"] = True - return torch.zeros_like(x) - - def fake_fp64(hidden_states, routing_weights, selected_experts): - called["fp64"] = True - return torch.zeros_like(hidden_states) - - monkeypatch.setattr(blk.experts, "sglang_fused_experts_forward", fake_sglang) - monkeypatch.setattr(blk, "_eager_forward_fp64", fake_fp64) - monkeypatch.setenv(FLAG, "1") - monkeypatch.setenv("XORL_MOE_FP64_ACCUM", "1") - blk.forward(x.view(1, TOKENS, HID)) - assert called.get("fp64") and not called.get("sglang") - - -def test_guards_reject_unsupported_experts(routed_inputs, monkeypatch): +def _assert_fused_expert_admission_and_trainable_dispatch_policy(routed_inputs, monkeypatch): monkeypatch.setenv(FLAG, "1") blk = _block("eager") x, w, ids = routed_inputs @@ -267,16 +271,19 @@ def test_guards_reject_unsupported_experts(routed_inputs, monkeypatch): with pytest.raises(ValueError): blk.experts.sglang_fused_experts_forward(x, None, None) + _assert_positive_swiglu_limit_fails_before_load(routed_inputs, monkeypatch) + _assert_trainable_guards(routed_inputs, monkeypatch) + _assert_trainable_dispatch_uses_autograd_function(routed_inputs, monkeypatch) -def test_positive_swiglu_limit_fails_before_sglang_load_or_kernel(routed_inputs, monkeypatch): - """Every SGLang fused-kernel/runner entry must reject XoRL's + +def _assert_positive_swiglu_limit_fails_before_load(routed_inputs, monkeypatch): + """Every supported SGLang fused-kernel entry must reject XoRL's semantically different clamp before importing or invoking SGLang.""" blk = _block("triton") blk.experts.swiglu_limit = 1.0 x, w, ids = routed_inputs cumsum = torch.arange(1, NUM_EXPERTS + 1, dtype=torch.int64) local_ids = ids.to(torch.int32) - parallel_state = SimpleNamespace(tp_size=1) loaded = [] def forbidden_loader(): @@ -284,14 +291,11 @@ def forbidden_loader(): raise AssertionError("positive swiglu_limit must fail before loading SGLang") monkeypatch.setattr(type(blk.experts), "_load_sglang_fused_experts_impl", staticmethod(forbidden_loader)) - monkeypatch.setattr(type(blk.experts), "_load_sglang_moe_runner_stack", staticmethod(forbidden_loader)) entrypoints = ( lambda: blk.experts.sglang_fused_experts_forward(x, w, ids), lambda: blk.experts.sglang_fused_experts_ep_compute(x[:NUM_EXPERTS], cumsum, w[:NUM_EXPERTS, 0]), lambda: blk.experts.sglang_ep_native_routed_partial(x, w, local_ids), - lambda: blk.experts._sglang_moe_tp_sim_sglang_forward(x, w, ids, parallel_state), - lambda: blk.experts._sglang_moe_tp_sim_sglang_runner_forward(x, w, ids, parallel_state), ) for entrypoint in entrypoints: with pytest.raises(NotImplementedError, match="positive swiglu_limit"): @@ -299,19 +303,7 @@ def forbidden_loader(): assert loaded == [] -def test_missing_sglang_raises_import_error_naming_flag(routed_inputs, monkeypatch): - import importlib.util # noqa: PLC0415 - - if importlib.util.find_spec("sglang") is not None: - pytest.skip("sglang installed; import-error guard not testable here") - monkeypatch.setenv(FLAG, "1") - blk = _block("eager") - x, w, ids = routed_inputs - with pytest.raises(ImportError, match=FLAG): - blk.experts.sglang_fused_experts_forward(x, w, ids) - - -def test_trainable_dispatch_uses_autograd_function(routed_inputs, monkeypatch): +def _assert_trainable_dispatch_uses_autograd_function(routed_inputs, monkeypatch): """When gradients are required, the flag path must route through the autograd Function; the no-grad path must keep using the plain kernel call.""" from xorl.models.layers.moe import experts as experts_mod # noqa: PLC0415 @@ -351,7 +343,7 @@ def fake_kernel_call(hidden_flat, *args, **kwargs): assert called == {"plain": True} -def test_trainable_guards(routed_inputs, monkeypatch): +def _assert_trainable_guards(routed_inputs, monkeypatch): blk = _block("eager") x, w, ids = routed_inputs monkeypatch.setenv(FLAG, "1") @@ -371,8 +363,15 @@ def test_trainable_guards(routed_inputs, monkeypatch): blk.experts.hidden_act = "silu" -@pytest.mark.gpu -def test_trainable_grads_match_stock_triton(monkeypatch): +def _assert_trainable_gradient_numerics_policy(monkeypatch): + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + with monkeypatch.context() as case_patch: + _assert_trainable_grads_match_stock_triton(case_patch) + _assert_masked_trainable_gradient_policy() + + +def _assert_trainable_grads_match_stock_triton(monkeypatch): """dX / dW13 / dW2 / d_topk_weights must be bit-identical to the stock triton path's gradients given the same drawn token permutation.""" if not torch.cuda.is_available(): @@ -446,8 +445,7 @@ def _masked_problem(device, seed=0, all_valid=False, all_masked=False): return x, wts, local_ids, gu, dn, e_local -@pytest.mark.gpu -def test_masked_trainable_grads_match_compacted_stock(): +def _assert_masked_trainable_gradient_policy(): """filter_expert=True grads must be bit-identical to the stock (unmasked) Function run directly on the compacted valid-pair topk=1 presentation, and masked slots' d(topk_weights) must be exactly zero.""" @@ -495,9 +493,11 @@ def test_masked_trainable_grads_match_compacted_stock(): dx_ref = dx_full.reshape(x.shape[0], wts.shape[1], -1).sum(dim=1) assert torch.equal(a[0].grad, dx_ref), "dX mismatch vs compacted stock presentation" + _assert_masked_all_valid_matches_unmasked() + _assert_all_masked_output_and_gradients_are_zero() -@pytest.mark.gpu -def test_masked_all_valid_bitwise_matches_unmasked_path(): + +def _assert_masked_all_valid_matches_unmasked(): """With zero masked slots, the filter_expert lane must produce grads bit-identical to the stock (filter_expert=False) lane.""" if not torch.cuda.is_available(): @@ -524,8 +524,7 @@ def test_masked_all_valid_bitwise_matches_unmasked_path(): assert torch.equal(mine, stock), f"{name}: all-valid masked lane diverged from stock lane" -@pytest.mark.gpu -def test_masked_all_masked_zero_output_and_grads(): +def _assert_all_masked_output_and_gradients_are_zero(): """A rank that owns none of the routed experts: zero forward, zero grads.""" if not torch.cuda.is_available(): pytest.skip("needs CUDA") @@ -548,7 +547,7 @@ def test_masked_all_masked_zero_output_and_grads(): assert g is not None and g.shape == ref.shape and (g == 0).all(), f"{name} must be exact zeros" -def test_weight_cache_reuses_and_invalidates(routed_inputs, monkeypatch): +def _assert_weight_mode_cache_and_kernel_layout_policy(routed_inputs, monkeypatch): """Cache mode: transposes are reused across forwards, invalidated on in-place parameter updates; transient mode makes fresh copies each forward.""" blk = _block("eager") @@ -565,15 +564,13 @@ def fake_impl(hidden, w13, w2, topk_weights, topk_ids, **kwargs): # transient mode -> fresh transpose copies each forward monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE", "transient") - monkeypatch.delenv("XORL_MOE_SGLANG_FUSED_EXPERTS_CACHE_WEIGHTS", raising=False) blk.experts.sglang_fused_experts_forward(x, w, ids) blk.experts.sglang_fused_experts_forward(x, w, ids) assert seen[0][0] is not seen[1][0] assert seen[0][0].is_contiguous() - # legacy cache alias (no explicit mode) -> same transposed tensors reused - monkeypatch.delenv("XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE", raising=False) - monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS_CACHE_WEIGHTS", "1") + # cached mode -> same transposed tensors reused + monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE", "cached") seen.clear() blk.experts.sglang_fused_experts_forward(x, w, ids) blk.experts.sglang_fused_experts_forward(x, w, ids) @@ -592,20 +589,22 @@ def fake_impl(hidden, w13, w2, topk_weights, topk_ids, **kwargs): blk.experts.invalidate_sglang_fused_weight_cache() assert blk.experts._sglang_fused_weight_cache == {} + _assert_weight_mode_selection(monkeypatch) + _assert_strided_mode_passes_zero_copy_views(routed_inputs, monkeypatch) + _assert_kernel_receives_sglang_layout(routed_inputs, monkeypatch) + _assert_strided_adapter_layout_policy(monkeypatch) -def test_weight_mode_selection(monkeypatch): - """Default is strided; explicit WEIGHT_MODE wins; legacy cache env aliases - cached; invalid rejects.""" + +def _assert_weight_mode_selection(monkeypatch): + """Default is strided; explicit modes resolve; invalid values reject.""" from xorl.models.layers.moe.experts import moe_sglang_fused_experts_weight_mode # noqa: PLC0415 monkeypatch.delenv("XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE", raising=False) - monkeypatch.delenv("XORL_MOE_SGLANG_FUSED_EXPERTS_CACHE_WEIGHTS", raising=False) assert moe_sglang_fused_experts_weight_mode() == "strided" - monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS_CACHE_WEIGHTS", "1") + monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE", "cached") assert moe_sglang_fused_experts_weight_mode() == "cached" - # explicit mode overrides the legacy cache alias monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE", "transient") assert moe_sglang_fused_experts_weight_mode() == "transient" monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE", "strided") @@ -616,7 +615,7 @@ def test_weight_mode_selection(monkeypatch): moe_sglang_fused_experts_weight_mode() -def test_strided_mode_passes_zero_copy_views(routed_inputs, monkeypatch): +def _assert_strided_mode_passes_zero_copy_views(routed_inputs, monkeypatch): """Strided mode must hand the kernel transpose-VIEWS of the GKN parameters (same storage, non-contiguous, serving element order) and never populate the cache.""" blk = _block("eager") @@ -644,15 +643,8 @@ def fake_impl(hidden, w13, w2, topk_weights, topk_ids, **kwargs): assert seen["gemm1_limit"] is None assert getattr(blk.experts, "_sglang_fused_weight_cache", None) in (None, {}) - # strided mode ignores the legacy cache env (explicit mode wins) - monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS_CACHE_WEIGHTS", "1") - seen.clear() - blk.experts.sglang_fused_experts_forward(x, w, ids) - assert seen["w13"].data_ptr() == blk.experts.gate_up_proj.data_ptr() - assert getattr(blk.experts, "_sglang_fused_weight_cache", None) in (None, {}) - -def test_strided_vendored_impl_layout_guard(): +def _assert_strided_adapter_layout_policy(monkeypatch): """The vendored strided impl accepts serving-contiguous and GKN transpose-view layouts and nothing else (importable without sglang).""" from xorl.ops.moe.sglang_fused_moe_strided import serving_layout_or_gkn_view # noqa: PLC0415 @@ -662,8 +654,10 @@ def test_strided_vendored_impl_layout_guard(): assert serving_layout_or_gkn_view(gkn.transpose(1, 2).contiguous()) # serving layout assert not serving_layout_or_gkn_view(gkn[:, ::2, :].transpose(1, 2)) # sliced: neither + _assert_strided_impl_split_gate_up_layout(monkeypatch) + -def test_strided_impl_delegates_with_split_gate_up_layout(monkeypatch): +def _assert_strided_impl_split_gate_up_layout(monkeypatch): import xorl.ops.moe.sglang_fused_moe_strided as strided_mod # noqa: PLC0415 seen = {} @@ -737,7 +731,7 @@ def get_server_args(): return created, published -def test_ensure_sglang_runtime_publishes_xorl_deterministic_context(monkeypatch): +def _assert_sglang_runtime_context_policy(monkeypatch): from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 created, published = _install_fake_sglang_runtime(monkeypatch, existing=False) @@ -754,8 +748,11 @@ def test_ensure_sglang_runtime_publishes_xorl_deterministic_context(monkeypatch) assert len(published) == 1 assert published[0][1] == "scheduler" + _assert_sglang_runtime_preserves_compatible_context(monkeypatch) + _assert_sglang_runtime_rejects_incompatible_context(monkeypatch) -def test_ensure_sglang_runtime_preserves_compatible_context(monkeypatch): + +def _assert_sglang_runtime_preserves_compatible_context(monkeypatch): from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 created, published = _install_fake_sglang_runtime(monkeypatch, existing=True) @@ -765,31 +762,33 @@ def test_ensure_sglang_runtime_preserves_compatible_context(monkeypatch): assert published == [] -@pytest.mark.parametrize( - ("deterministic", "fused_sum_all_reduce"), - [(False, False), (True, True)], -) -def test_ensure_sglang_runtime_rejects_incompatible_context(monkeypatch, deterministic, fused_sum_all_reduce): +def _assert_sglang_runtime_rejects_incompatible_context(monkeypatch): from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 - _install_fake_sglang_runtime( - monkeypatch, - existing=True, - deterministic=deterministic, - fused_sum_all_reduce=fused_sum_all_reduce, - ) - with pytest.raises(RuntimeError, match="SGLang MoE parity requires"): - MoEExperts._ensure_sglang_server_args() + for deterministic, fused_sum_all_reduce in ((False, False), (True, True)): + _install_fake_sglang_runtime( + monkeypatch, + existing=True, + deterministic=deterministic, + fused_sum_all_reduce=fused_sum_all_reduce, + ) + with pytest.raises(RuntimeError, match="SGLang MoE parity requires"): + MoEExperts._ensure_sglang_server_args() -def test_sglang_runtime_api_does_not_regress_to_legacy_globals(): - source = inspect.getsource(experts_mod.MoEExperts._ensure_sglang_server_args) - assert "get_global_server_args" not in source - assert "set_global_server_args_for_scheduler" not in source +@pytest.mark.gpu +def test_real_sglang_parity_policy(monkeypatch, auto_state): + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + with monkeypatch.context() as case_patch: + _assert_trainable_gradient_numerics_policy(case_patch) + with monkeypatch.context() as case_patch: + _assert_strided_mode_bit_identical_to_transient(case_patch) + with monkeypatch.context() as case_patch: + _assert_auto_parity_forward_deterministic_and_matches_explicit(case_patch, auto_state) -@pytest.mark.gpu -def test_strided_mode_bit_identical_to_transient(monkeypatch): +def _assert_strided_mode_bit_identical_to_transient(monkeypatch): """Forward and all gradients under WEIGHT_MODE=strided must be bit-identical to the transient-transpose mode (same kernels, view-strided addressing).""" if not torch.cuda.is_available(): @@ -841,7 +840,7 @@ def pinned_index_compute(expert_index, cumsum_t): assert torch.equal(a, b), f"{name} differs between strided and transient weight modes" -def test_kernel_receives_sglang_layout_and_fp32_weights(routed_inputs, monkeypatch): +def _assert_kernel_receives_sglang_layout(routed_inputs, monkeypatch): """Fake the kernel to check the exact tensors the SGLang path hands over.""" blk = _block("eager") x, w, ids = routed_inputs @@ -873,35 +872,7 @@ def fake_impl(hidden, w13, w2, topk_weights, topk_ids, **kwargs): assert seen["kwargs"]["apply_router_weight_on_input"] is False -@pytest.mark.gpu -def test_auto_ep1_forward_dispatches_parity_path(monkeypatch, auto_state): - """Unset env at ep=1 on CUDA must dispatch MoEBlock.forward to the parity - path; explicit 0 is the escape hatch back to the stock tree.""" - if not torch.cuda.is_available(): - pytest.skip("needs CUDA") - monkeypatch.delenv(FLAG, raising=False) - _stack_available(monkeypatch, True) - blk = _block("triton").to("cuda") - torch.manual_seed(1) - x = torch.randn(TOKENS, HID, dtype=torch.bfloat16, device="cuda") - called = {} - - def fake(hidden_states, routing_weights, selected_experts): - called["hit"] = True - return torch.zeros_like(hidden_states) - - monkeypatch.setattr(blk.experts, "sglang_fused_experts_forward", fake) - blk.forward(x.view(1, TOKENS, HID)) - assert called.get("hit"), "unset env at ep=1 on CUDA must auto-enable the parity path" - - called.clear() - monkeypatch.setenv(FLAG, "0") - blk.forward(x.view(1, TOKENS, HID)) - assert not called.get("hit"), "explicit 0 must keep the stock path" - - -@pytest.mark.gpu -def test_auto_parity_forward_deterministic_and_matches_explicit(monkeypatch, auto_state): +def _assert_auto_parity_forward_deterministic_and_matches_explicit(monkeypatch, auto_state): """Real-kernel ep=1 sanity: the auto-enabled forward is bit-identical across two runs (determinism) and to the explicit FLAG=1 path.""" if not torch.cuda.is_available(): diff --git a/tests/models/test_moe_sglang_fused_experts_ep.py b/tests/models/test_moe_sglang_fused_experts_ep.py index 3ecbaeb4..62496802 100644 --- a/tests/models/test_moe_sglang_fused_experts_ep.py +++ b/tests/models/test_moe_sglang_fused_experts_ep.py @@ -11,18 +11,13 @@ slices in serving layout per the weight mode (zero-copy strided views by default; transient/cached transpose-copies as escape hatches), no post-hoc multiply (kernel call is faked; numerical parity is covered separately). -4. Slot-combine sub-flag: the (token, slot) -> pair-row mapping is correct against a - brute-force oracle, the full sub-flag combine path reproduces the slot-ordered - fp32 reduction, and duplicate expert selections are rejected. -5. Trainable path: grad-requiring inputs route through +4. Trainable path: grad-requiring inputs route through ``_SglangFusedExpertsEPTrainFunction`` (no-grad keeps the plain kernel call), - training guards raise, empty ranks flow exact-zero weight grads, slot-combine - rejects grad-requiring expert outputs, and (GPU) the backward is bit-identical - to the stock ``TritonEPGroupGemm`` gradients on the same post-dispatch inputs. + training guards raise, and empty ranks flow exact-zero weight grads. Numerical + backward correctness is owned by the independent FP32-routing oracle report. """ -import sys -import types +import builtins import pytest import torch @@ -36,7 +31,6 @@ NUM_EXPERTS, TOP_K, HID, INTER, TOKENS = 8, 2, 32, 24, 16 FLAG = "XORL_MOE_SGLANG_FUSED_EXPERTS" -SLOT_COMBINE_FLAG = "XORL_MOE_SGLANG_FUSED_EXPERTS_SLOT_COMBINE" def _experts(moe_implementation: str = "eager") -> MoEExperts: @@ -91,16 +85,7 @@ def _world1_all_gather_into_tensor(output, input, group=None, async_op=False): dist.destroy_process_group() -def test_deepep_dispatch_raises_with_flag(routed_inputs, monkeypatch): - monkeypatch.setenv(FLAG, "1") - experts = _experts("eager") - experts.ep_dispatch = "deepep" - x, w, ids = routed_inputs - with pytest.raises(NotImplementedError, match="alltoall"): - experts._ep_forward(x, w, ids, _FakeParallelState()) - - -def test_deepep_exclusion_asserts_no_unverifiable_mechanism(routed_inputs, monkeypatch): +def _assert_deepep_exclusion_asserts_no_unverifiable_mechanism(routed_inputs, monkeypatch): """The reason states what is known; the wheel ships no CUDA source to check against.""" monkeypatch.setenv(FLAG, "1") experts = _experts("eager") @@ -109,6 +94,7 @@ def test_deepep_exclusion_asserts_no_unverifiable_mechanism(routed_inputs, monke with pytest.raises(NotImplementedError) as exc: experts._ep_forward(x, w, ids, _FakeParallelState()) message = str(exc.value) + assert "alltoall" in message assert "order and rounding schedule differ" in message assert "No mechanism is asserted" in message # The retired explanation did not distinguish the paths it gated: the @@ -116,7 +102,7 @@ def test_deepep_exclusion_asserts_no_unverifiable_mechanism(routed_inputs, monke assert "bf16 cast of the per-rank partials" not in message -def test_fp8_ep_compute_raises_with_flag(routed_inputs, monkeypatch): +def _assert_fp8_ep_compute_raises_with_flag(routed_inputs, monkeypatch): monkeypatch.setenv(FLAG, "1") experts = _experts("eager") experts.fp8_training_enabled = True @@ -125,7 +111,7 @@ def test_fp8_ep_compute_raises_with_flag(routed_inputs, monkeypatch): experts._ep_forward(x, w, ids, _FakeParallelState()) -def test_flag_off_ep_forward_keeps_stock_path(routed_inputs, monkeypatch, world1_ep_group): +def _assert_flag_off_ep_forward_keeps_stock_path(routed_inputs, monkeypatch, world1_ep_group): from xorl.models.layers.moe.backend import EP_EXPERT_COMPUTE # noqa: PLC0415 monkeypatch.delenv(FLAG, raising=False) @@ -148,7 +134,7 @@ def stock_compute(permute_tokens, cumsum, *args, **kwargs): assert torch.isfinite(out.float()).all() -def test_flag_off_scores_cast_to_token_dtype(routed_inputs, monkeypatch, world1_ep_group): +def _assert_flag_off_scores_cast_to_token_dtype(routed_inputs, monkeypatch, world1_ep_group): from xorl.distributed.moe.alltoall import alltoall_pre_dispatch # noqa: PLC0415 x, w, ids = routed_inputs @@ -161,7 +147,7 @@ def test_flag_off_scores_cast_to_token_dtype(routed_inputs, monkeypatch, world1_ assert ctx.expert_scores.dtype == torch.float32 -def test_flag_on_ep_compute_topk1_presentation(routed_inputs, monkeypatch, world1_ep_group): +def _assert_flag_on_ep_compute_topk1_presentation(routed_inputs, monkeypatch, world1_ep_group): """The wrapper must hand fused_experts_impl the exact EP topk=1 contract.""" monkeypatch.setenv(FLAG, "1") experts = _experts("eager") @@ -209,7 +195,7 @@ def fake_impl(hidden, w13, w2, topk_weights, topk_ids, **kwargs): assert seen["kwargs"]["filter_expert"] is False -def test_empty_rank_short_circuits_kernel(monkeypatch): +def _assert_empty_rank_short_circuits_kernel(monkeypatch): monkeypatch.setenv(FLAG, "1") experts = _experts("eager") @@ -224,10 +210,9 @@ def fake_impl(*args, **kwargs): assert out.shape == (0, HID) -def test_ep_compute_guards(routed_inputs, monkeypatch): +def _assert_ep_compute_guards(monkeypatch): monkeypatch.setenv(FLAG, "1") experts = _experts("eager") - x, w, ids = routed_inputs permute_tokens = torch.randn(4, HID, dtype=torch.bfloat16) cumsum = torch.tensor([1, 2, 2, 3, 3, 4, 4, 4], dtype=torch.int64) scores = torch.rand(4, dtype=torch.float32) @@ -251,11 +236,15 @@ def test_ep_compute_guards(routed_inputs, monkeypatch): experts.hidden_act = "silu" -def test_missing_sglang_raises_import_error_naming_flag(monkeypatch): - import importlib.util # noqa: PLC0415 +def _assert_missing_sglang_raises_import_error_naming_flag(monkeypatch): + real_import = builtins.__import__ + + def import_without_sglang(name, *args, **kwargs): + if name == "sglang" or name.startswith("sglang."): + raise ModuleNotFoundError("simulated missing sglang", name=name) + return real_import(name, *args, **kwargs) - if importlib.util.find_spec("sglang") is not None: - pytest.skip("sglang installed; import-error guard not testable here") + monkeypatch.setattr(builtins, "__import__", import_without_sglang) monkeypatch.setenv(FLAG, "1") experts = _experts("eager") permute_tokens = torch.randn(4, HID, dtype=torch.bfloat16) @@ -266,82 +255,6 @@ def test_missing_sglang_raises_import_error_naming_flag(monkeypatch): experts.sglang_fused_experts_ep_compute(permute_tokens, cumsum, scores) -def test_pair_slot_order_mapping_matches_bruteforce(): - torch.manual_seed(3) - num_tokens, topk = 11, 3 - selected = torch.stack([torch.randperm(NUM_EXPERTS)[:topk] for _ in range(num_tokens)]) - order = MoEExperts._sglang_fused_experts_pair_slot_order(selected) - - # brute-force oracle: arrival rows are (expert, token) pairs sorted by (e, t); - # arrival row r must land at flat slot index token * topk + slot. - pairs = sorted((int(selected[t, j]), t, j) for t in range(num_tokens) for j in range(topk)) - expected = torch.tensor([t * topk + j for (_, t, j) in pairs], dtype=order.dtype) - assert torch.equal(order, expected) - - -def test_slot_combine_matches_slot_ordered_reduction(routed_inputs, monkeypatch, world1_ep_group): - """Full sub-flag path: gather into [T, topk, H] slot order + reduce.""" - monkeypatch.setenv(FLAG, "1") - monkeypatch.setenv(SLOT_COMBINE_FLAG, "1") - experts = _experts("eager") - x, w, ids = routed_inputs - - def fake_impl(hidden, w13, w2, topk_weights, topk_ids, **kwargs): - # deterministic pseudo-weighted rows: distinguishable per pair row - return (hidden.float() * topk_weights).to(hidden.dtype) - - def fake_moe_sum_reduce(slots, out, scaling): - assert slots.dim() == 3 - out.copy_((slots.float().sum(dim=1) * scaling).to(out.dtype)) - - monkeypatch.setattr(type(experts), "_load_sglang_fused_experts_impl", staticmethod(lambda: fake_impl)) - monkeypatch.setattr(type(experts), "_sglang_fused_experts_ep_config_logged", True, raising=False) - monkeypatch.setitem(sys.modules, "sgl_kernel", types.SimpleNamespace(moe_sum_reduce=fake_moe_sum_reduce)) - - # Scoring contract: the slot-combine variant is inference-only (no autograd - # through moe_sum_reduce), matching the production logprob-replay context. - with torch.no_grad(): - out = experts._ep_forward(x, w, ids, _FakeParallelState(world1_ep_group)) - - # independent expectation: per-slot rows in [T, topk, H] slot order, fp32 sum. - slots = (x.float().unsqueeze(1) * w.unsqueeze(-1)).to(x.dtype) # [T, topk, H] - expected = slots.float().sum(dim=1).to(x.dtype) - assert torch.equal(out, expected) - - -def test_slot_combine_pair_count_guard(monkeypatch, world1_ep_group): - """Collapsed pair rows (duplicate expert selections) must fail loudly, not misalign. - - The alltoall dispatch itself already rejects duplicates upstream (split - mismatch), so exercise the combine-side guard directly with a crafted - context that returns fewer pair rows than ``num_tokens * topk`` slots. - """ - from xorl.distributed.moe.alltoall import AllToAllDispatchContext # noqa: PLC0415 - - experts = _experts("eager") - monkeypatch.setitem(sys.modules, "sgl_kernel", types.SimpleNamespace(moe_sum_reduce=lambda *a: None)) - - num_tokens, topk, rows = 4, 2, 7 # one pair row collapsed - counts = torch.zeros(1, NUM_EXPERTS, dtype=torch.int64) - counts[0, 0] = rows - ctx = AllToAllDispatchContext( - input_splits=[rows], - output_splits=[rows], - num_tokens_per_expert=counts, - routing_map=None, - perm_mapping=None, - expert_scores=None, - orig_shape=torch.Size((num_tokens, HID)), - num_experts=NUM_EXPERTS, - ) - expert_output = torch.randn(rows, HID, dtype=torch.bfloat16) - dispatch_kwargs = {"selected_experts": torch.randint(0, NUM_EXPERTS, (num_tokens, topk))} - with pytest.raises(NotImplementedError, match="unique expert selections"): - experts._sglang_fused_experts_slot_combine( - expert_output, ctx, dispatch_kwargs, _FakeParallelState(world1_ep_group) - ) - - def _dispatched_inputs(requires_grad: bool = False): """Post-dispatch pair rows for a direct sglang_fused_experts_ep_compute call.""" torch.manual_seed(2) @@ -354,7 +267,7 @@ def _dispatched_inputs(requires_grad: bool = False): return permute_tokens, cumsum, scores -def test_trainable_ep_dispatch_uses_autograd_function(monkeypatch): +def _assert_trainable_ep_dispatch_uses_autograd_function(monkeypatch): """When gradients are required, the EP compute must route through the autograd Function; the no-grad path must keep using the plain kernel call.""" from xorl.models.layers.moe import experts as experts_mod # noqa: PLC0415 @@ -373,7 +286,7 @@ def fake_kernel_call(permute_tokens, *args, **kwargs): monkeypatch.setattr(experts_mod._SglangFusedExpertsEPTrainFunction, "apply", staticmethod(fake_apply)) monkeypatch.setattr(experts_mod, "_sglang_fused_experts_ep_kernel_call", fake_kernel_call) - monkeypatch.setattr(type(experts), "_load_sglang_fused_experts_impl", staticmethod(lambda: (lambda *a, **k: None))) + monkeypatch.setattr(type(experts), "_load_sglang_fused_experts_impl", staticmethod(lambda: lambda *a, **k: None)) permute_tokens, cumsum, scores = _dispatched_inputs() experts.gate_up_proj.requires_grad_(True) @@ -393,10 +306,10 @@ def fake_kernel_call(permute_tokens, *args, **kwargs): assert called == {"plain": True} -def test_trainable_ep_guards(monkeypatch): +def _assert_trainable_ep_guards(monkeypatch): monkeypatch.setenv(FLAG, "1") experts = _experts("eager") - monkeypatch.setattr(type(experts), "_load_sglang_fused_experts_impl", staticmethod(lambda: (lambda *a, **k: None))) + monkeypatch.setattr(type(experts), "_load_sglang_fused_experts_impl", staticmethod(lambda: lambda *a, **k: None)) experts.gate_up_proj.requires_grad_(True) permute_tokens, cumsum, scores = _dispatched_inputs() @@ -411,7 +324,7 @@ def test_trainable_ep_guards(monkeypatch): experts.hidden_act = "silu" -def test_trainable_ep_empty_rank_flows_zero_weight_grads(monkeypatch): +def _assert_trainable_ep_empty_rank_flows_zero_weight_grads(monkeypatch): """An empty rank must skip the kernel but still flow exact-zero weight grads (keeps FSDP grad reduction uniform across ranks).""" monkeypatch.setenv(FLAG, "1") @@ -437,85 +350,12 @@ def fake_impl(*args, **kwargs): assert scores.grad is not None and scores.grad.shape == scores.shape -def test_slot_combine_rejects_grad_requiring_output(monkeypatch, world1_ep_group): - """moe_sum_reduce has no autograd; a grad-requiring expert output must fail - loudly instead of silently detaching the graph.""" - from xorl.distributed.moe.alltoall import AllToAllDispatchContext # noqa: PLC0415 - - experts = _experts("eager") - monkeypatch.setitem(sys.modules, "sgl_kernel", types.SimpleNamespace(moe_sum_reduce=lambda *a: None)) - - num_tokens, topk = 4, 2 - rows = num_tokens * topk - counts = torch.zeros(1, NUM_EXPERTS, dtype=torch.int64) - counts[0, 0] = rows - ctx = AllToAllDispatchContext( - input_splits=[rows], - output_splits=[rows], - num_tokens_per_expert=counts, - routing_map=None, - perm_mapping=None, - expert_scores=None, - orig_shape=torch.Size((num_tokens, HID)), - num_experts=NUM_EXPERTS, - ) - expert_output = torch.randn(rows, HID, dtype=torch.bfloat16, requires_grad=True) - dispatch_kwargs = {"selected_experts": torch.randint(0, NUM_EXPERTS, (num_tokens, topk))} - with pytest.raises(NotImplementedError, match="scoring-only"): - experts._sglang_fused_experts_slot_combine( - expert_output, ctx, dispatch_kwargs, _FakeParallelState(world1_ep_group) - ) - - -@pytest.mark.gpu -def test_ep_trainable_grads_match_stock_triton(): - """dX / d_pair_scores / dW13 / dW2 must be bit-identical to the stock - TritonEPGroupGemm gradients on the same post-dispatch pair rows (the row - permutation is pinned by the dispatch presentation, so no pinning shim).""" - if not torch.cuda.is_available(): - pytest.skip("needs CUDA") - pytest.importorskip("sglang") - pytest.importorskip("sgl_kernel") - from xorl.models.layers.moe.experts import MoEExperts, _SglangFusedExpertsEPTrainFunction # noqa: PLC0415 - from xorl.ops.moe.triton import TritonEPGroupGemm # noqa: PLC0415 - - device = torch.device("cuda") - torch.manual_seed(0) - E, H2, I2, rows = 8, 64, 48, 24 - gu = (torch.randn(E, H2, 2 * I2, device=device) * 0.3).to(torch.bfloat16) - dn = (torch.randn(E, I2, H2, device=device) * 0.3).to(torch.bfloat16) - x = (torch.randn(rows, H2, device=device) * 0.5).to(torch.bfloat16) - counts = torch.tensor([4, 2, 0, 5, 3, 4, 2, 4], dtype=torch.int64, device=device) - assert int(counts.sum()) == rows - cumsum = torch.cumsum(counts, dim=0) - scores = torch.rand(rows, device=device, dtype=torch.float32) - - impl = MoEExperts._load_sglang_fused_experts_impl() - a = [t.clone().requires_grad_(True) for t in (x, scores, gu, dn)] - out = _SglangFusedExpertsEPTrainFunction.apply(a[0], a[1], a[2], a[3], cumsum, impl, "silu", "silu", 0.0, True) - grad_out = (torch.randn_like(out.float()) * 0.1).to(out.dtype) - out.backward(grad_out) - - b = [t.clone().requires_grad_(True) for t in (x, scores, gu, dn)] - out_stock = TritonEPGroupGemm.apply(b[0], cumsum, b[2], b[3], I2, b[1], "silu", 0.0, True, 0) - out_stock.backward(grad_out) - - for name, mine, stock in ( - ("dX_pair_rows", a[0].grad, b[0].grad), - ("d_pair_scores", a[1].grad, b[1].grad), - ("dW13_gkn", a[2].grad, b[2].grad), - ("dW2_gkn", a[3].grad, b[3].grad), - ): - assert torch.equal(mine, stock), f"{name} gradient mismatch vs stock triton EP path" - - -def test_ep_strided_mode_passes_zero_copy_weight_views(monkeypatch): +def _assert_ep_strided_mode_passes_zero_copy_weight_views(monkeypatch): """Default (strided) mode must hand the EP kernel transpose-VIEWS of the local GKN weight slices (same storage, non-contiguous, serving element order) and never populate the weight cache.""" monkeypatch.setenv(FLAG, "1") monkeypatch.delenv("XORL_MOE_SGLANG_FUSED_EXPERTS_WEIGHT_MODE", raising=False) - monkeypatch.delenv("XORL_MOE_SGLANG_FUSED_EXPERTS_CACHE_WEIGHTS", raising=False) experts = _experts("eager") seen = {} @@ -539,7 +379,7 @@ def fake_impl(hidden, w13, w2, topk_weights, topk_ids, **kwargs): assert getattr(experts, "_sglang_fused_weight_cache", None) in (None, {}) -def test_ep_weight_cache_reuses_and_invalidates(monkeypatch): +def _assert_ep_weight_cache_reuses_and_invalidates(monkeypatch): """Escape hatches: transient mode makes fresh contiguous transpose-copies per call; cached mode reuses the same transposed tensors across scoring AND trainable calls (one cache per module, shared with the local path), drops @@ -586,3 +426,40 @@ def fake_impl(hidden, w13, w2, topk_weights, topk_ids, **kwargs): # explicit invalidation drops entries (same hook as the local path) experts.invalidate_sglang_fused_weight_cache() assert experts._sglang_fused_weight_cache == {} + + +def test_sglang_ep_policy(routed_inputs, monkeypatch, world1_ep_group): + with monkeypatch.context() as case_patch: + _assert_flag_on_ep_compute_topk1_presentation(routed_inputs, case_patch, world1_ep_group) + with monkeypatch.context() as case_patch: + _assert_deepep_exclusion_asserts_no_unverifiable_mechanism(routed_inputs, case_patch) + with monkeypatch.context() as case_patch: + _assert_fp8_ep_compute_raises_with_flag(routed_inputs, case_patch) + with monkeypatch.context() as case_patch: + _assert_missing_sglang_raises_import_error_naming_flag(case_patch) + with monkeypatch.context() as case_patch: + _assert_flag_off_ep_forward_keeps_stock_path(routed_inputs, case_patch, world1_ep_group) + with monkeypatch.context() as case_patch: + _assert_flag_off_scores_cast_to_token_dtype(routed_inputs, case_patch, world1_ep_group) + with monkeypatch.context() as case_patch: + _assert_empty_rank_short_circuits_kernel(case_patch) + with monkeypatch.context() as case_patch: + _assert_ep_compute_guards(case_patch) + with monkeypatch.context() as trainable_patch: + _assert_sglang_ep_trainable_dispatch_contract(trainable_patch) + with monkeypatch.context() as weight_patch: + _assert_sglang_ep_weight_presentation_modes(weight_patch) + + +def _assert_sglang_ep_trainable_dispatch_contract(monkeypatch): + _assert_trainable_ep_dispatch_uses_autograd_function(monkeypatch) + monkeypatch.undo() + _assert_trainable_ep_guards(monkeypatch) + monkeypatch.undo() + _assert_trainable_ep_empty_rank_flows_zero_weight_grads(monkeypatch) + + +def _assert_sglang_ep_weight_presentation_modes(monkeypatch): + _assert_ep_strided_mode_passes_zero_copy_weight_views(monkeypatch) + monkeypatch.undo() + _assert_ep_weight_cache_reuses_and_invalidates(monkeypatch) diff --git a/tests/models/test_moe_tp_parity.py b/tests/models/test_moe_tp_parity.py deleted file mode 100644 index 6950e07f..00000000 --- a/tests/models/test_moe_tp_parity.py +++ /dev/null @@ -1,894 +0,0 @@ -from types import SimpleNamespace - -import torch - -import xorl.distributed.parallel_state as parallel_state_module -from xorl.models.layers.moe import MoEBlock, MoEExperts -from xorl.ops.moe.activations import apply_moe_activation - - -def _fake_tp_state(tp_size: int = 2): - return SimpleNamespace(ep_enabled=False, tp_size=tp_size) - - -def test_sglang_moe_tp_sim_env_enables_no_ep_tp1(monkeypatch): - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - tp1_state = _fake_tp_state(tp_size=1) - - monkeypatch.delenv("XORL_SGLANG_MOE_TP_SIM", raising=False) - assert not experts.sglang_moe_tp_sim_enabled(tp1_state) - - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "sglang_runner") - assert experts.sglang_moe_tp_sim_enabled(tp1_state) - - ep_state = SimpleNamespace(ep_enabled=True, tp_size=1) - assert not experts.sglang_moe_tp_sim_enabled(ep_state) - - -def test_sglang_moe_tp_sim_layer_filter(monkeypatch): - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - tp1_state = _fake_tp_state(tp_size=1) - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "sglang") - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM_LAYERS", "2,4") - - assert not experts.sglang_moe_tp_sim_enabled(tp1_state) - - experts.layer_idx = 3 - assert not experts.sglang_moe_tp_sim_enabled(tp1_state) - - experts.layer_idx = 4 - assert experts.sglang_moe_tp_sim_enabled(tp1_state) - - -def _manual_sglang_tp_sim_direct( - experts: MoEExperts, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - tp_size: int, -) -> torch.Tensor: - hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) - selected_flat = selected_experts.reshape(hidden_flat.shape[0], -1) - routing_flat = routing_weights.reshape(hidden_flat.shape[0], -1) - output = hidden_flat.new_zeros(hidden_flat.shape) - shard_intermediate = experts.intermediate_size // tp_size - - for tp_rank in range(tp_size): - start = tp_rank * shard_intermediate - end = start + shard_intermediate - shard_output = hidden_flat.new_zeros(hidden_flat.shape) - for expert_idx in range(experts.num_experts): - mask = selected_flat == expert_idx - if not bool(mask.any().item()): - continue - token_rows, topk_slots = mask.nonzero(as_tuple=True) - tokens = hidden_flat.index_select(0, token_rows) - gate = tokens.matmul(experts.gate_up_proj[expert_idx, :, start:end]) - up = tokens.matmul( - experts.gate_up_proj[ - expert_idx, - :, - experts.intermediate_size + start : experts.intermediate_size + end, - ] - ) - activated = apply_moe_activation(experts.hidden_act, gate, up) - expert_out = activated.matmul(experts.down_proj[expert_idx, start:end, :]) - expert_out = expert_out * routing_flat[token_rows, topk_slots].unsqueeze(-1) - shard_output.index_add_(0, token_rows, expert_out) - output = output + shard_output - - return output.reshape(hidden_states.shape) - - -def _manual_sglang_tp_sim_direct_bf16_reduce( - experts: MoEExperts, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - tp_size: int, -) -> torch.Tensor: - hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) - selected_flat = selected_experts.reshape(hidden_flat.shape[0], -1) - routing_flat = routing_weights.reshape(hidden_flat.shape[0], -1) - output = hidden_flat.new_zeros(hidden_flat.shape) - shard_intermediate = experts.intermediate_size // tp_size - - for tp_rank in range(tp_size): - start = tp_rank * shard_intermediate - end = start + shard_intermediate - shard_output = hidden_flat.new_zeros(hidden_flat.shape) - for expert_idx in range(experts.num_experts): - mask = selected_flat == expert_idx - if not bool(mask.any().item()): - continue - token_rows, topk_slots = mask.nonzero(as_tuple=True) - tokens = hidden_flat.index_select(0, token_rows) - gate = tokens.matmul(experts.gate_up_proj[expert_idx, :, start:end]) - up = tokens.matmul( - experts.gate_up_proj[ - expert_idx, - :, - experts.intermediate_size + start : experts.intermediate_size + end, - ] - ) - activated = apply_moe_activation(experts.hidden_act, gate, up) - expert_out = activated.matmul(experts.down_proj[expert_idx, start:end, :]) - expert_out = expert_out * routing_flat[token_rows, topk_slots].unsqueeze(-1) - shard_output.index_add_(0, token_rows, expert_out) - output = output.to(torch.bfloat16) + shard_output.to(torch.bfloat16) - - return output.reshape(hidden_states.shape) - - -def _manual_sglang_tp_sim_cache( - experts: MoEExperts, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - selected_experts: torch.Tensor, - tp_size: int, -) -> torch.Tensor: - hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) - selected_flat = selected_experts.reshape(hidden_flat.shape[0], -1) - routing_flat = routing_weights.reshape(hidden_flat.shape[0], -1) - output = hidden_flat.new_zeros(hidden_flat.shape) - shard_intermediate = experts.intermediate_size // tp_size - - for tp_rank in range(tp_size): - start = tp_rank * shard_intermediate - end = start + shard_intermediate - topk = selected_flat.shape[1] - gate_up_cache = hidden_flat.new_zeros(hidden_flat.shape[0] * topk, 2 * shard_intermediate) - down_cache = hidden_flat.new_zeros(hidden_flat.shape[0], topk, hidden_flat.shape[-1]) - - for expert_idx in range(experts.num_experts): - mask = selected_flat == expert_idx - if not bool(mask.any().item()): - continue - token_rows, topk_slots = mask.nonzero(as_tuple=True) - assignment_rows = token_rows * topk + topk_slots - tokens = hidden_flat.index_select(0, token_rows) - gate = tokens.matmul(experts.gate_up_proj[expert_idx, :, start:end]) - up = tokens.matmul( - experts.gate_up_proj[ - expert_idx, - :, - experts.intermediate_size + start : experts.intermediate_size + end, - ] - ) - gate_up_cache[assignment_rows, :shard_intermediate] = gate - gate_up_cache[assignment_rows, shard_intermediate:] = up - - gate = gate_up_cache[:, :shard_intermediate] - up = gate_up_cache[:, shard_intermediate:] - activated_cache = apply_moe_activation(experts.hidden_act, gate, up) - - for expert_idx in range(experts.num_experts): - mask = selected_flat == expert_idx - if not bool(mask.any().item()): - continue - token_rows, topk_slots = mask.nonzero(as_tuple=True) - assignment_rows = token_rows * topk + topk_slots - activated = activated_cache.index_select(0, assignment_rows) - expert_out = activated.matmul(experts.down_proj[expert_idx, start:end, :]) - expert_out = expert_out * routing_flat[token_rows, topk_slots].unsqueeze(-1) - down_cache[token_rows, topk_slots, :] = expert_out - shard_output = down_cache.to(torch.float32).sum(dim=1).to(down_cache.dtype) - output = output + shard_output - - return output.reshape(hidden_states.shape) - - -def _patch_fake_group_gemm(monkeypatch): - import xorl.ops.group_gemm.kernel.group_gemm as group_gemm # noqa: PLC0415 - import xorl.ops.group_gemm.kernel.moe as moe_kernel # noqa: PLC0415 - - def fake_expert_histogram(expert_index, num_experts): - return torch.bincount(expert_index.reshape(-1), minlength=num_experts).to(torch.int32) - - def fake_moe_index_compute(expert_index, cumsum_t): - starts = torch.cat([cumsum_t.new_zeros(1), cumsum_t[:-1]]).to(torch.long) - offsets = starts.clone() - scatter_index = torch.empty_like(expert_index, dtype=torch.long) - for token_idx in range(expert_index.shape[0]): - for topk_idx in range(expert_index.shape[1]): - expert_idx = int(expert_index[token_idx, topk_idx].item()) - scatter_index[token_idx, topk_idx] = offsets[expert_idx] - offsets[expert_idx] += 1 - return scatter_index - - def fake_moe_scatter(hidden_states, scatter_index): - output = hidden_states.new_empty(scatter_index.numel(), hidden_states.shape[-1]) - for token_idx in range(scatter_index.shape[0]): - for topk_idx in range(scatter_index.shape[1]): - output[int(scatter_index[token_idx, topk_idx].item())] = hidden_states[token_idx] - return output - - def fake_group_gemm_same_nk(*, a, b, cumsum_M, max_M, transpose_b=False, **kwargs): - del kwargs - output_dim = b.shape[1] if transpose_b else b.shape[2] - output = a.new_empty(max_M, output_dim) - start = 0 - for expert_idx, end_value in enumerate(cumsum_M.tolist()): - end = int(end_value) - if end > start: - weight = b[expert_idx].transpose(0, 1) if transpose_b else b[expert_idx] - output[start:end] = a[start:end].matmul(weight) - start = end - return output - - monkeypatch.setattr(moe_kernel, "expert_histogram", fake_expert_histogram) - monkeypatch.setattr(moe_kernel, "moe_index_compute", fake_moe_index_compute) - monkeypatch.setattr(moe_kernel, "moe_scatter", fake_moe_scatter) - monkeypatch.setattr(group_gemm, "group_gemm_same_nk", fake_group_gemm_same_nk) - - -def test_sglang_moe_tp_sim_matches_manual_shard_reduce(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "1") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - - torch.manual_seed(0) - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - with torch.no_grad(): - experts.gate_up_proj.copy_(torch.randn_like(experts.gate_up_proj) * 0.1) - experts.down_proj.copy_(torch.randn_like(experts.down_proj) * 0.1) - - hidden_states = torch.randn(5, 4) - selected_experts = torch.tensor( - [ - [0, 1], - [2, 0], - [1, 2], - [2, 1], - [0, 2], - ], - dtype=torch.long, - ) - routing_weights = torch.tensor( - [ - [0.75, 0.25], - [0.60, 0.40], - [0.55, 0.45], - [0.80, 0.20], - [0.50, 0.50], - ], - dtype=hidden_states.dtype, - ) - - actual = experts(hidden_states, routing_weights, selected_experts) - expected = _manual_sglang_tp_sim_direct(experts, hidden_states, routing_weights, selected_experts, tp_size=2) - - torch.testing.assert_close(actual, expected) - - -def test_sglang_moe_tp_sim_bf16_reduce_forces_shard_accumulation_dtype(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "1") - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM_BF16_REDUCE", "1") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - - torch.manual_seed(0) - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - with torch.no_grad(): - experts.gate_up_proj.copy_(torch.randn_like(experts.gate_up_proj) * 0.1) - experts.down_proj.copy_(torch.randn_like(experts.down_proj) * 0.1) - - hidden_states = torch.randn(5, 4) - selected_experts = torch.tensor( - [ - [0, 1], - [2, 0], - [1, 2], - [2, 1], - [0, 2], - ], - dtype=torch.long, - ) - routing_weights = torch.tensor( - [ - [0.75, 0.25], - [0.60, 0.40], - [0.55, 0.45], - [0.80, 0.20], - [0.50, 0.50], - ], - dtype=hidden_states.dtype, - ) - - actual = experts(hidden_states, routing_weights, selected_experts) - expected = _manual_sglang_tp_sim_direct_bf16_reduce( - experts, - hidden_states, - routing_weights, - selected_experts, - tp_size=2, - ) - - assert actual.dtype == torch.bfloat16 - torch.testing.assert_close(actual, expected) - - -def test_sglang_moe_tp_cache_mode_matches_manual_cache_reduce(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "cache") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - - torch.manual_seed(0) - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - with torch.no_grad(): - experts.gate_up_proj.copy_(torch.randn_like(experts.gate_up_proj) * 0.1) - experts.down_proj.copy_(torch.randn_like(experts.down_proj) * 0.1) - - hidden_states = torch.randn(5, 4) - selected_experts = torch.tensor( - [ - [0, 1], - [2, 0], - [1, 2], - [2, 1], - [0, 2], - ], - dtype=torch.long, - ) - routing_weights = torch.tensor( - [ - [0.75, 0.25], - [0.60, 0.40], - [0.55, 0.45], - [0.80, 0.20], - [0.50, 0.50], - ], - dtype=hidden_states.dtype, - ) - - actual = experts(hidden_states, routing_weights, selected_experts) - expected = _manual_sglang_tp_sim_cache(experts, hidden_states, routing_weights, selected_experts, tp_size=2) - - torch.testing.assert_close(actual, expected) - - -def test_sglang_moe_tp_sim_size_override_under_tp1(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "cache") - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM_SIZE", "2") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=1)) - - torch.manual_seed(0) - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - with torch.no_grad(): - experts.gate_up_proj.copy_(torch.randn_like(experts.gate_up_proj) * 0.1) - experts.down_proj.copy_(torch.randn_like(experts.down_proj) * 0.1) - - hidden_states = torch.randn(5, 4) - selected_experts = torch.tensor( - [ - [0, 1], - [2, 0], - [1, 2], - [2, 1], - [0, 2], - ], - dtype=torch.long, - ) - routing_weights = torch.tensor( - [ - [0.75, 0.25], - [0.60, 0.40], - [0.55, 0.45], - [0.80, 0.20], - [0.50, 0.50], - ], - dtype=hidden_states.dtype, - ) - - actual = experts(hidden_states, routing_weights, selected_experts) - expected = _manual_sglang_tp_sim_cache(experts, hidden_states, routing_weights, selected_experts, tp_size=2) - - torch.testing.assert_close(actual, expected) - - -def test_sglang_moe_tp_triton_mode_shards_backend_call(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "triton") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - - torch.manual_seed(0) - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - with torch.no_grad(): - experts.gate_up_proj.copy_(torch.randn_like(experts.gate_up_proj) * 0.1) - experts.down_proj.copy_(torch.randn_like(experts.down_proj) * 0.1) - - hidden_states = torch.randn(5, 4) - selected_experts = torch.tensor( - [ - [0, 1], - [2, 0], - [1, 2], - [2, 1], - [0, 2], - ], - dtype=torch.long, - ) - routing_weights = torch.tensor( - [ - [0.75, 0.25], - [0.60, 0.40], - [0.55, 0.45], - [0.80, 0.20], - [0.50, 0.50], - ], - dtype=hidden_states.dtype, - ) - - calls = [] - - def fake_triton_moe_forward( - *, - num_experts, - routing_weights, - selected_experts, - hidden_states, - gate_proj, - up_proj, - down_proj, - gate_up_proj, - hidden_act, - swiglu_limit, - **kwargs, - ): - del kwargs, gate_proj, up_proj, swiglu_limit - calls.append( - { - "gate_up_shape": tuple(gate_up_proj.shape), - "down_shape": tuple(down_proj.shape), - "hidden_shape": tuple(hidden_states.shape), - } - ) - shard_intermediate = gate_up_proj.shape[-1] // 2 - output = hidden_states.new_zeros(hidden_states.shape) - for expert_idx in range(num_experts): - mask = selected_experts == expert_idx - if not bool(mask.any().item()): - continue - token_rows, topk_slots = mask.nonzero(as_tuple=True) - tokens = hidden_states.index_select(0, token_rows) - gate_up = tokens.matmul(gate_up_proj[expert_idx]) - gate, up = gate_up.split(shard_intermediate, dim=-1) - activated = apply_moe_activation(hidden_act, gate, up) - expert_out = activated.matmul(down_proj[expert_idx]) - expert_out = expert_out * routing_weights[token_rows, topk_slots].unsqueeze(-1) - output.index_add_(0, token_rows, expert_out) - return output - - import xorl.ops.moe.triton as triton_moe # noqa: PLC0415 - - monkeypatch.setattr(triton_moe, "triton_moe_forward", fake_triton_moe_forward) - - actual = experts(hidden_states, routing_weights, selected_experts) - expected = _manual_sglang_tp_sim_direct(experts, hidden_states, routing_weights, selected_experts, tp_size=2) - - torch.testing.assert_close(actual, expected) - assert calls == [ - {"gate_up_shape": (3, 4, 6), "down_shape": (3, 3, 4), "hidden_shape": (5, 4)}, - {"gate_up_shape": (3, 4, 6), "down_shape": (3, 3, 4), "hidden_shape": (5, 4)}, - ] - - -def test_sglang_moe_tp_triton_sgl_reduce_mode_matches_manual_cache_reduce(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "triton_sgl_reduce") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - _patch_fake_group_gemm(monkeypatch) - - torch.manual_seed(0) - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - with torch.no_grad(): - experts.gate_up_proj.copy_(torch.randn_like(experts.gate_up_proj) * 0.1) - experts.down_proj.copy_(torch.randn_like(experts.down_proj) * 0.1) - - hidden_states = torch.randn(5, 4) - selected_experts = torch.tensor( - [ - [0, 1], - [2, 0], - [1, 2], - [2, 1], - [0, 2], - ], - dtype=torch.long, - ) - routing_weights = torch.tensor( - [ - [0.75, 0.25], - [0.60, 0.40], - [0.55, 0.45], - [0.80, 0.20], - [0.50, 0.50], - ], - dtype=hidden_states.dtype, - ) - - actual = experts(hidden_states, routing_weights, selected_experts) - expected = _manual_sglang_tp_sim_cache(experts, hidden_states, routing_weights, selected_experts, tp_size=2) - - torch.testing.assert_close(actual, expected) - - -def test_sglang_moe_tp_deep_gemm_mode_matches_manual_cache_reduce(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "deep_gemm") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - _patch_fake_group_gemm(monkeypatch) - - calls = [] - - def fake_deep_gemm_group_gemm_same_nk(*, a, b, cumsum_M): - calls.append({"a_shape": tuple(a.shape), "b_shape": tuple(b.shape), "cumsum": tuple(cumsum_M.tolist())}) - output = a.new_empty(a.shape[0], b.shape[2]) - start = 0 - for expert_idx, end_value in enumerate(cumsum_M.tolist()): - end = int(end_value) - if end > start: - output[start:end] = a[start:end].matmul(b[expert_idx]) - start = end - return output - - monkeypatch.setattr( - MoEExperts, - "_deep_gemm_group_gemm_same_nk", - staticmethod(fake_deep_gemm_group_gemm_same_nk), - ) - - torch.manual_seed(0) - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - with torch.no_grad(): - experts.gate_up_proj.copy_(torch.randn_like(experts.gate_up_proj) * 0.1) - experts.down_proj.copy_(torch.randn_like(experts.down_proj) * 0.1) - - hidden_states = torch.randn(5, 4) - selected_experts = torch.tensor( - [ - [0, 1], - [2, 0], - [1, 2], - [2, 1], - [0, 2], - ], - dtype=torch.long, - ) - routing_weights = torch.tensor( - [ - [0.75, 0.25], - [0.60, 0.40], - [0.55, 0.45], - [0.80, 0.20], - [0.50, 0.50], - ], - dtype=hidden_states.dtype, - ) - - actual = experts(hidden_states, routing_weights, selected_experts) - expected = _manual_sglang_tp_sim_cache(experts, hidden_states, routing_weights, selected_experts, tp_size=2) - - torch.testing.assert_close(actual, expected) - assert calls == [ - {"a_shape": (10, 4), "b_shape": (3, 4, 6), "cumsum": (3, 6, 10)}, - {"a_shape": (10, 3), "b_shape": (3, 3, 4), "cumsum": (3, 6, 10)}, - {"a_shape": (10, 4), "b_shape": (3, 4, 6), "cumsum": (3, 6, 10)}, - {"a_shape": (10, 3), "b_shape": (3, 3, 4), "cumsum": (3, 6, 10)}, - ] - - -def test_sglang_moe_tp_sglang_mode_uses_sglang_weight_layout(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "sglang") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - - torch.manual_seed(0) - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - with torch.no_grad(): - experts.gate_up_proj.copy_(torch.randn_like(experts.gate_up_proj) * 0.1) - experts.down_proj.copy_(torch.randn_like(experts.down_proj) * 0.1) - - hidden_states = torch.randn(5, 4) - selected_experts = torch.tensor( - [ - [0, 1], - [2, 0], - [1, 2], - [2, 1], - [0, 2], - ], - dtype=torch.long, - ) - routing_weights = torch.tensor( - [ - [0.75, 0.25], - [0.60, 0.40], - [0.55, 0.45], - [0.80, 0.20], - [0.50, 0.50], - ], - dtype=hidden_states.dtype, - ) - - calls = [] - - def fake_fused_experts_impl( - hidden_states, - w1, - w2, - topk_weights, - topk_ids, - *, - activation, - filter_expert, - **kwargs, - ): - calls.append( - { - "w1_shape": tuple(w1.shape), - "w2_shape": tuple(w2.shape), - "activation": activation, - "gemm1_limit": kwargs["gemm1_limit"], - "filter_expert": filter_expert, - } - ) - shard_intermediate = w1.shape[1] // 2 - output = hidden_states.new_zeros(hidden_states.shape) - for expert_idx in range(w1.shape[0]): - mask = topk_ids == expert_idx - if not bool(mask.any().item()): - continue - token_rows, topk_slots = mask.nonzero(as_tuple=True) - tokens = hidden_states.index_select(0, token_rows) - gate_up = tokens.matmul(w1[expert_idx].transpose(0, 1)) - gate, up = gate_up.split(shard_intermediate, dim=-1) - activated = apply_moe_activation(activation, gate, up) - expert_out = activated.matmul(w2[expert_idx].transpose(0, 1)) - expert_out = expert_out * topk_weights[token_rows, topk_slots].unsqueeze(-1) - output.index_add_(0, token_rows, expert_out) - return output - - monkeypatch.setattr( - MoEExperts, - "_load_sglang_fused_experts_impl", - staticmethod(lambda: fake_fused_experts_impl), - ) - - actual = experts(hidden_states, routing_weights, selected_experts) - expected = _manual_sglang_tp_sim_direct(experts, hidden_states, routing_weights, selected_experts, tp_size=2) - - torch.testing.assert_close(actual, expected) - assert calls == [ - { - "w1_shape": (3, 6, 4), - "w2_shape": (3, 4, 3), - "activation": "silu", - "gemm1_limit": None, - "filter_expert": False, - }, - { - "w1_shape": (3, 6, 4), - "w2_shape": (3, 4, 3), - "activation": "silu", - "gemm1_limit": None, - "filter_expert": False, - }, - ] - - -def test_sglang_moe_tp_sglang_runner_mode_uses_runner_contract(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "sglang_runner") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - - torch.manual_seed(0) - experts = MoEExperts(num_experts=3, hidden_dim=4, intermediate_size=6, moe_implementation="eager") - with torch.no_grad(): - experts.gate_up_proj.copy_(torch.randn_like(experts.gate_up_proj) * 0.1) - experts.down_proj.copy_(torch.randn_like(experts.down_proj) * 0.1) - - hidden_states = torch.randn(5, 4) - selected_experts = torch.tensor( - [ - [0, 1], - [2, 0], - [1, 2], - [2, 1], - [0, 2], - ], - dtype=torch.long, - ) - routing_weights = torch.tensor( - [ - [0.75, 0.25], - [0.60, 0.40], - [0.55, 0.45], - [0.80, 0.20], - [0.50, 0.50], - ], - dtype=hidden_states.dtype, - ) - - calls = [] - - class FakeMoeRunnerBackend: - TRITON = "triton" - - class FakeMoeRunnerConfig(SimpleNamespace): - pass - - class FakeTritonMoeQuantInfo(SimpleNamespace): - pass - - class FakeStandardTopKOutput(SimpleNamespace): - pass - - class FakeStandardDispatchOutput(SimpleNamespace): - pass - - class FakeMoeRunner: - def __init__(self, backend, config): - calls.append( - { - "backend": backend, - "activation": config.activation, - "hidden_size": config.hidden_size, - "intermediate_size_per_partition": config.intermediate_size_per_partition, - "top_k": config.top_k, - "inplace": config.inplace, - "gemm1_clamp_limit": config.gemm1_clamp_limit, - "gate_up_interleaved": config.gate_up_interleaved, - } - ) - - def run(self, dispatch_output, quant_info): - hidden = dispatch_output.hidden_states - topk_weights = dispatch_output.topk_output.topk_weights - topk_ids = dispatch_output.topk_output.topk_ids - w13 = quant_info.w13_weight - w2 = quant_info.w2_weight - shard_intermediate = w13.shape[1] // 2 - output = hidden.new_zeros(hidden.shape) - for expert_idx in range(w13.shape[0]): - mask = topk_ids == expert_idx - if not bool(mask.any().item()): - continue - token_rows, topk_slots = mask.nonzero(as_tuple=True) - tokens = hidden.index_select(0, token_rows) - gate_up = tokens.matmul(w13[expert_idx].transpose(0, 1)) - gate, up = gate_up.split(shard_intermediate, dim=-1) - activated = apply_moe_activation("silu", gate, up) - expert_out = activated.matmul(w2[expert_idx].transpose(0, 1)) - expert_out = expert_out * topk_weights[token_rows, topk_slots].unsqueeze(-1) - output.index_add_(0, token_rows, expert_out) - return SimpleNamespace(hidden_states=output) - - monkeypatch.setattr( - MoEExperts, - "_load_sglang_moe_runner_stack", - staticmethod( - lambda: ( - FakeMoeRunner, - FakeMoeRunnerBackend, - FakeMoeRunnerConfig, - FakeTritonMoeQuantInfo, - FakeStandardDispatchOutput, - FakeStandardTopKOutput, - ) - ), - ) - - actual = experts(hidden_states, routing_weights, selected_experts) - expected = _manual_sglang_tp_sim_direct(experts, hidden_states, routing_weights, selected_experts, tp_size=2) - - torch.testing.assert_close(actual, expected) - assert calls == [ - { - "backend": "triton", - "activation": "silu", - "hidden_size": 4, - "intermediate_size_per_partition": 3, - "top_k": 2, - "inplace": False, - "gemm1_clamp_limit": None, - "gate_up_interleaved": False, - }, - { - "backend": "triton", - "activation": "silu", - "hidden_size": 4, - "intermediate_size_per_partition": 3, - "top_k": 2, - "inplace": False, - "gemm1_clamp_limit": None, - "gate_up_interleaved": False, - }, - ] - - -def test_eager_moe_block_uses_tp_sim_without_per_expert_bypass(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "1") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - - torch.manual_seed(1) - block = MoEBlock( - hidden_size=4, - num_experts=3, - top_k=2, - intermediate_size=6, - moe_implementation="eager", - ) - with torch.no_grad(): - block.experts.gate_up_proj.copy_(torch.randn_like(block.experts.gate_up_proj) * 0.1) - block.experts.down_proj.copy_(torch.randn_like(block.experts.down_proj) * 0.1) - hidden_states = torch.randn(2, 3, 4) - - flat_hidden = hidden_states.reshape(-1, hidden_states.shape[-1]) - routing_weights, selected_experts, _ = block.route(flat_hidden) - expected = _manual_sglang_tp_sim_direct( - block.experts, - flat_hidden, - routing_weights, - selected_experts, - tp_size=2, - ).reshape(hidden_states.shape) - - actual, _ = block(hidden_states) - - torch.testing.assert_close(actual, expected) - - -def test_sglang_moe_tp_sim_carry_shards_survives_moe_block_reshape(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "1") - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM_CARRY_SHARDS", "1") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - - torch.manual_seed(2) - block = MoEBlock( - hidden_size=4, - num_experts=3, - top_k=2, - intermediate_size=6, - moe_implementation="eager", - ) - with torch.no_grad(): - block.experts.gate_up_proj.copy_(torch.randn_like(block.experts.gate_up_proj) * 0.1) - block.experts.down_proj.copy_(torch.randn_like(block.experts.down_proj) * 0.1) - hidden_states = torch.randn(2, 3, 4) - - actual, _ = block(hidden_states) - carried_shards = getattr(actual, "_xorl_sglang_moe_tp_shards", None) - - assert carried_shards is not None - assert len(carried_shards) == 2 - assert all(shard.shape == actual.shape for shard in carried_shards) - expected = carried_shards[0] + carried_shards[1] - torch.testing.assert_close(actual, expected) - - -def test_sglang_moe_tp_sim_captures_flat_shards(monkeypatch): - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM", "1") - monkeypatch.setenv("XORL_SGLANG_MOE_TP_SIM_CARRY_SHARDS", "1") - monkeypatch.setattr(parallel_state_module, "get_parallel_state", lambda: _fake_tp_state(tp_size=2)) - - torch.manual_seed(3) - block = MoEBlock( - hidden_size=4, - num_experts=3, - top_k=2, - intermediate_size=6, - moe_implementation="eager", - ) - with torch.no_grad(): - block.experts.gate_up_proj.copy_(torch.randn_like(block.experts.gate_up_proj) * 0.1) - block.experts.down_proj.copy_(torch.randn_like(block.experts.down_proj) * 0.1) - captures = {} - block._diagnostic_capture_component = lambda name, tensor: captures.setdefault(name, tensor.detach().clone()) - hidden_states = torch.randn(2, 3, 4) - - actual, _ = block(hidden_states) - - assert "moe_experts_output_tp_shard_0" in captures - assert "moe_experts_output_tp_shard_1" in captures - assert "moe_experts_output_tp_shard_sum" in captures - assert captures["moe_experts_output_tp_shard_0"].shape == (6, 4) - assert captures["moe_experts_output_tp_shard_1"].shape == (6, 4) - torch.testing.assert_close( - captures["moe_experts_output_tp_shard_sum"], - captures["moe_experts_output_tp_shard_0"] + captures["moe_experts_output_tp_shard_1"], - ) - torch.testing.assert_close(captures["moe_experts_output_tp_shard_sum"], captures["moe_experts_output"]) - torch.testing.assert_close(actual, captures["moe_experts_output"].reshape_as(actual)) diff --git a/tests/models/test_moe_train_router_dispatch.py b/tests/models/test_moe_train_router_dispatch.py deleted file mode 100644 index a587a66c..00000000 --- a/tests/models/test_moe_train_router_dispatch.py +++ /dev/null @@ -1,125 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch -import torch.nn as nn - -from xorl.arguments import ModelArguments -from xorl.models.layers.moe.moe_block import MoEBlock -from xorl.models.layers.moe.router import TopKRouter -from xorl.server.server_arguments import ServerArguments - - -pytestmark = [pytest.mark.cpu] - - -def test_train_router_true_allowed_with_alltoall(): - moe = MoEBlock( - hidden_size=16, - num_experts=4, - top_k=2, - intermediate_size=32, - moe_implementation="eager", - train_router=True, - ) - moe.experts.ep_dispatch = "alltoall" - moe.train() - for p in moe.parameters(): - nn.init.normal_(p, std=0.01) - - x = torch.randn(1, 4, 16, requires_grad=True) - out, _ = moe(x) - out.sum().backward() - - assert moe.gate.weight.grad is not None - assert torch.isfinite(moe.gate.weight.grad).all() - assert moe.gate.weight.grad.abs().sum() > 0 - - -def test_train_router_true_rejected_with_deepep(): - moe = MoEBlock( - hidden_size=16, - num_experts=4, - top_k=2, - intermediate_size=32, - moe_implementation="eager", - train_router=True, - ) - moe.experts.ep_dispatch = "deepep" - - x = torch.randn(1, 4, 16) - with pytest.raises(AssertionError, match="ep_dispatch='deepep'"): - moe(x) - - -def test_model_arguments_default_train_router_false(): - args = ModelArguments(config_path="Qwen/Qwen3-8B") - - assert args.train_router is False - - -def test_server_arguments_default_train_router_false(): - args = ServerArguments(model_path="Qwen/Qwen3-8B") - - assert args.train_router is False - assert args.to_config_dict()["model"]["train_router"] is False - - -def test_from_config_defaults_train_router_false(): - config = SimpleNamespace( - hidden_size=16, - num_experts=4, - num_experts_per_tok=2, - moe_intermediate_size=32, - hidden_act="silu", - norm_topk_prob=True, - ) - - moe = MoEBlock.from_config(config, moe_implementation="eager") - - assert moe.train_router is False - - -def test_balanced_synthetic_routing_env(monkeypatch): - monkeypatch.setenv("XORL_MOE_SYNTHETIC_ROUTING", "balanced") - router = TopKRouter(num_experts=4, top_k=2) - - logits = torch.randn(8, 4) - routing_weights, selected_experts = router(logits, torch.bfloat16) - - expected_experts = torch.tensor( - [ - [0, 1], - [2, 3], - [0, 1], - [2, 3], - [0, 1], - [2, 3], - [0, 1], - [2, 3], - ] - ) - assert torch.equal(selected_experts, expected_experts) - assert routing_weights.dtype == torch.bfloat16 - torch.testing.assert_close(routing_weights.float(), torch.full((8, 2), 0.5)) - - counts = torch.bincount(selected_experts.flatten(), minlength=4) - assert torch.equal(counts, torch.full((4,), 4, dtype=counts.dtype)) - - -def test_balanced_synthetic_routing_replay_regather_uses_uniform_weights(monkeypatch): - monkeypatch.setenv("XORL_MOE_SYNTHETIC_ROUTING", "balanced") - moe = MoEBlock( - hidden_size=16, - num_experts=4, - top_k=2, - intermediate_size=32, - moe_implementation="eager", - ) - - router_logits = torch.randn(3, 4) - cached_experts = torch.tensor([[3, 2], [1, 0], [2, 1]]) - selected_experts, routing_weights = moe._regather_routing(router_logits, cached_experts, torch.float32) - - assert torch.equal(selected_experts, cached_experts) - torch.testing.assert_close(routing_weights, torch.full((3, 2), 0.5)) diff --git a/tests/models/test_moe_weight_auto_merge.py b/tests/models/test_moe_weight_auto_merge.py deleted file mode 100644 index 6986832d..00000000 --- a/tests/models/test_moe_weight_auto_merge.py +++ /dev/null @@ -1,373 +0,0 @@ -""" -Tests for MoE weight auto-merging functionality. - -Tests the ExpertWeightBuffer class and related functions that handle -automatic merging of per-expert HuggingFace weights into stacked format -during model loading. - -Note: This test file re-implements the classes/functions under test to avoid -import dependency issues. The actual implementation in module_utils.py should -be kept in sync with these copies. -""" - -import re -from collections import defaultdict -from typing import Dict, List, Optional, Set, Tuple - -import pytest -import torch - - -pytestmark = [pytest.mark.cpu] - - -# ============================================================================= -# Copy of the implementation for testing (avoids heavy import dependencies) -# ============================================================================= - -_EXPERT_KEY_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$") - -_FUSED_EXPERT_PATTERN = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate|up|down)_proj$") - - -def parse_expert_key(key: str) -> Optional[Tuple[int, int, str]]: - match = _EXPERT_KEY_PATTERN.match(key) - if match: - return int(match.group(1)), int(match.group(2)), match.group(3) - return None - - -def _model_needs_expert_merging(parameter_names: Set[str]) -> bool: - for name in parameter_names: - if _FUSED_EXPERT_PATTERN.match(name): - return True - return False - - -class ExpertWeightBuffer: - def __init__(self, num_experts: int): - self.num_experts = num_experts - self._stacked_buffers: Dict[Tuple[int, str], torch.Tensor] = {} - self._filled_experts: Dict[Tuple[int, str], Set[int]] = defaultdict(set) - - def add(self, layer_idx: int, expert_idx: int, proj: str, tensor: torch.Tensor) -> None: - key = (layer_idx, proj) - if key not in self._stacked_buffers: - stacked_shape = (self.num_experts,) + tensor.shape - self._stacked_buffers[key] = torch.empty(stacked_shape, dtype=tensor.dtype, device="cpu") - self._stacked_buffers[key][expert_idx].copy_(tensor) - self._filled_experts[key].add(expert_idx) - - def is_complete(self, layer_idx: int, proj: str) -> bool: - key = (layer_idx, proj) - return len(self._filled_experts.get(key, set())) == self.num_experts - - def pop_stacked(self, layer_idx: int, proj: str) -> torch.Tensor: - key = (layer_idx, proj) - if key not in self._stacked_buffers: - raise KeyError(f"No buffered experts for layer {layer_idx}, projection {proj}") - filled = self._filled_experts.pop(key) - if len(filled) != self.num_experts: - raise ValueError( - f"Incomplete experts for layer {layer_idx}, {proj}_proj: got {len(filled)}, expected {self.num_experts}" - ) - return self._stacked_buffers.pop(key) - - @staticmethod - def get_fused_name(layer_idx: int, proj: str) -> str: - return f"model.layers.{layer_idx}.mlp.experts.{proj}_proj" - - def get_pending_keys(self) -> List[Tuple[int, str]]: - return list(self._stacked_buffers.keys()) - - def get_pending_counts(self) -> Dict[Tuple[int, str], int]: - return {key: len(experts) for key, experts in self._filled_experts.items()} - - -# ============================================================================= -# Tests -# ============================================================================= - - -class TestParseAndMergingDetection: - """Tests for parse_expert_key and _model_needs_expert_merging.""" - - def test_parse_expert_key_and_merging_detection(self): - """Test valid/invalid key parsing, fused/non-MoE/empty detection.""" - # Valid keys - assert parse_expert_key("model.layers.0.mlp.experts.5.gate_proj.weight") == (0, 5, "gate") - assert parse_expert_key("model.layers.10.mlp.experts.127.up_proj.weight") == (10, 127, "up") - assert parse_expert_key("model.layers.35.mlp.experts.0.down_proj.weight") == (35, 0, "down") - - # Non-expert keys return None - for key in [ - "model.layers.0.self_attn.q_proj.weight", - "model.layers.0.mlp.gate_proj.weight", - "model.embed_tokens.weight", - "lm_head.weight", - ]: - assert parse_expert_key(key) is None, f"Expected None for {key}" - - # Already-fused format keys return None - for key in [ - "model.layers.0.mlp.experts.gate_proj", - "model.layers.0.mlp.experts.up_proj", - "model.layers.0.mlp.experts.down_proj", - ]: - assert parse_expert_key(key) is None, f"Expected None for {key}" - - # Invalid projection name and missing suffix - assert parse_expert_key("model.layers.0.mlp.experts.5.other_proj.weight") is None - assert parse_expert_key("model.layers.0.mlp.experts.5.gate_proj") is None - - # Fused format model needs merging - assert ( - _model_needs_expert_merging( - { - "model.layers.0.mlp.experts.gate_proj", - "model.layers.0.mlp.experts.up_proj", - "model.layers.0.mlp.experts.down_proj", - "model.layers.0.self_attn.q_proj.weight", - } - ) - is True - ) - - # Non-MoE model does not - assert ( - _model_needs_expert_merging( - { - "model.layers.0.mlp.gate_proj.weight", - "model.layers.0.mlp.up_proj.weight", - "model.layers.0.mlp.down_proj.weight", - "model.layers.0.self_attn.q_proj.weight", - } - ) - is False - ) - - # Empty - assert _model_needs_expert_merging(set()) is False - - -class TestExpertWeightBuffer: - """Tests for ExpertWeightBuffer: lifecycle, layers, projections, errors, edge cases.""" - - def test_buffer_lifecycle_layers_errors_and_edge_cases(self): - """Test full lifecycle, independent layers/projections, errors, overwrite, - single/large expert count, fused name, CPU storage, and full workflow.""" - num_experts = 4 - buffer = ExpertWeightBuffer(num_experts=num_experts) - - # Empty buffer - assert len(buffer.get_pending_keys()) == 0 - - # Add experts in reverse order - tensors = {} - for i in range(num_experts - 1, -1, -1): - tensor = torch.full((2, 2), fill_value=float(i)) - tensors[i] = tensor - buffer.add(0, i, "gate", tensor) - if i > 0: - assert not buffer.is_complete(0, "gate") - - assert buffer.is_complete(0, "gate") - - # Pop and verify order preserved - stacked = buffer.pop_stacked(0, "gate") - assert stacked.shape == (num_experts, 2, 2) - for i in range(num_experts): - assert torch.all(stacked[i] == float(i)), f"Expert {i} not in correct position" - - assert (0, "gate") not in buffer.get_pending_counts() - assert not buffer.is_complete(0, "gate") - - # Independent layers and projections - buffer2 = ExpertWeightBuffer(num_experts=2) - buffer2.add(0, 0, "gate", torch.randn(4, 4)) - buffer2.add(0, 1, "gate", torch.randn(4, 4)) - buffer2.add(1, 0, "gate", torch.randn(4, 4)) - assert buffer2.is_complete(0, "gate") - assert not buffer2.is_complete(1, "gate") - - buffer2.add(0, 0, "up", torch.randn(4, 4)) - assert buffer2.is_complete(0, "gate") - assert not buffer2.is_complete(0, "up") - - pending_keys = buffer2.get_pending_keys() - assert (0, "gate") in pending_keys - assert (1, "gate") in pending_keys - assert (0, "up") in pending_keys - - # Pop nonexistent key - buffer3 = ExpertWeightBuffer(num_experts=4) - with pytest.raises(KeyError): - buffer3.pop_stacked(0, "gate") - - # Pop incomplete - buffer3.add(0, 0, "gate", torch.randn(4, 4)) - buffer3.add(0, 1, "gate", torch.randn(4, 4)) - with pytest.raises(ValueError, match="Incomplete experts"): - buffer3.pop_stacked(0, "gate") - - # Overwrite expert - buffer4 = ExpertWeightBuffer(num_experts=2) - buffer4.add(0, 0, "gate", torch.full((2, 2), 1.0)) - buffer4.add(0, 0, "gate", torch.full((2, 2), 2.0)) - buffer4.add(0, 1, "gate", torch.full((2, 2), 1.0)) - stacked = buffer4.pop_stacked(0, "gate") - assert torch.all(stacked[0] == 2.0) - - # Single expert edge case - buffer5 = ExpertWeightBuffer(num_experts=1) - tensor = torch.randn(4, 4) - buffer5.add(0, 0, "gate", tensor) - assert buffer5.is_complete(0, "gate") - stacked = buffer5.pop_stacked(0, "gate") - assert stacked.shape == (1, 4, 4) - assert torch.allclose(stacked[0].cpu(), tensor.cpu()) - - # Large number of experts - large_num = 128 - buffer6 = ExpertWeightBuffer(num_experts=large_num) - for i in range(large_num): - buffer6.add(0, i, "gate", torch.randn(4, 4)) - assert buffer6.is_complete(0, "gate") - stacked = buffer6.pop_stacked(0, "gate") - assert stacked.shape == (large_num, 4, 4) - - # get_fused_name - assert ExpertWeightBuffer.get_fused_name(0, "gate") == "model.layers.0.mlp.experts.gate_proj" - assert ExpertWeightBuffer.get_fused_name(10, "up") == "model.layers.10.mlp.experts.up_proj" - assert ExpertWeightBuffer.get_fused_name(35, "down") == "model.layers.35.mlp.experts.down_proj" - - # CPU storage - if torch.cuda.is_available(): - buffer7 = ExpertWeightBuffer(num_experts=2) - t = torch.randn(4, 4, device="cuda") - buffer7.add(0, 0, "gate", t) - buffer7.add(0, 1, "gate", t) - stacked = buffer7.pop_stacked(0, "gate") - assert stacked.device.type == "cpu" - - # Full workflow - num_e = 8 - num_layers = 2 - hidden_size = 32 - intermediate_size = 64 - buffer8 = ExpertWeightBuffer(num_experts=num_e) - for layer_idx in range(num_layers): - for proj in ["gate", "up", "down"]: - shape = (hidden_size, intermediate_size) if proj == "down" else (intermediate_size, hidden_size) - for expert_idx in range(num_e): - buffer8.add(layer_idx, expert_idx, proj, torch.randn(*shape)) - if expert_idx == num_e - 1: - assert buffer8.is_complete(layer_idx, proj) - stacked = buffer8.pop_stacked(layer_idx, proj) - assert stacked.shape == (num_e, *shape) - assert len(buffer8.get_pending_keys()) == 0 - - -# ============================================================================= -# Tests for checkpoint format detection and loading logic -# ============================================================================= - - -def _checkpoint_has_per_expert_weights(checkpoint_keys): - for key in checkpoint_keys: - if _EXPERT_KEY_PATTERN.match(key): - return True - return False - - -class TestCheckpointFormatAndLoading: - """Tests for checkpoint format detection and loading both formats.""" - - def test_checkpoint_format_detection_and_loading(self): - """Test detection of per-expert/fused/non-MoE/empty/mixed formats, and loading both.""" - # Per-expert format - assert ( - _checkpoint_has_per_expert_weights( - { - "model.layers.0.mlp.experts.0.gate_proj.weight", - "model.layers.0.mlp.experts.1.gate_proj.weight", - "model.layers.0.self_attn.q_proj.weight", - } - ) - is True - ) - - # Fused format - assert ( - _checkpoint_has_per_expert_weights( - { - "model.layers.0.mlp.experts.gate_proj", - "model.layers.0.mlp.experts.up_proj", - } - ) - is False - ) - - # Non-MoE - assert ( - _checkpoint_has_per_expert_weights( - { - "model.layers.0.mlp.gate_proj.weight", - "model.layers.0.self_attn.q_proj.weight", - } - ) - is False - ) - - # Empty - assert _checkpoint_has_per_expert_weights(set()) is False - - # Mixed - assert ( - _checkpoint_has_per_expert_weights( - { - "model.layers.0.mlp.experts.gate_proj", - "model.layers.1.mlp.experts.0.gate_proj.weight", - } - ) - is True - ) - - # --- Loading both formats --- - model_params = { - "model.layers.0.mlp.experts.gate_proj", - "model.layers.0.mlp.experts.up_proj", - "model.layers.0.mlp.experts.down_proj", - "model.layers.0.self_attn.q_proj.weight", - } - - # Fused checkpoint matches directly - fused_checkpoint_keys = { - "model.layers.0.mlp.experts.gate_proj", - "model.layers.0.mlp.experts.up_proj", - "model.layers.0.mlp.experts.down_proj", - "model.layers.0.self_attn.q_proj.weight", - } - assert _model_needs_expert_merging(model_params) is True - assert _checkpoint_has_per_expert_weights(fused_checkpoint_keys) is False - for key in fused_checkpoint_keys: - assert key in model_params - - # Per-expert checkpoint needs merging - per_expert_keys = { - "model.layers.0.mlp.experts.0.gate_proj.weight", - "model.layers.0.mlp.experts.1.gate_proj.weight", - "model.layers.0.mlp.experts.0.up_proj.weight", - "model.layers.0.mlp.experts.1.up_proj.weight", - } - assert _checkpoint_has_per_expert_weights(per_expert_keys) is True - for key in per_expert_keys: - assert key not in model_params - - # After merging, fused names match model params - for key in per_expert_keys: - parsed = parse_expert_key(key) - if parsed: - layer_idx, _, proj = parsed - fused_name = ExpertWeightBuffer.get_fused_name(layer_idx, proj) - assert fused_name in model_params diff --git a/tests/models/test_moe_weight_loading_integration.py b/tests/models/test_moe_weight_loading_integration.py deleted file mode 100644 index 6b00a675..00000000 --- a/tests/models/test_moe_weight_loading_integration.py +++ /dev/null @@ -1,337 +0,0 @@ -""" -Integration tests for MoE weight auto-merging during model loading. - -These tests simulate the full flow of loading per-expert HuggingFace weights -into a model that expects fused (stacked) expert format. -""" - -import random -from typing import Dict, Iterator, Tuple - -import pytest -import torch -import torch.nn as nn - - -pytestmark = [pytest.mark.cpu] - -# Re-implement the core logic to test without full xorl dependencies -import re -from collections import defaultdict -from typing import Optional, Set - - -# ============================================================================= -# Copy of core implementation for testing -# ============================================================================= - -_EXPERT_KEY_PATTERN = re.compile(r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$") - -_FUSED_EXPERT_PATTERN = re.compile(r"^model\.layers\.\d+\.mlp\.experts\.(gate|up|down)_proj$") - - -def parse_expert_key(key: str) -> Optional[Tuple[int, int, str]]: - match = _EXPERT_KEY_PATTERN.match(key) - if match: - return int(match.group(1)), int(match.group(2)), match.group(3) - return None - - -def _model_needs_expert_merging(parameter_names: Set[str]) -> bool: - for name in parameter_names: - if _FUSED_EXPERT_PATTERN.match(name): - return True - return False - - -class ExpertWeightBuffer: - """Buffer with streaming copy (copy on add, not batch copy).""" - - def __init__(self, num_experts: int): - self.num_experts = num_experts - self._stacked_buffers: Dict[Tuple[int, str], torch.Tensor] = {} - self._filled_experts: Dict[Tuple[int, str], Set[int]] = defaultdict(set) - - def add(self, layer_idx: int, expert_idx: int, proj: str, tensor: torch.Tensor) -> None: - key = (layer_idx, proj) - # Transpose from HF nn.Linear [out, in] to (K, N) = [in, out] format - tensor = tensor.t().contiguous() - if key not in self._stacked_buffers: - stacked_shape = (self.num_experts,) + tensor.shape - self._stacked_buffers[key] = torch.empty(stacked_shape, dtype=tensor.dtype, device="cpu") - # Copy directly into the slice (streaming) - self._stacked_buffers[key][expert_idx].copy_(tensor) - self._filled_experts[key].add(expert_idx) - - def is_complete(self, layer_idx: int, proj: str) -> bool: - key = (layer_idx, proj) - return len(self._filled_experts.get(key, set())) == self.num_experts - - def pop_stacked(self, layer_idx: int, proj: str) -> torch.Tensor: - key = (layer_idx, proj) - if key not in self._stacked_buffers: - raise KeyError(f"No buffered experts for layer {layer_idx}, projection {proj}") - filled = self._filled_experts.pop(key) - if len(filled) != self.num_experts: - raise ValueError( - f"Incomplete experts for layer {layer_idx}, {proj}_proj: got {len(filled)}, expected {self.num_experts}" - ) - return self._stacked_buffers.pop(key) - - @staticmethod - def get_fused_name(layer_idx: int, proj: str) -> str: - return f"model.layers.{layer_idx}.mlp.experts.{proj}_proj" - - def get_pending_counts(self) -> Dict[Tuple[int, str], int]: - return {key: len(experts) for key, experts in self._filled_experts.items()} - - -# ============================================================================= -# Mock model classes for testing -# ============================================================================= - - -class MockFusedMoeExperts(nn.Module): - """Mock MoE experts module with fused/stacked weight format.""" - - def __init__(self, num_experts: int, hidden_size: int, intermediate_size: int): - super().__init__() - self.num_experts = num_experts - # Fused format (G, K, N): [num_experts, in_features, out_features] - self.gate_proj = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size)) - self.up_proj = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size)) - self.down_proj = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) - - -class MockMoeLayer(nn.Module): - """Mock MoE layer.""" - - def __init__(self, num_experts: int, hidden_size: int, intermediate_size: int): - super().__init__() - self.experts = MockFusedMoeExperts(num_experts, hidden_size, intermediate_size) - - -class MockModelInner(nn.Module): - """Inner model component that holds the layers.""" - - def __init__(self, num_layers: int, num_experts: int, hidden_size: int, intermediate_size: int): - super().__init__() - self.layers = nn.ModuleList( - [ - nn.ModuleDict({"mlp": MockMoeLayer(num_experts, hidden_size, intermediate_size)}) - for _ in range(num_layers) - ] - ) - - -class MockMoeModel(nn.Module): - """Mock model with MoE layers expecting fused expert format. - - Structure matches HuggingFace: model.layers.{i}.mlp.experts.{proj}_proj - """ - - def __init__(self, num_layers: int, num_experts: int, hidden_size: int, intermediate_size: int): - super().__init__() - self.num_experts = num_experts - # Use 'model' as the attribute name to match HuggingFace structure - self.model = MockModelInner(num_layers, num_experts, hidden_size, intermediate_size) - - class Config: - def __init__(self, num_experts): - self.num_experts = num_experts - - @property - def config(self): - return self.Config(self.num_experts) - - -# ============================================================================= -# Test utilities -# ============================================================================= - - -def create_per_expert_state_dict( - num_layers: int, - num_experts: int, - hidden_size: int, - intermediate_size: int, -) -> Dict[str, torch.Tensor]: - """Create a state dict in per-expert HuggingFace format.""" - state_dict = {} - for layer_idx in range(num_layers): - for expert_idx in range(num_experts): - # gate_proj and up_proj: [intermediate_size, hidden_size] - state_dict[f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.gate_proj.weight"] = torch.randn( - intermediate_size, hidden_size - ) - state_dict[f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.up_proj.weight"] = torch.randn( - intermediate_size, hidden_size - ) - # down_proj: [hidden_size, intermediate_size] - state_dict[f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.down_proj.weight"] = torch.randn( - hidden_size, intermediate_size - ) - return state_dict - - -def simulate_load_model_weights( - model: nn.Module, - state_dict_iterator: Iterator[Tuple[str, torch.Tensor]], - num_experts: int, -) -> Dict[str, torch.Tensor]: - """ - Simulate the load_model_weights function with auto-merge logic. - - Returns the loaded state dict (fused format) for verification. - """ - parameter_names_to_load = {name for name, _ in model.named_parameters()} - loaded_weights = {} - - # Check if model needs merging - expert_buffer = None - if _model_needs_expert_merging(parameter_names_to_load): - expert_buffer = ExpertWeightBuffer(num_experts) - - for name, tensor in state_dict_iterator: - # Check if this is a per-expert key that needs merging - if expert_buffer is not None: - parsed = parse_expert_key(name) - if parsed is not None: - layer_idx, expert_idx, proj = parsed - expert_buffer.add(layer_idx, expert_idx, proj, tensor) - - # If all experts collected, stack and dispatch - if expert_buffer.is_complete(layer_idx, proj): - fused_name = ExpertWeightBuffer.get_fused_name(layer_idx, proj) - stacked = expert_buffer.pop_stacked(layer_idx, proj) - - if fused_name in parameter_names_to_load: - parameter_names_to_load.remove(fused_name) - loaded_weights[fused_name] = stacked - continue - - # Normal key handling - if name in parameter_names_to_load: - parameter_names_to_load.remove(name) - loaded_weights[name] = tensor - - # Check for incomplete buffers - if expert_buffer is not None: - pending = expert_buffer.get_pending_counts() - if pending: - raise RuntimeError(f"Incomplete expert weights: {pending}") - - return loaded_weights - - -# ============================================================================= -# Integration Tests -# ============================================================================= - - -class TestMoeWeightLoadingIntegration: - """Integration tests for MoE weight loading with auto-merge.""" - - def test_load_and_verify_shapes_and_order(self): - """Test loading per-expert weights, verifying fused shapes, expert order, and projection sizes.""" - num_layers = 2 - num_experts = 4 - hidden_size = 32 - intermediate_size = 128 - - model = MockMoeModel(num_layers, num_experts, hidden_size, intermediate_size) - per_expert_state_dict = create_per_expert_state_dict(num_layers, num_experts, hidden_size, intermediate_size) - - loaded_weights = simulate_load_model_weights(model, iter(per_expert_state_dict.items()), num_experts) - - # All fused parameters created with correct shapes - for layer_idx in range(num_layers): - for proj in ["gate", "up", "down"]: - param_name = f"model.layers.{layer_idx}.mlp.experts.{proj}_proj" - assert param_name in loaded_weights, f"Missing {param_name}" - - gate = loaded_weights[f"model.layers.{layer_idx}.mlp.experts.gate_proj"] - up = loaded_weights[f"model.layers.{layer_idx}.mlp.experts.up_proj"] - down = loaded_weights[f"model.layers.{layer_idx}.mlp.experts.down_proj"] - - # gate and up: [num_experts, hidden_size, intermediate_size] - assert gate.shape == (num_experts, hidden_size, intermediate_size) - assert up.shape == (num_experts, hidden_size, intermediate_size) - # down: [num_experts, intermediate_size, hidden_size] - assert down.shape == (num_experts, intermediate_size, hidden_size) - - # Expert order preserved with identifiable values - model2 = MockMoeModel(1, num_experts, 8, 16) - state_dict2 = {} - for expert_idx in range(num_experts): - state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.gate_proj.weight"] = torch.full( - (16, 8), float(expert_idx) - ) - state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.up_proj.weight"] = torch.randn(16, 8) - state_dict2[f"model.layers.0.mlp.experts.{expert_idx}.down_proj.weight"] = torch.randn(8, 16) - - items = list(state_dict2.items()) - random.seed(42) - random.shuffle(items) - - loaded2 = simulate_load_model_weights(model2, iter(items), num_experts) - gate_proj = loaded2["model.layers.0.mlp.experts.gate_proj"] - for expert_idx in range(num_experts): - assert torch.all(gate_proj[expert_idx] == float(expert_idx)), f"Expert {expert_idx} not in correct position" - - def test_streaming_and_edge_cases(self): - """Test sharded streaming load, large expert count, single expert, and random order.""" - # Streaming across shards - num_layers = 2 - num_experts = 8 - hidden_size = 16 - intermediate_size = 32 - - model = MockMoeModel(num_layers, num_experts, hidden_size, intermediate_size) - state_dict = create_per_expert_state_dict(num_layers, num_experts, hidden_size, intermediate_size) - items = list(state_dict.items()) - shard1 = [(k, v) for k, v in items if any(f".experts.{i}." in k for i in range(4))] - shard2 = [(k, v) for k, v in items if any(f".experts.{i}." in k for i in range(4, 8))] - combined_iterator = iter(shard1 + shard2) - - loaded_weights = simulate_load_model_weights(model, combined_iterator, num_experts) - for layer_idx in range(num_layers): - gate = loaded_weights[f"model.layers.{layer_idx}.mlp.experts.gate_proj"] - assert gate.shape == (num_experts, hidden_size, intermediate_size) - - # Large expert count (128) - model_large = MockMoeModel(1, 128, 16, 32) - state_dict_large = create_per_expert_state_dict(1, 128, 16, 32) - loaded_large = simulate_load_model_weights(model_large, iter(state_dict_large.items()), 128) - assert loaded_large["model.layers.0.mlp.experts.gate_proj"].shape == (128, 16, 32) - - # Single expert - model_single = MockMoeModel(1, 1, 8, 16) - state_dict_single = create_per_expert_state_dict(1, 1, 8, 16) - loaded_single = simulate_load_model_weights(model_single, iter(state_dict_single.items()), 1) - assert loaded_single["model.layers.0.mlp.experts.gate_proj"].shape == (1, 8, 16) - - # Random order loading - - model_rand = MockMoeModel(2, 4, 8, 16) - state_dict_rand = create_per_expert_state_dict(2, 4, 8, 16) - items_rand = list(state_dict_rand.items()) - random.seed(123) - random.shuffle(items_rand) - loaded_rand = simulate_load_model_weights(model_rand, iter(items_rand), 4) - for layer_idx in range(2): - for proj in ["gate", "up", "down"]: - assert f"model.layers.{layer_idx}.mlp.experts.{proj}_proj" in loaded_rand - - def test_incomplete_experts_raises_error(self): - """Test that incomplete expert set raises an error.""" - model = MockMoeModel(1, 4, 8, 16) - - state_dict = {} - for expert_idx in range(3): # Missing expert 3 - state_dict[f"model.layers.0.mlp.experts.{expert_idx}.gate_proj.weight"] = torch.randn(16, 8) - state_dict[f"model.layers.0.mlp.experts.{expert_idx}.up_proj.weight"] = torch.randn(16, 8) - state_dict[f"model.layers.0.mlp.experts.{expert_idx}.down_proj.weight"] = torch.randn(8, 16) - - with pytest.raises(RuntimeError, match="Incomplete expert weights"): - simulate_load_model_weights(model, iter(state_dict.items()), 4) diff --git a/tests/models/test_mtp_low_precision_policy.py b/tests/models/test_mtp_low_precision_policy.py deleted file mode 100644 index 990a077d..00000000 --- a/tests/models/test_mtp_low_precision_policy.py +++ /dev/null @@ -1,49 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from xorl.models.module_utils import _matches_checkpoint_skip_key_pattern -from xorl.models.transformers.glm4_moe.checkpoint_handler import Glm4MoeCheckpointHandler -from xorl.models.transformers.qwen3_5_shared import QWEN3_5_CHECKPOINT_SKIP_KEY_PATTERNS - - -pytestmark = [pytest.mark.cpu] - - -def test_qwen35_checkpoint_policy_skips_top_level_mtp_keys(): - model = SimpleNamespace(_checkpoint_skip_key_patterns=QWEN3_5_CHECKPOINT_SKIP_KEY_PATTERNS) - - assert _matches_checkpoint_skip_key_pattern("mtp.layers.0.mlp.gate_proj.weight", model) - assert _matches_checkpoint_skip_key_pattern("mtp.pre_fc_norm_embedding.weight", model) - assert not _matches_checkpoint_skip_key_pattern("model.layers.0.mlp.gate_proj.weight", model) - - -def test_glm4_moe_checkpoint_policy_remaps_only_shared_mtp_tail(): - tensor = torch.randn(2, 3) - handler = Glm4MoeCheckpointHandler( - num_experts=2, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=4, - num_hidden_layers=3, - ) - - remapped = handler.on_load_weight("model.layers.3.embed_tokens.weight", tensor) - assert len(remapped) == 1 - assert remapped[0][0] == "model.embed_tokens.weight" - torch.testing.assert_close(remapped[0][1], tensor) - - remapped = handler.on_load_weight("model.layers.3.shared_head.norm.weight", tensor) - assert len(remapped) == 1 - assert remapped[0][0] == "model.norm.weight" - torch.testing.assert_close(remapped[0][1], tensor) - - remapped = handler.on_load_weight("model.layers.3.shared_head.head.weight", tensor) - assert len(remapped) == 1 - assert remapped[0][0] == "lm_head.weight" - torch.testing.assert_close(remapped[0][1], tensor) - - assert handler.on_load_weight("model.layers.3.eh_proj.weight", tensor) == [] - assert handler.on_load_weight("model.layers.3.enorm.weight", tensor) == [] - assert handler.on_load_weight("model.layers.3.transformer_layer.self_attn.q_proj.weight", tensor) == [] diff --git a/tests/models/test_nemotron_h_checkpoint.py b/tests/models/test_nemotron_h_checkpoint.py index 5a1f5328..5435dd4e 100644 --- a/tests/models/test_nemotron_h_checkpoint.py +++ b/tests/models/test_nemotron_h_checkpoint.py @@ -108,18 +108,22 @@ def _run_handler(handler: NemotronHCheckpointHandler, checkpoint: dict[str, torc return loaded -@pytest.mark.parametrize( - ("num_experts", "ep_rank", "ep_size"), - [(7, 0, 2), (8, 2, 2), (8, 0, 0)], -) -def test_checkpoint_handler_rejects_invalid_expert_parallelism(num_experts, ep_rank, ep_size): - with pytest.raises(ValueError): - NemotronHCheckpointHandler(num_experts=num_experts, ep_rank=ep_rank, ep_size=ep_size) +def _assert_nemotron_h_ep_checkpoint_ownership_policy(): + for num_experts, ep_rank, ep_size in ((7, 0, 2), (8, 2, 2), (8, 0, 0)): + with pytest.raises(ValueError): + NemotronHCheckpointHandler(num_experts=num_experts, ep_rank=ep_rank, ep_size=ep_size) + + _assert_ep_skip_key_policy() + _assert_ep_aware_loading_slices_and_counts_skips() + _assert_ep_plan_targets_expert_params() + +def test_nemotron_h_published_layout_load_save_and_hf_parity_contract(): + _assert_nemotron_h_ep_checkpoint_ownership_policy() -def test_nemotron_h_hf_parity_and_strict_load_accounting(): hf_model = _build_hf_model() - checkpoint = _published_checkpoint_layout(hf_model) + published_checkpoint = _published_checkpoint_layout(hf_model) + checkpoint = dict(published_checkpoint) checkpoint["mtp.layers.0.mixer.q_proj.weight"] = torch.zeros(2, 2) # must be ignored model = _build_xorl_model(hf_model.config) @@ -147,26 +151,19 @@ def test_nemotron_h_hf_parity_and_strict_load_accounting(): torch.testing.assert_close(outputs.last_hidden_state.float(), hf_hidden.float(), atol=1e-4, rtol=1e-4) torch.testing.assert_close(logits.float(), hf_logits.float(), atol=1e-4, rtol=1e-4) - -def test_nemotron_h_save_round_trips_to_published_layout(): - hf_model = _build_hf_model() - checkpoint = _published_checkpoint_layout(hf_model) - - model = _build_xorl_model(hf_model.config) - handler = model.get_checkpoint_handler(checkpoint_keys=set(checkpoint)) - model.load_state_dict(_run_handler(handler, checkpoint), strict=True) - saved = {} for name, tensor in model.state_dict().items(): for key, out_tensor in handler.on_save_weight(name, tensor): saved[key] = out_tensor - assert set(saved) == set(checkpoint) - for key in checkpoint: - torch.testing.assert_close(saved[key], checkpoint[key], atol=0.0, rtol=0.0) + assert set(saved) == set(published_checkpoint) + for key in published_checkpoint: + torch.testing.assert_close(saved[key], published_checkpoint[key], atol=0.0, rtol=0.0) + + _assert_nemotron_h_handler_accepts_hf_stacked_expert_layout() -def test_nemotron_h_handler_accepts_hf_stacked_expert_layout(): +def _assert_nemotron_h_handler_accepts_hf_stacked_expert_layout(): """The transformers 5.x in-memory format stores experts as stacked 3D [E, out, in].""" hf_model = _build_hf_model() @@ -184,7 +181,7 @@ def test_nemotron_h_handler_accepts_hf_stacked_expert_layout(): ) -def test_nemotron_h_ep_skip_key_fn(): +def _assert_ep_skip_key_policy(): handler = NemotronHCheckpointHandler(num_experts=NUM_EXPERTS, ep_rank=0, ep_size=2) skip_fn = handler.get_skip_key_fn() assert skip_fn is not None @@ -200,7 +197,7 @@ def test_nemotron_h_ep_skip_key_fn(): assert NemotronHCheckpointHandler(num_experts=NUM_EXPERTS).get_skip_key_fn() is None -def test_nemotron_h_ep_aware_loading_slices_and_counts_skips(): +def _assert_ep_aware_loading_slices_and_counts_skips(): hf_model = _build_hf_model() checkpoint = _published_checkpoint_layout(hf_model) @@ -224,7 +221,7 @@ def test_nemotron_h_ep_aware_loading_slices_and_counts_skips(): assert torch.equal(gate_up, stacked_up[local:].transpose(1, 2).contiguous()) -def test_nemotron_h_ep_plan_targets_expert_params(): +def _assert_ep_plan_targets_expert_params(): plan = get_ep_plan() assert isinstance(plan, ParallelPlan) assert plan._is_expert_parameter("model.layers.2.mixer.experts.gate_up_proj") diff --git a/tests/models/test_nemotron_h_model.py b/tests/models/test_nemotron_h_model.py index fe4765c9..2561cbfb 100644 --- a/tests/models/test_nemotron_h_model.py +++ b/tests/models/test_nemotron_h_model.py @@ -1,7 +1,11 @@ +import json + import pytest import torch +from xorl.models.auto import _load_local_xorl_config from xorl.models.module_utils import compute_loss +from xorl.models.registry import get_registry from xorl.models.transformers.nemotron_h.configuration_nemotron_h import NemotronHConfig from xorl.models.transformers.nemotron_h.modeling_nemotron_h import NemotronHForCausalLM @@ -49,7 +53,8 @@ def _build_model() -> NemotronHForCausalLM: return NemotronHForCausalLM(_tiny_config()) -def test_nemotron_h_forward_shape_and_router_logits(): +def test_nemotron_h_runtime_and_gradient_checkpointing_contract(tmp_path): + _assert_nemotron_h_registry_and_local_config_policy(tmp_path) model = _build_model() model.eval() input_ids = torch.randint(0, model.config.vocab_size, (2, SEQ_LEN)) @@ -63,14 +68,24 @@ def test_nemotron_h_forward_shape_and_router_logits(): assert len(outputs.router_logits) == num_moe_layers assert outputs.router_logits[0].shape == (2 * SEQ_LEN, model.config.n_routed_experts) - -def test_nemotron_h_backward_reaches_all_mixer_types(): - model = _build_model() model.train() - input_ids = torch.randint(0, model.config.vocab_size, (2, SEQ_LEN)) - - outputs = model(input_ids=input_ids) - outputs.last_hidden_state.float().pow(2).mean().backward() + input_ids = input_ids[:1] + labels = torch.randint(0, model.config.vocab_size, (1, SEQ_LEN)) + labels[:, :3] = -100 + cu_seqlens = torch.tensor([0, 7, SEQ_LEN], dtype=torch.int32) + outputs = model(input_ids=input_ids, cu_seq_lens_q=cu_seqlens, cu_seq_lens_k=cu_seqlens) + assert torch.isfinite(outputs.last_hidden_state).all() + result = compute_loss( + model.lm_head, + outputs.last_hidden_state, + loss_fn_name=None, + loss_fn_inputs={"labels": labels}, + loss_fn_params={"ce_mode": "eager"}, + logits_to_keep=0, + ) + assert result.loss.ndim == 0 + assert torch.isfinite(result.loss) + result.loss.backward() layers = model.model.layers grads = { @@ -96,30 +111,68 @@ def test_nemotron_h_backward_reaches_all_mixer_types(): assert torch.isfinite(grad).all(), f"non-finite grad for {name}" assert grad.abs().sum() > 0, f"zero grad for {name}" + _assert_nemotron_h_gradient_checkpointing_full_layer() + + +def _assert_nemotron_h_registry_and_local_config_policy(tmp_path): + registry = get_registry() + assert "NemotronHForCausalLM" in registry.supported_models + assert registry.get_model_cls_from_model_arch("NemotronHForCausalLM") is NemotronHForCausalLM + + config_dir = tmp_path / "nemotron-3-ultra" + config_dir.mkdir() + payload = { + "model_type": "nemotron_h", + "architectures": ["NemotronHForCausalLM"], + "vocab_size": 131072, + "hidden_size": 64, + "layers_block_type": ["mamba", "moe", "attention", "moe"], + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "mamba_num_heads": 8, + "mamba_head_dim": 16, + "n_groups": 4, + "ssm_state_size": 32, + "conv_kernel": 4, + "chunk_size": 64, + "mlp_hidden_act": "relu2", + "mamba_hidden_act": "silu", + "n_routed_experts": 16, + "num_experts_per_tok": 4, + "moe_intermediate_size": 48, + "moe_shared_expert_intermediate_size": 96, + "moe_latent_size": 32, + "routed_scaling_factor": 5.0, + "n_group": 1, + "topk_group": 1, + "norm_topk_prob": True, + "time_step_floor": 1e-4, + "time_step_min": 1e-3, + "time_step_max": 0.1, + "time_step_limit": [1e-4, 0.1], + "num_nextn_predict_layers": 1, + "mtp_layers_block_type": ["attention", "moe"], + "rescale_prenorm_residual": True, + "tie_word_embeddings": False, + } + (config_dir / "config.json").write_text(json.dumps(payload), encoding="utf-8") -def test_nemotron_h_loss_with_labels(): - model = _build_model() - model.train() - input_ids = torch.randint(0, model.config.vocab_size, (2, SEQ_LEN)) - labels = torch.randint(0, model.config.vocab_size, (2, SEQ_LEN)) - labels[:, :3] = -100 + config = _load_local_xorl_config(str(config_dir), {}) - outputs = model(input_ids=input_ids) - result = compute_loss( - model.lm_head, - outputs.last_hidden_state, - loss_fn_name=None, - loss_fn_inputs={"labels": labels}, - loss_fn_params={"ce_mode": "eager"}, - logits_to_keep=0, - ) - assert result.loss.ndim == 0 - assert torch.isfinite(result.loss) - result.loss.backward() - assert model.model.layers[0].mixer.in_proj.weight.grad is not None + assert isinstance(config, NemotronHConfig) + assert config.model_type == "nemotron_h" + assert config.architectures == ["NemotronHForCausalLM"] + assert config.layers_block_type == ["mamba", "moe", "attention", "moe"] + assert config.num_hidden_layers == 4 + assert config.moe_latent_size == 32 + assert config.routed_scaling_factor == 5.0 + assert config.time_step_limit == (1e-4, 0.1) + assert config.n_routed_experts == 16 + assert config.tie_word_embeddings is False -def test_nemotron_h_gradient_checkpointing_full_layer(): +def _assert_nemotron_h_gradient_checkpointing_full_layer(): model = _build_model() model.train() model.gradient_checkpointing_enable() @@ -154,18 +207,3 @@ def test_nemotron_h_packed_varlen_matches_per_sequence(): start += length separate = torch.cat(pieces, dim=1) torch.testing.assert_close(packed.last_hidden_state, separate, atol=1e-5, rtol=1e-4) - - -def test_nemotron_h_packed_varlen_smoke_all_block_types(): - """Packed kwargs flow through mamba + attention + moe blocks (forward + backward).""" - model = _build_model() - model.train() - input_ids = torch.randint(0, model.config.vocab_size, (1, SEQ_LEN)) - cu_seqlens = torch.tensor([0, 7, SEQ_LEN], dtype=torch.int32) - - outputs = model(input_ids=input_ids, cu_seq_lens_q=cu_seqlens, cu_seq_lens_k=cu_seqlens) - assert outputs.last_hidden_state.shape == (1, SEQ_LEN, model.config.hidden_size) - assert torch.isfinite(outputs.last_hidden_state).all() - outputs.last_hidden_state.pow(2).mean().backward() - assert model.model.layers[0].mixer.in_proj.weight.grad is not None - assert torch.isfinite(model.model.layers[0].mixer.in_proj.weight.grad).all() diff --git a/tests/models/test_nemotron_h_registry.py b/tests/models/test_nemotron_h_registry.py deleted file mode 100644 index 012b88b4..00000000 --- a/tests/models/test_nemotron_h_registry.py +++ /dev/null @@ -1,74 +0,0 @@ -import json - -import pytest - -from xorl.models.auto import _load_local_xorl_config -from xorl.models.registry import get_registry -from xorl.models.transformers.nemotron_h.configuration_nemotron_h import NemotronHConfig -from xorl.models.transformers.nemotron_h.modeling_nemotron_h import NemotronHForCausalLM - - -pytestmark = [pytest.mark.cpu] - - -def _ultra_style_config_dict() -> dict: - return { - "model_type": "nemotron_h", - "architectures": ["NemotronHForCausalLM"], - "vocab_size": 131072, - "hidden_size": 64, - "layers_block_type": ["mamba", "moe", "attention", "moe"], - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 16, - "mamba_num_heads": 8, - "mamba_head_dim": 16, - "n_groups": 4, - "ssm_state_size": 32, - "conv_kernel": 4, - "chunk_size": 64, - "mlp_hidden_act": "relu2", - "mamba_hidden_act": "silu", - "n_routed_experts": 16, - "num_experts_per_tok": 4, - "moe_intermediate_size": 48, - "moe_shared_expert_intermediate_size": 96, - "moe_latent_size": 32, - "routed_scaling_factor": 5.0, - "n_group": 1, - "topk_group": 1, - "norm_topk_prob": True, - "time_step_floor": 1e-4, - "time_step_min": 1e-3, - "time_step_max": 0.1, - "time_step_limit": [1e-4, 0.1], - "num_nextn_predict_layers": 1, - "mtp_layers_block_type": ["attention", "moe"], - "rescale_prenorm_residual": True, - "tie_word_embeddings": False, - } - - -def test_nemotron_h_registered(): - registry = get_registry() - assert "NemotronHForCausalLM" in registry.supported_models - assert registry.get_model_cls_from_model_arch("NemotronHForCausalLM") is NemotronHForCausalLM - - -def test_local_auto_config_builds_nemotron_h_config(tmp_path): - config_dir = tmp_path / "nemotron-3-ultra" - config_dir.mkdir() - (config_dir / "config.json").write_text(json.dumps(_ultra_style_config_dict())) - - config = _load_local_xorl_config(str(config_dir), {}) - - assert isinstance(config, NemotronHConfig) - assert config.model_type == "nemotron_h" - assert config.architectures == ["NemotronHForCausalLM"] - assert config.layers_block_type == ["mamba", "moe", "attention", "moe"] - assert config.num_hidden_layers == 4 - assert config.moe_latent_size == 32 - assert config.routed_scaling_factor == 5.0 - assert config.time_step_limit == (1e-4, 0.1) - assert config.n_routed_experts == 16 - assert config.tie_word_embeddings is False diff --git a/tests/models/test_olmo2_support.py b/tests/models/test_olmo2_support.py index 9791dd6f..f5e69b73 100644 --- a/tests/models/test_olmo2_support.py +++ b/tests/models/test_olmo2_support.py @@ -1,14 +1,11 @@ import pytest import torch -from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel from transformers.models.olmo2.configuration_olmo2 import Olmo2Config as HFOlmo2Config from transformers.models.olmo2.modeling_olmo2 import Olmo2ForCausalLM as HFOlmo2ForCausalLM from xorl.models.auto import build_foundation_model from xorl.models.transformers.olmo2.configuration_olmo2 import Olmo2Config as XOlmo2Config from xorl.models.transformers.olmo2.modeling_olmo2 import Olmo2ForCausalLM -from xorl.models.transformers.olmo2.parallelize import MODEL_TP_PLAN, TP_PLAN -from xorl.models.transformers.olmo2.tp_styles import LocalAxisRMSNormShard pytestmark = [pytest.mark.cpu] @@ -65,36 +62,11 @@ def test_build_foundation_model_accepts_hf_olmo2_config_object(): assert layer.self_attn.qkv_proj.bias is None assert layer.self_attn.o_proj.bias is None + _assert_olmo2_unfuse_for_tp_matches_hf_parameter_layout() + _assert_olmo2_checkpoint_handler_bidirectional_policy() -def test_olmo2_tp_plan_uses_local_axis_qk_norm(): - # OLMo-2's full-axis q_norm/k_norm doesn't compose with SequenceParallel - # or stock ColwiseParallel — under colwise q/k_proj the input arrives - # hidden-sharded and a full-hidden weight can't be applied directly. - # LocalAxisRMSNormShard shards the 1-D weight on dim 0 so each rank's - # slice matches its local q/k slice. This is the actual root cause of - # what was originally reported (the issue thought it was post-norm; the trace - # was at q_norm in _project_qkv). - assert isinstance(TP_PLAN["layers.*.self_attn.q_norm"], LocalAxisRMSNormShard) - assert isinstance(TP_PLAN["layers.*.self_attn.k_norm"], LocalAxisRMSNormShard) - # Post-norms see a Replicate input after the rowwise all-reduce in - # o_proj/down_proj — they should NOT be in the plan (no TP wrapping). - assert "layers.*.post_attention_layernorm" not in TP_PLAN - assert "layers.*.post_feedforward_layernorm" not in TP_PLAN - assert "norm" not in TP_PLAN - - # Standard colwise/rowwise everywhere else. - assert isinstance(TP_PLAN["layers.*.self_attn.q_proj"], ColwiseParallel) - assert isinstance(TP_PLAN["layers.*.self_attn.o_proj"], RowwiseParallel) - assert isinstance(TP_PLAN["layers.*.mlp.gate_proj"], ColwiseParallel) - assert isinstance(TP_PLAN["layers.*.mlp.down_proj"], RowwiseParallel) - - # lm_head: vanilla colwise (Replicate input from the model's final norm, - # vocab-parallel output for vocab_parallel_cross_entropy). - assert MODEL_TP_PLAN["lm_head"] == "colwise" or isinstance(MODEL_TP_PLAN["lm_head"], ColwiseParallel) - - -def test_olmo2_unfuse_for_tp_matches_hf_parameter_layout(): +def _assert_olmo2_unfuse_for_tp_matches_hf_parameter_layout(): model = Olmo2ForCausalLM(_make_xorl_olmo2_config()) model.unfuse_for_tp() @@ -111,10 +83,17 @@ def test_olmo2_unfuse_for_tp_matches_hf_parameter_layout(): assert not hasattr(layer.mlp, "gate_up_proj") assert hasattr(layer.mlp, "gate_proj") assert hasattr(layer.mlp, "up_proj") - assert model.get_checkpoint_handler() is None + # Unfused: the handler is returned with both merges disabled; its pre-quantized + # loading paths stay active. + handler = model.get_checkpoint_handler() + assert handler is not None + # HF's already-split keys pass straight through to matching parameters. + key = "model.layers.0.self_attn.q_proj.weight" + passthrough = handler.on_load_weight(key, layer.self_attn.q_proj.weight.detach()) + assert [name for name, _ in passthrough] == [key] -def test_olmo2_checkpoint_handler_exports_hf_compatible_attention_keys(): +def _assert_olmo2_checkpoint_handler_bidirectional_policy(): model = Olmo2ForCausalLM(_make_xorl_olmo2_config()) handler = model.get_checkpoint_handler() @@ -137,8 +116,10 @@ def test_olmo2_checkpoint_handler_exports_hf_compatible_attention_keys(): assert "model.layers.0.self_attn.qkv_proj.weight" not in transformed assert "model.layers.0.mlp.gate_up_proj.weight" not in transformed + _assert_olmo2_checkpoint_handler_loads_hf_weights_into_fused_model() + -def test_olmo2_checkpoint_handler_loads_hf_weights_into_fused_model(): +def _assert_olmo2_checkpoint_handler_loads_hf_weights_into_fused_model(): hf_config = _make_hf_olmo2_config() hf_config._attn_implementation = "eager" xorl_config = _make_xorl_olmo2_config() diff --git a/tests/models/test_op_parity_dense.py b/tests/models/test_op_parity_dense.py deleted file mode 100644 index 24bb743f..00000000 --- a/tests/models/test_op_parity_dense.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Op-parity regression guard for the dense Qwen3-1.7B lockstep ops. - -These are the non-aten operators that must stay bit-for-bit aligned with SGLang -for the static K3 (train/serve logprob) parity to hold: the eager RoPE apply and -the eager SwiGLU. (RMSNorm is guarded by ``test_rmsnorm_sglang_fused.py``; -attention is the shared FA kernel / irreducible paged-vs-contiguous floor and is -not asserted here.) The SGLang reference is inlined verbatim from its -``forward_native`` so the test needs no SGLang install. - -The K3 recipe uses the *eager* RoPE (flash rope unavailable -> naive path, -``rope_native``) and the *eager* SwiGLU (``activation_native``); the fused Triton -SwiGLU is intentionally NOT asserted bit-exact here — it is rejected for K3 at -network scale. -""" - -import pytest -import torch -import torch.nn.functional as F - -from xorl.models.layers import rope as xrope -from xorl.ops.fused_silu_and_mul import _native_silu_and_mul - - -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - -DEV = "cuda" -DT = torch.bfloat16 -HEAD_DIM = 128 -N_Q_HEADS = 16 -N_KV_HEADS = 8 -ROPE_THETA = 1_000_000 -SEQ = 96 -MAX_POS = 40960 - - -def _sg_cos_sin_cache(): - inv_freq = 1.0 / (ROPE_THETA ** (torch.arange(0, HEAD_DIM, 2, dtype=torch.float, device=DEV) / HEAD_DIM)) - t = torch.arange(MAX_POS, dtype=torch.float, device=DEV) - freqs = torch.einsum("i,j -> ij", t, inv_freq) - return torch.cat((freqs.cos(), freqs.sin()), dim=-1) # fp32 cache (SGLang keeps fp32 on CUDA) - - -def _sg_forward_native(positions, query, key, cache): - """Inlined SGLang RotaryEmbedding.forward_native (is_neox_style=True).""" - - def apply(x, cos, sin): - cos = cos.unsqueeze(-2).to(x.dtype) - sin = sin.unsqueeze(-2).to(x.dtype) - x1, x2 = torch.chunk(x, 2, dim=-1) - o1 = x1 * cos - x2 * sin - o2 = x2 * cos + x1 * sin - return torch.cat((o1, o2), dim=-1) - - cos_sin = cache.index_select(0, positions) - cos, sin = cos_sin.chunk(2, dim=-1) - qs = query.shape - q_rot = apply(query.view(SEQ, -1, HEAD_DIM), cos, sin).reshape(qs) - ks = key.shape - k_rot = apply(key.view(SEQ, -1, HEAD_DIM), cos, sin).reshape(ks) - return q_rot, k_rot - - -@requires_cuda -@pytest.mark.gpu -def test_rope_xorl_eager_bit_exact_vs_sglang_native(): - torch.manual_seed(0) - - q = torch.randn(1, SEQ, N_Q_HEADS, HEAD_DIM, device=DEV, dtype=DT) - k = torch.randn(1, SEQ, N_KV_HEADS, HEAD_DIM, device=DEV, dtype=DT) - positions = torch.arange(SEQ, device=DEV) - - class Cfg: - rope_scaling = None - head_dim = HEAD_DIM - hidden_size = N_Q_HEADS * HEAD_DIM - num_attention_heads = N_Q_HEADS - max_position_embeddings = MAX_POS - rope_theta = ROPE_THETA - rope_parameters = {} - - xrot = xrope.RotaryEmbedding(Cfg(), device=DEV) - # K3 recipe runs the naive (eager) rope path (flash apply unavailable). - assert xrope._flash_apply_rotary_emb is None - cos, sin = xrot.forward(q.view(1, SEQ, -1), positions.unsqueeze(0)) - xq, xk = xrope.apply_rotary_pos_emb(q, k, cos, sin) - - cache = _sg_cos_sin_cache() - sq = q.view(SEQ, N_Q_HEADS, HEAD_DIM).reshape(SEQ, -1).to(DT) - sk = k.view(SEQ, N_KV_HEADS, HEAD_DIM).reshape(SEQ, -1).to(DT) - sq_out, sk_out = _sg_forward_native(positions, sq.clone(), sk.clone(), cache) - sq_out = sq_out.view(SEQ, N_Q_HEADS, HEAD_DIM).unsqueeze(0) - sk_out = sk_out.view(SEQ, N_KV_HEADS, HEAD_DIM).unsqueeze(0) - - assert torch.equal(xq, sq_out), "xorl eager RoPE(Q) diverged from SGLang forward_native" - assert torch.equal(xk, sk_out), "xorl eager RoPE(K) diverged from SGLang forward_native" - - -@requires_cuda -@pytest.mark.gpu -def test_swiglu_xorl_native_bit_exact_vs_sglang_native(): - torch.manual_seed(1) - - inter = 6144 - gate_up = torch.randn(SEQ, 2 * inter, device=DEV, dtype=DT) - - x_xorl_native = _native_silu_and_mul(gate_up) # F.silu(gate) * up in bf16 - d = gate_up.shape[-1] // 2 - x_sg = F.silu(gate_up[..., :d]) * gate_up[..., d:] # SGLang forward_native - - assert torch.equal(x_xorl_native, x_sg), "xorl eager SwiGLU diverged from SGLang forward_native" diff --git a/tests/models/test_qwen2_support.py b/tests/models/test_qwen2_support.py index f8906d2e..b04b9ff3 100644 --- a/tests/models/test_qwen2_support.py +++ b/tests/models/test_qwen2_support.py @@ -60,8 +60,11 @@ def test_build_foundation_model_accepts_hf_qwen2_config_object(): assert model.model.layers[0].self_attn.qkv_proj.bias is not None assert model.model.layers[0].self_attn.o_proj.bias is None + _assert_qwen2_unfuse_for_tp_matches_hf_parameter_layout() + _assert_qwen2_checkpoint_handler_bidirectional_policy() -def test_qwen2_unfuse_for_tp_matches_hf_parameter_layout(): + +def _assert_qwen2_unfuse_for_tp_matches_hf_parameter_layout(): model = Qwen2ForCausalLM(_make_xorl_qwen2_config()) model.unfuse_for_tp() @@ -78,10 +81,17 @@ def test_qwen2_unfuse_for_tp_matches_hf_parameter_layout(): assert not hasattr(layer.mlp, "gate_up_proj") assert hasattr(layer.mlp, "gate_proj") assert hasattr(layer.mlp, "up_proj") - assert model.get_checkpoint_handler() is None + # Unfused: the handler is returned with both merges disabled; its pre-quantized + # loading paths stay active. + handler = model.get_checkpoint_handler() + assert handler is not None + # HF's already-split keys pass straight through to matching parameters. + key = "model.layers.0.self_attn.q_proj.weight" + passthrough = handler.on_load_weight(key, layer.self_attn.q_proj.weight.detach()) + assert [name for name, _ in passthrough] == [key] -def test_qwen2_checkpoint_handler_exports_hf_compatible_attention_keys(): +def _assert_qwen2_checkpoint_handler_bidirectional_policy(): model = Qwen2ForCausalLM(_make_xorl_qwen2_config()) handler = model.get_checkpoint_handler() @@ -104,8 +114,10 @@ def test_qwen2_checkpoint_handler_exports_hf_compatible_attention_keys(): assert "model.layers.0.self_attn.qkv_proj.weight" not in transformed assert "model.layers.0.mlp.gate_up_proj.weight" not in transformed + _assert_qwen2_checkpoint_handler_loads_hf_weights_into_fused_model() + -def test_qwen2_checkpoint_handler_loads_hf_weights_into_fused_model(): +def _assert_qwen2_checkpoint_handler_loads_hf_weights_into_fused_model(): hf_config = _make_hf_qwen2_config() hf_config._attn_implementation = "eager" xorl_config = _make_xorl_qwen2_config() diff --git a/tests/models/test_qwen35_lora_projection_topology.py b/tests/models/test_qwen35_lora_projection_topology.py new file mode 100644 index 00000000..cc088fcc --- /dev/null +++ b/tests/models/test_qwen35_lora_projection_topology.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch +from torch import nn +from torch.nn import functional as F + +from xorl.lora.fold import canonical_lora_fold_linear +from xorl.lora.modules.base import LoraModule +from xorl.lora.modules.delta_linear import LoraDeltaLinear +from xorl.lora.modules.linear import LoraLinear +from xorl.lora.target_manifest import collect_lora_runtime_modules +from xorl.lora.utils import ( + _get_default_target_modules, + inject_lora_into_model, + load_lora_checkpoint, + save_lora_checkpoint, +) +from xorl.models.transformers.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeMLP +from xorl.ops.batch_invariant_ops import set_trunk_linear_contract, wrap_trunk_linears_batch_invariant +from xorl.ops.linear_attention.layers.gated_deltanet import GatedDeltaNet + + +RANK = 16 + + +def _mlp_config() -> SimpleNamespace: + return SimpleNamespace( + hidden_size=8, + intermediate_size=6, + hidden_act="silu", + _activation_native=True, + ) + + +class _Layer(nn.Module): + def __init__(self, *, with_gdn: bool = True) -> None: + super().__init__() + if with_gdn: + self.linear_attn = GatedDeltaNet( + hidden_size=8, + expand_v=1, + head_dim=2, + num_heads=2, + num_v_heads=2, + use_short_conv=False, + layer_idx=0, + ) + self.mlp = nn.Module() + self.mlp.shared_expert = Qwen3_5MoeMLP(_mlp_config()) + + +class _Model(nn.Module): + def __init__(self, layers: int = 1, *, with_gdn: bool = True) -> None: + super().__init__() + self.config = SimpleNamespace(model_type="xorl_qwen3_5_moe") + self.model = nn.Module() + self.model.layers = nn.ModuleList([_Layer(with_gdn=with_gdn) for _ in range(layers)]) + + +def _projection_targets() -> list[str]: + return ["q_proj", "k_proj", "v_proj", "g_proj", "gate_proj", "up_proj", "down_proj"] + + +def test_qwen35_gdn_and_shared_expert_projection_topology_and_gradients(tmp_path) -> None: + torch.manual_seed(7) + model = _Model() + assert _get_default_target_modules(model) == [ + "q_proj", + "k_proj", + "v_proj", + "g_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] + inject_lora_into_model(model, r=RANK, lora_alpha=RANK, target_modules=_projection_targets()) + + gdn = model.model.layers[0].linear_attn + assert not hasattr(gdn, "in_proj_qkvz") + gdn_inputs = (gdn.q_proj, gdn.k_proj, gdn.v_proj, gdn.g_proj) + assert all(isinstance(module, LoraLinear) and module.r == RANK for module in gdn_inputs) + assert len({module.lora_A.data_ptr() for module in gdn_inputs}) == 4 + + shared = model.model.layers[0].mlp.shared_expert + assert isinstance(shared.gate_proj, LoraDeltaLinear) + assert isinstance(shared.up_proj, LoraDeltaLinear) + assert isinstance(shared.down_proj, LoraLinear) + assert all(module.r == RANK for module in (shared.gate_proj, shared.up_proj, shared.down_proj)) + + with torch.no_grad(): + for module in (*gdn_inputs, shared.gate_proj, shared.up_proj, shared.down_proj): + module.lora_B.normal_() + module.exact_merged_forward = True + + gate_base, up_base = shared.gate_up_proj.weight.chunk(2, dim=0) + gate_weight = canonical_lora_fold_linear(gate_base, shared.gate_proj.lora_A, shared.gate_proj.lora_B, 1.0) + up_weight = canonical_lora_fold_linear(up_base, shared.up_proj.lora_A, shared.up_proj.lora_B, 1.0) + down_weight = canonical_lora_fold_linear( + shared.down_proj.weight, + shared.down_proj.lora_A, + shared.down_proj.lora_B, + 1.0, + ) + inputs = torch.randn(2, 3, 8) + gate, up = F.linear(inputs, torch.cat((gate_weight, up_weight), dim=0)).chunk(2, dim=-1) + expected = F.linear(F.silu(gate) * up, down_weight) + actual = shared(inputs) + assert torch.equal(actual, expected) + + actual.sum().backward() + for module in (shared.gate_proj, shared.up_proj, shared.down_proj): + assert module.lora_A.grad is not None and torch.count_nonzero(module.lora_A.grad) > 0 + assert module.lora_B.grad is not None and torch.count_nonzero(module.lora_B.grad) > 0 + + save_lora_checkpoint( + model, + str(tmp_path), + base_model_name="Qwen/Qwen3.6-35B-A3B", + r=RANK, + lora_alpha=RANK, + preserve_lora_dtype=True, + ) + restored = _Model() + inject_lora_into_model(restored, r=RANK, lora_alpha=RANK, target_modules=_projection_targets()) + load_lora_checkpoint(restored, str(tmp_path), strict=True) + expected_state = {name: value for name, value in model.state_dict().items() if "lora_" in name} + actual_state = {name: value for name, value in restored.state_dict().items() if "lora_" in name} + assert actual_state.keys() == expected_state.keys() + for name, expected_value in expected_state.items(): + assert torch.equal(actual_state[name], expected_value) + + +def test_qwen36_forty_layer_shared_expert_inventory() -> None: + model = _Model(layers=40, with_gdn=False) + inject_lora_into_model( + model, + r=RANK, + lora_alpha=RANK, + target_modules=["gate_proj", "up_proj", "down_proj"], + ) + inventory = collect_lora_runtime_modules(model) + assert len(inventory) == 40 * 3 + for layer_idx in range(40): + prefix = f"model.layers.{layer_idx}.mlp.shared_expert" + assert inventory[f"{prefix}.gate_proj"] == RANK + assert inventory[f"{prefix}.up_proj"] == RANK + assert inventory[f"{prefix}.down_proj"] == RANK + + +def test_qwen_shared_expert_separate_factors_compose_with_exact_trunk_wrap() -> None: + model = _Model() + inject_lora_into_model(model, r=RANK, lora_alpha=RANK, target_modules=_projection_targets()) + for module in model.modules(): + if isinstance(module, LoraModule): + module.exact_merged_forward = True + try: + wrapped = wrap_trunk_linears_batch_invariant(model) + shared = model.model.layers[0].mlp.shared_expert + assert wrapped["gate_up_proj"] == 1 + assert getattr(shared.gate_up_proj, "_xorl_bi_trunk_wrapped", False) + assert getattr(shared.down_proj, "_xorl_bi_trunk_wrapped", False) + assert not getattr(shared.gate_proj, "_xorl_bi_trunk_wrapped", False) + assert not getattr(shared.up_proj, "_xorl_bi_trunk_wrapped", False) + finally: + set_trunk_linear_contract(False) + + +def test_qwen_shared_expert_supports_independent_fused_gate_up_adapters() -> None: + for target in ("gate_proj", "up_proj"): + model = _Model(with_gdn=False) + inject_lora_into_model(model, r=RANK, lora_alpha=RANK, target_modules=[target]) + shared = model.model.layers[0].mlp.shared_expert + assert isinstance(getattr(shared, target), LoraDeltaLinear) + other = "up_proj" if target == "gate_proj" else "gate_proj" + assert not hasattr(shared, other) diff --git a/tests/models/test_qwen35_rmsnorm_v2_candidate.py b/tests/models/test_qwen35_rmsnorm_v2_candidate.py deleted file mode 100644 index a8a03fe7..00000000 --- a/tests/models/test_qwen35_rmsnorm_v2_candidate.py +++ /dev/null @@ -1,81 +0,0 @@ -import torch - -from xorl.models.layers import normalization - - -EPS = 1e-6 - - -def _cpu_v2_forward(x, weight, eps, *, residual=None, zero_centered=False): - norm_input = x if residual is None else x + residual - fp32 = norm_input.float() - inv_rms = torch.rsqrt(fp32.square().mean(dim=-1, keepdim=True) + eps) - scale = weight.float() + 1.0 if zero_centered else weight.float() - out = (fp32 * inv_rms * scale).to(x.dtype) - return out if residual is None else (out, norm_input) - - -def _cpu_rms_backward(normed_input, weight, eps, grad_output, grad_residual_out=None): - with torch.enable_grad(): - x = normed_input.detach().float().requires_grad_(True) - w = weight.detach().float().requires_grad_(True) - inv_rms = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps) - out = x * inv_rms * w - objective = (out * grad_output.float()).sum() - if grad_residual_out is not None: - objective = objective + (x * grad_residual_out.float()).sum() - return torch.autograd.grad(objective, (x, w)) - - -def test_qwen_v2_zero_centered_backward_uses_effective_weight(monkeypatch): - monkeypatch.setattr(normalization, "rms_norm_v2", _cpu_v2_forward) - monkeypatch.setattr(normalization, "fused_rms_norm_backward", _cpu_rms_backward) - torch.manual_seed(11) - x = torch.randn(3, 8, dtype=torch.bfloat16, requires_grad=True) - weight = torch.randn(8, dtype=torch.float32, requires_grad=True) - grad = torch.randn_like(x) - - out = normalization._FamiliesV2ZeroCenteredRMSNorm.apply(x, weight, EPS) - out.backward(grad) - candidate_dx = x.grad.detach().clone() - candidate_dw = weight.grad.detach().clone() - - x_ref = x.detach().requires_grad_(True) - weight_ref = weight.detach().requires_grad_(True) - ref = _cpu_v2_forward(x_ref, weight_ref, EPS, zero_centered=True) - ref.backward(grad) - - assert torch.allclose(candidate_dx.float(), x_ref.grad.float(), atol=2e-2, rtol=2e-2) - assert torch.allclose(candidate_dw, weight_ref.grad, atol=2e-5, rtol=2e-5) - - -def test_qwen_v2_residual_backward_preserves_both_gradient_paths(monkeypatch): - monkeypatch.setattr(normalization, "rms_norm_v2", _cpu_v2_forward) - monkeypatch.setattr(normalization, "fused_rms_norm_backward", _cpu_rms_backward) - torch.manual_seed(13) - x = torch.randn(2, 8, dtype=torch.bfloat16, requires_grad=True) - residual = torch.randn(2, 8, dtype=torch.bfloat16, requires_grad=True) - weight = torch.randn(8, dtype=torch.float32, requires_grad=True) - grad_out = torch.randn_like(x) - grad_residual = torch.randn_like(residual) - - out, residual_out = normalization._FamiliesV2ZeroCenteredResidualRMSNorm.apply(x, residual, weight, EPS) - torch.autograd.backward((out, residual_out), (grad_out, grad_residual)) - candidate = (x.grad.detach().clone(), residual.grad.detach().clone(), weight.grad.detach().clone()) - - x_ref = x.detach().requires_grad_(True) - residual_ref = residual.detach().requires_grad_(True) - weight_ref = weight.detach().requires_grad_(True) - ref_out, ref_residual = _cpu_v2_forward( - x_ref, - weight_ref, - EPS, - residual=residual_ref, - zero_centered=True, - ) - torch.autograd.backward((ref_out, ref_residual), (grad_out, grad_residual)) - - assert torch.allclose(candidate[0].float(), x_ref.grad.float(), atol=2e-2, rtol=2e-2) - assert torch.equal(candidate[0], candidate[1]) - assert torch.allclose(candidate[1].float(), residual_ref.grad.float(), atol=2e-2, rtol=2e-2) - assert torch.allclose(candidate[2], weight_ref.grad, atol=2e-5, rtol=2e-5) diff --git a/tests/models/test_qwen3_5_apply_rotary.py b/tests/models/test_qwen3_5_apply_rotary.py index aa05a7bb..8d000421 100644 --- a/tests/models/test_qwen3_5_apply_rotary.py +++ b/tests/models/test_qwen3_5_apply_rotary.py @@ -1,6 +1,3 @@ -import ast -import inspect - import pytest import torch @@ -65,22 +62,7 @@ def to_interleaved(x: torch.Tensor) -> torch.Tensor: return to_interleaved(q_embed_h), to_interleaved(k_embed_h) -def test_default_matches_hf_half_rotate(): - """Default (interleaved=False) is the Qwen3.5/3.6 + HF/SGLang convention.""" - torch.manual_seed(0) - batch, seq, num_heads, head_dim = 2, 4, 3, 8 - q = torch.randn(batch, seq, num_heads, head_dim, dtype=torch.float32) - k = torch.randn(batch, seq, num_heads, head_dim, dtype=torch.float32) - cos, sin = _build_halved_cos_sin(batch, seq, head_dim) - - q_ours, k_ours = qwen3_5_apply_rotary_pos_emb(q, k, cos, sin) - q_ref, k_ref = _hf_reference_half_rotate(q, k, cos, sin) - - torch.testing.assert_close(q_ours, q_ref, atol=1e-6, rtol=1e-6) - torch.testing.assert_close(k_ours, k_ref, atol=1e-6, rtol=1e-6) - - -def test_interleaved_matches_pairwise_reference(): +def _assert_interleaved_matches_pairwise_reference(): """interleaved=True is the DSv3 MLA decoupled-RoPE pairwise convention.""" torch.manual_seed(0) batch, seq, num_heads, head_dim = 2, 5, 3, 8 @@ -95,21 +77,7 @@ def test_interleaved_matches_pairwise_reference(): torch.testing.assert_close(k_ours, k_ref, atol=1e-6, rtol=1e-6) -def test_interleaved_pairwise_rotation_d8(): - """Hand-worked d=8 sanity: pair i must be rotated by angle θ_i, not θ_{i+1}.""" - torch.manual_seed(0) - batch, seq, num_heads, head_dim = 1, 1, 1, 8 - q = torch.randn(batch, seq, num_heads, head_dim, dtype=torch.float32) - k = torch.zeros_like(q) - cos, sin = _build_halved_cos_sin(batch, seq, head_dim) - - q_out, _ = qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, interleaved=True) - # cos_unique[t=0] = [1,1,1,1] and sin_unique[t=0] = [0,0,0,0], so rotation - # at position 0 is the identity. - torch.testing.assert_close(q_out, q, atol=1e-6, rtol=1e-6) - - -def test_class_b_dispatches_before_eager_qwen_math(monkeypatch): +def _assert_class_b_rotary_admission_policy(monkeypatch): sentinel = (object(), object()) calls = [] @@ -131,64 +99,70 @@ def _stock(q, k, cos, sin, *, interleaved): assert qwen3_5_shared.qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, class_b=True) is sentinel assert calls == [(q, k, cos, sin, False)] + _assert_class_b_fails_loudly_outside_cuda_contract() -@pytest.mark.parametrize( - ("q", "k", "cos", "sin", "message"), - [ - ( - torch.zeros((1, 2, 1, 4), dtype=torch.bfloat16), - torch.zeros((1, 2, 1, 4), dtype=torch.bfloat16), - torch.zeros((1, 2, 4), dtype=torch.float32), - torch.zeros((1, 2, 4), dtype=torch.float32), - "requires CUDA", - ), - ( - torch.zeros((1, 2, 1, 4), dtype=torch.float32), - torch.zeros((1, 2, 1, 4), dtype=torch.float32), - torch.zeros((1, 2, 4), dtype=torch.float32), - torch.zeros((1, 2, 4), dtype=torch.float32), - "requires CUDA", - ), - ], -) -def test_class_b_fails_loudly_outside_cuda_contract(q, k, cos, sin, message): - with pytest.raises(RuntimeError, match=message): - qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, class_b=True) - - -@pytest.mark.parametrize("modeling_module", [modeling_qwen3_5, modeling_qwen3_5_moe]) -def test_qwen35_modeling_does_not_pass_interleaved_to_rotary(modeling_module): - """Regression: Qwen3.5/3.6 attention must NOT pass `interleaved=` into - `qwen3_5_apply_rotary_pos_emb`. q/k features use the standard half-rotate - convention (HF/SGLang); `mrope_interleaved` controls T/H/W frequency mixing - in cos/sin construction upstream, not the q/k rotation convention. Plumbing - it in would silently switch q/k to pairwise rotation for any HF Qwen3.5/3.6 - config that ships with `mrope_interleaved=true`. - """ - tree = ast.parse(inspect.getsource(modeling_module)) - offenders: list[int] = [] - for node in ast.walk(tree): - if ( - isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "qwen3_5_apply_rotary_pos_emb" - and any(kw.arg == "interleaved" for kw in node.keywords) - ): - offenders.append(node.lineno) - assert not offenders, ( - f"{modeling_module.__name__} calls `qwen3_5_apply_rotary_pos_emb(..., interleaved=...)` " - f"at line(s) {offenders}; this must not be plumbed from `mrope_interleaved`." + +def _assert_class_b_fails_loudly_outside_cuda_contract(): + cos = torch.zeros((1, 2, 4), dtype=torch.float32) + sin = torch.zeros_like(cos) + for dtype in (torch.bfloat16, torch.float32): + q = torch.zeros((1, 2, 1, 4), dtype=dtype) + k = torch.zeros_like(q) + with pytest.raises(RuntimeError, match="requires CUDA"): + qwen3_5_apply_rotary_pos_emb(q, k, cos, sin, class_b=True) + + +def _assert_qwen35_attention_keeps_half_rotate_when_mrope_is_interleaved(attention_type, config_type): + """mRoPE interleaves frequencies, not the q/k feature rotation convention.""" + torch.manual_seed(23) + config = config_type( + hidden_size=8, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=4, + num_hidden_layers=1, + layer_types=["full_attention"], + mrope_interleaved=True, + ) + attention = attention_type(config, layer_idx=0) + hidden = torch.randn(1, 3, 8) + cos, sin = _build_halved_cos_sin(batch=1, seq=3, head_dim=4) + + query, key, _ = attention._project_qkv(hidden, (cos, sin)) + + input_shape = hidden.shape[:-1] + hidden_shape = (*input_shape, -1, attention.head_dim) + query_pre_rope, _ = torch.chunk( + attention.q_proj(hidden).view(*input_shape, -1, attention.head_dim * 2), + 2, + dim=-1, ) + query_pre_rope = attention.q_norm(query_pre_rope.view(hidden_shape)) + key_pre_rope = attention.k_norm(attention.k_proj(hidden).view(hidden_shape)) + expected_query, expected_key = _hf_reference_half_rotate(query_pre_rope, key_pre_rope, cos, sin) + pairwise_query, pairwise_key = _hf_reference_pairwise(query_pre_rope, key_pre_rope, cos, sin) + + assert config.mrope_interleaved is True + torch.testing.assert_close(query, expected_query, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(key, expected_key, atol=1e-6, rtol=1e-6) + assert not torch.allclose(query, pairwise_query) + assert not torch.allclose(key, pairwise_key) -@pytest.mark.parametrize( - ("attention_type", "config_type"), - [ +def test_qwen35_rotary_numerics_admission_and_attention_policy(monkeypatch): + _assert_interleaved_matches_pairwise_reference() + _assert_class_b_rotary_admission_policy(monkeypatch) + + for attention_type, config_type in ( (modeling_qwen3_5.Qwen3_5Attention, Qwen3_5Config), (modeling_qwen3_5_moe.Qwen3_5MoeAttention, Qwen3_5MoeConfig), - ], -) -def test_qwen35_exact_attention_casts_post_rope_qk_to_bf16(attention_type, config_type): + ): + _assert_qwen35_attention_keeps_half_rotate_when_mrope_is_interleaved(attention_type, config_type) + + _assert_qwen35_exact_attention_casts_post_rope_qk_to_bf16_policy() + + +def _assert_qwen35_exact_attention_casts_post_rope_qk_to_bf16(attention_type, config_type): config = config_type( hidden_size=8, num_attention_heads=2, @@ -207,3 +181,11 @@ def test_qwen35_exact_attention_casts_post_rope_qk_to_bf16(attention_type, confi assert query.dtype is torch.bfloat16 assert key.dtype is torch.bfloat16 assert value.dtype is torch.float32 + + +def _assert_qwen35_exact_attention_casts_post_rope_qk_to_bf16_policy(): + for attention_type, config_type in ( + (modeling_qwen3_5.Qwen3_5Attention, Qwen3_5Config), + (modeling_qwen3_5_moe.Qwen3_5MoeAttention, Qwen3_5MoeConfig), + ): + _assert_qwen35_exact_attention_casts_post_rope_qk_to_bf16(attention_type, config_type) diff --git a/tests/models/test_qwen3_5_moe_rmsnorm.py b/tests/models/test_qwen3_5_moe_rmsnorm.py deleted file mode 100644 index 4a5a5bb7..00000000 --- a/tests/models/test_qwen3_5_moe_rmsnorm.py +++ /dev/null @@ -1,473 +0,0 @@ -import pytest -import torch - -from xorl.models.layers.normalization import set_rmsnorm_mode -from xorl.models.transformers.qwen3_5_moe import modeling_qwen3_5_moe -from xorl.models.transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig -from xorl.models.transformers.qwen3_5_moe.modeling_qwen3_5_moe import ( - Qwen3_5MoeDecoderLayer, - Qwen3_5MoeModel, - Qwen3_5MoeRMSNorm, -) -from xorl.ops.batch_invariant_ops import ( - rms_norm_batch_invariant, - set_batch_invariant_mode, -) - - -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - -HIDDEN = 2048 -N_TOKENS = 512 -EPS = 1e-6 - - -@pytest.mark.cpu -def test_qwen35_exact_norm_selection_is_structural(monkeypatch): - calls = [] - monkeypatch.setattr( - modeling_qwen3_5_moe, - "fast_zero_centered_batch_invariant_residual_rms_norm", - lambda hidden, _weight, _eps: calls.append("exact") or hidden + 1, - ) - monkeypatch.setattr( - modeling_qwen3_5_moe, - "native_zero_centered_rms_norm_without_batch_invariant", - lambda hidden, _weight, _eps: calls.append("legacy") or hidden + 2, - ) - - set_rmsnorm_mode("sglang") - try: - norm = Qwen3_5MoeRMSNorm(4, exact_contract=True) - x = torch.ones(2, 4) - assert torch.equal(norm(x, force_sglang_residual=True), x + 1) - assert calls == ["exact"] - finally: - set_rmsnorm_mode("native") - - -def test_qwen3_5_moe_sglang_rmsnorm_keeps_no_residual_native(monkeypatch): - calls = [] - - def fake_native(hidden_states, weight, variance_epsilon): - calls.append(("native", hidden_states.clone())) - return hidden_states + 1 - - def fake_eager(hidden_states, weight, variance_epsilon): - calls.append(("eager", hidden_states.clone())) - return hidden_states + 2 - - def fake_native_no_batch_invariant(hidden_states, weight, variance_epsilon): - calls.append(("native_no_batch_invariant", hidden_states.clone())) - return hidden_states + 3 - - monkeypatch.setattr(modeling_qwen3_5_moe, "native_zero_centered_rms_norm", fake_native) - monkeypatch.setattr(modeling_qwen3_5_moe, "eager_zero_centered_rms_norm", fake_eager) - monkeypatch.setattr( - modeling_qwen3_5_moe, - "native_zero_centered_rms_norm_without_batch_invariant", - fake_native_no_batch_invariant, - ) - - set_rmsnorm_mode("sglang") - try: - norm = modeling_qwen3_5_moe.Qwen3_5MoeRMSNorm(4, exact_contract=False) - - x = torch.ones(2, 4) - out = norm(x) - assert torch.equal(out, x + 1) - assert calls[-1][0] == "native" - - residual = torch.full((2, 4), 3.0) - out, residual_out = norm(x, residual=residual, prenorm=True) - assert torch.equal(residual_out, x + residual) - assert torch.equal(out, x + residual + 3) - assert calls[-1][0] == "native_no_batch_invariant" - - out = norm(x, force_sglang_residual=True) - assert torch.equal(out, x + 3) - assert calls[-1][0] == "native_no_batch_invariant" - finally: - set_rmsnorm_mode("native") - - -def test_qwen3_5_moe_sglang_fused_rmsnorm_routes_same_families_as_sglang(monkeypatch): - """Norm-seed contract (§14) port gate: ``sglang_fused`` must dispatch the same - family functions as ``sglang`` (family-2 = no-batch-invariant residual tree for - residual/forced norms, family-1 = native for plain no-residual norms).""" - calls = [] - - def fake_native(hidden_states, weight, variance_epsilon): - calls.append("native") - return hidden_states + 1 - - def fake_native_no_batch_invariant(hidden_states, weight, variance_epsilon): - calls.append("native_no_batch_invariant") - return hidden_states + 3 - - def fake_contract(hidden_states, weight, variance_epsilon): - calls.append("contract_family1") - return hidden_states + 5 - - monkeypatch.setattr(modeling_qwen3_5_moe, "native_zero_centered_rms_norm", fake_native) - monkeypatch.setattr( - modeling_qwen3_5_moe, - "native_zero_centered_rms_norm_without_batch_invariant", - fake_native_no_batch_invariant, - ) - monkeypatch.setattr( - modeling_qwen3_5_moe, - "fast_zero_centered_batch_invariant_rms_norm", - fake_contract, - ) - monkeypatch.setattr( - modeling_qwen3_5_moe, - "fast_zero_centered_batch_invariant_residual_rms_norm", - fake_native_no_batch_invariant, - ) - - set_rmsnorm_mode("sglang_fused") - try: - norm = modeling_qwen3_5_moe.Qwen3_5MoeRMSNorm(4, exact_contract=False) - exact_norm = modeling_qwen3_5_moe.Qwen3_5MoeRMSNorm(4, exact_contract=True) - x = torch.ones(2, 4) - - out = norm(x) - assert torch.equal(out, x + 1) - assert calls[-1] == "native" - - residual = torch.full((2, 4), 3.0) - out, residual_out = norm(x, residual=residual, prenorm=True) - assert torch.equal(residual_out, x + residual) - assert torch.equal(out, x + residual + 3) - assert calls[-1] == "native_no_batch_invariant" - - out = norm(x, force_sglang_residual=True) - assert torch.equal(out, x + 3) - assert calls[-1] == "native_no_batch_invariant" - - # Exact selection is module-owned: no-residual dispatch swaps to the - # batch-invariant family while residual/forced remains family-2. - out = exact_norm(x) - assert torch.equal(out, x + 5) - assert calls[-1] == "contract_family1" - - out = exact_norm(x, force_sglang_residual=True) - assert torch.equal(out, x + 3) - assert calls[-1] == "native_no_batch_invariant" - - # Exact and ordinary models may coexist without process-state leakage. - out = norm(x) - assert torch.equal(out, x + 1) - assert calls[-1] == "native" - finally: - set_rmsnorm_mode("native") - - -@pytest.mark.cpu -def test_qwen3_5_moe_v2_candidate_dispatches_without_changing_default(monkeypatch): - calls = [] - - def fake_v2(hidden, _weight, _eps, *, residual=None): - calls.append(residual is not None) - if residual is None: - return hidden + 7 - residual_out = hidden + residual - return residual_out + 7, residual_out - - monkeypatch.setattr(modeling_qwen3_5_moe, "fast_zero_centered_families_v2_rms_norm", fake_v2) - set_rmsnorm_mode("sglang_fused") - try: - default = Qwen3_5MoeRMSNorm(4, exact_contract=True) - candidate = Qwen3_5MoeRMSNorm(4, exact_contract=True, rmsnorm_family="v2") - x = torch.ones(2, 4) - residual = torch.full_like(x, 3) - assert default.rmsnorm_family == "v1" - assert torch.equal(candidate(x), x + 7) - out, residual_out = candidate(x, residual=residual, prenorm=True) - assert torch.equal(out, x + residual + 7) - assert torch.equal(residual_out, x + residual) - assert calls == [False, True] - finally: - set_rmsnorm_mode("native") - - with pytest.raises(RuntimeError, match="only in the exact training lane"): - Qwen3_5MoeRMSNorm(4, exact_contract=False, rmsnorm_family="v2") - - -# --------------------------------------------------------------------------- # -# Call sites: layer>0 input norm and the final norm must force family-2 in both -# sglang and sglang_fused modes (§14 family assignment, ported from qwen3_moe). -# --------------------------------------------------------------------------- # -class CaptureInputNorm(torch.nn.Module): - def __init__(self, mode: str): - super().__init__() - self.mode = mode - self.force_sglang_residual_values = [] - - def forward(self, hidden_states, *, force_sglang_residual=False): - self.force_sglang_residual_values.append(force_sglang_residual) - return hidden_states - - -class IdentityAttention(torch.nn.Module): - def forward(self, hidden_states, **kwargs): - return hidden_states, None - - -class IdentityPostAttentionNorm(torch.nn.Module): - def forward(self, hidden_states, residual=None, prenorm=False, **kwargs): - return hidden_states, residual - - -def _tiny_config(**overrides) -> Qwen3_5MoeConfig: - kwargs = dict( - vocab_size=32, - hidden_size=8, - intermediate_size=16, - moe_intermediate_size=4, - num_hidden_layers=2, - num_attention_heads=2, - num_key_value_heads=1, - num_experts=0, - num_experts_per_tok=1, - max_position_embeddings=16, - layer_types=["full_attention", "full_attention"], - _attn_implementation="eager", - pad_token_id=0, - ) - kwargs.update(overrides) - return Qwen3_5MoeConfig(**kwargs) - - -@pytest.mark.cpu -def test_qwen3_5_moe_v2_resolves_every_zero_centered_norm_site(): - config = _tiny_config() - config._qwen35_exact_contract = True - config._qwen35_rmsnorm_family = "v2" - set_rmsnorm_mode("sglang_fused") - try: - model = Qwen3_5MoeModel(config) - finally: - set_rmsnorm_mode("native") - - resolved = { - name: module.rmsnorm_family for name, module in model.named_modules() if isinstance(module, Qwen3_5MoeRMSNorm) - } - assert resolved - assert set(resolved.values()) == {"v2"} - assert "norm" in resolved - for layer_idx in range(config.num_hidden_layers): - prefix = f"layers.{layer_idx}" - assert resolved[f"{prefix}.input_layernorm"] == "v2" - assert resolved[f"{prefix}.post_attention_layernorm"] == "v2" - assert resolved[f"{prefix}.self_attn.q_norm"] == "v2" - assert resolved[f"{prefix}.self_attn.k_norm"] == "v2" - - -@pytest.mark.parametrize( - ("layer_idx", "mode", "expected_force"), - [ - (0, "sglang", False), - (1, "native", False), - (1, "sglang", True), - (0, "sglang_fused", False), - (1, "sglang_fused", True), - ], -) -def test_qwen3_5_moe_layer_input_norm_forces_sglang_residual_after_layer0(layer_idx, mode, expected_force): - layer = Qwen3_5MoeDecoderLayer(_tiny_config(), layer_idx=layer_idx) - assert layer.layer_idx == layer_idx - - input_norm = CaptureInputNorm(mode) - layer.input_layernorm = input_norm - layer.self_attn = IdentityAttention() - layer.post_attention_layernorm = IdentityPostAttentionNorm() - - hidden_states = torch.ones(1, 2, 8) - layer._pre_mlp_forward(hidden_states, position_embeddings=(hidden_states, hidden_states)) - - assert input_norm.force_sglang_residual_values == [expected_force] - - -@pytest.mark.parametrize( - ("mode", "expected_force"), - [ - ("native", False), - ("sglang", True), - ("sglang_fused", True), - ], -) -def test_qwen3_5_moe_model_final_norm_forces_sglang_residual(mode, expected_force): - class StubLayer(torch.nn.Module): - layer_type = "full_attention" - - def forward(self, hidden_states, *args, **kwargs): - return (hidden_states,) - - model = Qwen3_5MoeModel(_tiny_config()) - model.layers = torch.nn.ModuleList([StubLayer()]) - final_norm = CaptureInputNorm(mode) - model.norm = final_norm - - model(input_ids=torch.tensor([[0, 1]])) - - assert final_norm.force_sglang_residual_values == [expected_force] - - -# --------------------------------------------------------------------------- # -# GPU: module-level bitwise verify — sglang_fused == sglang on every dispatch -# shape, and the trunk-contract family-1 kernel == the aten interpose lane. -# --------------------------------------------------------------------------- # -@requires_cuda -@pytest.mark.gpu -def test_qwen3_5_module_sglang_fused_equals_sglang_bitwise(): - torch.manual_seed(3) - device = "cuda" - hidden = torch.randn(N_TOKENS, HIDDEN, device=device, dtype=torch.bfloat16) - residual = torch.randn(N_TOKENS, HIDDEN, device=device, dtype=torch.bfloat16) - - set_rmsnorm_mode("sglang") - try: - sg = Qwen3_5MoeRMSNorm(HIDDEN, eps=EPS).to(device) - set_rmsnorm_mode("sglang_fused") - sf = Qwen3_5MoeRMSNorm(HIDDEN, eps=EPS).to(device) - finally: - set_rmsnorm_mode("native") - with torch.no_grad(): - sg.weight.copy_(torch.randn(HIDDEN, device=device)) - sf.weight.copy_(sg.weight) - assert sg.mode == "sglang" and sf.mode == "sglang_fused" - - with set_batch_invariant_mode(True), torch.no_grad(): - # Residual (post-attention layernorm) path — family-2. - out_sg, rout_sg = sg(hidden, residual=residual, prenorm=True) - out_sf, rout_sf = sf(hidden, residual=residual, prenorm=True) - assert torch.equal(out_sg, out_sf) - assert torch.equal(rout_sg, rout_sf) - - # force_sglang_residual (layer>0 input norm / final norm) — family-2. - out_sg2 = sg(hidden, force_sglang_residual=True) - out_sf2 = sf(hidden, force_sglang_residual=True) - assert torch.equal(out_sg2, out_sf2) - - # No-residual, no-force (qk-norm / layer-0 input) — family-1 (interpose). - out_sg3 = sg(hidden) - out_sf3 = sf(hidden) - assert torch.equal(out_sg3, out_sf3) - - -@requires_cuda -@pytest.mark.gpu -def test_qwen3_5_trunk_contract_family1_bit_matches_interpose_kernel(): - """Family-1 term of the Q3.5 norm-seed contract: under the trunk-contract - lane, the no-residual dispatch (qk-norm shape) must equal the aten::rms_norm - interpose lane bit-for-bit — the same guarantee qwen3_moe has post-§14.""" - torch.manual_seed(21) - head_dim = 128 - x = torch.randn(256, 16, head_dim, device="cuda", dtype=torch.bfloat16) - set_rmsnorm_mode("sglang_fused") - try: - norm = Qwen3_5MoeRMSNorm(head_dim, eps=EPS, exact_contract=True).to("cuda") - finally: - set_rmsnorm_mode("native") - with torch.no_grad(): - norm.weight.copy_(torch.randn(head_dim, device="cuda")) - - with torch.no_grad(): - out = norm(x) - ref = rms_norm_batch_invariant(x.float(), 1.0 + norm.weight.float(), eps=EPS).to(x.dtype) - with set_batch_invariant_mode(True): - ref_interpose = torch.nn.functional.rms_norm(x.float(), (head_dim,), 1.0 + norm.weight.float(), eps=EPS).to( - x.dtype - ) - - assert torch.equal(out, ref) - assert torch.equal(out, ref_interpose), "contract-lane family-1 must equal the aten interpose lane bit-for-bit" - - -@requires_cuda -@pytest.mark.gpu -def test_qwen3_5_layer_forward_bit_exact_sglang_vs_fused(): - """Full Q3.5 full-attention decoder-layer forward must be bit-identical - between sglang and sglang_fused (model-level §14 gate for the hybrid). - Exercises the layer>0 input norm (forced family-2) and the post-attention - residual norm (family-2) with the real layer modules; family-1 (qk-norm) - is covered by the module-level tests above.""" - torch.manual_seed(7) - device = "cuda" - cfg = _tiny_config( - hidden_size=HIDDEN, - intermediate_size=1024, - num_attention_heads=16, - num_key_value_heads=8, - head_dim=128, - ) - set_rmsnorm_mode("sglang") - try: - layer = Qwen3_5MoeDecoderLayer(cfg, layer_idx=1).to(device=device, dtype=torch.bfloat16) - finally: - set_rmsnorm_mode("native") - layer.self_attn = IdentityAttention() - with torch.no_grad(): - layer.input_layernorm.weight.copy_(torch.randn(HIDDEN, device=device).to(torch.bfloat16)) - layer.post_attention_layernorm.weight.copy_(torch.randn(HIDDEN, device=device).to(torch.bfloat16)) - - hidden = torch.randn(1, 128, HIDDEN, device=device, dtype=torch.bfloat16) - pos = torch.zeros(1, 128, HIDDEN, device=device, dtype=torch.bfloat16) - - with set_batch_invariant_mode(True), torch.no_grad(): - for m in (layer.input_layernorm, layer.post_attention_layernorm): - m.mode = "sglang" - (out_sg,) = layer(hidden, position_embeddings=(pos, pos)) - for m in (layer.input_layernorm, layer.post_attention_layernorm): - m.mode = "sglang_fused" - (out_sf,) = layer(hidden, position_embeddings=(pos, pos)) - - assert torch.equal(out_sg, out_sf), "Q3.5 layer forward diverged between sglang and sglang_fused" - - -@requires_cuda -@pytest.mark.gpu -def test_qwen3_5_family2_residual_norm_contract(monkeypatch): - """The exact residual-tree norm matches the sampler's BI-mean composition.""" - from xorl.models.layers.normalization import ( # noqa: PLC0415 - fast_zero_centered_batch_invariant_residual_rms_norm, - native_zero_centered_rms_norm, - ) - from xorl.ops.batch_invariant_ops import mean_dim # noqa: PLC0415 - - torch.manual_seed(11) - device = "cuda" - x = torch.randn(513, HIDDEN, device=device, dtype=torch.bfloat16) - residual = torch.randn(513, HIDDEN, device=device, dtype=torch.bfloat16) - weight = (torch.randn(HIDDEN, device=device) * 0.02).to(torch.bfloat16) - - set_rmsnorm_mode("sglang_fused") - try: - exact_norm = modeling_qwen3_5_moe.Qwen3_5MoeRMSNorm(HIDDEN, eps=1e-6, exact_contract=True).to(device) - ordinary_norm = modeling_qwen3_5_moe.Qwen3_5MoeRMSNorm(HIDDEN, eps=1e-6, exact_contract=False).to(device) - finally: - set_rmsnorm_mode("native") - with torch.no_grad(): - exact_norm.weight.copy_(weight) - ordinary_norm.weight.copy_(weight) - - # reference: eager fp32 composition with the BI mean kernel (the sampler's - # forward_native under SGLANG_BATCH_INVARIANT_OPS=all) - y = (x + residual).float() - var = mean_dim(y * y, dim=-1, keepdim=True) - ref = (y * torch.rsqrt(var + 1e-6) * (1.0 + weight.float())).to(torch.bfloat16) - - with torch.no_grad(): - out = exact_norm(x, residual=residual) - assert torch.equal(out, ref), "family-2 contract output != eager-with-BI-mean composition" - - with torch.no_grad(): - out_native = ordinary_norm(x, residual=residual) - expected_native = native_zero_centered_rms_norm(x + residual, weight, 1e-6) - assert torch.equal(out_native, expected_native), "ordinary execution must preserve the native family-2 path" - - # the standalone helper equals the contract dispatch - helper = fast_zero_centered_batch_invariant_residual_rms_norm(x + residual, weight, 1e-6) - assert torch.equal(helper, ref) diff --git a/tests/models/test_qwen3_5_registry.py b/tests/models/test_qwen3_5_registry.py index d2c3d9e7..e06c3590 100644 --- a/tests/models/test_qwen3_5_registry.py +++ b/tests/models/test_qwen3_5_registry.py @@ -1,104 +1,113 @@ -from types import SimpleNamespace +import json +from xorl.models.auto import _load_local_xorl_config from xorl.models.registry import get_registry from xorl.models.transformers.qwen3_5.configuration_qwen3_5 import Qwen3_5Config from xorl.models.transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig -def test_qwen3_5_conditional_generation_registered(): - registry = get_registry() - assert "Qwen3_5ForConditionalGeneration" in registry.supported_models - assert "Qwen3_5MoeForConditionalGeneration" in registry.supported_models - - -def test_qwen3_5_moe_config_from_hf_config(): - rope_parameters = { - "rope_type": "default", - "rope_theta": 10_000_000, - "partial_rotary_factor": 0.25, - "mrope_interleaved": True, - } +def _moe_config_dict(): num_hidden_layers = 40 full_attention_interval = 4 layer_types = [ "full_attention" if (i + 1) % full_attention_interval == 0 else "linear_attention" for i in range(num_hidden_layers) ] - text_config = SimpleNamespace( - vocab_size=248320, - hidden_size=2048, - intermediate_size=2048, - shared_expert_intermediate_size=512, - num_hidden_layers=num_hidden_layers, - num_attention_heads=16, - num_key_value_heads=2, - head_dim=256, - hidden_act="silu", - max_position_embeddings=262144, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - attention_bias=False, - attention_dropout=0.0, - layer_types=layer_types, - full_attention_interval=full_attention_interval, - linear_num_key_heads=16, - linear_num_value_heads=32, - linear_key_head_dim=128, - linear_value_head_dim=128, - attn_output_gate=True, - linear_conv_kernel_dim=4, - moe_intermediate_size=512, - num_experts_per_tok=8, - num_experts=256, - router_aux_loss_coef=0.001, - mlp_only_layers=[], - rope_parameters=rope_parameters, - ) - hf_config = SimpleNamespace( - text_config=text_config, - tie_word_embeddings=False, - ) + return { + "model_type": "qwen3_5_moe", + "text_config": { + "vocab_size": 248320, + "hidden_size": 2048, + "intermediate_size": 2048, + "shared_expert_intermediate_size": 512, + "num_hidden_layers": num_hidden_layers, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "hidden_act": "silu", + "max_position_embeddings": 262144, + "initializer_range": 0.02, + "rms_norm_eps": 1e-6, + "use_cache": True, + "attention_bias": False, + "attention_dropout": 0.0, + "layer_types": layer_types, + "full_attention_interval": full_attention_interval, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "attn_output_gate": True, + "linear_conv_kernel_dim": 4, + "moe_intermediate_size": 512, + "num_experts_per_tok": 8, + "num_experts": 256, + "router_aux_loss_coef": 0.001, + "mlp_only_layers": [], + "rope_parameters": { + "rope_type": "default", + "rope_theta": 10_000_000, + "partial_rotary_factor": 0.25, + "mrope_interleaved": True, + }, + }, + "tie_word_embeddings": False, + } + - config = Qwen3_5MoeConfig.from_hf_config(hf_config) +def _dense_config_dict(): + return { + "model_type": "qwen3_5", + "text_config": { + "vocab_size": 248320, + "hidden_size": 4096, + "intermediate_size": 12288, + "num_hidden_layers": 32, + "num_attention_heads": 16, + "num_key_value_heads": 4, + "head_dim": 256, + "hidden_act": "silu", + "max_position_embeddings": 262144, + "initializer_range": 0.02, + "rms_norm_eps": 1e-6, + "use_cache": True, + "attention_bias": False, + "attention_dropout": 0.0, + "layer_types": ["linear_attention", "full_attention"], + "full_attention_interval": 4, + "rope_parameters": {"rope_type": "default", "rope_theta": 10_000_000}, + }, + "tie_word_embeddings": False, + } + + +def test_local_auto_config_builds_qwen3_5_family_configs(tmp_path): + """Exercise dense and MoE conversion through the real local-config loader.""" + registry = get_registry() + assert "Qwen3_5ForConditionalGeneration" in registry.supported_models + assert "Qwen3_5MoeForConditionalGeneration" in registry.supported_models + + for name, raw_config, expected_type in ( + ("dense", _dense_config_dict(), Qwen3_5Config), + ("moe", _moe_config_dict(), Qwen3_5MoeConfig), + ): + config_dir = tmp_path / name + config_dir.mkdir() + (config_dir / "config.json").write_text(json.dumps(raw_config)) + config = _load_local_xorl_config(str(config_dir), {}) + + assert isinstance(config, expected_type) + assert config.layer_types == [ + "full_attention" + if (layer_idx + 1) % raw_config["text_config"]["full_attention_interval"] == 0 + else "linear_attention" + for layer_idx in range(raw_config["text_config"]["num_hidden_layers"]) + ] + assert config.head_dim == 256 - assert config.layer_types == layer_types assert config.linear_num_key_heads == 16 assert config.linear_num_value_heads == 32 assert config.linear_key_head_dim == 128 assert config.linear_value_head_dim == 128 assert config.mrope_interleaved is True assert config.num_experts == 256 - - -def test_qwen3_5_config_from_hf_config(): - text_config = SimpleNamespace( - vocab_size=248320, - hidden_size=4096, - intermediate_size=12288, - num_hidden_layers=32, - num_attention_heads=16, - num_key_value_heads=4, - head_dim=256, - hidden_act="silu", - max_position_embeddings=262144, - initializer_range=0.02, - rms_norm_eps=1e-6, - use_cache=True, - attention_bias=False, - attention_dropout=0.0, - layer_types=["linear_attention", "full_attention"], - full_attention_interval=4, - rope_parameters={"rope_type": "default", "rope_theta": 10_000_000}, - ) - hf_config = SimpleNamespace(text_config=text_config, tie_word_embeddings=False) - - config = Qwen3_5Config.from_hf_config(hf_config) - - assert config.vocab_size == 248320 - assert config.hidden_size == 4096 - assert config.head_dim == 256 - assert config.layer_types == [ - "full_attention" if (layer_idx + 1) % text_config.full_attention_interval == 0 else "linear_attention" - for layer_idx in range(text_config.num_hidden_layers) - ] diff --git a/tests/models/test_qwen3_5_rmsnorm.py b/tests/models/test_qwen3_5_rmsnorm.py index 444458a5..3e7d7735 100644 --- a/tests/models/test_qwen3_5_rmsnorm.py +++ b/tests/models/test_qwen3_5_rmsnorm.py @@ -1,92 +1,151 @@ +"""Dense and MoE Qwen3.5 RMSNorm contract owners. + +The two model families carry parallel zero-centered RMSNorm implementations. +Their shared dispatch and site-assignment policy is exercised once here, with +both implementations passed through the same cases. GPU arithmetic is covered +by one representative model lifecycle because both classes call the same +normalization kernels. +""" + import pytest import torch +from xorl.models.layers import normalization from xorl.models.layers.normalization import set_rmsnorm_mode from xorl.models.transformers.qwen3_5 import modeling_qwen3_5 from xorl.models.transformers.qwen3_5.configuration_qwen3_5 import Qwen3_5Config -from xorl.models.transformers.qwen3_5.modeling_qwen3_5 import Qwen3_5DecoderLayer, Qwen3_5TextModel +from xorl.models.transformers.qwen3_5.modeling_qwen3_5 import ( + Qwen3_5DecoderLayer, + Qwen3_5RMSNorm, + Qwen3_5TextModel, +) +from xorl.models.transformers.qwen3_5_moe import modeling_qwen3_5_moe +from xorl.models.transformers.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig +from xorl.models.transformers.qwen3_5_moe.modeling_qwen3_5_moe import ( + Qwen3_5MoeDecoderLayer, + Qwen3_5MoeModel, + Qwen3_5MoeRMSNorm, +) +from xorl.ops.batch_invariant_ops import rms_norm_batch_invariant, set_batch_invariant_mode -def test_qwen3_5_exact_norm_selection_is_structural(monkeypatch): - calls = [] - monkeypatch.setattr( - modeling_qwen3_5, - "fast_zero_centered_batch_invariant_residual_rms_norm", - lambda hidden, _weight, _eps: calls.append("exact") or hidden + 1, +HIDDEN = 2048 +N_TOKENS = 512 +EPS = 1e-6 + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _tiny_dense_config(**overrides): + kwargs = dict( + vocab_size=32, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + max_position_embeddings=16, + layer_types=["full_attention", "full_attention"], + _attn_implementation="eager", + pad_token_id=0, ) - monkeypatch.setattr( - modeling_qwen3_5, - "native_zero_centered_rms_norm_without_batch_invariant", - lambda hidden, _weight, _eps: calls.append("legacy") or hidden + 2, + kwargs.update(overrides) + return Qwen3_5Config(**kwargs) + + +def _tiny_moe_config(**overrides): + kwargs = dict( + vocab_size=32, + hidden_size=8, + intermediate_size=16, + moe_intermediate_size=4, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + num_experts=0, + num_experts_per_tok=1, + max_position_embeddings=16, + layer_types=["full_attention", "full_attention"], + _attn_implementation="eager", + pad_token_id=0, ) + kwargs.update(overrides) + return Qwen3_5MoeConfig(**kwargs) - set_rmsnorm_mode("sglang") - try: - norm = modeling_qwen3_5.Qwen3_5RMSNorm(4, exact_contract=True) - x = torch.ones(2, 4) - assert torch.equal(norm(x, force_sglang_residual=True), x + 1) - assert calls == ["exact"] - finally: - set_rmsnorm_mode("native") + +VARIANTS = ( + ( + "dense", + modeling_qwen3_5, + Qwen3_5RMSNorm, + _tiny_dense_config, + Qwen3_5TextModel, + Qwen3_5DecoderLayer, + ), + ( + "moe", + modeling_qwen3_5_moe, + Qwen3_5MoeRMSNorm, + _tiny_moe_config, + Qwen3_5MoeModel, + Qwen3_5MoeDecoderLayer, + ), +) -def test_qwen3_5_sglang_fused_rmsnorm_routes_serving_families(monkeypatch): +def _assert_dispatch_policy(module, norm_cls, monkeypatch): calls = [] - def fake_native(hidden_states, weight, variance_epsilon): + def fake_native(hidden, _weight, _eps): calls.append("native") - return hidden_states + 1 + return hidden + 1 - def fake_residual(hidden_states, weight, variance_epsilon): + def fake_residual(hidden, _weight, _eps): calls.append("residual") - return hidden_states + 3 + return hidden + 3 - def fake_family1(hidden_states, weight, variance_epsilon): + def fake_family1(hidden, _weight, _eps): calls.append("family1") - return hidden_states + 5 - - monkeypatch.setattr(modeling_qwen3_5, "native_zero_centered_rms_norm", fake_native) - monkeypatch.setattr( - modeling_qwen3_5, - "native_zero_centered_rms_norm_without_batch_invariant", - fake_residual, - ) - monkeypatch.setattr( - modeling_qwen3_5, - "fast_zero_centered_batch_invariant_rms_norm", - fake_family1, - ) - monkeypatch.setattr( - modeling_qwen3_5, - "fast_zero_centered_batch_invariant_residual_rms_norm", - fake_residual, - ) - - set_rmsnorm_mode("sglang_fused") - try: - norm = modeling_qwen3_5.Qwen3_5RMSNorm(4, exact_contract=False) - exact_norm = modeling_qwen3_5.Qwen3_5RMSNorm(4, exact_contract=True) - x = torch.ones(2, 4) - - assert torch.equal(norm(x), x + 1) - assert calls[-1] == "native" - assert torch.equal(norm(x, force_sglang_residual=True), x + 3) - assert calls[-1] == "residual" - - assert torch.equal(exact_norm(x), x + 5) - assert calls[-1] == "family1" - assert torch.equal(exact_norm(x, force_sglang_residual=True), x + 3) - assert calls[-1] == "residual" - - # Exact and ordinary models may coexist in one process. Running the - # exact module must not alter the ordinary module's dispatch. - assert torch.equal(norm(x), x + 1) - assert calls[-1] == "native" - finally: - set_rmsnorm_mode("native") + return hidden + 5 + monkeypatch.setattr(module, "native_zero_centered_rms_norm", fake_native) + monkeypatch.setattr(module, "native_zero_centered_rms_norm_without_batch_invariant", fake_residual) + monkeypatch.setattr(module, "fast_zero_centered_batch_invariant_rms_norm", fake_family1) + monkeypatch.setattr(module, "fast_zero_centered_batch_invariant_residual_rms_norm", fake_residual) -def test_qwen3_5_v2_candidate_is_explicit_and_fail_loud(monkeypatch): + x = torch.ones(2, 4) + residual = torch.full_like(x, 3) + for mode in ("sglang", "sglang_fused"): + set_rmsnorm_mode(mode) + try: + ordinary = norm_cls(4, exact_contract=False) + exact = norm_cls(4, exact_contract=True) + + assert torch.equal(ordinary(x), x + 1) + assert calls[-1] == "native" + out, residual_out = ordinary(x, residual=residual, prenorm=True) + assert torch.equal(out, x + residual + 3) + assert torch.equal(residual_out, x + residual) + assert calls[-1] == "residual" + assert torch.equal(ordinary(x, force_sglang_residual=True), x + 3) + assert calls[-1] == "residual" + + expected = "family1" if mode == "sglang_fused" else "native" + expected_delta = 5 if mode == "sglang_fused" else 1 + assert torch.equal(exact(x), x + expected_delta) + assert calls[-1] == expected + assert torch.equal(exact(x, force_sglang_residual=True), x + 3) + assert calls[-1] == "residual" + + # Exact selection is instance-owned and cannot leak into an + # ordinary module living in the same process. + assert torch.equal(ordinary(x), x + 1) + assert calls[-1] == "native" + finally: + set_rmsnorm_mode("native") + + +def _assert_v2_admission_and_dispatch(module, norm_cls, monkeypatch): calls = [] def fake_v2(hidden, _weight, _eps, *, residual=None): @@ -96,98 +155,118 @@ def fake_v2(hidden, _weight, _eps, *, residual=None): residual_out = hidden + residual return residual_out + 7, residual_out - monkeypatch.setattr(modeling_qwen3_5, "fast_zero_centered_families_v2_rms_norm", fake_v2) + monkeypatch.setattr(module, "fast_zero_centered_families_v2_rms_norm", fake_v2) x = torch.ones(2, 4) residual = torch.full_like(x, 3) set_rmsnorm_mode("sglang_fused") try: - v1 = modeling_qwen3_5.Qwen3_5RMSNorm(4, exact_contract=True) - v2 = modeling_qwen3_5.Qwen3_5RMSNorm(4, exact_contract=True, rmsnorm_family="v2") - assert v1.rmsnorm_family == "v1" - assert torch.equal(v2(x), x + 7) - out, residual_out = v2(x, residual=residual, prenorm=True) - assert torch.equal(residual_out, x + residual) + default = norm_cls(4, exact_contract=True) + candidate = norm_cls(4, exact_contract=True, rmsnorm_family="v2") + assert default.rmsnorm_family == "v1" + assert torch.equal(candidate(x), x + 7) + out, residual_out = candidate(x, residual=residual, prenorm=True) assert torch.equal(out, x + residual + 7) + assert torch.equal(residual_out, x + residual) assert calls == [False, True] finally: set_rmsnorm_mode("native") with pytest.raises(RuntimeError, match="only in the exact training lane"): - modeling_qwen3_5.Qwen3_5RMSNorm(4, exact_contract=False, rmsnorm_family="v2") + norm_cls(4, exact_contract=False, rmsnorm_family="v2") - set_rmsnorm_mode("native") - try: - rejected = modeling_qwen3_5.Qwen3_5RMSNorm(4, exact_contract=True, rmsnorm_family="v2") - with pytest.raises(RuntimeError, match="requires rmsnorm_mode='sglang_fused'"): - rejected(x) - finally: - set_rmsnorm_mode("native") + rejected = norm_cls(4, exact_contract=True, rmsnorm_family="v2") + with pytest.raises(RuntimeError, match="requires rmsnorm_mode='sglang_fused'"): + rejected(x) class CaptureNorm(torch.nn.Module): - def __init__(self, mode: str): + def __init__(self, mode): super().__init__() self.mode = mode self.force_values = [] - def forward(self, hidden_states, *, force_sglang_residual=False, **kwargs): + def forward(self, hidden_states, residual=None, prenorm=False, *, force_sglang_residual=False, **_kwargs): self.force_values.append(force_sglang_residual) - if kwargs.get("prenorm"): - return hidden_states, kwargs.get("residual") + if prenorm: + return hidden_states, residual return hidden_states class IdentityAttention(torch.nn.Module): - def forward(self, hidden_states, **kwargs): + def forward(self, hidden_states, **_kwargs): return hidden_states, None -def _tiny_config(**overrides) -> Qwen3_5Config: - kwargs = dict( - vocab_size=32, - hidden_size=8, - intermediate_size=16, - num_hidden_layers=2, - num_attention_heads=2, - num_key_value_heads=1, - max_position_embeddings=16, - layer_types=["full_attention", "full_attention"], - _attn_implementation="eager", - pad_token_id=0, - ) - kwargs.update(overrides) - return Qwen3_5Config(**kwargs) - - -def test_qwen3_5_v2_resolves_every_zero_centered_norm_site(): - config = _tiny_config() +def _assert_v2_reaches_every_norm_site(variant): + name, _module, norm_cls, config_factory, model_cls, _layer_cls = variant + config = config_factory() config._qwen35_exact_contract = True config._qwen35_rmsnorm_family = "v2" set_rmsnorm_mode("sglang_fused") try: - model = Qwen3_5TextModel(config) + model = model_cls(config) finally: set_rmsnorm_mode("native") resolved = { - name: module.rmsnorm_family - for name, module in model.named_modules() - if isinstance(module, modeling_qwen3_5.Qwen3_5RMSNorm) + module_name: module.rmsnorm_family + for module_name, module in model.named_modules() + if isinstance(module, norm_cls) } - assert resolved - assert set(resolved.values()) == {"v2"} + assert resolved, name + assert set(resolved.values()) == {"v2"}, name assert "norm" in resolved for layer_idx in range(config.num_hidden_layers): prefix = f"layers.{layer_idx}" - assert resolved[f"{prefix}.input_layernorm"] == "v2" - assert resolved[f"{prefix}.post_attention_layernorm"] == "v2" - assert resolved[f"{prefix}.self_attn.q_norm"] == "v2" - assert resolved[f"{prefix}.self_attn.k_norm"] == "v2" + for suffix in ("input_layernorm", "post_attention_layernorm", "self_attn.q_norm", "self_attn.k_norm"): + assert resolved[f"{prefix}.{suffix}"] == "v2", (name, suffix) + + +def _run_to_post_attention_norm(name, layer, hidden): + layer.input_layernorm = CaptureNorm(layer.input_layernorm.mode) + captured = layer.input_layernorm + layer.self_attn = IdentityAttention() + layer.post_attention_layernorm = CaptureNorm("native") + if name == "dense": + layer.mlp = torch.nn.Identity() + layer(hidden, position_embeddings=(hidden, hidden)) + else: + layer._pre_mlp_forward(hidden, position_embeddings=(hidden, hidden)) + return captured + + +def _assert_layer_and_final_norm_site_policy(variant): + name, _module, _norm_cls, config_factory, model_cls, layer_cls = variant + hidden = torch.ones(1, 2, 8) + for layer_idx, mode, expected in ( + (0, "sglang", False), + (1, "native", False), + (1, "sglang", True), + (0, "sglang_fused", False), + (1, "sglang_fused", True), + ): + layer = layer_cls(config_factory(), layer_idx=layer_idx) + layer.input_layernorm.mode = mode + captured = _run_to_post_attention_norm(name, layer, hidden) + assert captured.force_values == [expected], (name, layer_idx, mode) + + class StubLayer(torch.nn.Module): + layer_type = "full_attention" + def forward(self, hidden_states, *_args, **_kwargs): + return (hidden_states,) + + for mode, expected in (("native", False), ("sglang", True), ("sglang_fused", True)): + model = model_cls(config_factory()) + model.layers = torch.nn.ModuleList([StubLayer()]) + model.norm = CaptureNorm(mode) + model(input_ids=torch.tensor([[0, 1]])) + assert model.norm.force_values == [expected], (name, mode) -def test_qwen3_5_gdn_gated_norm_remains_a_separate_exact_surface(): - config = _tiny_config(layer_types=["linear_attention", "full_attention"]) + +def _assert_dense_gdn_norm_remains_separate(): + config = _tiny_dense_config(layer_types=["linear_attention", "full_attention"]) config._qwen35_exact_contract = True config._qwen35_rmsnorm_family = "v2" set_rmsnorm_mode("sglang_fused") @@ -201,37 +280,208 @@ def test_qwen3_5_gdn_gated_norm_remains_a_separate_exact_surface(): assert not hasattr(layer.linear_attn.o_norm, "rmsnorm_family") -@pytest.mark.parametrize( - ("layer_idx", "mode", "expected"), - [(0, "sglang_fused", False), (1, "native", False), (1, "sglang", True), (1, "sglang_fused", True)], -) -def test_qwen3_5_layer_input_norm_selects_residual_family(layer_idx, mode, expected): - layer = Qwen3_5DecoderLayer(_tiny_config(), layer_idx=layer_idx) - norm = CaptureNorm(mode) - layer.input_layernorm = norm +def _cpu_v2_forward(x, weight, eps, *, residual=None, zero_centered=False): + norm_input = x if residual is None else x + residual + fp32 = norm_input.float() + inv_rms = torch.rsqrt(fp32.square().mean(dim=-1, keepdim=True) + eps) + scale = weight.float() + 1.0 if zero_centered else weight.float() + out = (fp32 * inv_rms * scale).to(x.dtype) + return out if residual is None else (out, norm_input) + + +def _cpu_rms_backward(normed_input, weight, eps, grad_output, grad_residual_out=None): + with torch.enable_grad(): + x = normed_input.detach().float().requires_grad_(True) + w = weight.detach().float().requires_grad_(True) + inv_rms = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps) + out = x * inv_rms * w + objective = (out * grad_output.float()).sum() + if grad_residual_out is not None: + objective = objective + (x * grad_residual_out.float()).sum() + return torch.autograd.grad(objective, (x, w)) + + +def _assert_v2_zero_centered_backward(monkeypatch): + monkeypatch.setattr(normalization, "rms_norm_v2", _cpu_v2_forward) + monkeypatch.setattr(normalization, "fused_rms_norm_backward", _cpu_rms_backward) + torch.manual_seed(11) + + x = torch.randn(3, 8, dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(8, dtype=torch.float32, requires_grad=True) + grad = torch.randn_like(x) + normalization._FamiliesV2ZeroCenteredRMSNorm.apply(x, weight, EPS).backward(grad) + + x_ref = x.detach().requires_grad_(True) + weight_ref = weight.detach().requires_grad_(True) + _cpu_v2_forward(x_ref, weight_ref, EPS, zero_centered=True).backward(grad) + assert torch.allclose(x.grad.float(), x_ref.grad.float(), atol=2e-2, rtol=2e-2) + assert torch.allclose(weight.grad, weight_ref.grad, atol=2e-5, rtol=2e-5) + + x = torch.randn(2, 8, dtype=torch.bfloat16, requires_grad=True) + residual = torch.randn(2, 8, dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(8, dtype=torch.float32, requires_grad=True) + grad_out = torch.randn_like(x) + grad_residual = torch.randn_like(residual) + out, residual_out = normalization._FamiliesV2ZeroCenteredResidualRMSNorm.apply(x, residual, weight, EPS) + torch.autograd.backward((out, residual_out), (grad_out, grad_residual)) + + x_ref = x.detach().requires_grad_(True) + residual_ref = residual.detach().requires_grad_(True) + weight_ref = weight.detach().requires_grad_(True) + ref_out, ref_residual = _cpu_v2_forward( + x_ref, + weight_ref, + EPS, + residual=residual_ref, + zero_centered=True, + ) + torch.autograd.backward((ref_out, ref_residual), (grad_out, grad_residual)) + assert torch.allclose(x.grad.float(), x_ref.grad.float(), atol=2e-2, rtol=2e-2) + assert torch.equal(x.grad, residual.grad) + assert torch.allclose(residual.grad.float(), residual_ref.grad.float(), atol=2e-2, rtol=2e-2) + assert torch.allclose(weight.grad, weight_ref.grad, atol=2e-5, rtol=2e-5) + + +@pytest.mark.cpu +def test_qwen3_5_norm_dispatch_site_and_backward_contract(monkeypatch): + for variant in VARIANTS: + _name, module, norm_cls, *_rest = variant + with monkeypatch.context() as dispatch_patch: + _assert_dispatch_policy(module, norm_cls, dispatch_patch) + with monkeypatch.context() as v2_patch: + _assert_v2_admission_and_dispatch(module, norm_cls, v2_patch) + _assert_v2_reaches_every_norm_site(variant) + _assert_layer_and_final_norm_site_policy(variant) + + _assert_dense_gdn_norm_remains_separate() + with monkeypatch.context() as backward_patch: + _assert_v2_zero_centered_backward(backward_patch) + + +def _assert_module_realization_matches_for_both_variants(): + torch.manual_seed(3) + hidden = torch.randn(N_TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) + residual = torch.randn(N_TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) + + for norm_cls in (Qwen3_5RMSNorm, Qwen3_5MoeRMSNorm): + set_rmsnorm_mode("sglang") + try: + sglang = norm_cls(HIDDEN, eps=EPS).cuda() + set_rmsnorm_mode("sglang_fused") + fused = norm_cls(HIDDEN, eps=EPS).cuda() + finally: + set_rmsnorm_mode("native") + with torch.no_grad(): + sglang.weight.copy_(torch.randn(HIDDEN, device="cuda")) + fused.weight.copy_(sglang.weight) + + with set_batch_invariant_mode(True), torch.no_grad(): + sglang_out, sglang_residual = sglang(hidden, residual=residual, prenorm=True) + fused_out, fused_residual = fused(hidden, residual=residual, prenorm=True) + assert torch.equal(sglang_out, fused_out) + assert torch.equal(sglang_residual, fused_residual) + assert torch.equal( + sglang(hidden, force_sglang_residual=True), + fused(hidden, force_sglang_residual=True), + ) + assert torch.equal(sglang(hidden), fused(hidden)) + + +def _assert_family1_matches_interpose(): + torch.manual_seed(21) + head_dim = 128 + x = torch.randn(256, 16, head_dim, device="cuda", dtype=torch.bfloat16) + set_rmsnorm_mode("sglang_fused") + try: + norm = Qwen3_5MoeRMSNorm(head_dim, eps=EPS, exact_contract=True).cuda() + finally: + set_rmsnorm_mode("native") + with torch.no_grad(): + norm.weight.copy_(torch.randn(head_dim, device="cuda")) + out = norm(x) + reference = rms_norm_batch_invariant(x.float(), 1.0 + norm.weight.float(), eps=EPS).to(x.dtype) + with set_batch_invariant_mode(True): + interpose = torch.nn.functional.rms_norm( + x.float(), + (head_dim,), + 1.0 + norm.weight.float(), + eps=EPS, + ).to(x.dtype) + assert torch.equal(out, reference) + assert torch.equal(out, interpose) + + +def _assert_moe_layer_forward_matches_realizations(): + torch.manual_seed(7) + config = _tiny_moe_config( + hidden_size=HIDDEN, + intermediate_size=1024, + num_attention_heads=16, + num_key_value_heads=8, + head_dim=128, + ) + set_rmsnorm_mode("sglang") + try: + layer = Qwen3_5MoeDecoderLayer(config, layer_idx=1).to(device="cuda", dtype=torch.bfloat16) + finally: + set_rmsnorm_mode("native") layer.self_attn = IdentityAttention() - layer.post_attention_layernorm = CaptureNorm(mode) - layer.mlp = torch.nn.Identity() - - hidden = torch.ones(1, 2, 8) - layer(hidden, position_embeddings=(hidden, hidden)) - - assert norm.force_values == [expected] - - -@pytest.mark.parametrize(("mode", "expected"), [("native", False), ("sglang", True), ("sglang_fused", True)]) -def test_qwen3_5_final_norm_selects_residual_family(mode, expected): - class StubLayer(torch.nn.Module): - layer_type = "full_attention" - - def forward(self, hidden_states, *args, **kwargs): - return (hidden_states,) - - model = Qwen3_5TextModel(_tiny_config()) - model.layers = torch.nn.ModuleList([StubLayer()]) - norm = CaptureNorm(mode) - model.norm = norm + with torch.no_grad(): + layer.input_layernorm.weight.copy_(torch.randn(HIDDEN, device="cuda", dtype=torch.bfloat16)) + layer.post_attention_layernorm.weight.copy_(torch.randn(HIDDEN, device="cuda", dtype=torch.bfloat16)) + + hidden = torch.randn(1, 128, HIDDEN, device="cuda", dtype=torch.bfloat16) + position = torch.zeros_like(hidden) + with set_batch_invariant_mode(True), torch.no_grad(): + for norm in (layer.input_layernorm, layer.post_attention_layernorm): + norm.mode = "sglang" + (sglang_out,) = layer(hidden, position_embeddings=(position, position)) + for norm in (layer.input_layernorm, layer.post_attention_layernorm): + norm.mode = "sglang_fused" + (fused_out,) = layer(hidden, position_embeddings=(position, position)) + assert torch.equal(sglang_out, fused_out) + + +def _assert_family2_residual_matches_serving_tree(): + from xorl.models.layers.normalization import ( # noqa: PLC0415 + fast_zero_centered_batch_invariant_residual_rms_norm, + native_zero_centered_rms_norm, + ) + from xorl.ops.batch_invariant_ops import mean_dim # noqa: PLC0415 - model(input_ids=torch.tensor([[0, 1]])) + torch.manual_seed(11) + x = torch.randn(513, HIDDEN, device="cuda", dtype=torch.bfloat16) + residual = torch.randn(513, HIDDEN, device="cuda", dtype=torch.bfloat16) + weight = (torch.randn(HIDDEN, device="cuda") * 0.02).to(torch.bfloat16) - assert norm.force_values == [expected] + set_rmsnorm_mode("sglang_fused") + try: + exact = Qwen3_5MoeRMSNorm(HIDDEN, eps=EPS, exact_contract=True).cuda() + ordinary = Qwen3_5MoeRMSNorm(HIDDEN, eps=EPS, exact_contract=False).cuda() + finally: + set_rmsnorm_mode("native") + with torch.no_grad(): + exact.weight.copy_(weight) + ordinary.weight.copy_(weight) + + summed = (x + residual).float() + variance = mean_dim(summed * summed, dim=-1, keepdim=True) + reference = (summed * torch.rsqrt(variance + EPS) * (1.0 + weight.float())).to(torch.bfloat16) + assert torch.equal(exact(x, residual=residual), reference) + assert torch.equal( + ordinary(x, residual=residual), + native_zero_centered_rms_norm(x + residual, weight, EPS), + ) + assert torch.equal( + fast_zero_centered_batch_invariant_residual_rms_norm(x + residual, weight, EPS), + reference, + ) + + +@requires_cuda +@pytest.mark.gpu +def test_qwen3_5_norm_bit_exact_model_contract(): + _assert_module_realization_matches_for_both_variants() + _assert_family1_matches_interpose() + _assert_moe_layer_forward_matches_realizations() + _assert_family2_residual_matches_serving_tree() diff --git a/tests/models/test_qwen3_5_trunk_wrap.py b/tests/models/test_qwen3_5_trunk_wrap.py index 646eb1fb..46db059f 100644 --- a/tests/models/test_qwen3_5_trunk_wrap.py +++ b/tests/models/test_qwen3_5_trunk_wrap.py @@ -11,7 +11,6 @@ router gate (contracted separately by the exact model program) and lm_head/embed. """ -import pytest import torch from xorl.lora.modules.linear import LoraLinear @@ -23,9 +22,7 @@ set_trunk_linear_contract, wrap_trunk_linears_batch_invariant, ) - - -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +from xorl.ops.bi_families_v2 import families_v2_enabled def _hybrid_config(**overrides) -> Qwen3_5MoeConfig: @@ -61,27 +58,40 @@ def _build(dtype=torch.bfloat16) -> Qwen3_5MoeForCausalLM: return Qwen3_5MoeForCausalLM(_hybrid_config()).to(dtype) -def test_exact_qwen_hook_enables_merged_lora_before_trunk_wrap(): +def _assert_exact_qwen_hook_enables_merged_lora_before_trunk_wrap(): + from xorl.ops.bi_families_v2 import _select_nonexact_families # noqa: PLC0415 + model = torch.nn.Module() - model.config = type("Config", (), {"model_type": "xorl_qwen3_5"})() + model.config = type( + "Config", + (), + {"model_type": "xorl_qwen3_5", "_qwen35_rmsnorm_family": "v1"}, + )() model.q_proj = LoraLinear(16, 16, r=2, lora_alpha=4, dtype=torch.bfloat16) + model.norm = torch.nn.Module() + model.norm.rmsnorm_family = "v1" try: + assert families_v2_enabled() is True assert model.q_proj.exact_merged_forward is False wrapped = _apply_qwen35_gdn_exact(model) + assert families_v2_enabled() is False assert model.q_proj.exact_merged_forward is True assert wrapped == {"q_proj": 1} assert model.q_proj._xorl_bi_trunk_wrapped is True finally: set_trunk_linear_contract(False) + _select_nonexact_families() -def test_qwen3_5_hybrid_trunk_wrap_selection(): +def test_qwen3_5_trunk_wrap_selection_policy(): + _assert_exact_qwen_hook_enables_merged_lora_before_trunk_wrap() + model = _build() try: wrapped = wrap_trunk_linears_batch_invariant(model) - assert not is_trunk_linear_contract_enabled(), "model wrapping must not leak numerical state process-wide" + assert is_trunk_linear_contract_enabled(), "model wrapping must arm the RMSNorm and trunk contract lane" # Leaf-name counts: q/k/v/o match BOTH the full-attn layer and the # GatedDeltaNet layer (2 each); gate_up/down match the dense-layer MLP @@ -132,24 +142,3 @@ def _is_wrapped(module): assert not _is_wrapped(model.model.embed_tokens) finally: set_trunk_linear_contract(False) - - -@requires_cuda -@pytest.mark.gpu -def test_qwen3_5_full_attn_forward_runs_under_trunk_wrap(): - """Wrapped full-attention + shared-expert + dense projections must run the - bf16 contract GEMM end-to-end (the runtime guard raises on any non-bf16 - operand).""" - torch.manual_seed(1) - config = _hybrid_config(layer_types=["full_attention", "full_attention"], _moe_implementation="eager") - model = Qwen3_5MoeForCausalLM(config).to(device="cuda", dtype=torch.bfloat16).eval() - try: - wrapped = wrap_trunk_linears_batch_invariant(model) - assert wrapped["shared_expert_gate"] == 1 - input_ids = torch.randint(0, config.vocab_size, (1, 8), device="cuda") - with torch.no_grad(): - out = model(input_ids=input_ids) - assert out.last_hidden_state.dtype == torch.bfloat16 - assert torch.isfinite(out.last_hidden_state.float()).all() - finally: - set_trunk_linear_contract(False) diff --git a/tests/models/test_qwen3_moe_cache.py b/tests/models/test_qwen3_moe_cache.py index fd593c15..68a04213 100644 --- a/tests/models/test_qwen3_moe_cache.py +++ b/tests/models/test_qwen3_moe_cache.py @@ -28,6 +28,8 @@ def _tiny_qwen3_moe_config() -> Qwen3MoeConfig: def test_qwen3_moe_decode_cache_matches_full_forward_with_r3_replay(): + _assert_flash_attention_decode_cache_uses_prefix_attention() + torch.manual_seed(0) set_replay_stage(None) RoutingReplay.clear_all() @@ -131,7 +133,7 @@ def test_qwen3_moe_decode_cache_matches_full_forward_with_r3_replay(): RoutingReplay._target_device = original_target_device -def test_flash_attention_decode_cache_uses_prefix_attention(): +def _assert_flash_attention_decode_cache_uses_prefix_attention(): query = torch.randn(1, 1, 2, 4) key = torch.randn(1, 5, 2, 4) diff --git a/tests/models/test_qwen3_moe_fused_lora.py b/tests/models/test_qwen3_moe_fused_lora.py deleted file mode 100644 index f8e109f4..00000000 --- a/tests/models/test_qwen3_moe_fused_lora.py +++ /dev/null @@ -1,602 +0,0 @@ -"""Tests for MoE experts with LoRA across all backends (eager, triton, native, quack).""" - -import pytest -import torch -import torch.nn as nn - -from xorl.lora import LoraLinear, inject_lora_into_model -from xorl.lora.mapping import can_apply_lora, get_lora_class_for_module -from xorl.models.layers.moe import MOE_EXPERT_BACKENDS, MoEBlock, MoEExperts, MoEExpertsLoRA, MoELoRAConfig -from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( - Qwen3MoeSparseExperts, - Qwen3MoeTritonExperts, -) - - -class MockConfig: - """Mock config for testing.""" - - def __init__( - self, - num_experts=4, - hidden_size=64, - moe_intermediate_size=128, - hidden_act="silu", - num_experts_per_tok=2, - norm_topk_prob=True, - ): - self.num_experts = num_experts - self.hidden_size = hidden_size - self.moe_intermediate_size = moe_intermediate_size - self.hidden_act = hidden_act - self.num_experts_per_tok = num_experts_per_tok - self.norm_topk_prob = norm_topk_prob - - -# --------------------------------------------------------------------------- -# 1. Base MoE experts: init, shapes, backends, LoRA mapping, registry, -# LoRA init (all backends), triton forward -# --------------------------------------------------------------------------- - - -class TestMoEExpertsBaseAndLoRAInit: - """Comprehensive tests for base MoEExperts and LoRA initialization across all backends.""" - - def test_init_shapes_backends_lora_mapping_and_lora_init(self): - """Test init fields, weight shapes, LoRA mapping, backend registry, and LoRA init for all backends.""" - config = MockConfig() - - # Qwen3 subclass inits - triton_exp = Qwen3MoeTritonExperts(config) - assert triton_exp.num_experts == config.num_experts - assert triton_exp.hidden_dim == config.hidden_size - assert triton_exp.intermediate_size == config.moe_intermediate_size - assert triton_exp.moe_implementation == "triton" - - sparse_exp = Qwen3MoeSparseExperts(config) - assert sparse_exp.moe_implementation == "eager" - - # Direct MoEExperts with all backends - for backend in ["eager", "triton", "native", "quack"]: - experts = MoEExperts( - num_experts=config.num_experts, - hidden_dim=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - moe_implementation=backend, - ) - assert experts.moe_implementation == backend - - # Weight shapes on a single instance - experts = MoEExperts( - num_experts=config.num_experts, - hidden_dim=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - ) - assert experts.gate_proj.shape == (config.num_experts, config.hidden_size, config.moe_intermediate_size) - assert experts.up_proj.shape == (config.num_experts, config.hidden_size, config.moe_intermediate_size) - assert experts.down_proj.shape == (config.num_experts, config.moe_intermediate_size, config.hidden_size) - - # LoRA mapping registered - assert can_apply_lora(experts) - assert get_lora_class_for_module(experts) is MoEExpertsLoRA - - # Backend registry - assert "eager" in MOE_EXPERT_BACKENDS - assert "fused" not in MOE_EXPERT_BACKENDS - - # LoRA init for all backends: frozen/trainable, shapes, repr - for backend in ["eager", "triton", "native", "quack"]: - lora_config = MoELoRAConfig(r=8, lora_alpha=16, target_modules=["gate_proj", "up_proj", "down_proj"]) - r = lora_config.r - lora_experts = MoEExpertsLoRA( - num_experts=config.num_experts, - hidden_dim=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - moe_implementation=backend, - lora_config=lora_config, - ) - assert lora_experts.num_experts == config.num_experts - assert lora_experts.moe_implementation == backend - assert lora_experts.lora_config == lora_config - - # Base weights frozen - assert not lora_experts.gate_proj.requires_grad - assert not lora_experts.up_proj.requires_grad - assert not lora_experts.down_proj.requires_grad - - # LoRA weights trainable, B initialized to zeros - for name in lora_config.target_modules: - lora_A = getattr(lora_experts, f"{name}_lora_A") - lora_B = getattr(lora_experts, f"{name}_lora_B") - assert isinstance(lora_A, nn.Parameter) and lora_A.requires_grad - assert isinstance(lora_B, nn.Parameter) and lora_B.requires_grad - assert torch.allclose(lora_B, torch.zeros_like(lora_B)) - - # LoRA weight shapes - assert lora_experts.gate_proj_lora_A.shape == (config.num_experts, config.hidden_size, r) - assert lora_experts.gate_proj_lora_B.shape == (config.num_experts, r, config.moe_intermediate_size) - assert lora_experts.down_proj_lora_A.shape == (config.num_experts, config.moe_intermediate_size, r) - assert lora_experts.down_proj_lora_B.shape == (config.num_experts, r, config.hidden_size) - - repr_str = lora_experts.extra_repr() - assert f"num_experts={config.num_experts}" in repr_str - assert f"r={lora_config.r}" in repr_str - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for triton MoE") - def test_triton_forward(self): - """Test forward pass with triton backend on GPU.""" - config = MockConfig() - experts = Qwen3MoeTritonExperts(config) - nn.init.xavier_normal_(experts.gate_proj.data) - nn.init.xavier_normal_(experts.up_proj.data) - nn.init.xavier_normal_(experts.down_proj.data) - - device = "cuda" - experts = experts.to(device).to(torch.bfloat16) - - num_tokens, top_k = 16, 2 - hidden_states = torch.randn(num_tokens, config.hidden_size, device=device, dtype=torch.bfloat16) - routing_weights = torch.softmax(torch.randn(num_tokens, top_k, device=device, dtype=torch.bfloat16), dim=-1) - selected_experts = torch.randint(0, config.num_experts, (num_tokens, top_k), device=device) - - output = experts(hidden_states, routing_weights, selected_experts) - assert output.shape == hidden_states.shape - - -# --------------------------------------------------------------------------- -# 2. Eager LoRA forward/backward (CPU) + MoEBlock end-to-end -# --------------------------------------------------------------------------- - - -class TestMoEExpertsLoRAEager: - """Test eager LoRA forward/backward on CPU, including via MoEBlock.""" - - def test_eager_forward_backward_and_moe_block(self): - """Test eager per-expert forward, backward gradients, and end-to-end MoEBlock.""" - lora_config = MoELoRAConfig(r=4, lora_alpha=8) - experts = MoEExpertsLoRA( - num_experts=4, - hidden_dim=32, - intermediate_size=64, - moe_implementation="eager", - lora_config=lora_config, - ) - nn.init.xavier_normal_(experts.gate_proj.data) - nn.init.xavier_normal_(experts.up_proj.data) - nn.init.xavier_normal_(experts.down_proj.data) - - # Forward - hidden = torch.randn(8, 32) - out = experts(hidden, expert_idx=0) - assert out.shape == (8, 32) - - # Backward - hidden = torch.randn(8, 32, requires_grad=True) - out = experts(hidden, expert_idx=0) - out.sum().backward() - assert experts.gate_proj_lora_A.grad is not None - assert experts.gate_proj_lora_B.grad is not None - assert experts.gate_proj.grad is None # base frozen - - # MoEBlock end-to-end - block = MoEBlock( - hidden_size=32, - num_experts=4, - top_k=2, - intermediate_size=64, - moe_implementation="eager", - ) - nn.init.xavier_normal_(block.experts.gate_proj.data) - nn.init.xavier_normal_(block.experts.up_proj.data) - nn.init.xavier_normal_(block.experts.down_proj.data) - nn.init.xavier_normal_(block.gate.weight.data) - - block.inject_lora(r=4, lora_alpha=8) - assert isinstance(block.experts, MoEExpertsLoRA) - - hidden = torch.randn(2, 4, 32) - output, router_logits = block(hidden) - assert output.shape == hidden.shape - - output.sum().backward() - assert block.experts.gate_proj_lora_A.grad is not None - - -# --------------------------------------------------------------------------- -# 3. GPU LoRA forward/backward (triton + native) + MoEBlock -# --------------------------------------------------------------------------- - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -class TestMoEExpertsLoRAGPU: - """Test triton and native LoRA forward/backward on GPU, including via MoEBlock.""" - - @pytest.mark.parametrize("backend", ["triton", "native"]) - def test_forward_backward_and_moe_block(self, backend): - """Test GPU LoRA forward, backward, and end-to-end MoEBlock for triton and native.""" - lora_config = MoELoRAConfig(r=4, lora_alpha=8) - exp = MoEExpertsLoRA( - num_experts=4, - hidden_dim=32, - intermediate_size=64, - moe_implementation=backend, - lora_config=lora_config, - ) - nn.init.xavier_normal_(exp.gate_proj.data) - nn.init.xavier_normal_(exp.up_proj.data) - nn.init.xavier_normal_(exp.down_proj.data) - - device = "cuda" - exp = exp.to(device).to(torch.bfloat16) - - num_tokens, top_k = 16, 2 - hidden = torch.randn(num_tokens, 32, device=device, dtype=torch.bfloat16, requires_grad=True) - weights = torch.softmax(torch.randn(num_tokens, top_k, device=device, dtype=torch.bfloat16), dim=-1) - selected = torch.randint(0, 4, (num_tokens, top_k), device=device) - - # Forward - output = exp(hidden, weights, selected) - assert output.shape == hidden.shape - - # Backward - output.sum().backward() - assert exp.gate_proj_lora_A.grad is not None - assert exp.down_proj_lora_B.grad is not None - assert exp.gate_proj.grad is None # base frozen - - # MoEBlock end-to-end (test only native to avoid duplicate) - if backend == "native": - block = MoEBlock( - hidden_size=32, - num_experts=4, - top_k=2, - intermediate_size=64, - moe_implementation="native", - ) - nn.init.xavier_normal_(block.experts.gate_proj.data) - nn.init.xavier_normal_(block.experts.up_proj.data) - nn.init.xavier_normal_(block.experts.down_proj.data) - nn.init.xavier_normal_(block.gate.weight.data) - block = block.to(device).to(torch.bfloat16) - - block.inject_lora(r=4, lora_alpha=8) - assert isinstance(block.experts, MoEExpertsLoRA) - assert block.experts.moe_implementation == "native" - - hidden2 = torch.randn(2, 4, 32, device=device, dtype=torch.bfloat16) - output2, router_logits = block(hidden2) - assert output2.shape == hidden2.shape - output2.sum().backward() - assert block.experts.gate_proj_lora_A.grad is not None - - -# --------------------------------------------------------------------------- -# 4. Cross-backend numerical correctness (all combined) -# --------------------------------------------------------------------------- - - -def _make_lora_block( - backend, - num_experts, - hidden_dim, - intermediate, - r, - lora_alpha, - device, - dtype, -): - """Create a MoEBlock with LoRA on the given backend, with deterministic init.""" - block = MoEBlock( - hidden_size=hidden_dim, - num_experts=num_experts, - top_k=2, - intermediate_size=intermediate, - moe_implementation=backend, - ) - torch.manual_seed(42) - nn.init.xavier_normal_(block.experts.gate_proj.data) - nn.init.xavier_normal_(block.experts.up_proj.data) - nn.init.xavier_normal_(block.experts.down_proj.data) - nn.init.xavier_normal_(block.gate.weight.data) - block = block.to(device).to(dtype) - - torch.manual_seed(123) - block.inject_lora(r=r, lora_alpha=lora_alpha) - return block - - -def _copy_block_weights(src_block, dst_block): - """Copy all weights from src to dst block (base + LoRA + gate).""" - with torch.no_grad(): - dst_block.gate.weight.copy_(src_block.gate.weight) - src_exp = src_block.experts - dst_exp = dst_block.experts - dst_exp.gate_proj.copy_(src_exp.gate_proj) - dst_exp.up_proj.copy_(src_exp.up_proj) - dst_exp.down_proj.copy_(src_exp.down_proj) - for proj in ["gate_proj", "up_proj", "down_proj"]: - getattr(dst_exp, f"{proj}_lora_A").copy_(getattr(src_exp, f"{proj}_lora_A")) - getattr(dst_exp, f"{proj}_lora_B").copy_(getattr(src_exp, f"{proj}_lora_B")) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -class TestCrossBackendConsistency: - """Test numerical agreement across backends: zero-LoRA, cross-backend, nonzero-LoRA.""" - - NUM_EXPERTS = 4 - HIDDEN_DIM = 64 - INTERMEDIATE = 128 - R = 8 - LORA_ALPHA = 16 - DTYPE = torch.bfloat16 - - def _make_pair(self, ref_backend, test_backend, device): - ref = _make_lora_block( - ref_backend, - self.NUM_EXPERTS, - self.HIDDEN_DIM, - self.INTERMEDIATE, - self.R, - self.LORA_ALPHA, - device, - self.DTYPE, - ) - test = _make_lora_block( - test_backend, - self.NUM_EXPERTS, - self.HIDDEN_DIM, - self.INTERMEDIATE, - self.R, - self.LORA_ALPHA, - device, - self.DTYPE, - ) - _copy_block_weights(ref, test) - return ref, test - - def test_zero_lora_and_nonzero_lora(self): - """With lora_B=0, LoRA output must equal base; nonzero LoRA must change output.""" - device = "cuda" - for backend in ["eager", "triton", "native"]: - # --- zero LoRA matches base --- - base_block = MoEBlock( - hidden_size=self.HIDDEN_DIM, - num_experts=self.NUM_EXPERTS, - top_k=2, - intermediate_size=self.INTERMEDIATE, - moe_implementation=backend, - ) - torch.manual_seed(42) - nn.init.xavier_normal_(base_block.experts.gate_proj.data) - nn.init.xavier_normal_(base_block.experts.up_proj.data) - nn.init.xavier_normal_(base_block.experts.down_proj.data) - nn.init.xavier_normal_(base_block.gate.weight.data) - base_block = base_block.to(device).to(self.DTYPE) - - lora_block = MoEBlock( - hidden_size=self.HIDDEN_DIM, - num_experts=self.NUM_EXPERTS, - top_k=2, - intermediate_size=self.INTERMEDIATE, - moe_implementation=backend, - ) - torch.manual_seed(42) - nn.init.xavier_normal_(lora_block.experts.gate_proj.data) - nn.init.xavier_normal_(lora_block.experts.up_proj.data) - nn.init.xavier_normal_(lora_block.experts.down_proj.data) - nn.init.xavier_normal_(lora_block.gate.weight.data) - lora_block = lora_block.to(device).to(self.DTYPE) - lora_block.inject_lora(r=self.R, lora_alpha=self.LORA_ALPHA) - - torch.manual_seed(999) - hidden = torch.randn(2, 8, self.HIDDEN_DIM, device=device, dtype=self.DTYPE) - base_out, _ = base_block(hidden) - lora_out, _ = lora_block(hidden) - torch.testing.assert_close( - lora_out, - base_out, - atol=1e-3, - rtol=1e-2, - msg=f"[{backend}] Zero-LoRA output should match base model", - ) - - # --- nonzero LoRA changes output --- - block = _make_lora_block( - backend, - self.NUM_EXPERTS, - self.HIDDEN_DIM, - self.INTERMEDIATE, - self.R, - self.LORA_ALPHA, - device, - self.DTYPE, - ) - with torch.no_grad(): - for proj in ["gate_proj", "up_proj", "down_proj"]: - lora_B = getattr(block.experts, f"{proj}_lora_B") - nn.init.xavier_normal_(lora_B) - - base_block2 = ( - MoEBlock( - hidden_size=self.HIDDEN_DIM, - num_experts=self.NUM_EXPERTS, - top_k=2, - intermediate_size=self.INTERMEDIATE, - moe_implementation=backend, - ) - .to(device) - .to(self.DTYPE) - ) - with torch.no_grad(): - base_block2.gate.weight.copy_(block.gate.weight) - base_block2.experts.gate_proj.copy_(block.experts.gate_proj) - base_block2.experts.up_proj.copy_(block.experts.up_proj) - base_block2.experts.down_proj.copy_(block.experts.down_proj) - - torch.manual_seed(999) - hidden2 = torch.randn(2, 8, self.HIDDEN_DIM, device=device, dtype=self.DTYPE) - base_out2, _ = base_block2(hidden2) - lora_out2, _ = block(hidden2) - - diff = (lora_out2 - base_out2).abs().max().item() - assert diff > 1e-3, f"[{backend}] Non-zero LoRA should change the output, but max diff={diff}" - - def test_cross_backend_output_and_gradients(self): - """Cross-backend outputs and LoRA gradients should match.""" - for ref_backend, test_backend in [("eager", "native"), ("eager", "triton"), ("triton", "native")]: - ref, test = self._make_pair(ref_backend, test_backend, "cuda") - - torch.manual_seed(999) - h1 = torch.randn(2, 8, self.HIDDEN_DIM, device="cuda", dtype=self.DTYPE) - h2 = h1.clone() - - ref_out, _ = ref(h1) - ref_out.sum().backward() - test_out, _ = test(h2) - test_out.sum().backward() - - torch.testing.assert_close( - test_out, - ref_out, - atol=0.05, - rtol=0.02, - msg=f"{ref_backend} vs {test_backend} output mismatch", - ) - - for proj in ["gate_proj", "up_proj", "down_proj"]: - ref_grad_A = getattr(ref.experts, f"{proj}_lora_A").grad - test_grad_A = getattr(test.experts, f"{proj}_lora_A").grad - ref_grad_B = getattr(ref.experts, f"{proj}_lora_B").grad - test_grad_B = getattr(test.experts, f"{proj}_lora_B").grad - - assert ref_grad_A is not None, f"ref {proj}_lora_A grad is None" - assert test_grad_A is not None, f"test {proj}_lora_A grad is None" - - torch.testing.assert_close( - test_grad_A, - ref_grad_A, - atol=0.05, - rtol=0.05, - msg=f"Gradient mismatch: {proj}_lora_A ({ref_backend} vs {test_backend})", - ) - torch.testing.assert_close( - test_grad_B, - ref_grad_B, - atol=0.05, - rtol=0.05, - msg=f"Gradient mismatch: {proj}_lora_B ({ref_backend} vs {test_backend})", - ) - - -# --------------------------------------------------------------------------- -# 5. from_module + LoRA injection + error handling (all combined) -# --------------------------------------------------------------------------- - - -class TestFromModuleAndInjection: - """Test from_module, inject_lora, and error handling.""" - - def test_from_module_inject_lora_and_errors(self): - """Test from_module for all backends, inject_lora via both APIs, Qwen3 subclass, and error handling.""" - config = MockConfig() - - for backend in ["eager", "triton", "native", "quack"]: - # from_module - base = MoEExperts( - num_experts=config.num_experts, - hidden_dim=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - moe_implementation=backend, - ) - nn.init.xavier_normal_(base.gate_proj.data) - nn.init.xavier_normal_(base.up_proj.data) - nn.init.xavier_normal_(base.down_proj.data) - - lora_exp = MoEExpertsLoRA.from_module(base, r=8, lora_alpha=16) - assert lora_exp.moe_implementation == backend - assert torch.allclose(lora_exp.gate_proj, base.gate_proj) - - # inject_lora_into_model - - class SimpleModel(nn.Module): - def __init__(self, config, backend): - super().__init__() - self.experts = MoEExperts( - num_experts=config.num_experts, - hidden_dim=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - moe_implementation=backend, - ) - nn.init.xavier_normal_(self.experts.gate_proj.data) - nn.init.xavier_normal_(self.experts.up_proj.data) - nn.init.xavier_normal_(self.experts.down_proj.data) - - model = SimpleModel(config, backend) - inject_lora_into_model(model, r=8, lora_alpha=16, target_modules=["experts"]) - assert isinstance(model.experts, MoEExpertsLoRA) - assert model.experts.moe_implementation == backend - assert hasattr(model.experts, "gate_proj_lora_A") - - # MoEBlock.inject_lora - block = MoEBlock( - hidden_size=config.hidden_size, - num_experts=config.num_experts, - top_k=config.num_experts_per_tok, - intermediate_size=config.moe_intermediate_size, - moe_implementation=backend, - ) - block.inject_lora(r=8, lora_alpha=16) - assert isinstance(block.experts, MoEExpertsLoRA) - assert block.experts.moe_implementation == backend - assert block.lora_adapter == "injected" - - # from_module with Qwen3 subclass - base = Qwen3MoeTritonExperts(config) - nn.init.xavier_normal_(base.gate_proj.data) - nn.init.xavier_normal_(base.up_proj.data) - nn.init.xavier_normal_(base.down_proj.data) - - lora_exp = MoEExpertsLoRA.from_module(base, r=8, lora_alpha=16) - assert isinstance(lora_exp, MoEExpertsLoRA) - assert torch.allclose(lora_exp.gate_proj, base.gate_proj) - - # Error handling - - class ModelA(nn.Module): - def __init__(self): - super().__init__() - self.layer1 = nn.Linear(64, 64) - - with pytest.raises(ValueError, match="No modules found matching target_modules"): - inject_lora_into_model(ModelA(), r=8, lora_alpha=16, target_modules=["nonexistent_proj"]) - - class UnsupportedModule(nn.Module): - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.randn(64, 64)) - - def forward(self, x): - return x @ self.weight - - class ModelB(nn.Module): - def __init__(self): - super().__init__() - self.custom_layer = UnsupportedModule() - - with pytest.raises(ValueError, match="No modules found matching target_modules"): - inject_lora_into_model(ModelB(), r=8, lora_alpha=16, target_modules=["custom_layer"]) - - class ModelC(nn.Module): - def __init__(self): - super().__init__() - self.q_proj = nn.Linear(64, 64) - self.v_proj = nn.Linear(64, 64) - - model = ModelC() - inject_lora_into_model(model, r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"]) - assert isinstance(model.q_proj, LoraLinear) - assert isinstance(model.v_proj, LoraLinear) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/models/test_qwen3_moe_rmsnorm.py b/tests/models/test_qwen3_moe_rmsnorm.py deleted file mode 100644 index 02c6001b..00000000 --- a/tests/models/test_qwen3_moe_rmsnorm.py +++ /dev/null @@ -1,340 +0,0 @@ -import pytest -import torch - -from xorl.models.layers.normalization import ( - RMS_NORM_FAMILY_NO_RESIDUAL, - RMS_NORM_FAMILY_RESIDUAL_TREE, -) -from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig -from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import ( - Qwen3MoeDecoderLayer, - _materialize_moe_tp_shards_with_residual, - _materialize_o_proj_partial_residual, -) - - -class CaptureInputNorm(torch.nn.Module): - def __init__(self, mode: str): - super().__init__() - self.mode = mode - self.family_values = [] - - def forward(self, hidden_states, *, force_sglang_residual=False, family=None): - assert force_sglang_residual is False, "call sites declare family, not force flags" - self.family_values.append(family) - return hidden_states - - -class CaptureDelayedInputNorm(torch.nn.Module): - mode = "native" - - def __init__(self): - super().__init__() - self.calls = [] - - def forward( - self, - hidden_states, - residual=None, - prenorm=False, - *, - force_sglang_residual=False, - force_sglang_residual_kernel=False, - family=None, - ): - self.calls.append( - { - "hidden_states": hidden_states.detach().clone(), - "residual": None if residual is None else residual.detach().clone(), - "prenorm": prenorm, - "force_sglang_residual": force_sglang_residual, - "force_sglang_residual_kernel": force_sglang_residual_kernel, - "family": family, - } - ) - if residual is None: - return hidden_states + 10.0 - return hidden_states + 10.0, hidden_states + residual - - -class IdentityAttention(torch.nn.Module): - def forward(self, hidden_states, **kwargs): - return hidden_states, None - - -class PartialOutputAttention(torch.nn.Module): - def __init__(self, partials): - super().__init__() - self.partials = partials - - def forward(self, hidden_states, **kwargs): - del hidden_states, kwargs - output = self.partials[0] + self.partials[1] - output._xorl_o_proj_tp_partials = self.partials - return output, None - - -class IdentityPostAttentionNorm(torch.nn.Module): - def forward(self, hidden_states, residual=None, prenorm=False, **kwargs): - return hidden_states, residual - - -def _small_config() -> Qwen3MoeConfig: - return Qwen3MoeConfig( - hidden_size=4, - intermediate_size=8, - num_attention_heads=2, - num_key_value_heads=2, - num_hidden_layers=2, - num_experts=0, - use_qk_norm=False, - _attn_implementation="eager", - ) - - -@pytest.mark.parametrize( - ("layer_idx", "mode", "expected_family"), - [ - (0, "sglang", RMS_NORM_FAMILY_NO_RESIDUAL), - (1, "native", RMS_NORM_FAMILY_RESIDUAL_TREE), - (1, "sglang", RMS_NORM_FAMILY_RESIDUAL_TREE), - (0, "sglang_fused", RMS_NORM_FAMILY_NO_RESIDUAL), - (1, "sglang_fused", RMS_NORM_FAMILY_RESIDUAL_TREE), - ], -) -def test_qwen3_moe_layer_input_norm_declares_family_by_layer(layer_idx, mode, expected_family): - """The input-norm call site declares the serving family explicitly: layer-0 is a - no-residual site, layer>0 is a pre-summed residual-tree site (the 2026-07-04 - norm-seed contract). The declaration is mode-independent; RMSNorm keeps the - residual-tree dispatch confined to the sglang modes.""" - layer = Qwen3MoeDecoderLayer(_small_config(), layer_idx=layer_idx) - assert layer.layer_idx == layer_idx - - input_norm = CaptureInputNorm(mode) - layer.input_layernorm = input_norm - layer.self_attn = IdentityAttention() - layer.post_attention_layernorm = IdentityPostAttentionNorm() - - hidden_states = torch.ones(1, 2, 4) - layer._pre_mlp_forward(hidden_states, position_embeddings=(hidden_states, hidden_states)) - - assert input_norm.family_values == [expected_family] - - -def test_qwen3_moe_model_final_norm_declares_residual_tree_family(): - from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import Qwen3MoeModel - - config = _small_config() - config.pad_token_id = 0 - model = Qwen3MoeModel(config) - - assert model.norm.family == RMS_NORM_FAMILY_RESIDUAL_TREE - - class StubLayer(torch.nn.Module): - def forward(self, hidden_states, *args, **kwargs): - return (hidden_states,) - - model.layers = torch.nn.ModuleList([StubLayer()]) - final_norm = CaptureInputNorm("sglang_fused") - model.norm = final_norm - - model(input_ids=torch.tensor([[0, 1]])) - - # The call is bare: the family lives on the module (declared at construction). - assert final_norm.family_values == [None] - - -def test_qwen3_moe_layer_consumes_delayed_residual_pair_at_input_norm(): - layer = Qwen3MoeDecoderLayer(_small_config(), layer_idx=1) - - input_norm = CaptureDelayedInputNorm() - layer.input_layernorm = input_norm - layer.self_attn = IdentityAttention() - layer.post_attention_layernorm = IdentityPostAttentionNorm() - captures = {} - layer._diagnostic_capture_component = lambda name, tensor: captures.setdefault(name, tensor.detach().clone()) - - hidden_delta = torch.full((1, 2, 4), 2.0) - residual = torch.full((1, 2, 4), 3.0) - hidden_states, post_attention_residual = layer._pre_mlp_forward( - (hidden_delta, residual), - position_embeddings=(hidden_delta, hidden_delta), - ) - - assert len(input_norm.calls) == 1 - call = input_norm.calls[0] - assert call["prenorm"] is True - assert call["force_sglang_residual"] is False - torch.testing.assert_close(call["hidden_states"], hidden_delta) - torch.testing.assert_close(call["residual"], residual) - torch.testing.assert_close(hidden_states, hidden_delta + 10.0) - torch.testing.assert_close(post_attention_residual, hidden_delta + residual) - torch.testing.assert_close(captures["delayed_pair_delta"], hidden_delta) - torch.testing.assert_close(captures["delayed_pair_residual"], residual) - assert "delayed_pair_shard_sum" not in captures - assert "delayed_pair_shard_materialized" not in captures - torch.testing.assert_close(captures["materialized_layer_input"], hidden_delta + residual) - torch.testing.assert_close(captures["input_norm_residual"], hidden_delta + residual) - torch.testing.assert_close(captures["input_norm"], hidden_delta + 10.0) - torch.testing.assert_close(captures["post_attention_norm_input"], hidden_delta + 10.0) - torch.testing.assert_close(captures["post_attention_norm_residual"], hidden_delta + residual) - torch.testing.assert_close(captures["post_attention_norm"], hidden_delta + 10.0) - torch.testing.assert_close(captures["post_attention_residual"], hidden_delta + residual) - - -def test_qwen3_moe_layer_consumes_delayed_tp_shards_at_input_norm(monkeypatch): - monkeypatch.setenv("XORL_QWEN3_MOE_DELAYED_RESIDUAL_PAIR_TP_SHARD_CARRY", "1") - layer = Qwen3MoeDecoderLayer(_small_config(), layer_idx=1) - - input_norm = CaptureDelayedInputNorm() - layer.input_layernorm = input_norm - layer.self_attn = IdentityAttention() - layer.post_attention_layernorm = IdentityPostAttentionNorm() - captures = {} - layer._diagnostic_capture_component = lambda name, tensor: captures.setdefault(name, tensor.detach().clone()) - - hidden_delta = torch.full((1, 2, 4), 99.0) - residual = torch.full((1, 2, 4), 3.0) - shard0 = torch.full((1, 2, 4), 2.0) - shard1 = torch.full((1, 2, 4), 5.0) - hidden_delta._xorl_sglang_moe_tp_shards = (shard0, shard1) - - hidden_states, post_attention_residual = layer._pre_mlp_forward( - (hidden_delta, residual), - position_embeddings=(hidden_delta, hidden_delta), - ) - - expected_materialized = residual + shard0 + shard1 - assert len(input_norm.calls) == 1 - call = input_norm.calls[0] - assert call["prenorm"] is False - assert call["residual"] is None - torch.testing.assert_close(call["hidden_states"], expected_materialized) - torch.testing.assert_close(hidden_states, expected_materialized + 10.0) - torch.testing.assert_close(post_attention_residual, expected_materialized) - torch.testing.assert_close(captures["delayed_pair_delta"], hidden_delta) - torch.testing.assert_close(captures["delayed_pair_residual"], residual) - torch.testing.assert_close(captures["delayed_pair_shard_sum"], shard0 + shard1) - torch.testing.assert_close(captures["delayed_pair_shard_materialized"], expected_materialized) - torch.testing.assert_close(captures["materialized_layer_input"], expected_materialized) - torch.testing.assert_close(captures["input_norm_residual"], expected_materialized) - torch.testing.assert_close(captures["input_norm"], expected_materialized + 10.0) - torch.testing.assert_close(captures["post_attention_norm_input"], expected_materialized + 10.0) - torch.testing.assert_close(captures["post_attention_norm_residual"], expected_materialized) - torch.testing.assert_close(captures["post_attention_residual"], expected_materialized) - - -def test_qwen3_moe_tp_shard_carry_reduces_shards_before_residual(monkeypatch): - monkeypatch.setenv("XORL_QWEN3_MOE_DELAYED_RESIDUAL_PAIR_TP_SHARD_CARRY", "1") - - residual = torch.tensor([1.0], dtype=torch.bfloat16) - hidden_delta = torch.tensor([99.0], dtype=torch.bfloat16) - shard0 = torch.tensor([0.002], dtype=torch.bfloat16) - shard1 = torch.tensor([0.002], dtype=torch.bfloat16) - hidden_delta._xorl_sglang_moe_tp_shards = (shard0, shard1) - - actual = _materialize_moe_tp_shards_with_residual(hidden_delta, residual) - expected = residual + (shard0 + shard1) - residual_first = (residual + shard0) + shard1 - - torch.testing.assert_close(actual, expected) - assert actual.item() != residual_first.item() - - -def test_qwen3_moe_layer_can_consume_o_proj_partials_at_post_attention_boundary(monkeypatch): - monkeypatch.setenv("XORL_QWEN3_MOE_POST_ATTENTION_O_PROJ_PARTIAL_RESIDUAL", "1") - - layer = Qwen3MoeDecoderLayer(_small_config(), layer_idx=1) - - input_norm = CaptureDelayedInputNorm() - post_norm = CaptureDelayedInputNorm() - partial0 = torch.full((1, 2, 4), 2.0) - partial1 = torch.full((1, 2, 4), 5.0) - layer.input_layernorm = input_norm - layer.self_attn = PartialOutputAttention((partial0, partial1)) - layer.post_attention_layernorm = post_norm - captures = {} - layer._diagnostic_capture_component = lambda name, tensor: captures.setdefault(name, tensor.detach().clone()) - - hidden_states = torch.full((1, 2, 4), 3.0) - output, post_attention_residual = layer._pre_mlp_forward( - hidden_states, - position_embeddings=(hidden_states, hidden_states), - ) - - expected_residual = hidden_states + partial0 + partial1 - assert len(post_norm.calls) == 1 - call = post_norm.calls[0] - assert call["residual"] is None - assert call["prenorm"] is False - torch.testing.assert_close(call["hidden_states"], expected_residual) - torch.testing.assert_close(output, expected_residual + 10.0) - torch.testing.assert_close(post_attention_residual, expected_residual) - torch.testing.assert_close(captures["post_attention_o_proj_partial_sum"], partial0 + partial1) - torch.testing.assert_close(captures["post_attention_partial_residual"], expected_residual) - torch.testing.assert_close(captures["post_attention_norm_input"], expected_residual) - torch.testing.assert_close(captures["post_attention_residual"], expected_residual) - - -def test_qwen3_moe_can_capture_o_proj_partial_residual_candidates_without_applying(monkeypatch): - monkeypatch.setenv("XORL_QWEN3_MOE_CAPTURE_O_PROJ_PARTIAL_RESIDUAL_CANDIDATES", "1") - monkeypatch.setenv("XORL_QWEN3_MOE_CAPTURE_O_PROJ_PARTIAL_RESIDUAL_CANDIDATES_LAYERS", "1") - - layer = Qwen3MoeDecoderLayer(_small_config(), layer_idx=1) - - input_norm = CaptureDelayedInputNorm() - post_norm = CaptureDelayedInputNorm() - partial0 = torch.full((1, 2, 4), 0.002, dtype=torch.bfloat16) - partial1 = torch.full((1, 2, 4), 0.002, dtype=torch.bfloat16) - layer.input_layernorm = input_norm - layer.self_attn = PartialOutputAttention((partial0, partial1)) - layer.post_attention_layernorm = post_norm - captures = {} - layer._diagnostic_capture_component = lambda name, tensor: captures.setdefault(name, tensor.detach().clone()) - - hidden_states = torch.ones((1, 2, 4), dtype=torch.bfloat16) - output, post_attention_residual = layer._pre_mlp_forward( - hidden_states, - position_embeddings=(hidden_states, hidden_states), - ) - - expected_sum_then_residual = hidden_states + (partial0 + partial1) - expected_residual_then_partials = (hidden_states + partial0) + partial1 - assert "post_attention_o_proj_partial_sum_sum_then_residual" in captures - assert "post_attention_partial_residual_sum_then_residual" in captures - assert "post_attention_partial_residual_residual_then_partials" in captures - assert "post_attention_partial_residual_fp32_sum_then_residual" in captures - torch.testing.assert_close( - captures["post_attention_partial_residual_sum_then_residual"], expected_sum_then_residual - ) - torch.testing.assert_close( - captures["post_attention_partial_residual_residual_then_partials"], - expected_residual_then_partials, - ) - assert expected_sum_then_residual.flatten()[0].item() != expected_residual_then_partials.flatten()[0].item() - - assert len(post_norm.calls) == 1 - torch.testing.assert_close(post_norm.calls[0]["hidden_states"], partial0 + partial1) - torch.testing.assert_close(output, partial0 + partial1 + 10.0) - torch.testing.assert_close(post_attention_residual, hidden_states + partial0 + partial1) - - -def test_qwen3_moe_o_proj_partial_residual_modes(monkeypatch): - hidden_states = torch.tensor([1.0], dtype=torch.bfloat16) - residual = torch.tensor([1.0], dtype=torch.bfloat16) - partial0 = torch.tensor([0.002], dtype=torch.bfloat16) - partial1 = torch.tensor([0.002], dtype=torch.bfloat16) - - monkeypatch.setenv("XORL_QWEN3_MOE_POST_ATTENTION_O_PROJ_PARTIAL_RESIDUAL_MODE", "sum_then_residual") - _, sum_then_residual = _materialize_o_proj_partial_residual(hidden_states, residual, (partial0, partial1)) - expected_sum_then_residual = residual + (partial0 + partial1) - torch.testing.assert_close(sum_then_residual, expected_sum_then_residual) - - monkeypatch.setenv( - "XORL_QWEN3_MOE_POST_ATTENTION_O_PROJ_PARTIAL_RESIDUAL_MODE", - "residual_then_partials", - ) - _, residual_then_partials = _materialize_o_proj_partial_residual(hidden_states, residual, (partial0, partial1)) - expected_residual_then_partials = (residual + partial0) + partial1 - torch.testing.assert_close(residual_then_partials, expected_residual_then_partials) - assert sum_then_residual.item() != residual_then_partials.item() diff --git a/tests/models/test_regather_routing_sqrtsoftplus.py b/tests/models/test_regather_routing_sqrtsoftplus.py deleted file mode 100644 index a7d59dd2..00000000 --- a/tests/models/test_regather_routing_sqrtsoftplus.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Regression test for ``MoEBlock._regather_routing`` sqrtsoftplus dispatch. - -DSv4 / GLM-style MoE uses ``scoring_func="sqrtsoftplus"`` (with optional -``noaux_tc`` selection-time bias and a ``routed_scaling_factor``). The -routing-replay regather path must recover the same per-token routing -weights as the eager router, otherwise gradient-checkpoint recompute -silently diverges from the forward pass. -""" - -import pytest -import torch -import torch.nn.functional as F - -from xorl.models.layers.moe.moe_block import MoEBlock -from xorl.models.layers.moe.router import TopKRouter - - -pytestmark = pytest.mark.cpu - - -def _make_block(scoring_func: str, routed_scaling_factor=None, num_experts=8, top_k=2): - """Build a minimal CPU MoEBlock with a chosen TopKRouter scoring func.""" - block = MoEBlock( - hidden_size=16, - num_experts=num_experts, - top_k=top_k, - intermediate_size=16, - moe_implementation="eager", - ) - block.router = TopKRouter( - num_experts=num_experts, - top_k=top_k, - norm_topk_prob=True, - scoring_func=scoring_func, - topk_method="noaux_tc" if scoring_func == "sqrtsoftplus" else None, - routed_scaling_factor=routed_scaling_factor, - ) - return block - - -def test_sqrtsoftplus_regather_matches_eager_router(): - """Regather from cached top-k must equal the V4 router's eager output.""" - torch.manual_seed(0) - num_tokens, num_experts, top_k = 6, 8, 2 - block = _make_block("sqrtsoftplus", num_experts=num_experts, top_k=top_k) - - router_logits = torch.randn(num_tokens, num_experts) - bias = torch.randn(num_experts) * 0.1 - - eager_w, eager_idx = block.router(router_logits, input_dtype=torch.float32, expert_bias=bias) - - regather_idx, regather_w = block._regather_routing(router_logits, eager_idx, input_dtype=torch.float32) - - assert torch.equal(regather_idx, eager_idx) - torch.testing.assert_close(regather_w, eager_w, rtol=1e-5, atol=1e-6) - - -def test_sqrtsoftplus_regather_applies_routed_scaling_factor(): - """Regather must apply ``routed_scaling_factor`` (DSv4 sets it to 1.0+).""" - torch.manual_seed(1) - num_tokens, num_experts, top_k = 4, 8, 2 - scaling = 2.5 - block = _make_block("sqrtsoftplus", routed_scaling_factor=scaling, num_experts=num_experts, top_k=top_k) - - router_logits = torch.randn(num_tokens, num_experts) - bias = torch.zeros(num_experts) - eager_w, eager_idx = block.router(router_logits, input_dtype=torch.float32, expert_bias=bias) - _, regather_w = block._regather_routing(router_logits, eager_idx, input_dtype=torch.float32) - - torch.testing.assert_close(regather_w, eager_w, rtol=1e-5, atol=1e-6) - # Cross-check against an unscaled regather to make sure scaling is applied. - block.router.routed_scaling_factor = None - _, unscaled_w = block._regather_routing(router_logits, eager_idx, input_dtype=torch.float32) - torch.testing.assert_close(regather_w, unscaled_w * scaling, rtol=1e-5, atol=1e-6) - - -def test_sqrtsoftplus_regather_preserves_input_dtype(): - """Regathered weights must come back in the requested input dtype.""" - torch.manual_seed(2) - block = _make_block("sqrtsoftplus") - router_logits = torch.randn(4, 8, dtype=torch.float32) - cached = torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7]], dtype=torch.long) - _, w = block._regather_routing(router_logits, cached, input_dtype=torch.bfloat16) - assert w.dtype == torch.bfloat16 - - -def test_softmax_regather_unchanged(): - """The softmax path must keep its existing semantics — we only added a - sqrtsoftplus branch above it.""" - torch.manual_seed(3) - num_tokens, num_experts, top_k = 5, 8, 2 - block = _make_block("softmax", num_experts=num_experts, top_k=top_k) - - router_logits = torch.randn(num_tokens, num_experts) - eager_w, eager_idx = block.router(router_logits, input_dtype=torch.float32) - _, regather_w = block._regather_routing(router_logits, eager_idx, input_dtype=torch.float32) - - # Hand-roll the expected: softmax → gather → renorm. - probs = F.softmax(router_logits, dim=1, dtype=torch.float32) - expected = torch.gather(probs, 1, eager_idx) - expected = expected / expected.sum(dim=-1, keepdim=True) - torch.testing.assert_close(regather_w, expected, rtol=1e-5, atol=1e-6) - torch.testing.assert_close(regather_w, eager_w, rtol=1e-5, atol=1e-6) diff --git a/tests/models/test_rmsnorm_family_contract.py b/tests/models/test_rmsnorm_family_contract.py index f8c207e6..e4a1e2a4 100644 --- a/tests/models/test_rmsnorm_family_contract.py +++ b/tests/models/test_rmsnorm_family_contract.py @@ -32,6 +32,8 @@ ) from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config from xorl.models.transformers.qwen3.modeling_qwen3 import Qwen3DecoderLayer +from xorl.models.transformers.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig +from xorl.models.transformers.qwen3_moe.modeling_qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeModel from xorl.ops.batch_invariant_ops import ( bi_fused_add_rms_norm, bi_rms_norm, @@ -53,9 +55,15 @@ @pytest.fixture(autouse=True) -def _pin_legacy_family_trees(monkeypatch): - """Keep the legacy-family contract independent of test collection order.""" - monkeypatch.setenv("XORL_FAMILIES_V2", "0") +def _pin_qualified_v1_family(): + """Keep the v1-family contract independent of test collection order.""" + from xorl.ops.bi_families_v2 import _select_nonexact_families, _select_qwen35_families_v1 + + _select_qwen35_families_v1() + try: + yield + finally: + _select_nonexact_families() def _make(shape, seed, device="cuda", dtype=torch.bfloat16): @@ -66,21 +74,28 @@ def _make(shape, seed, device="cuda", dtype=torch.bfloat16): # --------------------------------------------------------------------------- # # Structural guards (CPU). # --------------------------------------------------------------------------- # -def test_unknown_family_rejected_at_construction_and_call(): +def test_rmsnorm_family_structure_policy(): with pytest.raises(ValueError, match="Unknown RMSNorm family"): RMSNorm(4, eps=EPS, family="serving_qk") norm = RMSNorm(4, eps=EPS) with pytest.raises(ValueError, match="Unknown RMSNorm family"): norm(torch.ones(1, 4), family="family-3") + _assert_no_residual_family_rejects_residual_stream() + _assert_no_residual_family_rejects_residual_tree_force() + _assert_fused_add_through_no_residual_family_raises() + _assert_funnel_rejects_unknown_family() + _assert_zero_centered_rejects_residual_tree_family() + _assert_qwen_rmsnorm_site_declaration_policy() -def test_no_residual_family_rejects_residual_stream(): + +def _assert_no_residual_family_rejects_residual_stream(): norm = RMSNorm(4, eps=EPS, family=RMS_NORM_FAMILY_NO_RESIDUAL) with pytest.raises(ValueError, match="residual stream"): norm(torch.ones(1, 4), residual=torch.ones(1, 4), prenorm=True) -def test_no_residual_family_rejects_residual_tree_force(): +def _assert_no_residual_family_rejects_residual_tree_force(): norm = RMSNorm(4, eps=EPS, family=RMS_NORM_FAMILY_NO_RESIDUAL) with pytest.raises(ValueError, match="family flip"): norm(torch.ones(1, 4), force_sglang_residual=True) @@ -88,19 +103,19 @@ def test_no_residual_family_rejects_residual_tree_force(): norm(torch.ones(1, 4), force_sglang_residual_kernel=True) -def test_fused_add_through_no_residual_family_raises(): +def _assert_fused_add_through_no_residual_family_raises(): x = torch.ones(2, 4) with pytest.raises(ValueError, match="no fused-add kernel"): bi_fused_add_rms_norm(x, x, torch.ones(4), EPS, family=RMS_NORM_FAMILY_NO_RESIDUAL) -def test_funnel_rejects_unknown_family(): +def _assert_funnel_rejects_unknown_family(): x = torch.ones(2, 4) with pytest.raises(ValueError, match="Unknown RMSNorm family"): bi_rms_norm(x, torch.ones(4), EPS, family="serving") -def test_zero_centered_rejects_residual_tree_family(): +def _assert_zero_centered_rejects_residual_tree_family(): x = torch.ones(2, 4) with pytest.raises(ValueError, match="only exists in the 'serving_no_residual' family"): bi_rms_norm(x, torch.ones(4), EPS, family=RMS_NORM_FAMILY_RESIDUAL_TREE, zero_centered=True) @@ -109,7 +124,7 @@ def test_zero_centered_rejects_residual_tree_family(): # --------------------------------------------------------------------------- # # Modeling declarations: the K3 parity models must declare their families. # --------------------------------------------------------------------------- # -def test_dense_qwen3_declares_site_families(): +def _assert_qwen_rmsnorm_site_declaration_policy(): cfg = Qwen3Config( hidden_size=64, intermediate_size=128, @@ -127,8 +142,11 @@ def test_dense_qwen3_declares_site_families(): assert layer0.self_attn.q_norm.family == RMS_NORM_FAMILY_NO_RESIDUAL assert layer0.self_attn.k_norm.family == RMS_NORM_FAMILY_NO_RESIDUAL + _assert_shared_attention_qk_norms_declare_no_residual_family() + _assert_qwen3_moe_norm_family_declaration_policy() + -def test_shared_attention_qk_norms_declare_no_residual_family(): +def _assert_shared_attention_qk_norms_declare_no_residual_family(): cfg = Qwen3Config( hidden_size=64, intermediate_size=128, @@ -143,12 +161,69 @@ def test_shared_attention_qk_norms_declare_no_residual_family(): assert attn.k_norm.family == RMS_NORM_FAMILY_NO_RESIDUAL +def _assert_qwen3_moe_norm_family_declaration_policy(): + class CaptureNorm(torch.nn.Module): + def __init__(self): + super().__init__() + self.family_values = [] + + def forward(self, hidden_states, *, force_sglang_residual=False, family=None): + assert force_sglang_residual is False + self.family_values.append(family) + return hidden_states + + class IdentityAttention(torch.nn.Module): + def forward(self, hidden_states, **kwargs): + return hidden_states, None + + class IdentityPostAttentionNorm(torch.nn.Module): + def forward(self, hidden_states, residual=None, prenorm=False, **kwargs): + return hidden_states, residual + + config = Qwen3MoeConfig( + hidden_size=4, + intermediate_size=8, + num_attention_heads=2, + num_key_value_heads=2, + num_hidden_layers=2, + num_experts=0, + use_qk_norm=False, + _attn_implementation="eager", + pad_token_id=0, + ) + for layer_idx, expected_family in ( + (0, RMS_NORM_FAMILY_NO_RESIDUAL), + (1, RMS_NORM_FAMILY_RESIDUAL_TREE), + ): + layer = Qwen3MoeDecoderLayer(config, layer_idx=layer_idx) + input_norm = CaptureNorm() + layer.input_layernorm = input_norm + layer.self_attn = IdentityAttention() + layer.post_attention_layernorm = IdentityPostAttentionNorm() + hidden_states = torch.ones(1, 2, 4) + + layer._pre_mlp_forward(hidden_states, position_embeddings=(hidden_states, hidden_states)) + + assert input_norm.family_values == [expected_family] + + model = Qwen3MoeModel(config) + assert model.norm.family == RMS_NORM_FAMILY_RESIDUAL_TREE + + class StubLayer(torch.nn.Module): + def forward(self, hidden_states, *args, **kwargs): + return (hidden_states,) + + model.layers = torch.nn.ModuleList([StubLayer()]) + final_norm = CaptureNorm() + model.norm = final_norm + model(input_ids=torch.tensor([[0, 1]])) + assert final_norm.family_values == [None] + + # --------------------------------------------------------------------------- # # Loud tripwire for undeclared parity-lane calls. # --------------------------------------------------------------------------- # -@requires_cuda -@pytest.mark.gpu -def test_undeclared_family_warns_in_parity_lane(): +def _assert_family_declaration_tripwire_policy(monkeypatch): normalization._WARNED_UNDECLARED_FAMILY.clear() norm = RMSNorm(64, eps=EPS, mode="sglang_fused").to("cuda") x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) @@ -158,10 +233,11 @@ def test_undeclared_family_warns_in_parity_lane(): # Warned once per (mode, call-shape); a second call is silent. normalization._WARNED_UNDECLARED_FAMILY.clear() + _assert_declared_and_legacy_calls_do_not_warn() + _assert_undeclared_family_raises_when_required(monkeypatch) -@requires_cuda -@pytest.mark.gpu -def test_undeclared_family_raises_when_required(monkeypatch): + +def _assert_undeclared_family_raises_when_required(monkeypatch): monkeypatch.setenv("XORL_RMSNORM_REQUIRE_FAMILY", "1") norm = RMSNorm(64, eps=EPS, mode="sglang_fused").to("cuda") x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) @@ -173,9 +249,7 @@ def test_undeclared_family_raises_when_required(monkeypatch): norm(x, family=RMS_NORM_FAMILY_RESIDUAL_TREE) -@requires_cuda -@pytest.mark.gpu -def test_declared_and_legacy_calls_do_not_warn(): +def _assert_declared_and_legacy_calls_do_not_warn(): norm = RMSNorm(64, eps=EPS, mode="sglang_fused", family=RMS_NORM_FAMILY_NO_RESIDUAL).to("cuda") tree = RMSNorm(64, eps=EPS, mode="sglang_fused").to("cuda") x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16) @@ -190,10 +264,7 @@ def test_declared_and_legacy_calls_do_not_warn(): # family-declared module calls match the legacy call shapes, and the two # families genuinely differ on the seed shape. # --------------------------------------------------------------------------- # -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("shape", [SEED_SHAPE, HIDDEN_SHAPE]) -def test_funnel_matches_legacy_kernels_bitwise(shape): +def _assert_funnel_matches_legacy_kernels_bitwise(shape): x = _make(shape, 0) w = _make((shape[-1],), 300) with torch.no_grad(): @@ -214,7 +285,18 @@ def test_funnel_matches_legacy_kernels_bitwise(shape): @requires_cuda @pytest.mark.gpu -def test_zero_centered_twin_is_family1_with_fold(): +def test_rmsnorm_family_funnel_policy(monkeypatch): + with monkeypatch.context() as case_patch: + _assert_family_declaration_tripwire_policy(case_patch) + for shape in (SEED_SHAPE, HIDDEN_SHAPE): + _assert_funnel_matches_legacy_kernels_bitwise(shape) + + _assert_zero_centered_twin_is_family1_with_fold() + _assert_families_differ_on_seed_shape() + _assert_family_module_dispatch_policy() + + +def _assert_zero_centered_twin_is_family1_with_fold(): """The Qwen3.5 zero-centered twin (#468) registered in the family API: the funnel's zero_centered form, the differentiable wrapper, the raw family-1 kernel on the folded operands, and the interpose lane all agree bitwise.""" @@ -231,9 +313,7 @@ def test_zero_centered_twin_is_family1_with_fold(): assert torch.equal(funnel, interpose), "zero-centered twin diverged from the interpose lane" -@requires_cuda -@pytest.mark.gpu -def test_trunk_contract_lane_dispatches_family1_and_warns_undeclared(): +def _assert_family_module_dispatch_policy(): """#467 composition: under the scoped trunk-contract lane (global interpose off), no-residual sglang_fused dispatch is the family-1 kernel with real gradients -- for declared serving_no_residual sites without a warning, and @@ -261,10 +341,10 @@ def test_trunk_contract_lane_dispatches_family1_and_warns_undeclared(): assert torch.equal(out_declared, ref), "trunk-lane declared qk-norm left family-1" assert torch.equal(out_undeclared, ref) + _assert_family_declared_module_calls_match_legacy_bitwise_policy() -@requires_cuda -@pytest.mark.gpu -def test_families_differ_on_seed_shape(): + +def _assert_families_differ_on_seed_shape(): """The tripwire vitality check: if the two families ever collapse to the same bits on the seed shape, every bitwise family gate is vacuous and this contract should be re-evaluated.""" @@ -279,13 +359,9 @@ def test_families_differ_on_seed_shape(): assert n_diff < x.numel() * 1e-3 -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("mode", ["native", "sglang", "sglang_fused"]) -@pytest.mark.parametrize("bi", [False, True]) -def test_family_declared_module_calls_match_legacy_bitwise(mode, bi): +def _assert_family_declared_module_calls_match_legacy_bitwise(mode): """Family declarations replace the legacy force_sglang_residual call-site - exprs without changing a single bit, in every (mode, batch-invariant) lane. + expressions without changing a single bit in each mode branch. The legacy trunk marker is set explicitly here; exact Qwen production no longer mutates it as a model-build side effect, so test order must not @@ -303,7 +379,7 @@ def make_norm(**kwargs): set_trunk_linear_contract(True) try: - with set_batch_invariant_mode(bi), torch.no_grad(), warnings.catch_warnings(): + with set_batch_invariant_mode(False), torch.no_grad(), warnings.catch_warnings(): warnings.simplefilter("ignore") # qk-norm / layer-0 input: declared no-residual == legacy bare call. assert torch.equal(make_norm(family=RMS_NORM_FAMILY_NO_RESIDUAL)(x), make_norm()(x)) @@ -321,3 +397,8 @@ def make_norm(**kwargs): assert torch.equal(rout_new, rout_old) finally: set_trunk_linear_contract(False) + + +def _assert_family_declared_module_calls_match_legacy_bitwise_policy(): + for mode in ("native", "sglang", "sglang_fused"): + _assert_family_declared_module_calls_match_legacy_bitwise(mode) diff --git a/tests/models/test_rmsnorm_family_cross_engine.py b/tests/models/test_rmsnorm_family_cross_engine.py index 28c31253..f3953b7d 100644 --- a/tests/models/test_rmsnorm_family_cross_engine.py +++ b/tests/models/test_rmsnorm_family_cross_engine.py @@ -24,14 +24,11 @@ RMS_NORM_FAMILY_NO_RESIDUAL, RMS_NORM_FAMILY_RESIDUAL_TREE, RMSNorm, - fast_batch_invariant_rms_norm, fast_zero_centered_batch_invariant_rms_norm, ) from xorl.ops.batch_invariant_ops import ( # noqa: E402 - bi_fused_add_rms_norm, bi_rms_norm, set_batch_invariant_mode, - set_trunk_linear_contract, ) from xorl.ops.bi_families_v2 import rms_norm_v2 as xorl_rms_norm_v2 # noqa: E402 @@ -57,145 +54,118 @@ def _xorl_module(hidden, family, weight): return norm -@pytest.mark.parametrize("shape", SHAPES) -def test_qk_norm_site_class_bitwise(shape): +def test_rmsnorm_site_class_cross_engine_bitwise_policy(): """qk-norm / layer-0 input layernorm: xorl's parity-lane dispatch (native -> aten::rms_norm interpose) must bit-match SGLang's residual-is-None dispatch (family-1 ``rms_norm_batch_invariant``).""" - x = _make(shape, 0) - w = _make((shape[-1],), 300) - norm = _xorl_module(shape[-1], RMS_NORM_FAMILY_NO_RESIDUAL, w) - with set_batch_invariant_mode(True), torch.no_grad(): - xorl_out = norm(x) - serving_out = sgl_bio.rms_norm_batch_invariant(x, w, EPS) - assert torch.equal(xorl_out, serving_out), "qk-norm site-class diverged from serving family-1" + for shape in SHAPES: + x = _make(shape, 0) + w = _make((shape[-1],), 300) + norm = _xorl_module(shape[-1], RMS_NORM_FAMILY_NO_RESIDUAL, w) + with set_batch_invariant_mode(True), torch.no_grad(): + xorl_out = norm(x) + serving_out = sgl_bio.rms_norm_batch_invariant(x, w, EPS) + assert torch.equal(xorl_out, serving_out), f"qk-norm {shape} diverged from serving family-1" + + serving_funnel = sgl_bio.bi_rms_norm(x, w, EPS, family=RMS_NORM_FAMILY_NO_RESIDUAL) + assert torch.equal(xorl_out, serving_funnel), f"qk-norm funnel {shape} diverged" + + if shape == SHAPES[0]: + family2 = sgl_bio.rms_norm_residual_tree_batch_invariant(x, w, EPS) + assert not torch.equal(serving_out, family2), "serving families agree on the seed shape; gate is vacuous" + _assert_presummed_residual_tree_site_class_bitwise() + _assert_post_attention_residual_site_class_bitwise() + _assert_zero_centered_family1_twin_bitwise() + _assert_zero_centered_families_v2_candidate_bitwise() -@pytest.mark.parametrize("shape", SHAPES) -def test_presummed_residual_tree_site_class_bitwise(shape): + +def _assert_presummed_residual_tree_site_class_bitwise(): """Input layernorm at layer>0 / final norm: xorl normalizes the pre-summed single tensor through the residual tree; SGLang fuses the add. On the same summed value both must produce identical bits (gate via a zero residual and via SGLang's single-tensor residual-tree kernel).""" - x = _make(shape, 0) - w = _make((shape[-1],), 300) - norm = _xorl_module(shape[-1], RMS_NORM_FAMILY_RESIDUAL_TREE, w) - with set_batch_invariant_mode(True), torch.no_grad(): - xorl_out = norm(x) - serving_single = sgl_bio.rms_norm_residual_tree_batch_invariant(x, w, EPS) - serving_fused, serving_residual = sgl_bio.fused_add_rms_norm_batch_invariant(x, torch.zeros_like(x), w, EPS) - assert torch.equal(serving_residual, x) - assert torch.equal(xorl_out, serving_single), "pre-summed site-class diverged from serving residual tree" - assert torch.equal(xorl_out, serving_fused), "pre-summed site-class diverged from serving fused-add tree" - - -@pytest.mark.parametrize("shape", SHAPES) -def test_post_attention_residual_site_class_bitwise(shape): + for shape in SHAPES: + x = _make(shape, 0) + w = _make((shape[-1],), 300) + norm = _xorl_module(shape[-1], RMS_NORM_FAMILY_RESIDUAL_TREE, w) + with set_batch_invariant_mode(True), torch.no_grad(): + xorl_out = norm(x) + serving_single = sgl_bio.rms_norm_residual_tree_batch_invariant(x, w, EPS) + serving_fused, serving_residual = sgl_bio.fused_add_rms_norm_batch_invariant(x, torch.zeros_like(x), w, EPS) + serving_funnel = sgl_bio.bi_rms_norm(x, w, EPS, family=RMS_NORM_FAMILY_RESIDUAL_TREE) + assert torch.equal(serving_residual, x) + assert torch.equal(xorl_out, serving_single), f"pre-summed {shape} diverged from residual tree" + assert torch.equal(xorl_out, serving_fused), f"pre-summed {shape} diverged from fused-add tree" + assert torch.equal(xorl_out, serving_funnel), f"pre-summed funnel {shape} diverged" + + +def _assert_post_attention_residual_site_class_bitwise(): """Post-attention layernorm: xorl's fused residual dispatch must bit-match SGLang's fused residual dispatch, on both the normed output and the carried residual stream.""" - x = _make(shape, 0) - r = _make(shape, 1) - w = _make((shape[-1],), 300) - norm = _xorl_module(shape[-1], RMS_NORM_FAMILY_RESIDUAL_TREE, w) - with set_batch_invariant_mode(True), torch.no_grad(): - xorl_out, xorl_residual = norm(x, residual=r, prenorm=True) - serving_out, serving_residual = sgl_bio.fused_add_rms_norm_batch_invariant(x, r, w, EPS) - assert torch.equal(xorl_residual, serving_residual), "residual carry diverged from serving" - assert torch.equal(xorl_out, serving_out), "post-attention site-class diverged from serving" - - -@pytest.mark.parametrize("family", [RMS_NORM_FAMILY_NO_RESIDUAL, RMS_NORM_FAMILY_RESIDUAL_TREE]) -@pytest.mark.parametrize("shape", SHAPES) -def test_family_funnels_agree_cross_engine(family, shape): - """The vendored family funnels themselves must agree bitwise per family.""" - x = _make(shape, 7) - w = _make((shape[-1],), 301) - with torch.no_grad(): - assert torch.equal( - bi_rms_norm(x, w, EPS, family=family), - sgl_bio.bi_rms_norm(x, w, EPS, family=family), + for shape in SHAPES: + x = _make(shape, 0) + r = _make(shape, 1) + w = _make((shape[-1],), 300) + norm = _xorl_module(shape[-1], RMS_NORM_FAMILY_RESIDUAL_TREE, w) + with set_batch_invariant_mode(True), torch.no_grad(): + xorl_out, xorl_residual = norm(x, residual=r, prenorm=True) + serving_out, serving_residual = sgl_bio.fused_add_rms_norm_batch_invariant(x, r, w, EPS) + funnel_out, funnel_residual = sgl_bio.bi_fused_add_rms_norm( + x, + r, + w, + EPS, + family=RMS_NORM_FAMILY_RESIDUAL_TREE, ) + assert torch.equal(xorl_residual, serving_residual), f"residual carry {shape} diverged from serving" + assert torch.equal(xorl_out, serving_out), f"post-attention {shape} diverged from serving" + assert torch.equal(xorl_residual, funnel_residual), f"residual funnel carry {shape} diverged" + assert torch.equal(xorl_out, funnel_out), f"post-attention funnel {shape} diverged" -def test_fused_add_funnels_agree_cross_engine(): - x = _make((4096, 128), 7) - r = _make((4096, 128), 8) - w = _make((128,), 301) - with torch.no_grad(): - x_out, x_rout = bi_fused_add_rms_norm(x, r, w, EPS, family=RMS_NORM_FAMILY_RESIDUAL_TREE) - s_out, s_rout = sgl_bio.bi_fused_add_rms_norm(x, r, w, EPS, family=RMS_NORM_FAMILY_RESIDUAL_TREE) - assert torch.equal(x_out, s_out) - assert torch.equal(x_rout, s_rout) - - -def test_families_differ_cross_engine_gate_has_teeth(): - """If the two serving families ever agree bitwise on the seed shape, the - site-class gates above cannot catch a family flip; fail loudly instead.""" - x = _make((4096, 128), 0) - w = _make((128,), 300) - f1 = sgl_bio.rms_norm_batch_invariant(x, w, EPS) - f2 = sgl_bio.rms_norm_residual_tree_batch_invariant(x, w, EPS) - assert not torch.equal(f1, f2), "serving families agree on the seed shape; gates are vacuous" - - -@pytest.mark.parametrize("shape", SHAPES) -def test_trunk_contract_lane_qk_norm_site_class_bitwise(shape): - """The scoped trunk-contract lane (XORL_BI_TRUNK_LINEAR, no global interpose) - must dispatch the same family-1 kernel as serving's residual-is-None path -- - the inverse case (a qk-norm family flip) is what this - gate would catch.""" - x = _make(shape, 0) - w = _make((shape[-1],), 300) - norm = _xorl_module(shape[-1], RMS_NORM_FAMILY_NO_RESIDUAL, w) - set_trunk_linear_contract(True) - try: - with torch.no_grad(): - xorl_out = norm(x) - finally: - set_trunk_linear_contract(False) - serving_out = sgl_bio.rms_norm_batch_invariant(x, w, EPS) - assert torch.equal(xorl_out, serving_out), "trunk-lane qk-norm diverged from serving family-1" - # The differentiable family-1 wrapper is the same dispatch. - with torch.no_grad(): - assert torch.equal(fast_batch_invariant_rms_norm(x, w, EPS), serving_out) - - -@pytest.mark.parametrize("shape", SHAPES) -def test_zero_centered_family1_twin_bitwise(shape): +def _assert_zero_centered_family1_twin_bitwise(): """The Qwen3.5 zero-centered (Gemma-style) twin is family-1 with a fp32 ``1 + weight`` fold; both engines' funnels and xorl's differentiable wrapper must agree bitwise.""" - x = _make(shape, 5) - w = _make((shape[-1],), 302) - with torch.no_grad(): - xorl_fn = fast_zero_centered_batch_invariant_rms_norm(x, w, EPS) - xorl_funnel = bi_rms_norm(x, w, EPS, family=RMS_NORM_FAMILY_NO_RESIDUAL, zero_centered=True) - serving_funnel = sgl_bio.bi_rms_norm(x, w, EPS, family=RMS_NORM_FAMILY_NO_RESIDUAL, zero_centered=True) - assert torch.equal(xorl_fn, xorl_funnel), "zero-centered wrapper diverged from the xorl funnel" - assert torch.equal(xorl_fn, serving_funnel), "zero-centered twin diverged cross-engine" - - -@pytest.mark.parametrize("shape", [(64, 128), (64, 3840)]) -@pytest.mark.parametrize("with_residual", [False, True]) -def test_zero_centered_families_v2_candidate_bitwise(shape, with_residual): + for shape in SHAPES: + x = _make(shape, 5) + w = _make((shape[-1],), 302) + with torch.no_grad(): + xorl_fn = fast_zero_centered_batch_invariant_rms_norm(x, w, EPS) + xorl_funnel = bi_rms_norm(x, w, EPS, family=RMS_NORM_FAMILY_NO_RESIDUAL, zero_centered=True) + serving_funnel = sgl_bio.bi_rms_norm( + x, + w, + EPS, + family=RMS_NORM_FAMILY_NO_RESIDUAL, + zero_centered=True, + ) + assert torch.equal(xorl_fn, xorl_funnel), f"zero-centered wrapper {shape} diverged from xorl" + assert torch.equal(xorl_fn, serving_funnel), f"zero-centered twin {shape} diverged cross-engine" + + +def _assert_zero_centered_families_v2_candidate_bitwise(): """The opt-in Qwen families-v2 candidate is one shared arithmetic tree. Cover both the q/k-style no-residual form and the decoder-layer fused residual form. The small shape selects the fused realization; the hidden size exercises the production-width epilogue. """ - x = _make(shape, 11) - w = _make((shape[-1],), 312) - residual = _make(shape, 12) if with_residual else None - with torch.no_grad(): - xorl_result = xorl_rms_norm_v2(x, w, EPS, residual=residual, zero_centered=True) - serving_result = sgl_bio.rms_norm_v2(x, w, EPS, residual=residual, zero_centered=True) - - if with_residual: - xorl_out, xorl_residual = xorl_result - serving_out, serving_residual = serving_result - assert torch.equal(xorl_residual, serving_residual) - assert torch.equal(xorl_out, serving_out) - else: - assert torch.equal(xorl_result, serving_result) + for shape in ((64, 128), (64, 3840)): + for with_residual in (False, True): + x = _make(shape, 11) + w = _make((shape[-1],), 312) + residual = _make(shape, 12) if with_residual else None + with torch.no_grad(): + xorl_result = xorl_rms_norm_v2(x, w, EPS, residual=residual, zero_centered=True) + serving_result = sgl_bio.rms_norm_v2(x, w, EPS, residual=residual, zero_centered=True) + + if with_residual: + xorl_out, xorl_residual = xorl_result + serving_out, serving_residual = serving_result + assert torch.equal(xorl_residual, serving_residual), f"residual {shape} diverged" + assert torch.equal(xorl_out, serving_out), f"output {shape} diverged" + else: + assert torch.equal(xorl_result, serving_result), f"no-residual {shape} diverged" diff --git a/tests/models/test_rmsnorm_sglang_fused.py b/tests/models/test_rmsnorm_sglang_fused.py index a8e37e98..31ae9108 100644 --- a/tests/models/test_rmsnorm_sglang_fused.py +++ b/tests/models/test_rmsnorm_sglang_fused.py @@ -8,8 +8,6 @@ of the eager reference. The CPU tests exercise the eager fallback. """ -import os - import pytest import torch @@ -42,7 +40,7 @@ # --------------------------------------------------------------------------- # # CPU fallback (no Triton): sglang_fused must equal the eager sglang path. # --------------------------------------------------------------------------- # -def test_sglang_fused_cpu_residual_matches_eager(): +def _assert_sglang_fused_cpu_residual_matches_eager(): set_rmsnorm_mode("sglang_fused") try: norm = RMSNorm(4, eps=EPS) @@ -61,13 +59,19 @@ def test_sglang_fused_cpu_residual_matches_eager(): set_rmsnorm_mode("native") -# This suite pins the v1 family kernels through the fast_* dispatchers; with -# families-v2 default-on those dispatchers route to the v2 tree, so pin the -# kill switch (env is read per call). v2 has its own suite (test_bi_families_v2.py). -os.environ["XORL_FAMILIES_V2"] = "0" +@pytest.fixture(autouse=True) +def _pin_qualified_v1_family(): + """This suite owns the qualified v1 fast-dispatch path.""" + from xorl.ops.bi_families_v2 import _select_nonexact_families, _select_qwen35_families_v1 + _select_qwen35_families_v1() + try: + yield + finally: + _select_nonexact_families() -def test_sglang_fused_cpu_force_no_residual_matches_eager(): + +def _assert_sglang_fused_cpu_force_no_residual_matches_eager(): set_rmsnorm_mode("sglang_fused") try: norm = RMSNorm(4, eps=EPS) @@ -85,45 +89,39 @@ def test_sglang_fused_cpu_force_no_residual_matches_eager(): # --------------------------------------------------------------------------- # # GPU bit-exactness under batch-invariant mode (the K3 regime). # --------------------------------------------------------------------------- # -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -def test_fused_residual_bit_exact_vs_eager(dtype): - torch.manual_seed(0) - device = "cuda" - hidden = torch.randn(N_TOKENS, HIDDEN, device=device, dtype=dtype) - residual = torch.randn(N_TOKENS, HIDDEN, device=device, dtype=dtype) - weight = torch.randn(HIDDEN, device=device, dtype=dtype) +def _assert_fused_residual_bit_exact_vs_eager(): + for dtype in (torch.bfloat16, torch.float32): + torch.manual_seed(0) + device = "cuda" + hidden = torch.randn(N_TOKENS, HIDDEN, device=device, dtype=dtype) + residual = torch.randn(N_TOKENS, HIDDEN, device=device, dtype=dtype) + weight = torch.randn(HIDDEN, device=device, dtype=dtype) - with set_batch_invariant_mode(True): - expected_residual = hidden + residual - expected = sglang_residual_rms_norm(expected_residual, weight, EPS) - out, residual_out = fast_sglang_residual_rms_norm(hidden, residual, weight, EPS) + with set_batch_invariant_mode(True): + expected_residual = hidden + residual + expected = sglang_residual_rms_norm(expected_residual, weight, EPS) + out, residual_out = fast_sglang_residual_rms_norm(hidden, residual, weight, EPS) - # Residual carry must be bit-identical (it feeds the next layer's stream). - assert torch.equal(residual_out, expected_residual), "fused residual add diverged from torch add" - assert torch.equal(out, expected), "fused residual RMSNorm diverged from eager" + # Residual carry must be bit-identical (it feeds the next layer's stream). + assert torch.equal(residual_out, expected_residual), f"fused residual add diverged for {dtype}" + assert torch.equal(out, expected), f"fused residual RMSNorm diverged for {dtype}" -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -def test_fused_no_residual_bit_exact_vs_eager(dtype): - torch.manual_seed(1) - device = "cuda" - hidden = torch.randn(N_TOKENS, HIDDEN, device=device, dtype=dtype) - weight = torch.randn(HIDDEN, device=device, dtype=dtype) +def _assert_fused_no_residual_bit_exact_vs_eager(): + for dtype in (torch.bfloat16, torch.float32): + torch.manual_seed(1) + device = "cuda" + hidden = torch.randn(N_TOKENS, HIDDEN, device=device, dtype=dtype) + weight = torch.randn(HIDDEN, device=device, dtype=dtype) - with set_batch_invariant_mode(True): - expected = sglang_residual_rms_norm(hidden, weight, EPS) - out = fast_sglang_rms_norm(hidden, weight, EPS) + with set_batch_invariant_mode(True): + expected = sglang_residual_rms_norm(hidden, weight, EPS) + out = fast_sglang_rms_norm(hidden, weight, EPS) - assert torch.equal(out, expected), "fused no-residual RMSNorm diverged from eager" + assert torch.equal(out, expected), f"fused no-residual RMSNorm diverged for {dtype}" -@requires_cuda -@pytest.mark.gpu -def test_fused_residual_matches_3d_packed_shape(): +def _assert_fused_residual_matches_3d_packed_shape(): torch.manual_seed(2) device = "cuda" hidden = torch.randn(2, 96, HIDDEN, device=device, dtype=torch.bfloat16) @@ -144,9 +142,7 @@ def test_fused_residual_matches_3d_packed_shape(): # --------------------------------------------------------------------------- # # GPU: sglang_fused RMSNorm module == sglang module, bit-for-bit. # --------------------------------------------------------------------------- # -@requires_cuda -@pytest.mark.gpu -def test_module_sglang_fused_equals_sglang(): +def _assert_module_sglang_fused_equals_sglang(): torch.manual_seed(3) device = "cuda" hidden = torch.randn(N_TOKENS, HIDDEN, device=device, dtype=torch.bfloat16) @@ -189,9 +185,7 @@ def run(mode, **kwargs): # --------------------------------------------------------------------------- # # GPU: closed-form backward matches autograd of the eager reference. # --------------------------------------------------------------------------- # -@requires_cuda -@pytest.mark.gpu -def test_fused_residual_backward_matches_autograd(): +def _assert_fused_residual_backward_matches_autograd(): torch.manual_seed(4) device = "cuda" dtype = torch.bfloat16 @@ -222,9 +216,7 @@ def leaf(t): assert torch.allclose(w_f.grad.float(), w_ref.grad.float(), rtol=2e-2, atol=2e-2) -@requires_cuda -@pytest.mark.gpu -def test_dense_qwen3_layer_forward_bit_exact_sglang_vs_fused(): +def _assert_dense_qwen3_layer_forward_bit_exact_sglang_vs_fused(): """Full dense Qwen3 decoder-layer forward must be bit-identical between sglang and sglang_fused (the model-level K3-preservation gate). Exercises input_layernorm (force_sglang_residual path at layer>0) and @@ -267,9 +259,7 @@ def forward(self, hidden_states, **kwargs): assert torch.equal(out_sg, out_sf), "dense layer forward diverged between sglang and sglang_fused" -@requires_cuda -@pytest.mark.gpu -def test_single_tensor_force_call_bit_matches_serving_fused_residual_tree(): +def _assert_single_tensor_force_call_bit_matches_serving_fused_residual_tree(): """The layer>0 input-norm / final-norm call shape (pre-summed single tensor, force_sglang_residual=True) must be bit-identical to serving's fused residual tree (``fused_add_rms_norm_batch_invariant``). Guards the norm-seed trap where @@ -290,9 +280,7 @@ def test_single_tensor_force_call_bit_matches_serving_fused_residual_tree(): assert torch.equal(out_module, ref) -@requires_cuda -@pytest.mark.gpu -def test_fused_no_residual_backward_matches_autograd(): +def _assert_fused_no_residual_backward_matches_autograd(): torch.manual_seed(5) device = "cuda" dtype = torch.bfloat16 @@ -322,9 +310,7 @@ def leaf(t): # aten::rms_norm interpose kernel, NOT the fused sglang residual tree (the two # disagree at 1 ulp on rare bf16 boundary values). # --------------------------------------------------------------------------- # -@requires_cuda -@pytest.mark.gpu -def test_trunk_contract_no_residual_bit_matches_interpose_kernel(): +def _assert_trunk_contract_no_residual_bit_matches_interpose_kernel(): torch.manual_seed(21) # qk-norm call shape: [tokens, heads, head_dim] with a head_dim-sized weight. head_dim = 128 @@ -347,9 +333,7 @@ def test_trunk_contract_no_residual_bit_matches_interpose_kernel(): assert torch.equal(out, ref_interpose), "contract-lane qk-norm must equal the aten interpose lane bit-for-bit" -@requires_cuda -@pytest.mark.gpu -def test_no_residual_dispatch_unchanged_without_contract(): +def _assert_no_residual_dispatch_unchanged_without_contract(): torch.manual_seed(22) x = torch.randn(N_TOKENS, HIDDEN, device="cuda", dtype=torch.bfloat16) norm = RMSNorm(HIDDEN, eps=EPS, mode="sglang_fused").to(device="cuda", dtype=torch.bfloat16) @@ -359,9 +343,7 @@ def test_no_residual_dispatch_unchanged_without_contract(): assert torch.equal(out, ref) -@requires_cuda -@pytest.mark.gpu -def test_trunk_contract_no_residual_backward_matches_eager(): +def _assert_trunk_contract_no_residual_backward_matches_eager(): # Same convention as test_fused_no_residual_backward_matches_autograd: fp32 # weight leaf and the fp32-multiply eager reference (the kernel's semantics), # so the comparison is not dominated by bf16 grad-accumulation rounding. @@ -385,3 +367,28 @@ def leaf(t): assert torch.isfinite(h_c.grad.float()).all() and torch.isfinite(w_c.grad.float()).all() assert torch.allclose(h_c.grad.float(), h_ref.grad.float(), rtol=2e-2, atol=2e-2) assert torch.allclose(w_c.grad.float(), w_ref.grad.float(), rtol=2e-2, atol=2e-2) + + +def test_sglang_fused_rmsnorm_cpu_fallback_contract(): + _assert_sglang_fused_cpu_residual_matches_eager() + _assert_sglang_fused_cpu_force_no_residual_matches_eager() + + +@requires_cuda +@pytest.mark.gpu +def test_sglang_fused_rmsnorm_forward_backward_and_model_integration_contract(): + _assert_fused_residual_bit_exact_vs_eager() + _assert_fused_no_residual_bit_exact_vs_eager() + _assert_fused_residual_matches_3d_packed_shape() + _assert_module_sglang_fused_equals_sglang() + _assert_fused_residual_backward_matches_autograd() + _assert_fused_no_residual_backward_matches_autograd() + _assert_dense_qwen3_layer_forward_bit_exact_sglang_vs_fused() + _assert_single_tensor_force_call_bit_matches_serving_fused_residual_tree() + _assert_sglang_fused_rmsnorm_trunk_contract() + + +def _assert_sglang_fused_rmsnorm_trunk_contract(): + _assert_trunk_contract_no_residual_bit_matches_interpose_kernel() + _assert_no_residual_dispatch_unchanged_without_contract() + _assert_trunk_contract_no_residual_backward_matches_eager() diff --git a/tests/models/test_rmsnorm_sglang_jit.py b/tests/models/test_rmsnorm_sglang_jit.py deleted file mode 100644 index 878ec304..00000000 --- a/tests/models/test_rmsnorm_sglang_jit.py +++ /dev/null @@ -1,62 +0,0 @@ -import torch - -from xorl.models.layers.normalization import RMSNorm, native_rms_norm, set_rmsnorm_mode - - -def test_sglang_jit_rmsnorm_mode_cpu_falls_back_to_native_residual(): - set_rmsnorm_mode("sglang_jit") - try: - norm = RMSNorm(4, eps=1e-6) - with torch.no_grad(): - norm.weight.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0])) - - hidden_states = torch.tensor([[0.25, -0.5, 0.75, -1.0]], dtype=torch.float32) - residual = torch.tensor([[1.0, 0.5, -0.25, 0.125]], dtype=torch.float32) - - out, residual_out = norm(hidden_states, residual=residual, prenorm=True) - - expected_residual = hidden_states + residual - expected = native_rms_norm(expected_residual, norm.weight, norm.variance_epsilon) - assert torch.equal(residual_out, expected_residual) - assert torch.equal(out, expected) - finally: - set_rmsnorm_mode("native") - - -def test_sglang_jit_rmsnorm_mode_cpu_accepts_packed_shape(): - set_rmsnorm_mode("sglang_jit") - try: - norm = RMSNorm(4, eps=1e-6) - hidden_states = torch.randn(2, 3, 4, dtype=torch.float32) - residual = torch.randn(2, 3, 4, dtype=torch.float32) - - out, residual_out = norm(hidden_states, residual=residual, prenorm=True) - - expected_residual = hidden_states + residual - expected = native_rms_norm(expected_residual, norm.weight, norm.variance_epsilon) - assert out.shape == hidden_states.shape - assert residual_out.shape == hidden_states.shape - assert torch.equal(residual_out, expected_residual) - assert torch.allclose(out, expected) - finally: - set_rmsnorm_mode("native") - - -def test_sglang_kernel_rmsnorm_mode_cpu_falls_back_to_native_residual(): - set_rmsnorm_mode("sglang_kernel") - try: - norm = RMSNorm(4, eps=1e-6) - with torch.no_grad(): - norm.weight.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0])) - - hidden_states = torch.tensor([[0.25, -0.5, 0.75, -1.0]], dtype=torch.float32) - residual = torch.tensor([[1.0, 0.5, -0.25, 0.125]], dtype=torch.float32) - - out, residual_out = norm(hidden_states, residual=residual, prenorm=True) - - expected_residual = hidden_states + residual - expected = native_rms_norm(expected_residual, norm.weight, norm.variance_epsilon) - assert torch.equal(residual_out, expected_residual) - assert torch.equal(out, expected) - finally: - set_rmsnorm_mode("native") diff --git a/tests/models/test_rope_inv_freq_fp32.py b/tests/models/test_rope_inv_freq_fp32.py index 69d1b270..e68e7013 100644 --- a/tests/models/test_rope_inv_freq_fp32.py +++ b/tests/models/test_rope_inv_freq_fp32.py @@ -10,6 +10,7 @@ import torch from xorl.models.auto import build_foundation_model +from xorl.models.layers import rope as rope_module from xorl.models.layers.rope import ROPE_INIT_FUNCTIONS, RotaryEmbedding from xorl.models.transformers.qwen3.configuration_qwen3 import Qwen3Config @@ -83,51 +84,54 @@ def _build(rope_type: str, dtype: str, rope_native: bool = False): ) -def test_registry_is_fully_covered(): +def test_rope_registry_fp32_precision_and_contract_lane_policy(): assert set(ROPE_SCALINGS) == set(ROPE_INIT_FUNCTIONS) + for rope_type in sorted(ROPE_INIT_FUNCTIONS): + rotary = _build(rope_type, "bfloat16").model.rotary_emb + table = rotary._resolve_inv_freq(torch.device("cpu")) -@pytest.mark.parametrize("rope_type", sorted(ROPE_INIT_FUNCTIONS)) -def test_inv_freq_survives_model_wide_bf16_cast(rope_type: str): - rotary = _build(rope_type, "bfloat16").model.rotary_emb - table = rotary._resolve_inv_freq(torch.device("cpu")) + assert table.dtype == torch.float32, f"{rope_type}: rope reads a {table.dtype} frequency table" - assert table.dtype == torch.float32, f"{rope_type}: rope reads a {table.dtype} frequency table" + reference, _ = ROPE_INIT_FUNCTIONS[rope_type](_config(rope_type), "cpu") + assert torch.equal(table, reference.float()), f"{rope_type}: frequency table is not the fp32 CPU reference" - reference, _ = ROPE_INIT_FUNCTIONS[rope_type](_config(rope_type), "cpu") - assert torch.equal(table, reference.float()), f"{rope_type}: frequency table is not the fp32 CPU reference" + _assert_bf16_built_cos_sin_matches_fp32() + _assert_contract_lane_bits_unchanged() + _assert_native_default_cache_is_lazy_and_follows_execution_device() -@pytest.mark.parametrize("rope_type", sorted(ROPE_INIT_FUNCTIONS)) -def test_bf16_built_cos_sin_matches_fp32_built(rope_type: str): +def _assert_bf16_built_cos_sin_matches_fp32(): """The bf16-built model's rope table must produce the fp32-built model's cos/sin, bitwise.""" position_ids = torch.arange(MAX_POS)[None, :] x_bf16 = torch.zeros(1, MAX_POS, HEAD_DIM, dtype=torch.bfloat16) x_fp32 = torch.zeros(1, MAX_POS, HEAD_DIM, dtype=torch.float32) - with torch.no_grad(): - cos_bf16, sin_bf16 = _build(rope_type, "bfloat16").model.rotary_emb(x_bf16, position_ids) - cos_fp32, sin_fp32 = _build(rope_type, "float32").model.rotary_emb(x_fp32, position_ids) + # Forward consumption is shared across registry entries. Default covers the + # unscaled path; YaRN covers a non-unit attention scaling factor. + for rope_type in ("default", "yarn"): + with torch.no_grad(): + cos_bf16, sin_bf16 = _build(rope_type, "bfloat16").model.rotary_emb(x_bf16, position_ids) + cos_fp32, sin_fp32 = _build(rope_type, "float32").model.rotary_emb(x_fp32, position_ids) - assert torch.equal(cos_bf16.float(), cos_fp32.to(torch.bfloat16).float()), f"{rope_type}: cos differs" - assert torch.equal(sin_bf16.float(), sin_fp32.to(torch.bfloat16).float()), f"{rope_type}: sin differs" + assert torch.equal(cos_bf16.float(), cos_fp32.to(torch.bfloat16).float()), f"{rope_type}: cos differs" + assert torch.equal(sin_bf16.float(), sin_fp32.to(torch.bfloat16).float()), f"{rope_type}: sin differs" -@pytest.mark.parametrize("rope_type", sorted(ROPE_INIT_FUNCTIONS)) -def test_contract_lane_bits_unchanged(rope_type: str): +def _assert_contract_lane_bits_unchanged(): """rope_native (the zero-K3 contract lane) built fp32 reads exactly what it read before.""" position_ids = torch.arange(MAX_POS)[None, :] x = torch.zeros(1, MAX_POS, HEAD_DIM, dtype=torch.float32) with torch.no_grad(): - contract = _build(rope_type, "float32", rope_native=True).model.rotary_emb(x, position_ids) - stock = _build(rope_type, "float32", rope_native=False).model.rotary_emb(x, position_ids) + contract = _build("default", "float32", rope_native=True).model.rotary_emb(x, position_ids) + stock = _build("default", "float32", rope_native=False).model.rotary_emb(x, position_ids) - assert torch.equal(contract[0], stock[0]), f"{rope_type}: contract-lane cos moved" - assert torch.equal(contract[1], stock[1]), f"{rope_type}: contract-lane sin moved" + assert torch.equal(contract[0], stock[0]), "contract-lane cos moved" + assert torch.equal(contract[1], stock[1]), "contract-lane sin moved" -def test_native_default_cache_is_lazy_and_follows_execution_device(): +def _assert_native_default_cache_is_lazy_and_follows_execution_device(): rotary = _build("default", "float32", rope_native=True).model.rotary_emb assert rotary._sglang_default_cache is None @@ -142,8 +146,10 @@ def test_native_default_cache_is_lazy_and_follows_execution_device(): assert torch.equal(cos[..., : HEAD_DIM // 2].reshape_as(cached_cos), cached_cos) assert torch.equal(sin[..., : HEAD_DIM // 2].reshape_as(cached_sin), cached_sin) + _assert_qwen_class_b_cache_growth_preserves_cpu_fp32_recipe() + -def test_qwen_class_b_candidate_default_cache_growth_preserves_cpu_fp32_recipe(): +def _assert_qwen_class_b_cache_growth_preserves_cpu_fp32_recipe(): config = _config("default") config.max_position_embeddings = 8 config._rope_native = True @@ -196,3 +202,59 @@ def test_exact_architectures_build_default_rope_tables_on_their_serving_devices( assert torch.equal(glm_rotary._build_sglang_default_cache(MAX_POS, device), cuda_table) assert torch.equal(qwen_rotary._build_sglang_default_cache(MAX_POS, device), cpu_table) + + _assert_dense_qwen_eager_rope_matches_serving_bits() + + +def _assert_dense_qwen_eager_rope_matches_serving_bits(): + """Keep the dense-Qwen zero-K3 apply contract beside its table owner.""" + serving_rope_theta = 1_000_000 + serving_max_position = 40960 + sequence = 96 + q_heads = 16 + kv_heads = 8 + + class _Config: + rope_scaling = None + head_dim = HEAD_DIM + hidden_size = q_heads * HEAD_DIM + num_attention_heads = q_heads + max_position_embeddings = serving_max_position + rope_theta = serving_rope_theta + rope_parameters = {} + + torch.manual_seed(0) + q = torch.randn(1, sequence, q_heads, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, sequence, kv_heads, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + positions = torch.arange(sequence, device="cuda") + + rotary = RotaryEmbedding(_Config(), device="cuda") + assert rope_module._flash_apply_rotary_emb is None + cos, sin = rotary(q.view(1, sequence, -1), positions.unsqueeze(0)) + q_out, k_out = rope_module.apply_rotary_pos_emb(q, k, cos, sin) + + inv_freq = 1.0 / ( + serving_rope_theta ** (torch.arange(0, HEAD_DIM, 2, dtype=torch.float32, device="cuda") / HEAD_DIM) + ) + frequencies = torch.einsum( + "i,j->ij", + torch.arange(serving_max_position, dtype=torch.float32, device="cuda"), + inv_freq, + ) + cache = torch.cat((frequencies.cos(), frequencies.sin()), dim=-1) + + def serving_apply(value, selected_cos, selected_sin): + selected_cos = selected_cos.unsqueeze(-2).to(value.dtype) + selected_sin = selected_sin.unsqueeze(-2).to(value.dtype) + first, second = torch.chunk(value, 2, dim=-1) + return torch.cat( + (first * selected_cos - second * selected_sin, second * selected_cos + first * selected_sin), + dim=-1, + ) + + selected_cos, selected_sin = cache.index_select(0, positions).chunk(2, dim=-1) + q_reference = serving_apply(q.view(sequence, q_heads, HEAD_DIM), selected_cos, selected_sin).unsqueeze(0) + k_reference = serving_apply(k.view(sequence, kv_heads, HEAD_DIM), selected_cos, selected_sin).unsqueeze(0) + + assert torch.equal(q_out, q_reference) + assert torch.equal(k_out, k_reference) diff --git a/tests/models/test_sglang_rmsnorm.py b/tests/models/test_sglang_rmsnorm.py deleted file mode 100644 index 068f8b90..00000000 --- a/tests/models/test_sglang_rmsnorm.py +++ /dev/null @@ -1,26 +0,0 @@ -import torch - -from xorl.models.layers.normalization import RMSNorm, get_rmsnorm_mode, set_rmsnorm_mode - - -def test_sglang_rmsnorm_force_residual_uses_fp32_weight_multiply(): - hidden = torch.tensor([[1.5, -2.0, 0.25, 4.0]], dtype=torch.bfloat16) - norm = RMSNorm(hidden_size=4, eps=1e-6, mode="sglang") - norm.weight.data = torch.tensor([1.0, 1.25, -0.5, 2.0], dtype=torch.float32) - - out = norm(hidden, force_sglang_residual=True) - - hidden_f = hidden.float() - expected = hidden_f * torch.rsqrt(hidden_f.pow(2).mean(dim=-1, keepdim=True) + 1e-6) - expected = (expected * norm.weight.float()).to(hidden.dtype) - torch.testing.assert_close(out, expected) - - -def test_global_sglang_rmsnorm_mode_is_accepted(): - previous = get_rmsnorm_mode() - try: - set_rmsnorm_mode("sglang") - assert get_rmsnorm_mode() == "sglang" - assert RMSNorm(hidden_size=2).mode == "sglang" - finally: - set_rmsnorm_mode(previous) diff --git a/tests/models/test_topk_router.py b/tests/models/test_topk_router.py index 6b8c628d..1f655b80 100644 --- a/tests/models/test_topk_router.py +++ b/tests/models/test_topk_router.py @@ -11,13 +11,15 @@ * ``routed_scaling_factor`` multiplies the post-renorm weights. """ +from types import SimpleNamespace + import pytest import torch import torch.nn.functional as F from torch import nn from xorl.models.layers.moe import moe_block as moe_block_module -from xorl.models.layers.moe.moe_block import MoEBlock, _router_fp32_layers_enabled +from xorl.models.layers.moe.moe_block import MoEBlock from xorl.models.layers.moe.router import TopKRouter @@ -29,7 +31,7 @@ # --------------------------------------------------------------------------- -def test_softmax_path_matches_reference(): +def _assert_softmax_policy_matches_reference_and_ignores_v4_inputs(): torch.manual_seed(0) num_tokens, num_experts, top_k = 5, 8, 2 logits = torch.randn(num_tokens, num_experts) @@ -44,25 +46,21 @@ def test_softmax_path_matches_reference(): assert torch.equal(experts, ref_experts) torch.testing.assert_close(weights, ref_weights) - -def test_softmax_path_no_renorm(): + # Disabling normalization preserves the selected probability mass. torch.manual_seed(1) - logits = torch.randn(3, 6) - router = TopKRouter(num_experts=6, top_k=2, norm_topk_prob=False) - weights, _ = router(logits, input_dtype=torch.float32) - # When renorm is off, weights should NOT sum to 1 - sums = weights.sum(dim=-1) + unnormalized_logits = torch.randn(3, 6) + unnormalized = TopKRouter(num_experts=6, top_k=2, norm_topk_prob=False) + unnormalized_weights, _ = unnormalized(unnormalized_logits, input_dtype=torch.float32) + sums = unnormalized_weights.sum(dim=-1) assert (sums - 1).abs().max() > 1e-3 - -def test_softmax_path_invariant_to_v4_kwargs(): - """Passing V4-only kwargs to a softmax router has no effect.""" + # V4-only inputs are inert on the legacy softmax path. torch.manual_seed(2) - logits = torch.randn(4, 8) - router = TopKRouter(num_experts=8, top_k=2) - w_a, e_a = router(logits, input_dtype=torch.float32) - w_b, e_b = router( - logits, + legacy_logits = torch.randn(4, 8) + legacy = TopKRouter(num_experts=8, top_k=2) + w_a, e_a = legacy(legacy_logits, input_dtype=torch.float32) + w_b, e_b = legacy( + legacy_logits, input_dtype=torch.float32, expert_bias=torch.randn(8), # ignored tid2eid=torch.randint(0, 8, (100, 2)), # ignored @@ -72,7 +70,7 @@ def test_softmax_path_invariant_to_v4_kwargs(): torch.testing.assert_close(w_a, w_b) -def test_synthetic_balanced_routing_overrides_softmax_selection(monkeypatch): +def _assert_synthetic_balanced_routing_overrides_softmax_hash_and_bias(monkeypatch): monkeypatch.setenv("XORL_MOE_SYNTHETIC_ROUTING", "balanced") torch.manual_seed(7) num_tokens, num_experts, top_k = 7, 8, 2 @@ -94,76 +92,30 @@ def test_synthetic_balanced_routing_overrides_softmax_selection(monkeypatch): expected_weights = torch.full((num_tokens, top_k), 1.0 / top_k) torch.testing.assert_close(weights, expected_weights) - -@pytest.mark.parametrize( - ("policy", "expected_experts"), - [ - ("stable_low_id", [[0, 1], [0, 1]]), - ("tie_low_id", [[0, 1], [0, 1]]), - ("tie_high_id", [[2, 1], [3, 1]]), - ], -) -def test_softmax_topk_diagnostic_policy_breaks_ties_without_changing_weights( - monkeypatch, - policy, - expected_experts, -): - monkeypatch.setenv("XORL_MOE_ROUTER_TOPK_POLICY", policy) - logits = torch.tensor( - [ - [1.0, 1.0, 1.0, 0.0], - [0.5, 0.5, 0.0, 0.5], - ] + table = torch.zeros(16, top_k, dtype=torch.int32) + input_ids = torch.arange(num_tokens, dtype=torch.long) + bias = torch.full((num_experts,), 1000.0) + v4_router = TopKRouter( + num_experts=num_experts, + top_k=top_k, + scoring_func="sqrtsoftplus", + topk_method="noaux_tc", ) - router = TopKRouter(num_experts=4, top_k=2, norm_topk_prob=True) - - weights, experts = router(logits, input_dtype=torch.float32) - - expected_experts = torch.tensor(expected_experts, dtype=torch.long) - assert torch.equal(experts, expected_experts) - probs = F.softmax(logits, dim=1, dtype=torch.float32) - expected_weights = torch.gather(probs, dim=1, index=expected_experts) - expected_weights = expected_weights / expected_weights.sum(dim=-1, keepdim=True) - torch.testing.assert_close(weights, expected_weights) - - -def test_softmax_topk_logits_policy_selects_from_logits_and_gathers_softmax(monkeypatch): - monkeypatch.setenv("XORL_MOE_ROUTER_TOPK_POLICY", "logits") - logits = torch.tensor([[0.0, 1.0, 3.0, 2.0]]) - router = TopKRouter(num_experts=4, top_k=2, norm_topk_prob=False) - - weights, experts = router(logits, input_dtype=torch.float32) - - expected_experts = torch.tensor([[2, 3]]) - assert torch.equal(experts, expected_experts) - probs = F.softmax(logits, dim=1, dtype=torch.float32) - torch.testing.assert_close(weights, torch.gather(probs, dim=1, index=expected_experts)) - - -def test_invalid_softmax_topk_diagnostic_policy_raises(monkeypatch): - monkeypatch.setenv("XORL_MOE_ROUTER_TOPK_POLICY", "surprise") - router = TopKRouter(num_experts=4, top_k=2) - - with pytest.raises(ValueError, match="XORL_MOE_ROUTER_TOPK_POLICY"): - router(torch.randn(2, 4), input_dtype=torch.float32) - - -def test_router_fp32_layers_selector(monkeypatch): - monkeypatch.delenv("XORL_MOE_ROUTER_FP32_LAYERS", raising=False) - assert not _router_fp32_layers_enabled(15) - - monkeypatch.setenv("XORL_MOE_ROUTER_FP32_LAYERS", "14,15-16") - assert not _router_fp32_layers_enabled(13) - assert _router_fp32_layers_enabled(14) - assert _router_fp32_layers_enabled(15) - assert _router_fp32_layers_enabled(16) - assert not _router_fp32_layers_enabled(17) - - monkeypatch.setenv("XORL_MOE_ROUTER_FP32_LAYERS", "all") - assert _router_fp32_layers_enabled(None) + v4_weights, v4_experts = v4_router( + logits, + input_dtype=torch.float32, + expert_bias=bias, + tid2eid=table, + input_ids=input_ids, + ) + assert torch.equal(v4_experts, expected_experts) + scores = F.softplus(logits.float()).sqrt().type_as(logits) + expected_v4_weights = torch.gather(scores, dim=1, index=expected_experts) + expected_v4_weights = expected_v4_weights / (expected_v4_weights.sum(dim=-1, keepdim=True) + 1e-20) + torch.testing.assert_close(v4_weights, expected_v4_weights) -def test_moe_block_uses_layer_scoped_router_fp32(monkeypatch): +def _assert_moe_block_uses_configured_router_fp32(monkeypatch): class RecordingGate(nn.Module): def __init__(self): super().__init__() @@ -182,7 +134,7 @@ def forward(self, hidden_states): moe_implementation="eager", train_router=False, ) - block.layer_idx = 15 + block.config = SimpleNamespace(_router_fp32=True) block.gate = RecordingGate() calls = [] @@ -190,7 +142,6 @@ def fake_linear(hidden_states, weight, bias=None): calls.append((hidden_states.dtype, weight.dtype, bias)) return torch.tensor([[0.0, 1.0], [1.0, 0.0]], dtype=torch.float32) - monkeypatch.setenv("XORL_MOE_ROUTER_FP32_LAYERS", "15") monkeypatch.setattr(moe_block_module.F, "linear", fake_linear) _, selected_experts, router_logits = block.route(torch.ones(2, 2, dtype=torch.bfloat16)) @@ -201,12 +152,27 @@ def fake_linear(hidden_states, weight, bias=None): assert torch.equal(selected_experts, torch.tensor([[1], [0]])) +def _assert_trainable_router_rejects_deepep_dispatch(): + block = MoEBlock( + hidden_size=16, + num_experts=4, + top_k=2, + intermediate_size=32, + moe_implementation="eager", + train_router=True, + ) + block.experts.ep_dispatch = "deepep" + + with pytest.raises(AssertionError, match="ep_dispatch='deepep'"): + block(torch.randn(1, 4, 16)) + + # --------------------------------------------------------------------------- # DSv4 sqrtsoftplus + noaux_tc # --------------------------------------------------------------------------- -def test_sqrtsoftplus_noaux_selects_via_biased_scores(): +def _assert_sqrtsoftplus_noaux_uses_unbiased_weights_and_requires_bias(): """Selection comes from ``scores + bias``, weights from unbiased scores.""" torch.manual_seed(3) num_tokens, num_experts, top_k = 4, 6, 2 @@ -232,11 +198,9 @@ def test_sqrtsoftplus_noaux_selects_via_biased_scores(): expected = expected / (expected.sum(dim=-1, keepdim=True) + 1e-20) torch.testing.assert_close(weights, expected) - -def test_sqrtsoftplus_noaux_requires_bias(): - router = TopKRouter(num_experts=4, top_k=2, scoring_func="sqrtsoftplus", topk_method="noaux_tc") + missing_bias = TopKRouter(num_experts=4, top_k=2, scoring_func="sqrtsoftplus", topk_method="noaux_tc") with pytest.raises(AssertionError, match="noaux_tc requires expert_bias"): - router(torch.randn(3, 4), input_dtype=torch.float32) + missing_bias(torch.randn(3, 4), input_dtype=torch.float32) # --------------------------------------------------------------------------- @@ -244,7 +208,7 @@ def test_sqrtsoftplus_noaux_requires_bias(): # --------------------------------------------------------------------------- -def test_hash_routing_uses_tid2eid_for_selection(): +def _assert_hash_routing_uses_tid2eid_ignores_bias_and_requires_input_ids(): """Top-k indices come from ``tid2eid[input_ids]``; weights from gate.""" torch.manual_seed(4) vocab_size, num_experts, top_k = 16, 8, 2 @@ -275,70 +239,20 @@ def test_hash_routing_uses_tid2eid_for_selection(): expected_w = expected_w / (expected_w.sum(dim=-1, keepdim=True) + 1e-20) torch.testing.assert_close(weights, expected_w) - -def test_hash_routing_ignores_bias(): - """When tid2eid is set, bias does not affect anything.""" - torch.manual_seed(5) - vocab_size, num_experts, top_k = 16, 8, 2 - logits = torch.randn(3, num_experts) - table = torch.zeros(vocab_size, top_k, dtype=torch.int32) # all to expert 0 - input_ids = torch.tensor([0, 1, 2], dtype=torch.long) - - router = TopKRouter(num_experts=num_experts, top_k=top_k, scoring_func="sqrtsoftplus") - _, experts_no_bias = router(logits, input_dtype=torch.float32, tid2eid=table, input_ids=input_ids) big_bias = torch.full((num_experts,), 1000.0) big_bias[0] = -1000.0 - _, experts_with_bias = router( + weights_with_bias, experts_with_bias = router( logits, input_dtype=torch.float32, tid2eid=table, input_ids=input_ids, expert_bias=big_bias, ) - assert torch.equal(experts_no_bias, experts_with_bias) - assert (experts_no_bias == 0).all() - + assert torch.equal(experts, experts_with_bias) + torch.testing.assert_close(weights, weights_with_bias) -def test_hash_routing_requires_input_ids(): - table = torch.zeros(8, 2, dtype=torch.int32) - router = TopKRouter(num_experts=4, top_k=2, scoring_func="sqrtsoftplus") with pytest.raises(AssertionError, match="requires input_ids"): - router(torch.randn(3, 4), input_dtype=torch.float32, tid2eid=table) - - -def test_synthetic_balanced_routing_overrides_hash_and_bias(monkeypatch): - monkeypatch.setenv("XORL_MOE_SYNTHETIC_ROUTING", "balanced") - torch.manual_seed(8) - num_tokens, vocab_size, num_experts, top_k = 6, 16, 8, 2 - logits = torch.randn(num_tokens, num_experts) - table = torch.zeros(vocab_size, top_k, dtype=torch.int32) - input_ids = torch.arange(num_tokens, dtype=torch.long) - bias = torch.full((num_experts,), 1000.0) - - router = TopKRouter( - num_experts=num_experts, - top_k=top_k, - scoring_func="sqrtsoftplus", - topk_method="noaux_tc", - ) - weights, experts = router( - logits, - input_dtype=torch.float32, - expert_bias=bias, - tid2eid=table, - input_ids=input_ids, - ) - - expected_experts = torch.tensor( - [[0, 1], [2, 3], [4, 5], [6, 7], [0, 1], [2, 3]], - dtype=torch.long, - ) - assert torch.equal(experts, expected_experts) - - scores = F.softplus(logits.float()).sqrt().type_as(logits) - expected_weights = torch.gather(scores, dim=1, index=expected_experts) - expected_weights = expected_weights / (expected_weights.sum(dim=-1, keepdim=True) + 1e-20) - torch.testing.assert_close(weights, expected_weights) + router(logits, input_dtype=torch.float32, tid2eid=table) # --------------------------------------------------------------------------- @@ -346,7 +260,7 @@ def test_synthetic_balanced_routing_overrides_hash_and_bias(monkeypatch): # --------------------------------------------------------------------------- -def test_routed_scaling_factor_multiplies_weights(): +def _assert_routed_scaling_factor_multiplies_v4_and_rejects_softmax(): torch.manual_seed(6) logits = torch.randn(3, 6) bias = torch.zeros(6) @@ -362,21 +276,67 @@ def test_routed_scaling_factor_multiplies_weights(): w_scaled, _ = scaled(logits, input_dtype=torch.float32, expert_bias=bias) torch.testing.assert_close(w_scaled, w_base * 1.5) - -def test_routed_scaling_factor_rejected_on_softmax_path(): - """``routed_scaling_factor`` is V4-only; constructing a softmax router with - one set raises rather than silently ignoring it. - """ with pytest.raises(ValueError, match="routed_scaling_factor is only used"): TopKRouter(num_experts=6, top_k=2, routed_scaling_factor=1.5) +def _assert_moe_block_regather_matches_router_policy(): + def make_block(scoring_func: str, routed_scaling_factor=None): + block = MoEBlock( + hidden_size=16, + num_experts=8, + top_k=2, + intermediate_size=16, + moe_implementation="eager", + ) + block.router = TopKRouter( + num_experts=8, + top_k=2, + norm_topk_prob=True, + scoring_func=scoring_func, + topk_method="noaux_tc" if scoring_func == "sqrtsoftplus" else None, + routed_scaling_factor=routed_scaling_factor, + ) + return block + + torch.manual_seed(0) + block = make_block("sqrtsoftplus", routed_scaling_factor=2.5) + router_logits = torch.randn(6, 8) + eager_weights, eager_experts = block.router( + router_logits, + input_dtype=torch.float32, + expert_bias=torch.randn(8) * 0.1, + ) + + regathered_experts, regathered_weights = block._regather_routing( + router_logits, + eager_experts, + input_dtype=torch.float32, + ) + + assert torch.equal(regathered_experts, eager_experts) + torch.testing.assert_close(regathered_weights, eager_weights, rtol=1e-5, atol=1e-6) + + block.router.routed_scaling_factor = None + _, unscaled_weights = block._regather_routing(router_logits, eager_experts, input_dtype=torch.float32) + torch.testing.assert_close(regathered_weights, unscaled_weights * 2.5, rtol=1e-5, atol=1e-6) + _, bf16_weights = block._regather_routing(router_logits, eager_experts, input_dtype=torch.bfloat16) + assert bf16_weights.dtype is torch.bfloat16 + + torch.manual_seed(3) + block = make_block("softmax") + router_logits = torch.randn(5, 8) + eager_weights, eager_experts = block.router(router_logits, input_dtype=torch.float32) + _, regathered_weights = block._regather_routing(router_logits, eager_experts, input_dtype=torch.float32) + torch.testing.assert_close(regathered_weights, eager_weights, rtol=1e-5, atol=1e-6) + + # --------------------------------------------------------------------------- # from_config # --------------------------------------------------------------------------- -def test_from_config_reads_v4_fields(): +def _assert_from_config_selects_v4_and_legacy_policies(): """from_config picks up sqrtsoftplus + noaux_tc + scaling_factor from a V4 config.""" from xorl.models.transformers.deepseek_v4 import DeepseekV4Config # noqa: PLC0415 @@ -388,16 +348,25 @@ def test_from_config_reads_v4_fields(): assert router.top_k == cfg.num_experts_per_tok assert router.num_experts == cfg.n_routed_experts - -def test_from_config_legacy_softmax(): - """Legacy non-V4 configs land on the softmax path.""" - class _LegacyCfg: num_experts = 16 num_experts_per_tok = 2 norm_topk_prob = True - router = TopKRouter.from_config(_LegacyCfg()) - assert router.scoring_func == "softmax" - assert router.topk_method is None - assert router.routed_scaling_factor is None + legacy = TopKRouter.from_config(_LegacyCfg()) + assert legacy.scoring_func == "softmax" + assert legacy.topk_method is None + assert legacy.routed_scaling_factor is None + + +def test_topk_router_legacy_and_configuration_contract(monkeypatch): + with monkeypatch.context() as case_patch: + _assert_synthetic_balanced_routing_overrides_softmax_hash_and_bias(case_patch) + _assert_sqrtsoftplus_noaux_uses_unbiased_weights_and_requires_bias() + _assert_hash_routing_uses_tid2eid_ignores_bias_and_requires_input_ids() + _assert_softmax_policy_matches_reference_and_ignores_v4_inputs() + _assert_routed_scaling_factor_multiplies_v4_and_rejects_softmax() + _assert_moe_block_regather_matches_router_policy() + _assert_from_config_selects_v4_and_legacy_policies() + _assert_moe_block_uses_configured_router_fp32(monkeypatch) + _assert_trainable_router_rejects_deepep_dispatch() diff --git a/tests/models/test_unfuse_projections.py b/tests/models/test_unfuse_projections.py new file mode 100644 index 00000000..d2fcfbd8 --- /dev/null +++ b/tests/models/test_unfuse_projections.py @@ -0,0 +1,218 @@ +"""Unfusing ``qkv_proj`` / ``gate_up_proj``: checkpoint handlers and the MoE shared expert. + +The load-time contract when a model is unfused is narrow: the checkpoint handler stops +merging the projections the model no longer fuses, and keeps doing everything else -- +pre-quantized loading paths, and on Qwen3.5 the GatedDeltaNet ``in_proj_qkv`` remapping. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from xorl.models.transformers.qwen3.checkpoint_handler import Qwen3CheckpointHandler +from xorl.models.transformers.qwen3_5 import parallelize as qwen3_5_parallelize +from xorl.models.transformers.qwen3_5.checkpoint_handler import Qwen3_5CheckpointHandler +from xorl.models.transformers.qwen3_5_moe import parallelize as qwen3_5_moe_parallelize + + +pytestmark = [pytest.mark.cpu] + + +GATE_KEY = "model.layers.0.mlp.gate_proj.weight" +UP_KEY = "model.layers.0.mlp.up_proj.weight" +Q_KEY = "model.layers.0.self_attn.q_proj.weight" +K_KEY = "model.layers.0.self_attn.k_proj.weight" +V_KEY = "model.layers.0.self_attn.v_proj.weight" + +LINEAR_KEY_DIM = 6 +LINEAR_VALUE_DIM = 10 + + +def _qwen3_handler(**kwargs: object) -> Qwen3CheckpointHandler: + return Qwen3CheckpointHandler(num_attention_heads=4, num_key_value_heads=2, head_dim=8, **kwargs) + + +def _qwen3_5_handler(**kwargs: object) -> Qwen3_5CheckpointHandler: + return Qwen3_5CheckpointHandler( + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + linear_key_dim=LINEAR_KEY_DIM, + linear_value_dim=LINEAR_VALUE_DIM, + **kwargs, + ) + + +def _assert_passthrough(result: list, key: str, tensor: torch.Tensor) -> None: + [(name, passed)] = result + assert name == key + assert passed is tensor + + +class TestQwen3HandlerSkipFlags: + def test_merges_gate_and_up_by_default(self): + handler = _qwen3_handler() + + assert handler.on_load_weight(GATE_KEY, torch.ones(3, 2)) == [] + [(name, merged)] = handler.on_load_weight(UP_KEY, torch.zeros(3, 2)) + + assert name == "model.layers.0.mlp.gate_up_proj.weight" + assert merged.shape == (6, 2) + + def test_merges_qkv_by_default(self): + handler = _qwen3_handler() + + assert handler.on_load_weight(Q_KEY, torch.ones(32, 2)) == [] + assert handler.on_load_weight(K_KEY, torch.ones(16, 2)) == [] + [(name, _merged)] = handler.on_load_weight(V_KEY, torch.ones(16, 2)) + + assert name == "model.layers.0.self_attn.qkv_proj.weight" + + def test_gate_and_up_pass_through_when_skipped(self): + handler = _qwen3_handler(skip_gate_up_merge=True) + gate = torch.ones(3, 2) + + _assert_passthrough(handler.on_load_weight(GATE_KEY, gate), GATE_KEY, gate) + + def test_qkv_passes_through_when_skipped(self): + handler = _qwen3_handler(skip_qkv_merge=True) + q = torch.ones(32, 2) + + _assert_passthrough(handler.on_load_weight(Q_KEY, q), Q_KEY, q) + + def test_skipping_one_merge_leaves_the_other_active(self): + """The flags are independent: an architecture can fuse one and not the other.""" + handler = _qwen3_handler(skip_qkv_merge=True) + q = torch.ones(32, 2) + + _assert_passthrough(handler.on_load_weight(Q_KEY, q), Q_KEY, q) + assert handler.on_load_weight(GATE_KEY, torch.ones(3, 2)) == [] + [(name, _merged)] = handler.on_load_weight(UP_KEY, torch.zeros(3, 2)) + assert name == "model.layers.0.mlp.gate_up_proj.weight" + + def test_no_pending_warning_when_merges_are_skipped(self): + """A skipped merge must not report the keys it deliberately never buffered.""" + handler = _qwen3_handler(skip_qkv_merge=True, skip_gate_up_merge=True) + handler.on_load_weight(GATE_KEY, torch.ones(3, 2)) + handler.on_load_weight(Q_KEY, torch.ones(32, 2)) + + with warnings.catch_warnings(record=True) as recorded: + warnings.simplefilter("always") + handler.on_load_complete() + + assert [w for w in recorded if "Incomplete" in str(w.message)] == [] + + +class TestQwen3_5HandlerKeepsLinearAttentionMapping: + """The regression the granular flags exist to prevent. + + ``Qwen3_5CheckpointHandler`` also splits the GatedDeltaNet ``in_proj_qkv`` packing, + which has nothing to do with how the MLP is stored. Skipping the gate/up merge must + not take that with it. + """ + + def test_in_proj_qkv_still_splits_when_gate_up_merge_is_skipped(self): + handler = _qwen3_5_handler(skip_gate_up_merge=True) + rows = 2 * LINEAR_KEY_DIM + LINEAR_VALUE_DIM + tensor = torch.arange(rows * 2, dtype=torch.float32).reshape(rows, 2) + + result = handler.on_load_weight("model.layers.0.linear_attn.in_proj_qkv.weight", tensor) + + assert [name for name, _ in result] == [ + "model.layers.0.linear_attn.q_proj.weight", + "model.layers.0.linear_attn.k_proj.weight", + "model.layers.0.linear_attn.v_proj.weight", + ] + (_, q), (_, k), (_, v) = result + assert torch.equal(q, tensor[:LINEAR_KEY_DIM]) + assert torch.equal(k, tensor[LINEAR_KEY_DIM : 2 * LINEAR_KEY_DIM]) + assert torch.equal(v, tensor[2 * LINEAR_KEY_DIM :]) + + def test_gate_and_up_pass_through_when_skipped(self): + handler = _qwen3_5_handler(skip_gate_up_merge=True) + gate = torch.ones(3, 2) + + _assert_passthrough(handler.on_load_weight(GATE_KEY, gate), GATE_KEY, gate) + + def test_gate_and_up_still_merge_by_default(self): + handler = _qwen3_5_handler() + + assert handler.on_load_weight(GATE_KEY, torch.ones(3, 2)) == [] + [(name, _merged)] = handler.on_load_weight(UP_KEY, torch.zeros(3, 2)) + + assert name == "model.layers.0.mlp.gate_up_proj.weight" + + +class _StubFusedMLP(nn.Module): + """Stands in for ``Qwen3_5MoeMLP`` / ``Qwen3_5MLP``: fused until unfused.""" + + def __init__(self, hidden: int = 4, intermediate: int = 4): + super().__init__() + self.hidden = hidden + self.intermediate = intermediate + self.gate_up_proj = nn.Linear(hidden, 2 * intermediate, bias=False) + self.down_proj = nn.Linear(intermediate, hidden, bias=False) + + def unfuse_for_tp(self): + self.gate_proj = nn.Linear(self.hidden, self.intermediate, bias=False) + self.up_proj = nn.Linear(self.hidden, self.intermediate, bias=False) + del self.gate_up_proj + + @property + def is_unfused(self) -> bool: + return not hasattr(self, "gate_up_proj") + + +class _StubMoEBlock(nn.Module): + """Stands in for ``Qwen3_5MoeSparseMoeBlock``: routed experts plus a shared expert.""" + + def __init__(self, shared_expert: _StubFusedMLP) -> None: + super().__init__() + self.experts = nn.Linear(4, 4, bias=False) + self.shared_expert = shared_expert + + +class _Layer(nn.Module): + def __init__(self, mlp: nn.Module) -> None: + super().__init__() + self.self_attn = None + self.mlp = mlp + + +class _Model(nn.Module): + def __init__(self, layers: list[nn.Module]) -> None: + super().__init__() + self.model = nn.Module() + self.model.layers = nn.ModuleList(layers) + self.config = SimpleNamespace(base_model_tp_plan=None) + + +class TestSharedExpertUnfusing: + def test_unfuses_dense_mlps_and_shared_experts_alike(self): + """Routed expert weights are EP-sharded, not TP-sharded, so only the block's + shared expert and the dense MLPs are split. + """ + shared = _StubFusedMLP() + dense = _StubFusedMLP() + moe_block = _StubMoEBlock(shared) + model = _Model([_Layer(moe_block), _Layer(dense)]) + + qwen3_5_moe_parallelize.unfuse_for_tp(model) + + assert shared.is_unfused + assert dense.is_unfused + assert isinstance(moe_block.experts, nn.Linear) + + +class TestDenseQwen3_5Unfusing: + def test_layers_without_self_attn_are_tolerated(self): + """Qwen3.5's linear-attention layers carry no ``self_attn``.""" + dense = _StubFusedMLP() + model = _Model([_Layer(dense)]) + + qwen3_5_parallelize.unfuse_for_tp(model) + + assert dense.is_unfused diff --git a/tests/ops/data/bi_golden_trees.json b/tests/ops/data/bi_golden_trees.json index 8dd9e3b1..a6c478e1 100644 --- a/tests/ops/data/bi_golden_trees.json +++ b/tests/ops/data/bi_golden_trees.json @@ -400,18 +400,6 @@ "tail": "99407941b1429fc2cec11ac04dc36cc1" } }, - "qk_v2_strided_t256_h8_d128": { - "out": { - "sha256": "d251cf917d759561c0e6c6f0ec7b7c63d5682f21b5bb17a2100f77e7e894832f", - "shape": [ - 256, - 1024 - ], - "dtype": "torch.bfloat16", - "head": "5d4077c024c0cabf76bf6e3f553f4f3f553fbfbf083f48c0b73f96bf573f70bf", - "tail": "8e3b4c3ad43f07ba15bc923e2440c33f" - } - }, "head_v2_scoring_n64_real": { "logprob": { "sha256": "aed8ac4fe4ea61e054cac382841b16449c609a7f5674ddd33cf35499c22dfc34", diff --git a/tests/ops/dsv4/test_compressor.py b/tests/ops/dsv4/test_compressor.py index 0bfceeee..3e5f6b51 100644 --- a/tests/ops/dsv4/test_compressor.py +++ b/tests/ops/dsv4/test_compressor.py @@ -5,6 +5,7 @@ import pytest import torch +from xorl.ops.dsv4 import utils from xorl.ops.dsv4.compressor import DeepSeekV4Compressor from xorl.ops.dsv4.rope import precompute_freqs_cis @@ -36,8 +37,7 @@ def _compressor_config(max_position_embeddings=768): ) -def test_c128_context_parallel_only_requires_ratio_divisibility(monkeypatch): - monkeypatch.delenv("XORL_DSV4_ROPE_MAX_SEQ_LEN", raising=False) +def test_context_parallel_compression_ratio_admission_policy(monkeypatch): precompute_freqs_cis.cache_clear() compressor = DeepSeekV4Compressor( @@ -56,9 +56,22 @@ def test_c128_context_parallel_only_requires_ratio_divisibility(monkeypatch): assert out.shape == (1, 3, 16) + # Exercise the cache-capacity guard through its production CP consumer. + short_cache_compressor = DeepSeekV4Compressor( + _compressor_config(max_position_embeddings=600), + head_dim=16, + compress_ratio=128, + rotate=False, + cp_group=_FakeCPGroup(), + ) + with pytest.raises(ValueError, match="RoPE cache is too short"): + short_cache_compressor.forward_raw(torch.ones(1, 384, 32)) -def test_c4_context_parallel_keeps_overlap_divisibility_guard(monkeypatch): - monkeypatch.delenv("XORL_DSV4_ROPE_MAX_SEQ_LEN", raising=False) + _assert_c4_context_parallel_keeps_overlap_divisibility_guard(monkeypatch) + _assert_rotate_activation_fallback_is_orthonormal(monkeypatch) + + +def _assert_c4_context_parallel_keeps_overlap_divisibility_guard(monkeypatch): precompute_freqs_cis.cache_clear() compressor = DeepSeekV4Compressor( @@ -71,3 +84,18 @@ def test_c4_context_parallel_keeps_overlap_divisibility_guard(monkeypatch): with pytest.raises(AssertionError, match="overlap=True"): compressor.forward_raw(torch.ones(1, 12, 32)) + + +def _assert_rotate_activation_fallback_is_orthonormal(monkeypatch): + monkeypatch.setattr(utils, "_fast_hadamard_transform", None) + pattern = torch.tensor([[1.0, 0.0, 0.0, 0.0]], dtype=torch.bfloat16) + torch.testing.assert_close(utils.rotate_activation(pattern), torch.full_like(pattern, 0.5)) + + torch.manual_seed(0) + for width in (4, 16, 64, 256): + values = torch.randn(3, width, dtype=torch.bfloat16) + rotated = utils.rotate_activation(values) + torch.testing.assert_close(utils.rotate_activation(rotated), values, atol=8e-3, rtol=8e-3) + input_norm = values.float().pow(2).sum(dim=-1).sqrt() + output_norm = rotated.float().pow(2).sum(dim=-1).sqrt() + torch.testing.assert_close(input_norm, output_norm, atol=1e-2, rtol=1e-2) diff --git a/tests/ops/dsv4/test_rope.py b/tests/ops/dsv4/test_rope.py deleted file mode 100644 index 8b71c597..00000000 --- a/tests/ops/dsv4/test_rope.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Unit tests for DSv4 RoPE cache sizing.""" - -from types import SimpleNamespace - -import pytest -import torch - -from xorl.ops.dsv4.cp_utils import get_freqs_cis_for_cp -from xorl.ops.dsv4.rope import precompute_freqs_cis, wrapped_precompute_freqs_cis - - -pytestmark = pytest.mark.cpu - - -def _rope_config(max_position_embeddings=32): - return SimpleNamespace( - max_position_embeddings=max_position_embeddings, - rope_parameters={ - "factor": 4.0, - "original_max_position_embeddings": 16, - "beta_fast": 32.0, - "beta_slow": 1.0, - }, - ) - - -def test_wrapped_precompute_freqs_cis_uses_config_max_position_embeddings(monkeypatch): - monkeypatch.delenv("XORL_DSV4_ROPE_MAX_SEQ_LEN", raising=False) - precompute_freqs_cis.cache_clear() - - freqs = wrapped_precompute_freqs_cis(_rope_config(max_position_embeddings=40), rope_head_dim=8, base=10000.0) - - assert freqs.shape == (40, 4) - - -def test_wrapped_precompute_freqs_cis_env_override_wins(monkeypatch): - monkeypatch.setenv("XORL_DSV4_ROPE_MAX_SEQ_LEN", "12") - precompute_freqs_cis.cache_clear() - - freqs = wrapped_precompute_freqs_cis(_rope_config(max_position_embeddings=40), rope_head_dim=8, base=10000.0) - - assert freqs.shape == (12, 4) - - -def test_get_freqs_cis_for_cp_errors_when_cache_too_short(): - class _FakeGroup: - def rank(self): - return 2 - - freqs = torch.empty(10, 4) - - with pytest.raises(ValueError, match="RoPE cache is too short"): - get_freqs_cis_for_cp(freqs, seqlen_local=8, cp_size=4, cp_group=_FakeGroup()) diff --git a/tests/ops/dsv4/test_rotate_activation.py b/tests/ops/dsv4/test_rotate_activation.py deleted file mode 100644 index 8315a3e1..00000000 --- a/tests/ops/dsv4/test_rotate_activation.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Pure-torch FWHT fallback for ``rotate_activation``. - -The on-disk Flash routed-expert weights need ``rotate_activation`` to -work on every C4 layer (DSA indexer with ``rotate=True`` compressor). -The kernel ships in the ``fast_hadamard_transform`` package; we maintain -a pure-torch fallback so CPU CI and lean dev images don't need that -package to import the model. -""" - -import pytest -import torch - - -pytestmark = pytest.mark.cpu - - -def test_fwht_round_trip_orthonormal(): - """Orthonormal Hadamard squared = identity: H @ H @ x == x (modulo - rounding) when ``scale = 1 / sqrt(D)`` is applied at each FWHT call. - """ - from xorl.ops.dsv4.utils import _fwht_torch - - torch.manual_seed(0) - for D in (4, 16, 64, 256): - x = torch.randn(3, D, dtype=torch.bfloat16) - scale = D**-0.5 - once = _fwht_torch(x, scale) - twice = _fwht_torch(once, scale) - # Two ortho-Hadamards = identity, but bf16 accumulates ~5e-3 error. - torch.testing.assert_close(twice, x, atol=8e-3, rtol=8e-3) - - -def test_fwht_known_pattern(): - """``H_4 @ [1, 0, 0, 0] = [1, 1, 1, 1]`` (unnormalized; scale=1).""" - from xorl.ops.dsv4.utils import _fwht_torch - - x = torch.tensor([[1.0, 0.0, 0.0, 0.0]], dtype=torch.bfloat16) - out = _fwht_torch(x, scale=1.0) - expected = torch.ones(1, 4, dtype=torch.bfloat16) - torch.testing.assert_close(out, expected) - - -def test_fwht_preserves_l2_norm_when_orthonormal(): - """An orthonormal transform preserves L2 norm.""" - from xorl.ops.dsv4.utils import _fwht_torch - - torch.manual_seed(1) - D = 128 - x = torch.randn(2, 5, D, dtype=torch.bfloat16) - out = _fwht_torch(x, scale=D**-0.5) - n_x = x.float().pow(2).sum(dim=-1).sqrt() - n_out = out.float().pow(2).sum(dim=-1).sqrt() - torch.testing.assert_close(n_x, n_out, atol=1e-2, rtol=1e-2) - - -def test_fwht_rejects_non_power_of_two(): - from xorl.ops.dsv4.utils import _fwht_torch - - x = torch.randn(3, 6, dtype=torch.bfloat16) - with pytest.raises(ValueError, match="power-of-2"): - _fwht_torch(x, 1.0) - - -def test_rotate_activation_dispatches_to_fallback_when_kernel_missing(monkeypatch): - """When ``fast_hadamard_transform`` isn't installed, ``rotate_activation`` - must transparently use the torch FWHT (no AssertionError).""" - from xorl.ops.dsv4 import utils - - monkeypatch.setattr(utils, "_fast_hadamard_transform", None) - torch.manual_seed(2) - x = torch.randn(2, 4, 64, dtype=torch.bfloat16) # last dim = 64 (power of 2) - out = utils.rotate_activation(x) - assert out.shape == x.shape - assert out.dtype == torch.bfloat16 - assert torch.isfinite(out).all() diff --git a/tests/ops/dsv4/test_v4_tilelang_indexer.py b/tests/ops/dsv4/test_v4_tilelang_indexer.py index e1625e06..06cac426 100644 --- a/tests/ops/dsv4/test_v4_tilelang_indexer.py +++ b/tests/ops/dsv4/test_v4_tilelang_indexer.py @@ -113,20 +113,6 @@ def ref_apply_causal_mask(index_scores, compress_ratio): return index_scores -def ref_fused_qk_topk(q, k, weights, compress_ratio, topk): - """Full reference: scores + causal mask + topk. - - Returns: - index_scores: [batch, seqlen_q, seqlen_kv] fp32 (masked) - topk_indices: [batch, seqlen_q, topk] int64 - """ - index_scores = ref_compute_index_scores(q, weights, k) - index_scores = ref_apply_causal_mask(index_scores, compress_ratio) - actual_topk = min(topk, index_scores.shape[-1]) - topk_indices = index_scores.topk(actual_topk, dim=-1)[1] - return index_scores, topk_indices - - # --------------------------------------------------------------------------- # Test fixtures # --------------------------------------------------------------------------- @@ -159,26 +145,12 @@ def make_inputs(seqlen_q, batch, heads, dim, compress_ratio, device="cuda"): # --------------------------------------------------------------------------- # Test configurations: (seqlen_q, batch, heads, dim, compress_ratio, topk) FORWARD_CONFIGS = [ - # Short sequences, small batch — basic correctness + # One representative for each kernel geometry that changes execution: + # basic, batched, production heads/top-k, and the C128 compression family. (128, 1, 8, 128, 4, 32), - (128, 1, 16, 128, 4, 32), - (128, 2, 8, 128, 4, 32), - # Medium sequences — typical training - (512, 1, 8, 128, 4, 64), - (512, 2, 16, 128, 4, 128), - (512, 1, 64, 128, 4, 128), - # Long sequences - (2048, 1, 8, 128, 4, 128), + (512, 4, 16, 128, 4, 64), (2048, 1, 64, 128, 4, 512), - # C128 layer type (2048, 1, 8, 128, 128, 16), - (1024, 1, 16, 128, 128, 8), - # Larger batch - (512, 4, 8, 128, 4, 64), - (256, 4, 16, 128, 4, 64), - # Edge: seqlen just above compress_ratio (small KV) - (16, 1, 8, 128, 4, 4), - (256, 1, 8, 128, 128, 2), ] FORWARD_CONFIG_IDS = [f"sq{sq}_b{b}_h{h}_d{d}_cr{cr}_top{tk}" for sq, b, h, d, cr, tk in FORWARD_CONFIGS] @@ -209,6 +181,7 @@ def test_indexer_forward_scores(seqlen_q, batch, heads, dim, compress_ratio, top valid_mask = ref_scores != float("-inf") ref_valid = ref_scores[valid_mask] tl_valid = tl_scores[valid_mask] + tl_masked = tl_scores[~valid_mask] diff = compute_diff(ref_valid, tl_valid) print(f"\n[FWD] sq={seqlen_q}, b={batch}, h={heads}, cr={compress_ratio}, topk={topk}") @@ -218,209 +191,18 @@ def test_indexer_forward_scores(seqlen_q, batch, heads, dim, compress_ratio, top assert diff.rel_diff < 1e-3, f"rel_diff too large: {diff.rel_diff:.2e}" assert diff.max_abs_diff < 1.0, f"max_abs_diff too large: {diff.max_abs_diff:.2e}" assert diff.mean_abs_diff < 0.05, f"mean_abs_diff too large: {diff.mean_abs_diff:.2e}" + assert tl_masked.numel() > 0 + assert torch.isneginf(tl_masked).all(), "future compressed groups were not masked" - -@requires_cuda() -@requires_tilelang() -@pytest.mark.parametrize("seqlen_q,batch,heads,dim,compress_ratio,topk", FORWARD_CONFIGS, ids=FORWARD_CONFIG_IDS) -def test_indexer_forward_topk(seqlen_q, batch, heads, dim, compress_ratio, topk): - """Verify topk self-consistency: selected scores >= non-selected scores. - - The kernel is non-deterministic across calls (GPU GEMM accumulation order), so - instead of comparing two separate kernel calls, we verify the output of a single - call is internally consistent. - """ - from xorl.ops.dsv4.kernel.tilelang_indexer import v4_lighting_indexer - from xorl.ops.dsv4.kernel.tilelang_indexer_fwd import ( - _make_causal_cu_seqlens, - batched_indexer_fwd, - ) - - q, k, weights = make_inputs(seqlen_q, batch, heads, dim, compress_ratio) - seqlen_kv = seqlen_q // compress_ratio - - tl_score, tl_topk = v4_lighting_indexer(q, k, weights, compress_ratio, topk) - - # Verify: selected scores should be >= min of selected scores at each position - # Get full logits from the same kernel call is not possible (autograd Function - # only returns topk scores). So re-run forward to get full logits. - cu_ks, cu_ke = _make_causal_cu_seqlens(seqlen_q, seqlen_kv, compress_ratio, q.device) - full_logits = batched_indexer_fwd(q, k, weights, cu_ks, cu_ke) - - b, sq, _ = tl_topk.shape - violations = 0 - total = 0 - for bi in range(b): - for qi in range(sq): - selected = set(tl_topk[bi, qi].cpu().tolist()) - {-1} - if not selected: - continue - selected_scores = full_logits[bi, qi, list(selected)] - min_selected = selected_scores.min().item() - # Check non-selected valid positions have scores <= min_selected - valid_end = (qi + 1) // compress_ratio - for ki in range(valid_end): - if ki not in selected: - total += 1 - if full_logits[bi, qi, ki].item() > min_selected + 1e-5: - violations += 1 - - print( - f"\n[TOPK] sq={seqlen_q}, b={batch}, h={heads}, cr={compress_ratio}, topk={topk}: " - f"violations={violations}/{total}" - ) - if total > 0: - violation_rate = violations / total - # Allow small violation rate due to kernel non-determinism on tied scores - assert violation_rate < 0.05, f"topk violation rate too high: {violation_rate:.4f}" - - -# --------------------------------------------------------------------------- -# Backward tests -# --------------------------------------------------------------------------- -BACKWARD_CONFIGS = [ - # (seqlen_q, batch, heads, dim, compress_ratio, topk) - # Keep topk as power of 2 (required by backward kernel) - (128, 1, 8, 128, 4, 32), - (256, 1, 16, 128, 4, 64), - (512, 1, 8, 128, 4, 64), - (512, 2, 16, 128, 4, 128), - (1024, 1, 8, 128, 4, 128), - (1024, 1, 64, 128, 4, 512), - # C128 - (2048, 1, 8, 128, 128, 16), - # Larger batch - (256, 4, 8, 128, 4, 64), -] - -BACKWARD_CONFIG_IDS = [f"sq{sq}_b{b}_h{h}_d{d}_cr{cr}_top{tk}" for sq, b, h, d, cr, tk in BACKWARD_CONFIGS] - - -def ref_indexer_backward_dense(q, k, weights, compress_ratio, topk, topk_indices): - """Dense PyTorch autograd reference for backward. - - Computes forward scores using dense einsum, gathers at given topk_indices, - then backpropagates. Uses fp32 throughout for the reference. - - Returns: - grad_q, grad_w, grad_k - """ - q = q.clone().float().requires_grad_(True) - k = k.clone().float().requires_grad_(True) - weights = weights.clone().float().requires_grad_(True) - - # Dense forward: q @ k^T -> [sq, b, h, sk] - index_scores = torch.einsum("sbhd,tbd->sbht", q, k) - index_scores = torch.relu(index_scores) - index_scores = index_scores * weights.unsqueeze(-1) - index_scores = index_scores.sum(dim=2) # [sq, b, sk] - index_scores = index_scores.transpose(0, 1) # [b, sq, sk] - - # Gather at topk positions and compute loss - valid_mask = topk_indices != -1 - safe_indices = topk_indices.clamp(min=0).to(torch.int64) - gathered_scores = torch.gather(index_scores, dim=-1, index=safe_indices) - gathered_scores = torch.where(valid_mask, gathered_scores, torch.tensor(0.0, device=q.device)) - - loss = gathered_scores.sum() - loss.backward() - - return q.grad, weights.grad, k.grad - - -@requires_cuda() -@requires_tilelang() -@pytest.mark.parametrize("seqlen_q,batch,heads,dim,compress_ratio,topk", BACKWARD_CONFIGS, ids=BACKWARD_CONFIG_IDS) -def test_indexer_backward(seqlen_q, batch, heads, dim, compress_ratio, topk): - """Compare tilelang backward gradients against dense PyTorch autograd reference. - - The main precision gap comes from bf16 GEMM (kernel) vs fp32 einsum (reference) - producing different ReLU boundaries for scores near zero. This is structural: - ~30% of Q@K^T products land near the ReLU boundary where bf16 truncation flips - the sign, causing binary gradient differences at those positions. - """ - from xorl.ops.dsv4.kernel.tilelang_indexer import v4_lighting_indexer - - q, k, weights = make_inputs(seqlen_q, batch, heads, dim, compress_ratio) - - # --- TileLang forward + backward --- - q_tl = q.clone().requires_grad_(True) - k_tl = k.clone().requires_grad_(True) - w_tl = weights.clone().requires_grad_(True) - - tl_score, tl_topk = v4_lighting_indexer(q_tl, k_tl, w_tl, compress_ratio, topk) - - valid_mask = tl_topk != -1 - tl_score_masked = torch.where(valid_mask, tl_score, torch.tensor(0.0, device=q.device)) - loss = tl_score_masked.sum() - loss.backward() - - # --- Dense fp32 reference backward using SAME topk_indices --- - ref_grad_q, ref_grad_w, ref_grad_k = ref_indexer_backward_dense(q, k, weights, compress_ratio, topk, tl_topk) - - print(f"\n[BWD] sq={seqlen_q}, b={batch}, h={heads}, cr={compress_ratio}, topk={topk}") - - for name, ref_g, tl_g in [ - ("grad_q", ref_grad_q, q_tl.grad), - ("grad_weights", ref_grad_w, w_tl.grad), - ("grad_k", ref_grad_k, k_tl.grad), - ]: - if ref_g is None or tl_g is None: - print(f" {name}: SKIPPED (None)") - continue - diff = compute_diff(ref_g.float(), tl_g.float()) - print_diff(name, diff) - - # bf16 GEMM vs fp32 einsum: ~30% of ReLU boundaries flip, causing - # binary gradient mismatches. rel_diff reflects this structural gap. - assert diff.rel_diff < 0.5, f"{name} rel_diff too large: {diff.rel_diff:.2e}" - assert diff.mean_abs_diff < 1.0, f"{name} mean_abs_diff too large: {diff.mean_abs_diff:.2e}" - - -# --------------------------------------------------------------------------- -# Masking correctness test -# --------------------------------------------------------------------------- -@requires_cuda() -@requires_tilelang() -@pytest.mark.parametrize("compress_ratio", [4, 128]) -def test_causal_mask_correctness(compress_ratio): - """Verify that the tilelang kernel correctly masks future compressed groups.""" - from xorl.ops.dsv4.kernel.tilelang_indexer_fwd import ( - _make_causal_cu_seqlens, - batched_indexer_fwd, - ) - - seqlen_q = 512 - batch = 1 - heads = 8 - dim = 128 - seqlen_kv = seqlen_q // compress_ratio - - q, k, weights = make_inputs(seqlen_q, batch, heads, dim, compress_ratio) - - cu_ks, cu_ke = _make_causal_cu_seqlens(seqlen_q, seqlen_kv, compress_ratio, q.device) - logits = batched_indexer_fwd(q, k, weights, cu_ks, cu_ke) # [batch, sq, sk] - - # Check that future positions are -inf - violations = 0 - total_checked = 0 - for qi in range(seqlen_q): - valid_end = (qi + 1) // compress_ratio - for ki in range(valid_end, seqlen_kv): - total_checked += 1 - if logits[0, qi, ki].item() != float("-inf"): - violations += 1 - - print(f"\n[MASK] cr={compress_ratio}: checked {total_checked} future positions, violations={violations}") - assert violations == 0, f"Found {violations} future-position violations (should be -inf)" + if (seqlen_q, batch, heads, dim, compress_ratio, topk) == FORWARD_CONFIGS[0]: + _assert_large_values() + _assert_zero_inputs() # --------------------------------------------------------------------------- # Numerical stability test: large values # --------------------------------------------------------------------------- -@requires_cuda() -@requires_tilelang() -def test_large_values(): +def _assert_large_values(): """Test with large input values to check for overflow/underflow.""" from xorl.ops.dsv4.kernel.tilelang_indexer_fwd import ( _make_causal_cu_seqlens, @@ -454,9 +236,7 @@ def test_large_values(): # --------------------------------------------------------------------------- # Zero input test # --------------------------------------------------------------------------- -@requires_cuda() -@requires_tilelang() -def test_zero_inputs(): +def _assert_zero_inputs(): """Test that zero inputs produce zero scores.""" from xorl.ops.dsv4.kernel.tilelang_indexer_fwd import ( _make_causal_cu_seqlens, @@ -476,115 +256,3 @@ def test_zero_inputs(): valid_mask = tl_scores != float("-inf") valid_scores = tl_scores[valid_mask] assert (valid_scores == 0).all(), f"Expected all zeros for valid positions, got max={valid_scores.max():.2e}" - - -# --------------------------------------------------------------------------- -# V4 real-world config test -# --------------------------------------------------------------------------- -@requires_cuda() -@requires_tilelang() -@pytest.mark.parametrize( - "seqlen_q", - [256, 512, 1024, 2048], - ids=["sq256", "sq512", "sq1024", "sq2048"], -) -def test_v4_real_config(seqlen_q): - """Test with V4's actual indexer configuration: heads=64, dim=128, topk=512, cr=4.""" - from xorl.ops.dsv4.kernel.tilelang_indexer import v4_lighting_indexer - from xorl.ops.dsv4.kernel.tilelang_indexer_fwd import ( - _make_causal_cu_seqlens, - batched_indexer_fwd, - ) - - batch = 1 - heads = 64 - dim = 128 - compress_ratio = 4 - topk = min(512, seqlen_q // compress_ratio) - seqlen_kv = seqlen_q // compress_ratio - - q, k, weights = make_inputs(seqlen_q, batch, heads, dim, compress_ratio) - - # Forward scores comparison - ref_scores = ref_compute_index_scores(q, weights, k) - ref_scores = ref_apply_causal_mask(ref_scores, compress_ratio) - - cu_ks, cu_ke = _make_causal_cu_seqlens(seqlen_q, seqlen_kv, compress_ratio, q.device) - tl_scores = batched_indexer_fwd(q, k, weights, cu_ks, cu_ke) - - valid_mask = ref_scores != float("-inf") - diff = compute_diff(ref_scores[valid_mask], tl_scores[valid_mask]) - print(f"\n[V4-REAL] sq={seqlen_q}, h=64, d=128, cr=4, topk={topk}") - print_diff("logits", diff) - - assert diff.rel_diff < 1e-3, f"rel_diff too large: {diff.rel_diff:.2e}" - - # Topk self-consistency: verify selected scores >= non-selected scores - tl_score, tl_topk = v4_lighting_indexer(q, k, weights, compress_ratio, topk) - violations = 0 - total = 0 - for qi in range(seqlen_q): - selected = set(tl_topk[0, qi].cpu().tolist()) - {-1} - if not selected: - continue - min_sel = tl_scores[0, qi, list(selected)].min().item() - valid_end = (qi + 1) // compress_ratio - for ki in range(valid_end): - if ki not in selected: - total += 1 - if tl_scores[0, qi, ki].item() > min_sel + 1e-5: - violations += 1 - - print(f" topk self-consistency: violations={violations}/{total}") - if total > 0: - assert violations / total < 0.05, f"topk violation rate too high: {violations}/{total}" - - -# --------------------------------------------------------------------------- -# Comprehensive diff summary (not a test, useful for manual inspection) -# --------------------------------------------------------------------------- -@requires_cuda() -@requires_tilelang() -def test_diff_summary(): - """Print a comprehensive diff summary across all configurations.""" - from xorl.ops.dsv4.kernel.tilelang_indexer_fwd import ( - _make_causal_cu_seqlens, - batched_indexer_fwd, - ) - - configs = [ - (128, 1, 8, 128, 4), - (512, 1, 16, 128, 4), - (512, 2, 64, 128, 4), - (1024, 1, 64, 128, 4), - (2048, 1, 8, 128, 128), - ] - - print("\n" + "=" * 90) - print(f"{'Config':<40} {'rel_diff':>10} {'max_abs':>10} {'mean_abs':>10} {'p99':>10}") - print("=" * 90) - - for seqlen_q, batch, heads, dim, compress_ratio in configs: - seqlen_kv = seqlen_q // compress_ratio - - q, k, weights = make_inputs(seqlen_q, batch, heads, dim, compress_ratio) - ref_scores = ref_compute_index_scores(q, weights, k) - ref_scores = ref_apply_causal_mask(ref_scores, compress_ratio) - - cu_ks, cu_ke = _make_causal_cu_seqlens(seqlen_q, seqlen_kv, compress_ratio, q.device) - tl_scores = batched_indexer_fwd(q, k, weights, cu_ks, cu_ke) - - valid_mask = ref_scores != float("-inf") - diff = compute_diff(ref_scores[valid_mask], tl_scores[valid_mask]) - - label = f"sq{seqlen_q}_b{batch}_h{heads}_cr{compress_ratio}" - print( - f"{label:<40} {diff.rel_diff:>10.2e} {diff.max_abs_diff:>10.2e} " - f"{diff.mean_abs_diff:>10.2e} {diff.p99_abs_diff:>10.2e}" - ) - - print("=" * 90) - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/ops/dsv4/test_v4_tilelang_sparse_mla.py b/tests/ops/dsv4/test_v4_tilelang_sparse_mla.py index eaf2f50e..3bd7d896 100644 --- a/tests/ops/dsv4/test_v4_tilelang_sparse_mla.py +++ b/tests/ops/dsv4/test_v4_tilelang_sparse_mla.py @@ -15,7 +15,7 @@ - Different sequence lengths: 128, 256, 512, 1024, 2048 - Different batch sizes: 1, 2 - Different topk: 64, 128, 256, 512 - - attn_sink values: zero, positive, negative, mixed + - attn_sink values: zero boundary and mixed-sign random - Edge cases: all indices valid, some indices -1 """ @@ -140,12 +140,8 @@ def make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk, device="cuda", sink_ attn_sink = torch.randn(heads, device=device, dtype=torch.float32) elif sink_mode == "zero": attn_sink = torch.zeros(heads, device=device, dtype=torch.float32) - elif sink_mode == "positive": - attn_sink = torch.rand(heads, device=device, dtype=torch.float32) * 2 - elif sink_mode == "negative": - attn_sink = -torch.rand(heads, device=device, dtype=torch.float32) * 2 else: - attn_sink = torch.randn(heads, device=device, dtype=torch.float32) + raise ValueError(f"unsupported test sink mode: {sink_mode}") # Generate valid random topk indices (no duplicates per query, all valid) actual_topk = min(topk, seqlen_kv) @@ -169,20 +165,12 @@ def make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk, device="cuda", sink_ # --------------------------------------------------------------------------- FORWARD_CONFIGS = [ # (batch, seqlen, heads, dim, seqlen_kv, topk) + # Cover each compiled top-k specialization while folding batch and the + # production 64-head geometry into the same four cases. (1, 128, 8, 512, 160, 64), - (1, 256, 8, 512, 320, 128), - (1, 256, 16, 512, 320, 128), - (2, 128, 8, 512, 160, 64), - (1, 512, 8, 512, 640, 256), - (1, 512, 16, 512, 640, 128), - # V4 real config: H=64 - (1, 256, 64, 512, 320, 128), - (1, 512, 64, 512, 640, 256), + (2, 256, 8, 512, 320, 128), + (1, 512, 16, 512, 640, 256), (1, 1024, 64, 512, 1280, 512), - # Larger topk - (1, 256, 8, 512, 320, 256), - # Small topk - (1, 256, 8, 512, 320, 64), ] FORWARD_IDS = [f"b{b}_s{s}_h{h}_d{d}_kv{kv}_top{tk}" for b, s, h, d, kv, tk in FORWARD_CONFIGS] @@ -191,7 +179,7 @@ def make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk, device="cuda", sink_ @requires_cuda() @requires_tilelang() @pytest.mark.parametrize("batch,seqlen,heads,dim,seqlen_kv,topk", FORWARD_CONFIGS, ids=FORWARD_IDS) -def test_sparse_mla_forward(batch, seqlen, heads, dim, seqlen_kv, topk): +def test_sparse_mla_forward_policy(batch, seqlen, heads, dim, seqlen_kv, topk): """Compare tilelang sparse MLA forward against PyTorch reference.""" from xorl.ops.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 @@ -208,34 +196,42 @@ def test_sparse_mla_forward(batch, seqlen, heads, dim, seqlen_kv, topk): assert diff.rel_diff < 1e-3, f"rel_diff too large: {diff.rel_diff:.2e}" assert diff.max_abs_diff < 0.1, f"max_abs_diff too large: {diff.max_abs_diff:.2e}" + if (batch, seqlen, heads, dim, seqlen_kv, topk) == FORWARD_CONFIGS[0]: + _assert_attn_sink_policy() + # --------------------------------------------------------------------------- # attn_sink correctness tests # --------------------------------------------------------------------------- -@requires_cuda() -@requires_tilelang() -@pytest.mark.parametrize("sink_mode", ["zero", "positive", "negative", "random"]) -def test_attn_sink_modes(sink_mode): - """Test that attn_sink is correctly incorporated for different value ranges.""" +def _assert_attn_sink_policy(): + """Test attn_sink reference parity and prove that the kernel does not ignore it.""" from xorl.ops.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 - batch, seqlen, heads, dim, seqlen_kv, topk = 1, 256, 8, 512, 320, 128 - q, kv, attn_sink, topk_idxs = make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk, sink_mode=sink_mode) - sm_scale = (1.0 / dim) ** 0.5 - - ref_o = ref_dense_attn(q, kv, attn_sink, topk_idxs, sm_scale) - tl_o, _ = sparse_mqa_fwd_interface(q, kv, attn_sink, topk_idxs, sm_scale=sm_scale) + # Zero is the boundary; random contains both signs. Positive-only and + # negative-only tensors use the identical arithmetic, while the separate + # effect test proves the sink is not ignored. + for sink_mode in ("zero", "random"): + batch, seqlen, heads, dim, seqlen_kv, topk = 1, 256, 8, 512, 320, 128 + q, kv, attn_sink, topk_idxs = make_inputs( + batch, + seqlen, + heads, + dim, + seqlen_kv, + topk, + sink_mode=sink_mode, + ) + sm_scale = (1.0 / dim) ** 0.5 - diff = compute_diff(ref_o.float(), tl_o.float()) - print(f"\n[SINK-{sink_mode}]") - print_diff("output", diff) + ref_o = ref_dense_attn(q, kv, attn_sink, topk_idxs, sm_scale) + tl_o, _ = sparse_mqa_fwd_interface(q, kv, attn_sink, topk_idxs, sm_scale=sm_scale) + diff = compute_diff(ref_o.float(), tl_o.float()) + assert diff.rel_diff < 1e-3, f"rel_diff too large for sink_mode={sink_mode}: {diff.rel_diff:.2e}" - assert diff.rel_diff < 1e-3, f"rel_diff too large for sink_mode={sink_mode}: {diff.rel_diff:.2e}" + _assert_attn_sink_changes_output(sparse_mqa_fwd_interface) -@requires_cuda() -@requires_tilelang() -def test_attn_sink_effect(): +def _assert_attn_sink_changes_output(sparse_mqa_fwd_interface): """Verify attn_sink actually changes output (not ignored).""" from xorl.ops.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 @@ -262,7 +258,6 @@ def test_attn_sink_effect(): BACKWARD_CONFIGS = [ # (batch, seqlen, heads, dim, seqlen_kv, topk) (1, 128, 8, 512, 160, 64), - (1, 256, 16, 512, 320, 128), (2, 128, 8, 512, 160, 64), (1, 256, 64, 512, 320, 128), (1, 512, 8, 512, 640, 256), @@ -312,7 +307,7 @@ def ref_dense_attn_with_grad(q, kv, attn_sink, topk_idxs, sm_scale): @requires_cuda() @requires_tilelang() @pytest.mark.parametrize("batch,seqlen,heads,dim,seqlen_kv,topk", BACKWARD_CONFIGS, ids=BACKWARD_IDS) -def test_sparse_mla_backward(batch, seqlen, heads, dim, seqlen_kv, topk): +def test_sparse_mla_backward_policy(batch, seqlen, heads, dim, seqlen_kv, topk): """Compare tilelang backward gradients against PyTorch autograd reference.""" from xorl.ops.dsv4.attention_core import sparse_attn_tilelang # noqa: PLC0415 @@ -395,8 +390,8 @@ def test_sparse_mla_backward_deterministic_dkv(): @requires_cuda() @requires_tilelang() -def test_sparse_mla_backward_partial_invalid_indices(): - """Backward must ignore -1 sparse slots without touching invalid dKV rows.""" +def test_sparse_mla_partial_invalid_indices_forward_backward_policy(): + """Forward and backward must ignore -1 sparse slots without touching invalid dKV rows.""" from xorl.ops.dsv4.attention_core import sparse_attn_tilelang # noqa: PLC0415 batch, seqlen, heads, dim, seqlen_kv, topk = 1, 128, 8, 512, 160, 64 @@ -424,13 +419,13 @@ def test_sparse_mla_backward_partial_invalid_indices(): print_diff(name, diff) assert diff.rel_diff < 0.05, f"{name} rel_diff too large with invalid indices: {diff.rel_diff:.2e}" + _assert_partial_invalid_forward_interface() + # --------------------------------------------------------------------------- # Index masking test: some indices are -1 # --------------------------------------------------------------------------- -@requires_cuda() -@requires_tilelang() -def test_partial_invalid_indices(): +def _assert_partial_invalid_forward_interface(): """Test with some indices set to -1 (invalid).""" from xorl.ops.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 @@ -450,44 +445,3 @@ def test_partial_invalid_indices(): assert diff.rel_diff < 1e-3, f"rel_diff too large with partial invalid: {diff.rel_diff:.2e}" assert not torch.isnan(tl_o).any(), "NaN in output with partial invalid indices" - - -# --------------------------------------------------------------------------- -# Comprehensive diff summary -# --------------------------------------------------------------------------- -@requires_cuda() -@requires_tilelang() -def test_diff_summary(): - """Print a comprehensive diff summary across all forward configs.""" - from xorl.ops.dsv4.kernel.tilelang_sparse_mla_fwd import sparse_mqa_fwd_interface # noqa: PLC0415 - - configs = [ - (1, 128, 8, 512, 160, 64), - (1, 256, 16, 512, 320, 128), - (1, 256, 64, 512, 320, 128), - (1, 512, 64, 512, 640, 256), - (1, 1024, 64, 512, 1280, 512), - ] - - print("\n" + "=" * 100) - print(f"{'Config':<45} {'rel_diff':>10} {'max_abs':>10} {'mean_abs':>10} {'p99':>10}") - print("=" * 100) - - for batch, seqlen, heads, dim, seqlen_kv, topk in configs: - q, kv, attn_sink, topk_idxs = make_inputs(batch, seqlen, heads, dim, seqlen_kv, topk) - sm_scale = (1.0 / dim) ** 0.5 - - ref_o = ref_dense_attn(q, kv, attn_sink, topk_idxs, sm_scale) - tl_o, _ = sparse_mqa_fwd_interface(q, kv, attn_sink, topk_idxs, sm_scale=sm_scale) - - diff = compute_diff(ref_o.float(), tl_o.float()) - label = f"b{batch}_s{seqlen}_h{heads}_d{dim}_kv{seqlen_kv}_top{topk}" - print( - f"{label:<45} {diff.rel_diff:>10.2e} {diff.max_abs_diff:>10.2e} {diff.mean_abs_diff:>10.2e} {diff.p99_abs_diff:>10.2e}" - ) - - print("=" * 100) - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/ops/loss/test_causallm_z_loss.py b/tests/ops/loss/test_causallm_z_loss.py index b5616a70..29dbd388 100644 --- a/tests/ops/loss/test_causallm_z_loss.py +++ b/tests/ops/loss/test_causallm_z_loss.py @@ -46,8 +46,10 @@ def inputs(): return hidden_states, weight, labels -def test_eager_z_loss_matches_reference(inputs): +def test_eager_z_loss_matches_reference_and_backpropagates(inputs): hidden_states, weight, labels = inputs + hidden_states = hidden_states.detach().requires_grad_(True) + weight = weight.detach().requires_grad_(True) ce_ref = _reference_ce_loss(hidden_states, weight, labels) z_ref = _reference_z_loss(hidden_states, weight, labels) coef = 1e-3 @@ -64,32 +66,15 @@ def test_eager_z_loss_matches_reference(inputs): assert_close(out.metrics["ce_loss"], ce_ref) assert_close(out.metrics["z_loss"], z_ref) assert_close(out.loss, ce_ref + coef * z_ref) + out.loss.backward() + assert hidden_states.grad is not None and torch.isfinite(hidden_states.grad).all() + assert weight.grad is not None and torch.isfinite(weight.grad).all() - -def test_causallm_logprob_temperature_matches_scaled_logits(inputs): - hidden_states, weight, labels = inputs - temperature = 0.7 - - out = causallm_loss_function( - hidden_states=hidden_states, - weight=weight, - labels=labels, - ce_mode="eager", - return_per_token=True, - logprob_temperature=temperature, - ) - - logits = (hidden_states.reshape(-1, hidden_states.size(-1)) @ weight.t()).float() / temperature - expected_ce = torch.nn.functional.cross_entropy( - logits, - labels.reshape(-1), - reduction="none", - ignore_index=-100, - ).view_as(labels) - assert_close(out.per_token_logprobs, -expected_ce) + _assert_eager_no_z_loss_when_coef_zero(inputs) + _assert_tp_path_rejects_z_loss() -def test_eager_no_z_loss_when_coef_zero(inputs): +def _assert_eager_no_z_loss_when_coef_zero(inputs): hidden_states, weight, labels = inputs ce_ref = _reference_ce_loss(hidden_states, weight, labels) @@ -105,43 +90,6 @@ def test_eager_no_z_loss_when_coef_zero(inputs): assert_close(out.loss, ce_ref) -def test_eager_z_loss_grad_flows(inputs): - hidden_states, weight, labels = inputs - hidden_states = hidden_states.detach().requires_grad_(True) - weight = weight.detach().requires_grad_(True) - - out = causallm_loss_function( - hidden_states=hidden_states, - weight=weight, - labels=labels, - ce_mode="eager", - z_loss_coef=1.0, - ) - out.loss.backward() - - assert hidden_states.grad is not None and torch.isfinite(hidden_states.grad).all() - assert weight.grad is not None and torch.isfinite(weight.grad).all() - - -def test_eager_z_loss_zero_when_logits_centered(): - """If logits are all zeros, logsumexp = log(V) (constant) so Z-loss = log(V)^2.""" - torch.manual_seed(1) - B, S, V, H = 1, 3, 8, 4 - hidden_states = torch.zeros(B, S, H) - weight = torch.zeros(V, H) - labels = torch.zeros(B, S, dtype=torch.long) - - out = causallm_loss_function( - hidden_states=hidden_states, - weight=weight, - labels=labels, - ce_mode="eager", - z_loss_coef=1.0, - ) - expected_z = torch.tensor(float(torch.log(torch.tensor(V)).item() ** 2)) - assert_close(out.metrics["z_loss"], expected_z) - - @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="compiled CE+LSE^2 path requires CUDA") def test_compiled_z_loss_matches_eager(inputs): @@ -173,7 +121,7 @@ def test_compiled_z_loss_matches_eager(inputs): assert_close(out_compiled.loss, out_eager.loss) -def test_tp_path_rejects_z_loss(): +def _assert_tp_path_rejects_z_loss(): """TP path must error out clearly when Z-loss is requested.""" torch.manual_seed(2) B, S, V, H = 1, 2, 8, 4 diff --git a/tests/ops/loss/test_drgrpo_loss.py b/tests/ops/loss/test_drgrpo_loss.py index 1876906e..47de3897 100644 --- a/tests/ops/loss/test_drgrpo_loss.py +++ b/tests/ops/loss/test_drgrpo_loss.py @@ -85,8 +85,8 @@ class TestDRGRPOLoss: Then update the assert_close(...) calls with the new values. """ - def test_forward(self, inputs): - """Forward pass produces expected loss value (regression test).""" + def test_forward_backward_and_metrics(self, inputs): + """Forward value, gradients, and metric schema form one numerical contract.""" d = inputs hidden_states = d["hidden_states"].clone().requires_grad_(True) @@ -109,8 +109,27 @@ def test_forward(self, inputs): # Default loss_reducer is TokenPartial(scale=mask.sum()); fixture has 4 # active tokens of 8 → 2× the previous numel-scaled value. assert_close(output.loss, torch.tensor(0.727356)) + expected_keys = { + "loss/ratio/mean", + "loss/kl_policy/mean", + "loss/clip/clipped_ratio/mean", + "loss/clip/high_fraction", + "loss/clip/low_fraction", + "loss/kl_ref/mean", + } + assert expected_keys <= output.metrics.keys() + + output.loss.backward() + assert hidden_states.grad is not None + assert hidden_states.grad.isfinite().all() + assert_close(hidden_states.grad.norm(), torch.tensor(3.028616)) - def test_logprob_temperature_changes_behavior_k3(self): + self._assert_zero_loss_boundaries(inputs) + self._assert_positive_advantages_encourage_high_prob(inputs) + self._assert_kl_penalty_requires_reference_and_affects_loss(inputs) + self._assert_logprob_temperature_changes_behavior_k3() + + def _assert_logprob_temperature_changes_behavior_k3(self): hidden_states = torch.tensor([[[1.0, -0.5], [0.25, 0.75]]]) weight = torch.tensor([[0.5, -1.0], [-0.25, 0.75], [1.0, 0.5]]) labels = torch.tensor([[0, 2]]) @@ -146,33 +165,8 @@ def test_logprob_temperature_changes_behavior_k3(self): torch.testing.assert_close(raw.metrics["loss/kl_policy/mean"], torch.tensor(0.0), atol=1e-6, rtol=0.0) assert behavior.metrics["loss/kl_policy/mean"].abs() > 1e-6 - def test_backward(self, inputs): - """Backward pass produces expected gradient norm (regression test).""" - d = inputs - hidden_states = d["hidden_states"].clone().requires_grad_(True) - - output = drgrpo_loss_function( - hidden_states=hidden_states, - weight=d["weight"], - labels=d["labels_with_mask"], - old_logprobs=d["old_logprobs"], - advantages=d["advantages"], - ref_logprobs=d["ref_logprobs"], - ignore_index=d["ignore_index"], - clip_low=0.2, - clip_high=0.2, - beta=0.1, - ) - - output.loss.backward() - assert hidden_states.grad is not None - assert hidden_states.grad.isfinite().all() - # Regression test: expected value computed with seed=42 fixture inputs. - # Loss scaled 2× under new TokenPartial(scale=mask.sum()) default → grad scaled 2×. - assert_close(hidden_states.grad.norm(), torch.tensor(3.028616)) - - def test_zero_advantages(self, inputs): - """Zero advantages produce finite (near-zero) loss.""" + def _assert_zero_loss_boundaries(self, inputs): + """Zero advantages, no trainable labels, and empty sequences remain finite and zero.""" d = inputs advantages = torch.zeros_like(d["advantages"]) @@ -187,15 +181,9 @@ def test_zero_advantages(self, inputs): ) assert output.loss.isfinite() - # With zero advantages, policy gradient loss should be zero assert output.loss.abs() < 1e-5 - def test_all_ignored_labels(self, inputs): - """Loss should be finite (zero) when all labels are ignored (no trainable tokens).""" - d = inputs - # Set all labels to ignore_index all_ignored = torch.full_like(d["labels"], d["ignore_index"]) - output = drgrpo_loss_function( hidden_states=d["hidden_states"], weight=d["weight"], @@ -209,8 +197,6 @@ def test_all_ignored_labels(self, inputs): assert output.loss.isfinite() assert output.loss == 0.0 - def test_empty_sequence(self): - """Loss should be zero when sequence length is 0.""" B, V, H = 2, 10, 16 hidden_states = torch.empty(B, 0, H) weight = torch.randn(V, H) @@ -230,23 +216,7 @@ def test_empty_sequence(self): assert output.loss.isfinite() assert output.loss == 0.0 - def test_requires_ref_logprobs_when_beta_positive(self, inputs): - """ValueError raised when beta > 0 but ref_logprobs is None.""" - d = inputs - - with pytest.raises(ValueError, match="ref_logprobs required"): - drgrpo_loss_function( - hidden_states=d["hidden_states"], - weight=d["weight"], - labels=d["labels_with_mask"], - old_logprobs=d["old_logprobs"], - advantages=d["advantages"], - ref_logprobs=None, - ignore_index=d["ignore_index"], - beta=0.1, - ) - - def test_positive_advantages_encourage_high_prob(self, inputs): + def _assert_positive_advantages_encourage_high_prob(self, inputs): """With positive advantages, higher target probability yields lower loss.""" d = inputs B, S, V, H = d["B"], d["S"], d["V"], d["H"] @@ -294,10 +264,22 @@ def test_positive_advantages_encourage_high_prob(self, inputs): # Higher probability should yield lower (more negative) loss assert loss_high.loss < loss_low.loss - def test_kl_penalty_affects_loss(self, inputs): - """KL penalty modifies loss when beta > 0.""" + def _assert_kl_penalty_requires_reference_and_affects_loss(self, inputs): + """Positive KL weight requires reference logprobs and changes the loss.""" d = inputs + with pytest.raises(ValueError, match="ref_logprobs required"): + drgrpo_loss_function( + hidden_states=d["hidden_states"], + weight=d["weight"], + labels=d["labels_with_mask"], + old_logprobs=d["old_logprobs"], + advantages=d["advantages"], + ref_logprobs=None, + ignore_index=d["ignore_index"], + beta=0.1, + ) + loss_no_kl = drgrpo_loss_function( hidden_states=d["hidden_states"], weight=d["weight"], @@ -324,35 +306,6 @@ def test_kl_penalty_affects_loss(self, inputs): # KL metrics should be present when beta > 0 assert "loss/kl_ref/mean" in loss_with_kl.metrics - def test_metrics_present(self, inputs): - """Output includes expected metrics.""" - d = inputs - - output = drgrpo_loss_function( - hidden_states=d["hidden_states"], - weight=d["weight"], - labels=d["labels_with_mask"], - old_logprobs=d["old_logprobs"], - advantages=d["advantages"], - ref_logprobs=d["ref_logprobs"], - ignore_index=d["ignore_index"], - clip_low=0.2, - clip_high=0.2, - beta=0.1, - ) - - expected_keys = [ - "loss/ratio/mean", - "loss/kl_policy/mean", - "loss/clip/clipped_ratio/mean", - "loss/clip/high_fraction", - "loss/clip/low_fraction", - "loss/kl_ref/mean", - ] - - for key in expected_keys: - assert key in output.metrics, f"Missing metric: {key}" - def test_microbatch_composition(self, inputs): """Per-mb partial shares sum to single-batch values for both loss and metrics. diff --git a/tests/ops/loss/test_fp8_lm_head_ce.py b/tests/ops/loss/test_fp8_lm_head_ce.py index 291aaad8..0433d4b2 100644 --- a/tests/ops/loss/test_fp8_lm_head_ce.py +++ b/tests/ops/loss/test_fp8_lm_head_ce.py @@ -23,67 +23,81 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return F.linear(x, self.weight) -def test_compute_per_token_ce_uses_lm_head_module_when_provided(): +def test_compute_per_token_ce_selects_lm_head_path_locally_and_under_tp(monkeypatch): + collective_calls = _patch_identity_tp_collectives(monkeypatch) torch.manual_seed(0) - hidden = torch.randn(5, 4) - labels = torch.tensor([0, 1, -100, 2, 3]) - module_weight = torch.randn(6, 4) - raw_weight = torch.zeros_like(module_weight) + hidden = torch.randn(2, 4) + labels = torch.tensor([0, 1]) + module_weight = torch.randn(3, 4) + raw_weight = torch.randn(3, 4) lm_head = CountingHead(module_weight) - got = compute_per_token_ce( + module_ce = compute_per_token_ce( hidden, raw_weight, labels, ignore_index=-100, ce_mode="compiled", - num_chunks=2, + num_chunks=1, lm_head=lm_head, ) - expected = F.cross_entropy(F.linear(hidden, module_weight), labels, reduction="none", ignore_index=-100) - + expected_module_ce = F.cross_entropy(F.linear(hidden, module_weight), labels, reduction="none") assert lm_head.calls == 1 - torch.testing.assert_close(got, expected) - + torch.testing.assert_close(module_ce, expected_module_ce) -def test_compute_per_token_ce_lm_head_fp32_bypasses_module(): - # lm_head_fp32 takes precedence over an FP8 lm_head module: the module must - # NOT be called, and logits come from the fp32 (master) raw weight. - torch.manual_seed(0) - hidden = torch.randn(5, 4) - labels = torch.tensor([0, 1, -100, 2, 3]) - module_weight = torch.randn(6, 4) - raw_weight = torch.randn(6, 4) # distinct master weight - lm_head = CountingHead(module_weight) - - got = compute_per_token_ce( + master_ce = compute_per_token_ce( hidden, raw_weight, labels, ignore_index=-100, ce_mode="eager", - num_chunks=2, + num_chunks=1, lm_head=lm_head, lm_head_fp32=True, ) - expected = F.cross_entropy( - (hidden.float() @ raw_weight.float().t()).float(), + expected_master_ce = F.cross_entropy(hidden.float() @ raw_weight.float().t(), labels, reduction="none") + assert lm_head.calls == 1 + torch.testing.assert_close(master_ce, expected_master_ce) + + tp_module_ce = compute_per_token_ce( + hidden, + torch.zeros_like(raw_weight), + labels, + ignore_index=-100, + ce_mode="compiled", + num_chunks=1, + tp_group=object(), + lm_head=lm_head, + ) + assert lm_head.calls == 2 + assert (2, 4) not in collective_calls + torch.testing.assert_close(tp_module_ce, expected_module_ce) + + tp_master_ce = compute_per_token_ce( + hidden, + raw_weight, labels, - reduction="none", ignore_index=-100, + ce_mode="eager", + num_chunks=1, + tp_group=object(), + lm_head=lm_head, + lm_head_fp32=True, ) + assert lm_head.calls == 2 + assert torch.isfinite(tp_master_ce).all() - assert lm_head.calls == 0 # FP8 module bypassed - torch.testing.assert_close(got, expected) + _assert_logprob_temperature_threads_through_per_token_and_causallm_losses() + _assert_loss_dispatchers_lm_head_module_policy(monkeypatch) -def test_compute_per_token_ce_applies_logprob_temperature(): +def _assert_logprob_temperature_threads_through_per_token_and_causallm_losses(): torch.manual_seed(2) hidden = torch.randn(5, 4) labels = torch.tensor([0, 1, -100, 2, 3]) weight = torch.randn(6, 4) - got = compute_per_token_ce( + per_token_ce = compute_per_token_ce( hidden, weight, labels, @@ -91,67 +105,35 @@ def test_compute_per_token_ce_applies_logprob_temperature(): ce_mode="eager", logprob_temperature=0.7, ) - expected = F.cross_entropy( + expected_ce = F.cross_entropy( (hidden @ weight.t()).float() / 0.7, labels, reduction="none", ignore_index=-100, ) - torch.testing.assert_close(got, expected) - - -def test_causallm_loss_applies_logprob_temperature_to_per_token_logprobs(): - torch.manual_seed(3) - hidden = torch.randn(1, 5, 4) - labels = torch.tensor([[0, 1, -100, 2, 3]]) - weight = torch.randn(6, 4) + torch.testing.assert_close(per_token_ce, expected_ce) result = causallm_loss_function( - hidden, + hidden.unsqueeze(0), weight, - labels, + labels.unsqueeze(0), ignore_index=-100, ce_mode="eager", return_per_token=True, logprob_temperature=0.7, ) expected_ce = F.cross_entropy( - (hidden.reshape(-1, 4) @ weight.t()).float() / 0.7, - labels.reshape(-1), - reduction="none", - ignore_index=-100, - ) - - torch.testing.assert_close(result.per_token_logprobs, -expected_ce.view_as(labels)) - - -def test_compute_per_token_ce_lm_head_fp32_bypasses_module_with_tp_group(monkeypatch): - _patch_identity_tp_collectives(monkeypatch) - torch.manual_seed(4) - hidden = torch.randn(2, 4) - labels = torch.tensor([0, 1]) - module_weight = torch.randn(3, 4) - raw_weight = torch.randn(3, 4, requires_grad=True) - lm_head = CountingHead(module_weight) - - got = compute_per_token_ce( - hidden, - raw_weight, + (hidden @ weight.t()).float() / 0.7, labels, + reduction="none", ignore_index=-100, - ce_mode="eager", - num_chunks=1, - tp_group=object(), - lm_head=lm_head, - lm_head_fp32=True, ) - assert lm_head.calls == 0 # FP8 module bypassed even under TP - assert torch.isfinite(got).all() + torch.testing.assert_close(result.per_token_logprobs.squeeze(0), -expected_ce) -def test_importance_sampling_loss_threads_lm_head_module_to_ce(): +def _assert_importance_sampling_loss_threads_lm_head_module_to_ce(): torch.manual_seed(1) hidden = torch.randn(1, 5, 4) labels = torch.tensor([[0, 1, -100, 2, 3]]) @@ -196,32 +178,7 @@ def fake_all_reduce(tensor, *args, **kwargs): return calls -def test_compute_per_token_ce_uses_lm_head_module_with_tp_group(monkeypatch): - calls = _patch_identity_tp_collectives(monkeypatch) - torch.manual_seed(2) - hidden = torch.randn(2, 4) - labels = torch.tensor([0, 1]) - weight = torch.randn(3, 4, requires_grad=True) - lm_head = CountingHead(weight) - - got = compute_per_token_ce( - hidden, - torch.zeros_like(weight), - labels, - ignore_index=-100, - ce_mode="compiled", - num_chunks=1, - tp_group=object(), - lm_head=lm_head, - ) - expected = F.cross_entropy(F.linear(hidden, weight), labels, reduction="none", ignore_index=-100) - - assert lm_head.calls == 1 - assert (2, 4) not in calls - torch.testing.assert_close(got, expected) - - -def test_causallm_loss_lm_head_fp32_bypasses_module(): +def _assert_causallm_loss_lm_head_fp32_bypasses_module(): # causallm_loss_function has its OWN use_lm_head_module path (not via # compute_per_token_ce); lm_head_fp32 must bypass the FP8 module here too so # the per-token logprobs (which drive the K3 metric) come from fp32 weights. @@ -253,7 +210,10 @@ def test_causallm_loss_lm_head_fp32_bypasses_module(): torch.testing.assert_close(result.per_token_logprobs.reshape(-1), -expected_ce) -def test_causallm_loss_uses_lm_head_module_with_tp_group_and_reduces_hidden_grad(monkeypatch): +def _assert_loss_dispatchers_lm_head_module_policy(monkeypatch): + _assert_importance_sampling_loss_threads_lm_head_module_to_ce() + _assert_causallm_loss_lm_head_fp32_bypasses_module() + calls = _patch_identity_tp_collectives(monkeypatch) torch.manual_seed(3) hidden = torch.randn(1, 2, 4) diff --git a/tests/ops/loss/test_fused_linear_logprob.py b/tests/ops/loss/test_fused_linear_logprob.py index 79a45aec..6de7e8d9 100644 --- a/tests/ops/loss/test_fused_linear_logprob.py +++ b/tests/ops/loss/test_fused_linear_logprob.py @@ -43,49 +43,44 @@ def _make_inputs(N, H, V, dtype, has_bias, seed=0): return h, w, b, labels -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("has_bias", [False, True]) -@pytest.mark.parametrize("temperature", [1.0, 0.7]) -def test_forward_matches_cross_entropy(dtype, has_bias, temperature): - h, w, b, labels = _make_inputs(128, 256, 1024, dtype, has_bias) - ce = fused_selected_logprob_ce(h, w, labels, bias=b, ignore_index=-100, temperature=temperature) - ref = _ref_ce(h, w, b, labels, -100, temperature) - err = (ce - ref).abs().max().item() - tol = 5e-2 if dtype == torch.bfloat16 else 3e-3 - assert err < tol, f"forward mismatch: {err}" - # selected-token log-probability is exactly -CE - logp = -ce - assert torch.isfinite(logp[labels != -100]).all() - - -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("has_bias", [False, True]) -@pytest.mark.parametrize("temperature", [1.0, 0.7]) -def test_backward_matches_autograd(dtype, has_bias, temperature): - h, w, b, labels = _make_inputs(128, 256, 1024, dtype, has_bias) - valid = (labels != -100).sum().clamp(min=1).float() - - hf = h.clone().requires_grad_(True) - wf = w.clone().requires_grad_(True) - bf = b.clone().requires_grad_(True) if has_bias else None - ce = fused_selected_logprob_ce(hf, wf, labels, bias=bf, ignore_index=-100, temperature=temperature) - (ce.sum() / valid).backward() - - hr = h.clone().requires_grad_(True) - wr = w.clone().requires_grad_(True) - br = b.clone().requires_grad_(True) if has_bias else None - ref = _ref_ce(hr, wr, br, labels, -100, temperature) - (ref.sum() / valid).backward() - - tol = 5e-2 if dtype == torch.bfloat16 else 3e-3 - assert (hf.grad - hr.grad).abs().max().item() < tol, "grad_h mismatch" - assert (wf.grad - wr.grad).abs().max().item() < tol, "grad_W mismatch" - if has_bias: - assert (bf.grad - br.grad).abs().max().item() < tol, "grad_b mismatch" - - -def test_needs_input_grad_frozen_output_layer(): - """LoRA RL: weight (and bias) frozen -> grad_W/grad_b are None; grad_h correct.""" +def test_fused_selected_logprob_matches_eager_forward_and_backward(): + # Dtype, bias, and temperature are independent branches; pairwise cases + # cover both values without a Cartesian product. + for dtype, has_bias, temperature in ( + (torch.bfloat16, False, 1.0), + (torch.float32, True, 0.7), + ): + h, w, b, labels = _make_inputs(128, 256, 1024, dtype, has_bias) + valid = (labels != -100).sum().clamp(min=1).float() + + hf = h.clone().requires_grad_(True) + wf = w.clone().requires_grad_(True) + bf = b.clone().requires_grad_(True) if has_bias else None + ce = fused_selected_logprob_ce(hf, wf, labels, bias=bf, ignore_index=-100, temperature=temperature) + (ce.sum() / valid).backward() + + hr = h.clone().requires_grad_(True) + wr = w.clone().requires_grad_(True) + br = b.clone().requires_grad_(True) if has_bias else None + ref = _ref_ce(hr, wr, br, labels, -100, temperature) + (ref.sum() / valid).backward() + + context = f"dtype={dtype}, bias={has_bias}, temperature={temperature}" + tol = 5e-2 if dtype == torch.bfloat16 else 3e-3 + assert torch.isfinite((-ce)[labels != -100]).all() + assert (ce - ref).abs().max().item() < tol, f"forward mismatch: {context}" + assert (hf.grad - hr.grad).abs().max().item() < tol, f"grad_h mismatch: {context}" + assert (wf.grad - wr.grad).abs().max().item() < tol, f"grad_W mismatch: {context}" + if has_bias: + assert (bf.grad - br.grad).abs().max().item() < tol, f"grad_b mismatch: {context}" + + _assert_input_gradient_policy_for_frozen_output_layer() + _assert_irregular_tail_shape_matches_eager() + _assert_production_vocab_paths_are_finite_and_match_eager() + _assert_loss_dispatchers_match_eager() + + +def _assert_input_gradient_policy_for_frozen_output_layer(): h, w, b, labels = _make_inputs(128, 256, 1024, torch.bfloat16, has_bias=True) hf = h.clone().requires_grad_(True) wf = w.clone().requires_grad_(False) # frozen output head @@ -104,18 +99,13 @@ def test_needs_input_grad_frozen_output_layer(): ref.sum().backward() assert (hf.grad - hr.grad).abs().max().item() < 5e-2, "grad_h must match even with frozen head" - -def test_no_input_needs_grad_is_detached(): - """All inputs frozen -> output is detached and backward does no work.""" h, w, _, labels = _make_inputs(64, 128, 512, torch.bfloat16, has_bias=False) out = fused_selected_logprob_ce(h, w, labels, ignore_index=-100) assert not out.requires_grad -@pytest.mark.parametrize("shape", [(37, 130, 777), (1, 64, 200), (200, 512, 50000)]) -def test_irregular_shapes(shape): - N, H, V = shape - h, w, b, labels = _make_inputs(N, H, V, torch.bfloat16, has_bias=True) +def _assert_irregular_tail_shape_matches_eager(): + h, w, b, labels = _make_inputs(37, 130, 777, torch.bfloat16, has_bias=True) hf = h.clone().requires_grad_(True) wf = w.clone().requires_grad_(True) ce = fused_selected_logprob_ce(hf, wf, labels, bias=b, ignore_index=-100) @@ -124,34 +114,20 @@ def test_irregular_shapes(shape): assert (ce - ref).abs().max().item() < 5e-2 -def test_dispatch_via_compute_per_token_ce(): - """The "fused_quack" ce_mode routes through compute_per_token_ce and matches eager.""" +def _assert_loss_dispatchers_match_eager(): + from xorl.ops.loss.causallm_loss import causallm_loss_function # noqa: PLC0415 + h, w, _, labels = _make_inputs(96, 192, 800, torch.bfloat16, has_bias=False) fused = compute_per_token_ce(h, w, labels, ignore_index=-100, ce_mode="fused_quack") eager = compute_per_token_ce(h, w, labels, ignore_index=-100, ce_mode="eager") assert (fused - eager).abs().max().item() < 5e-2 - -def test_dispatch_via_causallm_loss_function(): - """ce_mode='fused_quack' routes through the chunked fused path in - ``causallm_loss_function`` and matches the eager loss.""" - from xorl.ops.loss.causallm_loss import causallm_loss_function # noqa: PLC0415 - - h, w, _, labels = _make_inputs(96, 192, 800, torch.bfloat16, has_bias=False) h3, lab3 = h.view(1, 96, 192), labels.view(1, 96) fused = causallm_loss_function(h3, w, lab3, ignore_index=-100, ce_mode="fused_quack") eager = causallm_loss_function(h3, w, lab3, ignore_index=-100, ce_mode="eager") assert torch.isfinite(fused.loss).all() assert (fused.loss - eager.loss).abs().item() < 5e-2 - -def test_quack_linear_return_per_token_dispatch_via_causallm_loss_function(): - """quack_linear should support per-token returns without changing the scalar - training path.""" - from xorl.ops.loss.causallm_loss import causallm_loss_function # noqa: PLC0415 - - h, w, _, labels = _make_inputs(96, 192, 800, torch.bfloat16, has_bias=False) - h3, lab3 = h.view(1, 96, 192), labels.view(1, 96) quack = causallm_loss_function(h3, w, lab3, ignore_index=-100, ce_mode="quack_linear", return_per_token=True) eager = causallm_loss_function(h3, w, lab3, ignore_index=-100, ce_mode="eager", return_per_token=True) @@ -160,29 +136,17 @@ def test_quack_linear_return_per_token_dispatch_via_causallm_loss_function(): assert (quack.per_token_loss - eager.per_token_loss).abs().max().item() < 5e-2 assert (quack.per_token_logprobs - eager.per_token_logprobs).abs().max().item() < 5e-2 + _assert_importance_sampling_loss_with_fused_mode() + -@pytest.mark.parametrize("vocab", [100000, 151936, 201088]) -def test_production_vocab_sizes_no_nan(vocab): +def _assert_production_vocab_paths_are_finite_and_match_eager(): """Regression: the quack CE fwd kernel launched cluster blocks whose column tile was entirely out of range at these vocab sizes ((cluster_n-1)*tile_n >= V), reducing max over zero elements (-inf) and NaN-ing the LSE for every row. 151936/201088 are the Qwen3 / GPT-OSS vocabs; tests previously only covered V <= 65536-class shapes where every cluster block owns columns.""" - N, H = 512, 1024 - h, w, b, labels = _make_inputs(N, H, vocab, torch.bfloat16, has_bias=False) - hf = h.clone().requires_grad_(True) - wf = w.clone().requires_grad_(True) - ce = fused_selected_logprob_ce(hf, wf, labels, bias=b, ignore_index=-100, chunk_size=256) - assert torch.isfinite(ce[labels != -100]).all(), "NaN/inf per-token CE at production vocab size" - ref = _ref_ce(h, w, b, labels, -100, 1.0) - assert (ce - ref).abs().max().item() < 5e-2 - ce.sum().backward() - assert torch.isfinite(hf.grad).all() and torch.isfinite(wf.grad).all() - - -def test_quack_linear_per_token_matches_eager_at_qwen_vocab(): - """quack_linear per-token dispatch at the Qwen3 vocab (the shape long-context - recipes actually run); small-vocab dispatch tests missed the cluster NaN.""" + # The Qwen integration case and the largest direct GPT-OSS case cover the + # two production boundaries that the small-vocabulary dispatcher case misses. from xorl.ops.loss.causallm_loss import causallm_loss_function # noqa: PLC0415 N, H, V = 512, 1024, 151936 @@ -194,6 +158,19 @@ def test_quack_linear_per_token_matches_eager_at_qwen_vocab(): assert (quack.loss - eager.loss).abs().item() < 5e-2 assert (quack.per_token_loss - eager.per_token_loss).abs().max().item() < 5e-2 + del h3, lab3, quack, eager, h, w, labels + + V = 201088 + h, w, b, labels = _make_inputs(N, H, V, torch.bfloat16, has_bias=False) + hf = h.clone().requires_grad_(True) + wf = w.clone().requires_grad_(True) + ce = fused_selected_logprob_ce(hf, wf, labels, bias=b, ignore_index=-100, chunk_size=256) + assert torch.isfinite(ce[labels != -100]).all() + ref = _ref_ce(h, w, b, labels, -100, 1.0) + assert (ce - ref).abs().max().item() < 5e-2 + ce.sum().backward() + assert torch.isfinite(hf.grad).all() and torch.isfinite(wf.grad).all() + def test_causallm_fused_quack_does_not_materialize_full_logits(): """Regression for the OOM bug: ``causallm_loss_function`` had no fused_quack @@ -227,7 +204,7 @@ def test_causallm_fused_quack_does_not_materialize_full_logits(): ) -def test_importance_sampling_loss_with_fused_mode(): +def _assert_importance_sampling_loss_with_fused_mode(): """End-to-end RL loss with ce_mode="fused_quack": gradient flows through the surrogate (frozen output head — the LoRA RL case).""" N, H, V = 64, 128, 1000 @@ -244,30 +221,3 @@ def test_importance_sampling_loss_with_fused_mode(): out.loss.backward() assert h.grad is not None and torch.isfinite(h.grad).all() assert w.grad is None, "frozen output head must receive no gradient" - - -def test_logits_tile_bounded_by_chunk_when_frozen(): - """Chunking bounds peak activation: with N >> chunk_size, the frozen-W peak - stays well below a full [N, V_local] logits tile (it is never materialized).""" - N, H, V = 16384, 2048, 50000 - chunk = 2048 - h = torch.randn(N, H, device="cuda", dtype=torch.bfloat16, requires_grad=True) - w = (torch.randn(V, H, device="cuda", dtype=torch.bfloat16) / (H**0.5)).requires_grad_(False) - labels = torch.randint(0, V, (N,), device="cuda") - full_tile_mb = N * V * 4 / 1024 / 1024 # full [N, V] fp32 logits - - # warmup (quack compile / caches) - fused_selected_logprob_ce(h, w, labels, ignore_index=-100, chunk_size=chunk).sum().backward() - h.grad = None - torch.cuda.synchronize() - torch.cuda.reset_peak_memory_stats() - base = torch.cuda.memory_allocated() - fused_selected_logprob_ce(h, w, labels, ignore_index=-100, chunk_size=chunk).sum().backward() - torch.cuda.synchronize() - peak_mb = (torch.cuda.max_memory_allocated() - base) / 1024 / 1024 - - # peak should be on the order of a few [chunk, V] tiles, far below the full tile - assert peak_mb < full_tile_mb / 2, ( - f"fused frozen-W activation {peak_mb:.0f} MB should be well below the full " - f"[N, V_local] logits tile {full_tile_mb:.0f} MB" - ) diff --git a/tests/ops/loss/test_importance_sampling_loss.py b/tests/ops/loss/test_importance_sampling_loss.py deleted file mode 100644 index 8597633e..00000000 --- a/tests/ops/loss/test_importance_sampling_loss.py +++ /dev/null @@ -1,145 +0,0 @@ -import pytest -import torch - -from tests.ops.loss.conftest import assert_close -from xorl.ops.loss import TokenPartial, importance_sampling_loss_function - - -_IGNORE = -100 - - -@pytest.fixture -def inputs(): - torch.manual_seed(11) - B, S, V, H = 3, 5, 12, 16 - - hidden_states = torch.randn(B, S, H) / (H**0.5) - weight = torch.randn(V, H) - labels = torch.randint(0, V, (B, S)) - - mask_pattern = torch.tensor( - [ - [1, 1, 0, 1, 0], - [1, 0, 0, 0, 0], - [1, 1, 1, 1, 1], - ], - dtype=torch.bool, - ) - labels_with_mask = labels.clone() - labels_with_mask[~mask_pattern] = _IGNORE - - return { - "B": B, - "hidden_states": hidden_states, - "weight": weight, - "labels": labels_with_mask, - "old_logprobs": torch.randn(B, S) * 0.3 - 1.5, - "advantages": torch.randn(B, S), - } - - -def _call(d, slc, *, loss_reducer=None, metric_reducer=None, **kwargs): - return importance_sampling_loss_function( - hidden_states=d["hidden_states"][slc], - weight=d["weight"], - labels=d["labels"][slc], - old_logprobs=d["old_logprobs"][slc], - advantages=d["advantages"][slc], - ignore_index=_IGNORE, - ce_mode="eager", - loss_reducer=loss_reducer, - metric_reducer=metric_reducer, - **kwargs, - ) - - -@pytest.mark.parametrize( - "extra", - [ - pytest.param({}, id="basic"), - pytest.param({"compute_kl_stats": True}, id="kl_stats"), - ], -) -def test_identity_against_legacy(inputs, extra): - d = inputs - legacy = _call(d, slice(None), **extra) - - mask = (d["labels"] != _IGNORE).float() - reducer = TokenPartial(scale=mask.sum()) - explicit = _call(d, slice(None), loss_reducer=reducer, metric_reducer=reducer, **extra) - - assert_close(explicit.loss, legacy.loss) - for key, expected in legacy.metrics.items(): - assert key in explicit.metrics - assert_close( - torch.as_tensor(explicit.metrics[key], dtype=torch.float64), - torch.as_tensor(expected, dtype=torch.float64), - ) - - -@pytest.mark.parametrize( - "extra", - [ - pytest.param({}, id="basic"), - pytest.param({"compute_kl_stats": True}, id="kl_stats"), - ], -) -def test_microbatch_composition(inputs, extra): - d = inputs - B = d["B"] - - mask = (d["labels"] != _IGNORE).float() - loss_reducer = TokenPartial(scale=mask.sum()) - metric_reducer = TokenPartial(scale=mask.sum()) - - single = _call(d, slice(None), loss_reducer=loss_reducer, metric_reducer=metric_reducer, **extra) - mbs = [ - _call(d, slice(b, b + 1), loss_reducer=loss_reducer, metric_reducer=metric_reducer, **extra) for b in range(B) - ] - - assert_close(sum(mb.loss for mb in mbs), single.loss) - - composing = {"ratio_mean", "kl_sample_train_k3", "entropy_sample"} - for key, expected in single.metrics.items(): - if key not in composing: - continue - assert_close( - torch.as_tensor(sum(mb.metrics[key] for mb in mbs), dtype=torch.float64), - torch.as_tensor(expected, dtype=torch.float64), - ) - - -def test_kl_debug_tail_metrics(inputs): - d = inputs - out = _call(d, slice(None), compute_kl_stats=True) - - valid = d["labels"] != _IGNORE - log_ratio = out.per_token_logprobs - d["old_logprobs"] - k3 = torch.exp(log_ratio) - log_ratio - 1.0 - - assert_close(torch.as_tensor(out.metrics["kl_k3_debug_max"]), k3[valid].max()) - assert_close(torch.as_tensor(out.metrics["kl_k3_debug_logratio_min"]), log_ratio[valid].min()) - assert_close(torch.as_tensor(out.metrics["kl_k3_debug_logratio_max"]), log_ratio[valid].max()) - assert_close(torch.as_tensor(out.metrics["kl_k3_debug_abs_logratio_max"]), log_ratio[valid].abs().max()) - assert out.metric_ops["kl_k3_debug_max"] == "max" - assert out.metric_ops["kl_k3_debug_logratio_min"] == "min" - assert out.metric_ops["kl_k3_debug_logratio_max"] == "max" - - -def test_logprob_temperature_drives_behavior_k3(inputs): - d = inputs - temperature = 0.7 - labels = d["labels"] - logits = (d["hidden_states"].reshape(-1, d["hidden_states"].size(-1)) @ d["weight"].t()).float() - behavior_ce = torch.nn.functional.cross_entropy( - logits / temperature, - labels.reshape(-1), - reduction="none", - ignore_index=_IGNORE, - ).view_as(labels) - d = {**d, "old_logprobs": -behavior_ce} - - out = _call(d, slice(None), compute_kl_stats=True, logprob_temperature=temperature) - - assert_close(out.per_token_logprobs, -behavior_ce) - assert_close(torch.as_tensor(out.metrics["kl_sample_train_k3"]), torch.tensor(0.0)) diff --git a/tests/ops/loss/test_opd_loss.py b/tests/ops/loss/test_opd_loss.py index 4d708ce0..5beb2611 100644 --- a/tests/ops/loss/test_opd_loss.py +++ b/tests/ops/loss/test_opd_loss.py @@ -27,8 +27,7 @@ def inputs(): return hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights -@pytest.mark.parametrize("num_chunks_case", [1, 2, 7, "n_valid_plus_one", 0]) -def test_opd_loss_matches_reference(inputs, num_chunks_case): +def _assert_opd_loss_matches_reference(inputs, num_chunks_case): hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs n_valid = int((labels != -100).sum().item()) num_chunks = n_valid + 1 if num_chunks_case == "n_valid_plus_one" else int(num_chunks_case) @@ -55,8 +54,12 @@ def test_opd_loss_matches_reference(inputs, num_chunks_case): assert out.metrics["valid_tokens"] == int((labels != -100).sum().item()) -@pytest.mark.parametrize("backend", ["streaming", "tilelang"]) -def test_opd_streaming_backends_match_reference(inputs, backend): +def _assert_opd_loss_reference_chunking(inputs): + for num_chunks_case in (0, 1, "n_valid_plus_one"): + _assert_opd_loss_matches_reference(inputs, num_chunks_case) + + +def _assert_opd_streaming_backend_matches_reference(inputs, backend): hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs hidden_states = hidden_states.detach().requires_grad_(True) weight = weight.detach().requires_grad_(True) @@ -90,7 +93,12 @@ def test_opd_streaming_backends_match_reference(inputs, backend): assert teacher_weight.grad is None -def test_opd_streaming_lowmem_matches_streaming(inputs): +def _assert_opd_streaming_backends_match_reference(inputs): + for backend in ("streaming", "tilelang"): + _assert_opd_streaming_backend_matches_reference(inputs, backend) + + +def _assert_opd_streaming_lowmem_matches_streaming(inputs): """streaming_lowmem must be loss- and gradient-identical to plain streaming. The lowmem path keeps lm-head weights in native dtype and upcasts each vocab @@ -133,7 +141,7 @@ def run(lowmem): assert_close(gh_b, gh_a) -def test_opd_streaming_backend_reads_sharded_teacher_store(inputs, tmp_path): +def _assert_opd_streaming_backend_reads_sharded_teacher_store(inputs, tmp_path): hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs model_dir = tmp_path / "teacher_model" model_dir.mkdir() @@ -161,7 +169,7 @@ def test_opd_streaming_backend_reads_sharded_teacher_store(inputs, tmp_path): assert_close(out.loss, expected) -def test_opd_loss_backward(inputs): +def _assert_opd_loss_backward(inputs): hidden_states, weight, labels, teacher_hidden_states, teacher_weight, _ = inputs hidden_states = hidden_states.detach().requires_grad_(True) weight = weight.detach().requires_grad_(True) @@ -184,7 +192,7 @@ def test_opd_loss_backward(inputs): assert teacher_weight.grad is None -def test_opd_loss_respects_token_partial_reducer(inputs): +def _assert_opd_loss_respects_token_partial_reducer(inputs): hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs out = opd_loss_function( hidden_states=hidden_states, @@ -202,7 +210,7 @@ def test_opd_loss_respects_token_partial_reducer(inputs): assert_close(out.loss, expected * n_valid) -def test_opd_loss_all_ignored_is_finite(inputs): +def _assert_opd_loss_all_ignored_is_finite(inputs): hidden_states, weight, labels, teacher_hidden_states, teacher_weight, _ = inputs labels = torch.full_like(labels, -100) hidden_states = hidden_states.to(torch.bfloat16).detach().requires_grad_(True) @@ -228,7 +236,7 @@ def test_opd_loss_all_ignored_is_finite(inputs): assert torch.count_nonzero(weight.grad) == 0 -def test_opd_loss_bf16_inputs_return_fp32_loss(inputs): +def _assert_opd_loss_bf16_inputs_return_fp32_loss(inputs): hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs hidden_states = hidden_states.to(torch.bfloat16).detach().requires_grad_(True) weight = weight.to(torch.bfloat16).detach().requires_grad_(True) @@ -253,7 +261,7 @@ def test_opd_loss_bf16_inputs_return_fp32_loss(inputs): assert weight.grad is not None and weight.grad.isfinite().all() -def test_opd_loss_return_per_token(inputs): +def _assert_opd_loss_return_per_token(inputs): hidden_states, weight, labels, teacher_hidden_states, teacher_weight, teacher_weights = inputs out = opd_loss_function( hidden_states=hidden_states, @@ -275,7 +283,7 @@ def test_opd_loss_return_per_token(inputs): assert_close(out.per_token_loss.sum() / denom, expected) -def test_opd_loss_hidden_only_mse_with_zero_kl_weight_returns_hidden_term(): +def _assert_opd_loss_hidden_only_mse_with_zero_kl_weight_returns_hidden_term(): hidden_states = torch.tensor( [[[1.0, 2.0], [3.0, 5.0], [7.0, 11.0]]], requires_grad=True, @@ -325,7 +333,7 @@ def test_opd_loss_hidden_only_mse_with_zero_kl_weight_returns_hidden_term(): assert "opd_kl" in out.metrics -def test_oprd_hidden_distance_matches_full_materialization_and_detaches_teacher(): +def _assert_oprd_hidden_distance_matches_full_materialization_and_detaches_teacher(): torch.manual_seed(23) student = torch.randn(3, 5, 7, dtype=torch.bfloat16).requires_grad_(True) teacher = torch.randn(3, 5, 7, dtype=torch.bfloat16).requires_grad_(True) @@ -341,7 +349,7 @@ def test_oprd_hidden_distance_matches_full_materialization_and_detaches_teacher( assert teacher.grad is None -def test_oprd_hidden_distance_from_fetcher_matches_full_materialization(): +def _assert_oprd_hidden_distance_from_fetcher_matches_full_materialization(): torch.manual_seed(24) student = torch.randn(3, 5, 7, dtype=torch.bfloat16).requires_grad_(True) teacher = torch.randn(3, 5, 7, dtype=torch.bfloat16).requires_grad_(True) @@ -367,3 +375,29 @@ def fetcher(start: int, end: int) -> torch.Tensor: distance.sum().backward() assert student.grad is not None and student.grad.isfinite().all() assert teacher.grad is None + + +def test_opd_numerical_backend_contract(inputs, tmp_path): + _assert_opd_loss_reference_chunking(inputs) + _assert_opd_streaming_backends_match_reference(inputs) + _assert_opd_streaming_lowmem_matches_streaming(inputs) + _assert_opd_streaming_backend_reads_sharded_teacher_store(inputs, tmp_path) + _assert_opd_gradient_and_reduction_contract(inputs) + + +def _assert_opd_gradient_and_reduction_contract(inputs): + _assert_opd_loss_backward(inputs) + _assert_opd_loss_respects_token_partial_reducer(inputs) + _assert_opd_loss_bf16_inputs_return_fp32_loss(inputs) + + +def test_opd_output_and_hidden_only_edge_contract(inputs): + _assert_opd_loss_all_ignored_is_finite(inputs) + _assert_opd_loss_return_per_token(inputs) + _assert_opd_loss_hidden_only_mse_with_zero_kl_weight_returns_hidden_term() + _assert_oprd_hidden_distance_contract() + + +def _assert_oprd_hidden_distance_contract(): + _assert_oprd_hidden_distance_matches_full_materialization_and_detaches_teacher() + _assert_oprd_hidden_distance_from_fetcher_matches_full_materialization() diff --git a/tests/ops/loss/test_opd_verl_parity.py b/tests/ops/loss/test_opd_verl_parity.py index 5b03f5b6..5c1136bc 100644 --- a/tests/ops/loss/test_opd_verl_parity.py +++ b/tests/ops/loss/test_opd_verl_parity.py @@ -16,9 +16,6 @@ import torch.nn.functional as F from xorl.ops.loss.compiled_cross_entropy import ( - compiled_forward_kl_full_function, - compiled_forward_kl_full_with_diag_function, - compiled_reverse_kl_with_diag_function, compiled_sampled_token_logprobs_function, ) from xorl.ops.loss.opd_loss import ( @@ -55,11 +52,10 @@ def _make_synthetic_inputs( # --------------------------------------------------------------------------- -def test_reverse_kl_full_default_matches_existing_behavior(): - """Default loss_mode='reverse_kl_full' produces the same loss as the legacy - direct call to compiled_reverse_kl_function (modulo metric emission).""" +def test_full_vocab_modes_diagnostics_weighting_clamp_and_dispatch_policy(): + """Full-vocab modes expose one stable policy across diagnostics and weighting branches.""" sh, sw, th, tw, labels = _make_synthetic_inputs() - result = opd_loss_function( + reverse = opd_loss_function( hidden_states=sh, weight=sw, labels=labels, @@ -67,19 +63,16 @@ def test_reverse_kl_full_default_matches_existing_behavior(): teacher_lm_head_weight=tw, ) # Diagnostic fields default to 0.0 because emit_full_vocab_diagnostics=False. - assert result.metrics["opd_teacher_entropy"] == 0.0 - assert result.metrics["opd_top1_agreement"] == 0.0 - assert result.metrics["opd_pg_clipfrac"] == 0.0 + assert reverse.metrics["opd_teacher_entropy"] == 0.0 + assert reverse.metrics["opd_top1_agreement"] == 0.0 + assert reverse.metrics["opd_pg_clipfrac"] == 0.0 # PG-mode-only metrics default to 0.0 when use_policy_gradient=False. - assert result.metrics["opd_pg_clipfrac_lower"] == 0.0 - assert result.metrics["opd_ppo_kl"] == 0.0 - assert result.metrics["valid_tokens"] > 0 - assert torch.isfinite(result.loss).item() + assert reverse.metrics["opd_pg_clipfrac_lower"] == 0.0 + assert reverse.metrics["opd_ppo_kl"] == 0.0 + assert reverse.metrics["valid_tokens"] > 0 + assert torch.isfinite(reverse.loss).item() - -def test_reverse_kl_full_with_diagnostics(): - sh, sw, th, tw, labels = _make_synthetic_inputs() - result = opd_loss_function( + reverse_diagnostics = opd_loss_function( hidden_states=sh, weight=sw, labels=labels, @@ -87,16 +80,9 @@ def test_reverse_kl_full_with_diagnostics(): teacher_lm_head_weight=tw, emit_full_vocab_diagnostics=True, ) - # Entropy is non-negative. - assert result.metrics["opd_teacher_entropy"] >= 0.0 - assert result.metrics["opd_student_entropy"] >= 0.0 - # top1_agreement is in [0, 1]. - assert 0.0 <= result.metrics["opd_top1_agreement"] <= 1.0 - - -def test_forward_kl_full_loss_matches_reference_formula(): - """forward_kl_full == sum_v p_T(v) * (log p_T(v) - log p_S(v)) over full vocab.""" - sh, sw, th, tw, labels = _make_synthetic_inputs() + assert reverse_diagnostics.metrics["opd_teacher_entropy"] >= 0.0 + assert reverse_diagnostics.metrics["opd_student_entropy"] >= 0.0 + assert 0.0 <= reverse_diagnostics.metrics["opd_top1_agreement"] <= 1.0 # Reference: eager computation. s_logits = sh @ sw.t() @@ -108,7 +94,7 @@ def test_forward_kl_full_loss_matches_reference_formula(): ref_token_kl = ref_token_kl * valid ref_mean = ref_token_kl[labels != -100].mean().item() - result = opd_loss_function( + forward = opd_loss_function( hidden_states=sh, weight=sw, labels=labels, @@ -116,12 +102,9 @@ def test_forward_kl_full_loss_matches_reference_formula(): teacher_lm_head_weight=tw, loss_mode=LOSS_MODE_FORWARD_KL_FULL, ) - assert result.metrics["opd_kl"] == pytest.approx(ref_mean, rel=1e-5, abs=1e-5) - + assert forward.metrics["opd_kl"] == pytest.approx(ref_mean, rel=1e-5, abs=1e-5) -def test_forward_kl_full_with_diagnostics(): - sh, sw, th, tw, labels = _make_synthetic_inputs() - result = opd_loss_function( + forward_diagnostics = opd_loss_function( hidden_states=sh, weight=sw, labels=labels, @@ -130,12 +113,9 @@ def test_forward_kl_full_with_diagnostics(): loss_mode=LOSS_MODE_FORWARD_KL_FULL, emit_full_vocab_diagnostics=True, ) - assert result.metrics["opd_teacher_entropy"] >= 0.0 - assert 0.0 <= result.metrics["opd_top1_agreement"] <= 1.0 - + assert forward_diagnostics.metrics["opd_teacher_entropy"] >= 0.0 + assert 0.0 <= forward_diagnostics.metrics["opd_top1_agreement"] <= 1.0 -def test_unsupported_loss_mode_raises(): - sh, sw, th, tw, labels = _make_synthetic_inputs() with pytest.raises(ValueError, match="forward_kl_topk"): opd_loss_function( hidden_states=sh, @@ -146,53 +126,110 @@ def test_unsupported_loss_mode_raises(): loss_mode="forward_kl_topk", ) + clamped = opd_loss_function( + hidden_states=sh, + weight=sw, + labels=labels, + teacher_hidden_states=th, + teacher_lm_head_weight=tw, + loss_max_clamp=0.1, + ) + assert clamped.metrics["opd_loss_max"] <= 0.1 + 1e-6 + assert clamped.metrics["opd_loss_min"] >= -0.1 - 1e-6 + + base = opd_loss_function( + hidden_states=sh.clone().detach().requires_grad_(True), + weight=sw.clone().detach().requires_grad_(True), + labels=labels, + teacher_hidden_states=th, + teacher_lm_head_weight=tw, + ) + scaled = opd_loss_function( + hidden_states=sh.clone().detach().requires_grad_(True), + weight=sw.clone().detach().requires_grad_(True), + labels=labels, + teacher_hidden_states=th, + teacher_lm_head_weight=tw, + use_task_rewards=True, + distillation_loss_coef=2.5, + ) + assert scaled.loss.item() == pytest.approx(base.loss.item() * 2.5, rel=1e-5) + + with_disabled_coef = opd_loss_function( + hidden_states=sh.clone().detach().requires_grad_(True), + weight=sw.clone().detach().requires_grad_(True), + labels=labels, + teacher_hidden_states=th, + teacher_lm_head_weight=tw, + use_task_rewards=False, + distillation_loss_coef=999.0, + ) + assert with_disabled_coef.loss.item() == pytest.approx(base.loss.item(), rel=1e-5) + + common_args = dict( + hidden_states=sh, + weight=sw, + labels=labels, + teacher_hidden_states=th, + teacher_lm_head_weight=tw, + ) + keys_reverse = set(opd_loss_function(**common_args).metrics) + keys_reverse_diag = set(opd_loss_function(**common_args, emit_full_vocab_diagnostics=True).metrics) + keys_forward = set(opd_loss_function(**common_args, loss_mode=LOSS_MODE_FORWARD_KL_FULL).metrics) + keys_forward_diag = set( + opd_loss_function( + **common_args, + loss_mode=LOSS_MODE_FORWARD_KL_FULL, + emit_full_vocab_diagnostics=True, + ).metrics + ) + keys_k3 = set(opd_loss_function(**common_args, loss_mode="k3").metrics) + assert keys_reverse == keys_reverse_diag == keys_forward == keys_forward_diag == keys_k3 + _assert_kl_estimator_values_straight_through_gradient_and_dispatch() + _assert_policy_gradient_mode_requires_inputs_and_emits_finite_metrics() + _assert_compiled_sampled_token_logprobs_safe_with_ignored_labels() + # --------------------------------------------------------------------------- # KL estimators (k1/k2/k3/abs/mse/low_var_kl) — byte-for-byte against VERL formula # --------------------------------------------------------------------------- -@pytest.mark.parametrize("mode", ["k1", "kl", "abs", "mse", "k2", "k3", "low_var_kl"]) -def test_kl_estimator_matches_verl_kl_penalty(mode): - """_kl_penalty_estimator(...) reproduces VERL's kl_penalty_forward formulae.""" +def _assert_kl_estimator_values_straight_through_gradient_and_dispatch(): + """Estimator modes match VERL, including k3+ gradients and OPD dispatch.""" torch.manual_seed(7) logp = torch.randn(32) * 0.5 ref = torch.randn(32) * 0.5 - out = _kl_penalty_estimator(logp, ref, mode) - - if mode in ("kl", "k1"): - expected = logp - ref - elif mode == "abs": - expected = (logp - ref).abs() - elif mode in ("mse", "k2"): - expected = 0.5 * (logp - ref).square() - elif mode in ("k3", "low_var_kl"): - kl = (ref - logp).clamp(min=-20, max=20) - expected = (kl.exp() - kl - 1).clamp(min=-10, max=10) - else: - raise AssertionError(mode) - torch.testing.assert_close(out, expected) + for mode in ("k1", "kl", "abs", "mse", "k2", "k3", "low_var_kl"): + out = _kl_penalty_estimator(logp, ref, mode) + + if mode in ("kl", "k1"): + expected = logp - ref + elif mode == "abs": + expected = (logp - ref).abs() + elif mode in ("mse", "k2"): + expected = 0.5 * (logp - ref).square() + elif mode in ("k3", "low_var_kl"): + kl = (ref - logp).clamp(min=-20, max=20) + expected = (kl.exp() - kl - 1).clamp(min=-10, max=10) + else: + raise AssertionError(mode) + torch.testing.assert_close(out, expected) - -def test_kl_estimator_plus_suffix_gives_k2_straight_through_gradient(): - """The "+" suffix preserves the forward value but routes gradient through 0.5*(logp-ref)^2.""" torch.manual_seed(11) - logp = torch.randn(8, requires_grad=True) - ref = torch.randn(8) + grad_logp = torch.randn(8, requires_grad=True) + grad_ref = torch.randn(8) - out = _kl_penalty_estimator(logp, ref, "k3+") + out = _kl_penalty_estimator(grad_logp, grad_ref, "k3+") # Forward value matches plain k3. - out_k3 = _kl_penalty_estimator(logp.detach(), ref, "k3") + out_k3 = _kl_penalty_estimator(grad_logp.detach(), grad_ref, "k3") torch.testing.assert_close(out.detach(), out_k3) # Backward gradient matches k2. out.sum().backward() - expected_grad = logp.detach() - ref # d/dlogp of 0.5*(logp-ref)^2 - torch.testing.assert_close(logp.grad, expected_grad) - + expected_grad = grad_logp.detach() - grad_ref # d/dlogp of 0.5*(logp-ref)^2 + torch.testing.assert_close(grad_logp.grad, expected_grad) -def test_estimator_loss_mode_dispatch(): - """opd_loss_function with loss_mode='k3' uses the estimator path + emits opd_abs_loss.""" sh, sw, th, tw, labels = _make_synthetic_inputs() result = opd_loss_function( hidden_states=sh, @@ -209,32 +246,12 @@ def test_estimator_loss_mode_dispatch(): assert result.metrics["opd_teacher_entropy"] == 0.0 -# --------------------------------------------------------------------------- -# Clamps -# --------------------------------------------------------------------------- - - -def test_loss_max_clamp_applied(): - """loss_max_clamp bounds the per-token loss before weighting.""" - sh, sw, th, tw, labels = _make_synthetic_inputs() - clamped = opd_loss_function( - hidden_states=sh, - weight=sw, - labels=labels, - teacher_hidden_states=th, - teacher_lm_head_weight=tw, - loss_max_clamp=0.1, - ) - assert clamped.metrics["opd_loss_max"] <= 0.1 + 1e-6 - assert clamped.metrics["opd_loss_min"] >= -0.1 - 1e-6 - - # --------------------------------------------------------------------------- # Policy-gradient mode # --------------------------------------------------------------------------- -def test_pg_mode_requires_old_logprobs(): +def _assert_policy_gradient_mode_requires_inputs_and_emits_finite_metrics(): sh, sw, th, tw, labels = _make_synthetic_inputs() with pytest.raises(ValueError, match="old_logprobs"): opd_loss_function( @@ -247,9 +264,6 @@ def test_pg_mode_requires_old_logprobs(): use_policy_gradient=True, ) - -def test_pg_mode_smoke_returns_finite_loss_and_pg_clipfrac(): - sh, sw, th, tw, labels = _make_synthetic_inputs() # Fake old logprobs ~ student's current logprob to keep ratio near 1. old_lp = torch.zeros_like(labels, dtype=torch.float32) - 2.5 result = opd_loss_function( @@ -276,146 +290,7 @@ def test_pg_mode_smoke_returns_finite_loss_and_pg_clipfrac(): assert "opd_kl" in result.metrics -# --------------------------------------------------------------------------- -# Task-reward mixing (coef scaling) -# --------------------------------------------------------------------------- - - -def test_use_task_rewards_scales_loss_by_coef(): - sh, sw, th, tw, labels = _make_synthetic_inputs() - base = opd_loss_function( - hidden_states=sh.clone().detach().requires_grad_(True), - weight=sw.clone().detach().requires_grad_(True), - labels=labels, - teacher_hidden_states=th, - teacher_lm_head_weight=tw, - ) - scaled = opd_loss_function( - hidden_states=sh.clone().detach().requires_grad_(True), - weight=sw.clone().detach().requires_grad_(True), - labels=labels, - teacher_hidden_states=th, - teacher_lm_head_weight=tw, - use_task_rewards=True, - distillation_loss_coef=2.5, - ) - assert scaled.loss.item() == pytest.approx(base.loss.item() * 2.5, rel=1e-5) - - -def test_use_task_rewards_false_ignores_coef(): - sh, sw, th, tw, labels = _make_synthetic_inputs() - base = opd_loss_function( - hidden_states=sh.clone().detach().requires_grad_(True), - weight=sw.clone().detach().requires_grad_(True), - labels=labels, - teacher_hidden_states=th, - teacher_lm_head_weight=tw, - ) - with_coef = opd_loss_function( - hidden_states=sh.clone().detach().requires_grad_(True), - weight=sw.clone().detach().requires_grad_(True), - labels=labels, - teacher_hidden_states=th, - teacher_lm_head_weight=tw, - use_task_rewards=False, - distillation_loss_coef=999.0, # ignored when use_task_rewards=False - ) - assert with_coef.loss.item() == pytest.approx(base.loss.item(), rel=1e-5) - - -# --------------------------------------------------------------------------- -# Metric dict shape invariant (always-emit, no conditional keys) -# --------------------------------------------------------------------------- - - -def test_metrics_dict_has_stable_key_set_across_loss_modes(): - """Dict-keyed all_reduce deadlocks if ranks differ on which keys are emitted; - every loss_mode must produce the same set of keys.""" - sh, sw, th, tw, labels = _make_synthetic_inputs() - common_args = dict( - hidden_states=sh, - weight=sw, - labels=labels, - teacher_hidden_states=th, - teacher_lm_head_weight=tw, - ) - keys_reverse = set(opd_loss_function(**common_args).metrics.keys()) - keys_reverse_diag = set(opd_loss_function(**common_args, emit_full_vocab_diagnostics=True).metrics.keys()) - keys_forward = set(opd_loss_function(**common_args, loss_mode=LOSS_MODE_FORWARD_KL_FULL).metrics.keys()) - keys_forward_diag = set( - opd_loss_function( - **common_args, - loss_mode=LOSS_MODE_FORWARD_KL_FULL, - emit_full_vocab_diagnostics=True, - ).metrics.keys() - ) - keys_k3 = set(opd_loss_function(**common_args, loss_mode="k3").metrics.keys()) - assert keys_reverse == keys_reverse_diag == keys_forward == keys_forward_diag == keys_k3 - - -# --------------------------------------------------------------------------- -# Backend functions (smoke) -# --------------------------------------------------------------------------- - - -def test_compiled_reverse_kl_with_diag_returns_4tuple(): - sh, sw, th, tw, labels = _make_synthetic_inputs() - sh_flat = sh.reshape(-1, sh.size(-1)) - th_flat = th.reshape(-1, th.size(-1)) - lab_flat = labels.reshape(-1) - valid = lab_flat != -100 - out = compiled_reverse_kl_with_diag_function( - student_hidden_states=sh_flat[valid], - student_weight=sw, - teacher_hidden_states=th_flat[valid], - teacher_weight=tw, - labels=lab_flat[valid], - ignore_index=-100, - ) - assert len(out) == 4 - token_kl, teacher_entropy, student_entropy, top1_agreement = out - assert token_kl.shape == teacher_entropy.shape == student_entropy.shape == top1_agreement.shape - - -def test_compiled_forward_kl_full_smoke(): - sh, sw, th, tw, labels = _make_synthetic_inputs() - sh_flat = sh.reshape(-1, sh.size(-1)) - th_flat = th.reshape(-1, th.size(-1)) - lab_flat = labels.reshape(-1) - valid = lab_flat != -100 - token_kl = compiled_forward_kl_full_function( - student_hidden_states=sh_flat[valid], - student_weight=sw, - teacher_hidden_states=th_flat[valid], - teacher_weight=tw, - labels=lab_flat[valid], - ignore_index=-100, - log_prob_min_clamp=None, - ) - assert token_kl.shape == lab_flat[valid].shape - # forward KL with full distribution is always >= 0. - assert (token_kl >= -1e-5).all() - - -def test_compiled_forward_kl_full_with_diag_returns_4tuple(): - sh, sw, th, tw, labels = _make_synthetic_inputs() - sh_flat = sh.reshape(-1, sh.size(-1)) - th_flat = th.reshape(-1, th.size(-1)) - lab_flat = labels.reshape(-1) - valid = lab_flat != -100 - out = compiled_forward_kl_full_with_diag_function( - student_hidden_states=sh_flat[valid], - student_weight=sw, - teacher_hidden_states=th_flat[valid], - teacher_weight=tw, - labels=lab_flat[valid], - ignore_index=-100, - log_prob_min_clamp=None, - ) - assert len(out) == 4 - - -def test_compiled_sampled_token_logprobs_safe_with_ignored_labels(): +def _assert_compiled_sampled_token_logprobs_safe_with_ignored_labels(): sh, sw, th, tw, labels = _make_synthetic_inputs() sh_flat = sh.reshape(-1, sh.size(-1)) th_flat = th.reshape(-1, th.size(-1)) diff --git a/tests/ops/loss/test_policy_loss.py b/tests/ops/loss/test_policy_loss.py deleted file mode 100644 index 90da8c4a..00000000 --- a/tests/ops/loss/test_policy_loss.py +++ /dev/null @@ -1,165 +0,0 @@ -import pytest -import torch - -from tests.ops.loss.conftest import assert_close -from xorl.ops.loss import TokenPartial, policy_loss_function - - -_IGNORE = -100 - - -@pytest.fixture -def inputs(): - torch.manual_seed(7) - B, S, V, H = 3, 5, 12, 16 - - hidden_states = torch.randn(B, S, H) / (H**0.5) - weight = torch.randn(V, H) - labels = torch.randint(0, V, (B, S)) - - # Non-uniform mask across rows so each mb has a different valid count. - mask_pattern = torch.tensor( - [ - [1, 1, 0, 1, 0], - [1, 0, 0, 0, 0], - [1, 1, 1, 1, 1], - ], - dtype=torch.bool, - ) - labels_with_mask = labels.clone() - labels_with_mask[~mask_pattern] = _IGNORE - - return { - "B": B, - "hidden_states": hidden_states, - "weight": weight, - "labels": labels_with_mask, - "old_logprobs": torch.randn(B, S) * 0.3 - 1.5, - "rollout_logprobs": torch.randn(B, S) * 0.3 - 1.5, - "advantages": torch.randn(B, S), - } - - -def _call(d, slc, *, loss_reducer=None, metric_reducer=None, **kwargs): - return policy_loss_function( - hidden_states=d["hidden_states"][slc], - weight=d["weight"], - labels=d["labels"][slc], - old_logprobs=d["old_logprobs"][slc], - advantages=d["advantages"][slc], - rollout_logprobs=d["rollout_logprobs"][slc], - ignore_index=_IGNORE, - ce_mode="eager", - loss_reducer=loss_reducer, - metric_reducer=metric_reducer, - **kwargs, - ) - - -@pytest.mark.parametrize( - "extra", - [ - pytest.param({}, id="vanilla_ppo"), - pytest.param({"use_tis": True}, id="tis"), - pytest.param({"icepop_beta": 1.5}, id="icepop"), - pytest.param({"compute_kl_stats": True}, id="kl_stats"), - ], -) -def test_identity_against_legacy(inputs, extra): - """``TokenPartial(scale=mask.sum())`` reproduces the legacy local-mean result.""" - d = inputs - legacy = _call(d, slice(None), **extra) - - mask = (d["labels"] != _IGNORE).float() - reducer = TokenPartial(scale=mask.sum()) - explicit = _call(d, slice(None), loss_reducer=reducer, metric_reducer=reducer, **extra) - - assert_close(explicit.loss, legacy.loss) - for key, expected in legacy.metrics.items(): - assert key in explicit.metrics - # Ratio min/max are local reductions — identical regardless of reducer. - # Mean metrics match because TokenPartial(scale=mask.sum()) ≡ legacy mean. - assert_close( - torch.as_tensor(explicit.metrics[key], dtype=torch.float64), - torch.as_tensor(expected, dtype=torch.float64), - ) - - -@pytest.mark.parametrize( - "extra", - [ - pytest.param({}, id="vanilla_ppo"), - pytest.param({"use_tis": True}, id="tis"), - pytest.param({"icepop_beta": 1.5}, id="icepop"), - pytest.param({"compute_kl_stats": True}, id="kl_stats"), - ], -) -def test_microbatch_composition(inputs, extra): - """Sum of partial shares across non-uniform mbs equals the single-batch value.""" - d = inputs - B = d["B"] - - mask = (d["labels"] != _IGNORE).float() - loss_reducer = TokenPartial(scale=mask.sum()) - metric_reducer = TokenPartial(scale=mask.sum()) - - single = _call(d, slice(None), loss_reducer=loss_reducer, metric_reducer=metric_reducer, **extra) - mbs = [ - _call(d, slice(b, b + 1), loss_reducer=loss_reducer, metric_reducer=metric_reducer, **extra) for b in range(B) - ] - - assert_close(sum(mb.loss for mb in mbs), single.loss) - - # Reducer-routed metrics should compose under sum. - composing = { - "pg_clipfrac", - "icepop_maskfrac", - "tis_mean", - "tis_clipfrac", - "kl_sample_train_k3", - "entropy_sample", - "ratio_mean", - } - for key, expected in single.metrics.items(): - if key not in composing: - continue - assert_close( - torch.as_tensor(sum(mb.metrics[key] for mb in mbs), dtype=torch.float64), - torch.as_tensor(expected, dtype=torch.float64), - ) - - -def test_logprob_temperature_drives_behavior_k3(inputs): - d = inputs - temperature = 0.7 - labels = d["labels"] - logits = (d["hidden_states"].reshape(-1, d["hidden_states"].size(-1)) @ d["weight"].t()).float() - behavior_ce = torch.nn.functional.cross_entropy( - logits / temperature, - labels.reshape(-1), - reduction="none", - ignore_index=_IGNORE, - ).view_as(labels) - d = {**d, "old_logprobs": -behavior_ce} - - out = _call(d, slice(None), compute_kl_stats=True, logprob_temperature=temperature) - - assert_close(out.per_token_logprobs, -behavior_ce) - assert_close(torch.as_tensor(out.metrics["kl_sample_train_k3"]), torch.tensor(0.0)) - - -def test_kl_debug_tail_metrics(inputs): - d = inputs - out = _call(d, slice(None), compute_kl_stats=True) - - valid = d["labels"] != _IGNORE - log_ratio = out.per_token_logprobs - d["old_logprobs"] - k3 = torch.exp(log_ratio) - log_ratio - 1.0 - - assert_close(torch.as_tensor(out.metrics["kl_k3_debug_max"]), k3[valid].max()) - assert_close(torch.as_tensor(out.metrics["kl_k3_debug_logratio_min"]), log_ratio[valid].min()) - assert_close(torch.as_tensor(out.metrics["kl_k3_debug_logratio_max"]), log_ratio[valid].max()) - assert_close(torch.as_tensor(out.metrics["kl_k3_debug_abs_logratio_max"]), log_ratio[valid].abs().max()) - assert out.metric_ops["kl_k3_debug_max"] == "max" - assert out.metric_ops["kl_k3_debug_logratio_min"] == "min" - assert out.metric_ops["kl_k3_debug_logratio_max"] == "max" diff --git a/tests/ops/loss/test_reducers.py b/tests/ops/loss/test_reducers.py deleted file mode 100644 index 77028fa0..00000000 --- a/tests/ops/loss/test_reducers.py +++ /dev/null @@ -1,197 +0,0 @@ -"""Direct tests for the Reducer abstraction. - -The contract: a reducer is a closure over a caller-supplied denominator, so its -outputs are partial shares — they sum across micro-batches (and across ranks -under all_reduce(SUM)) to the globally-correct value. -""" - -import pytest -import torch - -from tests.ops.loss.conftest import assert_close -from xorl.ops.loss import SequencePartial, TokenPartial - - -@pytest.mark.parametrize( - "scale_fn", - [ - pytest.param(lambda mask: mask.sum(), id="active_count"), - pytest.param(lambda mask: torch.tensor(float(mask.numel())), id="numel"), - pytest.param(lambda mask: torch.tensor(1.0), id="ones_raw_sum"), - ], -) -def test_token_partial_shares_sum(scale_fn): - """Per-microbatch TokenPartial shares sum to the single-batch result.""" - torch.manual_seed(0) - B, S = 4, 6 - values = torch.randn(B, S) - mask = torch.randint(0, 2, (B, S)).float() - - reducer = TokenPartial(scale=scale_fn(mask)) - single = reducer(values, mask) - summed = sum(reducer(values[b : b + 1], mask[b : b + 1]) for b in range(B)) - - assert_close(summed, single) - - -def test_token_partial_scale_one_equals_raw_sum(): - """``TokenPartial(scale=1)`` is the raw masked sum (deferred-divide form).""" - torch.manual_seed(0) - B, S = 4, 6 - values = torch.randn(B, S) - mask = torch.randint(0, 2, (B, S)).float() - - out = TokenPartial(scale=torch.tensor(1.0))(values, mask) - assert_close(out, (values * mask).sum()) - - -def test_sequence_partial_shares_sum(): - """Per-microbatch SequencePartial shares (with sliced cu_seqlens_local) sum to the single-batch result.""" - torch.manual_seed(0) - B, S = 4, 6 - values = torch.randn(B, S) - mask = torch.randint(0, 2, (B, S)).float() - seq_lengths = mask.sum(dim=-1) - seq_count = torch.tensor(float(B)) - full_cu_seqlens = torch.arange(0, B * S + 1, S) - mb_cu_seqlens = torch.arange(0, S + 1, S) - - single = SequencePartial( - scale=seq_count, - cu_seqlens_local=full_cu_seqlens, - seq_lengths_global=seq_lengths, - )(values, mask) - summed = sum( - SequencePartial( - scale=seq_count, - cu_seqlens_local=mb_cu_seqlens, - seq_lengths_global=seq_lengths[b : b + 1], - )(values[b : b + 1], mask[b : b + 1]) - for b in range(B) - ) - - assert_close(summed, single) - - -def test_sequence_partial_scale_one_equals_sum_of_per_seq_means(): - """``SequencePartial(scale=1)`` is the deferred-outer-divide form of SequencePartial(scale=n_seqs).""" - torch.manual_seed(0) - B, S = 4, 6 - values = torch.randn(B, S) - mask = torch.randint(0, 2, (B, S)).float() - seq_lengths = mask.sum(dim=-1) - seq_count = torch.tensor(float(B)) - cu_seqlens_local = torch.arange(0, B * S + 1, S) - - deferred = SequencePartial( - scale=torch.tensor(1.0), - cu_seqlens_local=cu_seqlens_local, - seq_lengths_global=seq_lengths, - )(values, mask) - finalized = SequencePartial( - scale=seq_count, - cu_seqlens_local=cu_seqlens_local, - seq_lengths_global=seq_lengths, - )(values, mask) - - assert_close(deferred / seq_count, finalized) - - -@pytest.mark.parametrize( - "reducer", - [ - pytest.param(TokenPartial(scale=torch.tensor(0.0)), id="token"), - pytest.param( - SequencePartial( - scale=torch.tensor(0.0), - cu_seqlens_local=torch.tensor([0, 4, 8]), - seq_lengths_global=torch.zeros(2), - ), - id="sequence", - ), - ], -) -def test_empty_mask_yields_zero(reducer): - """Zero denominators clamp to 1 and produce 0, not NaN.""" - values = torch.randn(2, 4) - mask = torch.zeros(2, 4) - assert reducer(values, mask) == 0.0 - - -def test_sequence_partial_packed_row_matches_per_segment_mean(): - """One row, three packed segments described by ``cu_seqlens`` — sum of per-segment means / n_seqs.""" - torch.manual_seed(0) - # Row of 10 tokens packing three segments of lengths [3, 4, 3]. - values = torch.randn(1, 10) - mask = torch.ones(1, 10) - cu_seqlens = torch.tensor([0, 3, 7, 10]) - seg_lengths = torch.tensor([3, 4, 3]) - n_seqs = torch.tensor(float(seg_lengths.numel())) - - out = SequencePartial( - scale=n_seqs, - cu_seqlens_local=cu_seqlens, - seq_lengths_global=seg_lengths, - )(values, mask) - - flat = (values * mask).flatten() - expected = (flat[0:3].sum() / 3 + flat[3:7].sum() / 4 + flat[7:10].sum() / 3) / n_seqs - - assert_close(out, expected) - - -def test_sequence_partial_packed_cp_shares_sum(): - """Packed row split across two CP shards: per-shard partials sum to the single-batch result. - - Layout: one row of 10 tokens with three packed segments of pre-shard lengths - [3, 4, 3]. CP=2 splits the row at column 5: - - - Shard 0 covers columns [0, 5): segment 0 lives wholly here ([0, 3)), - and the first half of segment 1 ([3, 5), local length 2 of the - pre-shard 4). - - Shard 1 covers columns [5, 10): the second half of segment 1 - ([5, 7), local length 2 of 4) and segment 2 wholly ([7, 10)). - - Both shards reference the same ``seq_lengths_global`` for any segment - they touch. Segment 1's contributions from the two shards each divide - by 4, then sum to the correct full-segment mean. - """ - torch.manual_seed(0) - values = torch.randn(1, 10) - mask = torch.ones(1, 10) - seg_lengths_global = torch.tensor([3, 4, 3]) - n_seqs = torch.tensor(float(seg_lengths_global.numel())) - - full = SequencePartial( - scale=n_seqs, - cu_seqlens_local=torch.tensor([0, 3, 7, 10]), - seq_lengths_global=seg_lengths_global, - )(values, mask) - - shard_0 = SequencePartial( - scale=n_seqs, - cu_seqlens_local=torch.tensor([0, 3, 5]), - seq_lengths_global=seg_lengths_global[:2], - )(values[:, :5], mask[:, :5]) - - shard_1 = SequencePartial( - scale=n_seqs, - cu_seqlens_local=torch.tensor([0, 2, 5]), - seq_lengths_global=seg_lengths_global[1:], - )(values[:, 5:], mask[:, 5:]) - - assert_close(shard_0 + shard_1, full) - - -def test_token_partial_with_n_seqs_scale_equals_seq_mean_token_sum(): - """``TokenPartial(scale=n_seqs)`` expresses verl's seq-mean-token-sum policy.""" - torch.manual_seed(0) - B, S = 4, 6 - values = torch.randn(B, S) - mask = torch.randint(0, 2, (B, S)).float() - seq_count = torch.tensor(float(B)) - - via_token_partial = TokenPartial(scale=seq_count)(values, mask) - direct = (values * mask).sum(dim=-1).sum() / seq_count - - assert_close(via_token_partial, direct) diff --git a/tests/ops/loss/test_rl_primitives.py b/tests/ops/loss/test_rl_primitives.py deleted file mode 100644 index f941e624..00000000 --- a/tests/ops/loss/test_rl_primitives.py +++ /dev/null @@ -1,115 +0,0 @@ -import pytest -import torch - -from xorl.rl import ( - compute_gspo_kl, - compute_kl_estimate, - compute_opsm_mask, - compute_policy_clip_loss, - compute_sequence_kl, - reduce_token_or_sample_mean, -) - - -pytestmark = [pytest.mark.cpu] - - -def _slime_compute_approx_kl(log_probs, log_probs_base, kl_loss_type, importance_ratio=None): - """Reference formula from Slime's slime/utils/ppo_utils.py.""" - log_ratio = log_probs.float() - log_probs_base.float() - if kl_loss_type == "k1": - kl = log_ratio - elif kl_loss_type == "k2": - kl = log_ratio**2 / 2.0 - elif kl_loss_type in ["k3", "low_var_kl"]: - log_ratio = -log_ratio - kl = log_ratio.exp() - 1 - log_ratio - else: - raise ValueError(f"Unknown kl_loss_type: {kl_loss_type}") - if importance_ratio is not None: - kl = importance_ratio * kl - if kl_loss_type == "low_var_kl": - kl = torch.clamp(kl, min=-10, max=10) - return kl - - -@pytest.mark.parametrize("kind", ["k1", "k2", "k3", "low_var_kl"]) -def test_compute_kl_estimate_matches_slime_reference(kind): - policy = torch.tensor([[-2.0, -1.0, 12.0], [-3.0, -4.5, -25.0]]) - base = torch.tensor([[-2.5, -0.25, -3.0], [-4.0, -4.0, 3.0]]) - importance_ratio = torch.tensor([[1.0, 0.5, 2.0], [1.5, 1.0, 0.25]]) - - expected = _slime_compute_approx_kl(policy, base, kind, importance_ratio=importance_ratio) - - torch.testing.assert_close(compute_kl_estimate(policy, base, kind, importance_ratio), expected) - - -def test_compute_sequence_and_gspo_kl_match_slime_reference(): - current = torch.tensor([[-2.0, -1.0, -3.0], [-4.0, -2.5, -1.0]]) - old = torch.tensor([[-1.5, -1.25, -2.0], [-3.0, -3.5, -1.0]]) - masks = torch.tensor([[1, 1, 0], [1, 1, 1]], dtype=torch.float32) - - expected_seq = torch.tensor([(0.5 - 0.25) / 2.0, (1.0 - 1.0 + 0.0) / 3.0]) - expected_gspo = expected_seq.unsqueeze(-1).expand_as(current) - - torch.testing.assert_close(compute_sequence_kl(current, old, masks), expected_seq) - torch.testing.assert_close(compute_gspo_kl(current, old, masks), expected_gspo) - - -def test_compute_policy_clip_loss_matches_slime_reference_with_dual_clip(): - ppo_kl = torch.tensor([-0.3, 0.4, -0.1, 0.2]) - advantages = torch.tensor([1.0, 1.0, -2.0, -0.5]) - eps_clip = 0.2 - eps_clip_high = 0.25 - eps_clip_c = 3.0 - - ratio = (-ppo_kl).exp() - pg_losses1 = -ratio * advantages - pg_losses2 = -ratio.clamp(1 - eps_clip, 1 + eps_clip_high) * advantages - clip_pg_losses1 = torch.maximum(pg_losses1, pg_losses2) - pg_losses3 = -eps_clip_c * advantages - expected_loss = torch.where(advantages < 0, torch.minimum(pg_losses3, clip_pg_losses1), clip_pg_losses1) - expected_clipfrac = torch.gt(pg_losses2, pg_losses1).float() - - actual_loss, actual_clipfrac, actual_ratio = compute_policy_clip_loss( - ppo_kl, - advantages, - eps_clip, - eps_clip_high, - eps_clip_c, - ) - - torch.testing.assert_close(actual_loss, expected_loss) - torch.testing.assert_close(actual_clipfrac, expected_clipfrac) - torch.testing.assert_close(actual_ratio, ratio) - - -def test_compute_opsm_mask_matches_slime_reference(): - current = torch.tensor([[-2.0, -2.5, -2.0], [-1.0, -1.0, -1.0]]) - old = torch.tensor([[-1.0, -1.0, -2.0], [-1.2, -1.4, -1.0]]) - advantages = torch.tensor([[-0.5, 0.2, -1.0], [-0.1, -0.2, 0.3]]) - masks = torch.tensor([[1, 1, 0], [1, 1, 1]], dtype=torch.float32) - - opsm_mask, clipfrac = compute_opsm_mask(current, old, advantages, masks, delta=0.3) - - expected_mask = torch.tensor([[0.0, 1.0, 1.0], [1.0, 1.0, 1.0]]) - expected_clipfrac = torch.tensor(1.0 / 2.0) - - torch.testing.assert_close(opsm_mask, expected_mask) - torch.testing.assert_close(clipfrac, expected_clipfrac) - - -def test_reduce_token_or_sample_mean_modes_are_explicit(): - values = torch.tensor([[1.0, 3.0, 100.0], [2.0, 8.0, 10.0], [7.0, 11.0, 13.0]]) - masks = torch.tensor([[1, 1, 0], [1, 1, 1], [0, 0, 0]], dtype=torch.float32) - - torch.testing.assert_close(reduce_token_or_sample_mean(values, masks, "token_sum"), torch.tensor(24.0)) - torch.testing.assert_close(reduce_token_or_sample_mean(values, masks, "token_mean"), torch.tensor(24.0 / 5.0)) - torch.testing.assert_close( - reduce_token_or_sample_mean(values, masks, "slime_sum_of_sample_mean"), - torch.tensor(2.0 + 20.0 / 3.0), - ) - torch.testing.assert_close( - reduce_token_or_sample_mean(values, masks, "sample_mean"), - torch.tensor((2.0 + 20.0 / 3.0) / 2.0), - ) diff --git a/tests/ops/loss/test_shared_loss_contracts.py b/tests/ops/loss/test_shared_loss_contracts.py new file mode 100644 index 00000000..cb2140b6 --- /dev/null +++ b/tests/ops/loss/test_shared_loss_contracts.py @@ -0,0 +1,158 @@ +import pytest +import torch + +from tests.ops.loss.conftest import assert_close +from xorl.ops.loss import TokenPartial, importance_sampling_loss_function, policy_loss_function + + +_IGNORE = -100 +_IMPLEMENTATIONS = ("importance_sampling", "policy") + + +@pytest.fixture +def inputs(): + torch.manual_seed(17) + batch, sequence, vocab, hidden = 3, 5, 12, 16 + hidden_states = torch.randn(batch, sequence, hidden) / (hidden**0.5) + labels = torch.randint(0, vocab, (batch, sequence)) + labels[ + ~torch.tensor( + [ + [1, 1, 0, 1, 0], + [1, 0, 0, 0, 0], + [1, 1, 1, 1, 1], + ], + dtype=torch.bool, + ) + ] = _IGNORE + return { + "hidden_states": hidden_states, + "weight": torch.randn(vocab, hidden), + "labels": labels, + "old_logprobs": torch.randn(batch, sequence) * 0.3 - 1.5, + "rollout_logprobs": torch.randn(batch, sequence) * 0.3 - 1.5, + "advantages": torch.randn(batch, sequence), + } + + +def _call(implementation, inputs, slc=slice(None), *, loss_reducer=None, metric_reducer=None, **kwargs): + common = { + "hidden_states": inputs["hidden_states"][slc], + "weight": inputs["weight"], + "labels": inputs["labels"][slc], + "old_logprobs": inputs["old_logprobs"][slc], + "advantages": inputs["advantages"][slc], + "ignore_index": _IGNORE, + "ce_mode": "eager", + "loss_reducer": loss_reducer, + "metric_reducer": metric_reducer, + **kwargs, + } + if implementation == "policy": + return policy_loss_function(rollout_logprobs=inputs["rollout_logprobs"][slc], **common) + return importance_sampling_loss_function(**common) + + +def test_token_partial_denominator_composition_and_loss_identity(inputs): + cases = ( + ("importance_sampling", {}), + ("importance_sampling", {"compute_kl_stats": True}), + ("policy", {}), + ("policy", {"use_tis": True}), + ("policy", {"icepop_beta": 1.5}), + ("policy", {"compute_kl_stats": True}), + ) + reducer = TokenPartial(scale=(inputs["labels"] != _IGNORE).sum()) + for implementation, extra in cases: + legacy = _call(implementation, inputs, **extra) + explicit = _call(implementation, inputs, loss_reducer=reducer, metric_reducer=reducer, **extra) + + assert_close(explicit.loss, legacy.loss) + for key, expected in legacy.metrics.items(): + assert key in explicit.metrics + assert_close( + torch.as_tensor(explicit.metrics[key], dtype=torch.float64), + torch.as_tensor(expected, dtype=torch.float64), + ) + + microbatches = [ + _call( + implementation, + inputs, + slice(row, row + 1), + loss_reducer=reducer, + metric_reducer=reducer, + **extra, + ) + for row in range(inputs["labels"].size(0)) + ] + assert_close(sum(result.loss for result in microbatches), explicit.loss) + + composing_metrics = {"ratio_mean", "kl_sample_train_k3", "entropy_sample"} + if implementation == "policy": + composing_metrics.update({"pg_clipfrac", "icepop_maskfrac", "tis_mean", "tis_clipfrac"}) + for key in composing_metrics & explicit.metrics.keys(): + assert_close( + torch.as_tensor(sum(result.metrics[key] for result in microbatches), dtype=torch.float64), + torch.as_tensor(explicit.metrics[key], dtype=torch.float64), + ) + + torch.manual_seed(0) + batch, sequence = 4, 6 + values = torch.randn(batch, sequence) + mask = torch.randint(0, 2, (batch, sequence)).float() + + reducer = TokenPartial(scale=mask.sum()) + single = reducer(values, mask) + shares = sum(reducer(values[row : row + 1], mask[row : row + 1]) for row in range(batch)) + assert_close(shares, single) + assert_close(TokenPartial(scale=torch.tensor(1.0))(values, mask), (values * mask).sum()) + + sequence_count = torch.tensor(float(batch)) + assert_close( + TokenPartial(scale=sequence_count)(values, mask), + (values * mask).sum(dim=-1).sum() / sequence_count, + ) + assert TokenPartial(scale=torch.tensor(0.0))(torch.randn(2, 4), torch.zeros(2, 4)) == 0.0 + + +def test_behavior_k3_observability_and_temperature_policy(inputs): + for implementation in _IMPLEMENTATIONS: + out = _call(implementation, inputs, compute_kl_stats=True) + valid = inputs["labels"] != _IGNORE + log_ratio = out.per_token_logprobs - inputs["old_logprobs"] + k3 = torch.exp(log_ratio) - log_ratio - 1.0 + + assert_close(torch.as_tensor(out.metrics["kl_k3_debug_max"]), k3[valid].max()) + assert_close(torch.as_tensor(out.metrics["kl_k3_debug_logratio_min"]), log_ratio[valid].min()) + assert_close(torch.as_tensor(out.metrics["kl_k3_debug_logratio_max"]), log_ratio[valid].max()) + assert_close(torch.as_tensor(out.metrics["kl_k3_debug_abs_logratio_max"]), log_ratio[valid].abs().max()) + assert out.metric_ops["kl_k3_debug_max"] == "max" + assert out.metric_ops["kl_k3_debug_logratio_min"] == "min" + assert out.metric_ops["kl_k3_debug_logratio_max"] == "max" + + _assert_logprob_temperature_drives_behavior_k3(inputs) + + +def _assert_logprob_temperature_drives_behavior_k3(inputs): + temperature = 0.7 + labels = inputs["labels"] + logits = (inputs["hidden_states"].reshape(-1, inputs["hidden_states"].size(-1)) @ inputs["weight"].t()).float() + behavior_ce = torch.nn.functional.cross_entropy( + logits / temperature, + labels.reshape(-1), + reduction="none", + ignore_index=_IGNORE, + ).view_as(labels) + temperature_inputs = {**inputs, "old_logprobs": -behavior_ce} + + for implementation in _IMPLEMENTATIONS: + out = _call( + implementation, + temperature_inputs, + compute_kl_stats=True, + logprob_temperature=temperature, + ) + + assert_close(out.per_token_logprobs, -behavior_ce) + assert_close(torch.as_tensor(out.metrics["kl_sample_train_k3"]), torch.tensor(0.0)) diff --git a/tests/ops/loss/test_streaming_forward_kl.py b/tests/ops/loss/test_streaming_forward_kl.py index 5dbd3bd3..ad89a627 100644 --- a/tests/ops/loss/test_streaming_forward_kl.py +++ b/tests/ops/loss/test_streaming_forward_kl.py @@ -14,7 +14,6 @@ from tests.ops.loss.conftest import assert_close from xorl.ops.loss import opd_loss_function -from xorl.ops.loss.compiled_cross_entropy import compiled_forward_kl_full_function from xorl.ops.loss.opd_streaming_kl import ( streaming_forward_kl_function, streaming_forward_kl_lowmem_function, @@ -56,15 +55,6 @@ def inputs(): return student_hidden, student_weight, teacher_hidden, teacher_weight, labels -def test_forward_value_matches_dense_reference(inputs): - student_hidden, student_weight, teacher_hidden, teacher_weight, labels = inputs - got = streaming_forward_kl_function( - student_hidden, student_weight, teacher_hidden, teacher_weight, labels, vocab_chunk_size=7 - ) - expected = _dense_forward_kl(student_hidden, student_weight, teacher_hidden, teacher_weight, labels) - assert_close(got, expected) - - def test_backward_matches_dense_reference(inputs): student_hidden, student_weight, teacher_hidden, teacher_weight, labels = inputs @@ -93,42 +83,25 @@ def test_backward_matches_dense_reference(inputs): assert th_b.grad is None assert tw_b.grad is None + _assert_chunking_invariance(inputs) + _assert_ignore_index_zero_loss_and_grad(inputs) + _assert_lowmem_matches_streaming(inputs) + _assert_opd_loss_function_forward_kl_streaming_matches_compiled() -def test_matches_compiled_forward_kl_reference(inputs): - student_hidden, student_weight, teacher_hidden, teacher_weight, labels = inputs - got = streaming_forward_kl_function( - student_hidden, student_weight, teacher_hidden, teacher_weight, labels, vocab_chunk_size=7 - ) - # CPU path of the compiled reference is the eager dense implementation. - expected = compiled_forward_kl_full_function( - student_hidden_states=student_hidden, - student_weight=student_weight, - teacher_hidden_states=teacher_hidden, - teacher_weight=teacher_weight, - labels=labels, - ignore_index=IGNORE_INDEX, - lm_head_fp32=False, - teacher_lm_head_fp32=False, - ) - assert_close(got, expected) - - -@pytest.mark.parametrize("chunk", [7, 40, 40000]) -def test_chunking_invariance(inputs, chunk): + +def _assert_chunking_invariance(inputs): """Online accumulation must be identical across chunkings (multi-chunk vs one).""" student_hidden, student_weight, teacher_hidden, teacher_weight, labels = inputs sh = student_hidden.clone().requires_grad_(True) sw = student_weight.clone().requires_grad_(True) - kl = streaming_forward_kl_function(sh, sw, teacher_hidden, teacher_weight, labels, vocab_chunk_size=chunk) + kl = streaming_forward_kl_function(sh, sw, teacher_hidden, teacher_weight, labels, vocab_chunk_size=7) kl.sum().backward() - # Single-chunk (vocab_chunk_size >= vocab) reference. + # Single-chunk reference (chunk size equals the fixture's vocabulary size). sh_ref = student_hidden.clone().requires_grad_(True) sw_ref = student_weight.clone().requires_grad_(True) - kl_ref = streaming_forward_kl_function( - sh_ref, sw_ref, teacher_hidden, teacher_weight, labels, vocab_chunk_size=40000 - ) + kl_ref = streaming_forward_kl_function(sh_ref, sw_ref, teacher_hidden, teacher_weight, labels, vocab_chunk_size=40) kl_ref.sum().backward() assert_close(kl, kl_ref) @@ -136,7 +109,7 @@ def test_chunking_invariance(inputs, chunk): assert_close(sw.grad, sw_ref.grad) -def test_ignore_index_zero_loss_and_grad(inputs): +def _assert_ignore_index_zero_loss_and_grad(inputs): student_hidden, student_weight, teacher_hidden, teacher_weight, labels = inputs sh = student_hidden.clone().requires_grad_(True) sw = student_weight.clone().requires_grad_(True) @@ -153,7 +126,7 @@ def test_ignore_index_zero_loss_and_grad(inputs): assert kl[~ignored].abs().sum() > 0 -def test_lowmem_matches_streaming(inputs): +def _assert_lowmem_matches_streaming(inputs): """The lowmem forward-KL path must be loss- and gradient-identical to plain streaming.""" student_hidden, student_weight, teacher_hidden, teacher_weight, labels = inputs @@ -172,8 +145,7 @@ def run(fn): assert_close(gw_b, gw_a) -@pytest.mark.parametrize("backend", ["streaming", "tilelang"]) -def test_opd_loss_function_forward_kl_streaming_matches_compiled(backend): +def _assert_opd_loss_function_forward_kl_streaming_matches_compiled(): """End-to-end dispatch: forward_kl_full on the streaming backend matches the compiled (torch_compile) backend through opd_loss_function, with finite grads. """ @@ -205,15 +177,18 @@ def run(kl_backend): return out.loss, hs.grad, w.grad loss_c, gh_c, gw_c = run("torch_compile") - loss_s, gh_s, gw_s = run(backend) + for backend in ("streaming", "tilelang"): + loss_s, gh_s, gw_s = run(backend) + + assert_close(loss_s, loss_c) + assert_close(gh_s, gh_c) + assert_close(gw_s, gw_c) + assert gh_s.isfinite().all() and gw_s.isfinite().all() - assert_close(loss_s, loss_c) - assert_close(gh_s, gh_c) - assert_close(gw_s, gw_c) - assert gh_s.isfinite().all() and gw_s.isfinite().all() + _assert_opd_loss_function_forward_kl_streaming_rejects_logprob_clamp() -def test_opd_loss_function_forward_kl_streaming_rejects_logprob_clamp(): +def _assert_opd_loss_function_forward_kl_streaming_rejects_logprob_clamp(): """The streaming forward-KL backend never materializes student log-probs, so log_prob_min_clamp must fail loud rather than be silently ignored. """ diff --git a/tests/ops/loss/test_vocab_parallel_reverse_kl.py b/tests/ops/loss/test_vocab_parallel_reverse_kl.py deleted file mode 100755 index 7594b141..00000000 --- a/tests/ops/loss/test_vocab_parallel_reverse_kl.py +++ /dev/null @@ -1,117 +0,0 @@ -#!/usr/bin/env python3 -"""Multi-process correctness test for vocab_parallel_reverse_kl_function. - -Shards the lm_head along the vocab dim across `world` gloo/CPU ranks and checks -the vocab-parallel reverse-KL (kl, grad_hidden, grad_weight) against a -single-process full-vocab brute-force reference. Gradient-identity is the bar. - -Run (no GPU needed): - python -m pytest tests/ops/loss/test_vocab_parallel_reverse_kl.py -""" - -from __future__ import annotations - -import os - -import torch -import torch.distributed as dist -import torch.multiprocessing as mp -import torch.nn.functional as F - - -WORLD = 4 -N = 48 # tokens -H = 96 # hidden -V = 4 * 130 # vocab (divisible by WORLD; 520) -IGNORE = -100 - - -# The kernel computes in float32 (the real lm_head_fp32 path), so compare -# against a float32 reference; residual is float32 summation-order noise. -def _full_inputs(): - torch.manual_seed(1234) - sh = torch.randn(N, H, dtype=torch.float32) * 0.3 - th = torch.randn(N, H, dtype=torch.float32) * 0.3 - sw = torch.randn(V, H, dtype=torch.float32) * 0.05 - tw = torch.randn(V, H, dtype=torch.float32) * 0.05 - labels = torch.randint(0, V, (N,)) - labels[: N // 6] = IGNORE - return sh, th, sw, tw, labels - - -def reference(sh, sw, th, tw, labels): - sh = sh.clone().requires_grad_(True) - sw = sw.clone().requires_grad_(True) - s_logits = sh @ sw.t() - t_logits = th @ tw.t() - s_logp = F.log_softmax(s_logits, dim=-1) - t_logp = F.log_softmax(t_logits, dim=-1) - kl = (s_logp.exp() * (s_logp - t_logp)).sum(dim=-1) - valid = (labels != IGNORE).to(kl.dtype) - kl = kl * valid - kl.sum().backward() - return kl.detach(), sh.grad.detach(), sw.grad.detach() - - -def _worker(rank, world, ret): - os.environ.setdefault("MASTER_ADDR", "127.0.0.1") - os.environ.setdefault("MASTER_PORT", "29577") - dist.init_process_group("gloo", rank=rank, world_size=world) - from xorl.ops.loss.vocab_parallel_reverse_kl import vocab_parallel_reverse_kl_function # noqa: PLC0415 - - sh, th, sw, tw, labels = _full_inputs() - shard = V // world - lo, hi = rank * shard, (rank + 1) * shard - - sh_v = sh.clone().requires_grad_(True) - sw_local = sw[lo:hi].clone().requires_grad_(True) - tw_local = tw[lo:hi].clone() - - kl = vocab_parallel_reverse_kl_function( - student_hidden_states=sh_v, - student_weight_local=sw_local, - teacher_hidden_states=th, - teacher_weight_local=tw_local, - labels=labels, - ignore_index=IGNORE, - group=None, - ) - kl.sum().backward() - - if rank == 0: - kl_ref, gh_ref, gw_ref = reference(sh, sw, th, tw, labels) - # gather all ranks' local grad_weight shards - gws = [torch.zeros_like(sw_local) for _ in range(world)] - dist.all_gather(gws, sw_local.grad.contiguous()) - gw_full = torch.cat(gws, dim=0) - ret["kl"] = (kl.detach() - kl_ref).abs().max().item() - ret["gh"] = (sh_v.grad.detach() - gh_ref).abs().max().item() - ret["gw"] = (gw_full - gw_ref).abs().max().item() - ret["kl_scale"] = kl_ref.abs().max().item() - ret["gh_scale"] = gh_ref.abs().max().item() - ret["gw_scale"] = gw_ref.abs().max().item() - else: - dist.all_gather([torch.zeros_like(sw_local) for _ in range(world)], sw_local.grad.contiguous()) - dist.barrier() - dist.destroy_process_group() - - -def main(): - mgr = mp.Manager() - ret = mgr.dict() - mp.spawn(_worker, args=(WORLD, ret), nprocs=WORLD, join=True) - print(f"vocab-parallel reverse-KL vs full-vocab reference (world={WORLD}, V={V}, N={N}, H={H}):") - kl_rel = ret["kl"] / max(ret["kl_scale"], 1e-30) - gh_rel = ret["gh"] / max(ret["gh_scale"], 1e-30) - gw_rel = ret["gw"] / max(ret["gw_scale"], 1e-30) - print(f" kl max|abs|={ret['kl']:.3e} rel={kl_rel:.3e} (scale {ret['kl_scale']:.3e})") - print(f" grad_hidden max|abs|={ret['gh']:.3e} rel={gh_rel:.3e} (scale {ret['gh_scale']:.3e})") - print(f" grad_weight max|abs|={ret['gw']:.3e} rel={gw_rel:.3e} (scale {ret['gw_scale']:.3e})") - # float32 summation-order tolerance. - ok = kl_rel < 1e-3 and gh_rel < 1e-4 and gw_rel < 1e-4 - print(f" => {'PASS (matches full-vocab reference to float32 precision)' if ok else 'FAIL'}") - return 0 if ok else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/ops/test_attention.py b/tests/ops/test_attention.py index 5d5b30e1..f6ec7a4f 100644 --- a/tests/ops/test_attention.py +++ b/tests/ops/test_attention.py @@ -1,8 +1,5 @@ """Tests for attention backend functions.""" -import importlib -import sys -import types from unittest.mock import Mock, patch import pytest @@ -10,9 +7,7 @@ import xorl.models.layers.attention.backend as backend_module import xorl.models.layers.attention.backend.flash_attention as flash_module -from xorl.models.layers.attention.backend import ATTENTION_FUNCTIONS, is_flash_attention from xorl.models.layers.attention.backend.eager import eager_attention_forward -from xorl.models.layers.attention.utils import repeat_kv try: @@ -32,85 +27,29 @@ class TestAttentionBackendRegistry: """flash_attention_forward picks FA3/FA4 internally, so registration must not depend on FA3 alone: an unregistered key silently dispatches to eager.""" - def test_flash_attention_2_registry_and_mask_detection_are_consistent(self): - assert is_flash_attention("flash_attention_2") == ("flash_attention_2" in ATTENTION_FUNCTIONS) - if "flash_attention_2" in ATTENTION_FUNCTIONS: - assert ATTENTION_FUNCTIONS["flash_attention_2"] is ATTENTION_FUNCTIONS["flash_attention_3"] - - def test_flash_attention_registers_without_flash_attention_3_dependency(self, monkeypatch): - fake_flash_attn = types.ModuleType("flash_attn") - fake_flash_attn.__path__ = [] - fake_cute = types.ModuleType("flash_attn.cute") - fake_cute.flash_attn_func = Mock() - fake_cute.flash_attn_varlen_func = Mock() - - monkeypatch.setitem(sys.modules, "flash_attn_interface", None) - monkeypatch.setitem(sys.modules, "flash_attn", fake_flash_attn) - monkeypatch.setitem(sys.modules, "flash_attn.cute", fake_cute) - - try: - reloaded_flash = importlib.reload(flash_module) - reloaded_backend = importlib.reload(backend_module) - - assert not reloaded_flash.FA3_AVAILABLE - assert reloaded_flash.FA4_AVAILABLE - assert "flash_attention_3" in reloaded_backend.ATTENTION_FUNCTIONS - assert "flash_attention_4" in reloaded_backend.ATTENTION_FUNCTIONS - finally: - monkeypatch.undo() - importlib.reload(flash_module) - importlib.reload(backend_module) - - # Resolved through backend_module rather than the names imported above: - # the reload test rebinds the module's registry, and a stale reference would - # make these assert against a dict get_attention_fn no longer reads. - def test_unavailable_flash_backend_raises_instead_of_running_eager(self, monkeypatch): - monkeypatch.delitem(backend_module.ATTENTION_FUNCTIONS, "flash_attention_3", raising=False) - with pytest.raises(ImportError, match="flash_attention_3"): - backend_module.get_attention_fn("flash_attention_3") + def test_attention_backend_registry_and_resolution_policy(self, monkeypatch): + assert flash_module.FA3_AVAILABLE or flash_module.FA4_AVAILABLE + assert backend_module.ATTENTION_FUNCTIONS["flash_attention_2"] is flash_module.flash_attention_forward + assert backend_module.ATTENTION_FUNCTIONS["flash_attention_3"] is flash_module.flash_attention_forward + assert backend_module.is_flash_attention("flash_attention_2") + assert backend_module.is_flash_attention("flash_attention_3") + if flash_module.FA4_AVAILABLE: + assert "flash_attention_4" in backend_module.ATTENTION_FUNCTIONS + assert backend_module.is_flash_attention("flash_attention_4") - def test_get_attention_fn_returns_registered_and_non_flash_backends(self): + self._assert_resolution_keeps_non_flash_fallback_and_rejects_unavailable_flash(monkeypatch) + TestEagerAttentionForward()._assert_eager_attention_head_layout() + + def _assert_resolution_keeps_non_flash_fallback_and_rejects_unavailable_flash(self, monkeypatch): assert backend_module.get_attention_fn("eager") is eager_attention_forward assert backend_module.get_attention_fn("native") is backend_module.ATTENTION_FUNCTIONS["native"] # A non-flash implementation has no attention contract to drop, so it # may still default to eager; only the flash family raises. assert backend_module.get_attention_fn("flex_attention") is eager_attention_forward - -class TestRepeatKV: - """Test suite for repeat_kv function.""" - - def test_repeat_kv_shapes_values_and_gpu(self): - """repeat_kv: identity for n_rep=1, correct shapes for 2x/3x, value replication, GPU support.""" - batch, num_heads, seqlen, head_dim = 2, 4, 10, 64 - hidden_states = torch.randn(batch, num_heads, seqlen, head_dim) - - # n_rep=1: unchanged - result1 = repeat_kv(hidden_states, n_rep=1) - assert result1.shape == hidden_states.shape - assert torch.allclose(result1, hidden_states) - - # n_rep=2: doubled heads - result2 = repeat_kv(hidden_states, n_rep=2) - assert result2.shape == (batch, num_heads * 2, seqlen, head_dim) - - # n_rep=3: tripled heads - result3 = repeat_kv(hidden_states, n_rep=3) - assert result3.shape == (batch, num_heads * 3, seqlen, head_dim) - - # Value replication correctness - small = torch.arange(1 * 2 * 3 * 4).reshape(1, 2, 3, 4).float() - rep = repeat_kv(small, n_rep=2) - assert torch.allclose(rep[:, 0], rep[:, 1]) - assert torch.allclose(rep[:, 2], rep[:, 3]) - assert not torch.allclose(rep[:, 0], rep[:, 2]) - - # GPU - if torch.cuda.is_available(): - gpu_states = torch.randn(batch, num_heads, seqlen, head_dim).cuda() - gpu_result = repeat_kv(gpu_states, n_rep=2) - assert gpu_result.device.type == "cuda" - assert gpu_result.shape == (batch, num_heads * 2, seqlen, head_dim) + monkeypatch.delitem(backend_module.ATTENTION_FUNCTIONS, "flash_attention_3", raising=False) + with pytest.raises(ImportError, match="flash_attention_3"): + backend_module.get_attention_fn("flash_attention_3") class TestFlashAttentionForward: @@ -219,7 +158,9 @@ def test_flash_attention_api_behavior(self): ) assert mock_fa.call_args[1]["deterministic"] is True - def test_varlen_path_with_cu_seqlens(self): + self._assert_varlen_path_with_cu_seqlens() + + def _assert_varlen_path_with_cu_seqlens(self): """cu_seqlens kwargs trigger the varlen path.""" module = Mock() module.is_causal = True @@ -254,13 +195,16 @@ def test_varlen_path_with_cu_seqlens(self): class TestEagerAttentionForward: """Regression tests for eager attention head handling.""" - def test_eager_attention_head_layout(self): + def _assert_eager_attention_head_layout(self): """Ulysses-sync head layout handling and invalid head layout error.""" module = Mock() module.num_key_value_groups = 8 module.training = False - # Valid: local Q=4, KV=1 -> repeat=4 + # Valid: local Q=4, KV=1 -> repeat=4. Compare the complete attention + # transaction against an independent repeat_interleave reference so + # incorrect KV values cannot hide behind shape-only assertions. + torch.manual_seed(0) batch, seq, q_heads, kv_heads, head_dim = 1, 8, 4, 1, 16 query = torch.randn(batch, seq, q_heads, head_dim) key = torch.randn(batch, seq, kv_heads, head_dim) @@ -277,6 +221,17 @@ def test_eager_attention_head_layout(self): ) assert attn_output.shape == (batch, seq, q_heads, head_dim) assert attn_weights.shape == (batch, q_heads, seq, seq) + query_heads = query.transpose(1, 2) + key_heads = key.transpose(1, 2).repeat_interleave(q_heads // kv_heads, dim=1) + value_heads = value.transpose(1, 2).repeat_interleave(q_heads // kv_heads, dim=1) + expected_weights = torch.softmax( + torch.matmul(query_heads, key_heads.transpose(2, 3)) * head_dim**-0.5, + dim=-1, + dtype=torch.float32, + ).to(query.dtype) + expected_output = torch.matmul(expected_weights, value_heads).transpose(1, 2).contiguous() + torch.testing.assert_close(attn_weights, expected_weights) + torch.testing.assert_close(attn_output, expected_output) # Invalid: q_heads not divisible by kv_heads with pytest.raises(RuntimeError, match="query_heads=3 is not divisible by kv_heads=2"): @@ -307,7 +262,7 @@ def _module(): module.config._flash_attention_deterministic = False return module - def test_varlen_path_routes_through_sgl_kernel(self): + def test_sgl_page_size1_kv_cache_policy(self): total_tokens, num_heads, head_dim = 32, 8, 64 query = torch.randn(1, total_tokens, num_heads, head_dim) key = torch.randn(1, total_tokens, num_heads, head_dim) @@ -343,7 +298,11 @@ def test_varlen_path_routes_through_sgl_kernel(self): assert kwargs["causal"] is True assert result.shape == (1, total_tokens, num_heads, head_dim) - def test_single_sequence_batched_path_synthesizes_cu_seqlens(self): + self._assert_single_sequence_synthesizes_cu_seqlens() + self._assert_rejects_cross_attention_cu_seqlens() + self._assert_alternate_flash_attention_path_selection_policy() + + def _assert_single_sequence_synthesizes_cu_seqlens(self): seqlen, num_heads, head_dim = 16, 8, 64 query = torch.randn(1, seqlen, num_heads, head_dim) key = torch.randn(1, seqlen, num_heads, head_dim) @@ -367,7 +326,7 @@ def test_single_sequence_batched_path_synthesizes_cu_seqlens(self): assert kwargs["max_seqlen_q"] == seqlen assert result.shape == (1, seqlen, num_heads, head_dim) - def test_paged_kvcache_flag_pins_num_splits(self): + def _assert_alternate_flash_attention_path_selection_policy(self): total_tokens, num_heads, head_dim = 32, 8, 64 query = torch.randn(1, total_tokens, num_heads, head_dim) key = torch.randn(1, total_tokens, num_heads, head_dim) @@ -398,7 +357,10 @@ def test_paged_kvcache_flag_pins_num_splits(self): assert kwargs["num_splits"] == 1 assert kwargs["k_cache"].shape == (total_tokens, 1, num_heads, head_dim) - def test_flags_off_keep_default_varlen_path(self): + self._assert_flags_off_keep_default_varlen_path() + self._assert_fa4_path_pins_num_splits_and_forwards_scale() + + def _assert_flags_off_keep_default_varlen_path(self): total_tokens, num_heads, head_dim = 32, 8, 64 query = torch.randn(1, total_tokens, num_heads, head_dim) key = torch.randn(1, total_tokens, num_heads, head_dim) @@ -428,7 +390,7 @@ def test_flags_off_keep_default_varlen_path(self): assert not mock_sgl.called assert mock_varlen.called - def test_fa4_path_pins_num_splits_and_forwards_scale(self): + def _assert_fa4_path_pins_num_splits_and_forwards_scale(self): total_tokens, num_heads, head_dim = 32, 8, 64 query = torch.randn(1, total_tokens, num_heads, head_dim) key = torch.randn(1, total_tokens, num_heads, head_dim) @@ -460,7 +422,7 @@ def test_fa4_path_pins_num_splits_and_forwards_scale(self): assert kwargs["softmax_scale"] == 0.125 assert kwargs["causal"] is True - def test_rejects_cross_attention_cu_seqlens(self): + def _assert_rejects_cross_attention_cu_seqlens(self): num_heads, head_dim = 8, 64 q = torch.randn(32, num_heads, head_dim) k = torch.randn(48, num_heads, head_dim) diff --git a/tests/ops/test_bi_families_v2.py b/tests/ops/test_bi_families_v2.py deleted file mode 100644 index 3b6736ae..00000000 --- a/tests/ops/test_bi_families_v2.py +++ /dev/null @@ -1,93 +0,0 @@ -"""Functional tests for the trainer's families-v2 reduction trees.""" - -import os -from pathlib import Path - -import pytest -import torch - -from xorl.ops import bi_families_v2 as v2 -from xorl.ops.batch_invariant_ops import bi_lm_head_selected_logprob - - -H, V = 1024, 20480 - -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - - -def _payload(shape, seed): - generator = torch.Generator(device="cpu").manual_seed(seed) - return torch.randn(shape, generator=generator, dtype=torch.float32).to(torch.bfloat16).cuda() - - -def test_vendored_module_is_self_contained(): - """No engine imports: the same bytes have to run under either engine.""" - source = Path(v2.__file__).read_text() - for banned in ("import xorl", "from xorl", "import sglang", "from sglang"): - assert banned not in source, f"engine import {banned!r} in the vendored module" - - -def test_nonexact_family_selection_preserves_legacy_rollback(): - saved = {name: os.environ.get(name) for name in v2.FAMILIES_V2_ENV_VARS} - selected = v2._EXACT_FAMILIES_VERSION - try: - v2._EXACT_FAMILIES_VERSION = None - for name in v2.FAMILIES_V2_ENV_VARS: - os.environ.pop(name, None) - assert v2.families_v2_enabled() is True - for name in v2.FAMILIES_V2_ENV_VARS: - os.environ[name] = "0" - assert v2.families_v2_enabled() is False - del os.environ[name] - finally: - v2._EXACT_FAMILIES_VERSION = selected - for name, value in saved.items(): - if value is None: - os.environ.pop(name, None) - else: - os.environ[name] = value - - -def test_exact_family_selection_ignores_legacy_rollback(monkeypatch): - selected = v2._EXACT_FAMILIES_VERSION - try: - monkeypatch.setenv("XORL_FAMILIES_V2", "0") - v2._select_glm52_families_v2() - assert v2.families_v2_enabled() is True - - monkeypatch.setenv("XORL_FAMILIES_V2", "1") - v2._select_qwen35_families_v1() - assert v2.families_v2_enabled() is False - finally: - v2._EXACT_FAMILIES_VERSION = selected - - -@requires_cuda -@pytest.mark.gpu -def test_head_v2_selected_logit_is_bitwise_equal_to_v1(): - """The head redefinition keeps v1's pinned GEMM K chain, so the selected - logit is bitwise equal to v1; the vocabulary statistics layout is what - changed. The log-sum-exp is NOT claimed equal across the two generations — - a v2 trainer must be paired with a v2 sampler.""" - hidden, weight = _payload((16, H), 9), _payload((V, H), 10) - tokens = torch.arange(16, device="cuda") * 733 % V - logits, lse_decode = v2.head_v2_full_logits_with_lse(hidden, weight) - logprob, lse_scoring, selected = v2.head_v2_selected_logprob(hidden, weight, tokens) - - assert torch.equal(lse_decode, lse_scoring), "decode and scoring must share one tree" - assert torch.equal(selected, logits.gather(1, tokens[:, None]).squeeze(1)) - _, _, selected_v1 = bi_lm_head_selected_logprob(hidden, weight, tokens) - assert torch.equal(selected, selected_v1) - expected = torch.clamp_max(logits.gather(1, tokens[:, None]).squeeze(1) - lse_decode, 0.0) - assert torch.equal(expected, logprob) - - -@requires_cuda -@pytest.mark.gpu -def test_head_v2_is_batch_composition_invariant(): - hidden, weight = _payload((32, H), 11), _payload((V, H), 12) - tokens = torch.arange(32, device="cuda") * 601 % V - logprob, lse, _ = v2.head_v2_selected_logprob(hidden, weight, tokens) - for rows in (1, 4): - lp, ls, _ = v2.head_v2_selected_logprob(hidden[:rows].contiguous(), weight, tokens[:rows].contiguous()) - assert torch.equal(lp, logprob[:rows]) and torch.equal(ls, lse[:rows]) diff --git a/tests/ops/test_bi_families_v2_dispatch.py b/tests/ops/test_bi_families_v2_dispatch.py deleted file mode 100644 index 3cc2573b..00000000 --- a/tests/ops/test_bi_families_v2_dispatch.py +++ /dev/null @@ -1,92 +0,0 @@ -"""The v2 norm structure switch: which realization runs, and that it cannot matter. - -Both realizations compute the same tree, so the switch is a speed choice — but a -speed choice that was pointed the wrong way sent every shipped hidden size at -prefill row counts to the slower realization. These gates pin the direction and -re-pin bit-neutrality, so the switch stays free to move on measurement. -""" - -import pytest -import torch - - -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - -SHIPPED_H = (2048, 3840, 4096) -DEEP_H = 16384 - - -def _payload(rows, h): - g = torch.Generator(device="cpu").manual_seed(rows * 7 + h) - x = torch.randn((rows, h), generator=g, dtype=torch.float32).to(torch.bfloat16).cuda() - w = torch.randn((h,), generator=g, dtype=torch.float32).to(torch.bfloat16).cuda() - return x, w - - -def _split_spy(monkeypatch, v2): - calls = [] - real = v2._rms_norm_v2_split - - def spy(*args, **kwargs): - calls.append(args[0].shape) - return real(*args, **kwargs) - - monkeypatch.setattr(v2, "_rms_norm_v2_split", spy) - return calls - - -@requires_cuda -def test_dispatch_keeps_shipped_hidden_sizes_on_the_fused_realization(monkeypatch): - from xorl.ops import bi_families_v2 as v2 - - calls = _split_spy(monkeypatch, v2) - for h in SHIPPED_H: - for m in (1, 64, 512, 2048): - x, w = _payload(m, h) - v2.rms_norm_v2(x, w, 1e-6, residual=torch.zeros_like(x)) - assert calls == [], f"split realization ran at shipped hidden sizes: {calls}" - - -@requires_cuda -def test_dispatch_reaches_the_split_realization_at_deep_tile_shapes(monkeypatch): - from xorl.ops import bi_families_v2 as v2 - - calls = _split_spy(monkeypatch, v2) - x, w = _payload(8, DEEP_H) - v2.rms_norm_v2(x, w, 1e-6, residual=torch.zeros_like(x)) - assert calls, "split realization is unreachable; the cross-structure gates test nothing" - - -@requires_cuda -@pytest.mark.parametrize("h", (3840, DEEP_H)) -@pytest.mark.parametrize("rows", (1, 8, 512)) -def test_split_and_fused_realizations_are_bitwise_identical(h, rows): - """Structure is not bit-relevant — the premise the switch is allowed to move on.""" - from xorl.ops import bi_families_v2 as v2 - - x, w = _payload(rows, h) - r = torch.randn_like(x) - zx = x[:, :128].contiguous() - zw = w[:128].contiguous() - - def both(*args): - split = v2._rms_norm_v2_split(*args) - fused = getattr(v2, "_rms_norm_v2_fused", None) - if fused is not None: - return fused(*args), split - # base tree has no named fused entry: force it through the row switch - original = v2.V2_NORM_SPLIT_M - try: - v2.V2_NORM_SPLIT_M = 10**9 - x_, w_, eps_, res_, zc_ = args - out = v2.rms_norm_v2(x_, w_, eps_, residual=res_, zero_centered=zc_) - finally: - v2.V2_NORM_SPLIT_M = original - return out, split - - for args in ((x, w, 1e-6, r, False), (x, w, 1e-6, None, False), (zx, zw, 1e-6, None, True)): - fused_out, split_out = both(*args) - fused_out = fused_out if isinstance(fused_out, tuple) else (fused_out,) - split_out = split_out if isinstance(split_out, tuple) else (split_out,) - for a, b in zip(fused_out, split_out, strict=True): - assert torch.equal(a, b) diff --git a/tests/ops/test_bi_families_v2_norm.py b/tests/ops/test_bi_families_v2_norm.py index 1bfb5daa..9d24cb7c 100644 --- a/tests/ops/test_bi_families_v2_norm.py +++ b/tests/ops/test_bi_families_v2_norm.py @@ -1,4 +1,4 @@ -"""Contract gates for the families-v2 norm trees (hidden-dim RMSNorm + qk-norm). +"""Contract gates for the families-v2 hidden-dimension RMSNorm trees. The v2 trees are the frozen production contract for the batch-invariance lane. These gates pin the properties the design rests on: @@ -12,8 +12,6 @@ dispatch rule, so the gate keeps its teeth when that rule changes; - the dispatch rule itself is exercised directly, including the shipped hidden sizes where it must select the fused realization; -- the strided qk-norm entry normalizes packed qkv views in place without - touching the k/v bytes beside it; Correctness against an fp64 reference is a wrongness check, not a bit gate: v2 defines its own bits, and the reference cannot arbitrate between two trees that @@ -23,17 +21,17 @@ import pytest import torch +import xorl.models.layers.normalization as normalization from xorl.ops.bi_families_v2 import ( V2_NORM_SPLIT_MIN_TILES, V2_NORM_TILE, - qk_norm_v2, + families_v2_enabled, rms_norm_v2, ) EPS = 1e-6 H = 3840 -HEAD_DIM = 128 requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") pytestmark = pytest.mark.gpu @@ -69,8 +67,7 @@ def _force(monkeypatch, realization): # --- the tree: correctness and batch invariance ------------------------------ -@requires_cuda -def test_norm_v2_within_one_ulp_of_fp64_reference(): +def _assert_norm_v2_within_one_ulp_of_fp64_reference(): x, residual, weight = _payload((32, H), 1), _payload((32, H), 2), _payload((H,), 3) out, residual_out = rms_norm_v2(x, weight, EPS, residual=residual) assert _max_ulp(out, _reference(x, weight, residual=residual)) <= 1 @@ -94,9 +91,11 @@ def test_norm_v2_within_one_ulp_of_fp64_reference(): <= 1 ) + _assert_norm_v2_is_batch_composition_invariant() + _assert_norm_v2_is_run_to_run_bitwise() -@requires_cuda -def test_norm_v2_is_batch_composition_invariant(): + +def _assert_norm_v2_is_batch_composition_invariant(): x, residual, weight = _payload((128, H), 4), _payload((128, H), 5), _payload((H,), 6) full_out, full_residual = rms_norm_v2(x, weight, EPS, residual=residual) for rows in (1, 2, 17, 64): @@ -105,8 +104,7 @@ def test_norm_v2_is_batch_composition_invariant(): assert torch.equal(residual_out, full_residual[:rows]) -@requires_cuda -def test_norm_v2_is_run_to_run_bitwise(): +def _assert_norm_v2_is_run_to_run_bitwise(): x, residual, weight = _payload((64, H), 7), _payload((64, H), 8), _payload((H,), 9) first = rms_norm_v2(x, weight, EPS, residual=residual) for _ in range(3): @@ -122,87 +120,98 @@ def test_norm_v2_is_run_to_run_bitwise(): @requires_cuda -@pytest.mark.parametrize("hidden_size", [H, 4096, 12288]) -@pytest.mark.parametrize("rows", [1, 17, 64, 512]) -def test_fused_and_split_realizations_are_bitwise_identical(monkeypatch, hidden_size, rows): - x = _payload((rows, hidden_size), 600 + rows) - residual = _payload((rows, hidden_size), 700 + rows) - weight = _payload((hidden_size,), 800 + hidden_size) - - _force(monkeypatch, "fused") - fused_out, fused_residual = rms_norm_v2(x, weight, EPS, residual=residual) - fused_plain = rms_norm_v2(x, weight, EPS) - fused_zero_centered = rms_norm_v2(x, weight, EPS, zero_centered=True) - fused_zero_centered_residual = rms_norm_v2( - x, - weight, - EPS, - residual=residual, - zero_centered=True, - ) - - _force(monkeypatch, "split") - split_out, split_residual = rms_norm_v2(x, weight, EPS, residual=residual) - split_plain = rms_norm_v2(x, weight, EPS) - split_zero_centered = rms_norm_v2(x, weight, EPS, zero_centered=True) - split_zero_centered_residual = rms_norm_v2( - x, - weight, - EPS, - residual=residual, - zero_centered=True, - ) +def test_norm_v2_numerical_realization_and_dispatch_policy(monkeypatch): + _assert_norm_v2_within_one_ulp_of_fp64_reference() - assert torch.equal(fused_out, split_out) - assert torch.equal(fused_residual, split_residual) - assert torch.equal(fused_plain, split_plain) - assert torch.equal(fused_zero_centered, split_zero_centered) - assert torch.equal(fused_zero_centered_residual[0], split_zero_centered_residual[0]) - assert torch.equal(fused_zero_centered_residual[1], split_zero_centered_residual[1]) - - -@requires_cuda -def test_split_realization_is_actually_exercised(monkeypatch): - """Guard the guard: prove forcing 'split' reaches the split kernels. - - Without this, a refactor that dropped the split realization entirely would - leave the equality tests above passing vacuously. - """ import xorl.ops.bi_families_v2 as module - calls = [] - original = module._rms_norm_v2_split + split_calls = [] + original_split = module._rms_norm_v2_split def counting_split(*args, **kwargs): - calls.append(1) - return original(*args, **kwargs) + split_calls.append(1) + return original_split(*args, **kwargs) monkeypatch.setattr(module, "_rms_norm_v2_split", counting_split) - _force(monkeypatch, "split") - rms_norm_v2(_payload((512, H), 11), _payload((H,), 12), EPS) - assert calls, "forcing the split realization did not reach _rms_norm_v2_split" + + # Tail, aligned, and deep split-tile shapes at the smallest and largest + # useful row counts cover the kernel geometry; intermediate row literals + # do not select different code while each realization is forced. + for hidden_size in (H, 4096, 12288): + for rows in (1, 512): + x = _payload((rows, hidden_size), 600 + rows) + residual = _payload((rows, hidden_size), 700 + rows) + weight = _payload((hidden_size,), 800 + hidden_size) + + _force(monkeypatch, "fused") + fused_out, fused_residual = rms_norm_v2(x, weight, EPS, residual=residual) + fused_plain = rms_norm_v2(x, weight, EPS) + fused_zero_centered = rms_norm_v2(x, weight, EPS, zero_centered=True) + fused_zero_centered_residual = rms_norm_v2( + x, + weight, + EPS, + residual=residual, + zero_centered=True, + ) + + _force(monkeypatch, "split") + split_out, split_residual = rms_norm_v2(x, weight, EPS, residual=residual) + split_plain = rms_norm_v2(x, weight, EPS) + split_zero_centered = rms_norm_v2(x, weight, EPS, zero_centered=True) + split_zero_centered_residual = rms_norm_v2( + x, + weight, + EPS, + residual=residual, + zero_centered=True, + ) + + context = f"hidden_size={hidden_size}, rows={rows}" + assert torch.equal(fused_out, split_out), context + assert torch.equal(fused_residual, split_residual), context + assert torch.equal(fused_plain, split_plain), context + assert torch.equal(fused_zero_centered, split_zero_centered), context + assert torch.equal(fused_zero_centered_residual[0], split_zero_centered_residual[0]), context + assert torch.equal(fused_zero_centered_residual[1], split_zero_centered_residual[1]), context + assert split_calls, "forcing the split realization did not reach _rms_norm_v2_split" + + monkeypatch.undo() + with monkeypatch.context() as dispatch_patch: + _assert_norm_v2_dispatch_policy(dispatch_patch) + with monkeypatch.context() as reachability_patch: + _assert_norm_v2_reaches_trainer_dispatch(reachability_patch) # --- the dispatch rule ------------------------------------------------------- -@requires_cuda -@pytest.mark.parametrize("hidden_size", [2048, 3840, 4096]) -@pytest.mark.parametrize("rows", [1, 64, 256, 512, 2048]) -def test_shipped_hidden_sizes_always_take_the_fused_realization(hidden_size, rows): - """Common shipped hidden sizes stay below the measured split boundary.""" +def _assert_norm_v2_dispatch_policy(monkeypatch): import xorl.ops.bi_families_v2 as module - n_tiles = -(-hidden_size // V2_NORM_TILE) - assert n_tiles < V2_NORM_SPLIT_MIN_TILES - assert module._v2_norm_use_split(rows, n_tiles) is False + split_calls = [] + original_split = module._rms_norm_v2_split + def counting_split(*args, **kwargs): + split_calls.append(args[0].shape) + return original_split(*args, **kwargs) -@requires_cuda -def test_dispatch_rule_needs_a_deep_tile_chain_and_few_rows(): - """Split only when the split-kernel tile chain is deep and rows are few.""" - import xorl.ops.bi_families_v2 as module + monkeypatch.setattr(module, "_rms_norm_v2_split", counting_split) + # Common shipped hidden sizes stay below the measured split boundary. + for hidden_size in (2048, 3840, 4096): + n_tiles = -(-hidden_size // V2_NORM_TILE) + assert n_tiles < V2_NORM_SPLIT_MIN_TILES + # Once the tile count is below the threshold, row count cannot change + # the decision; retain only both row-count extremes. + for rows in (1, 2048): + assert module._v2_norm_use_split(rows, n_tiles) is False + x = _payload((rows, hidden_size), rows * 7 + hidden_size) + weight = _payload((hidden_size,), rows * 11 + hidden_size) + module.rms_norm_v2(x, weight, EPS, residual=torch.zeros_like(x)) + assert split_calls == [], f"split realization ran at shipped hidden sizes: {split_calls}" + + # Split only when the split-kernel tile chain is deep and rows are few. shallow = V2_NORM_SPLIT_MIN_TILES - 1 assert module._v2_norm_use_split(1, shallow) is False @@ -210,13 +219,8 @@ def test_dispatch_rule_needs_a_deep_tile_chain_and_few_rows(): assert module._v2_norm_use_split(deep, deep) is True assert module._v2_norm_use_split(deep + 1, deep) is False - -@requires_cuda -def test_dispatch_uses_the_split_kernels_tile_basis(): - """The rejected rule was passed the fused kernel's 4096-wide chunk count, - understating split parallelism by exactly 8x.""" - import xorl.ops.bi_families_v2 as module - + # The rejected rule used the fused kernel's 4096-wide chunk count, + # understating split parallelism by exactly 8x. hidden_size = 5120 split_tiles = -(-hidden_size // V2_NORM_TILE) fused_chunks = -(-hidden_size // 4096) @@ -225,24 +229,24 @@ def test_dispatch_uses_the_split_kernels_tile_basis(): assert module._v2_norm_use_split(1, split_tiles) is True assert module._v2_norm_use_split(1, fused_chunks) is False + # Prove the production dispatcher, not just its decision helper, reaches + # the split realization at a deep tile shape. + deep_hidden = 16384 + deep_x = _payload((8, deep_hidden), 8 * 7 + deep_hidden) + deep_weight = _payload((deep_hidden,), 8 * 11 + deep_hidden) + module.rms_norm_v2(deep_x, deep_weight, EPS, residual=torch.zeros_like(deep_x)) + assert split_calls == [(8, deep_hidden)] -# --- qk-norm ----------------------------------------------------------------- +def _assert_norm_v2_reaches_trainer_dispatch(monkeypatch): + assert families_v2_enabled() is True -@requires_cuda -def test_qk_norm_v2_strided_matches_contiguous_and_leaves_kv_untouched(): - tokens, n_q, n_kv = 256, 8, 2 - packed = _payload((tokens, (n_q + 2 * n_kv) * HEAD_DIM), 13) - weight = _payload((HEAD_DIM,), 14) - q_view = packed[:, : n_q * HEAD_DIM] - - out = qk_norm_v2(q_view, weight, EPS, head_dim=HEAD_DIM) - assert torch.equal(out, qk_norm_v2(q_view.contiguous(), weight, EPS, head_dim=HEAD_DIM)) - reference = _reference(q_view.contiguous().reshape(-1, HEAD_DIM), weight) - assert _max_ulp(out, reference.reshape(tokens, n_q * HEAD_DIM)) <= 1 - - scratch = packed.clone() - scratch_view = scratch[:, : n_q * HEAD_DIM] - qk_norm_v2(scratch_view, weight, EPS, head_dim=HEAD_DIM, out=scratch_view) - assert torch.equal(scratch_view, out), "in-place qk-norm diverged from out-of-place" - assert torch.equal(scratch[:, n_q * HEAD_DIM :], packed[:, n_q * HEAD_DIM :]), "k/v bytes moved" + x, residual, weight = _payload((64, H), 15), _payload((64, H), 16), _payload((H,), 17) + expected, expected_residual = rms_norm_v2(x, weight, EPS, residual=residual) + expected_plain = rms_norm_v2(x, weight, EPS) + + assert torch.equal(normalization.fast_sglang_rms_norm(x, weight, EPS), expected_plain) + assert torch.equal(normalization.fast_batch_invariant_rms_norm(x, weight, EPS), expected_plain) + fused_output, fused_residual = normalization.fast_sglang_residual_rms_norm(x, residual, weight, EPS) + assert torch.equal(fused_output, expected) + assert torch.equal(fused_residual, expected_residual) diff --git a/tests/ops/test_bi_families_v2_norm_dispatch.py b/tests/ops/test_bi_families_v2_norm_dispatch.py deleted file mode 100644 index 2f9fbc8c..00000000 --- a/tests/ops/test_bi_families_v2_norm_dispatch.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Reachability gate for the families-v2 norm adoption. - -The structure gates in tests/ops/test_bi_families_v2_norm.py prove the two -realizations of the v2 tree agree bitwise. This file proves the trainer's norm -entry points actually reach those kernels, and that the kill switch takes them -back to v1. Unreachable kernels gate nothing. -""" - -import pytest -import torch - -import xorl.models.layers.normalization as normalization -from xorl.ops.bi_families_v2 import families_v2_enabled, rms_norm_v2 - - -EPS = 1e-6 -H = 3840 - -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -pytestmark = pytest.mark.gpu - - -def _payload(shape, seed): - generator = torch.Generator(device="cpu").manual_seed(seed) - return torch.randn(shape, generator=generator, dtype=torch.float32).to(torch.bfloat16).cuda() - - -@requires_cuda -def test_norm_dispatcher_routes_to_families_v2(monkeypatch): - """The trainer's norm entry points must reach the v2 kernels, and the kill - switch must take them back to v1. Unreachable kernels gate nothing.""" - monkeypatch.delenv("XORL_FAMILIES_V2", raising=False) - monkeypatch.delenv("SGLANG_FAMILIES_V2", raising=False) - assert families_v2_enabled() is True, "families v2 is default on" - - x, residual, weight = _payload((64, H), 15), _payload((64, H), 16), _payload((H,), 17) - expected, expected_residual = rms_norm_v2(x, weight, EPS, residual=residual) - expected_plain = rms_norm_v2(x, weight, EPS) - - assert torch.equal(normalization.fast_sglang_rms_norm(x, weight, EPS), expected_plain) - assert torch.equal(normalization.fast_batch_invariant_rms_norm(x, weight, EPS), expected_plain) - - fused_out, fused_residual = normalization.fast_sglang_residual_rms_norm(x, residual, weight, EPS) - assert torch.equal(fused_out, expected) - assert torch.equal(fused_residual, expected_residual) - - # Rollback: assert on the kernel selected, not on the bits. The two trees - # agree on most values, so a bit comparison would be a weak signal. - import xorl.ops.bi_families_v2 as module # noqa: PLC0415 - - reached = [] - original = module.rms_norm_v2 - monkeypatch.setattr(normalization, "rms_norm_v2", lambda *a, **k: (reached.append(1), original(*a, **k))[1]) - monkeypatch.setenv("XORL_FAMILIES_V2", "0") - assert families_v2_enabled() is False, "the kill switch must roll this engine back" - normalization.fast_batch_invariant_rms_norm(x, weight, EPS) - normalization.fast_sglang_residual_rms_norm(x, residual, weight, EPS) - assert not reached, "the kill switch left the dispatcher on the v2 kernels" - - -def test_kill_switch_rolls_back_on_either_engine_variable(monkeypatch): - """One flag, both engines. A setting that rolls back only one of the two - would put the trainer and the sampler on different trees.""" - for variable in ("XORL_FAMILIES_V2", "SGLANG_FAMILIES_V2"): - monkeypatch.delenv("XORL_FAMILIES_V2", raising=False) - monkeypatch.delenv("SGLANG_FAMILIES_V2", raising=False) - assert families_v2_enabled() is True - monkeypatch.setenv(variable, "0") - assert families_v2_enabled() is False, f"{variable}=0 must roll back" diff --git a/tests/ops/test_bi_fused_lm_head.py b/tests/ops/test_bi_fused_lm_head.py index 7aa0da51..e6102341 100644 --- a/tests/ops/test_bi_fused_lm_head.py +++ b/tests/ops/test_bi_fused_lm_head.py @@ -1,6 +1,9 @@ import pytest import torch +from xorl.ops import batch_invariant_ops as v1 +from xorl.ops import bi_families_v2 as v2 +from xorl.ops.loss.bi_fused_lm_head import bi_fused_per_token_ce from xorl.ops.loss.causallm_loss import causallm_loss_function @@ -21,36 +24,34 @@ def _inputs(seed=0): @requires_cuda @pytest.mark.gpu -def test_bi_fused_matches_eager_fp32_reference(): - hidden, weight, labels = _inputs() - out = causallm_loss_function(hidden, weight, labels, ce_mode="bi_fused", lm_head_fp32=True, return_per_token=True) - ref = causallm_loss_function(hidden, weight, labels, ce_mode="eager", lm_head_fp32=True, return_per_token=True) - assert torch.allclose(out.per_token_logprobs, ref.per_token_logprobs, rtol=1e-4, atol=1e-5) - assert torch.allclose(out.loss.float(), ref.loss.float(), rtol=1e-4, atol=1e-5) - # ignored positions contribute exactly zero - assert (out.per_token_loss.view(-1)[:7] == 0).all() - - -@requires_cuda -@pytest.mark.gpu -def test_bi_fused_backward_matches_eager_autograd(): - hidden, weight, labels = _inputs(1) - - h1 = hidden.clone().requires_grad_(True) - w1 = weight.clone().requires_grad_(True) - causallm_loss_function(h1, w1, labels, ce_mode="bi_fused", lm_head_fp32=True).loss.backward() - - h2 = hidden.clone().requires_grad_(True) - w2 = weight.clone().requires_grad_(True) - causallm_loss_function(h2, w2, labels, ce_mode="eager", lm_head_fp32=True).loss.backward() - - assert torch.allclose(h1.grad.float(), h2.grad.float(), rtol=2e-2, atol=2e-4) - assert torch.allclose(w1.grad.float(), w2.grad.float(), rtol=2e-2, atol=2e-4) - - -@requires_cuda -@pytest.mark.gpu -def test_bi_fused_deterministic_and_batch_invariant(): +def test_bi_fused_forward_backward_and_kernel_edge_policy(): + for seed, temperature in ((0, None), (4, 0.7)): + hidden, weight, labels = _inputs(seed) + fused_hidden = hidden.clone().requires_grad_(True) + fused_weight = weight.clone().requires_grad_(True) + eager_hidden = hidden.clone().requires_grad_(True) + eager_weight = weight.clone().requires_grad_(True) + kwargs = {"lm_head_fp32": True, "return_per_token": True} + if temperature is not None: + kwargs["logprob_temperature"] = temperature + + out = causallm_loss_function(fused_hidden, fused_weight, labels, ce_mode="bi_fused", **kwargs) + ref = causallm_loss_function(eager_hidden, eager_weight, labels, ce_mode="eager", **kwargs) + assert torch.allclose(out.per_token_logprobs, ref.per_token_logprobs, rtol=1e-4, atol=1e-5) + assert torch.allclose(out.loss.float(), ref.loss.float(), rtol=1e-4, atol=1e-5) + assert (out.per_token_loss.view(-1)[:7] == 0).all() + out.loss.backward() + ref.loss.backward() + assert torch.allclose(fused_hidden.grad.float(), eager_hidden.grad.float(), rtol=2e-2, atol=2e-4) + assert torch.allclose(fused_weight.grad.float(), eager_weight.grad.float(), rtol=2e-2, atol=2e-4) + + _assert_bi_fused_deterministic_and_batch_invariant() + _assert_bi_fused_guards() + _assert_bi_kernel_unit_temperature_is_exact_identity() + _assert_head_v2_projection_stats_invariance_and_training_policy() + + +def _assert_bi_fused_deterministic_and_batch_invariant(): hidden, weight, labels = _inputs(2) kw = dict(ce_mode="bi_fused", lm_head_fp32=True, return_per_token=True) a = causallm_loss_function(hidden, weight, labels, **kw).per_token_logprobs @@ -60,9 +61,7 @@ def test_bi_fused_deterministic_and_batch_invariant(): assert torch.equal(sub, a[:, 64:128]) -@requires_cuda -@pytest.mark.gpu -def test_bi_fused_guards(): +def _assert_bi_fused_guards(): hidden, weight, labels = _inputs(3) with pytest.raises(NotImplementedError, match="lm_head_fp32"): causallm_loss_function(hidden, weight, labels, ce_mode="bi_fused", lm_head_fp32=False) @@ -70,39 +69,7 @@ def test_bi_fused_guards(): causallm_loss_function(hidden, weight, labels, ce_mode="bi_fused", lm_head_fp32=True, z_loss_coef=0.1) -@requires_cuda -@pytest.mark.gpu -def test_bi_fused_temperature_matches_eager_reference(): - hidden, weight, labels = _inputs(4) - kw = dict(lm_head_fp32=True, logprob_temperature=0.7, return_per_token=True) - out = causallm_loss_function(hidden, weight, labels, ce_mode="bi_fused", **kw) - ref = causallm_loss_function(hidden, weight, labels, ce_mode="eager", **kw) - assert torch.allclose(out.per_token_logprobs, ref.per_token_logprobs, rtol=1e-4, atol=1e-5) - assert (out.per_token_loss.view(-1)[:7] == 0).all() - - -@requires_cuda -@pytest.mark.gpu -def test_bi_fused_temperature_backward_matches_eager_autograd(): - hidden, weight, labels = _inputs(5) - - h1 = hidden.clone().requires_grad_(True) - w1 = weight.clone().requires_grad_(True) - causallm_loss_function( - h1, w1, labels, ce_mode="bi_fused", lm_head_fp32=True, logprob_temperature=0.7 - ).loss.backward() - - h2 = hidden.clone().requires_grad_(True) - w2 = weight.clone().requires_grad_(True) - causallm_loss_function(h2, w2, labels, ce_mode="eager", lm_head_fp32=True, logprob_temperature=0.7).loss.backward() - - assert torch.allclose(h1.grad.float(), h2.grad.float(), rtol=2e-2, atol=2e-4) - assert torch.allclose(w1.grad.float(), w2.grad.float(), rtol=2e-2, atol=2e-4) - - -@requires_cuda -@pytest.mark.gpu -def test_bi_kernel_unit_temperature_is_exact_identity(): +def _assert_bi_kernel_unit_temperature_is_exact_identity(): from xorl.ops.batch_invariant_ops import bi_lm_head_selected_logprob hidden, weight, _ = _inputs(6) @@ -117,10 +84,10 @@ def test_bi_kernel_unit_temperature_is_exact_identity(): assert torch.equal(lse_none, lse_ones) assert torch.equal(sel_none, sel_ones) + _assert_bi_kernel_p1_tokens_clamp_to_exact_zero() -@requires_cuda -@pytest.mark.gpu -def test_bi_kernel_p1_tokens_clamp_to_exact_zero(): + +def _assert_bi_kernel_p1_tokens_clamp_to_exact_zero(): from xorl.ops.batch_invariant_ops import bi_lm_head_selected_logprob torch.manual_seed(7) @@ -136,3 +103,43 @@ def test_bi_kernel_p1_tokens_clamp_to_exact_zero(): # observed live as +2**-18) must clamp to exactly 0.0, never positive. assert (lp <= 0).all() assert (lp == 0).any() + + +def _assert_head_v2_projection_stats_invariance_and_training_policy(): + def make(shape, seed): + generator = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(shape, generator=generator, dtype=torch.float32).to(torch.bfloat16).cuda() + + hidden, weight = make((16, 128), 1), make((512, 128), 2) + tokens = torch.arange(16, device="cuda") * 29 % 512 + + logits, decode_lse = v2.head_v2_full_logits_with_lse(hidden, weight) + reference_logits = torch.empty_like(logits) + v1._bi_lm_head_chunk_gemm_fp32(hidden, weight.t(), reference_logits) + assert torch.equal(logits, reference_logits) + + logprob, score_lse, selected = v2.head_v2_selected_logprob(hidden, weight, tokens) + assert torch.equal(decode_lse, score_lse) + assert torch.equal(selected, logits.gather(1, tokens[:, None]).squeeze(1)) + assert torch.equal(logprob, torch.clamp_max(selected - decode_lse, 0.0)) + + hidden, weight = make((32, 128), 3), make((512, 128), 4) + tokens = torch.arange(32, device="cuda") * 31 % 512 + full, _, _ = v2.head_v2_selected_logprob(hidden, weight, tokens) + sub, _, _ = v2.head_v2_selected_logprob(hidden[8:16].contiguous(), weight, tokens[8:16].contiguous()) + assert torch.equal(sub, full[8:16]) + + hidden.requires_grad_(True) + weight.requires_grad_(True) + ce = bi_fused_per_token_ce(hidden, weight, tokens) + assert torch.equal(ce, -full) + ce.mean().backward() + assert torch.isfinite(hidden.grad).all() + assert torch.isfinite(weight.grad).all() + + v2._select_qwen35_families_v1() + try: + qualified_v1 = bi_fused_per_token_ce(hidden.detach(), weight.detach(), tokens, vocab_chunk=256) + assert torch.isfinite(qualified_v1).all() + finally: + v2._select_nonexact_families() diff --git a/tests/ops/test_bi_gdn_contract.py b/tests/ops/test_bi_gdn_contract.py index e67aaad9..2e5f8e45 100644 --- a/tests/ops/test_bi_gdn_contract.py +++ b/tests/ops/test_bi_gdn_contract.py @@ -1,3 +1,5 @@ +import importlib + import pytest import torch import torch.nn.functional as F @@ -11,6 +13,17 @@ HV, DV = 32, 128 +class _FakeKernel: + def __init__(self): + self.calls = [] + + def __getitem__(self, grid): + def launch(**kwargs): + self.calls.append((grid, kwargs)) + + return launch + + def _rel_err(got: torch.Tensor, exp: torch.Tensor) -> float: return float((got.float() - exp.float()).norm() / exp.float().norm().clamp_min(1e-12)) @@ -24,13 +37,20 @@ def _gating_inputs(T, seed=0): return A_log, a, b, dt_bias -@requires_cuda -@pytest.mark.gpu -def test_gating_forward_matches_reference_composition(): - A_log, a, b, dt_bias = _gating_inputs(2048) - g, beta = bi_fused_gdn_gating(A_log, a, b, dt_bias) - g_ref = -A_log.exp().view(1, 1, -1) * F.softplus(a + dt_bias.view(1, 1, -1)) - beta_ref = b.float().sigmoid() +def _assert_gating_forward_and_backward_match_reference_composition(): + A_log, a, b, dt_bias = _gating_inputs(512, seed=1) + actual_A = A_log.clone().requires_grad_(True) + actual_a = a.clone().requires_grad_(True) + actual_b = b.clone().requires_grad_(True) + actual_dt = dt_bias.clone().requires_grad_(True) + g, beta = bi_fused_gdn_gating(actual_A, actual_a, actual_b, actual_dt) + + reference_A = A_log.clone().requires_grad_(True) + reference_a = a.clone().requires_grad_(True) + reference_b = b.clone().requires_grad_(True) + reference_dt = dt_bias.clone().requires_grad_(True) + g_ref = -reference_A.exp().view(1, 1, -1) * F.softplus(reference_a + reference_dt.view(1, 1, -1)) + beta_ref = reference_b.float().sigmoid() assert g.dtype == torch.float32 and g.shape == g_ref.shape # the serving kernel's tl.log(1+tl.exp) vs torch softplus is a 1-ulp fp32 term assert torch.allclose(g, g_ref, rtol=1e-5, atol=1e-5) @@ -39,28 +59,15 @@ def test_gating_forward_matches_reference_composition(): assert beta.dtype == torch.float32 assert torch.allclose(beta, beta_ref, rtol=1e-6, atol=1e-6) assert not torch.equal(beta, beta_ref.to(b.dtype).float()), "beta must not be bf16-rounded" - - -@requires_cuda -@pytest.mark.gpu -def test_gating_backward_matches_autograd_reference(): - A_log, a, b, dt_bias = _gating_inputs(512, seed=1) - ag = a.clone().requires_grad_(True) - bg = b.clone().requires_grad_(True) - Ag = A_log.clone().requires_grad_(True) - dg = dt_bias.clone().requires_grad_(True) - g, beta = bi_fused_gdn_gating(Ag, ag, bg, dg) (g.square() + beta.square()).sum().backward() - - ar = a.clone().requires_grad_(True) - br = b.clone().requires_grad_(True) - Ar = A_log.clone().requires_grad_(True) - dr = dt_bias.clone().requires_grad_(True) - g_ref = -Ar.exp().view(1, 1, -1) * F.softplus(ar + dr.view(1, 1, -1)) - beta_ref = br.float().sigmoid() (g_ref.square() + beta_ref.square()).sum().backward() - for got, exp in ((ag.grad, ar.grad), (bg.grad, br.grad), (Ag.grad, Ar.grad), (dg.grad, dr.grad)): + for got, exp in ( + (actual_a.grad, reference_a.grad), + (actual_b.grad, reference_b.grad), + (actual_A.grad, reference_A.grad), + (actual_dt.grad, reference_dt.grad), + ): assert _rel_err(got, exp) < 1e-2 @@ -74,45 +81,41 @@ def _norm_inputs(T, seed=0): @requires_cuda @pytest.mark.gpu -def test_gated_norm_forward_matches_reference_and_is_row_invariant(): - x, z, w = _norm_inputs(1024) +def test_gdn_gating_norm_and_exact_model_dispatch_policy(): + _assert_gating_forward_and_backward_match_reference_composition() + + x, z, w = _norm_inputs(256, seed=2) eps = 1e-6 - y = bi_rms_norm_gated(x, w, z, eps) - xf = x.float() + actual_x = x.clone().requires_grad_(True) + actual_z = z.clone().requires_grad_(True) + actual_w = w.clone().requires_grad_(True) + y = bi_rms_norm_gated(actual_x, actual_w, actual_z, eps) + + reference_x = x.clone().requires_grad_(True) + reference_z = z.clone().requires_grad_(True) + reference_w = w.clone().requires_grad_(True) + xf = reference_x.float() n = xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) - y_ref = ((n * w.float()) * (z.float() * torch.sigmoid(z.float()))).to(x.dtype) + y_ref = ((n * reference_w.float()) * (reference_z.float() * torch.sigmoid(reference_z.float()))).to(x.dtype) assert y.shape == x.shape and y.dtype == x.dtype # rare bf16-ULP tail vs the torch composition is the contracted serving term assert torch.allclose(y.float(), y_ref.float(), rtol=1e-2, atol=1e-2) # per-row numerics must not depend on the number of rows in the launch y_sub = bi_rms_norm_gated(x[:, :3], w, z[:, :3], eps) assert torch.equal(y[:, :3], y_sub) - - -@requires_cuda -@pytest.mark.gpu -def test_gated_norm_backward_matches_autograd_reference(): - x, z, w = _norm_inputs(256, seed=2) - eps = 1e-6 - xg = x.clone().requires_grad_(True) - zg = z.clone().requires_grad_(True) - wg = w.clone().requires_grad_(True) - bi_rms_norm_gated(xg, wg, zg, eps).float().square().sum().backward() - - xr = x.clone().requires_grad_(True) - zr = z.clone().requires_grad_(True) - wr = w.clone().requires_grad_(True) - xf = xr.float() - n = xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) - ((n * wr.float()) * (zr.float() * torch.sigmoid(zr.float()))).square().sum().backward() - - for got, exp in ((xg.grad, xr.grad), (zg.grad, zr.grad), (wg.grad, wr.grad)): + y.float().square().sum().backward() + y_ref.float().square().sum().backward() + for got, exp in ( + (actual_x.grad, reference_x.grad), + (actual_z.grad, reference_z.grad), + (actual_w.grad, reference_w.grad), + ): assert _rel_err(got, exp) < 1e-2 + _assert_fused_rms_norm_gated_module_routes_under_exact_model_program() -@requires_cuda -@pytest.mark.gpu -def test_fused_rms_norm_gated_module_routes_under_exact_model_program(): + +def _assert_fused_rms_norm_gated_module_routes_under_exact_model_program(): x, z, w = _norm_inputs(128, seed=3) module = FusedRMSNormGated(DV, eps=1e-6).to(device="cuda", dtype=torch.bfloat16) with torch.no_grad(): @@ -130,10 +133,10 @@ def test_fused_rms_norm_gated_module_routes_under_exact_model_program(): with gdn_contract(True), pytest.raises(NotImplementedError): module(x, z, residual=torch.zeros_like(x)) + _assert_gated_deltanet_gating_routes_under_exact_model_program() -@requires_cuda -@pytest.mark.gpu -def test_gated_deltanet_gating_routes_under_exact_model_program(): + +def _assert_gated_deltanet_gating_routes_under_exact_model_program(): from xorl.ops.linear_attention.layers.gated_deltanet import GatedDeltaNet # noqa: PLC0415 ordinary = ( @@ -195,3 +198,41 @@ def test_solve_tril_num_warps_pinned(): ): for cfg in kernel.fn.configs: assert cfg.num_warps == 2 + + _assert_kkt_reduction_geometry_matches_the_serving_contract() + + +def _assert_kkt_reduction_geometry_matches_the_serving_contract(): + module = importlib.import_module("xorl.ops.linear_attention.ops.common.chunk_scaled_dot_kkt") + contract_kernel = _FakeKernel() + autotuned_kernel = _FakeKernel() + original_contract_kernel = module._chunk_scaled_dot_kkt_fwd_kernel + original_autotuned_kernel = module.chunk_scaled_dot_kkt_fwd_kernel + module._chunk_scaled_dot_kkt_fwd_kernel = contract_kernel + module.chunk_scaled_dot_kkt_fwd_kernel = autotuned_kernel + try: + k = torch.empty(1, 64, 32, 128) + g = torch.empty(1, 64, 32) + beta = torch.empty(1, 64, 32) + + with gdn_contract(True): + module.chunk_scaled_dot_kkt_fwd(k=k, g=g, beta=beta) + assert not autotuned_kernel.calls + _, kwargs = contract_kernel.calls.pop() + assert kwargs["BK"] == 64 + assert kwargs["num_warps"] == 8 + assert kwargs["num_stages"] == 3 + assert kwargs["IS_VARLEN"] is False + assert kwargs["USE_G"] is True + assert kwargs["SAFE_EXP"] is True + + with gdn_contract(False): + module.chunk_scaled_dot_kkt_fwd(k=k, g=g, beta=beta) + assert not contract_kernel.calls + _, kwargs = autotuned_kernel.calls.pop() + assert "BK" not in kwargs + assert "num_warps" not in kwargs + assert "num_stages" not in kwargs + finally: + module._chunk_scaled_dot_kkt_fwd_kernel = original_contract_kernel + module.chunk_scaled_dot_kkt_fwd_kernel = original_autotuned_kernel diff --git a/tests/ops/test_bi_gemm_config_table.py b/tests/ops/test_bi_gemm_config_table.py index 41c56749..2a61a5b1 100644 --- a/tests/ops/test_bi_gemm_config_table.py +++ b/tests/ops/test_bi_gemm_config_table.py @@ -11,7 +11,6 @@ import torch import triton -from xorl.ops import batch_invariant_ops from xorl.ops.batch_invariant_ops import ( _deepgemm_ready, _matmul_persistent_deepgemm, @@ -34,19 +33,6 @@ ] -@pytest.mark.cpu -def test_ambient_legacy_envs_cannot_change_the_production_program(monkeypatch): - monkeypatch.setenv("XORL_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM", "0") - monkeypatch.setenv("SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT", "1") - monkeypatch.setenv("SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_COMPARISON_TEST", "1") - monkeypatch.setenv("XORL_BI_GEMM_CONFIG_TABLE", "0") - - assert batch_invariant_ops._ENABLE_MM_DEEPGEMM is True - assert not hasattr(batch_invariant_ops, "_ENABLE_MM_FALLBACK_VARIANT") - assert not hasattr(batch_invariant_ops, "_ENABLE_MM_COMPARISON_TEST") - assert lookup_mm_config(torch.bfloat16, 1, 3840, 3840) != dict(BASELINE_CONFIG["torch.bfloat16"], BLOCK_SIZE_K=64) - - def _launch(a, b, cfg, out_dtype=None): M, K = a.shape _, N = b.shape @@ -87,42 +73,36 @@ def _inputs(M, K, N, dtype, seed=0): return a, w.t() -def test_block_k_pinned_per_dtype(): - for dt_str, block_k in PINNED_BLOCK_K.items(): - dt = getattr(torch, dt_str.removeprefix("torch.")) - for M, N, K in [(1, 128, 128), (300, 3840, 3840), (32768, 11520, 3840)]: - assert lookup_mm_config(dt, M, N, K)["BLOCK_SIZE_K"] == block_k - - -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("shape", SHAPES) -def test_table_config_bitwise_equals_baseline(dtype, shape): - M, K, N = shape - a, b = _inputs(M, K, N, dtype) - base = dict(BASELINE_CONFIG[str(dtype)], BLOCK_SIZE_K=PINNED_BLOCK_K[str(dtype)]) - cfg = lookup_mm_config(dtype, M, N, K) - out_base = _launch(a, b, base) - out_cfg = _launch(a, b, cfg) - assert torch.equal(out_base, out_cfg), f"table config moved bits at {dtype} {shape}: {cfg}" - # fp32-out store path (the BI lm-head chunk form): mirror the production - # launcher — table config with OutOfResources->baseline fallback — and - # assert bit-neutrality down to raw accumulator bits - if dtype == torch.bfloat16: - from triton.runtime.errors import OutOfResources # noqa: PLC0415 - - cfg32 = lookup_mm_config(dtype, M, N, K, out_itemsize=4) - try: - out32 = _launch(a, b, cfg32, torch.float32) - except OutOfResources: - out32 = _launch(a, b, base, torch.float32) - assert torch.equal(_launch(a, b, base, torch.float32), out32) - - @requires_cuda @pytest.mark.gpu -def test_rows_invariant_across_m_and_buckets(): +def test_table_config_bitwise_equals_baseline(): + # This is one doctrine-level contract: every generated table entry must be + # bit-neutral relative to the dtype-pinned reduction tree. + for dtype in (torch.bfloat16, torch.float32): + for shape in SHAPES: + M, K, N = shape + a, b = _inputs(M, K, N, dtype) + base = dict(BASELINE_CONFIG[str(dtype)], BLOCK_SIZE_K=PINNED_BLOCK_K[str(dtype)]) + cfg = lookup_mm_config(dtype, M, N, K) + out_base = _launch(a, b, base) + out_cfg = _launch(a, b, cfg) + assert torch.equal(out_base, out_cfg), f"table config moved bits at {dtype} {shape}: {cfg}" + # fp32-out store path (the BI lm-head chunk form): mirror the production + # launcher — table config with OutOfResources->baseline fallback — and + # assert bit-neutrality down to raw accumulator bits + if dtype == torch.bfloat16: + from triton.runtime.errors import OutOfResources # noqa: PLC0415 + + cfg32 = lookup_mm_config(dtype, M, N, K, out_itemsize=4) + try: + out32 = _launch(a, b, cfg32, torch.float32) + except OutOfResources: + out32 = _launch(a, b, base, torch.float32) + assert torch.equal(_launch(a, b, base, torch.float32), out32), shape + _assert_rows_invariant_across_m_and_buckets() + + +def _assert_rows_invariant_across_m_and_buckets(): K, N = 3840, 3840 g = torch.Generator(device="cuda").manual_seed(3) w = torch.randn((N, K), device="cuda", generator=g, dtype=torch.float32).to(torch.bfloat16) diff --git a/tests/ops/test_bi_golden_gates.py b/tests/ops/test_bi_golden_gates.py index edf311a6..886a8812 100644 --- a/tests/ops/test_bi_golden_gates.py +++ b/tests/ops/test_bi_golden_gates.py @@ -83,7 +83,7 @@ def ids(n, seed, mod): return torch.from_numpy((_splitmix64(n, seed) % np.uint64(mod)).astype(np.int64)).cuda() -def _check(case, **outs): +def _assert_golden(case, **outs): for name, t in outs.items(): got = hashlib.sha256(t.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes()).hexdigest() want = GOLDENS[case][name]["sha256"] @@ -91,29 +91,32 @@ def _check(case, **outs): @requires_h100 -def test_v1_family1_and_zero_centered_goldens(): - _check( +def test_v1_contract_tree_goldens(): + _assert_golden( "family1_qk_4096x128", out=v1.bi_rms_norm(bf16((4096, DQK), 101), w(DQK, 102), EPS, family=v1.RMS_NORM_FAMILY_NO_RESIDUAL), ) - _check( + _assert_golden( "family1_h3840_m64", out=v1.bi_rms_norm(bf16((64, REAL_H), 111), w(REAL_H, 112), EPS, family=v1.RMS_NORM_FAMILY_NO_RESIDUAL), ) - _check( + _assert_golden( "family1_zero_centered_512x128", out=v1.bi_rms_norm( bf16((512, DQK), 121), w(DQK, 122), EPS, family=v1.RMS_NORM_FAMILY_NO_RESIDUAL, zero_centered=True ), ) + _assert_v1_family2_and_mean_goldens() + _assert_v1_log_softmax_and_mm_goldens() + _assert_v1_lmhead_scoring_goldens() -@requires_h100 -def test_v1_family2_and_mean_goldens(): + +def _assert_v1_family2_and_mean_goldens(): out, res = v1.bi_fused_add_rms_norm( bf16((64, REAL_H), 131), bf16((64, REAL_H), 132), w(REAL_H, 133), EPS, family=v1.RMS_NORM_FAMILY_RESIDUAL_TREE ) - _check("family2_residual_m64_h3840", out=out, residual_out=res) + _assert_golden("family2_residual_m64_h3840", out=out, residual_out=res) out, res = v1.bi_fused_add_rms_norm( bf16((1, REAL_H), 141, edge_rows=False), bf16((1, REAL_H), 142, edge_rows=False), @@ -121,46 +124,53 @@ def test_v1_family2_and_mean_goldens(): EPS, family=v1.RMS_NORM_FAMILY_RESIDUAL_TREE, ) - _check("family2_residual_m1_h3840", out=out, residual_out=res) - _check( + _assert_golden("family2_residual_m1_h3840", out=out, residual_out=res) + _assert_golden( "family2_noresidual_m64_h3840", out=v1.bi_rms_norm(bf16((64, REAL_H), 151), w(REAL_H, 152), EPS, family=v1.RMS_NORM_FAMILY_RESIDUAL_TREE), ) - _check("mean_dim_m64_h3840", out=v1.mean_dim(fp32((64, REAL_H), 161, exp_lo=110, exp_hi=132), -1, True)) + _assert_golden( + "mean_dim_m64_h3840", + out=v1.mean_dim(fp32((64, REAL_H), 161, exp_lo=110, exp_hi=132), -1, True), + ) -@requires_h100 -def test_v1_log_softmax_and_mm_goldens(): - _check("log_softmax_4x128256", out=v1.log_softmax(fp32((4, REAL_VOCAB), 171, exp_lo=124, exp_hi=131), dim=-1)) - _check( +def _assert_v1_log_softmax_and_mm_goldens(): + _assert_golden( + "log_softmax_4x128256", + out=v1.log_softmax(fp32((4, REAL_VOCAB), 171, exp_lo=124, exp_hi=131), dim=-1), + ) + _assert_golden( "mm_trunk_64x3840x5120", out=v1.matmul_persistent(bf16((64, REAL_H), 181), bf16((REAL_H, 5120), 182, edge_rows=False)), ) - _check( + _assert_golden( "mm_offtable_17x1024x1000", out=v1.matmul_persistent(bf16((17, 1024), 191), bf16((1024, 1000), 192, edge_rows=False)), ) -@requires_h100 -def test_v1_lmhead_scoring_goldens(): +def _assert_v1_lmhead_scoring_goldens(): h, wt, tk = bf16((8, REAL_H), 201), bf16((REAL_VOCAB, REAL_H), 202, edge_rows=False), ids(8, 203, REAL_VOCAB) lp, lse, sel = v1.bi_lm_head_selected_logprob(h, wt, tk) - _check("lmhead_scoring_n8_real", logprob=lp, lse=lse, selected=sel) + _assert_golden("lmhead_scoring_n8_real", logprob=lp, lse=lse, selected=sel) h, wt, tk = bf16((8, REAL_H), 211), bf16((REAL_VOCAB, REAL_H), 212, edge_rows=False), ids(8, 213, REAL_VOCAB) lp, lse, sel = v1.bi_lm_head_selected_logprob(h, wt, tk, torch.full((8,), 0.7, dtype=torch.float32, device="cuda")) - _check("lmhead_scoring_n8_temp0.7", logprob=lp, lse=lse, selected=sel) + _assert_golden("lmhead_scoring_n8_temp0.7", logprob=lp, lse=lse, selected=sel) h, wt, tk = bf16((64, 512), 221), bf16((20480, 512), 222, edge_rows=False), ids(64, 223, 20480) lp, lse, sel = v1.bi_lm_head_selected_logprob(h, wt, tk) - _check("lmhead_scoring_n64_small", logprob=lp, lse=lse, selected=sel) + _assert_golden("lmhead_scoring_n64_small", logprob=lp, lse=lse, selected=sel) @requires_h100 -def test_v2_norm_goldens(): +def test_v2_contract_tree_goldens(): out, res = v2.rms_norm_v2(bf16((64, REAL_H), 131), w(REAL_H, 133), EPS, residual=bf16((64, REAL_H), 132)) - _check("norm_v2_residual_m64_h3840", out=out, residual_out=res) - _check("norm_v2_noresidual_m64_h3840", out=v2.rms_norm_v2(bf16((64, REAL_H), 131), w(REAL_H, 133), EPS)) - _check( + _assert_golden("norm_v2_residual_m64_h3840", out=out, residual_out=res) + _assert_golden( + "norm_v2_noresidual_m64_h3840", + out=v2.rms_norm_v2(bf16((64, REAL_H), 131), w(REAL_H, 133), EPS), + ) + _assert_golden( "norm_v2_zero_centered", out=v2.rms_norm_v2(bf16((64, REAL_H), 131)[:, :DQK].contiguous(), w(DQK, 102), EPS, zero_centered=True), ) @@ -171,20 +181,19 @@ def test_v2_norm_goldens(): residual=bf16((64, REAL_H), 132), zero_centered=True, ) - _check("norm_v2_zero_centered_residual_m64_h3840", out=out, residual_out=res) - packed = bf16((256, (8 + 2 * 2) * DQK), 401) - _check("qk_v2_strided_t256_h8_d128", out=v2.qk_norm_v2(packed[:, : 8 * DQK], w(DQK, 102), EPS, head_dim=DQK)) + _assert_golden("norm_v2_zero_centered_residual_m64_h3840", out=out, residual_out=res) + _assert_v2_head_goldens() -@requires_h100 -def test_v2_head_goldens(): + +def _assert_v2_head_goldens(): h, wt, tk = bf16((64, REAL_H), 501), bf16((REAL_VOCAB, REAL_H), 502, edge_rows=False), ids(64, 503, REAL_VOCAB) lp, lse, sel = v2.head_v2_selected_logprob(h, wt, tk) - _check("head_v2_scoring_n64_real", logprob=lp, lse=lse, selected=sel) + _assert_golden("head_v2_scoring_n64_real", logprob=lp, lse=lse, selected=sel) lp, lse, _ = v2.head_v2_selected_logprob(h, wt, tk, torch.full((64,), 0.7, dtype=torch.float32, device="cuda")) - _check("head_v2_scoring_temp0.7", logprob=lp, lse=lse) + _assert_golden("head_v2_scoring_temp0.7", logprob=lp, lse=lse) logits, lse_d = v2.head_v2_full_logits_with_lse(h, wt) - _check("head_v2_decode_n64", logits=logits, lse=lse_d) + _assert_golden("head_v2_decode_n64", logits=logits, lse=lse_d) hs, ws2, ts = bf16((8, 512), 512), bf16((20580, 512), 511, edge_rows=False), ids(8, 513, 20580) lp, lse, _ = v2.head_v2_selected_logprob(hs, ws2, ts) - _check("head_v2_small_tail", logprob=lp, lse=lse) + _assert_golden("head_v2_small_tail", logprob=lp, lse=lse) diff --git a/tests/ops/test_bi_head_v2.py b/tests/ops/test_bi_head_v2.py deleted file mode 100644 index 823f1a6a..00000000 --- a/tests/ops/test_bi_head_v2.py +++ /dev/null @@ -1,55 +0,0 @@ -"""Focused gates for the one-launch final token-probability contract.""" - -import pytest -import torch - -from xorl.ops import batch_invariant_ops as v1 -from xorl.ops import bi_families_v2 as v2 -from xorl.ops.loss.bi_fused_lm_head import bi_fused_per_token_ce - - -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - - -def _make(shape: tuple[int, ...], seed: int) -> torch.Tensor: - generator = torch.Generator(device="cpu").manual_seed(seed) - return torch.randn(shape, generator=generator, dtype=torch.float32).to(torch.bfloat16).cuda() - - -@requires_cuda -@pytest.mark.gpu -def test_head_v2_keeps_projection_bits_and_one_stats_tree(): - hidden, weight = _make((16, 128), 1), _make((512, 128), 2) - tokens = torch.arange(16, device="cuda") * 29 % 512 - - logits, decode_lse = v2.head_v2_full_logits_with_lse(hidden, weight) - reference_logits = torch.empty_like(logits) - v1._bi_lm_head_chunk_gemm_fp32(hidden, weight.t(), reference_logits) - assert torch.equal(logits, reference_logits) - - logprob, score_lse, selected = v2.head_v2_selected_logprob(hidden, weight, tokens) - assert torch.equal(decode_lse, score_lse) - assert torch.equal(selected, logits.gather(1, tokens[:, None]).squeeze(1)) - assert torch.equal(logprob, torch.clamp_max(selected - decode_lse, 0.0)) - - -@requires_cuda -@pytest.mark.gpu -def test_head_v2_is_batch_invariant_and_trainable(monkeypatch): - hidden, weight = _make((32, 128), 3), _make((512, 128), 4) - tokens = torch.arange(32, device="cuda") * 31 % 512 - full, _, _ = v2.head_v2_selected_logprob(hidden, weight, tokens) - sub, _, _ = v2.head_v2_selected_logprob(hidden[8:16].contiguous(), weight, tokens[8:16].contiguous()) - assert torch.equal(sub, full[8:16]) - - hidden.requires_grad_(True) - weight.requires_grad_(True) - ce = bi_fused_per_token_ce(hidden, weight, tokens) - assert torch.equal(ce, -full) - ce.mean().backward() - assert torch.isfinite(hidden.grad).all() - assert torch.isfinite(weight.grad).all() - - monkeypatch.setenv("XORL_FAMILIES_V2", "0") - rollback = bi_fused_per_token_ce(hidden.detach(), weight.detach(), tokens, vocab_chunk=256) - assert torch.isfinite(rollback).all() diff --git a/tests/ops/test_bi_mean.py b/tests/ops/test_bi_mean.py deleted file mode 100644 index 21668d32..00000000 --- a/tests/ops/test_bi_mean.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Regression gate: BI-mode aten::mean. - -aten::mean full-reduce dispatches to the mean.dim override with dim=[]; the -empty n_elems product made it return the SUM. Locks: full-reduce == mean, -dim reductions bit-identical to the certified mean_dim kernel path. -""" - -import pytest -import torch - -from xorl.ops.batch_invariant_ops import mean_dim, set_batch_invariant_mode - - -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - - -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("shape", [(512,), (4, 8, 16), (33, 127)]) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -def test_full_reduce_mean_is_mean_not_sum(shape, dtype): - torch.manual_seed(0) - x = torch.randn(shape, device="cuda", dtype=dtype) - ref = x.double().mean() - with set_batch_invariant_mode(True): - got = x.mean() - assert got.shape == torch.Size([]) - tol = 1e-2 if dtype is torch.bfloat16 else 1e-5 - assert abs(got.double().item() - ref.item()) < tol - # the bug returned the full sum - assert abs(got.double().item() - x.double().sum().item()) > tol or x.numel() == 1 - - -@requires_cuda -@pytest.mark.gpu -def test_full_reduce_dtype_kwarg(): - torch.manual_seed(0) - x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) - with set_batch_invariant_mode(True): - got = x.mean(dtype=torch.float32) - assert got.dtype == torch.float32 - assert abs(got.item() - x.double().mean().item()) < 1e-2 - - -@requires_cuda -@pytest.mark.gpu -def test_dim_reductions_unchanged_bitwise(): - torch.manual_seed(0) - x = torch.randn(4, 8, 16, device="cuda", dtype=torch.bfloat16) - with set_batch_invariant_mode(True): - got_1d = x.mean(-1) - got_keep = x.mean(-1, keepdim=True) - got_2d = x.mean(dim=(0, 1)) - assert torch.equal(got_1d, mean_dim(x, 2)) - assert torch.equal(got_keep, mean_dim(x, 2, keepdim=True)) - n = x.shape[0] * x.shape[1] - assert torch.equal(got_2d, torch.sum(x, dim=(0, 1), dtype=torch.float32) / n) diff --git a/tests/ops/test_bi_router_gemm.py b/tests/ops/test_bi_router_gemm.py index f8a72578..d8cf0c30 100644 --- a/tests/ops/test_bi_router_gemm.py +++ b/tests/ops/test_bi_router_gemm.py @@ -19,7 +19,7 @@ def _inputs(n, seed=0): @requires_cuda @pytest.mark.gpu -def test_router_gemm_matches_fp32_upcast_reference(): +def test_router_gemm_linear_topk_and_moe_dispatch_policy(): # bf16xbf16 products are exact in fp32, so the contract kernel must equal an # fp32 GEMM over the (exact) fp32 upcasts up to fp32 reduction-order noise. hidden, weight = _inputs(256) @@ -28,30 +28,20 @@ def test_router_gemm_matches_fp32_upcast_reference(): ref = torch.nn.functional.linear(hidden.float(), weight.float()) assert torch.allclose(out, ref, rtol=1e-5, atol=1e-4) + empty = bi_router_gemm(torch.empty(0, H, device="cuda", dtype=torch.bfloat16), weight) + assert empty.shape == (0, E) and empty.dtype == torch.float32 + with pytest.raises(AssertionError): + bi_router_gemm(hidden.float(), weight) + with pytest.raises(AssertionError): + bi_router_gemm(hidden, weight.float()) -@requires_cuda -@pytest.mark.gpu -def test_router_gemm_batch_invariant_to_padding(): - # A row's logits must not depend on how many other rows share the batch - # (batch invariance): rows of a small batch equal the same rows of a big one. - hidden, weight = _inputs(300) - full = bi_router_gemm(hidden, weight) - sub = bi_router_gemm(hidden[:7].contiguous(), weight) - assert torch.equal(full[:7], sub) - + _assert_router_gemm_autograd_backward_matches_linear() + _assert_bf16_fp32_linear_preserves_leading_dims_and_backward() + _assert_topk_weights_renormalize_or_cast_and_require_fp32() + _assert_moe_block_route_selects_exact_contract_or_default_path() -@requires_cuda -@pytest.mark.gpu -def test_router_gemm_empty_tokens(): - _, weight = _inputs(1) - empty = torch.empty(0, H, device="cuda", dtype=torch.bfloat16) - out = bi_router_gemm(empty, weight) - assert out.shape == (0, E) and out.dtype == torch.float32 - -@requires_cuda -@pytest.mark.gpu -def test_router_gemm_autograd_backward_matches_linear(): +def _assert_router_gemm_autograd_backward_matches_linear(): hidden, weight = _inputs(64, seed=1) h1 = hidden.clone().requires_grad_(True) @@ -67,9 +57,7 @@ def test_router_gemm_autograd_backward_matches_linear(): assert torch.allclose(w1.grad.float(), w2.grad.float(), rtol=2e-2, atol=2e-2) -@requires_cuda -@pytest.mark.gpu -def test_bf16_fp32_linear_preserves_leading_dims_and_backward(): +def _assert_bf16_fp32_linear_preserves_leading_dims_and_backward(): hidden, weight = _inputs(6, seed=2) hidden = hidden.reshape(2, 3, H).requires_grad_(True) weight = weight.requires_grad_(True) @@ -84,19 +72,7 @@ def test_bf16_fp32_linear_preserves_leading_dims_and_backward(): assert weight.grad is not None and weight.grad.dtype is torch.bfloat16 -@requires_cuda -@pytest.mark.gpu -def test_router_gemm_rejects_non_bf16(): - hidden, weight = _inputs(8) - with pytest.raises(AssertionError): - bi_router_gemm(hidden.float(), weight) - with pytest.raises(AssertionError): - bi_router_gemm(hidden, weight.float()) - - -@requires_cuda -@pytest.mark.gpu -def test_topk_weights_fixed_order_renorm(): +def _assert_topk_weights_renormalize_or_cast_and_require_fp32(): torch.manual_seed(0) vals = torch.rand(128, TOP_K, device="cuda", dtype=torch.float32) + 0.1 w = bi_router_topk_weights(vals, norm_topk_prob=True, out_dtype=torch.bfloat16) @@ -106,19 +82,11 @@ def test_topk_weights_fixed_order_renorm(): # no-renorm path is a pure cast w2 = bi_router_topk_weights(vals, norm_topk_prob=False, out_dtype=torch.bfloat16) assert torch.equal(w2, vals.to(torch.bfloat16)) - - -@requires_cuda -@pytest.mark.gpu -def test_topk_weights_requires_fp32(): - vals = torch.rand(4, TOP_K, device="cuda", dtype=torch.bfloat16) with pytest.raises(AssertionError): - bi_router_topk_weights(vals) + bi_router_topk_weights(vals.to(torch.bfloat16)) -@requires_cuda -@pytest.mark.gpu -def test_moe_block_route_uses_structural_exact_contract(): +def _assert_moe_block_route_selects_exact_contract_or_default_path(): # Exact MoEBlock routing must produce # router logits equal to the standalone contract kernel, and selection/ # weights equal to the contract post-processing on those logits. @@ -137,21 +105,27 @@ def test_moe_block_route_uses_structural_exact_contract(): ) hidden = (torch.randn(140, H, device="cuda") * 0.5).to(torch.bfloat16) - rw, sel, logits = block.route(hidden) + with torch.no_grad(): + rw, sel, logits = block.route(hidden) - ref_logits = bi_router_gemm(hidden, block.gate.weight) - assert torch.equal(logits, ref_logits) + ref_logits = bi_router_gemm(hidden, block.gate.weight) + assert torch.equal(logits, ref_logits) - probs = torch.softmax(ref_logits, dim=1, dtype=torch.float) - ref_vals, ref_sel = torch.topk(probs, TOP_K, dim=-1) - ref_w = bi_router_topk_weights(ref_vals, True, torch.bfloat16) - assert torch.equal(sel, ref_sel) - assert torch.equal(rw, ref_w) + probs = torch.softmax(ref_logits, dim=1, dtype=torch.float) + ref_vals, ref_sel = torch.topk(probs, TOP_K, dim=-1) + ref_w = bi_router_topk_weights(ref_vals, True, torch.bfloat16) + assert torch.equal(sel, ref_sel) + assert torch.equal(rw, ref_w) + # Prove batch composition through the production router, not only the + # standalone GEMM helper. + sub_rw, sub_sel, sub_logits = block.route(hidden[:7].contiguous()) + assert torch.equal(sub_logits, logits[:7]) + assert torch.equal(sub_sel, sel[:7]) + assert torch.equal(sub_rw, rw[:7]) + + del block, rw, sel, logits, ref_logits, probs, ref_vals, ref_sel, ref_w -@requires_cuda -@pytest.mark.gpu -def test_moe_block_route_default_path_unchanged(): # Ordinary models retain the stock gate GEMM; logits are bf16, not fp32. block = ( MoEBlock( @@ -165,5 +139,6 @@ def test_moe_block_route_default_path_unchanged(): .to(torch.bfloat16) ) hidden = (torch.randn(16, H, device="cuda") * 0.5).to(torch.bfloat16) - _, _, logits = block.route(hidden) + with torch.no_grad(): + _, _, logits = block.route(hidden) assert logits.dtype == torch.bfloat16 diff --git a/tests/ops/test_bi_trunk_linear.py b/tests/ops/test_bi_trunk_linear.py index 1d6588bf..236e9588 100644 --- a/tests/ops/test_bi_trunk_linear.py +++ b/tests/ops/test_bi_trunk_linear.py @@ -20,6 +20,7 @@ from xorl.ops.batch_invariant_ops import ( is_trunk_linear_contract_enabled, matmul_persistent, + mean_dim, rms_norm_batch_invariant, set_batch_invariant_mode, set_trunk_linear_contract, @@ -64,6 +65,7 @@ def forward(self, ids): @pytest.fixture(autouse=True) def _reset_contract_state(): + set_trunk_linear_contract(False) yield set_trunk_linear_contract(False) @@ -72,17 +74,20 @@ def _reset_contract_state(): # Selection (CPU) # --------------------------------------------------------------------------- # @pytest.mark.cpu -def test_wrap_selection_counts_and_exclusions(): +def test_wrap_selection_and_admission_policy(): model = _TrunkModel(n_layers=2) wrapped = wrap_trunk_linears_batch_invariant(model) assert wrapped == dict.fromkeys(("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"), 2) assert not getattr(model.lm_head, "_xorl_bi_trunk_wrapped", False) assert all(getattr(layer.q_proj, "_xorl_bi_trunk_wrapped", False) for layer in model.layers) - assert not is_trunk_linear_contract_enabled(), "wrapping one model must not mutate process-global dispatch" + assert is_trunk_linear_contract_enabled(), "wrapping a model must arm RMSNorm dispatch for the same contract lane" + _assert_wrap_is_idempotent() + _assert_wrap_skips_routed_experts() + _assert_wrap_rejects_unsupported_module_types() -@pytest.mark.cpu -def test_wrap_is_idempotent(): + +def _assert_wrap_is_idempotent(): model = _TrunkModel(n_layers=1) first = wrap_trunk_linears_batch_invariant(model) assert sum(first.values()) == 7 @@ -90,8 +95,7 @@ def test_wrap_is_idempotent(): assert second == {}, "re-wrap must be a no-op, not a double-wrap" -@pytest.mark.cpu -def test_wrap_skips_routed_experts(): +def _assert_wrap_skips_routed_experts(): model = _TrunkModel(n_layers=1) experts = nn.ModuleList( nn.ModuleDict({"gate_proj": nn.Linear(HIDDEN, INTER, dtype=torch.bfloat16)}) for _ in range(2) @@ -102,45 +106,39 @@ def test_wrap_skips_routed_experts(): assert not getattr(model.layers[0].experts[0].gate_proj, "_xorl_bi_trunk_wrapped", False) -@pytest.mark.cpu -def test_wrap_raises_on_no_match(): - with pytest.raises(RuntimeError, match="matched no trunk linears"): - wrap_trunk_linears_batch_invariant(nn.Linear(4, 4)) - - -@pytest.mark.cpu -def test_wrap_raises_on_lora_wrapped_module(): - model = _TrunkModel(n_layers=1) - model.layers[0].q_proj = LoraLinear(HIDDEN, HIDDEN, r=4, lora_alpha=8, bias=False, dtype=torch.bfloat16) - with pytest.raises(NotImplementedError, match="adapter-wrapped"): - wrap_trunk_linears_batch_invariant(model) - - -@pytest.mark.cpu -def test_wrap_raises_on_linear_subclass(): +def _assert_wrap_rejects_unsupported_module_types(): class _FakeFP8Linear(nn.Linear): pass - model = _TrunkModel(n_layers=1) - model.layers[0].up_proj = _FakeFP8Linear(HIDDEN, INTER, bias=False, dtype=torch.bfloat16) - with pytest.raises(NotImplementedError, match="not a plain"): - wrap_trunk_linears_batch_invariant(model) - - -@pytest.mark.cpu -def test_wrap_raises_on_non_bf16_weight(): - # fp32 is tolerated at wrap time (mixed-precision master; mp_policy casts to - # bf16 before forward and the runtime guard enforces bf16 GEMM operands), but - # fp16 and other dtypes are outside the contract. - model = _TrunkModel(n_layers=1) - model.layers[0].down_proj = nn.Linear(INTER, HIDDEN, bias=False, dtype=torch.float16) - with pytest.raises(RuntimeError, match="bf16-only"): - wrap_trunk_linears_batch_invariant(model) + def lora_model(): + model = _TrunkModel(n_layers=1) + model.layers[0].q_proj = LoraLinear(HIDDEN, HIDDEN, r=4, lora_alpha=8, bias=False, dtype=torch.bfloat16) + return model + + def fp8_subclass_model(): + model = _TrunkModel(n_layers=1) + model.layers[0].up_proj = _FakeFP8Linear(HIDDEN, INTER, bias=False, dtype=torch.bfloat16) + return model + + def fp16_model(): + # fp32 masters are allowed because mixed precision casts before forward; + # fp16 and other runtime operand dtypes are outside the contract. + model = _TrunkModel(n_layers=1) + model.layers[0].down_proj = nn.Linear(INTER, HIDDEN, bias=False, dtype=torch.float16) + return model + + cases = [ + (nn.Linear(4, 4), RuntimeError, "matched no trunk linears"), + (lora_model(), NotImplementedError, "adapter-wrapped"), + (fp8_subclass_model(), NotImplementedError, "not a plain"), + (fp16_model(), RuntimeError, "bf16-only"), + ] + for model, error_type, error_pattern in cases: + with pytest.raises(error_type, match=error_pattern): + wrap_trunk_linears_batch_invariant(model) -@requires_cuda -@pytest.mark.gpu -def test_wrap_raises_under_global_interpose(): +def _assert_wrap_rejects_global_interpose(): model = _TrunkModel(n_layers=1) with set_batch_invariant_mode(True): with pytest.raises(RuntimeError, match="cannot be combined"): @@ -159,19 +157,26 @@ def _wrapped_linear(bias=False, seed=0): @requires_cuda @pytest.mark.gpu -@pytest.mark.parametrize("bias", [False, True]) -def test_forward_bitwise_matches_matmul_persistent(bias): - lin = _wrapped_linear(bias=bias) - x = torch.randn(4, 64, HIDDEN, device="cuda", dtype=torch.bfloat16) - with torch.no_grad(): - out = lin(x) - ref = matmul_persistent(x.reshape(-1, HIDDEN), lin.weight.t(), bias=lin.bias) - assert torch.equal(out, ref.reshape(4, 64, HIDDEN)) +def test_forward_and_backward_bitwise_and_admission_policy(): + for bias in (False, True): + lin = _wrapped_linear(bias=bias) + x = torch.randn(4, 64, HIDDEN, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + out = lin(x) + ref = matmul_persistent(x.reshape(-1, HIDDEN), lin.weight.t(), bias=lin.bias) + assert torch.equal(out, ref.reshape(4, 64, HIDDEN)) + + _assert_forward_matches_global_interpose_lane() + _assert_forward_is_batch_invariant() + _assert_global_interpose_mean_reduction_policy() + _assert_forward_rejects_non_bf16_input() + _assert_wrap_rejects_global_interpose() + _assert_backward_bitwise_matches_cublas_autograd_with_and_without_bias() + set_trunk_linear_contract(False) + _assert_global_interpose_gradient_policy() -@requires_cuda -@pytest.mark.gpu -def test_forward_bitwise_matches_global_interpose_lane(): +def _assert_forward_matches_global_interpose_lane(): # The wrapped forward must produce the SAME bits as F.linear under the global # aten::mm interpose (the serving/verification lane it replaces for training). lin = _wrapped_linear(bias=False, seed=1) @@ -183,9 +188,7 @@ def test_forward_bitwise_matches_global_interpose_lane(): assert torch.equal(out, ref) -@requires_cuda -@pytest.mark.gpu -def test_forward_is_batch_invariant(): +def _assert_forward_is_batch_invariant(): lin = _wrapped_linear(seed=2) x = torch.randn(300, HIDDEN, device="cuda", dtype=torch.bfloat16) with torch.no_grad(): @@ -194,47 +197,71 @@ def test_forward_is_batch_invariant(): assert torch.equal(full[:7], sub) -@requires_cuda -@pytest.mark.gpu -def test_forward_raises_on_non_bf16_input(): +def _assert_forward_rejects_non_bf16_input(): lin = _wrapped_linear(seed=3) with pytest.raises(RuntimeError, match="bf16-only"): lin(torch.randn(8, HIDDEN, device="cuda", dtype=torch.float32)) +def _assert_global_interpose_mean_reduction_policy(): + for shape, dtype in ( + ((512,), torch.bfloat16), + ((512,), torch.float32), + ((33, 127), torch.bfloat16), + ((33, 127), torch.float32), + ): + torch.manual_seed(0) + x = torch.randn(shape, device="cuda", dtype=dtype) + reference = x.double().mean() + with set_batch_invariant_mode(True): + output = x.mean() + tolerance = 1e-2 if dtype is torch.bfloat16 else 1e-5 + assert output.shape == torch.Size([]) + assert abs(output.double().item() - reference.item()) < tolerance + assert abs(output.double().item() - x.double().sum().item()) > tolerance or x.numel() == 1 + + x = torch.randn(4, 8, 16, device="cuda", dtype=torch.bfloat16) + with set_batch_invariant_mode(True): + fp32_output = x.mean(dtype=torch.float32) + output_1d = x.mean(-1) + output_keepdim = x.mean(-1, keepdim=True) + output_2d = x.mean(dim=(0, 1)) + assert fp32_output.dtype is torch.float32 + assert abs(fp32_output.item() - x.double().mean().item()) < 1e-2 + assert torch.equal(output_1d, mean_dim(x, 2)) + assert torch.equal(output_keepdim, mean_dim(x, 2, keepdim=True)) + assert torch.equal(output_2d, torch.sum(x, dim=(0, 1), dtype=torch.float32) / (x.shape[0] * x.shape[1])) + + # --------------------------------------------------------------------------- # # Backward contract (GPU): bitwise vs cuBLAS autograd # --------------------------------------------------------------------------- # -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("bias", [False, True]) -def test_backward_bitwise_matches_cublas_autograd(bias): - lin = _wrapped_linear(bias=bias, seed=4) - x0 = torch.randn(2, 96, HIDDEN, device="cuda", dtype=torch.bfloat16) - g_out = torch.randn(2, 96, HIDDEN, device="cuda", dtype=torch.bfloat16) +def _assert_backward_bitwise_matches_cublas_autograd_with_and_without_bias(): + for bias in (False, True): + lin = _wrapped_linear(bias=bias, seed=4) + x0 = torch.randn(2, 96, HIDDEN, device="cuda", dtype=torch.bfloat16) + g_out = torch.randn(2, 96, HIDDEN, device="cuda", dtype=torch.bfloat16) - x = x0.clone().requires_grad_(True) - out = lin(x) - out.backward(g_out) + x = x0.clone().requires_grad_(True) + out = lin(x) + out.backward(g_out) - x_ref = x0.clone().requires_grad_(True) - w_ref = lin.weight.detach().clone().requires_grad_(True) - b_ref = lin.bias.detach().clone().requires_grad_(True) if bias else None - out_ref = F.linear(x_ref, w_ref, b_ref) - out_ref.backward(g_out) + x_ref = x0.clone().requires_grad_(True) + w_ref = lin.weight.detach().clone().requires_grad_(True) + b_ref = lin.bias.detach().clone().requires_grad_(True) if bias else None + out_ref = F.linear(x_ref, w_ref, b_ref) + out_ref.backward(g_out) - assert torch.equal(x.grad, x_ref.grad), "grad_input must stay bitwise on the cuBLAS path" - assert torch.equal(lin.weight.grad, w_ref.grad), "grad_weight must stay bitwise on the cuBLAS path" - if bias: - assert torch.equal(lin.bias.grad, b_ref.grad) + assert torch.equal(x.grad, x_ref.grad), "grad_input must stay bitwise on the cuBLAS path" + assert torch.equal(lin.weight.grad, w_ref.grad), "grad_weight must stay bitwise on the cuBLAS path" + if bias: + assert torch.equal(lin.bias.grad, b_ref.grad) # --------------------------------------------------------------------------- # # The global interpose loud-fails on training forwards # --------------------------------------------------------------------------- # -@requires_cuda -@pytest.mark.gpu -def test_global_interpose_raises_on_grad_requiring_inputs(): +def _assert_global_interpose_gradient_policy(): x = torch.randn(32, HIDDEN, device="cuda", dtype=torch.bfloat16, requires_grad=True) w = torch.randn(HIDDEN, HIDDEN, device="cuda", dtype=torch.bfloat16, requires_grad=True) wn = torch.randn(HIDDEN, device="cuda", dtype=torch.bfloat16, requires_grad=True) @@ -251,10 +278,11 @@ def test_global_interpose_raises_on_grad_requiring_inputs(): with pytest.raises(RuntimeError, match="XORL_BI_TRUNK_LINEAR"): _ = x.float().mean(-1) + _assert_global_interpose_works_under_no_grad() + _assert_global_interpose_allows_grad_free_inputs() -@requires_cuda -@pytest.mark.gpu -def test_global_interpose_still_works_under_no_grad(): + +def _assert_global_interpose_works_under_no_grad(): torch.manual_seed(5) x = torch.randn(32, HIDDEN, device="cuda", dtype=torch.bfloat16, requires_grad=True) w = torch.randn(HIDDEN, HIDDEN, device="cuda", dtype=torch.bfloat16, requires_grad=True) @@ -267,22 +295,7 @@ def test_global_interpose_still_works_under_no_grad(): assert torch.equal(out_norm, rms_norm_batch_invariant(x, wn, eps=1e-6)) -@requires_cuda -@pytest.mark.gpu -def test_global_interpose_silent_before_now_loud(): - # Regression pin for the guarded bug: without the guard, the aten::rms_norm - # override returned an output disconnected from the autograd graph (q/k-norm grads - # silently vanished). The op must now refuse instead of detaching. - x = torch.randn(16, HIDDEN, device="cuda", dtype=torch.bfloat16, requires_grad=True) - wn = torch.randn(HIDDEN, device="cuda", dtype=torch.bfloat16) - with set_batch_invariant_mode(True): - with pytest.raises(RuntimeError, match="rms_norm"): - F.rms_norm(x, (HIDDEN,), wn, eps=1e-6) - - -@requires_cuda -@pytest.mark.gpu -def test_global_interpose_allows_grad_free_inputs_in_grad_context(): +def _assert_global_interpose_allows_grad_free_inputs(): # Verification flows that forward non-leaf, grad-free tensors inside a # grad-enabled context must keep working. x = torch.randn(32, HIDDEN, device="cuda", dtype=torch.bfloat16) diff --git a/tests/ops/test_block_fp8.py b/tests/ops/test_block_fp8.py deleted file mode 100644 index 3f7fe12b..00000000 --- a/tests/ops/test_block_fp8.py +++ /dev/null @@ -1,182 +0,0 @@ -"""Tests for FP8 block quantization operations. - -These tests verify the correctness of block-based FP8 quantization/dequantization -kernels and FP8 GEMM operations. -""" - -import pytest -import torch - - -# Try to import the block_fp8 module -try: - from xorl.ops.quantize import ( - block_fp8_dequantize as block_fp8_dequant, - ) - from xorl.ops.quantize import ( - block_fp8_dequantize_gkn as block_fp8_weight_dequant, - ) - from xorl.ops.quantize import ( - block_fp8_gemm, - ) - from xorl.ops.quantize import ( - block_fp8_quantize as block_fp8_quant, - ) - - HAS_BLOCK_FP8 = True -except ImportError: - HAS_BLOCK_FP8 = False - -# Skip all tests if block_fp8 is not available or if CUDA is not available -pytestmark = [ - pytest.mark.gpu, - pytest.mark.skipif(not HAS_BLOCK_FP8, reason="block_fp8 module not available"), - pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"), -] - - -class TestBlockFP8Quantization: - """Comprehensive tests for block_fp8 quantization and dequantization.""" - - def test_quantization_shapes_dtypes_and_scales(self): - """Quantization output shapes, dtypes, scale ranges, block sizes, and max value handling.""" - # Basic quantization shape and dtype - x = torch.randn(4, 128, device="cuda", dtype=torch.float32) - y, s = block_fp8_quant(x, block_size=128) - assert y.shape == x.shape - assert y.dtype == torch.float8_e4m3fn - assert s.shape == (4, 1) - assert s.dtype == torch.float32 - - # Multiple blocks - x2 = torch.randn(2, 512, device="cuda", dtype=torch.float32) - y2, s2 = block_fp8_quant(x2, block_size=128) - assert y2.shape == x2.shape - assert s2.shape == (2, 4) - - # Shape preservation across various dimensions - for shape in [(128,), (4, 256), (2, 3, 384)]: - x_s = torch.randn(*shape, device="cuda", dtype=torch.float32) - y_s, s_s = block_fp8_quant(x_s, block_size=128) - assert y_s.shape == x_s.shape - expected_s_shape = (*shape[:-1], shape[-1] // 128) - assert s_s.shape == expected_s_shape - - # Scale range for random normal data - x3 = torch.randn(4, 256, device="cuda", dtype=torch.float32) - _, s3 = block_fp8_quant(x3, block_size=128) - assert torch.all(s3 > 0) - assert torch.all(s3 < 1.0) - assert torch.all(s3 > 1e-6) - - # Max FP8 value (448.0) -> scale should be 1.0 - x_max = torch.full((2, 128), 448.0, device="cuda", dtype=torch.float32) - _, s_max = block_fp8_quant(x_max, block_size=128) - assert torch.allclose(s_max, torch.ones_like(s_max), atol=1e-5) - - # Different block sizes - x4 = torch.randn(2, 512, device="cuda", dtype=torch.float32) - for bs in [64, 128, 256]: - y4, s4 = block_fp8_quant(x4, block_size=bs) - assert y4.shape == x4.shape - assert s4.shape == (2, 512 // bs) - - def test_quantization_input_requirements(self): - """Contiguity and divisibility requirements.""" - # Non-contiguous tensor should fail - x = torch.randn(4, 256, device="cuda", dtype=torch.float32) - with pytest.raises(AssertionError): - block_fp8_quant(x.t(), block_size=128) - - # Non-divisible size should fail - x2 = torch.randn(4, 130, device="cuda", dtype=torch.float32) - with pytest.raises(AssertionError): - block_fp8_quant(x2, block_size=128) - - -class TestBlockFP8Dequantization: - """Comprehensive tests for dequantization accuracy and roundtrip.""" - - def test_dequantization_accuracy_and_roundtrip(self): - """Dequantization accuracy, shape preservation, and roundtrip across shapes.""" - # Basic accuracy - x_orig = torch.randn(4, 256, device="cuda", dtype=torch.float32) - y, s = block_fp8_quant(x_orig, block_size=128) - x_dequant = block_fp8_dequant(y, s, block_size=128) - assert x_dequant.shape == x_orig.shape - assert x_dequant.dtype == torch.float32 - relative_error = torch.abs(x_dequant - x_orig) / (torch.abs(x_orig) + 1e-6) - assert relative_error.mean().item() < 0.05 - - # Roundtrip across shapes - for shape in [(128,), (4, 256), (2, 3, 384)]: - x_s = torch.randn(*shape, device="cuda", dtype=torch.float32) - y_s, s_s = block_fp8_quant(x_s, block_size=128) - x_d = block_fp8_dequant(y_s, s_s, block_size=128) - assert x_d.shape == x_s.shape - assert torch.allclose(x_d, x_s, rtol=0.1, atol=0.05) - - # Non-contiguous dequant should fail - y2, s2 = block_fp8_quant(torch.randn(4, 256, device="cuda", dtype=torch.float32), block_size=128) - with pytest.raises(AssertionError): - block_fp8_dequant(y2.t(), s2, block_size=128) - - -class TestBlockFP8Integration: - """Integration tests: determinism, memory efficiency, edge cases.""" - - def test_determinism_memory_and_edge_cases(self): - """Determinism, memory efficiency, edge cases (small/large/mixed values, single block).""" - # Determinism - x = torch.randn(4, 256, device="cuda", dtype=torch.float32) - y1, s1 = block_fp8_quant(x, block_size=128) - y2, s2 = block_fp8_quant(x, block_size=128) - assert torch.equal(y1, y2) - assert torch.equal(s1, s2) - - # Memory efficiency - M, K = 1024, 2048 - x_fp32 = torch.randn(M, K, device="cuda", dtype=torch.float32) - y_fp8, s = block_fp8_quant(x_fp32, block_size=128) - fp32_bytes = x_fp32.element_size() * x_fp32.numel() - fp8_bytes = y_fp8.element_size() * y_fp8.numel() + s.element_size() * s.numel() - assert fp8_bytes < fp32_bytes - - # Very small values - x_small = torch.full((2, 128), 1e-6, device="cuda", dtype=torch.float32) - y_sm, s_sm = block_fp8_quant(x_small, block_size=128) - x_d_sm = block_fp8_dequant(y_sm, s_sm, block_size=128) - assert torch.allclose(x_d_sm, x_small, rtol=0.5, atol=1e-7) - - # Very large values - x_large = torch.full((2, 128), 400.0, device="cuda", dtype=torch.float32) - y_lg, s_lg = block_fp8_quant(x_large, block_size=128) - x_d_lg = block_fp8_dequant(y_lg, s_lg, block_size=128) - assert torch.allclose(x_d_lg, x_large, rtol=0.05, atol=1.0) - - # Mixed positive/negative - signs preserved - x_mixed = torch.randn(4, 256, device="cuda", dtype=torch.float32) - x_mixed[0] = torch.abs(x_mixed[0]) - x_mixed[1] = -torch.abs(x_mixed[1]) - y_m, s_m = block_fp8_quant(x_mixed, block_size=128) - x_d_m = block_fp8_dequant(y_m, s_m, block_size=128) - assert torch.all((x_mixed >= 0) == (x_d_m >= 0)) - - # Single block (minimum size) - x_single = torch.randn(1, 128, device="cuda", dtype=torch.float32) - y_sg, s_sg = block_fp8_quant(x_single, block_size=128) - assert y_sg.shape == (1, 128) - assert s_sg.shape == (1, 1) - x_d_sg = block_fp8_dequant(y_sg, s_sg, block_size=128) - assert torch.allclose(x_d_sg, x_single, rtol=0.1, atol=0.05) - - # 1D/2D consistency - x_2d = torch.randn(256, 512, device="cuda", dtype=torch.float32) - y_1d, s_1d = block_fp8_quant(x_2d, block_size=128) - x_d_1d = block_fp8_dequant(y_1d, s_1d, block_size=128) - assert x_d_1d.shape == x_2d.shape - assert torch.allclose(x_d_1d, x_2d, rtol=0.1, atol=0.05) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/ops/test_block_fp8_gkn.py b/tests/ops/test_block_fp8_gkn.py deleted file mode 100644 index 10a837bd..00000000 --- a/tests/ops/test_block_fp8_gkn.py +++ /dev/null @@ -1,159 +0,0 @@ -"""Tests for block_fp8 GKN (2D weight) quantization kernels. - -Correctness tests and bandwidth benchmarks for block_fp8_quantize_gkn -and block_fp8_dequantize_gkn. -""" - -import pytest -import torch -import triton - -from xorl.ops.quantize import block_fp8_dequantize_gkn, block_fp8_quantize_gkn - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -class TestBlockFP8QuantizeGKN: - """Correctness tests for block_fp8_quantize_gkn: shapes, accuracy, edge cases.""" - - def test_quantize_shapes_accuracy_and_edge_cases(self): - """Output shapes, roundtrip accuracy (f32/bf16), non-divisible shapes, zeros, scale ranges, large matrix.""" - # Basic output shapes - K, N = 512, 256 - x = torch.randn(K, N, device="cuda", dtype=torch.float32) - y, s = block_fp8_quantize_gkn(x, block_size=128) - assert y.shape == (K, N) - assert y.dtype == torch.float8_e4m3fn - assert s.shape == (triton.cdiv(K, 128), triton.cdiv(N, 128)) - assert s.dtype == torch.float32 - - # Roundtrip accuracy (f32) - x_deq = block_fp8_dequantize_gkn(y, s, block_size=128) - rel_err = (x - x_deq).abs().mean() / x.abs().mean() - assert rel_err < 0.03, f"f32 roundtrip rel error {rel_err:.4f} too high" - - # Roundtrip accuracy (bf16 input) - K2, N2 = 256, 256 - x_bf = torch.randn(K2, N2, device="cuda", dtype=torch.bfloat16) - y_bf, s_bf = block_fp8_quantize_gkn(x_bf.float(), block_size=128) - x_deq_bf = block_fp8_dequantize_gkn(y_bf, s_bf, block_size=128) - rel_err_bf = (x_bf.float() - x_deq_bf).abs().mean() / x_bf.float().abs().mean() - assert rel_err_bf < 0.03 - - # Non-divisible shapes (one dim) - K3, N3 = 384, 256 - x3 = torch.randn(K3, N3, device="cuda", dtype=torch.float32) - y3, s3 = block_fp8_quantize_gkn(x3, block_size=128) - x3_deq = block_fp8_dequantize_gkn(y3, s3, block_size=128) - assert (x3 - x3_deq).abs().mean() / x3.abs().mean() < 0.03 - - # Non-divisible both dims - K4, N4 = 300, 200 - x4 = torch.randn(K4, N4, device="cuda", dtype=torch.float32) - y4, s4 = block_fp8_quantize_gkn(x4, block_size=128) - assert y4.shape == (K4, N4) - assert s4.shape == (triton.cdiv(K4, 128), triton.cdiv(N4, 128)) - x4_deq = block_fp8_dequantize_gkn(y4, s4, block_size=128) - assert (x4 - x4_deq).abs().mean() / x4.abs().mean() < 0.03 - - # Zero block - x_z = torch.zeros(128, 128, device="cuda", dtype=torch.float32) - y_z, s_z = block_fp8_quantize_gkn(x_z, block_size=128) - x_z_deq = block_fp8_dequantize_gkn(y_z, s_z, block_size=128) - assert (x_z_deq == 0).all() - - # Scale value ranges - x5 = torch.randn(256, 256, device="cuda", dtype=torch.float32) * 3.0 - _, s5 = block_fp8_quantize_gkn(x5, block_size=128) - assert (s5 > 0).all() - assert s5.max().item() < 1.0 - assert s5.min().item() > 1e-12 - - # Large matrix - K6, N6 = 4096, 4096 - x6 = torch.randn(K6, N6, device="cuda", dtype=torch.float32) - y6, s6 = block_fp8_quantize_gkn(x6, block_size=128) - x6_deq = block_fp8_dequantize_gkn(y6, s6, block_size=128) - assert (x6 - x6_deq).abs().mean() / x6.abs().mean() < 0.03 - - -class TestBlockFP8DequantizeGKN: - """Correctness tests for block_fp8_dequantize_gkn: shape, dtype, input requirements.""" - - def test_dequantize_output_and_requirements(self): - """Output shape/dtype, contiguity requirement, 2D requirement.""" - K, N = 256, 512 - x = torch.randn(K, N, device="cuda", dtype=torch.float32) - y, s = block_fp8_quantize_gkn(x, block_size=128) - x_deq = block_fp8_dequantize_gkn(y, s, block_size=128) - assert x_deq.shape == (K, N) - assert x_deq.dtype == torch.get_default_dtype() - - # Non-contiguous input should fail - K2, N2 = 256, 256 - x2 = torch.randn(K2, N2, device="cuda", dtype=torch.float32) - y2, s2 = block_fp8_quantize_gkn(x2, block_size=128) - with pytest.raises(AssertionError): - block_fp8_dequantize_gkn(y2.t(), s2, block_size=128) - - # Non-2D input should fail - with pytest.raises(AssertionError): - block_fp8_dequantize_gkn( - torch.zeros(8, 128, 128, device="cuda", dtype=torch.float8_e4m3fn), - torch.ones(8, 1, 1, device="cuda", dtype=torch.float32), - block_size=128, - ) - - -# --------------------------------------------------------------------------- -# Bandwidth benchmarks -# --------------------------------------------------------------------------- - - -@pytest.mark.benchmark -class TestBlockFP8GKNBandwidth: - """Bandwidth tests targeting >2000 GB/s on H100.""" - - def test_quantize_dequantize_gkn_bandwidth(self): - """Measure quantize and dequantize bandwidth.""" - K, N = 4096, 4096 - x = torch.randn(K, N, device="cuda", dtype=torch.float32) - - # --- Quantize bandwidth --- - for _ in range(10): - block_fp8_quantize_gkn(x) - torch.cuda.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(100): - block_fp8_quantize_gkn(x) - end.record() - torch.cuda.synchronize() - elapsed_ms = start.elapsed_time(end) / 100 - total_bytes = K * N * (4 + 1) - bw_quant = total_bytes / (elapsed_ms * 1e-3) / 1e9 - print(f"\nblock_fp8_quantize_gkn: {bw_quant:.0f} GB/s ({elapsed_ms * 1000:.1f} us)") - - # --- Dequantize bandwidth --- - y, s = block_fp8_quantize_gkn(x) - for _ in range(10): - block_fp8_dequantize_gkn(y, s) - torch.cuda.synchronize() - start2 = torch.cuda.Event(enable_timing=True) - end2 = torch.cuda.Event(enable_timing=True) - start2.record() - for _ in range(100): - block_fp8_dequantize_gkn(y, s) - end2.record() - torch.cuda.synchronize() - elapsed_ms2 = start2.elapsed_time(end2) / 100 - total_bytes2 = K * N * (1 + 4) - bw_dequant = total_bytes2 / (elapsed_ms2 * 1e-3) / 1e9 - print(f"\nblock_fp8_dequantize_gkn: {bw_dequant:.0f} GB/s ({elapsed_ms2 * 1000:.1f} us)") - if bw_quant < 2000 or bw_dequant < 2000: - pytest.skip( - f"Bandwidth below 2000 GB/s target " - f"(quant={bw_quant:.0f} GB/s, dequant={bw_dequant:.0f} GB/s; may vary by GPU)" - ) diff --git a/tests/ops/test_block_fp8_native.py b/tests/ops/test_block_fp8_native.py index 78e9b5e1..ac2ef594 100644 --- a/tests/ops/test_block_fp8_native.py +++ b/tests/ops/test_block_fp8_native.py @@ -4,13 +4,11 @@ import torch import torch.distributed.checkpoint as dcp -from xorl.distributed.torch_parallelize import _expert_fsdp_kwargs_for_module from xorl.ops.block_fp8_native import ( NativeBlockFP8Linear, pack_fp8_as_float32, unpack_float32_as_fp8, validate_native_fp8_dcp_checkpoint, - validate_native_fp8_state_metadata, ) @@ -20,7 +18,7 @@ def _fp8_values(shape): return values.to(torch.float8_e4m3fn) -def test_pack_roundtrip_is_byte_exact(): +def _assert_pack_roundtrip_is_byte_exact(): weight = _fp8_values((256, 128)) packed = pack_fp8_as_float32(weight) restored = unpack_float32_as_fp8(packed, tuple(weight.shape)) @@ -30,7 +28,7 @@ def test_pack_roundtrip_is_byte_exact(): assert torch.equal(restored.view(torch.uint8), weight.contiguous().view(torch.uint8)) -def test_linear_state_is_frozen_reshardable_and_scale_exact(): +def _assert_linear_state_survives_dtype_apply_byte_exactly(): module = NativeBlockFP8Linear(256, 384) weight = _fp8_values((384, 256)) scales = torch.arange(6, dtype=torch.float32).reshape(3, 2) / 7 @@ -42,13 +40,6 @@ def test_linear_state_is_frozen_reshardable_and_scale_exact(): assert module.fsdp_requires_full_precision is True assert torch.equal(module.fp8_weight().view(torch.uint8), weight.view(torch.uint8)) assert torch.equal(module.weight_scale_inv.view(torch.uint8), scales.view(torch.uint8)) - - -def test_dtype_apply_preserves_packed_and_scale_bytes(): - module = NativeBlockFP8Linear(128, 128) - weight = _fp8_values((128, 128)) - scales = torch.tensor([[0.1234567]], dtype=torch.float32) - module.load_prequantized(weight, scales) weight_bytes = module.packed_weight_f32.view(torch.uint8).clone() scale_bytes = module.weight_scale_inv.view(torch.uint8).clone() @@ -60,7 +51,7 @@ def test_dtype_apply_preserves_packed_and_scale_bytes(): assert torch.equal(module.weight_scale_inv.view(torch.uint8), scale_bytes) -def test_no_eager_sglang_import_and_cpu_forward_fails_closed(): +def _assert_cpu_execution_and_materialization_fail_closed_without_eager_sglang_import(): before = {name for name in sys.modules if name == "sglang" or name.startswith("sglang.")} module = NativeBlockFP8Linear(128, 128) after = {name for name in sys.modules if name == "sglang" or name.startswith("sglang.")} @@ -68,10 +59,6 @@ def test_no_eager_sglang_import_and_cpu_forward_fails_closed(): assert after == before with pytest.raises(RuntimeError, match="requires CUDA"): module(torch.zeros(1, 128, dtype=torch.bfloat16)) - - -def test_weight_materialization_is_an_explicit_input_free_cuda_path(): - module = NativeBlockFP8Linear(128, 128) with pytest.raises(ValueError, match="does not accept activation or range inputs"): module(torch.zeros(1, 128, dtype=torch.bfloat16), return_dequantized_weight=True) with pytest.raises(ValueError, match="does not accept activation or range inputs"): @@ -80,7 +67,7 @@ def test_weight_materialization_is_an_explicit_input_free_cuda_path(): module(return_dequantized_weight=True) -def test_partition_ranges_cross_the_module_forward_hook_boundary(monkeypatch): +def _assert_partition_ranges_cross_the_module_forward_hook_boundary(monkeypatch): module = NativeBlockFP8Linear(256, 384) input = torch.zeros(2, 128, dtype=torch.bfloat16) expected = torch.ones(2, 128, dtype=torch.bfloat16) @@ -101,7 +88,7 @@ def fake_forward_partition(value, *, output_range=None, input_range=None): assert calls == [(input, (128, 256), (0, 128))] -def test_phase_one_forward_rejects_activation_or_base_gradients(): +def _assert_phase_one_forward_rejects_activation_or_base_gradients(): module = NativeBlockFP8Linear(128, 128) with pytest.raises(RuntimeError, match="scoring-only"): module(torch.zeros(1, 128, dtype=torch.bfloat16, requires_grad=True)) @@ -111,7 +98,7 @@ def test_phase_one_forward_rejects_activation_or_base_gradients(): module(torch.zeros(1, 128, dtype=torch.bfloat16)) -def test_linear_partition_contract_fails_before_kernel_dispatch(): +def _assert_linear_partition_contract_fails_before_kernel_dispatch(): module = NativeBlockFP8Linear(256, 384) input = torch.zeros(1, 128, dtype=torch.bfloat16) @@ -121,7 +108,7 @@ def test_linear_partition_contract_fails_before_kernel_dispatch(): module.forward_partition(input, input_range=(0, 256)) -def test_rejects_wrong_pair_dtype_shape_and_nonfinite_scale(): +def _assert_rejects_wrong_pair_dtype_shape_and_nonfinite_scale(): module = NativeBlockFP8Linear(128, 128) weight = _fp8_values((128, 128)) @@ -135,7 +122,7 @@ def test_rejects_wrong_pair_dtype_shape_and_nonfinite_scale(): module.load_prequantized(_fp8_values((128, 256)), torch.ones(1, 1, dtype=torch.float32)) -def test_state_dict_roundtrip_retains_both_parameter_byte_streams(): +def _assert_state_dict_roundtrip_retains_both_parameter_byte_streams(): source = NativeBlockFP8Linear(256, 384) target = NativeBlockFP8Linear(256, 384) weight = _fp8_values((384, 256)) @@ -149,22 +136,15 @@ def test_state_dict_roundtrip_retains_both_parameter_byte_streams(): assert torch.equal(target.weight_scale_inv.view(torch.uint8), source.weight_scale_inv.view(torch.uint8)) -def test_state_dict_and_dcp_metadata_reject_castable_payloads(): +def _assert_state_dict_rejects_castable_payloads(): module = NativeBlockFP8Linear(128, 128) state = module.state_dict() state["packed_weight_f32"] = state["packed_weight_f32"].to(torch.bfloat16) with pytest.raises(TypeError, match="refusing a load_state_dict cast"): module.load_state_dict(state, strict=True) - good_metadata = {name: (parameter.dtype, tuple(parameter.shape)) for name, parameter in module.named_parameters()} - validate_native_fp8_state_metadata(module, good_metadata) - bad_metadata = dict(good_metadata) - bad_metadata["weight_scale_inv"] = (torch.bfloat16, tuple(module.weight_scale_inv.shape)) - with pytest.raises(ValueError, match="DCP metadata mismatch"): - validate_native_fp8_state_metadata(module, bad_metadata) - -def test_apply_exception_never_strands_protected_parameters(): +def _assert_apply_exception_never_strands_protected_parameters(): module = NativeBlockFP8Linear(128, 128) module.child = torch.nn.Linear(1, 1) original = dict(module.named_parameters()) @@ -182,7 +162,7 @@ def fail_on_nonempty(tensor): assert all(restored[name] is parameter for name, parameter in original.items()) -def test_real_dcp_metadata_is_checked_before_load(tmp_path): +def _assert_real_dcp_metadata_is_checked_before_load(tmp_path): module = NativeBlockFP8Linear(128, 128) good_path = tmp_path / "good" dcp.save({"model": module.state_dict()}, checkpoint_id=good_path) @@ -197,7 +177,7 @@ def test_real_dcp_metadata_is_checked_before_load(tmp_path): validate_native_fp8_dcp_checkpoint(str(bad_path), module.state_dict()) -def test_dcp_preflight_uses_ep_restored_expected_shape(tmp_path): +def _assert_dcp_preflight_uses_ep_restored_expected_shape(tmp_path): # Simulate ModelState.state_dict(): live expert params may have E_local=2, # while the expected DCP view restores the global E=4 dimension. restored_state = { @@ -212,13 +192,16 @@ def test_dcp_preflight_uses_ep_restored_expected_shape(tmp_path): validate_native_fp8_dcp_checkpoint(str(checkpoint_path), restored_state) -def test_generic_expert_fsdp_helper_removes_mixed_precision_without_skipping_fsdp(): - class FullPrecisionExpertState: - fsdp_requires_full_precision = True - - kwargs = _expert_fsdp_kwargs_for_module( - {"mesh": "mesh", "mp_policy": "bf16", "reshard_after_forward": True}, - FullPrecisionExpertState(), - ) - - assert kwargs == {"mesh": "mesh", "reshard_after_forward": True} +def test_native_block_fp8_encoding_checkpoint_execution_and_admission_contract(tmp_path, monkeypatch): + _assert_pack_roundtrip_is_byte_exact() + _assert_linear_state_survives_dtype_apply_byte_exactly() + _assert_state_dict_roundtrip_retains_both_parameter_byte_streams() + _assert_state_dict_rejects_castable_payloads() + _assert_apply_exception_never_strands_protected_parameters() + _assert_real_dcp_metadata_is_checked_before_load(tmp_path) + _assert_dcp_preflight_uses_ep_restored_expected_shape(tmp_path) + _assert_cpu_execution_and_materialization_fail_closed_without_eager_sglang_import() + _assert_partition_ranges_cross_the_module_forward_hook_boundary(monkeypatch) + _assert_phase_one_forward_rejects_activation_or_base_gradients() + _assert_linear_partition_contract_fails_before_kernel_dispatch() + _assert_rejects_wrong_pair_dtype_shape_and_nonfinite_scale() diff --git a/tests/ops/test_eager_vs_native_moe.py b/tests/ops/test_eager_vs_native_moe.py index 758aea12..887d3948 100644 --- a/tests/ops/test_eager_vs_native_moe.py +++ b/tests/ops/test_eager_vs_native_moe.py @@ -20,14 +20,11 @@ def _import_moe(): - """Import MoE layers; returns (MoEBlock, MoEExperts) or skips.""" - try: - from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 - from xorl.models.layers.moe.moe_block import MoEBlock # noqa: PLC0415 + """Import MoE layers lazily so CPU-only collection stays lightweight.""" + from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 + from xorl.models.layers.moe.moe_block import MoEBlock # noqa: PLC0415 - return MoEBlock, MoEExperts - except Exception as e: - pytest.skip(f"Cannot import MoE layers: {e}") + return MoEBlock, MoEExperts def _make_pair(num_experts, hidden_dim, intermediate, top_k, seed=42): @@ -56,22 +53,18 @@ def _make_pair(num_experts, hidden_dim, intermediate, top_k, seed=42): # Test 1: Forward + backward agreement across all configs # --------------------------------------------------------------------------- -ALL_CONFIGS = [ +PARITY_CONFIGS = [ # (num_experts, hidden_dim, intermediate, top_k, batch, seq) (4, 64, 128, 2, 2, 8), - (8, 128, 256, 2, 4, 16), (4, 64, 128, 1, 2, 8), # top_k=1 (8, 128, 256, 4, 2, 16), # top_k=4 - (16, 64, 128, 2, 1, 4), # 16 experts (4, 64, 128, 2, 1, 1), # minimal seq ] -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("ne,hd,inter,topk,bs,seq", ALL_CONFIGS) -def test_forward_and_backward_agreement(ne, hd, inter, topk, bs, seq): - """Eager and native forward outputs and backward gradients should match.""" +def _assert_forward_and_backward_agreement(ne, hd, inter, topk, bs, seq): eager_block, native_block = _make_pair(ne, hd, inter, topk) + config = f"E={ne}, H={hd}, I={inter}, top_k={topk}, batch={bs}, seq={seq}" # --- Forward agreement --- torch.manual_seed(999) @@ -87,35 +80,55 @@ def test_forward_and_backward_agreement(ne, hd, inter, topk, bs, seq): eager_out, atol=0.05, rtol=0.02, - msg=f"Forward mismatch: max_diff={max_diff:.6f}", + msg=f"Forward mismatch ({config}): max_diff={max_diff:.6f}", ) - # --- Backward agreement (for larger configs) --- - if ne <= 8 and hd >= 64: - torch.manual_seed(999) - x_eager = torch.randn(bs, seq, hd, device=DEVICE, dtype=DTYPE, requires_grad=True) - x_native = x_eager.detach().clone().requires_grad_(True) - - eager_out2, _ = eager_block(x_eager) - eager_out2.sum().backward() - native_out2, _ = native_block(x_native) - native_out2.sum().backward() - - atol, rtol = 0.05, 0.05 - torch.testing.assert_close(x_native.grad, x_eager.grad, atol=atol, rtol=rtol, msg="Input gradient mismatch") - for name in ["gate_proj", "up_proj", "down_proj"]: - eager_grad = getattr(eager_block.experts, name).grad - native_grad = getattr(native_block.experts, name).grad - assert eager_grad is not None, f"eager {name} grad is None" - assert native_grad is not None, f"native {name} grad is None" - torch.testing.assert_close(native_grad, eager_grad, atol=atol, rtol=rtol, msg=f"{name} gradient mismatch") + # --- Backward agreement --- + torch.manual_seed(999) + x_eager = torch.randn(bs, seq, hd, device=DEVICE, dtype=DTYPE, requires_grad=True) + x_native = x_eager.detach().clone().requires_grad_(True) + + eager_out2, _ = eager_block(x_eager) + eager_out2.sum().backward() + native_out2, _ = native_block(x_native) + native_out2.sum().backward() + + atol, rtol = 0.05, 0.05 + torch.testing.assert_close( + x_native.grad, + x_eager.grad, + atol=atol, + rtol=rtol, + msg=f"Input gradient mismatch ({config})", + ) + for name in ["gate_proj", "up_proj", "down_proj"]: + eager_grad = getattr(eager_block.experts, name).grad + native_grad = getattr(native_block.experts, name).grad + assert eager_grad is not None, f"eager {name} grad is None ({config})" + assert native_grad is not None, f"native {name} grad is None ({config})" torch.testing.assert_close( - native_block.gate.weight.grad, - eager_block.gate.weight.grad, + native_grad, + eager_grad, atol=atol, rtol=rtol, - msg="Gate weight gradient mismatch", + msg=f"{name} gradient mismatch ({config})", ) + torch.testing.assert_close( + native_block.gate.weight.grad, + eager_block.gate.weight.grad, + atol=atol, + rtol=rtol, + msg=f"Gate weight gradient mismatch ({config})", + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_forward_and_backward_agreement(): + """Eager and native agree at routing and token-count boundaries.""" + for config in PARITY_CONFIGS: + _assert_forward_and_backward_agreement(*config) + + _assert_determinism_and_edge_cases() # --------------------------------------------------------------------------- @@ -123,8 +136,7 @@ def test_forward_and_backward_agreement(ne, hd, inter, topk, bs, seq): # --------------------------------------------------------------------------- -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_determinism_and_edge_cases(): +def _assert_determinism_and_edge_cases(): """Determinism: same input produces identical output. Edge case: all tokens to same expert.""" MoEBlock, MoEExperts = _import_moe() @@ -179,51 +191,3 @@ def test_determinism_and_edge_cases(): rtol=0.01, msg="Same-expert output mismatch", ) - - -# --------------------------------------------------------------------------- -# Test 3: Large scale forward + backward -# --------------------------------------------------------------------------- - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_large_scale(): - """Forward + backward agreement at larger dimensions (E=64, H=512, I=1024, K=8).""" - ne, hd, inter, topk = 64, 512, 1024, 8 - eager_block, native_block = _make_pair(ne, hd, inter, topk) - - # --- Forward --- - torch.manual_seed(42) - x = torch.randn(1, 128, hd, device=DEVICE, dtype=DTYPE) - with torch.no_grad(): - eager_out, _ = eager_block(x) - native_out, _ = native_block(x) - torch.testing.assert_close( - native_out, - eager_out, - atol=0.1, - rtol=0.05, - msg="Large scale forward mismatch", - ) - - # --- Backward --- - torch.manual_seed(42) - x_eager = torch.randn(1, 128, hd, device=DEVICE, dtype=DTYPE, requires_grad=True) - x_native = x_eager.detach().clone().requires_grad_(True) - - eager_out2, _ = eager_block(x_eager) - eager_out2.sum().backward() - native_out2, _ = native_block(x_native) - native_out2.sum().backward() - - atol, rtol = 0.1, 0.1 - torch.testing.assert_close(x_native.grad, x_eager.grad, atol=atol, rtol=rtol, msg="Large scale input grad mismatch") - - for name in ["gate_proj", "up_proj", "down_proj"]: - eg = getattr(eager_block.experts, name).grad - ng = getattr(native_block.experts, name).grad - torch.testing.assert_close(ng, eg, atol=atol, rtol=rtol, msg=f"Large scale {name} gradient mismatch") - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/ops/test_ep_adapter_wrappers.py b/tests/ops/test_ep_adapter_wrappers.py index f224b0d6..27f9d329 100644 --- a/tests/ops/test_ep_adapter_wrappers.py +++ b/tests/ops/test_ep_adapter_wrappers.py @@ -6,198 +6,28 @@ expert_scores through, and compare output to a naive reference. """ -import importlib.util import inspect -import sys -import types -from pathlib import Path import pytest import torch import torch.nn.functional as F import xorl.models.layers.moe.backend as moe_backend +from tests._helpers.moe import counts_from_cumsum, patch_ep_kernels from xorl.models.layers.moe.backend import EP_EXPERT_COMPUTE, EP_EXPERT_COMPUTE_MOE_ACT -from xorl.utils import import_utils +from xorl.ops.moe import quack as quack_moe +from xorl.ops.moe import triton as triton_moe pytestmark = pytest.mark.cpu -_MODULE_PATHS = { - "xorl.ops.moe.triton": Path(__file__).resolve().parents[2] / "src/xorl/ops/moe/triton.py", - "xorl.ops.moe.quack": Path(__file__).resolve().parents[2] / "src/xorl/ops/moe/quack.py", -} -_BACKEND_INIT_PATH = Path(__file__).resolve().parents[2] / "src/xorl/models/layers/moe/backend/__init__.py" -_BACKEND_PACKAGE = "xorl.models.layers.moe.backend" - - -def _counts_from_cumsum(cumsum: torch.Tensor) -> list[int]: - counts = torch.empty_like(cumsum) - counts[0] = cumsum[0] - counts[1:] = cumsum[1:] - cumsum[:-1] - return counts.tolist() - - -def _naive_group_gemm_same_nk(a, b, cumsum_M, max_M, transpose_a=False, transpose_b=False, **kwargs): - del max_M, kwargs - assert not transpose_a - outputs = [] - start = 0 - for expert_idx, count in enumerate(_counts_from_cumsum(cumsum_M)): - end = start + count - weight = b[expert_idx] - if transpose_b: - outputs.append(a[start:end] @ weight.transpose(0, 1)) - else: - outputs.append(a[start:end] @ weight) - start = end - return torch.cat(outputs, dim=0) - - -def _naive_group_gemm_same_mn(a, b, c, cumsum_K, max_K, transpose_a=False, transpose_b=False, **kwargs): - del max_K, kwargs - start = 0 - for expert_idx, count in enumerate(_counts_from_cumsum(cumsum_K)): - end = start + count - lhs = a[start:end].transpose(0, 1) if transpose_a else a[start:end] - rhs = b[start:end].transpose(0, 1) if transpose_b else b[start:end] - c[expert_idx].copy_(lhs @ rhs) - start = end - return c - - -def reference_ep_forward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores): - """Naive reference: per-expert matmul with SiLU + optional score scaling.""" - outputs = [] - start = 0 - for expert_idx, count in enumerate(_counts_from_cumsum(cumsum)): - end = start + count - x = permute_tokens[start:end] - h = F.silu(x @ gate_proj[expert_idx]) * (x @ up_proj[expert_idx]) - if expert_scores is not None: - h = h * expert_scores[start:end].to(h.dtype).unsqueeze(-1) - outputs.append(h @ down_proj[expert_idx]) - start = end - return torch.cat(outputs, dim=0) - - -def _patch_kernels_and_load_backend(monkeypatch, backend_type: str): - """Patch kernel modules and load backend/__init__.py to get the adapter wrappers.""" - moe_stub = types.ModuleType("xorl.ops.group_gemm.kernel.moe") - moe_stub.expert_histogram = None - moe_stub.moe_gather = None - moe_stub.moe_index_compute = None - moe_stub.moe_scatter = None - monkeypatch.setattr(import_utils, "is_fused_moe_available", lambda: True) - sys.modules.pop("xorl.ops.group_gemm.kernel.moe", None) - sys.modules.pop("xorl.ops.group_gemm.kernel.group_gemm", None) - sys.modules.pop("xorl.ops.group_gemm.kernel.quack", None) - monkeypatch.setitem(sys.modules, "xorl.ops.group_gemm.kernel.moe", moe_stub) - - if backend_type == "triton": - group_gemm_stub = types.ModuleType("xorl.ops.group_gemm.kernel.group_gemm") - group_gemm_stub.group_gemm_same_nk = _naive_group_gemm_same_nk - group_gemm_stub.group_gemm_same_mn = _naive_group_gemm_same_mn - monkeypatch.setitem(sys.modules, "xorl.ops.group_gemm.kernel.group_gemm", group_gemm_stub) - module_name = "xorl.ops.moe.triton" - else: - quack_stub = types.ModuleType("xorl.ops.group_gemm.kernel.quack") - quack_stub.cumsum_to_cu_seqlens = lambda cumsum: cumsum - quack_stub.quack_group_gemm_same_nk = _naive_group_gemm_same_nk - quack_stub.quack_group_gemm_same_mn = _naive_group_gemm_same_mn - monkeypatch.setitem(sys.modules, "xorl.ops.group_gemm.kernel.quack", quack_stub) - module_name = "xorl.ops.moe.quack" - - spec = importlib.util.spec_from_file_location(f"adapter_test_{backend_type}", _MODULE_PATHS[module_name]) - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - spec.loader.exec_module(module) - monkeypatch.setitem(sys.modules, module_name, module) - - return module - - -def _make_test_data(dtype=torch.float32): - torch.manual_seed(42) - num_local_experts = 2 - hidden_dim = 8 - intermediate_size = 12 - counts = torch.tensor([3, 2]) - cumsum = torch.cumsum(counts, dim=0) - num_tokens = int(cumsum[-1].item()) - - permute_tokens = torch.randn(num_tokens, hidden_dim, dtype=dtype) - gate_proj = torch.randn(num_local_experts, hidden_dim, intermediate_size, dtype=dtype) - up_proj = torch.randn(num_local_experts, hidden_dim, intermediate_size, dtype=dtype) - down_proj = torch.randn(num_local_experts, intermediate_size, hidden_dim, dtype=dtype) - gate_up_proj = torch.cat([gate_proj, up_proj], dim=-1) - expert_scores = torch.rand(num_tokens, dtype=dtype) - - return permute_tokens, cumsum, gate_proj, up_proj, down_proj, gate_up_proj, intermediate_size, expert_scores - - -def test_quack_ep_registers_without_optional_moe_act(monkeypatch): +def _assert_quack_ep_registers_without_optional_moe_act(): """The base Quack EP path must not depend on the optional MoE-act class.""" - for module_name in list(sys.modules): - if module_name == _BACKEND_PACKAGE or module_name.startswith(f"{_BACKEND_PACKAGE}."): - monkeypatch.delitem(sys.modules, module_name, raising=False) - - quack_stub = types.ModuleType("xorl.ops.moe.quack") - - class QuackEPGroupGemm: - @staticmethod - def apply(*args, **kwargs): - return args, kwargs - - quack_stub.QuackEPGroupGemm = QuackEPGroupGemm - monkeypatch.setitem(sys.modules, "xorl.ops.moe.quack", quack_stub) - - spec = importlib.util.spec_from_file_location( - _BACKEND_PACKAGE, - _BACKEND_INIT_PATH, - submodule_search_locations=[str(_BACKEND_INIT_PATH.parent)], - ) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, _BACKEND_PACKAGE, module) - assert spec.loader is not None - spec.loader.exec_module(module) - - assert "quack" in module.EP_EXPERT_COMPUTE - assert "quack" not in module.EP_EXPERT_COMPUTE_MOE_ACT - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton/Quack EP kernels require CUDA") -@pytest.mark.parametrize( - ("backend_type", "class_name"), - [ - pytest.param("triton", "TritonEPGroupGemm", id="triton-ep"), - pytest.param("quack", "QuackEPGroupGemm", id="quack-ep"), - ], -) -def test_adapter_forwards_expert_scores(monkeypatch, backend_type, class_name): - """Verify that EP adapter wrappers accept and forward expert_scores.""" - try: - kernel_module = _patch_kernels_and_load_backend(monkeypatch, backend_type) - except ImportError as exc: - pytest.skip(f"{backend_type} unavailable: {exc}") - - ( - permute_tokens, - cumsum, - gate_proj, - up_proj, - down_proj, - gate_up_proj, - intermediate_size, - expert_scores, - ) = _make_test_data() - - fn_cls = getattr(kernel_module, class_name) - output = fn_cls.apply(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores) - - ref = reference_ep_forward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores) - torch.testing.assert_close(output, ref) + assert hasattr(moe_backend, "_QuackEPGroupGemm") + assert not hasattr(moe_backend, "_QuackEPGroupGemmMoeAct") + assert "quack" in EP_EXPERT_COMPUTE + assert "quack" not in EP_EXPERT_COMPUTE_MOE_ACT # --------------------------------------------------------------------------- @@ -210,32 +40,33 @@ def test_adapter_forwards_expert_scores(monkeypatch, backend_type, class_name): _REQUIRED_EP_PARAMS = ("expert_scores", "hidden_act", "activation_native", "swiglu_limit") -@pytest.mark.parametrize("name,fn", list(EP_EXPERT_COMPUTE.items())) -def test_ep_compute_signature_contract(name, fn): - """Every EP compute function must explicitly accept expert_scores, hidden_act, - and a **kwargs catch-all for forward-compat extras (gate_up_bias, down_bias, etc.). - """ - sig = inspect.signature(fn) - params = sig.parameters +def _assert_ep_registry_signature_contracts(): + """Every registered EP function exposes the shared forward-compatible API.""" + registries = { + "EP_EXPERT_COMPUTE": EP_EXPERT_COMPUTE, + "EP_EXPERT_COMPUTE_MOE_ACT": EP_EXPERT_COMPUTE_MOE_ACT, + } + for registry_name, registry in registries.items(): + for name, fn in registry.items(): + sig = inspect.signature(fn) + params = sig.parameters - for required in _REQUIRED_EP_PARAMS: - assert required in params, ( - f"EP_EXPERT_COMPUTE['{name}'] is missing explicit '{required}' param. Signature: {sig}" - ) + for required in _REQUIRED_EP_PARAMS: + assert required in params, ( + f"{registry_name}['{name}'] is missing explicit '{required}' param. Signature: {sig}" + ) - has_var_keyword = any(p.kind == p.VAR_KEYWORD for p in params.values()) - assert has_var_keyword, ( - f"EP_EXPERT_COMPUTE['{name}'] has no **kwargs — new extras like " - f"gate_up_bias will break callers. Signature: {sig}" - ) + has_var_keyword = any(p.kind == p.VAR_KEYWORD for p in params.values()) + assert has_var_keyword, ( + f"{registry_name}['{name}'] has no **kwargs — new extras like " + f"gate_up_bias will break callers. Signature: {sig}" + ) -def test_native_ep_adapter_consumes_fp8_kwargs_when_fp8_compute_is_disabled(): +def _assert_native_ep_adapter_consumes_fp8_kwargs_when_fp8_compute_is_disabled(): """Native EP is a BF16 fallback path but receives the common EP FP8 kwargs.""" - fn = EP_EXPERT_COMPUTE.get("native") - if fn is None: - pytest.skip("native EP backend is unavailable") + fn = EP_EXPERT_COMPUTE["native"] permute_tokens = torch.empty(0, 4) cumsum = torch.zeros(1, dtype=torch.int32) @@ -260,10 +91,8 @@ def test_native_ep_adapter_consumes_fp8_kwargs_when_fp8_compute_is_disabled(): assert out.shape == permute_tokens.shape -def test_native_ep_adapter_rejects_fp8_expert_compute_explicitly(): - fn = EP_EXPERT_COMPUTE.get("native") - if fn is None: - pytest.skip("native EP backend is unavailable") +def _assert_native_ep_adapter_rejects_fp8_expert_compute_explicitly(): + fn = EP_EXPERT_COMPUTE["native"] permute_tokens = torch.empty(0, 4) cumsum = torch.zeros(1, dtype=torch.int32) @@ -285,10 +114,8 @@ def test_native_ep_adapter_rejects_fp8_expert_compute_explicitly(): ) -def test_triton_ep_adapter_consumes_fp8_kwargs_when_fp8_compute_is_disabled(monkeypatch): - fn = EP_EXPERT_COMPUTE.get("triton") - if fn is None or not hasattr(moe_backend, "TritonEPGroupGemm"): - pytest.skip("triton EP backend is unavailable") +def _assert_triton_ep_adapter_consumes_fp8_kwargs_when_fp8_compute_is_disabled(monkeypatch): + fn = EP_EXPERT_COMPUTE["triton"] def fake_apply(*args): return args[0].new_empty(args[0].shape) @@ -318,10 +145,8 @@ def fake_apply(*args): assert out.shape == permute_tokens.shape -def test_triton_ep_adapter_rejects_fp8_expert_compute_explicitly(): - fn = EP_EXPERT_COMPUTE.get("triton") - if fn is None: - pytest.skip("triton EP backend is unavailable") +def _assert_triton_ep_adapter_rejects_fp8_expert_compute_explicitly(): + fn = EP_EXPERT_COMPUTE["triton"] permute_tokens = torch.empty(0, 4) cumsum = torch.zeros(1, dtype=torch.int32) @@ -345,10 +170,8 @@ def test_triton_ep_adapter_rejects_fp8_expert_compute_explicitly(): ) -def test_triton_moe_act_ep_adapter_consumes_common_kwargs(monkeypatch): - fn = EP_EXPERT_COMPUTE_MOE_ACT.get("triton") - if fn is None or not hasattr(moe_backend, "TritonEPGroupGemmMoeAct"): - pytest.skip("triton moe_act EP backend is unavailable") +def _assert_triton_moe_act_ep_adapter_consumes_common_kwargs(monkeypatch): + fn = EP_EXPERT_COMPUTE_MOE_ACT["triton"] def fake_apply(*args): return args[0].new_empty(args[0].shape) @@ -381,27 +204,8 @@ def fake_apply(*args): assert out.shape == permute_tokens.shape -@pytest.mark.parametrize("name,fn", list(EP_EXPERT_COMPUTE_MOE_ACT.items())) -def test_moe_act_ep_compute_signature_contract(name, fn): - sig = inspect.signature(fn) - params = sig.parameters - - for required in _REQUIRED_EP_PARAMS: - assert required in params, ( - f"EP_EXPERT_COMPUTE_MOE_ACT['{name}'] is missing explicit '{required}' param. Signature: {sig}" - ) - - has_var_keyword = any(p.kind == p.VAR_KEYWORD for p in params.values()) - assert has_var_keyword, ( - f"EP_EXPERT_COMPUTE_MOE_ACT['{name}'] has no **kwargs — new extras like " - f"gate_up_bias will break callers. Signature: {sig}" - ) - - -def test_triton_moe_act_ep_adapter_rejects_fp8_expert_compute_explicitly(): - fn = EP_EXPERT_COMPUTE_MOE_ACT.get("triton") - if fn is None: - pytest.skip("triton moe_act EP backend is unavailable") +def _assert_triton_moe_act_ep_adapter_rejects_fp8_expert_compute_explicitly(): + fn = EP_EXPERT_COMPUTE_MOE_ACT["triton"] permute_tokens = torch.empty(0, 4) cumsum = torch.zeros(1, dtype=torch.int32) @@ -426,11 +230,9 @@ def test_triton_moe_act_ep_adapter_rejects_fp8_expert_compute_explicitly(): ) -def test_quack_ep_adapter_forwards_activation_native(monkeypatch): - fn = EP_EXPERT_COMPUTE.get("quack") - quack_cls = getattr(moe_backend, "_QuackEPGroupGemm", None) - if fn is None or quack_cls is None: - pytest.skip("quack EP backend is unavailable") +def _assert_quack_ep_adapter_forwards_activation_native(monkeypatch): + fn = EP_EXPERT_COMPUTE["quack"] + quack_cls = moe_backend._QuackEPGroupGemm seen = {} @@ -463,3 +265,79 @@ def fake_apply(*args): assert out.shape == permute_tokens.shape assert seen["activation_native"] is True + + +def _reference_ep_forward(permute_tokens, cumsum, gate_proj, up_proj, down_proj, expert_scores): + outputs = [] + start = 0 + for expert_idx, count in enumerate(counts_from_cumsum(cumsum)): + end = start + count + x = permute_tokens[start:end] + hidden = F.silu(x @ gate_proj[expert_idx]) * (x @ up_proj[expert_idx]) + hidden = hidden * expert_scores[start:end].to(hidden.dtype).unsqueeze(-1) + outputs.append(hidden @ down_proj[expert_idx]) + start = end + return torch.cat(outputs, dim=0) + + +def _assert_ep_group_gemm_propagates_routing_score_gradients(monkeypatch, module, class_name): + patch_ep_kernels(monkeypatch, module) + fn = getattr(module, class_name) + + torch.manual_seed(0) + num_local_experts = 2 + hidden_dim = 8 + intermediate_size = 12 + counts = torch.tensor([3, 2]) + cumsum = torch.cumsum(counts, dim=0) + num_tokens = int(cumsum[-1].item()) + + permute_tokens = torch.randn(num_tokens, hidden_dim) + gate_proj = torch.randn(num_local_experts, hidden_dim, intermediate_size) + up_proj = torch.randn(num_local_experts, hidden_dim, intermediate_size) + down_proj = torch.randn(num_local_experts, intermediate_size, hidden_dim) + expert_scores = torch.rand(num_tokens, requires_grad=True) + upstream = torch.randn(num_tokens, hidden_dim) + + gate_up_proj = torch.cat([gate_proj, up_proj], dim=-1) + output = fn.apply(permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores) + output.backward(upstream) + grad_scores = expert_scores.grad.detach().clone() + + expert_scores_ref = expert_scores.detach().clone().requires_grad_(True) + ref_output = _reference_ep_forward( + permute_tokens, + cumsum, + gate_proj, + up_proj, + down_proj, + expert_scores_ref, + ) + ref_output.backward(upstream) + + torch.testing.assert_close(output, ref_output) + torch.testing.assert_close(grad_scores, expert_scores_ref.grad) + + +def test_ep_adapter_registry_backend_arguments_and_fp8_boundary_contract(monkeypatch): + _assert_quack_ep_registers_without_optional_moe_act() + _assert_ep_registry_signature_contracts() + assert {"native", "triton", "quack"} <= EP_EXPERT_COMPUTE.keys() + assert "triton" in EP_EXPERT_COMPUTE_MOE_ACT + assert hasattr(moe_backend, "TritonEPGroupGemm") + assert hasattr(moe_backend, "TritonEPGroupGemmMoeAct") + + for module, class_name in ( + (triton_moe, "TritonEPGroupGemm"), + (quack_moe, "QuackEPGroupGemm"), + ): + with monkeypatch.context() as kernel_patch: + _assert_ep_group_gemm_propagates_routing_score_gradients(kernel_patch, module, class_name) + + _assert_native_ep_adapter_consumes_fp8_kwargs_when_fp8_compute_is_disabled() + _assert_native_ep_adapter_rejects_fp8_expert_compute_explicitly() + _assert_triton_ep_adapter_consumes_fp8_kwargs_when_fp8_compute_is_disabled(monkeypatch) + _assert_triton_ep_adapter_rejects_fp8_expert_compute_explicitly() + _assert_triton_moe_act_ep_adapter_consumes_common_kwargs(monkeypatch) + _assert_triton_moe_act_ep_adapter_rejects_fp8_expert_compute_explicitly() + _assert_quack_ep_adapter_forwards_activation_native(monkeypatch) diff --git a/tests/ops/test_ep_routing_scores.py b/tests/ops/test_ep_routing_scores.py deleted file mode 100644 index 4ce7e7b9..00000000 --- a/tests/ops/test_ep_routing_scores.py +++ /dev/null @@ -1,181 +0,0 @@ -import importlib.util -import sys -import types -from pathlib import Path - -import pytest -import torch -import torch.nn.functional as F - - -pytestmark = pytest.mark.cpu - -_MODULE_PATHS = { - "xorl.ops.moe.triton": Path(__file__).resolve().parents[2] / "src/xorl/ops/moe/triton.py", - "xorl.ops.moe.quack": Path(__file__).resolve().parents[2] / "src/xorl/ops/moe/quack.py", -} - - -def _counts_from_cumsum(cumsum: torch.Tensor) -> list[int]: - counts = torch.empty_like(cumsum) - counts[0] = cumsum[0] - counts[1:] = cumsum[1:] - cumsum[:-1] - return counts.tolist() - - -def _naive_group_gemm_same_nk(a, b, cumsum_M, max_M, transpose_a=False, transpose_b=False, **kwargs): - del max_M, kwargs - assert not transpose_a - - outputs = [] - start = 0 - for expert_idx, count in enumerate(_counts_from_cumsum(cumsum_M)): - end = start + count - weight = b[expert_idx] - if transpose_b: - outputs.append(a[start:end] @ weight.transpose(0, 1)) - else: - outputs.append(a[start:end] @ weight) - start = end - - return torch.cat(outputs, dim=0) - - -def _naive_group_gemm_same_mn(a, b, c, cumsum_K, max_K, transpose_a=False, transpose_b=False, **kwargs): - del max_K, kwargs - - start = 0 - for expert_idx, count in enumerate(_counts_from_cumsum(cumsum_K)): - end = start + count - lhs = a[start:end].transpose(0, 1) if transpose_a else a[start:end] - rhs = b[start:end].transpose(0, 1) if transpose_b else b[start:end] - c[expert_idx].copy_(lhs @ rhs) - start = end - - return c - - -def _patch_ep_kernels(monkeypatch, module_name: str): - import xorl.utils.import_utils as import_utils - - moe_stub = types.ModuleType("xorl.ops.group_gemm.kernel.moe") - moe_stub.expert_histogram = None - moe_stub.moe_gather = None - moe_stub.moe_index_compute = None - moe_stub.moe_scatter = None - moe_stub.moe_add_gather = None - monkeypatch.setattr(import_utils, "is_fused_moe_available", lambda: True) - sys.modules.pop("xorl.ops.group_gemm.kernel.moe", None) - sys.modules.pop("xorl.ops.group_gemm.kernel.group_gemm", None) - sys.modules.pop("xorl.ops.group_gemm.kernel.quack", None) - monkeypatch.setitem( - sys.modules, - "xorl.ops.group_gemm.kernel.moe", - moe_stub, - ) - if module_name.endswith("triton"): - group_gemm_stub = types.ModuleType("xorl.ops.group_gemm.kernel.group_gemm") - group_gemm_stub.group_gemm_same_nk = _naive_group_gemm_same_nk - group_gemm_stub.group_gemm_same_mn = _naive_group_gemm_same_mn - monkeypatch.setitem( - sys.modules, - "xorl.ops.group_gemm.kernel.group_gemm", - group_gemm_stub, - ) - else: - quack_stub = types.ModuleType("xorl.ops.group_gemm.kernel.quack") - quack_stub.cumsum_to_cu_seqlens = lambda cumsum: cumsum - quack_stub.quack_group_gemm_same_nk = _naive_group_gemm_same_nk - quack_stub.quack_group_gemm_same_mn = _naive_group_gemm_same_mn - monkeypatch.setitem( - sys.modules, - "xorl.ops.group_gemm.kernel.quack", - quack_stub, - ) - spec = importlib.util.spec_from_file_location( - f"codex_test_{module_name.rsplit('.', maxsplit=1)[-1]}", _MODULE_PATHS[module_name] - ) - module = importlib.util.module_from_spec(spec) - assert spec.loader is not None - spec.loader.exec_module(module) - return module - - -def reference_ep_forward( - permute_tokens: torch.Tensor, - cumsum: torch.Tensor, - gate_proj: torch.Tensor, - up_proj: torch.Tensor, - down_proj: torch.Tensor, - expert_scores: torch.Tensor, -) -> torch.Tensor: - outputs = [] - start = 0 - for expert_idx, count in enumerate(_counts_from_cumsum(cumsum)): - end = start + count - x = permute_tokens[start:end] - h = F.silu(x @ gate_proj[expert_idx]) * (x @ up_proj[expert_idx]) - h = h * expert_scores[start:end].to(h.dtype).unsqueeze(-1) - outputs.append(h @ down_proj[expert_idx]) - start = end - - return torch.cat(outputs, dim=0) - - -@pytest.mark.parametrize( - ("module_name", "class_name"), - [ - pytest.param("xorl.ops.moe.triton", "TritonEPGroupGemm", id="triton"), - pytest.param("xorl.ops.moe.quack", "QuackEPGroupGemm", id="quack"), - ], -) -def test_ep_group_gemm_propagates_routing_score_gradients(monkeypatch, module_name, class_name): - try: - module = _patch_ep_kernels(monkeypatch, module_name) - except ImportError as exc: - pytest.skip(f"{module_name} unavailable: {exc}") - - fn = getattr(module, class_name) - - torch.manual_seed(0) - dtype = torch.float32 - num_local_experts = 2 - hidden_dim = 8 - intermediate_size = 12 - counts = torch.tensor([3, 2]) - cumsum = torch.cumsum(counts, dim=0) - num_tokens = int(cumsum[-1].item()) - - permute_tokens = torch.randn(num_tokens, hidden_dim, dtype=dtype) - gate_proj = torch.randn(num_local_experts, hidden_dim, intermediate_size, dtype=dtype) - up_proj = torch.randn(num_local_experts, hidden_dim, intermediate_size, dtype=dtype) - down_proj = torch.randn(num_local_experts, intermediate_size, hidden_dim, dtype=dtype) - expert_scores = torch.rand(num_tokens, dtype=dtype, requires_grad=True) - upstream = torch.randn(num_tokens, hidden_dim, dtype=dtype) - - # Both TritonEPGroupGemm and QuackEPGroupGemm take a fused gate_up_proj + intermediate_size (int). - gate_up_proj = torch.cat([gate_proj, up_proj], dim=-1) - output = fn.apply( - permute_tokens, - cumsum, - gate_up_proj, - down_proj, - intermediate_size, - expert_scores, - ) - output.backward(upstream) - grad_scores = expert_scores.grad.detach().clone() - - expert_scores_ref = expert_scores.detach().clone().requires_grad_(True) - ref_output = reference_ep_forward( - permute_tokens, - cumsum, - gate_proj, - up_proj, - down_proj, - expert_scores_ref, - ) - ref_output.backward(upstream) - - torch.testing.assert_close(output, ref_output) - torch.testing.assert_close(grad_scores, expert_scores_ref.grad) diff --git a/tests/ops/test_flashqla_contract_pin.py b/tests/ops/test_flashqla_contract_pin.py index bf8d1df0..bc769369 100644 --- a/tests/ops/test_flashqla_contract_pin.py +++ b/tests/ops/test_flashqla_contract_pin.py @@ -1,4 +1,4 @@ -"""P5 GDN contract-pin regression tests for the FlashQLA backend. +"""P5 GDN backend selection and contract-pin regression tests for FlashQLA. Ports the FlashQLA certification gates 2 (M/batch-invariance) and 4 (chunk-chaining exactness through the fp32 pool-layout handoff) as regression @@ -13,9 +13,13 @@ fp32 g/beta, ``use_qk_l2norm_in_kernel=True``. """ +import warnings + import pytest import torch +import xorl.ops.linear_attention.layers.gated_deltanet as gated_deltanet +from xorl.ops.linear_attention import GatedDeltaNet from xorl.ops.linear_attention.backend import FLASHQLA_AUTOCP_ENV, resolve_flashqla_auto_cp from xorl.ops.linear_attention.modules.bi_contract import gdn_contract @@ -32,16 +36,12 @@ def _flashqla_chunk_or_skip(): if torch.cuda.get_device_capability() != (9, 0): pytest.skip("FlashQLA requires a Hopper (SM90) GPU") - try: - import tilelang.language as _tl # noqa: PLC0415 - - if "prefer_instruction" not in inspect.signature(_tl.copy).parameters: - pytest.skip("tilelang lacks prefer_instruction (PR #2303); FlashQLA TMA path unavailable") - from xorl.ops.linear_attention.flashqla import chunk_gated_delta_rule as flashqla_chunk # noqa: PLC0415 - except pytest.skip.Exception: - raise - except Exception as exc: # tilelang missing / SM90 import-time check / build failure - pytest.skip(f"FlashQLA backend unavailable: {exc}") + import tilelang.language as _tl # noqa: PLC0415 + + if "prefer_instruction" not in inspect.signature(_tl.copy).parameters: + pytest.skip("tilelang lacks prefer_instruction (PR #2303); FlashQLA TMA path unavailable") + from xorl.ops.linear_attention.flashqla import chunk_gated_delta_rule as flashqla_chunk # noqa: PLC0415 + return flashqla_chunk @@ -111,6 +111,39 @@ def test_resolve_auto_cp_precedence(monkeypatch): assert resolve_flashqla_auto_cp(None) is False assert resolve_flashqla_auto_cp(False) is False + _assert_gdn_backend_env_dispatches_to_flashqla(monkeypatch) + + +def _assert_gdn_backend_env_dispatches_to_flashqla(monkeypatch): + calls = [] + + def fake_flashqla_chunk(**kwargs): + calls.append(kwargs) + assert "cp_context" not in kwargs + return kwargs["v"], None + + monkeypatch.setenv("XORL_GDN_BACKEND", "flashqla") + monkeypatch.setattr(gated_deltanet, "flashqla_chunk_gated_delta_rule", fake_flashqla_chunk) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + layer = GatedDeltaNet( + hidden_size=128, + expand_v=1.0, + head_dim=128, + num_heads=1, + num_v_heads=1, + mode="chunk", + use_gate=False, + use_short_conv=False, + ) + layer.train() + out, _, _ = layer(torch.randn(1, 8, 128)) + + assert len(calls) == 1 + assert calls[0]["q"].shape == (1, 8, 1, 128) + assert out.shape == (1, 8, 128) + @requires_cuda @pytest.mark.gpu @@ -150,7 +183,7 @@ def spy(*args, **kwargs): @requires_cuda @pytest.mark.gpu -def test_gate2_packed_varlen_matches_individual_calls(armed_contract_lane): +def test_gate2_shape_invariance_policy(armed_contract_lane): """Gate 2a: packed varlen rows with fp32 initial states == per-request calls, bitwise.""" fn = _flashqla_chunk_or_skip() partial_lens = [1, 17, 64, 33] @@ -169,27 +202,26 @@ def test_gate2_packed_varlen_matches_individual_calls(armed_contract_lane): _assert_bitwise(o_pack[:, lo:hi], o_i, f"row {i} (len {L}) out") _assert_bitwise(s_pack[i : i + 1], s_i, f"row {i} (len {L}) state") + _assert_gate2_same_row_bits_invariant_to_total_m() + _assert_gate2_block_dv_tile_heuristic_bit_invariant() -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("T", [2048, 4096]) -def test_gate2_same_row_bits_invariant_to_total_M(armed_contract_lane, T): + +def _assert_gate2_same_row_bits_invariant_to_total_m(): """Gate 2b: the same row alone (CP-eligible bs1) vs packed 4xT — bitwise under the pin.""" fn = _flashqla_chunk_or_skip() - q, k, v, g, beta = _make_inputs(4 * T, seed=22) + for T in (2048, 4096): + q, k, v, g, beta = _make_inputs(4 * T, seed=22) - o_alone, s_alone = _run(fn, q[:, :T], k[:, :T], v[:, :T], g[:, :T], beta[:, :T]) + o_alone, s_alone = _run(fn, q[:, :T], k[:, :T], v[:, :T], g[:, :T], beta[:, :T]) - cu = torch.arange(0, 4 * T + 1, T, device="cuda", dtype=torch.long) - o_pack, s_pack = _run(fn, q, k, v, g, beta, cu_seqlens=cu) + cu = torch.arange(0, 4 * T + 1, T, device="cuda", dtype=torch.long) + o_pack, s_pack = _run(fn, q, k, v, g, beta, cu_seqlens=cu) - _assert_bitwise(o_pack[:, :T], o_alone, f"T={T} out") - _assert_bitwise(s_pack[0:1], s_alone, f"T={T} state") + _assert_bitwise(o_pack[:, :T], o_alone, f"T={T} out") + _assert_bitwise(s_pack[0:1], s_alone, f"T={T} state") -@requires_cuda -@pytest.mark.gpu -def test_gate2_block_dv_tile_heuristic_bit_invariant(armed_contract_lane): +def _assert_gate2_block_dv_tile_heuristic_bit_invariant(): """Gate 2c: identical row 0 at B=1/2/3 (block_DV tile heuristic 32/64/128) — bitwise.""" fn = _flashqla_chunk_or_skip() q, k, v, g, beta = _make_inputs(1024, seed=24, batch=3) @@ -204,24 +236,24 @@ def test_gate2_block_dv_tile_heuristic_bit_invariant(armed_contract_lane): @requires_cuda @pytest.mark.gpu -@pytest.mark.parametrize("T,step", [(256, CHUNK), (4096, CHUNK), (4096, 256)]) -def test_gate4_chunk_chaining_bitwise_through_pool_layout(armed_contract_lane, T, step): +def test_gate4_chunk_chaining_bitwise_through_pool_layout(armed_contract_lane): """Gate 4: one call == chained calls with fp32 state handoff through the sglang pool layout ([N, HV, V, K] transpose round-trip) — the recompute-decode prerequisite.""" fn = _flashqla_chunk_or_skip() - q, k, v, g, beta = _make_inputs(T, seed=42) - - o_ref, s_ref = _run(fn, q, k, v, g, beta) - - pool = None - outs = [] - for t0 in range(0, T, step): - sl = slice(t0, min(t0 + step, T)) - init = pool.transpose(-1, -2).contiguous() if pool is not None else None - o, s = _run(fn, q[:, sl], k[:, sl], v[:, sl], g[:, sl], beta[:, sl], initial_state=init) - pool = s.transpose(-1, -2).contiguous() - outs.append(o) - o_chain = torch.cat(outs, dim=1) - - _assert_bitwise(o_ref, o_chain, f"T={T} step={step} out") - _assert_bitwise(s_ref.transpose(-1, -2).contiguous(), pool, f"T={T} step={step} state") + for T, step in ((256, CHUNK), (4096, CHUNK), (4096, 256)): + q, k, v, g, beta = _make_inputs(T, seed=42) + + o_ref, s_ref = _run(fn, q, k, v, g, beta) + + pool = None + outs = [] + for t0 in range(0, T, step): + sl = slice(t0, min(t0 + step, T)) + init = pool.transpose(-1, -2).contiguous() if pool is not None else None + o, s = _run(fn, q[:, sl], k[:, sl], v[:, sl], g[:, sl], beta[:, sl], initial_state=init) + pool = s.transpose(-1, -2).contiguous() + outs.append(o) + o_chain = torch.cat(outs, dim=1) + + _assert_bitwise(o_ref, o_chain, f"T={T} step={step} out") + _assert_bitwise(s_ref.transpose(-1, -2).contiguous(), pool, f"T={T} step={step} state") diff --git a/tests/ops/test_flashqla_gdn.py b/tests/ops/test_flashqla_gdn.py index d7394a13..e59eb780 100644 --- a/tests/ops/test_flashqla_gdn.py +++ b/tests/ops/test_flashqla_gdn.py @@ -18,18 +18,14 @@ def _flashqla_chunk_or_skip(): if torch.cuda.get_device_capability() != (9, 0): pytest.skip("FlashQLA requires a Hopper (SM90) GPU") - try: - import tilelang.language as tl # noqa: PLC0415 - - if "prefer_instruction" not in inspect.signature(tl.copy).parameters: - pytest.skip("tilelang lacks the required prefer_instruction support") - from xorl.ops.linear_attention.flashqla import ( # noqa: PLC0415 - chunk_gated_delta_rule, - ) - except pytest.skip.Exception: - raise - except Exception as exc: - pytest.skip(f"FlashQLA backend unavailable: {exc}") + import tilelang.language as tl # noqa: PLC0415 + + if "prefer_instruction" not in inspect.signature(tl.copy).parameters: + pytest.skip("tilelang lacks the required prefer_instruction support") + from xorl.ops.linear_attention.flashqla import ( # noqa: PLC0415 + chunk_gated_delta_rule, + ) + return chunk_gated_delta_rule @@ -51,8 +47,7 @@ def _inputs(num_heads: int, *, requires_grad: bool): return q, k, v, g, beta -@pytest.mark.parametrize("num_heads", [4, 32]) -def test_flashqla_matches_fla_forward(num_heads): +def _assert_flashqla_matches_fla_forward(num_heads): flashqla_chunk = _flashqla_chunk_or_skip() q, k, v, g, beta = _inputs(num_heads, requires_grad=False) kwargs = { @@ -73,24 +68,26 @@ def test_flashqla_matches_fla_forward(num_heads): assert _cosine(fla_state, flashqla_state) > 0.99 -@pytest.mark.parametrize("num_heads", [4, 32]) -def test_flashqla_matches_fla_backward(num_heads): - flashqla_chunk = _flashqla_chunk_or_skip() - gradients = {} - for name, implementation in (("fla", fla_chunk), ("flashqla", flashqla_chunk)): - q, k, v, g, beta = _inputs(num_heads, requires_grad=True) - output, _ = implementation( - q=q, - k=k, - v=v, - g=g, - beta=beta, - output_final_state=False, - use_qk_l2norm_in_kernel=True, - ) - output.float().square().mean().backward() - gradients[name] = (q.grad, k.grad, v.grad, g.grad, beta.grad) - - for reference, actual in zip(gradients["fla"], gradients["flashqla"]): - assert actual is not None and torch.isfinite(actual).all() - assert _cosine(reference, actual) > 0.97 +def test_flashqla_forward_and_backward_match_fla(): + for num_heads in (4, 32): + _assert_flashqla_matches_fla_forward(num_heads) + + flashqla_chunk = _flashqla_chunk_or_skip() + gradients = {} + for name, implementation in (("fla", fla_chunk), ("flashqla", flashqla_chunk)): + q, k, v, g, beta = _inputs(num_heads, requires_grad=True) + output, _ = implementation( + q=q, + k=k, + v=v, + g=g, + beta=beta, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + ) + output.float().square().mean().backward() + gradients[name] = (q.grad, k.grad, v.grad, g.grad, beta.grad) + + for reference, actual in zip(gradients["fla"], gradients["flashqla"]): + assert actual is not None and torch.isfinite(actual).all(), num_heads + assert _cosine(reference, actual) > 0.97, num_heads diff --git a/tests/ops/test_gated_delta_rule.py b/tests/ops/test_gated_delta_rule.py deleted file mode 100644 index 0386c5db..00000000 --- a/tests/ops/test_gated_delta_rule.py +++ /dev/null @@ -1,41 +0,0 @@ -import pytest -import torch -import torch.nn.functional as F -from fla.ops.gated_delta_rule import chunk_gated_delta_rule - - -pytestmark = [ - pytest.mark.gpu, - pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), -] - - -def test_chunk_gated_delta_rule_backward_real_qwen35_shape(): - """Regression test for the Hopper illegal-memory-access autotune candidate.""" - torch.manual_seed(0) - device = "cuda" - batch, seq_len, num_heads, head_dim = 1, 4096, 4, 128 - - q = torch.randn(batch, seq_len, num_heads, head_dim, device=device, dtype=torch.bfloat16, requires_grad=True) - k = torch.randn(batch, seq_len, num_heads, head_dim, device=device, dtype=torch.bfloat16, requires_grad=True) - v = torch.randn(batch, seq_len, num_heads, head_dim, device=device, dtype=torch.bfloat16, requires_grad=True) - beta = torch.rand(batch, seq_len, num_heads, device=device, dtype=torch.float32).sigmoid().requires_grad_() - g = F.logsigmoid(torch.randn(batch, seq_len, num_heads, device=device, dtype=torch.float32)).requires_grad_() - h0 = torch.zeros(batch, num_heads, head_dim, head_dim, device=device, dtype=torch.float32, requires_grad=True) - - o, ht = chunk_gated_delta_rule( - q=q, - k=k, - v=v, - g=g, - beta=beta, - initial_state=h0, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - ) - loss = o.float().square().mean() + ht.float().square().mean() - loss.backward() - - for grad in (q.grad, k.grad, v.grad, beta.grad, g.grad, h0.grad): - assert grad is not None - assert torch.isfinite(grad).all() diff --git a/tests/ops/test_gated_deltanet_backend.py b/tests/ops/test_gated_deltanet_backend.py deleted file mode 100644 index 141569fe..00000000 --- a/tests/ops/test_gated_deltanet_backend.py +++ /dev/null @@ -1,45 +0,0 @@ -import warnings - -import pytest -import torch - -import xorl.ops.linear_attention.layers.gated_deltanet as gated_deltanet -from xorl.ops.linear_attention import GatedDeltaNet - - -pytestmark = [pytest.mark.cpu] - - -def _tiny_gdn() -> GatedDeltaNet: - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - layer = GatedDeltaNet( - hidden_size=128, - expand_v=1.0, - head_dim=128, - num_heads=1, - num_v_heads=1, - mode="chunk", - use_gate=False, - use_short_conv=False, - ) - layer.train() - return layer - - -def test_gdn_backend_env_dispatches_to_flashqla(monkeypatch): - calls = [] - - def fake_flashqla_chunk(**kwargs): - calls.append(kwargs) - assert "cp_context" not in kwargs - return kwargs["v"], None - - monkeypatch.setenv("XORL_GDN_BACKEND", "flashqla") - monkeypatch.setattr(gated_deltanet, "flashqla_chunk_gated_delta_rule", fake_flashqla_chunk) - - out, _, _ = _tiny_gdn()(torch.randn(1, 8, 128)) - - assert len(calls) == 1 - assert calls[0]["q"].shape == (1, 8, 1, 128) - assert out.shape == (1, 8, 128) diff --git a/tests/ops/test_gdn_conv_contract.py b/tests/ops/test_gdn_conv_contract.py index 332b10e0..796096b6 100644 --- a/tests/ops/test_gdn_conv_contract.py +++ b/tests/ops/test_gdn_conv_contract.py @@ -80,7 +80,7 @@ def _tiny_gdn(**overrides) -> GatedDeltaNet: @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") class TestConvContractGPU: - def test_forward_bitwise_matches_serving_invocation_q35_shapes(self): + def test_forward_backward_parity_and_determinism_policy(self): device, dtype = torch.device("cuda"), torch.bfloat16 convs = _make_convs(KEY_DIM, VALUE_DIM, device, dtype) torch.manual_seed(0) @@ -95,7 +95,12 @@ def test_forward_bitwise_matches_serving_invocation_q35_shapes(self): assert torch.equal(k, ref_k) assert torch.equal(v, ref_v) - def test_forward_bitwise_varlen_and_batch(self): + self._assert_forward_bitwise_varlen_and_batch() + self._assert_forward_determinism_double_run() + self._assert_backward_parity_and_determinism() + TestConvContractGPU()._assert_end_to_end_gdn_block_policy() + + def _assert_forward_bitwise_varlen_and_batch(self): device, dtype = torch.device("cuda"), torch.bfloat16 convs = _make_convs(KEY_DIM, VALUE_DIM, device, dtype) torch.manual_seed(1) @@ -115,7 +120,7 @@ def test_forward_bitwise_varlen_and_batch(self): ref, _ = _serving_reference(q_in, k_in, v_in, convs) assert torch.equal(torch.cat((q, k, v), dim=-1), ref) - def test_forward_determinism_double_run(self): + def _assert_forward_determinism_double_run(self): device, dtype = torch.device("cuda"), torch.bfloat16 convs = _make_convs(256, 512, device, dtype) torch.manual_seed(2) @@ -125,7 +130,7 @@ def test_forward_determinism_double_run(self): for a, b in zip(first, second, strict=True): assert torch.equal(a, b) - def test_backward_matches_torch_depthwise_autograd(self): + def _assert_backward_parity_and_determinism(self): """Same grad_output through both lanes: differences isolate the backward composition.""" device, dtype = torch.device("cuda"), torch.bfloat16 convs = _make_convs(256, 512, device, dtype) @@ -153,7 +158,9 @@ def contract(inputs): for got, ref in zip(contract_grads, eager_grads, strict=True): torch.testing.assert_close(got, ref) - def test_backward_determinism_double_run(self): + self._assert_backward_determinism_double_run() + + def _assert_backward_determinism_double_run(self): device, dtype = torch.device("cuda"), torch.bfloat16 convs = _make_convs(256, 512, device, dtype) @@ -171,7 +178,7 @@ def run(): for a, b in zip(run(), run(), strict=True): assert torch.equal(a, b) - def test_end_to_end_gdn_block_grad_vs_eager(self): + def _assert_end_to_end_gdn_block_policy(self): device, dtype = torch.device("cuda"), torch.bfloat16 torch.manual_seed(6) layer = _tiny_gdn().to(device=device, dtype=dtype) @@ -194,7 +201,9 @@ def run(exact: bool): for name, ref in eager_grads.items(): torch.testing.assert_close(contract_grads[name], ref, rtol=5e-2, atol=1e-2, msg=lambda m: f"{name}: {m}") - def test_end_to_end_gdn_block_determinism(self): + self._assert_end_to_end_gdn_block_determinism() + + def _assert_end_to_end_gdn_block_determinism(self): device, dtype = torch.device("cuda"), torch.bfloat16 torch.manual_seed(8) layer = _tiny_gdn(exact_contract=True).to(device=device, dtype=dtype) @@ -235,14 +244,19 @@ def test_matches_sglang_tree_kernel_if_available(self): @pytest.mark.cpu class TestConvContractGuards: - def test_weight_pack_is_bitwise_neutral(self): + def test_weight_pack_routing_admission_and_state_lifecycle_policy(self, monkeypatch): convs = _make_convs(64, 128, torch.device("cpu"), torch.float32) packed = _pack_conv_weight(*(conv.weight for conv in convs)) assert torch.equal(packed[:64], convs[0].weight.squeeze(1)) assert torch.equal(packed[64:128], convs[1].weight.squeeze(1)) assert torch.equal(packed[128:], convs[2].weight.squeeze(1)) - def test_forward_routes_through_contract_when_armed(self, monkeypatch): + with monkeypatch.context() as case_patch: + self._assert_forward_routes_through_contract_when_armed(case_patch) + self._assert_contract_admission_policy() + self._assert_contract_state_lifecycle_policy() + + def _assert_forward_routes_through_contract_when_armed(self, monkeypatch): calls = [] def fake_contract(q_in, k_in, v_in, *convs, cu_seqlens=None): @@ -263,30 +277,36 @@ def fake_gating(A_log, a, b, dt_bias): assert len(calls) == 1 assert out.shape == (1, 8, 256) - def test_use_cache_raises(self): - layer = _tiny_gdn(exact_contract=True) - layer.eval() - with pytest.raises(RuntimeError, match="prefill only"): + def _assert_contract_admission_policy(self): + def use_cache(): + layer = _tiny_gdn(exact_contract=True) + layer.eval() layer(torch.randn(1, 128, 256), use_cache=True) - def test_cp_context_raises(self): - layer = _tiny_gdn(exact_contract=True) - cp_context = SimpleNamespace(cu_seqlens=torch.tensor([0, 8]), group=object(), is_first_rank=True) - with pytest.raises(RuntimeError, match="does not support CP"): - layer(torch.randn(1, 8, 256), cp_context=cp_context) - - def test_no_short_conv_raises(self): - layer = _tiny_gdn(use_short_conv=False, exact_contract=True) - with pytest.raises(RuntimeError, match="requires short convolution"): - layer(torch.randn(1, 8, 256)) - - def test_conv_bias_raises(self): - convs = _make_convs(32, 64, torch.device("cpu"), torch.float32, bias=True) - inputs = [torch.randn(1, 8, dim) for dim in (32, 32, 64)] - with pytest.raises(NotImplementedError, match="bias"): + def cp_context(): + layer = _tiny_gdn(exact_contract=True) + context = SimpleNamespace(cu_seqlens=torch.tensor([0, 8]), group=object(), is_first_rank=True) + layer(torch.randn(1, 8, 256), cp_context=context) + + def no_short_conv(): + _tiny_gdn(use_short_conv=False, exact_contract=True)(torch.randn(1, 8, 256)) + + def conv_bias(): + convs = _make_convs(32, 64, torch.device("cpu"), torch.float32, bias=True) + inputs = [torch.randn(1, 8, dim) for dim in (32, 32, 64)] causal_conv1d_qkv_contract(*inputs, *convs) - def test_exact_and_ordinary_modules_do_not_leak_contract_state(self): + cases = [ + ("decode cache", use_cache, RuntimeError, "prefill only"), + ("context parallelism", cp_context, RuntimeError, "does not support CP"), + ("missing short convolution", no_short_conv, RuntimeError, "requires short convolution"), + ("convolution bias", conv_bias, NotImplementedError, "bias"), + ] + for _label, invoke, error_type, error_pattern in cases: + with pytest.raises(error_type, match=error_pattern): + invoke() + + def _assert_contract_state_lifecycle_policy(self): seen = [] exact = _tiny_gdn(exact_contract=True) ordinary = _tiny_gdn(exact_contract=False) @@ -306,7 +326,9 @@ def record(self, hidden_states, **_kwargs): assert seen == [True, False, True] assert not _is_gdn_contract_enabled() - def test_checkpoint_recompute_reestablishes_module_contract(self): + self._assert_checkpoint_recompute_reestablishes_module_contract() + + def _assert_checkpoint_recompute_reestablishes_module_contract(self): seen = [] exact = _tiny_gdn(exact_contract=True) diff --git a/tests/ops/test_gdn_decode_prep.py b/tests/ops/test_gdn_decode_prep.py deleted file mode 100644 index e0b8abe2..00000000 --- a/tests/ops/test_gdn_decode_prep.py +++ /dev/null @@ -1,247 +0,0 @@ -"""Bitwise gates for the decode-shaped GDN prep kernels (P5). - -The decode-scheduled triangular solve must be bit-identical to the pinned -(num_warps=2) ``solve_tril`` — the frozen prefill-contract kernel — and the -opt-in recompute-decode composition must reproduce the pinned FlashQLA prefill -bitwise through per-step partial-chunk recompute. -""" - -import pytest -import torch -import torch.nn.functional as F - - -requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -requires_hopper = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9, - reason="FlashQLA is SM90-only", -) - -HK, HV, DK, DV = 16, 32, 128, 128 -CHUNK = 64 -SCALE = DK**-0.5 - - -def _decode_shape_A(batch: int, seed: int, regime: str = "normal"): - """A = chunk_scaled_dot_kkt output on l2normed k at the padded decode shape.""" - from xorl.ops.linear_attention.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd - - gen = torch.Generator(device="cuda").manual_seed(seed) - total = batch * CHUNK - k = torch.randn(1, total, HV, DK, generator=gen, device="cuda", dtype=torch.bfloat16) - beta = torch.rand(1, total, HV, generator=gen, device="cuda", dtype=torch.float32) - if regime == "k_large": - k = (k.float() * 100).to(torch.bfloat16) - elif regime == "padded": - k = k.view(1, batch, CHUNK, HV, DK).clone() - k[:, :, 37:] = 0 - k = k.view(1, total, HV, DK) - beta = beta.view(1, batch, CHUNK, HV).clone() - beta[:, :, 37:] = 0 - beta = beta.view(1, total, HV) - k = F.normalize(k.float(), dim=-1).to(torch.bfloat16) - cu = torch.arange(0, total + 1, CHUNK, device="cuda", dtype=torch.long) - A = chunk_scaled_dot_kkt_fwd(k=k, g=None, beta=beta, cu_seqlens=cu, output_dtype=torch.float32) - return A, cu - - -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("batch", [1, 16, 64]) -@pytest.mark.parametrize("regime", ["normal", "k_large", "padded"]) -def test_solve_tril_decode_bitwise_vs_pinned(batch, regime): - from xorl.ops.linear_attention.ops.utils import solve_tril - from xorl.ops.linear_attention.ops.utils.solve_tril_decode import solve_tril_decode - - A, cu = _decode_shape_A(batch, seed=batch + 17, regime=regime) - ref = solve_tril(A=A, cu_seqlens=cu, output_dtype=torch.bfloat16) - got = solve_tril_decode(A=A, cu_seqlens=cu, output_dtype=torch.bfloat16) - assert torch.equal(ref, got) - - -@requires_cuda -@pytest.mark.gpu -def test_solve_tril_decode_bitwise_non_varlen_and_partial_chunks(): - from xorl.ops.linear_attention.ops.utils import solve_tril - from xorl.ops.linear_attention.ops.utils.solve_tril_decode import solve_tril_decode - - A, _ = _decode_shape_A(64, seed=3) - Ab = A.view(64, CHUNK, HV, CHUNK).contiguous() - assert torch.equal( - solve_tril(A=Ab, cu_seqlens=None, output_dtype=torch.bfloat16), - solve_tril_decode(A=Ab, cu_seqlens=None, output_dtype=torch.bfloat16), - ) - cu = torch.tensor([0, 64, 257, 450, 707, 1000], device="cuda", dtype=torch.long) - Ap = A[:, :1000].contiguous() - assert torch.equal( - solve_tril(A=Ap, cu_seqlens=cu, output_dtype=torch.bfloat16), - solve_tril_decode(A=Ap, cu_seqlens=cu, output_dtype=torch.bfloat16), - ) - - -@requires_cuda -@pytest.mark.gpu -@pytest.mark.parametrize("diag_group", [1, 4, 16]) -@pytest.mark.parametrize("diag_warps", [2, 8]) -@pytest.mark.parametrize("merge_warps", [2, 8]) -def test_solve_tril_decode_launch_config_invariance(diag_group, diag_warps, merge_warps): - # The reduction tree is spelled out structurally (fma + explicit adds), so - # bits must not move with the launch config; this pins that property. - from xorl.ops.linear_attention.ops.utils import solve_tril - from xorl.ops.linear_attention.ops.utils import solve_tril_decode as mod - from xorl.ops.linear_attention.ops.utils.index import prepare_chunk_indices - - A, cu = _decode_shape_A(64, seed=29) - ref = solve_tril(A=A, cu_seqlens=cu, output_dtype=torch.bfloat16) - ci = prepare_chunk_indices(cu, CHUNK) - B, T, H, BT = A.shape - Ai = torch.empty_like(A, dtype=torch.bfloat16) - Di = torch.empty(B, T, H, 16, dtype=torch.float32, device="cuda") - mod.solve_tril_64x64_diag_inv_grouped_kernel[len(ci) * 4, B * (H // diag_group)]( - A=A, - Di=Di, - cu_seqlens=cu, - chunk_indices=ci, - T=T, - H=H, - BT=BT, - G=diag_group, - IS_VARLEN=True, - num_warps=diag_warps, - num_stages=1, - ) - mod.solve_tril_64x64_merge_inv_kernel[len(ci), B * H]( - A=A, - Di=Di, - Ai=Ai, - cu_seqlens=cu, - chunk_indices=ci, - T=T, - H=H, - BT=BT, - DOT_PRECISION="ieee", - IS_VARLEN=True, - num_warps=merge_warps, - num_stages=1, - ) - assert torch.equal(ref, Ai) - - -def _make_inputs(T: int, seed: int): - gen = torch.Generator(device="cuda").manual_seed(seed) - - def rnd(*shape): - return torch.randn(*shape, generator=gen, device="cuda", dtype=torch.bfloat16) - - q = rnd(1, T, HK, DK).repeat_interleave(HV // HK, dim=2).contiguous() - k = rnd(1, T, HK, DK).repeat_interleave(HV // HK, dim=2).contiguous() - v = rnd(1, T, HV, DV) - a_in = rnd(1, T, HV) - b_in = rnd(1, T, HV) - A_log = torch.empty(HV, device="cuda", dtype=torch.float32).uniform_(0, 2, generator=gen).log() - dt_bias = torch.rand(HV, device="cuda", dtype=torch.float32, generator=gen) - g = -A_log.exp().view(1, 1, -1) * F.softplus(a_in.float() + dt_bias.view(1, 1, -1)) - beta = b_in.float().sigmoid().to(torch.bfloat16).float() - return q, k, v, g, beta - - -def _pinned_prefill(q, k, v, g, beta, initial_state=None, cu_seqlens=None): - from xorl.ops.linear_attention import tilelang_gemm_v1 - - tilelang_gemm_v1.patch() - from xorl.ops.linear_attention.flashqla.ops.gated_delta_rule.chunk import chunk_gated_delta_rule_fwd - from xorl.ops.linear_attention.flashqla.utils import l2norm - - q = l2norm(q) - k = l2norm(k) - _, _, o, _, final_state = chunk_gated_delta_rule_fwd( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=SCALE, - initial_state=initial_state, - cu_seqlens=cu_seqlens, - output_final_state=True, - output_h=False, - auto_cp=False, - ) - return o.to(q.dtype), final_state - - -@requires_hopper -@pytest.mark.gpu -def test_recompute_decode_bitwise_vs_pinned_prefill(): - # gate-5 protocol at T=193: padded fixed-shape per-step recompute from fp32 - # chunk checkpoints through the decode prep path reproduces the pinned - # prefill bitwise. - from xorl.ops.linear_attention.gdn_decode_prep import chunk_gated_delta_rule_fwd_decode - - T = 193 - q, k, v, g, beta = _make_inputs(T, seed=51) - o_ref, _ = _pinned_prefill(q, k, v, g, beta) - - checkpoint = None - outs = [] - for t in range(T): - t0 = (t // CHUNK) * CHUNK - L = t + 1 - t0 - pad = CHUNK - L - qp = F.pad(q[:, t0 : t + 1], (0, 0, 0, 0, 0, pad)) - kp = F.pad(k[:, t0 : t + 1], (0, 0, 0, 0, 0, pad)) - vp = F.pad(v[:, t0 : t + 1], (0, 0, 0, 0, 0, pad)) - gp = F.pad(g[:, t0 : t + 1], (0, 0, 0, pad)) - bp = F.pad(beta[:, t0 : t + 1], (0, 0, 0, pad)) - init = checkpoint.transpose(-1, -2).contiguous() if checkpoint is not None else None - o, s = chunk_gated_delta_rule_fwd_decode(qp, kp, vp, gp, bp, scale=SCALE, initial_state=init) - outs.append(o[:, L - 1 : L]) - if (t + 1) % CHUNK == 0: - checkpoint = s.transpose(-1, -2).contiguous() - o_dec = torch.cat(outs, dim=1) - assert torch.equal(o_ref, o_dec) - - -@requires_hopper -@pytest.mark.gpu -def test_padded_decode_call_bitwise_and_graph_capturable(): - from xorl.ops.linear_attention.gdn_decode_prep import chunk_gated_delta_rule_fwd_decode - from xorl.ops.linear_attention.ops.utils.index import prepare_chunk_indices - - batch = 16 - gen = torch.Generator(device="cuda").manual_seed(7) - total = batch * CHUNK - q = torch.randn(1, total, HV, DK, generator=gen, device="cuda", dtype=torch.bfloat16) - k = torch.randn(1, total, HV, DK, generator=gen, device="cuda", dtype=torch.bfloat16) - v = torch.randn(1, total, HV, DV, generator=gen, device="cuda", dtype=torch.bfloat16) - g = -torch.rand(1, total, HV, generator=gen, device="cuda", dtype=torch.float32) - beta = torch.rand(1, total, HV, generator=gen, device="cuda", dtype=torch.float32) - init = torch.randn(batch, HV, DK, DV, generator=gen, device="cuda", dtype=torch.float32) - cu = torch.arange(0, total + 1, CHUNK, device="cuda", dtype=torch.long) - ci = prepare_chunk_indices(cu, CHUNK) - - with torch.no_grad(): - o_pin, s_pin = _pinned_prefill(q, k, v, g, beta, initial_state=init, cu_seqlens=cu) - - def fn(): - return chunk_gated_delta_rule_fwd_decode( - q, k, v, g, beta, scale=SCALE, initial_state=init, cu_seqlens=cu, chunk_indices=ci - ) - - o_dec, s_dec = fn() - assert torch.equal(o_pin, o_dec) - assert torch.equal(s_pin, s_dec) - - torch.cuda.synchronize() - side = torch.cuda.Stream() - with torch.cuda.stream(side): - for _ in range(3): - fn() - torch.cuda.current_stream().wait_stream(side) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - o_g, s_g = fn() - graph.replay() - torch.cuda.synchronize() - assert torch.equal(o_pin, o_g) - assert torch.equal(s_pin, s_g) diff --git a/tests/ops/test_group_gemm.py b/tests/ops/test_group_gemm.py index f8b3fcde..eb224044 100644 --- a/tests/ops/test_group_gemm.py +++ b/tests/ops/test_group_gemm.py @@ -8,8 +8,12 @@ import torch -# Mark all tests as GPU since group_gemm requires CUDA -pytestmark = pytest.mark.gpu +# Grouped GEMM requires CUDA. Skip unsupported hosts before entering the test +# body, but let failures importing XORL's own kernel module surface normally. +pytestmark = [ + pytest.mark.gpu, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] def naive_group_gemm_same_nk( @@ -21,10 +25,7 @@ def naive_group_gemm_same_nk( ) -> torch.Tensor: """Naive PyTorch implementation of grouped GEMM with same N, K.""" G = b.shape[0] - if transpose_b: - N, K = b.shape[1], b.shape[2] - else: - K, N = b.shape[1], b.shape[2] + N = b.shape[1] if transpose_b else b.shape[2] total_M = a.shape[1] if transpose_a else a.shape[0] output = torch.zeros(total_M, N, dtype=a.dtype, device=a.device) @@ -85,16 +86,11 @@ def naive_group_gemm_same_mn( class TestGroupGemmSameNK: - """Test suite for group_gemm_same_nk: basic, unequal groups, transpose_b, large dims, single group.""" + """Numerical and input-policy contracts for ``group_gemm_same_nk``.""" - def test_same_nk_comprehensive(self): - """Basic forward, unequal groups, transpose_b, large dims, single group.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - try: - from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_nk # noqa: PLC0415 - except ImportError: - pytest.skip("group_gemm not available") + def test_same_nk_and_same_mn_comprehensive(self): + """Aligned/unaligned numerics, transpose-B, and input rejection.""" + from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_nk # noqa: PLC0415 # --- Basic forward --- G, K, N = 4, 128, 256 @@ -112,7 +108,7 @@ def test_same_nk_comprehensive(self): assert torch.allclose(output_kernel.float(), output_naive.float(), rtol=1e-2, atol=1e-2) # --- Unequal group sizes --- - G2, K2, N2 = 8, 64, 128 + G2, K2, N2 = 8, 70, 130 gs2 = [5, 100, 2, 50, 30, 8, 45, 20] total_M2 = sum(gs2) cumsum_M2 = torch.tensor([sum(gs2[: i + 1]) for i in range(G2)], dtype=torch.int32).cuda() @@ -129,38 +125,23 @@ def test_same_nk_comprehensive(self): naive3 = naive_group_gemm_same_nk(a3, b3, cumsum_M, transpose_a=False, transpose_b=True) assert torch.allclose(out3.float(), naive3.float(), rtol=1e-2, atol=1e-2) - # --- Large dimensions --- - G4, K4, N4 = 8, 4096, 14336 - gs4 = [512, 480, 520, 490, 510, 505, 495, 488] - total_M4 = sum(gs4) - cumsum_M4 = torch.tensor([sum(gs4[: i + 1]) for i in range(G4)], dtype=torch.int32).cuda() - a4 = torch.randn(total_M4, K4, dtype=torch.bfloat16).cuda() - b4 = torch.randn(G4, K4, N4, dtype=torch.bfloat16).cuda() - out4 = group_gemm_same_nk(a4, b4, cumsum_M4, max(gs4)) - naive4 = naive_group_gemm_same_nk(a4, b4, cumsum_M4) - assert torch.allclose(out4.float(), naive4.float(), rtol=5e-2, atol=5e-2) - - # --- Single group --- - M_sg, K_sg, N_sg = 64, 128, 256 - cumsum_sg = torch.tensor([M_sg], dtype=torch.int32).cuda() - a_sg = torch.randn(M_sg, K_sg, dtype=torch.float16).cuda() - b_sg = torch.randn(1, K_sg, N_sg, dtype=torch.float16).cuda() - out_sg = group_gemm_same_nk(a_sg, b_sg, cumsum_sg, M_sg) - expected_sg = torch.matmul(a_sg, b_sg[0]) - assert torch.allclose(out_sg.float(), expected_sg.float(), rtol=1e-2, atol=1e-2) + # Input-policy guards use the same admitted shape. + a_nc = torch.randn(total_M2, K2 * 2, dtype=torch.float16).cuda()[:, ::2] + with pytest.raises(AssertionError, match="Not implemented: Noncontiguous input"): + group_gemm_same_nk(a_nc, b2, cumsum_M2, max(gs2)) + + with pytest.raises(AssertionError, match="a.device.*b.device"): + group_gemm_same_nk(a2, b2.cpu(), cumsum_M2, max(gs2)) + + TestGroupGemmSameMN()._assert_same_mn_comprehensive() class TestGroupGemmSameMN: """Test suite for group_gemm_same_mn: basic, unequal groups, zero-K, single group.""" - def test_same_mn_comprehensive(self): + def _assert_same_mn_comprehensive(self): """Basic forward, unequal K dims, zero-K group, single group.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - try: - from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_mn # noqa: PLC0415 - except ImportError: - pytest.skip("group_gemm not available") + from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_mn # noqa: PLC0415 # --- Basic forward --- G, M, N = 4, 128, 256 @@ -200,52 +181,3 @@ def test_same_mn_comprehensive(self): naive3 = naive_group_gemm_same_mn(a3, b3, cumsum_K3, M3, N3) assert torch.all(c3[1] == 0) assert torch.allclose(c3.float(), naive3.float(), rtol=1e-2, atol=1e-2) - - # --- Single group --- - K_sg, M_sg, N_sg = 256, 64, 128 - cumsum_sg = torch.tensor([K_sg], dtype=torch.int32).cuda() - a_sg = torch.randn(K_sg, M_sg, dtype=torch.float16).cuda() - b_sg = torch.randn(K_sg, N_sg, dtype=torch.float16).cuda() - c_sg = torch.empty(1, M_sg, N_sg, dtype=torch.float16).cuda() - group_gemm_same_mn(a_sg, b_sg, c_sg, cumsum_sg, K_sg, transpose_a=True) - expected_sg = torch.matmul(a_sg.t(), b_sg).unsqueeze(0) - assert torch.allclose(c_sg.float(), expected_sg.float(), rtol=1e-2, atol=1e-2) - - -class TestGroupGemmProperties: - """Test mathematical properties and edge cases: dtype support, contiguity, device consistency.""" - - def test_properties(self): - """Dtype support (float16/bfloat16), contiguity requirement, device consistency.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - try: - from xorl.ops.group_gemm.kernel.group_gemm import group_gemm_same_nk # noqa: PLC0415 - except ImportError: - pytest.skip("group_gemm not available") - - G, K, N = 2, 64, 128 - group_sizes = [32, 32] - total_M = sum(group_sizes) - cumsum_M = torch.tensor([sum(group_sizes[: i + 1]) for i in range(G)], dtype=torch.int32).cuda() - max_M = max(group_sizes) - - # Dtype support - for dtype in [torch.float16, torch.bfloat16]: - a = torch.randn(total_M, K, dtype=dtype).cuda() - b = torch.randn(G, K, N, dtype=dtype).cuda() - output = group_gemm_same_nk(a, b, cumsum_M, max_M) - assert output.dtype == dtype - assert output.shape == (total_M, N) - - # Contiguity requirement - a_nc = torch.randn(total_M, K * 2, dtype=torch.bfloat16).cuda()[:, ::2] - b_ok = torch.randn(G, K, N, dtype=torch.bfloat16).cuda() - with pytest.raises(AssertionError, match="Not implemented: Noncontiguous input"): - group_gemm_same_nk(a_nc, b_ok, cumsum_M, max_M) - - # Device consistency - a_gpu = torch.randn(total_M, K, dtype=torch.bfloat16).cuda() - b_cpu = torch.randn(G, K, N, dtype=torch.bfloat16) - with pytest.raises(AssertionError, match="a.device.*b.device"): - group_gemm_same_nk(a_gpu, b_cpu, cumsum_M, max_M) diff --git a/tests/ops/test_kkt_contract.py b/tests/ops/test_kkt_contract.py deleted file mode 100644 index bccb7709..00000000 --- a/tests/ops/test_kkt_contract.py +++ /dev/null @@ -1,47 +0,0 @@ -import importlib - -import torch - -from xorl.ops.linear_attention.modules.bi_contract import gdn_contract - - -class _FakeKernel: - def __init__(self): - self.calls = [] - - def __getitem__(self, grid): - def launch(**kwargs): - self.calls.append((grid, kwargs)) - - return launch - - -def test_bi_contract_pins_serving_kkt_reduction_geometry(monkeypatch): - module = importlib.import_module("xorl.ops.linear_attention.ops.common.chunk_scaled_dot_kkt") - contract_kernel = _FakeKernel() - autotuned_kernel = _FakeKernel() - monkeypatch.setattr(module, "_chunk_scaled_dot_kkt_fwd_kernel", contract_kernel) - monkeypatch.setattr(module, "chunk_scaled_dot_kkt_fwd_kernel", autotuned_kernel) - - k = torch.empty(1, 64, 32, 128) - g = torch.empty(1, 64, 32) - beta = torch.empty(1, 64, 32) - - with gdn_contract(True): - module.chunk_scaled_dot_kkt_fwd(k=k, g=g, beta=beta) - assert not autotuned_kernel.calls - _, kwargs = contract_kernel.calls.pop() - assert kwargs["BK"] == 64 - assert kwargs["num_warps"] == 8 - assert kwargs["num_stages"] == 3 - assert kwargs["IS_VARLEN"] is False - assert kwargs["USE_G"] is True - assert kwargs["SAFE_EXP"] is True - - with gdn_contract(False): - module.chunk_scaled_dot_kkt_fwd(k=k, g=g, beta=beta) - assert not contract_kernel.calls - _, kwargs = autotuned_kernel.calls.pop() - assert "BK" not in kwargs - assert "num_warps" not in kwargs - assert "num_stages" not in kwargs diff --git a/tests/ops/test_lora_utils.py b/tests/ops/test_lora_utils.py deleted file mode 100644 index 0c34ca72..00000000 --- a/tests/ops/test_lora_utils.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Tests for stacked LoRA helper utilities.""" - -import pytest -import torch - -from xorl.ops.group_gemm.kernel.lora_utils import ( - get_lora_delta_weight_stacked, - init_lora_weights_stacked, - merge_lora_weights_stacked, - unmerge_lora_weights_stacked, -) - - -pytestmark = [pytest.mark.cpu] - - -def test_stacked_lora_helpers_use_gkn_layout(): - lora_A, lora_B = init_lora_weights_stacked( - num_experts=2, - r=3, - in_features=4, - out_features=5, - ) - - assert lora_A.shape == (2, 4, 3) - assert lora_B.shape == (2, 3, 5) - assert torch.equal(lora_B, torch.zeros_like(lora_B)) - - lora_A = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3) - lora_B = torch.arange(2 * 3 * 5, dtype=torch.float32).reshape(2, 3, 5) - base = torch.ones(2, 4, 5, dtype=torch.float32) - scaling = 0.25 - expected_delta = torch.bmm(lora_A, lora_B) * scaling - - delta = get_lora_delta_weight_stacked(lora_A, lora_B, scaling) - merged = merge_lora_weights_stacked(base, lora_A, lora_B, scaling) - unmerged = unmerge_lora_weights_stacked(merged, lora_A, lora_B, scaling) - - assert torch.equal(delta, expected_delta) - assert torch.equal(merged, base + expected_delta) - assert torch.equal(unmerged, base) diff --git a/tests/ops/test_moe_gkn_format.py b/tests/ops/test_moe_gkn_format.py deleted file mode 100644 index dc972111..00000000 --- a/tests/ops/test_moe_gkn_format.py +++ /dev/null @@ -1,337 +0,0 @@ -"""Tests for (G, K, N) weight format correctness in MoE experts. - -These tests verify that the MoE expert computation with weights stored in -(G, K, N) = [num_experts, in_features, out_features] format produces -results identical to naive per-expert nn.Linear computation. - -Weight format: - - HuggingFace nn.Linear: [out_features, in_features] - - Our (G,K,N) stacked: [num_experts, in_features, out_features] - - Conversion: weight_gkn[e] = nn_linear_weight[e].t() -""" - -import pytest -import torch -import torch.nn as nn -import torch.nn.functional as F - - -# --------------------------------------------------------------------------- -# Reference implementation -- per-expert loop with nn.Linear weights -# --------------------------------------------------------------------------- - - -def reference_moe_forward( - hidden_states, - routing_weights, - selected_experts, - gate_proj_gkn, - up_proj_gkn, - down_proj_gkn, - num_experts, -): - """Naive per-expert loop MoE forward using (G,K,N) weights directly.""" - num_tokens = hidden_states.shape[0] - hidden_dim = hidden_states.shape[1] - output = torch.zeros(num_tokens, hidden_dim, dtype=hidden_states.dtype, device=hidden_states.device) - - expert_mask = F.one_hot(selected_experts, num_classes=num_experts).permute(2, 1, 0) - for expert_idx in range(num_experts): - idx, top_x = torch.where(expert_mask[expert_idx]) - if top_x.numel() == 0: - continue - current_state = hidden_states[top_x] - current_weights = routing_weights[top_x, idx] - gate_out = torch.matmul(current_state, gate_proj_gkn[expert_idx]) - up_out = torch.matmul(current_state, up_proj_gkn[expert_idx]) - expert_out = torch.matmul(F.silu(gate_out) * up_out, down_proj_gkn[expert_idx]) - expert_out = expert_out * current_weights.unsqueeze(-1) - output.index_add_(0, top_x, expert_out.to(hidden_states.dtype)) - return output - - -def make_gkn_weights_from_linear(experts_list): - """Convert list of nn.Module experts to stacked (G,K,N) format.""" - gate = torch.stack([e.gate_proj.weight.t() for e in experts_list]) - up = torch.stack([e.up_proj.weight.t() for e in experts_list]) - down = torch.stack([e.down_proj.weight.t() for e in experts_list]) - return gate, up, down - - -class ExpertMLP(nn.Module): - """Single expert MLP for reference.""" - - def __init__(self, hidden_size, intermediate_size): - super().__init__() - self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) - self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) - self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) - - def forward(self, x): - return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) - - -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - -class TestGKNWeightFormat: - """Test (G,K,N) weight format: matmul equivalence, shapes, and reference MoE vs HF loop.""" - - def test_gkn_format_correctness(self): - """Single expert matmul equivalence, stacked shapes, and reference vs HF loop.""" - # --- Single expert matmul equivalence --- - torch.manual_seed(42) - hidden_size, intermediate_size = 64, 128 - linear = nn.Linear(hidden_size, intermediate_size, bias=False) - x = torch.randn(8, hidden_size) - ref_out = linear(x) - w_kn = linear.weight.t() - gkn_out = x @ w_kn - torch.testing.assert_close(gkn_out, ref_out, atol=1e-5, rtol=1e-5) - - # --- Stacked weight shapes --- - num_experts = 8 - experts = [ExpertMLP(hidden_size, intermediate_size) for _ in range(num_experts)] - gate, up, down = make_gkn_weights_from_linear(experts) - assert gate.shape == (num_experts, hidden_size, intermediate_size) - assert up.shape == (num_experts, hidden_size, intermediate_size) - assert down.shape == (num_experts, intermediate_size, hidden_size) - - # --- Reference MoE vs HF loop --- - torch.manual_seed(42) - num_experts2, hidden_size2, intermediate_size2 = 4, 64, 128 - num_tokens, top_k = 16, 2 - experts2 = [ExpertMLP(hidden_size2, intermediate_size2) for _ in range(num_experts2)] - gate_gkn, up_gkn, down_gkn = make_gkn_weights_from_linear(experts2) - hidden_states = torch.randn(num_tokens, hidden_size2) - selected_experts = torch.randint(0, num_experts2, (num_tokens, top_k)) - routing_weights = torch.softmax(torch.randn(num_tokens, top_k), dim=-1) - - output_gkn = reference_moe_forward( - hidden_states, - routing_weights, - selected_experts, - gate_gkn, - up_gkn, - down_gkn, - num_experts2, - ) - - output_hf = torch.zeros_like(hidden_states) - expert_mask = F.one_hot(selected_experts, num_classes=num_experts2).permute(2, 1, 0) - for expert_idx in range(num_experts2): - idx, top_x = torch.where(expert_mask[expert_idx]) - if top_x.numel() == 0: - continue - current_state = hidden_states[top_x] - current_weights = routing_weights[top_x, idx] - expert_out = experts2[expert_idx](current_state) * current_weights.unsqueeze(-1) - output_hf.index_add_(0, top_x, expert_out) - torch.testing.assert_close(output_gkn, output_hf, atol=1e-5, rtol=1e-5) - - -class TestCheckpointLoadingGKN: - """Test checkpoint loading produces correct (G,K,N) weights and output.""" - - @staticmethod - def _get_buffer_class(): - try: - from xorl.models.checkpoint_handlers.buffers import ExpertWeightBuffer # noqa: PLC0415 - - return ExpertWeightBuffer - except (ImportError, ModuleNotFoundError): - pytest.skip("checkpoint_handlers import requires transformers") - - def test_checkpoint_loading_and_output(self): - """ExpertWeightBuffer transposes correctly and loaded weights produce correct output.""" - ExpertWeightBuffer = self._get_buffer_class() - - torch.manual_seed(42) - num_experts, hidden_size, intermediate_size = 4, 32, 64 - experts = [ExpertMLP(hidden_size, intermediate_size) for _ in range(num_experts)] - - buf = ExpertWeightBuffer(num_experts) - for ei in range(num_experts): - buf.add(0, ei, "gate", experts[ei].gate_proj.weight.detach()) - buf.add(0, ei, "up", experts[ei].up_proj.weight.detach()) - buf.add(0, ei, "down", experts[ei].down_proj.weight.detach()) - - gate_stacked = buf.pop_stacked(0, "gate") - up_stacked = buf.pop_stacked(0, "up") - down_stacked = buf.pop_stacked(0, "down") - - # Shapes and values - assert gate_stacked.shape == (num_experts, hidden_size, intermediate_size) - assert up_stacked.shape == (num_experts, hidden_size, intermediate_size) - assert down_stacked.shape == (num_experts, intermediate_size, hidden_size) - - for e in range(num_experts): - torch.testing.assert_close(gate_stacked[e], experts[e].gate_proj.weight.t()) - torch.testing.assert_close(up_stacked[e], experts[e].up_proj.weight.t()) - torch.testing.assert_close(down_stacked[e], experts[e].down_proj.weight.t()) - - # Loaded weights produce correct output - num_tokens, top_k = 16, 2 - hidden_states = torch.randn(num_tokens, hidden_size) - selected_experts = torch.randint(0, num_experts, (num_tokens, top_k)) - routing_weights = torch.softmax(torch.randn(num_tokens, top_k), dim=-1) - - output_gkn = reference_moe_forward( - hidden_states, - routing_weights, - selected_experts, - gate_stacked, - up_stacked, - down_stacked, - num_experts, - ) - - output_hf = torch.zeros_like(hidden_states) - expert_mask = F.one_hot(selected_experts, num_classes=num_experts).permute(2, 1, 0) - for ei in range(num_experts): - idx, top_x = torch.where(expert_mask[ei]) - if top_x.numel() == 0: - continue - current_state = hidden_states[top_x] - cw = routing_weights[top_x, idx] - expert_out = experts[ei](current_state) * cw.unsqueeze(-1) - output_hf.index_add_(0, top_x, expert_out) - torch.testing.assert_close(output_gkn, output_hf, atol=1e-5, rtol=1e-5) - - -@pytest.mark.gpu -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") -class TestBackendGKN: - """Test eager, native, and triton backends with (G,K,N) format, plus MoEBlock cross-backend agreement.""" - - def test_backends_match_reference_and_agree(self): - """Eager, native, triton backends match reference; all MoEBlock backends agree.""" - # --- Eager backend --- - try: - from xorl.models.layers.moe.backend.eager import eager_expert_forward # noqa: PLC0415 - except (ImportError, ModuleNotFoundError): - pytest.skip("eager backend import requires transformers") - - torch.manual_seed(42) - num_experts, hidden_size, intermediate_size = 4, 64, 128 - num_tokens, top_k = 32, 2 - - experts = [ExpertMLP(hidden_size, intermediate_size) for _ in range(num_experts)] - gate_gkn, up_gkn, down_gkn = make_gkn_weights_from_linear(experts) - - hidden_states_cpu = torch.randn(num_tokens, hidden_size) - act_fn = "silu" - - for expert_idx in range(num_experts): - ref_out = experts[expert_idx](hidden_states_cpu) - eager_out = eager_expert_forward(hidden_states_cpu, expert_idx, gate_gkn, up_gkn, down_gkn, act_fn) - torch.testing.assert_close(eager_out, ref_out, atol=1e-5, rtol=1e-5) - - # --- Native backend --- - from xorl.models.layers.moe.backend.native import native_expert_forward # noqa: PLC0415 - - device, dtype = "cuda", torch.bfloat16 - gate_cuda = gate_gkn.to(device).to(dtype) - up_cuda = up_gkn.to(device).to(dtype) - down_cuda = down_gkn.to(device).to(dtype) - hidden_states = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype) - selected = torch.randint(0, num_experts, (num_tokens, top_k), device=device) - rw = torch.softmax(torch.randn(num_tokens, top_k, device=device, dtype=dtype), dim=-1) - - native_out = native_expert_forward( - hidden_states, - rw, - selected, - gate_cuda, - up_cuda, - down_cuda, - num_experts, - ) - ref_out = reference_moe_forward(hidden_states, rw, selected, gate_cuda, up_cuda, down_cuda, num_experts) - torch.testing.assert_close(native_out, ref_out, atol=0.02, rtol=0.02) - - # --- Triton backend --- - try: - from xorl.utils.import_utils import is_fused_moe_available # noqa: PLC0415 - - if not is_fused_moe_available(): - raise ImportError - from xorl.ops.moe.triton import TritonMoeExpertsFunction # noqa: PLC0415 - - gate_up_cuda = torch.cat([gate_cuda, up_cuda], dim=-1) - triton_out = TritonMoeExpertsFunction.apply( - num_experts, - rw, - selected, - hidden_states, - gate_cuda, - up_cuda, - down_cuda, - gate_up_cuda, - ) - torch.testing.assert_close(triton_out, ref_out, atol=0.01, rtol=0.01) - - # Triton backward: gradients exist and non-zero - gate_g = torch.randn( - num_experts, hidden_size, intermediate_size, device=device, dtype=dtype, requires_grad=True - ) - up_g = torch.randn( - num_experts, hidden_size, intermediate_size, device=device, dtype=dtype, requires_grad=True - ) - down_g = torch.randn( - num_experts, intermediate_size, hidden_size, device=device, dtype=dtype, requires_grad=True - ) - h_g = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype, requires_grad=True) - gate_up_g = torch.cat([gate_g, up_g], dim=-1) - out_g = TritonMoeExpertsFunction.apply(num_experts, rw, selected, h_g, gate_g, up_g, down_g, gate_up_g) - out_g.sum().backward() - assert h_g.grad is not None and h_g.grad.abs().max() > 0 - assert gate_g.grad is not None and gate_g.grad.abs().max() > 0 - except (ImportError, ModuleNotFoundError): - pass # triton not available, skip - - # --- MoEBlock all backends agree --- - try: - from xorl.models.layers.moe import MOE_EXPERT_BACKENDS, MoEBlock # noqa: PLC0415 - except (ImportError, ModuleNotFoundError): - return # skip if not importable - - experts_ref = [ExpertMLP(hidden_size, intermediate_size) for _ in range(num_experts)] - g_gkn, u_gkn, d_gkn = make_gkn_weights_from_linear(experts_ref) - gate_w = torch.randn(num_experts, hidden_size) - hidden_block = torch.randn(2, 8, hidden_size, device=device, dtype=dtype) - outputs = {} - - for backend in MOE_EXPERT_BACKENDS: - block = MoEBlock( - hidden_size=hidden_size, - num_experts=num_experts, - top_k=2, - intermediate_size=intermediate_size, - moe_implementation=backend, - ) - with torch.no_grad(): - block.experts.gate_proj.copy_(g_gkn) - block.experts.up_proj.copy_(u_gkn) - block.experts.down_proj.copy_(d_gkn) - block.gate.weight.copy_(gate_w) - block = block.to(device).to(dtype) - out, _ = block(hidden_block) - outputs[backend] = out - - backends = list(outputs.keys()) - for i in range(len(backends)): - for j in range(i + 1, len(backends)): - torch.testing.assert_close( - outputs[backends[i]], - outputs[backends[j]], - atol=0.05, - rtol=0.02, - msg=f"Backend mismatch: {backends[i]} vs {backends[j]}", - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) diff --git a/tests/ops/test_moe_nongated.py b/tests/ops/test_moe_nongated.py index 31721f18..5ffecec5 100644 --- a/tests/ops/test_moe_nongated.py +++ b/tests/ops/test_moe_nongated.py @@ -9,15 +9,7 @@ import torch import torch.nn.functional as F - -def _import_experts(): - """Import MoEExperts or skip (mirrors test_eager_vs_native_moe.py).""" - try: - from xorl.models.layers.moe.experts import MoEExperts # noqa: PLC0415 - - return MoEExperts - except Exception as e: - pytest.skip(f"Cannot import MoE layers: {e}") +from xorl.models.layers.moe.experts import MoEExperts def _make_routing(num_tokens, num_experts, top_k, device, dtype): @@ -70,7 +62,6 @@ def _reference_forward(hidden_states, routing_weights, selected_experts, up_proj @pytest.mark.cpu def test_eager_nongated_matches_reference(): """Eager non-gated forward + gradients match a plain per-expert torch loop.""" - MoEExperts = _import_experts() ne, hd, inter, top_k, num_tokens = 4, 16, 24, 2, 10 # hd is a latent dim, not a model hidden size experts = _make_nongated_experts(MoEExperts, ne, hd, inter, "eager", "cpu", torch.float32) @@ -97,37 +88,13 @@ def test_eager_nongated_matches_reference(): torch.testing.assert_close(experts.gate_up_proj.grad, up_ref.grad) torch.testing.assert_close(experts.down_proj.grad, down_ref.grad) + _assert_nongated_constructor_rejects_unsupported_policies() -@pytest.mark.cpu -def test_relu2_activation_registry(): - """relu2 is registered, normalized, and equals relu(x)**2 when gate ≡ up.""" - from xorl.ops.moe.activations import ( # noqa: PLC0415 - MOE_ACTIVATIONS, - UNGATED_HIDDEN_ACTS, - apply_moe_activation, - normalize_hidden_act, - ) - - assert normalize_hidden_act("relu2") == "relu2" - assert "relu2" in MOE_ACTIVATIONS - assert "relu2" in UNGATED_HIDDEN_ACTS - x = torch.randn(32) - torch.testing.assert_close(apply_moe_activation("relu2", x, x), torch.square(F.relu(x))) - - -@pytest.mark.cpu -def test_nongated_quack_raises(): - """quack backend rejects non-gated experts with a clear error.""" - MoEExperts = _import_experts() +def _assert_nongated_constructor_rejects_unsupported_policies(): + """Non-gated experts reject unsupported backends and gated activations.""" with pytest.raises(NotImplementedError, match="non-gated"): MoEExperts(4, 16, 24, hidden_act="relu2", moe_implementation="quack", gated=False) - - -@pytest.mark.cpu -def test_nongated_rejects_gated_activation(): - """Non-gated experts require an ungated activation (relu2).""" - MoEExperts = _import_experts() with pytest.raises(ValueError, match="non-gated"): MoEExperts(4, 16, 24, hidden_act="silu", moe_implementation="eager", gated=False) @@ -139,50 +106,52 @@ def test_nongated_rejects_gated_activation(): @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("backend", ["triton", "native"]) -def test_nongated_backend_matches_eager(backend): +def test_nongated_backends_match_eager(): """Triton/native non-gated forward + gradients match eager.""" - MoEExperts = _import_experts() from xorl.models.layers.moe.backend import MOE_EXPERT_BACKENDS # noqa: PLC0415 - if backend not in MOE_EXPERT_BACKENDS: - pytest.skip(f"{backend} backend not available") - - ne, hd, inter, top_k, num_tokens = 8, 64, 128, 2, 32 - device, dtype = "cuda", torch.bfloat16 - - eager = _make_nongated_experts(MoEExperts, ne, hd, inter, "eager", device, dtype) - other = MoEExperts(ne, hd, inter, hidden_act="relu2", moe_implementation=backend, gated=False).to(device, dtype) - with torch.no_grad(): - other.gate_up_proj.copy_(eager.gate_up_proj) - other.down_proj.copy_(eager.down_proj) - - torch.manual_seed(7) - x_eager = torch.randn(num_tokens, hd, device=device, dtype=dtype, requires_grad=True) - x_other = x_eager.detach().clone().requires_grad_(True) - routing_weights, selected_experts = _make_routing(num_tokens, ne, top_k, device, dtype) - - out_eager = _eager_moe_forward(eager, x_eager, routing_weights, selected_experts) - out_other = other(x_other, routing_weights, selected_experts) - torch.testing.assert_close(out_other, out_eager, atol=0.05, rtol=0.05, msg=f"{backend} forward mismatch") - - grad_out = torch.randn_like(out_eager) - out_eager.backward(grad_out) - out_other.backward(grad_out) - - atol, rtol = 0.05, 0.05 - torch.testing.assert_close(x_other.grad, x_eager.grad, atol=atol, rtol=rtol, msg=f"{backend} input grad mismatch") - torch.testing.assert_close( - other.gate_up_proj.grad, - eager.gate_up_proj.grad, - atol=atol, - rtol=rtol, - msg=f"{backend} gate_up_proj grad mismatch", - ) - torch.testing.assert_close( - other.down_proj.grad, eager.down_proj.grad, atol=atol, rtol=rtol, msg=f"{backend} down_proj grad mismatch" - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) + backends = ("triton", "native") + missing = set(backends) - MOE_EXPERT_BACKENDS.keys() + assert not missing, f"shipped non-gated backends failed to register: {sorted(missing)}" + + for backend in backends: + ne, hd, inter, top_k, num_tokens = 8, 64, 128, 2, 32 + device, dtype = "cuda", torch.bfloat16 + + eager = _make_nongated_experts(MoEExperts, ne, hd, inter, "eager", device, dtype) + other = MoEExperts(ne, hd, inter, hidden_act="relu2", moe_implementation=backend, gated=False).to(device, dtype) + with torch.no_grad(): + other.gate_up_proj.copy_(eager.gate_up_proj) + other.down_proj.copy_(eager.down_proj) + + torch.manual_seed(7) + x_eager = torch.randn(num_tokens, hd, device=device, dtype=dtype, requires_grad=True) + x_other = x_eager.detach().clone().requires_grad_(True) + routing_weights, selected_experts = _make_routing(num_tokens, ne, top_k, device, dtype) + + out_eager = _eager_moe_forward(eager, x_eager, routing_weights, selected_experts) + out_other = other(x_other, routing_weights, selected_experts) + torch.testing.assert_close(out_other, out_eager, atol=0.05, rtol=0.05, msg=f"{backend} forward mismatch") + + grad_out = torch.randn_like(out_eager) + out_eager.backward(grad_out) + out_other.backward(grad_out) + + atol, rtol = 0.05, 0.05 + torch.testing.assert_close( + x_other.grad, x_eager.grad, atol=atol, rtol=rtol, msg=f"{backend} input grad mismatch" + ) + torch.testing.assert_close( + other.gate_up_proj.grad, + eager.gate_up_proj.grad, + atol=atol, + rtol=rtol, + msg=f"{backend} gate_up_proj grad mismatch", + ) + torch.testing.assert_close( + other.down_proj.grad, + eager.down_proj.grad, + atol=atol, + rtol=rtol, + msg=f"{backend} down_proj grad mismatch", + ) diff --git a/tests/ops/test_moe_ops.py b/tests/ops/test_moe_ops.py index 0dd9e9d2..3cd7fa84 100644 --- a/tests/ops/test_moe_ops.py +++ b/tests/ops/test_moe_ops.py @@ -8,8 +8,12 @@ import torch -# Mark all tests as GPU since MoE ops require CUDA -pytestmark = pytest.mark.gpu +# MoE kernels require CUDA. Skip unsupported hosts before entering the test +# body, but let failures importing XORL's own kernel modules surface normally. +pytestmark = [ + pytest.mark.gpu, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] def naive_expert_histogram(input: torch.Tensor, num_bins: int) -> torch.Tensor: @@ -49,20 +53,9 @@ def naive_moe_add_gather(x: torch.Tensor, y: torch.Tensor, index: torch.Tensor) class TestExpertHistogramAndIndex: """Tests for expert_histogram and moe_index_compute.""" - def test_histogram_and_index_compute(self): + def _assert_histogram_and_index_compute(self): """Histogram: basic, 2D, large, int64 inputs. Index compute: ordering and uniqueness.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - try: - from xorl.ops.group_gemm.kernel.moe import expert_histogram, moe_index_compute # noqa: PLC0415 - except ImportError: - pytest.skip("moe ops not available") - - # Basic 1D histogram - input_1d = torch.tensor([0, 1, 2, 0, 1, 3, 2, 0], dtype=torch.int32).cuda() - output = expert_histogram(input_1d, 4) - assert torch.equal(output.cpu(), naive_expert_histogram(input_1d, 4).cpu()) - assert output.tolist() == [3, 2, 2, 1] + from xorl.ops.group_gemm.kernel.moe import expert_histogram, moe_index_compute # noqa: PLC0415 # 2D input (topk routing) input_2d = torch.tensor([[0, 1], [2, 0], [1, 3], [2, 0]], dtype=torch.int32).cuda() @@ -104,10 +97,7 @@ def test_histogram_and_index_compute(self): class TestDeterministicScatter: - """Deterministic scatter (default ON): run-invariant permutation from moe_index_compute. - - XORL_MOE_DETERMINISTIC_SCATTER=0 is the escape hatch back to the atomics kernel. - """ + """Run-invariant permutation from ``moe_index_compute``.""" @staticmethod def _naive_stable_index(experts_for_tokens: torch.Tensor, num_experts: int) -> torch.Tensor: @@ -127,15 +117,9 @@ def _naive_stable_index(experts_for_tokens: torch.Tensor, num_experts: int) -> t seen[e] += 1 return torch.tensor(out, dtype=torch.int64).view(experts_for_tokens.shape) - def test_deterministic_scatter(self, monkeypatch): - """Run-invariance, validity vs stock coverage, exact stable order, escape-hatch routing.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - try: - import xorl.ops.group_gemm.kernel.moe as moe_mod # noqa: PLC0415 - from xorl.ops.group_gemm.kernel.moe import expert_histogram, moe_index_compute # noqa: PLC0415 - except ImportError: - pytest.skip("moe ops not available") + def _assert_deterministic_scatter(self): + """Run-invariance, full slot coverage, and exact stable order.""" + from xorl.ops.group_gemm.kernel.moe import expert_histogram, moe_index_compute # noqa: PLC0415 torch.manual_seed(0) num_experts, M, topk = 16, 4096, 8 @@ -143,27 +127,22 @@ def test_deterministic_scatter(self, monkeypatch): hist = expert_histogram(experts, num_experts) cumsum = torch.cumsum(hist, dim=0) - # Stock atomics draw (escape hatch =0) for the coverage comparison. - monkeypatch.setenv("XORL_MOE_DETERMINISTIC_SCATTER", "0") - stock = moe_index_compute(experts, cumsum) - - # Default ON: flag unset routes through the deterministic path. - monkeypatch.delenv("XORL_MOE_DETERMINISTIC_SCATTER", raising=False) det_a = moe_index_compute(experts, cumsum) det_b = moe_index_compute(experts, cumsum) # Run-invariance: two independent calls produce the identical permutation. assert torch.equal(det_a, det_b) - # Validity: same shape/dtype contract as stock, full slot coverage, and per-expert - # slot SETS identical to the stock draw (both fill [cumsum[e-1], cumsum[e])). - assert det_a.shape == stock.shape - assert det_a.dtype == stock.dtype + # Validity: full slot coverage and each expert owns its cumsum region. + assert det_a.shape == experts.shape + assert det_a.dtype == torch.int64 full = torch.arange(M * topk, device=det_a.device) assert torch.equal(det_a.flatten().sort()[0], full) for e in range(num_experts): mask = experts == e - assert torch.equal(det_a[mask].sort()[0], stock[mask].sort()[0]) + start = 0 if e == 0 else int(cumsum[e - 1]) + stop = int(cumsum[e]) + assert torch.equal(det_a[mask].sort()[0], torch.arange(start, stop, device=det_a.device)) # Exact value: stable order = flattened (token, k) order within each expert. expected = self._naive_stable_index(experts.cpu(), num_experts).cuda() @@ -172,60 +151,18 @@ def test_deterministic_scatter(self, monkeypatch): # int64 ids too (the model path passes topk_ids as int64). assert torch.equal(moe_index_compute(experts.long().contiguous(), cumsum), expected) - # Truthy values keep routing through the deterministic path. - monkeypatch.setenv("XORL_MOE_DETERMINISTIC_SCATTER", "1") - assert torch.equal(moe_index_compute(experts, cumsum), expected) - - # Escape hatch: explicit falsey values must never route through the - # deterministic path (old atomics behavior). - def _boom(*args, **kwargs): - raise AssertionError("deterministic path used with escape hatch set") - - monkeypatch.setattr(moe_mod, "_moe_index_compute_deterministic", _boom) - for falsey in ("0", "false", "off", ""): - monkeypatch.setenv("XORL_MOE_DETERMINISTIC_SCATTER", falsey) - out = moe_index_compute(experts, cumsum) - assert torch.equal(out.flatten().sort()[0], full) - - # Default ON: unset must route through the deterministic path. - monkeypatch.delenv("XORL_MOE_DETERMINISTIC_SCATTER", raising=False) - with pytest.raises(AssertionError, match="deterministic path used"): - moe_index_compute(experts, cumsum) - class TestMoEGatherScatterAddGather: """Tests for moe_gather, moe_scatter, moe_add_gather, and scatter-gather roundtrip.""" def test_gather_scatter_add_gather(self): """Gather (basic, overlapping, large), scatter, scatter-gather roundtrip, add_gather, equivalence.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - try: - from xorl.ops.group_gemm.kernel.moe import moe_add_gather, moe_gather, moe_scatter # noqa: PLC0415 - except ImportError: - pytest.skip("moe ops not available") - - # --- Gather: basic --- - M, topk, N = 4, 2, 8 - x = torch.randn(M * topk, N, dtype=torch.float16).cuda() - index = torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7]], dtype=torch.int32).cuda() - output = moe_gather(x, index) - assert output.shape == (M, N) - assert torch.allclose(output.float(), naive_moe_gather(x, index).float(), rtol=1e-2, atol=1e-2) - - # Gather: overlapping indices - M2, topk2, N2 = 3, 2, 16 - x2 = torch.randn(M2 * topk2, N2, dtype=torch.bfloat16).cuda() - index2 = torch.tensor([[0, 1], [1, 2], [0, 2]], dtype=torch.int32).cuda() - assert torch.allclose( - moe_gather(x2, index2).float(), - naive_moe_gather(x2, index2).float(), - rtol=1e-2, - atol=1e-2, - ) + TestExpertHistogramAndIndex()._assert_histogram_and_index_compute() + TestDeterministicScatter()._assert_deterministic_scatter() + from xorl.ops.group_gemm.kernel.moe import moe_add_gather, moe_gather, moe_scatter # noqa: PLC0415 - # Gather: large dimensions - M3, topk3, N3 = 1024, 2, 4096 + # Gather across multiple hidden-dimension blocks. + M3, topk3, N3 = 65, 2, 2049 x3 = torch.randn(M3 * topk3, N3, dtype=torch.float16).cuda() index3 = torch.arange(M3 * topk3, dtype=torch.int32).cuda().reshape(M3, topk3) assert torch.allclose( @@ -235,13 +172,6 @@ def test_gather_scatter_add_gather(self): atol=1e-2, ) - # --- Scatter: basic --- - xs = torch.randn(M, N, dtype=torch.float16).cuda() - idx_s = torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7]], dtype=torch.int32).cuda() - out_s = moe_scatter(xs, idx_s) - assert out_s.shape == (M * topk, N) - assert torch.allclose(out_s.float(), naive_moe_scatter(xs, idx_s).float(), rtol=1e-2, atol=1e-2) - # --- Scatter-gather roundtrip --- M4, topk4, N4 = 8, 2, 16 x4 = torch.randn(M4, N4, dtype=torch.bfloat16).cuda() @@ -251,56 +181,14 @@ def test_gather_scatter_add_gather(self): expected = x4 * topk4 assert torch.allclose(gathered, expected, rtol=1e-2, atol=1e-2) - # --- Add-gather: basic --- - xa = torch.randn(M * topk, N, dtype=torch.float16).cuda() - ya = torch.randn(M * topk, N, dtype=torch.float16).cuda() - idx_a = torch.arange(M * topk, dtype=torch.int32).cuda().reshape(M, topk) - out_ag = moe_add_gather(xa, ya, idx_a) - assert out_ag.shape == (M, N) - assert torch.allclose(out_ag.float(), naive_moe_add_gather(xa, ya, idx_a).float(), rtol=1e-2, atol=1e-2) - - # Add-gather equivalence with manual add+gather + # Add-gather against an independent reference. M5, topk5, N5 = 8, 2, 16 x5 = torch.randn(M5 * topk5, N5, dtype=torch.bfloat16).cuda() y5 = torch.randn(M5 * topk5, N5, dtype=torch.bfloat16).cuda() idx5 = torch.arange(M5 * topk5, dtype=torch.int32).cuda().reshape(M5, topk5) assert torch.allclose( moe_add_gather(x5, y5, idx5), - moe_gather(x5 + y5, idx5), + naive_moe_add_gather(x5, y5, idx5), rtol=1e-3, atol=1e-3, ) - - -class TestMoEIntegration: - """Integration tests combining multiple MoE operations.""" - - def test_full_moe_pipeline(self): - """Complete MoE forward pipeline: histogram -> index -> scatter -> gather.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - try: - from xorl.ops.group_gemm.kernel.moe import ( # noqa: PLC0415 - expert_histogram, - moe_gather, - moe_index_compute, - moe_scatter, - ) - except ImportError: - pytest.skip("moe ops not available") - - M, topk, N, num_experts = 16, 2, 32, 4 - - hidden_states = torch.randn(M, N, dtype=torch.float16).cuda() - experts_for_tokens = torch.randint(0, num_experts, (M, topk), dtype=torch.int32).cuda() - - histogram = expert_histogram(experts_for_tokens.flatten(), num_experts) - cumsum = torch.cumsum(histogram, dim=0).int().cuda() - indices = moe_index_compute(experts_for_tokens, cumsum) - - scattered = moe_scatter(hidden_states, indices) - expert_outputs = scattered # identity for simplicity - final_output = moe_gather(expert_outputs, indices) - - assert final_output.shape == (M, N) - assert final_output.abs().max() > 0 diff --git a/tests/ops/test_moe_routing_weight_position.py b/tests/ops/test_moe_routing_weight_position.py index 30b68a1b..a00ea361 100644 --- a/tests/ops/test_moe_routing_weight_position.py +++ b/tests/ops/test_moe_routing_weight_position.py @@ -1,7 +1,6 @@ """Routing-weight position (before vs after the down GEMM) in ``TritonEPGroupGemm``. -Covers the ``moe_routing_weights_before_down`` config knob and its -``XORL_MOE_ROUTING_WEIGHTS_BEFORE_DOWN`` env override: +Covers the ``moe_routing_weights_before_down`` config knob: - Gradient parity of BOTH positions against an fp64 eager reference. The two positions are mathematically identical (a per-row scalar commutes through the @@ -18,8 +17,7 @@ import torch.nn.functional as F import xorl.ops.moe.triton as moe_triton -from tests.ops.test_ep_routing_scores import _counts_from_cumsum, _patch_ep_kernels -from xorl.arguments import ModelArguments +from tests._helpers.moe import counts_from_cumsum, patch_ep_kernels pytestmark = pytest.mark.cpu @@ -50,7 +48,7 @@ def _fp64_reference(permute_tokens, cumsum, gate_up_proj, down_proj, intermediat outputs = [] start = 0 - for expert_idx, count in enumerate(_counts_from_cumsum(cumsum)): + for expert_idx, count in enumerate(counts_from_cumsum(cumsum)): end = start + count xs = x[start:end] gate_up = xs @ gup[expert_idx] @@ -79,8 +77,9 @@ def _run_position(module, before_down, inputs, scores_require_grad=True): return out.detach(), grads -def test_both_routing_positions_match_fp64_reference(monkeypatch): - module = _patch_ep_kernels(monkeypatch, "xorl.ops.moe.triton") +def test_routing_weight_position_numerical_contract(monkeypatch): + patch_ep_kernels(monkeypatch, moe_triton) + module = moe_triton inputs = _make_inputs() ref_out, ref_grads = _fp64_reference(*inputs) @@ -101,10 +100,15 @@ def test_both_routing_positions_match_fp64_reference(monkeypatch): eb = max(per_position[True][key], 1e-30) assert max(ea / eb, eb / ea) < 3.0, f"{key}: error class diverged (after={ea:.3e}, before={eb:.3e})" + _assert_before_down_without_score_grad_matches_reference(monkeypatch) + with monkeypatch.context() as config_patch: + _assert_routing_weight_position_configuration_policy(config_patch) -def test_before_down_without_score_grad_matches_fp64_reference(monkeypatch): + +def _assert_before_down_without_score_grad_matches_reference(monkeypatch): """The in-place score fold on the recomputed intermediate (no router grad) is safe.""" - module = _patch_ep_kernels(monkeypatch, "xorl.ops.moe.triton") + patch_ep_kernels(monkeypatch, moe_triton) + module = moe_triton inputs = _make_inputs(seed=7) _, ref_grads = _fp64_reference(*inputs) @@ -113,8 +117,7 @@ def test_before_down_without_score_grad_matches_fp64_reference(monkeypatch): torch.testing.assert_close(grads[key].double(), ref_grads[key], rtol=1e-3, atol=1e-4) -def test_routing_weight_position_knob(monkeypatch): - monkeypatch.delenv("XORL_MOE_ROUTING_WEIGHTS_BEFORE_DOWN", raising=False) +def _assert_routing_weight_position_configuration_policy(monkeypatch): monkeypatch.setattr(moe_triton, "_ROUTING_WEIGHTS_BEFORE_DOWN_CONFIG", False) assert moe_triton.routing_weights_before_down() is False @@ -123,54 +126,49 @@ def test_routing_weight_position_knob(monkeypatch): moe_triton.set_routing_weights_before_down(False) assert moe_triton.routing_weights_before_down() is False - # Env var force-enables regardless of the config default (read lazily, so it - # keeps working when set after import). - monkeypatch.setenv("XORL_MOE_ROUTING_WEIGHTS_BEFORE_DOWN", "1") - assert moe_triton.routing_weights_before_down() is True - + _assert_auto_resolution_regimes(monkeypatch) + _assert_auto_resolution_disabled_under_parity_opt_in(monkeypatch) + _assert_explicit_true_overrides_regime(monkeypatch) + _assert_explicit_false_overrides_regime(monkeypatch) + _assert_invalid_setting_raises() -def test_model_arguments_field_default(): - args = ModelArguments(model_path="Qwen/Qwen3-Coder-30B-A3B-Instruct") - assert args.moe_routing_weights_before_down == "auto" - -@pytest.mark.parametrize( - ("train_router", "ep_dispatch", "expected"), - [ +def _assert_auto_resolution_regimes(monkeypatch): + # Pin the stock expert tree. Unset is an auto mode that enables the serving + # kernel on supported EP1 CUDA lanes, so it is not equivalent to opt-out. + monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS", "0") + for train_router, ep_dispatch, expected in ( (True, "alltoall", True), # the measured-win regime (True, "deepep", False), (False, "alltoall", False), (False, "deepep", False), - ], -) -def test_auto_resolution_regimes(monkeypatch, train_router, ep_dispatch, expected): - # Pin the stock expert tree. Unset is an auto mode that enables the serving - # kernel on supported EP1 CUDA lanes, so it is not equivalent to opt-out. - monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS", "0") - resolved = moe_triton.resolve_routing_weights_before_down( - "auto", train_router=train_router, ep_dispatch=ep_dispatch - ) - assert resolved is expected + ): + resolved = moe_triton.resolve_routing_weights_before_down( + "auto", train_router=train_router, ep_dispatch=ep_dispatch + ) + assert resolved is expected -def test_auto_resolution_disabled_under_parity_opt_in(monkeypatch): +def _assert_auto_resolution_disabled_under_parity_opt_in(monkeypatch): """The XORL_MOE_SGLANG_FUSED_EXPERTS parity lane keeps the historical after-down tree.""" monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS", "1") assert moe_triton.resolve_routing_weights_before_down("auto", train_router=True, ep_dispatch="alltoall") is False -@pytest.mark.parametrize("setting", [True, "true", "1"]) -def test_explicit_true_overrides_regime(monkeypatch, setting): +def _assert_explicit_true_overrides_regime(monkeypatch): monkeypatch.setenv("XORL_MOE_SGLANG_FUSED_EXPERTS", "1") - assert moe_triton.resolve_routing_weights_before_down(setting, train_router=False, ep_dispatch="deepep") is True + for setting in (True, "true"): + assert moe_triton.resolve_routing_weights_before_down(setting, train_router=False, ep_dispatch="deepep") is True -@pytest.mark.parametrize("setting", [False, "false", "0"]) -def test_explicit_false_overrides_regime(monkeypatch, setting): +def _assert_explicit_false_overrides_regime(monkeypatch): monkeypatch.delenv("XORL_MOE_SGLANG_FUSED_EXPERTS", raising=False) - assert moe_triton.resolve_routing_weights_before_down(setting, train_router=True, ep_dispatch="alltoall") is False + for setting in (False, "false"): + assert ( + moe_triton.resolve_routing_weights_before_down(setting, train_router=True, ep_dispatch="alltoall") is False + ) -def test_invalid_setting_raises(): +def _assert_invalid_setting_raises(): with pytest.raises(ValueError, match="moe_routing_weights_before_down"): moe_triton.resolve_routing_weights_before_down("maybe", train_router=True, ep_dispatch="alltoall") diff --git a/tests/ops/test_moe_torch_compile.py b/tests/ops/test_moe_torch_compile.py index df0f88e8..07d1112b 100644 --- a/tests/ops/test_moe_torch_compile.py +++ b/tests/ops/test_moe_torch_compile.py @@ -83,18 +83,12 @@ def _make_moe_block(moe_backend, hidden_size=128, num_experts=4, top_k=2, interm def _available_backends(): """Return list of available MoE backends on this system.""" + from xorl.utils.import_utils import is_fused_moe_available # noqa: PLC0415 + backends = ["native", "eager"] - try: - from xorl.utils.import_utils import is_fused_moe_available # noqa: PLC0415 - - if is_fused_moe_available(): - backends.append("triton") - except Exception: - pass - try: - backends.append("quack") - except Exception: - pass + if is_fused_moe_available(): + backends.append("triton") + backends.append("quack") return backends @@ -106,10 +100,8 @@ def _available_backends(): # --------------------------------------------------------------------------- -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("moe_backend", AVAILABLE_BACKENDS) -def test_moe_block_compile(moe_backend): - """MoEBlock compile: aot_eager tracing, inductor compile, fullgraph check, and correctness.""" +def _assert_moe_block_compile(moe_backend): + """MoEBlock compile: aot_eager tracing, inductor compile, and correctness.""" # --- aot_eager tracing (forward + backward) --- block = _make_moe_block(moe_backend) compiled_aot = torch.compile(block, fullgraph=False, backend="aot_eager") @@ -130,15 +122,6 @@ def test_moe_block_compile(moe_backend): out2.sum().backward() assert x2.grad is not None - # --- fullgraph=True (strictest -- detects graph breaks) --- - block3 = _make_moe_block(moe_backend) - compiled_fg = torch.compile(block3, fullgraph=True, backend="aot_eager") - x3 = torch.randn(2, 8, 128, device=DEVICE, dtype=DTYPE) - try: - compiled_fg(x3) - except Exception: - pass - # --- correctness: compiled vs uncompiled match --- torch.manual_seed(42) block4 = _make_moe_block(moe_backend) @@ -152,14 +135,25 @@ def test_moe_block_compile(moe_backend): torch.testing.assert_close(ref_logits, comp_logits, atol=0, rtol=0) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_moe_block_decoder_and_full_model_compile_policy(): + for moe_backend in AVAILABLE_BACKENDS: + _assert_moe_block_compile(moe_backend) + _assert_decoder_layer_compile(moe_backend) + + # Lower-level contracts already compile every available MoE backend with + # both compiler backends. Full-model composition only needs each compiler + # path once. + for moe_backend, compile_backend in (("native", "aot_eager"), ("eager", "inductor")): + _assert_full_model_per_layer_compile(moe_backend, compile_backend) + + # --------------------------------------------------------------------------- # Test 2: Qwen3MoeDecoderLayer compile (aot_eager + inductor) # --------------------------------------------------------------------------- -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("moe_backend", AVAILABLE_BACKENDS) -def test_decoder_layer_compile(moe_backend): +def _assert_decoder_layer_compile(moe_backend): """Decoder layer compile: aot_eager and inductor, forward + backward.""" seq_len = 8 @@ -190,19 +184,7 @@ def test_decoder_layer_compile(moe_backend): # --------------------------------------------------------------------------- -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize( - "moe_backend,compile_backend", - [ - ("native", "aot_eager"), - ("native", "inductor"), - ("eager", "aot_eager"), - ("eager", "inductor"), - ] - + ([("triton", "aot_eager"), ("triton", "inductor")] if "triton" in AVAILABLE_BACKENDS else []) - + ([("quack", "aot_eager"), ("quack", "inductor")] if "quack" in AVAILABLE_BACKENDS else []), -) -def test_full_model_per_layer_compile(moe_backend, compile_backend): +def _assert_full_model_per_layer_compile(moe_backend, compile_backend): """Apply torch.compile to each decoder layer, run forward + backward.""" config = _tiny_moe_config(_moe_implementation=moe_backend) @@ -223,412 +205,3 @@ def test_full_model_per_layer_compile(moe_backend, compile_backend): output.last_hidden_state.sum().backward() has_grad = any(p.grad is not None for p in model.parameters() if p.requires_grad) assert has_grad, "No gradients found" - - -# --------------------------------------------------------------------------- -# Benchmark: TFLOPS measurement compiled vs uncompiled -# --------------------------------------------------------------------------- - - -def _moe_flops(batch, seq, hidden, intermediate, num_experts, top_k): - """Estimate FLOPs for one MoE forward pass.""" - tokens = batch * seq - return tokens * top_k * 6 * hidden * intermediate - - -def _benchmark_moe_block(block, x, warmup=30, iters=50): - """Benchmark forward+backward and return median GPU time in seconds.""" - for _ in range(warmup): - out, _ = block(x) - out.sum().backward() - block.zero_grad() - - torch.cuda.synchronize() - times = [] - for _ in range(iters): - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - - start_event.record() - out, _ = block(x) - out.sum().backward() - end_event.record() - - torch.cuda.synchronize() - times.append(start_event.elapsed_time(end_event) / 1000.0) - block.zero_grad() - - times.sort() - return times[len(times) // 2] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("seq_len", [1024, 4096]) -def bench_tflops(seq_len): - """Benchmark TFLOPS for compiled vs uncompiled MoE across backends.""" - hidden = 1024 - intermediate = 2048 - num_experts = 8 - top_k = 2 - batch = 4 - - flops = _moe_flops(batch, seq_len, hidden, intermediate, num_experts, top_k) - x = torch.randn(batch, seq_len, hidden, device=DEVICE, dtype=DTYPE, requires_grad=True) - - results = {} - - for backend in AVAILABLE_BACKENDS: - torch.manual_seed(42) - block = _make_moe_block(backend, hidden, num_experts, top_k, intermediate) - - t_base = _benchmark_moe_block(block, x) - tflops_base = flops / t_base / 1e12 - - torch._dynamo.reset() - compiled_block = torch.compile( - block, - fullgraph=False, - backend="inductor", - dynamic=False, - ) - try: - t_compiled = _benchmark_moe_block(compiled_block, x) - except Exception as e: - print(f" {backend} inductor compile failed: {type(e).__name__}") - results[backend] = { - "base_ms": t_base * 1000, - "compiled_ms": float("nan"), - "base_tflops": tflops_base, - "compiled_tflops": float("nan"), - "speedup": float("nan"), - "compile_backend": "inductor (FAILED)", - } - torch._dynamo.reset() - continue - tflops_compiled = flops / t_compiled / 1e12 - speedup = t_base / t_compiled - results[backend] = { - "base_ms": t_base * 1000, - "compiled_ms": t_compiled * 1000, - "base_tflops": tflops_base, - "compiled_tflops": tflops_compiled, - "speedup": speedup, - "compile_backend": "inductor", - } - - print("\n" + "=" * 90) - print(f" MoE TFLOPS Benchmark (batch={batch}, seq={seq_len}, hidden={hidden}, E={num_experts}, top_k={top_k})") - print(f" FLOPs per fwd: {flops / 1e9:.1f} GFLOP | dynamic=False, warmup=30, iters=50") - print("=" * 95) - print( - f" {'Backend':<10} {'Compiler':<10} {'Base (ms)':>10} {'Compiled (ms)':>14} " - f"{'Base TFLOPS':>12} {'Comp TFLOPS':>12} {'Speedup':>8}" - ) - print("-" * 95) - for backend, r in results.items(): - print( - f" {backend:<10} {r['compile_backend']:<10} {r['base_ms']:>10.2f} {r['compiled_ms']:>14.2f} " - f"{r['base_tflops']:>12.2f} {r['compiled_tflops']:>12.2f} " - f"{r['speedup']:>7.2f}x" - ) - print("=" * 95) - - -# --------------------------------------------------------------------------- -# Benchmark: Peak memory usage compiled vs uncompiled -# --------------------------------------------------------------------------- - - -def _measure_peak_memory(fn, warmup=5): - """Run fn, return peak GPU memory allocated in bytes.""" - for _ in range(warmup): - fn() - - torch.cuda.synchronize() - torch.cuda.reset_peak_memory_stats() - torch.cuda.empty_cache() - - fn() - - torch.cuda.synchronize() - return torch.cuda.max_memory_allocated() - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("seq_len", [1024, 4096]) -def bench_memory(seq_len): - """Measure peak GPU memory for compiled vs uncompiled MoE fwd+bwd.""" - hidden = 1024 - intermediate = 2048 - num_experts = 8 - top_k = 2 - batch = 4 - - results = {} - - for backend in AVAILABLE_BACKENDS: - torch.manual_seed(42) - block = _make_moe_block(backend, hidden, num_experts, top_k, intermediate) - x = torch.randn(batch, seq_len, hidden, device=DEVICE, dtype=DTYPE, requires_grad=True) - - def run_base(): - out, _ = block(x) - out.sum().backward() - block.zero_grad() - - mem_base = _measure_peak_memory(run_base, warmup=3) - - torch._dynamo.reset() - compiled_block = torch.compile( - block, - fullgraph=False, - backend="inductor", - dynamic=False, - ) - - def run_compiled(): - out, _ = compiled_block(x) - out.sum().backward() - block.zero_grad() - - try: - mem_compiled = _measure_peak_memory(run_compiled, warmup=10) - except Exception as e: - print(f" {backend} inductor compile failed: {type(e).__name__}") - results[backend] = { - "base_mb": mem_base / 1024**2, - "compiled_mb": float("nan"), - "diff_mb": float("nan"), - "ratio": float("nan"), - "compile_backend": "inductor (FAILED)", - } - del compiled_block, x - torch._dynamo.reset() - torch.cuda.empty_cache() - continue - - results[backend] = { - "base_mb": mem_base / 1024**2, - "compiled_mb": mem_compiled / 1024**2, - "diff_mb": (mem_compiled - mem_base) / 1024**2, - "ratio": mem_compiled / mem_base if mem_base > 0 else float("inf"), - "compile_backend": "inductor", - } - - del compiled_block, x - torch._dynamo.reset() - torch.cuda.empty_cache() - - print("\n" + "=" * 95) - print(f" MoE Peak Memory (batch={batch}, seq={seq_len}, hidden={hidden}, E={num_experts}, top_k={top_k})") - print("=" * 95) - print(f" {'Backend':<10} {'Compiler':<10} {'Base (MB)':>10} {'Compiled (MB)':>14} {'Delta (MB)':>12} {'Ratio':>8}") - print("-" * 95) - for backend, r in results.items(): - print( - f" {backend:<10} {r['compile_backend']:<10} {r['base_mb']:>10.1f} " - f"{r['compiled_mb']:>14.1f} {r['diff_mb']:>+12.1f} " - f"{r['ratio']:>7.2f}x" - ) - print("=" * 95) - - -# --------------------------------------------------------------------------- -# Benchmark: Full decoder layer (attention + MoE + norms + residuals) -# --------------------------------------------------------------------------- - - -def _decoder_layer_flops(batch, seq, hidden, intermediate, num_heads, num_kv_heads, num_experts, top_k): - """Estimate FLOPs for one Qwen3MoeDecoderLayer forward pass.""" - tokens = batch * seq - head_dim = hidden // num_heads - attn_proj = 2 * tokens * hidden * (hidden + 2 * (num_kv_heads * head_dim) + hidden) - attn_core = 2 * 2 * batch * num_heads * seq * seq * head_dim - moe = tokens * top_k * 6 * hidden * intermediate - return attn_proj + attn_core + moe - - -def _benchmark_decoder_layer(layer, x, position_ids, position_embeddings, warmup=30, iters=50): - """Benchmark decoder layer fwd+bwd, return median GPU time in seconds.""" - for _ in range(warmup): - outputs = layer( - hidden_states=x, - position_ids=position_ids, - position_embeddings=position_embeddings, - ) - outputs[0].sum().backward() - layer.zero_grad() - if x.grad is not None: - x.grad = None - - torch.cuda.synchronize() - times = [] - for _ in range(iters): - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - start_event.record() - outputs = layer( - hidden_states=x, - position_ids=position_ids, - position_embeddings=position_embeddings, - ) - outputs[0].sum().backward() - end_event.record() - torch.cuda.synchronize() - times.append(start_event.elapsed_time(end_event) / 1000.0) - layer.zero_grad() - if x.grad is not None: - x.grad = None - - times.sort() - return times[len(times) // 2] - - -def _measure_decoder_layer_peak_memory(layer, x, position_ids, position_embeddings, warmup=10): - """Measure peak GPU memory for one decoder layer fwd+bwd.""" - for _ in range(warmup): - outputs = layer( - hidden_states=x, - position_ids=position_ids, - position_embeddings=position_embeddings, - ) - outputs[0].sum().backward() - layer.zero_grad() - if x.grad is not None: - x.grad = None - - torch.cuda.synchronize() - torch.cuda.reset_peak_memory_stats() - torch.cuda.empty_cache() - - outputs = layer( - hidden_states=x, - position_ids=position_ids, - position_embeddings=position_embeddings, - ) - outputs[0].sum().backward() - - torch.cuda.synchronize() - peak = torch.cuda.max_memory_allocated() - layer.zero_grad() - if x.grad is not None: - x.grad = None - return peak - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("seq_len", [1024, 4096]) -def test_decoder_layer_benchmark(seq_len): - """Benchmark full Qwen3MoeDecoderLayer: TFLOPS + peak memory, compiled vs uncompiled.""" - - hidden = 1024 - intermediate = 512 - num_heads = 16 - num_kv_heads = 4 - num_experts = 8 - top_k = 2 - batch = 4 - - flops = _decoder_layer_flops( - batch, - seq_len, - hidden, - intermediate, - num_heads, - num_kv_heads, - num_experts, - top_k, - ) - - results = {} - - for backend in AVAILABLE_BACKENDS: - config = _tiny_moe_config( - hidden_size=hidden, - intermediate_size=hidden * 4, - moe_intermediate_size=intermediate, - num_attention_heads=num_heads, - num_key_value_heads=num_kv_heads, - num_experts=num_experts, - num_experts_per_tok=top_k, - _moe_implementation=backend, - _attn_implementation="sdpa", - ) - - layer = Qwen3MoeDecoderLayer(config, layer_idx=0).to(DEVICE, DTYPE) - x = torch.randn(batch, seq_len, hidden, device=DEVICE, dtype=DTYPE, requires_grad=True) - position_ids = torch.arange(seq_len, device=DEVICE).unsqueeze(0).expand(batch, -1) - position_embeddings = _make_position_embeddings(config, seq_len, DEVICE, DTYPE) - - t_base = _benchmark_decoder_layer(layer, x, position_ids, position_embeddings) - tflops_base = flops / t_base / 1e12 - mem_base = _measure_decoder_layer_peak_memory(layer, x, position_ids, position_embeddings) - - torch._dynamo.reset() - compiled_layer = torch.compile(layer, fullgraph=False, backend="inductor", dynamic=False) - - try: - t_compiled = _benchmark_decoder_layer(compiled_layer, x, position_ids, position_embeddings) - tflops_compiled = flops / t_compiled / 1e12 - mem_compiled = _measure_decoder_layer_peak_memory( - compiled_layer, - x, - position_ids, - position_embeddings, - ) - results[backend] = { - "base_ms": t_base * 1000, - "compiled_ms": t_compiled * 1000, - "base_tflops": tflops_base, - "compiled_tflops": tflops_compiled, - "speedup": t_base / t_compiled, - "base_mb": mem_base / 1024**2, - "compiled_mb": mem_compiled / 1024**2, - "mem_diff_mb": (mem_compiled - mem_base) / 1024**2, - "compile_backend": "inductor", - } - except Exception as e: - print(f" {backend} inductor compile failed: {type(e).__name__}") - results[backend] = { - "base_ms": t_base * 1000, - "compiled_ms": float("nan"), - "base_tflops": tflops_base, - "compiled_tflops": float("nan"), - "speedup": float("nan"), - "base_mb": mem_base / 1024**2, - "compiled_mb": float("nan"), - "mem_diff_mb": float("nan"), - "compile_backend": "inductor (FAILED)", - } - - del compiled_layer, layer, x, position_embeddings - torch._dynamo.reset() - torch.cuda.empty_cache() - - print("\n" + "=" * 110) - print( - f" Decoder Layer Benchmark (batch={batch}, seq={seq_len}, hidden={hidden}, " - f"heads={num_heads}/{num_kv_heads}, E={num_experts}, top_k={top_k})" - ) - print(f" FLOPs per fwd: {flops / 1e9:.1f} GFLOP | dynamic=False, warmup=30, iters=50") - print("=" * 110) - print( - f" {'Backend':<10} {'Compiler':<10} {'Base ms':>8} {'Comp ms':>8} " - f"{'Base TF':>8} {'Comp TF':>8} {'Speed':>6} " - f"{'Base MB':>8} {'Comp MB':>8} {'Mem D':>8}" - ) - print("-" * 110) - for backend, r in results.items(): - print( - f" {backend:<10} {r['compile_backend']:<10} " - f"{r['base_ms']:>8.2f} {r['compiled_ms']:>8.2f} " - f"{r['base_tflops']:>8.2f} {r['compiled_tflops']:>8.2f} " - f"{r['speedup']:>5.2f}x " - f"{r['base_mb']:>8.1f} {r['compiled_mb']:>8.1f} " - f"{r['mem_diff_mb']:>+8.1f}" - ) - print("=" * 110) - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) diff --git a/tests/ops/test_nf4.py b/tests/ops/test_nf4.py deleted file mode 100644 index b914ff9b..00000000 --- a/tests/ops/test_nf4.py +++ /dev/null @@ -1,253 +0,0 @@ -"""Tests for NF4 (NormalFloat4) quantization kernels. - -Validates correctness (roundtrip accuracy, codec, shapes) and measures -effective memory bandwidth of the dequantization kernels. -""" - -import time - -import pytest -import torch - - -try: - from xorl.ops.quantize import nf4_dequantize, nf4_dequantize_gkn, nf4_quantize, nf4_quantize_gkn - from xorl.ops.quantize.nf4_codec import NF4_MIN_STEP, NF4_TABLE - - HAS_NF4 = True -except ImportError: - HAS_NF4 = False - -pytestmark = [ - pytest.mark.gpu, - pytest.mark.skipif(not HAS_NF4, reason="nf4 module not available"), - pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"), -] - - -class TestNF4Codec: - """Tests for NF4 encode/decode correctness.""" - - def test_nf4_table_properties(self): - """NF4 table: 16 values, sorted, symmetric about 0, bounded [-1, 1].""" - assert len(NF4_TABLE) == 16 - assert NF4_TABLE[0] == -1.0 - assert NF4_TABLE[-1] == 1.0 - assert NF4_TABLE[7] == 0.0 - # Sorted - for i in range(15): - assert NF4_TABLE[i] < NF4_TABLE[i + 1] - - def test_roundtrip_exact_table_values(self): - """Quantizing exact NF4 table values should produce exact roundtrip.""" - table = torch.tensor(NF4_TABLE, dtype=torch.float32, device="cuda") - # Create a weight tensor where each row has exactly one table value repeated - M, K = 16, 64 - x = table[:, None].expand(M, K).contiguous() - packed, scales = nf4_quantize(x, group_size=64) - out = nf4_dequantize(packed, scales, M * K, group_size=64).reshape(M, K) - # Each row should dequantize to a constant close to the original table value - for i in range(16): - expected = NF4_TABLE[i] - actual = out[i, 0].item() - assert abs(actual - expected) < 0.05, f"Code {i}: expected {expected:.4f}, got {actual:.4f}" - - -class TestNF4Quantize1D: - """Tests for 1D (flat) NF4 quantization/dequantization.""" - - @pytest.mark.parametrize("group_size", [32, 64, 128]) - def test_shapes_and_dtypes(self, group_size): - """Output shapes and dtypes are correct.""" - M, K = 128, 256 - x = torch.randn(M, K, device="cuda", dtype=torch.float32) - packed, scales = nf4_quantize(x, group_size=group_size) - n = M * K - assert packed.shape == (n // 2,) - assert packed.dtype == torch.uint8 - assert scales.shape == (n // group_size,) - assert scales.dtype == torch.float32 - - @pytest.mark.parametrize("group_size", [32, 64, 128]) - def test_roundtrip_accuracy(self, group_size): - """Roundtrip should produce reasonable reconstruction.""" - M, K = 256, 512 - x = torch.randn(M, K, device="cuda", dtype=torch.float32) - packed, scales = nf4_quantize(x, group_size=group_size) - out = nf4_dequantize(packed, scales, M * K, group_size=group_size) - out = out.reshape(M, K) - # NF4 is 4-bit: expect ~1-5% relative error for normal data - rel_err = (out.float() - x).abs().mean() / x.abs().mean() - assert rel_err < 0.10, f"Relative error {rel_err:.4f} too high for group_size={group_size}" - - def test_zero_input(self): - """Zero input should roundtrip to near-zero.""" - M, K = 64, 128 - x = torch.zeros(M, K, device="cuda", dtype=torch.float32) - packed, scales = nf4_quantize(x, group_size=64) - out = nf4_dequantize(packed, scales, M * K, group_size=64).reshape(M, K) - assert out.abs().max() < 1e-6 - - def test_output_dtype_bf16(self): - """Dequantized output should be bfloat16.""" - M, K = 64, 128 - x = torch.randn(M, K, device="cuda", dtype=torch.float32) - packed, scales = nf4_quantize(x, group_size=64) - out = nf4_dequantize(packed, scales, M * K, group_size=64) - assert out.dtype == torch.bfloat16 - - def test_scales_are_absmax(self): - """Scales should equal per-group absmax.""" - M, K = 4, 128 - gs = 64 - x = torch.randn(M, K, device="cuda", dtype=torch.float32) - packed, scales = nf4_quantize(x, group_size=gs) - # Compute expected absmax per group - x_groups = x.reshape(-1, gs) - expected_scales = x_groups.abs().max(dim=1).values - torch.testing.assert_close(scales, expected_scales, atol=1e-6, rtol=1e-6) - - def test_large_tensor(self): - """Test with realistic model weight dimensions.""" - M, K = 4096, 4096 - x = torch.randn(M, K, device="cuda", dtype=torch.bfloat16).float() - packed, scales = nf4_quantize(x, group_size=64) - out = nf4_dequantize(packed, scales, M * K, group_size=64).reshape(M, K) - rel_err = (out.float() - x).abs().mean() / x.abs().mean() - assert rel_err < 0.10 - - -class TestNF4QuantizeGKN: - """Tests for 2D GKN NF4 quantization/dequantization.""" - - @pytest.mark.parametrize("group_size", [32, 64, 128]) - def test_shapes_and_dtypes(self, group_size): - """Output shapes and dtypes are correct for GKN format.""" - K, N = 256, 128 - x = torch.randn(K, N, device="cuda", dtype=torch.float32) - packed, scales = nf4_quantize_gkn(x, group_size=group_size) - assert packed.shape == (K // 2, N) - assert packed.dtype == torch.uint8 - assert scales.shape == (K // group_size, N) - assert scales.dtype == torch.float32 - - @pytest.mark.parametrize("group_size", [32, 64, 128]) - def test_roundtrip_accuracy(self, group_size): - """GKN roundtrip should produce reasonable reconstruction.""" - K, N = 512, 256 - x = torch.randn(K, N, device="cuda", dtype=torch.float32) - packed, scales = nf4_quantize_gkn(x, group_size=group_size) - out = nf4_dequantize_gkn(packed, scales, K, N, group_size=group_size) - assert out.shape == (K, N) - assert out.dtype == torch.bfloat16 - rel_err = (out.float() - x).abs().mean() / x.abs().mean() - assert rel_err < 0.10, f"Relative error {rel_err:.4f} too high for group_size={group_size}" - - def test_1d_vs_gkn_consistency(self): - """1D and GKN quantization should produce equivalent results for [K, N] input.""" - K, N = 256, 128 - gs = 64 - x = torch.randn(K, N, device="cuda", dtype=torch.float32) - # 1D: flatten and process - packed_1d, scales_1d = nf4_quantize(x, group_size=gs) - out_1d = nf4_dequantize(packed_1d, scales_1d, K * N, group_size=gs).reshape(K, N) - # GKN: process as 2D - packed_gkn, scales_gkn = nf4_quantize_gkn(x, group_size=gs) - out_gkn = nf4_dequantize_gkn(packed_gkn, scales_gkn, K, N, group_size=gs) - # Both should have similar reconstruction error (not identical due to grouping direction) - err_1d = (out_1d.float() - x).abs().mean() - err_gkn = (out_gkn.float() - x).abs().mean() - # Both should have reasonable error - assert err_1d < 0.1 * x.abs().mean() - assert err_gkn < 0.1 * x.abs().mean() - - def test_large_tensor(self): - """Test with MoE expert weight dimensions.""" - K, N = 2048, 7168 # Typical MoE hidden -> intermediate - x = torch.randn(K, N, device="cuda", dtype=torch.bfloat16).float() - packed, scales = nf4_quantize_gkn(x, group_size=64) - out = nf4_dequantize_gkn(packed, scales, K, N, group_size=64) - rel_err = (out.float() - x).abs().mean() / x.abs().mean() - assert rel_err < 0.10 - - -class TestNF4Bandwidth: - """Measure effective memory bandwidth of NF4 dequantization kernels.""" - - def _measure_bandwidth(self, fn, packed, scales, num_elements, group_size, num_iters=200): - """Measure effective bandwidth of a dequant kernel.""" - # Warmup - for _ in range(20): - fn(packed, scales, num_elements, group_size) - torch.cuda.synchronize() - - start = time.perf_counter() - for _ in range(num_iters): - fn(packed, scales, num_elements, group_size) - torch.cuda.synchronize() - elapsed = time.perf_counter() - start - - # Total bytes: packed (read) + scales (read) + output bf16 (write) - packed_bytes = num_elements // 2 - scale_bytes = (num_elements // group_size) * 4 - output_bytes = num_elements * 2 - total_bytes = packed_bytes + scale_bytes + output_bytes - bw = total_bytes * num_iters / elapsed / 1e9 # GB/s - return bw - - def _measure_bandwidth_gkn(self, packed, scales, K, N, group_size, num_iters=200): - """Measure effective bandwidth of GKN dequant kernel.""" - for _ in range(20): - nf4_dequantize_gkn(packed, scales, K, N, group_size) - torch.cuda.synchronize() - - start = time.perf_counter() - for _ in range(num_iters): - nf4_dequantize_gkn(packed, scales, K, N, group_size) - torch.cuda.synchronize() - elapsed = time.perf_counter() - start - - num_elements = K * N - packed_bytes = num_elements // 2 - scale_bytes = (num_elements // group_size) * 4 - output_bytes = num_elements * 2 - total_bytes = packed_bytes + scale_bytes + output_bytes - bw = total_bytes * num_iters / elapsed / 1e9 - return bw - - @pytest.mark.parametrize( - "size", - [ - (4096, 4096), # ~16M elements - (8192, 8192), # ~67M elements - ], - ) - def test_dequant_1d_bandwidth(self, size): - """1D dequant should achieve >2000 GB/s on large tensors.""" - M, K = size - gs = 64 - x = torch.randn(M, K, device="cuda", dtype=torch.float32) - packed, scales = nf4_quantize(x, group_size=gs) - bw = self._measure_bandwidth(nf4_dequantize, packed, scales, M * K, gs) - print(f"\n[NF4 1D dequant] {M}x{K} gs={gs}: {bw:.0f} GB/s") - # Soft assertion: warn if below target - if bw < 2000: - pytest.skip(f"Bandwidth {bw:.0f} GB/s below 2000 GB/s target (may vary by GPU)") - - @pytest.mark.parametrize( - "size", - [ - (4096, 4096), - (2048, 7168), # MoE expert dimensions - ], - ) - def test_dequant_gkn_bandwidth(self, size): - """GKN dequant should achieve >2000 GB/s on large tensors.""" - K, N = size - gs = 64 - x = torch.randn(K, N, device="cuda", dtype=torch.float32) - packed, scales = nf4_quantize_gkn(x, group_size=gs) - bw = self._measure_bandwidth_gkn(packed, scales, K, N, gs) - print(f"\n[NF4 GKN dequant] {K}x{N} gs={gs}: {bw:.0f} GB/s") - if bw < 2000: - pytest.skip(f"Bandwidth {bw:.0f} GB/s below 2000 GB/s target (may vary by GPU)") diff --git a/tests/ops/test_nvfp4_fake_quant.py b/tests/ops/test_nvfp4_fake_quant.py index 50f3ed32..1105fd46 100644 --- a/tests/ops/test_nvfp4_fake_quant.py +++ b/tests/ops/test_nvfp4_fake_quant.py @@ -1,8 +1,8 @@ """Tests for the NVFP4 STE fake-quant op (weight-only + 3D MoE experts). -Validates the pure-PyTorch NVFP4 round-to-nearest forward (matches an independent -reference; values lie on the E2M1 grid) and the straight-through-identity backward, -including through an ``F.linear`` and the 3D expert helpers. Pure PyTorch — runs on CPU. +Validates the pure-PyTorch NVFP4 round-to-nearest forward against an independent +reference and the straight-through backward through ``F.linear`` and the 3D +expert helpers. Pure PyTorch — runs on CPU. """ import pytest @@ -15,9 +15,8 @@ FP8_E4M3_MAX, _fake_quantize_3d_experts, _fake_quantize_3d_fused_gate_up, - fake_quantize, + fake_quantize_activation_nvfp4, fake_quantize_nvfp4, - is_supported_format, ) @@ -44,62 +43,16 @@ def _ref_fake_quant(w: torch.Tensor, block_size: int = 16) -> torch.Tensor: class TestForward: - @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) - @pytest.mark.parametrize("shape", [(256, 256), (128, 512), (320, 64)]) - def test_forward_matches_reference(self, dtype, shape): - torch.manual_seed(0) - w = torch.randn(*shape, dtype=dtype) - out = fake_quantize_nvfp4(w, block_size=16) - ref = _ref_fake_quant(w, block_size=16) - assert out.shape == w.shape - assert out.dtype == w.dtype - torch.testing.assert_close(out, ref, rtol=0, atol=0) - - def test_values_lie_on_scaled_grid(self): - """Each dequantized element is a grid value times its block's scale.""" - torch.manual_seed(0) - w = torch.randn(64, 16, dtype=torch.float32) - out = fake_quantize_nvfp4(w, block_size=16) - grid = torch.tensor(_E2M1_ABS) - allowed = torch.cat([grid, -grid]).unique() - for r in range(out.shape[0]): - row = out[r] - scale = row.abs().max() / FP4_E2M1_MAX - if scale == 0: - continue - codes = row / scale - nearest = torch.min((codes[:, None] - allowed[None, :]).abs(), dim=1).values - assert nearest.max() < 1e-2, "dequantized values must lie on the E2M1 grid" - - def test_roundtrip_error_small(self): - """FP4 fake-quant of a Gaussian weight should be a low-error approximation.""" - torch.manual_seed(0) - w = torch.randn(256, 256, dtype=torch.float32) - out = fake_quantize_nvfp4(w, block_size=16) - rel_err = (out - w).norm() / w.norm() - assert rel_err < 0.15, f"relative fake-quant error too high: {rel_err:.4f}" - - def test_is_lossy(self): - """Fake quant must actually change the weights (not a no-op identity).""" - torch.manual_seed(0) - w = torch.randn(128, 128, dtype=torch.float32) - out = fake_quantize_nvfp4(w) - assert not torch.equal(out, w) - + def test_2d_quantization_reference_dispatch_and_ste(self): + for dtype in (torch.bfloat16, torch.float32): + torch.manual_seed(0) + w = torch.randn(64, 32, dtype=dtype) + out = fake_quantize_nvfp4(w, block_size=16) + ref = _ref_fake_quant(w, block_size=16) + assert out.shape == w.shape + assert out.dtype == w.dtype + torch.testing.assert_close(out, ref, rtol=0, atol=0) -class TestBackwardSTE: - def test_identity_gradient(self): - """Backward is the straight-through identity: w.grad == grad_out.""" - torch.manual_seed(0) - w = torch.randn(256, 256, dtype=torch.float32, requires_grad=True) - out = fake_quantize_nvfp4(w, block_size=16) - grad_out = torch.randn_like(out) - out.backward(grad_out) - assert w.grad is not None - torch.testing.assert_close(w.grad, grad_out, rtol=0, atol=0) - - def test_gradient_through_linear(self): - """STE flows the full upstream gradient through an F.linear weight.""" torch.manual_seed(0) x = torch.randn(8, 256, dtype=torch.float32) w = torch.randn(128, 256, dtype=torch.float32, requires_grad=True) @@ -110,40 +63,40 @@ def test_gradient_through_linear(self): assert w.grad is not None torch.testing.assert_close(w.grad, expected, rtol=1e-4, atol=1e-4) + _assert_activation_wrapper_preserves_leading_dimensions() + TestContract()._assert_input_shape_admission() + TestMoEExperts3D()._assert_projection_shapes_and_ste() + + +def _assert_activation_wrapper_preserves_leading_dimensions(): + torch.manual_seed(2) + x = torch.randn(2, 3, 32, requires_grad=True) + upstream = torch.randn_like(x) + + actual = fake_quantize_activation_nvfp4(x, block_size=16) + flattened = fake_quantize_nvfp4(x.detach().reshape(-1, 32), block_size=16).reshape_as(x) + + assert actual.shape == x.shape + torch.testing.assert_close(actual, flattened, rtol=0, atol=0) + (actual * upstream).sum().backward() + torch.testing.assert_close(x.grad, upstream, rtol=0, atol=0) + class TestContract: - def test_rejects_non_2d(self): + def _assert_input_shape_admission(self): w = torch.randn(4, 16, 16, dtype=torch.float32) with pytest.raises(AssertionError): fake_quantize_nvfp4(w) - def test_rejects_not_divisible_by_block(self): - w = torch.randn(17, 17, dtype=torch.float32) # 289 not divisible by 16 - with pytest.raises(AssertionError): + w = torch.randn(16, 24, dtype=torch.float32) + with pytest.raises(AssertionError, match="in_features"): fake_quantize_nvfp4(w, block_size=16) -class TestRegistry: - def test_nvfp4_supported(self): - assert is_supported_format("nvfp4") - assert not is_supported_format("int4") - - def test_dispatch_matches_direct(self): - torch.manual_seed(0) - w = torch.randn(128, 128, dtype=torch.bfloat16) - torch.testing.assert_close(fake_quantize(w, "nvfp4", 16), fake_quantize_nvfp4(w, 16), rtol=0, atol=0) - - def test_unsupported_format_raises(self): - w = torch.randn(128, 128, dtype=torch.float32) - with pytest.raises(ValueError): - fake_quantize(w, "int4", 16) - - class TestMoEExperts3D: """The 3D GKN helpers used by the MoE expert fake-quant wrap.""" - def test_down_proj_ste_and_shape(self): - """down_proj [E, K=I, N=H]: STE identity backward, shape preserved, lossy.""" + def _assert_projection_shapes_and_ste(self): torch.manual_seed(0) E, I, H = 4, 64, 128 w = torch.randn(E, I, H, dtype=torch.float32, requires_grad=True) @@ -154,7 +107,18 @@ def test_down_proj_ste_and_shape(self): w_fq.backward(g) torch.testing.assert_close(w.grad, g, rtol=0, atol=0) - def test_experts_quantized_independently(self): + w = torch.randn(E, H, 2 * I, dtype=torch.float32, requires_grad=True) + w_fq = _fake_quantize_3d_fused_gate_up(w, intermediate_size=I, block_size=16) + assert w_fq.shape == w.shape + assert not torch.equal(w_fq.detach(), w.detach()) + g = torch.randn_like(w_fq) + w_fq.backward(g) + torch.testing.assert_close(w.grad, g, rtol=0, atol=0) + + self._assert_experts_quantized_independently() + self._assert_fused_gate_up_uses_per_half_global_scale() + + def _assert_experts_quantized_independently(self): """Scaling one expert's weights must not change another expert's dequant.""" torch.manual_seed(0) E, I, H = 3, 32, 64 @@ -166,29 +130,7 @@ def test_experts_quantized_independently(self): torch.testing.assert_close(out2[1:], base[1:], rtol=0, atol=0) assert not torch.equal(out2[0], base[0]) - def test_fused_gate_up_ste_and_shape(self): - """gate_up_proj [E, H, 2I]: STE identity backward, shape preserved, lossy.""" - torch.manual_seed(0) - E, H, I = 4, 128, 64 - w = torch.randn(E, H, 2 * I, dtype=torch.float32, requires_grad=True) - w_fq = _fake_quantize_3d_fused_gate_up(w, intermediate_size=I, block_size=16) - assert w_fq.shape == w.shape - assert not torch.equal(w_fq.detach(), w.detach()) - g = torch.randn_like(w_fq) - w_fq.backward(g) - torch.testing.assert_close(w.grad, g, rtol=0, atol=0) - - def test_fused_gate_up_metadata_shapes(self): - """return_metadata exposes GKN-layout block_scales + per-(expert,half) scale.""" - torch.manual_seed(0) - E, H, I, bs = 2, 64, 48, 16 - w = torch.randn(E, H, 2 * I, dtype=torch.float32) - _, meta = _fake_quantize_3d_fused_gate_up(w, intermediate_size=I, block_size=bs, return_metadata=True) - assert meta["weight_scale_2"].shape == (E, 2) - assert meta["block_scales"].shape == (E, 2, H // bs, I) - assert meta["codes"].shape == (E, 2, I, H) - - def test_fused_gate_up_uses_per_half_global_scale(self): + def _assert_fused_gate_up_uses_per_half_global_scale(self): """gate and up of an expert are quantized with INDEPENDENT global scales. Per-half weight_scale_2 (REVERT of PR #399's shared scale): the strictly-matched @@ -211,14 +153,3 @@ def test_fused_gate_up_uses_per_half_global_scale(self): torch.testing.assert_close(ws2[:, 1], up_expected, rtol=0, atol=0) # ... and they are NOT shared: up (×100) >> gate (×0.01). assert (ws2[:, 1] > ws2[:, 0] * 10).all(), "per-half scales must differ for differently-scaled halves" - - def test_2d_requires_in_features_divisible_by_block(self): - """K (in_features) not divisible by block_size must fail loud, not silently - group elements across output rows (M*K-divisible-but-K-not is the trap).""" - # M*K = 16*24 = 384 is divisible by 16, but K=24 is NOT -> must raise. - w = torch.randn(16, 24, dtype=torch.float32) - with pytest.raises(AssertionError, match="in_features"): - fake_quantize_nvfp4(w, block_size=16) - # K divisible by block_size is fine. - ok = fake_quantize_nvfp4(torch.randn(16, 32), block_size=16) - assert ok.shape == (16, 32) diff --git a/tests/ops/test_prequant_gnk_to_gkn.py b/tests/ops/test_prequant_gnk_to_gkn.py deleted file mode 100644 index f36782ba..00000000 --- a/tests/ops/test_prequant_gnk_to_gkn.py +++ /dev/null @@ -1,188 +0,0 @@ -"""Tests for pre-quantized GNK -> GKN weight loading correctness. - -HuggingFace/modelopt checkpoints store expert weights in GNK format: - [N, K] per expert (N=out_features, K=in_features) - -Our internal format is GKN: - [K, N] per expert (K=in_features, N=out_features) - -The loading code transposes quantized data: fp8/packed.T and scales.T. -This only works if 2D block quantization is transposition-equivariant: - quant(W.T).T == quant(W) - -These tests verify that property for both block_fp8 and nvfp4 formats. -""" - -import pytest -import torch -import triton - -from xorl.ops.quantize import ( - block_fp8_dequantize_gkn, - block_fp8_quantize_gkn, - nvfp4_quantize, -) -from xorl.ops.quantize.nvfp4_gkn_quantize import nvfp4_dequantize_gkn, nvfp4_quantize_gkn - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") - - -# ========================================================================= -# Block FP8 -- GNK -> GKN -# ========================================================================= - - -class TestBlockFP8PrequantGNKtoGKN: - """Verify block_fp8 quantized data can be correctly transposed from GNK to GKN.""" - - def test_block_fp8_transpose_roundtrip_and_stacking(self): - """Transpose matches direct, roundtrip accuracy, path equivalence, non-square shapes, multi-expert stacking.""" - K, N = 512, 256 - W_gkn = torch.randn(K, N, device="cuda", dtype=torch.float32) - - # --- Transpose matches direct quantization --- - fp8_direct, scales_direct = block_fp8_quantize_gkn(W_gkn) - W_gnk = W_gkn.T.contiguous() - fp8_gnk, scales_gnk = block_fp8_quantize_gkn(W_gnk) - fp8_transposed = fp8_gnk.T.contiguous() - scales_transposed = scales_gnk.T.contiguous() - assert torch.equal(fp8_direct, fp8_transposed), "FP8 values differ after transpose" - assert torch.equal(scales_direct, scales_transposed), "Scales differ after transpose" - - # --- GNK->GKN roundtrip accuracy --- - fp8_loaded = fp8_gnk.T.contiguous() - scales_loaded = scales_gnk.float().T.contiguous() - W_recovered = block_fp8_dequantize_gkn(fp8_loaded, scales_loaded) - rel_err = (W_gkn - W_recovered).abs().mean() / W_gkn.abs().mean() - assert rel_err < 0.03, f"Roundtrip rel error {rel_err:.4f} > 0.03" - - # --- Direct path == transposed path dequantization --- - W_direct = block_fp8_dequantize_gkn(fp8_direct, scales_direct) - fp8_gnk2, scales_gnk2 = block_fp8_quantize_gkn(W_gkn.T.contiguous()) - W_trans = block_fp8_dequantize_gkn(fp8_gnk2.T.contiguous(), scales_gnk2.T.contiguous()) - assert torch.equal(W_direct, W_trans), "Direct and transposed paths diverge" - - # --- Non-square shapes --- - for K_ns, N_ns in [(256, 512), (512, 256), (384, 256)]: - W_ns = torch.randn(K_ns, N_ns, device="cuda", dtype=torch.float32) - fp8_ns, s_ns = block_fp8_quantize_gkn(W_ns.T.contiguous()) - fp8_l = fp8_ns.T.contiguous() - s_l = s_ns.T.contiguous() - assert fp8_l.shape == (K_ns, N_ns) - assert s_l.shape == (triton.cdiv(K_ns, 128), triton.cdiv(N_ns, 128)) - W_rec = block_fp8_dequantize_gkn(fp8_l, s_l) - re = (W_ns - W_rec).abs().mean() / W_ns.abs().mean() - assert re < 0.03, f"Shape ({K_ns},{N_ns}): rel error {re:.4f}" - - # --- Multiple experts stacked --- - G = 4 - experts_gkn = torch.randn(G, K, N, device="cuda", dtype=torch.float32) - fp8_list, scales_list = [], [] - for i in range(G): - w_gnk = experts_gkn[i].T.contiguous() - fp8_e, s_e = block_fp8_quantize_gkn(w_gnk) - fp8_list.append(fp8_e.T.contiguous()) - scales_list.append(s_e.T.contiguous()) - fp8_stacked = torch.stack(fp8_list) - scales_stacked = torch.stack(scales_list) - assert fp8_stacked.shape == (G, K, N) - assert scales_stacked.shape == (G, triton.cdiv(K, 128), triton.cdiv(N, 128)) - for i in range(G): - W_rec = block_fp8_dequantize_gkn(fp8_stacked[i], scales_stacked[i]) - re = (experts_gkn[i] - W_rec).abs().mean() / experts_gkn[i].abs().mean() - assert re < 0.03, f"Expert {i}: rel error {re:.4f}" - - -# ========================================================================= -# NVFP4 -- GNK -> GKN -# ========================================================================= - - -class TestNVFP4PrequantGNKtoGKN: - """Verify nvfp4 quantized data can be correctly transposed from GNK to GKN.""" - - def test_nvfp4_transpose_roundtrip_and_stacking(self): - """Roundtrip, direct match, dequant path agreement, non-square, multi-expert stacking, global scale absorption.""" - K, N = 256, 256 - block_size = 16 - W_gkn = torch.randn(K, N, device="cuda", dtype=torch.float32) - W_gnk = W_gkn.T.contiguous() - - # --- GNK -> GKN roundtrip --- - packed_flat, scales_flat, global_scale = nvfp4_quantize(W_gnk, block_size) - packed_gnk = packed_flat.reshape(N, K // 2) - scales_gnk = scales_flat.reshape(N, K // block_size) - packed_gkn = packed_gnk.T.contiguous() - scales_gkn = (scales_gnk.float() * global_scale.float()).T.contiguous() - gs_loaded = torch.ones(1, dtype=torch.float32, device=W_gkn.device) - W_recovered = nvfp4_dequantize_gkn(packed_gkn, scales_gkn, gs_loaded, K, N, block_size) - rel_err = (W_gkn - W_recovered.float()).abs().mean() / W_gkn.abs().mean() - assert rel_err < 0.15, f"Roundtrip rel error {rel_err:.4f} > 0.15" - - # --- Transposed matches direct quantization --- - packed_direct, scales_direct, gs_direct = nvfp4_quantize_gkn(W_gkn, block_size) - packed_flat2, _, _ = nvfp4_quantize(W_gnk, block_size) - packed_gnk2 = packed_flat2.reshape(N, K // 2) - packed_transposed = packed_gnk2.T.contiguous() - assert torch.equal(packed_direct, packed_transposed), "Packed data differs" - - # --- Transposed dequant matches direct dequant --- - W_a = nvfp4_dequantize_gkn(packed_direct, scales_direct, gs_direct, K, N, block_size) - - packed_flat_b, scales_flat_b, gs_b = nvfp4_quantize(W_gnk, block_size) - packed_gnk_b = packed_flat_b.reshape(N, K // 2) - scales_gnk_b = scales_flat_b.reshape(N, K // block_size) - packed_gkn_b = packed_gnk_b.T.contiguous() - scales_gkn_b = (scales_gnk_b.float() * gs_b.float()).T.contiguous() - gs_one = torch.ones(1, dtype=torch.float32, device=W_gkn.device) - W_b = nvfp4_dequantize_gkn(packed_gkn_b, scales_gkn_b, gs_one, K, N, block_size) - max_diff = (W_a.float() - W_b.float()).abs().max().item() - assert max_diff < 1e-3, f"Max diff between paths: {max_diff}" - - # --- Non-square shapes --- - for K_ns, N_ns in [(256, 512), (512, 256)]: - W_ns = torch.randn(K_ns, N_ns, device="cuda", dtype=torch.float32) - W_ns_gnk = W_ns.T.contiguous() - pf, sf, gs = nvfp4_quantize(W_ns_gnk, block_size) - p_gnk = pf.reshape(N_ns, K_ns // 2) - s_gnk = sf.reshape(N_ns, K_ns // block_size) - p_gkn = p_gnk.T.contiguous() - s_gkn = (s_gnk.float() * gs.float()).T.contiguous() - gs_l = torch.ones(1, dtype=torch.float32, device=W_ns.device) - assert p_gkn.shape == (K_ns // 2, N_ns) - assert s_gkn.shape == (K_ns // block_size, N_ns) - W_rec = nvfp4_dequantize_gkn(p_gkn, s_gkn, gs_l, K_ns, N_ns, block_size) - re = (W_ns - W_rec.float()).abs().mean() / W_ns.abs().mean() - assert re < 0.15, f"Shape ({K_ns},{N_ns}): rel error {re:.4f}" - - # --- Multiple experts stacked --- - G, K_e, N_e = 4, 256, 512 - experts_gkn = torch.randn(G, K_e, N_e, device="cuda", dtype=torch.float32) - packed_list, scales_list = [], [] - for i in range(G): - w_gnk = experts_gkn[i].T.contiguous() - pf, sf, gs = nvfp4_quantize(w_gnk, block_size) - p_gnk = pf.reshape(N_e, K_e // 2) - s_gnk = sf.reshape(N_e, K_e // block_size) - packed_list.append(p_gnk.T.contiguous()) - scales_list.append((s_gnk.float() * gs.float()).T.contiguous()) - packed_stacked = torch.stack(packed_list) - scales_stacked = torch.stack(scales_list) - assert packed_stacked.shape == (G, K_e // 2, N_e) - assert scales_stacked.shape == (G, K_e // block_size, N_e) - gs_l = torch.ones(1, dtype=torch.float32, device=experts_gkn.device) - for i in range(G): - W_rec = nvfp4_dequantize_gkn(packed_stacked[i], scales_stacked[i], gs_l, K_e, N_e, block_size) - re = (experts_gkn[i] - W_rec.float()).abs().mean() / experts_gkn[i].abs().mean() - assert re < 0.15, f"Expert {i}: rel error {re:.4f}" - - # --- Global scale absorption preserves precision --- - packed_abs, scales_fp8_abs, gs_abs = nvfp4_quantize_gkn(W_gkn, block_size) - W_non_absorbed = nvfp4_dequantize_gkn(packed_abs, scales_fp8_abs, gs_abs, K, N, block_size) - absorbed_scales = scales_fp8_abs.float() * gs_abs.float() - gs_one_abs = torch.ones(1, dtype=torch.float32, device=W_gkn.device) - W_absorbed = nvfp4_dequantize_gkn(packed_abs, absorbed_scales, gs_one_abs, K, N, block_size) - assert torch.equal(W_non_absorbed, W_absorbed), ( - f"Max diff: {(W_non_absorbed.float() - W_absorbed.float()).abs().max().item()}" - ) diff --git a/tests/ops/test_quack_ep_parity.py b/tests/ops/test_quack_ep_parity.py index 9a095351..7f5ce360 100644 --- a/tests/ops/test_quack_ep_parity.py +++ b/tests/ops/test_quack_ep_parity.py @@ -12,8 +12,6 @@ output-and-gradient parity before being trusted in a training config. """ -from types import SimpleNamespace - import pytest import torch import torch.nn.functional as F @@ -25,29 +23,11 @@ requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.cpu -def test_deepep_none_gradient_paths_match_forward_arity(): - try: - from xorl.ops.moe.quack import ( # noqa: PLC0415 - QuackEPDeepEPCombine, - QuackEPDeepEPNoPermute, - ) - except Exception as exc: # noqa: BLE001 - pytest.skip(f"cannot import DeepEP Quack ops: {exc}") - - ctx = SimpleNamespace(needs_input_grad=(False,) * 16) - assert QuackEPDeepEPCombine.backward(ctx, None) == (None,) * 16 - assert QuackEPDeepEPNoPermute.backward(ctx, None) == (None,) * 16 - - def _imports(): - try: - from xorl.ops.moe.quack import QuackEPGroupGemm # noqa: PLC0415 - from xorl.ops.moe.triton import TritonEPGroupGemm # noqa: PLC0415 + from xorl.ops.moe.quack import QuackEPGroupGemm # noqa: PLC0415 + from xorl.ops.moe.triton import TritonEPGroupGemm # noqa: PLC0415 - return QuackEPGroupGemm, TritonEPGroupGemm - except Exception as exc: # noqa: BLE001 - pytest.skip(f"cannot import EP group GEMM ops: {exc}") + return QuackEPGroupGemm, TritonEPGroupGemm def _cos(a: torch.Tensor, b: torch.Tensor) -> float: @@ -56,46 +36,46 @@ def _cos(a: torch.Tensor, b: torch.Tensor) -> float: @requires_gpu @pytest.mark.gpu -@pytest.mark.parametrize( - "tokens_per_expert", - [ - [13, 0, 7, 21], # uneven with an empty expert - [64, 64, 64, 64], # balanced - [1, 127, 0, 32], # extreme skew + empty expert - ], -) -@pytest.mark.parametrize("with_scores", [False, True]) -def test_quack_ep_group_gemm_matches_triton(tokens_per_expert, with_scores): +def test_quack_ep_group_gemm_and_halfconcat_forward_policy(): QuackEPGroupGemm, TritonEPGroupGemm = _imports() - torch.manual_seed(0) - cumsum = torch.tensor(tokens_per_expert, device="cuda").cumsum(0) - M = int(cumsum[-1]) - x = torch.randn(M, H, device="cuda", dtype=DTYPE) - gate_up = torch.randn(E, H, 2 * I, device="cuda", dtype=DTYPE) * 0.02 - down = torch.randn(E, I, H, device="cuda", dtype=DTYPE) * 0.02 - scores = torch.rand(M, device="cuda", dtype=torch.float32) if with_scores else None - - def run(cls): - x_ = x.clone().requires_grad_(True) - g = gate_up.clone().requires_grad_(True) - d = down.clone().requires_grad_(True) - out = cls.apply(x_, cumsum, g, d, I, scores) - # Non-trivial upstream gradient; also exercises the backward grad-arity - # contract (a missing grad raises "returned an incorrect number of - # gradients" here). - out.float().pow(2).sum().backward() - return out.detach(), x_.grad, g.grad, d.grad - - quack_tensors = run(QuackEPGroupGemm) - triton_tensors = run(TritonEPGroupGemm) - for name, q, t in zip(("out", "grad_x", "grad_gate_up", "grad_down"), quack_tensors, triton_tensors): - cos = _cos(q, t) - assert cos > 0.999, f"{name}: quack/triton cosine {cos:.6f} (tokens_per_expert={tokens_per_expert})" - - -@requires_gpu -@pytest.mark.gpu -def test_quack_ep_forward_matches_halfconcat_reference(): + cases = ( + ([13, 0, 7, 21], False), # uneven with an empty expert + ([64, 64, 64, 64], True), # balanced with score scaling + ([1, 127, 0, 32], True), # extreme skew + empty expert with score scaling + ) + for tokens_per_expert, with_scores in cases: + torch.manual_seed(0) + cumsum = torch.tensor(tokens_per_expert, device="cuda").cumsum(0) + M = int(cumsum[-1]) + x = torch.randn(M, H, device="cuda", dtype=DTYPE) + gate_up = torch.randn(E, H, 2 * I, device="cuda", dtype=DTYPE) * 0.02 + down = torch.randn(E, I, H, device="cuda", dtype=DTYPE) * 0.02 + scores = torch.rand(M, device="cuda", dtype=torch.float32) if with_scores else None + + def run(cls): + x_ = x.clone().requires_grad_(True) + g = gate_up.clone().requires_grad_(True) + d = down.clone().requires_grad_(True) + out = cls.apply(x_, cumsum, g, d, I, scores) + # Non-trivial upstream gradient; also exercises the backward grad-arity + # contract (a missing grad raises "returned an incorrect number of + # gradients" here). + out.float().pow(2).sum().backward() + return out.detach(), x_.grad, g.grad, d.grad + + quack_tensors = run(QuackEPGroupGemm) + triton_tensors = run(TritonEPGroupGemm) + for name, q, t in zip(("out", "grad_x", "grad_gate_up", "grad_down"), quack_tensors, triton_tensors): + cos = _cos(q, t) + assert cos > 0.999, ( + f"{name}: quack/triton cosine {cos:.6f} " + f"(tokens_per_expert={tokens_per_expert}, with_scores={with_scores})" + ) + + _assert_quack_ep_forward_matches_halfconcat_reference() + + +def _assert_quack_ep_forward_matches_halfconcat_reference(): """Pin the gate/up convention itself: silu(h[:, :I]) * h[:, I:] (half-concat).""" QuackEPGroupGemm, _ = _imports() torch.manual_seed(1) diff --git a/tests/ops/test_quack_process_safety.py b/tests/ops/test_quack_process_safety.py index 4fcfc802..95cde304 100644 --- a/tests/ops/test_quack_process_safety.py +++ b/tests/ops/test_quack_process_safety.py @@ -1,64 +1,39 @@ -import importlib.util import io import os import struct -import sys import time import types from pathlib import Path import pytest +from xorl.ops.quack import _worker_protocol as worker_protocol +from xorl.ops.quack import cache_utils, cute_dsl_ptxas -_QUACK_DIR = Path(__file__).parents[2] / "src" / "xorl" / "ops" / "quack" - -def _load_module(name: str, path: Path, monkeypatch: pytest.MonkeyPatch | None = None): - if monkeypatch is not None: - monkeypatch.setitem(sys.modules, "cutlass", types.ModuleType("cutlass")) - spec = importlib.util.spec_from_file_location(name, path) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -def _load_cache_utils(monkeypatch: pytest.MonkeyPatch): - cutlass = types.ModuleType("cutlass") - cutlass.__version__ = "test" - cute = types.ModuleType("cutlass.cute") - cute.runtime = types.SimpleNamespace() - cutlass.cute = cute - tvm_ffi = types.ModuleType("tvm_ffi") - tvm_ffi.__version__ = "test" - monkeypatch.setitem(sys.modules, "cutlass", cutlass) - monkeypatch.setitem(sys.modules, "cutlass.cute", cute) - monkeypatch.setitem(sys.modules, "tvm_ffi", tvm_ffi) - return _load_module("quack_cache_utils_test", _QUACK_DIR / "cache_utils.py") - - -def test_worker_protocol_times_out_on_silent_worker(): - protocol = _load_module("quack_worker_protocol_test", _QUACK_DIR / "_worker_protocol.py") +def test_quack_process_and_cache_safety_policy(tmp_path, monkeypatch): read_fd, write_fd = os.pipe() started = time.monotonic() try: with os.fdopen(read_fd, "rb", buffering=0) as stream: with pytest.raises(TimeoutError, match="did not respond"): - protocol.recv_message(stream, timeout_s=0.05) + worker_protocol.recv_message(stream, timeout_s=0.05) finally: os.close(write_fd) assert time.monotonic() - started < 1 + _assert_worker_protocol_rejects_truncated_body() + _assert_ptxas_uses_unique_outputs_and_a_timeout(tmp_path, monkeypatch) + _assert_cache_key_hash_policy() -def test_worker_protocol_rejects_truncated_body(): - protocol = _load_module("quack_worker_protocol_truncated_test", _QUACK_DIR / "_worker_protocol.py") + +def _assert_worker_protocol_rejects_truncated_body(): stream = io.BytesIO(struct.pack("