Add fail-closed PP byte-boundary contract for exact lanes - #52
Add fail-closed PP byte-boundary contract for exact lanes#52kiddyboots216 wants to merge 3 commits into
Conversation
Broly Security ScanNote ✅ Clean scan Note Re-scan this PR anytime with
|
9f74cd1 to
459f8a5
Compare
ecbab1e to
5da2e10
Compare
459f8a5 to
51f9665
Compare
5da2e10 to
28254bd
Compare
51f9665 to
4d2688b
Compare
Exact value programs engage a byte-boundary admission when split for pipeline parallelism: pipeline_module_split validates that cuts land only on decoder-layer boundaries (the trainer's materialized BF16 residual stream — a natural rounding boundary of the serving program), that pruning preserved global layer identity so layer-position-keyed kernel selection cannot flip, that embed/final-norm/head placement matches the uncut program, and that the model family is certified (Qwen3.5 dense; the exact Qwen3.6-35B-A3B MoE and GLM-5.2 programs RAISE). Admission validates resolved reality, not configuration metadata: the model must declare bfloat16 (the Qwen3.5 config adapter propagates the checkpoint's declaration) AND its per-stage parameters must be bf16 — fp32 exceptions only from an explicit pin allowlist (the exact GDN's A_log/dt_bias) or uniformly-fp32 masters under a declared bf16 mixed-precision compute intent threaded from _parallelize. As ground truth, _pp_forward asserts the actual wire dtype (received inter-stage tensor and emitted hidden must be bf16) on every scheduled forward. Marked stage forwards fail closed on per-microbatch metadata: the complete varlen set (position_ids, cu_seq_lens_q/k, max_length_q/k) is required and VALUE-validated (int32, strictly increasing, endpoints [0, total_tokens], q/k span equality, max_length >= longest span) — a missing or malformed cu_seq_lens would otherwise silently merge packed documents into one attention span. make_pp_loss_fn refuses ce_mode='bi_fused' explicitly instead of substituting a different lm-head program under PP (PP for exact lanes is a forward/scoring envelope until a bi_fused PP loss program exists). Every engagement logs per stage; validation is two-pass so a failure cannot leave earlier stages marked.
Add conventional CPU admission tests and two-GPU PP1-versus-PP2 byte gates. Unsupported-family errors and test documentation describe the enforced invariants directly.
exact_contract_family() now reads the config._exact_contract_family stamp written at model resolution; a present stamp is authoritative (including a stamped None on generic models). Configs that predate stamping fall back to the shared resolve_exact_contract_family() legacy-flag resolver, so direct-construction paths classify exactly as before. The local Qwen model-type set moves out of this module; adding a family is now a stamp/registry entry plus byte evidence rather than surgery on the classifier.
qywu
left a comment
qywu
left a comment
There was a problem hiding this comment.
The PP byte-boundary contract looks solid, and the SwiGLU scoping in this branch is the correct version — fused_silu_and_mul stays byte-identical to base (verified on H100: 0.00% of elements differ) while exact_fp32_silu_and_mul is gated behind _qwen35_exact_contract. That's the design #51's body describes but doesn't implement.
The blocker is purely structural.
This branch is not stacked on #51
git merge-base fix/one-round-fp32-swiglu-trainer pr/pp-byte-alignment is e7beb672, with 4 commits on each side. Two of the four commits here (459f8a5e, 261921d6) are a divergent duplicate of #51's SwiGLU work rather than a continuation of it, which is why git merge-tree reports add/add conflicts on:
tests/models/test_qwen35_exact_swiglu_grad_engagement.py
tests/models/test_qwen35_exact_swiglu_policy.py
tests/ops/test_exact_fp32_silu_and_mul.py
src/xorl/ops/fused_silu_and_mul.py
src/xorl/models/transformers/qwen3_5{,_moe}/modeling_qwen3_5{,_moe}.py
So #51 and this PR contain two different implementations of the same operator change, and merging in the declared order requires discarding one of them by hand.
This branch also inherits #49's diverged base — see my review there.
Also
pre-commit run --all-files fails ruff-format at this head on 3 files added by this PR:
src/xorl/distributed/pp_byte_contract.py
tests/distributed/test_pp_byte_alignment.py
tests/distributed/test_pp_byte_contract.py
The body says "production code and conventional tests only" with a clean validation section. main passes pre-commit cleanly, so this would fail lint CI on merge — it just hasn't run, because lint.yml triggers only on branches: [main].
Requested changes
- Rebase onto #51's actual head (or onto whatever survives if #51 is closed in favour of this branch's scoping — which I'd suggest, see my review on #51), dropping the duplicated SwiGLU commits.
- Run
pre-commit run --all-files.
I couldn't verify the PP1-vs-PP2 replay claims — those artifacts aren't in the diff and there's no test CI in the repo. Flagging that as a gap in what review can establish here, not as a doubt about the result.
4d2688b to
d19cd28
Compare
28254bd to
4fbd91c
Compare
|
Superseded by #57, which consolidates the complete exact trainer-serving stack directly onto main. |
2 participants
Add a fail-closed PP byte-boundary contract for exact lanes
Stack position: after the scoped one-round FP32 SwiGLU fix.
Problem
Pipeline parallelism must not change the byte program used by an exact
trainer/sampler pair. Unsupported cuts, dtypes, layer identities, or malformed
packed-sequence metadata could otherwise engage partially or change forward
semantics silently.
Change
required embedding/final-norm/head placement.
identity, and complete value-correct varlen metadata.
active, so a late failure cannot leave earlier stages partially engaged.
bi_fusedPPloss construction.
Validation
cases for rogue FP32 parameters, malformed metadata, invalid cuts, and
partial engagement.
qualified environment.
PP1 and PP2, K3 = 0.0.
Admission defaults remain conservative; this PR establishes the contract and
gates rather than broadening production topology by assertion. It contains
production code and conventional tests only; design drafts and campaign
receipts are not included.
Generality follow-up (2026-08-13)
exact_contract_family()now treats the resolution-time_exact_contract_familystamp as authoritative, including a stampedNonefor generic models. Direct-construction configs that predate stamping retain the shared legacy resolver. The classifier no longer owns a hard-coded Qwen model-type set, so admitting another family is a resolver/registry change backed by byte evidence rather than classifier surgery.