Add architecture-specific MoE and zero-K3 training support

[kiddyboots216]

Summary

Complete the MoE, GDN, and Qwen3.5 trainer surface needed for the validated zero-K3 path.

This adds serving-layout fused experts, deterministic routing and scatter, batch-invariant router behavior, native ordered EP combine, Qwen3.5 MoE and dense GDN wiring, FlashQLA packaging, decode and prefill kernels, residual normalization, and token diagnostics. It also performs the final shared-file reconciliation across model, loss, runner, trainer, and weight-synchronization paths.

The final reconciliation corrects Qwen3.5 MRoPE feature rotation at the dense and MoE attention call sites and restores the temperature-aware fused-head forward/backward contract required by policy losses.

The filesystem fallback for R3 routing side payloads now uses a versioned JSON manifest and fixed-dtype raw tensor files. The reader validates paths, shapes, dtypes, byte counts, and file sizes, refuses symlinks, and rejects legacy pickle references before opening payload data.

This head inherits the Foundation trust-boundary hardening and is explicitly reconciled with the repaired Numerical base. Its MoE-specific runner deltas retain Foundation's R3 slicing, diagnostics, rank-local row batching, and ZORL dispatch while adding cross-rank top-token diagnostics.

Raw experiment outputs, launch manifests, shared-environment operations, and machine-specific dependency mutations are excluded. This PR supersedes #16 and #20.

Validation

• All-files pre-commit, the repository privacy guard, and git diff --check passed.
• Every Python test path changed from the Numerical head to this head passed in the declared dependency environment: 282 passed, 11 declared skips, 0 failures.
• The focused dispatcher and request-processor suite passed after the filesystem-spool hardening: 42 passed. Dedicated legacy-pickle, symlink, safe-slice, cleanup, and repository-privacy checks also passed.
• The final focused trust-boundary/R3/privacy selection passed: 47 tests, 0 failures. The versioned runner protocol selection passed: 7 tests, 0 failures.
• After the final Broly remediation was inherited, the combined endpoint, diagnostic, adapter, runner, side-payload, request-processor, protocol, and privacy selection passed: 133 tests passed, 1 declared skip, 0 failures. Ruff and formatting checks passed.
• On the final stacked head, the expanded security, protocol, artifact-backed norm, and Qwen RMSNorm model selection passed: 306 tests passed, 1 declared skip, 0 failures.
• The final constrained evidence-harness follow-up passed: 51 tests passed, 1 declared skip, 0 failures.
• Optimizer resume now uses safetensors plus strict JSON metadata; bitwise resume, fail-closed legacy handling, and transactional rollback regressions passed.
• After the review-fix restack, every Python test path changed by this PR passed: 383 passed, 10 declared skips, 0 failures. All-files pre-commit, the privacy guard, and git diff --check passed.
• The paired 64-token trainer/sampler replay was exact at every token.
• A live two-rollout, 128-completion-token forward/backward and Adam step passed with behavior K3 0.0, ratio mean 1.0, policy KL 0.0, and both clip fractions 0.0.

Class-B RoPE closure

The apparent Class-B residual was an FSDP transport defect, not a remaining trainer/sampler arithmetic mismatch. Decoder FSDP2 units recursively downcast the model-root FP32 cos/sin tuple to BF16 before the Class-B apply function. The engagement tally covered the apply route but not table precision.

This rebased head inherits the generic fail-closed transport repair from PR #22 and retains Qwen3.5's shared Class-B dispatch.

The corrected one-GPU dense Qwen3.5 replay passed on 128 completion tokens with behavior K3 0.0, ratio mean 1.0, negative log-ratio mean 0.0, Class-B-only engagement, and a healthy backward/optimizer step. Focused Qwen3.5 and generic contract tests passed (9 tests). This certifies only the measured one-GPU dense lane.

Remaining gates

The exact paired result certifies the one-GPU dense Qwen3.5 lane. At-scale expert equality, routing vitality and replay, native EP8 forward/backward/combine order, and paired replay/update for additional production topologies remain required. A world-size-4 native-EP mechanics run is partial evidence, not EP8 certification.

Stack

This PR targets the Numerical-contract PR.

[broly-code-security-scanner]

Broly Security Scan

Note

✅ Clean scan
No vulnerabilities detected in this PR.

Note

Re-scan this PR anytime with /broly scan — useful after /broly undismiss, or to refresh findings without a new push.

Broly — SAST (zai-org/GLM-5.2) · Secrets · SCA · IaC · GH Actions · Base Images · Supply Chain Threats · Exploit Chains · Adversarial Verification

We're continuously improving Broly's accuracy and finding quality — your feedback is valuable. False positives, missed findings, bugs, and feature requests all welcome.

 

[qywu]

Reviewed: description is clear, CI passing, no suspicious file changes. LGTM.

[qywu]

Deep review summary

Read the full description and the ~7.5k line / 55-file diff (not just the first screen), cross-checked the MoE EP-combine/scatter-gather math, the norm-family dispatch, KV-cache threading, and the security claims against the actual diff.

What checks out

• _ep_combine_native / _ep_combine_sim (src/xorl/models/transformers/qwen3_5_moe/modeling_qwen3_5_moe.py) and exchange_and_chain_sum (src/xorl/models/layers/moe/ep_native_combine.py) use a consistent (n-1) -> 0 chain-sum order across the simulation path, the native training path, and the docstrings — no order mismatch.
• _AllGatherSumBackward (ep_native_combine.py:617-643): the reduce-scatter-sum backward is correct given that every EP rank's forward consumes the entire gathered token batch (each rank's local shard genuinely receives gradient contributions from every other rank's forward use of it).
• experts.py's filter_expert=True masked backward (compacting valid (token, slot) pairs before expert_histogram/moe_index_compute, then scattering grads back with valid_flat) correctly handles the numel()==0 all-masked case by returning exact zeros, and the L2Norm/RoPE/eager-mask fixes in eager.py and multi_head_attention.py are consistent with the added past_key_values incremental-decode support.
• model_runner.py:6530-6504 world-size guard on dist.all_reduce(loss_report, ...) matches the described single-GPU NCCL/libnvidia-ml.so.1 fix and is scoped correctly.

Non-blocking observations

• Description vs. diff: the "Filesystem fallback for R3 routing side payloads ... versioned JSON manifest ... refuses symlinks, rejects legacy pickle" paragraph describes code that already exists pre-PR (runner_dispatcher.py manifest/symlink checks, security.py::resolve_path_within) — none of that logic appears in this diff. It's disclosed as "inherited," but the wording ("now uses...") reads like new work in this PR; worth tightening so reviewers don't credit this PR for hardening it didn't add.
• qwen3_moe/modeling_qwen3_moe.py adds a large amount of env-var-gated diagnostic scaffolding (XORL_QWEN3_MOE_DELAYED_RESIDUAL_PAIR*, ..._O_PROJ_PARTIAL_RESIDUAL*, ..._CAPTURE_O_PROJ_PARTIAL_RESIDUAL_CANDIDATES*) with several near-duplicate boolean-env-parsing helper functions. It is off by default and does have dedicated tests, but it's a lot of surface area/branching to carry for a "complete MoE support" PR — consider consolidating the repeated {"1","true","yes","on"} parsing helpers and confirming these diagnostic-only paths are pruned once bring-up is done.

No logic bugs, race conditions, or trust-boundary issues were found in the reviewed diff that would justify blocking. Filing as comment-only given the above are non-blocking clarity/complexity notes rather than defects.