[S1ro1]

Nemotron-H now runs as a native PyTorch module through PrimeRL Mamba-2, shared attention, and unified MoE layers. Both Nano (regular MoE) and Super (latent MoE) layouts are retained. Packed convolution and SSD reset at document boundaries, and Mamba owns the full-sequence/head redistribution required by Ulysses context parallelism. Checkpoint conversion maps official per-expert and latent-expert weights into the canonical runtime layout.

Updated against main (daa72a91d), retaining current fusion, normalization, and CP interfaces. Restored LoRA adapter mixer names and installed Mamba SSD through the GPU extra. The pinned upstream source uses Triton without requiring the optional selective-scan CUDA extension. Configuration guards and HF execution comparisons were removed; existing packed-boundary tests exercise real GPU kernels. Router compute follows trainer.model.moe_router_dtype (float32 by default, optionally bfloat16); the model does not force an override.

NemotronHModel and NemotronHForCausalLM both inherit NemotronHPreTrainedModel, following the existing pattern used by other architectures. The shared Nemotron base retains the config type and architecture-specific conversion/CP hooks. Unused TP/PP plans, no-split metadata, checkpointing flags, redundant attributes, and obsolete use_mamba_kernels handling are removed. Shared model construction, LoRA, LM-head patching, and weight-transfer code remain unchanged. Broader removal of the model abstractions is deferred to a follow-up PR.

Garrett's config-docstring and unused-HF-attribute comments are addressed. The broader context-parallel setter refactor and additional gradient tests remain deferred.

Validation

All GPU validation used Slurm-allocated Hopper H200s. Current revision f6d7a4033 restores the shared Nemotron base while retaining the attribute cleanup. Ruff lint/format checks, factory construction/dtype/inherited-hook checks, and all 15 existing model tests passed (Slurm job 343).

• Full unit suite: 699 passed, 11 skipped, rerun after the nn.Module cleanup (Slurm job 336).
• Targeted Nemotron, packed-boundary, and KL tests: 15 passed, rerun after removing the router override.
• H200 checks verified gate GEMM and weight-gradient dtypes for both BF16 and FP32 router settings.
• Compiled two-GPU Mamba Ulysses: output, input-gradient, and reduced parameter-gradient checks passed.
• Fused QKV + compiled selective activation checkpointing: forward/backward and meta-buffer initialization passed.
• Eight two-GPU checkpoint-resume combinations: exact weights, CPU-offloaded AdamW state, and next update across fused/unfused QKV and FSDP shard placements.
• Official Nano checkpoint: exact key/shape conversion round trip, 6,243 → 401 → 6,243 tensors.
• uv sync --all-extras --all-packages, Ruff lint, and format checks passed.

The 20-step run below was completed at a7f798048 (Slurm job 339), before restoring the original shared base. The preceding revision 00746d6d2 was also validated: direct inheritance and the original factory passed construction checks; compiled two-GPU Mamba and fused-QKV selective-AC checks passed again (job 341). All 15 targeted unit tests passed (job 340). The two-step end-to-end Nano run completed (job 342), with mismatch KL 0.00475226 and 0.00522814, finite gradients, and three successful NCCL weight updates including startup. Compiled Mamba/selective-AC and eight checkpoint-resume probes also passed again on this code (jobs 337–338). The end-to-end run used nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 on eight Slurm-allocated H200s: four trainer GPUs with EP=4 and four inference replicas. The native trainer used FlashAttention 3, compilation, selective activation checkpointing, and NCCL weight updates. Reverse-text used batch size 64, group size 8, sequence length 1024, and maximum completion length 512. Optimization and reduction dtypes retained their defaults.

Launched with uv run rl @ rl-module-cleanup.toml. The config sets trainer.env_vars.TRITON_CACHE_DIR to node-local /tmp to avoid shared-home artifact read failures; no diagnostic hooks are present in this run.

The automated review requests a math-environment table. This run uses reverse-text and does not satisfy that environment-specific requirement: the current math task requires a policy-aware container runtime, which is unavailable on the validation nodes.

20/20 steps completed, with 21 successful NCCL weight reloads (startup plus every update). Mismatch KL: mean 0.0059826873, maximum 0.0076148757, final 0.0071559078. All step KL values and gradient norms were finite. Super latent behavior was exercised by small model tests and distributed probes; no full-size Super run was performed.

Step	Mean mismatch KL
1	0.0061262716
2	0.0044362689
3	0.0046579577
4	0.0057872585
5	0.0052399668
6	0.0067561148
7	0.0045683677
8	0.0047115656
9	0.0063753380
10	0.0065832101
11	0.0062584030
12	0.0072330278
13	0.0066353143
14	0.0060227909
15	0.0053012762
16	0.0066529727
17	0.0055223401
18	0.0060145170
19	0.0076148757
20	0.0071559078

CI at a7f798048: Ruff, slim install, and CodeQL passed; remaining jobs were running or queued at the last check. On the previous commit (a60d27578): all 15 Nemotron tests passed. The GPU unit job had five failures in unchanged DeepSeek-V4 tests (packed-batch assertions and unsupported FP8 kernels on that runner); CPU unit tests, Ruff, slim install, and CodeQL passed.

---

Note

High Risk
Large rewrite of a hybrid MoE/Mamba training path, checkpoint layout, and distributed Mamba CP, with dependency pins that affect the whole GPU stack.

Overview
Replaces the Hugging Face–wrapped Nemotron-H stack with a native hybrid model: one NemotronHDecoderLayer per block type (Mamba, attention, MoE), a first-party NemotronHMamba2 that uses FLA packed causal conv plus mamba_ssm Triton SSD, and Ulysses CP via sequence↔head all-to-all inside the Mamba module (the shared cp_mamba helper is deleted).

Checkpoint and config: NemotronHConfig is realigned to current HF field names and layer_types; conversion is explicit per layer type instead of inferring from key presence (the old router-bias MapValue shift is gone). Install/docs now pull mamba-ssm from a pinned git rev through the gpu extra, with tilelang==0.1.12 pinned for vLLM compatibility; trainer init no longer disables use_mamba_kernels or monkey-patches HF Mamba.

Tests drop HF parity checks and instead exercise native packed boundaries, CP wiring on NemotronHMamba2, and existing KL/conversion coverage.

^{Reviewed by Cursor Bugbot for commit f6d7a40. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 6c97e62. Configure here.}

[garrett361]

LGTM, just minor comments. I had forgotten that nemotron is nope.

[garrett361]

Moving to dropping the config doc strings everywhere? Would be happy to have one fewer thing to keep in sync.

[garrett361]

why? compatibilty?

[garrett361]

is anything actually using _no_split_modules anymore? Not a question specific to this PR.

[garrett361]

similar non-nemotron-specific question: are we using these anywhere?

[garrett361]

nit: we have a lot of very similar set_context_parallel_attributes type functions throughout, would be nice to simplify this (nemotron and beyond)

[garrett361]

worth testing the grads match where we expect them to, also? maybe not

[garrett361]

I think not, and actually don't know how much we need this test. More of a test of fla.modules.conv than anything else, which isn't our responsibility (though of course good to know if broken). Up to you.