Add common MoE, MiniMax-M3, and Nemotron-H support

[kiddyboots216]

Summary

Add the common MoE substrate and the MiniMax-M3 and Nemotron-H model families that consume it and the preceding SSM work. This includes MoE backends, routing/replay behavior, checkpoint classification, model registry wiring, and focused architecture/checkpoint tests.

Validation

• Common MoE, MiniMax-M3, and Nemotron-H imports passed.
• Focused suite: 48 passed.
• Review follow-up MoE, checkpoint, and routing-replay selection: 28 passed.
• Repository public-reference guard, all-files pre-commit, and git diff --check passed.

Stack

This is 4/9, stacked on the distributed cut.

[kiddyboots216]

/broly scan

[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

Reviewed the full diff (13.7k additions, 46 files: MoE backends, routing/replay, LoRA-for-MoE, checkpoint handlers for MiniMax-M3/Nemotron-H, ops kernels, tests). Overall the MoE substrate (router, aux-loss, backend dispatch) and the checkpoint round-trip tests are solid — test_nemotron_h_checkpoint.py and test_minimax_m3_support.py do real numerical validation (assert_close against HF references), not just "doesn't crash" checks. However, I found concrete, verifiable functional bugs that should block merge.

Blocking issues

1. NemotronHForCausalLM.forward never computes logits — the model cannot train or generate
src/xorl/models/transformers/nemotron_h/modeling_nemotron_h.py (new forward, ~line 8622-8636)

outputs = self.model(...)
return MoeCausalLMOutput(
    last_hidden_state=outputs.last_hidden_state,
    router_logits=outputs.router_logits,
)

self.lm_head is constructed in __init__ but is never called in forward. There's no labels parameter, no logits field, no loss computed — contrast with the sibling MiniMaxM3SparseForCausalLM.forward in this same PR, which does call self.lm_head(...) and computes cross-entropy loss. As written, NemotronHForCausalLM (the exported ModelClass) can only produce hidden states; any training or generation call will fail or silently operate on missing logits. This directly contradicts the PR description's claim of "Nemotron-H model families" support with "focused architecture/checkpoint tests" — the checkpoint/shape tests pass because they don't exercise the forward pass through lm_head.

2. Autograd backward() returns the wrong number of gradients when grad_output is None
src/xorl/ops/moe/quack.py, QuackEPDeepEPCombine.backward (~line 11347-11348) and QuackEPDeepEPNoPermute.backward (~line 11800-11801)

Both forward() signatures take 16 args after ctx (permute_tokens, cumsum, gate_up_proj, down_proj, intermediate_size, expert_scores, buffer, dispatch_ctx, async_combine, hidden_act, activation_native, fp8_compute, fp8_grouped_backend, fp8_block_size, gate_up_bias, down_bias), and the normal-path backward() correctly returns a 16-tuple. But the early-return guard:

if grad_output is None:
    return None, None, None, None, None, None, None, None, None, None, None, None, None, None, None

returns only 15 Nones. PyTorch's autograd engine will raise RuntimeError: function ... returned an incorrect number of gradients, expected 16, got 15 the first time either of these newly-added DeepEP combine/no-permute autograd Functions is invoked with grad_output is None (e.g., an unused output in some loss/branch configuration). This is an easy, real crash path in new code with no test coverage exercising it.

Non-blocking but worth fixing

3. Duplicate _env_flag definitions with incompatible signatures
src/xorl/ops/moe/quack.py, lines ~9103 and ~9266. _env_flag(name, default: str = "0") is defined first and used at import time (_memory_trace_enabled(), _BACKWARD_DEBUG_FINITE); a second, incompatible _env_flag(name, default: bool) is defined later in the same module and permanently shadows the first. This only "works" today because the first version's only callers execute eagerly before the second def runs. Any future lazy/re-entrant call to _memory_trace_enabled() will hit TypeError: _env_flag() missing 1 required positional argument. Should be a single function (or renamed) to remove the landmine.

4. Async D2H copy in routing replay may race
src/xorl/models/layers/moe/routing_replay.py (record/record_weights, ~lines 5869-5889). The device→host copy was changed to copy_(..., non_blocking=pin) with no subsequent stream synchronization or recorded CUDA event before the pinned buffer is appended to top_indices_list/top_weights_list and control returns to the caller. If any downstream consumer of the replay buffers reads them without first synchronizing (not visible in this diff range), this can silently read a partially-copied/stale buffer, corrupting recorded routing decisions used for replay. Please confirm there's a synchronization point before these buffers are consumed, or add one at the record site.

5. Unvalidated expert/EP-size divisibility in checkpoint handlers
src/xorl/models/transformers/nemotron_h/checkpoint_handler.py (~lines 7621-7624): self._local_num_experts = num_experts // ep_size with no assertion that num_experts % ep_size == 0. For non-divisible configs, trailing experts fall into no rank's [start, end) range and are silently dropped from every shard rather than raising. Worth an explicit validation/assert given checkpoint correctness is the whole point of this handler.

Verdict

Requesting changes primarily for issues #1 and #2, which are concrete functional breakages in the new Nemotron-H and Quack/DeepEP code paths, not stylistic concerns.

[qywu]

Approved per maintainer direction, superseding the change-request review above. See prior review comment for the technical issues found; these are not resolved in the diff, tracking as follow-up.