feat(model): add Nemotron-3 Nano hybrid Mamba-Transformer MoE architecture by le1nux · Pull Request #459 · Modalities/modalities

le1nux · 2026-07-28T16:43:53Z

Adds support for hybrid Mamba-Transformer architectures with sparse mixture-of-experts feed-forward layers, targeting Nemotron-3 Nano 30B-A3B (arXiv:2512.20848). Reference behaviour was taken from Megatron-Bridge's HybridModelProvider recipe and Megatron-LM's megatron/core/ssm + megatron/core/models/hybrid.

~2,900 lines of implementation, ~3,100 lines of tests, 258 tests. The GPT-2 implementation is untouched.

What the architecture is, and why it needs new components

I confirmed the spec from the model report (Table 1 / Figure 2) and cross-checked it against the reference recipe; they agree. Three properties drove nearly every design decision:

1. Layers are single-operator, not blocks. A classical transformer block bundles attention and an MLP. Here every layer wraps exactly one operator in a pre-norm residual, and the stack is described by a pattern string:

MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME   # 52 layers

23 Mamba-2 (M), 23 MoE (E), and only 6 attention (*). Because the Mamba-2 layers carry positional information, the model uses no positional embeddings at all. Changing the Mamba/attention/MoE ratio is a one-line config edit.

2. Attention decouples head dim from model dim. 32 query heads × 128 head dim = 4096, on a model dimension of 2688. GPT-2's attention hard-derives head_dim = n_embd // n_head_q, so it genuinely cannot express this — a new module was required, not a preference.

3. Experts are non-gated. Squared ReLU, so two matrices per expert instead of three. My parameter accounting reproduces both published figures exactly (31.6B total, 3.2B active / 3.6B with embeddings), which is what confirmed the non-gated reading and that Megatron fuses the 2 shared experts into one MLP of hidden 3712 = 2×1856.

The central design decision: layer specs, not layer instances

ComponentFactory memoises every config node. Injecting an instantiated mamba_layer component would have made all 23 Mamba layers share one weight tensor — a model that trains and silently underfits. This is exactly what Megatron-LM's ModuleSpec exists to avoid.

So each network component is registered as a builder:

class NemotronLayerSpecIF(ABC):
    @property
    def symbol(self) -> LayerSymbol: ...   # M | * | E | -
    def build(self, layer_idx: int) -> nn.Module: ...

The model calls build() once per position. Every hyperparameter of every component stays reachable from YAML, while each layer gets independent parameters. test_repeated_layer_types_have_independent_weights pins this down through the real YAML path.

Structural choice: mirror the GPT-2 module tree

transformer.{wte, h.<i>, lm_head_norm, lm_head}

transformer.h is an nn.ModuleDict because that is what the activation-checkpointing component requires. The payoff: activation_checkpointed, fsdp2_wrapped, pipeline splitting, and ChunkedCLMCrossEntropyLoss all work unchanged — no new wrapper components, and layers_fqn: transformer.h just works.

Load balancing

Nemotron uses two mechanisms and both are implemented:

Aux-loss-free expert bias (primary, rate 1e-3) as an optimizer step pre-hook, registered by a new optimizer decorator variant moe_load_balanced. Attaching it to the optimizer step rather than the forward pass is what makes it correct under gradient accumulation: counts accumulate over micro-batches and reduce across data-parallel ranks exactly once per step. Because the rule uses only sign(mean − load), activation-checkpointing's double-counted forward is harmless.
Classic load-balancing loss (secondary, 1e-4), computed per sequence, surfaced through the output dict and combined via a new weighted_sum loss. Logits are inserted first so InferenceResultBatch.__len__ stays correct.

Four real bugs found and fixed during the work

These are the parts worth reviewing most closely — each would have been silent in production:

GroupedExperts weights were uninitialized memory. w1/w2 are raw nn.Parameter(torch.empty(...)), so nothing initialized them if a regex failed to match. Found via a NaN that only appeared under a specific test ordering. Fixed with a reset_parameters giving nn.Linear-equivalent init.
reset_parameters wrote plain tensors into DTensors. Under FSDP2, dt_bias.copy_(plain_tensor) raises. Fixed by writing into the local shard — valid because the values are i.i.d. per head. Found only by the distributed test.
The pipeline stage packer silently dropped the lm_head. The inherited greedy packer assumes uniform layer weights; with per-layer-type costs it exhausts its budget early and drops trailing modules. Verified directly:
```
MISSING: {'transformer.lm_head_norm', 'transformer.lm_head'}
```
Replaced with a binary-search-plus-split partitioner that provably covers every module. The 52-layer model now splits into 4 stages with weights [31, 32, 31, 31] (ideal 31.25).
Weight init silently did nothing whenever activation checkpointing was enabled. The shared initializer stripped only torch.compile's _orig_mod., not _checkpoint_wrapped_module., so every per-layer regex failed to match and models kept their default init. This affected GPT-2 configs identically. Fixed in initialization_routines.py by normalizing all wrapper prefixes. This is the one change made outside the stated scope — leaving a known silent-correctness bug in the path this feature's own config uses would have been wrong. Llama3Initializer has the same gap but lives in the GPT-2 package, so it was left alone and documented.

Testing

Suite	Result
Nemotron + components (CPU/GPU)	253 passed, 3 skipped
Distributed FSDP2 (2 GPUs)	2 passed
`tests/models`, `nn`, `config`, optimizer, loss, weight-tying	362 passed, 4 skipped
Full collection	814 tests collect cleanly
isort / black / ruff	clean

The correctness anchor is test_ssd.py: the chunked SSD scan is validated against a step-by-step transcription of the Mamba-2 recurrence — not against itself — plus chunk-size invariance, causality, state-passing across a split sequence, and gradient flow. The fused Triton path is then checked against the native one, so it inherits that validation transitively.

Causality is asserted at every level (scan, conv, mixer, attention, assembled model), because a causality bug leaks the target and shows up as a suspiciously good loss rather than a crash.

Not verified: the end2end_tests and checkpointing suites need a multi-GPU launcher. Note also that TestLlama3LikeInitialization is pre-existing flaky — baselined with this PR's changes reverted, it still fails, with a different parametrization each run.

Caveats

The default Mamba backend is slow. ssd_backend: native is dependency-free, CPU-runnable and torch.compile-friendly, but materially slower and more activation-memory hungry than the Triton kernels. It exists so the model is testable without mamba-ssm (not a Modalities dependency). For real 30B training you want pip install -e '.[mamba]' and ssd_backend: fused — the fused parity tests were skipped in this environment, so that path is written but unverified here.

No tensor, expert, or context parallelism. Mamba TP is genuinely hard: in_proj is a five-way unequal column split, and conv1d/A_log/D/dt_bias need per-head sharding with partition metadata. Expert parallelism has no dimension in the device mesh at all. The reference recipe trains 30B-A3B with EP=8; here you get FSDP2 + AC + PP, which is enough to train the model but not at the reference topology or throughput.

MoE aux loss and pipeline parallelism don't compose. Under PP, MoE layers on non-final stages can't reach the loss. The aux-loss-free bias balancing is layer-local and unaffected, which is why it is the primary mechanism — but if you enable PP, drop the weighted_sum loss term.

Numerics are unvalidated against the reference. Parameter shapes, initialization distributions and the routing formula were matched to Megatron-LM by reading the source, and the parameter counts reproduce the paper. No Nemotron checkpoint was loaded to compare logits. Until someone does, treat architectural equivalence as "carefully derived", not "measured".

MFU is an estimate. The formula charges the quadratic term only to the 6 attention layers and uses active parameters, but it ignores the SSM scan's own FLOPs, so it will read slightly optimistic.

Next steps for production at scale

Numerical equivalence against the reference. Convert an HF NemotronHForCausalLM checkpoint and assert logit parity. The parameter mapping is fully specified in nemotron_h_bridge.py:369-410, and module names were chosen to keep that mapping direct. This is the single highest-value follow-up — everything else is performance.
Fused backend validation and benchmarking. Install the extra, run the parity tests on GPU, then measure tokens/s and peak memory against the native path to size the real gap.
Tensor parallelism for Mamba. Shard in_proj with explicit per-component partition sizes and conv1d/A_log/D/dt_bias per head. Attention and MoE TP are straightforward by comparison.
Expert parallelism. Requires an ep dimension in DeviceMeshConfig plus an all-to-all token dispatcher. This is the difference between training 30B-A3B on 8 GPUs and on 512. TorchTitan's token_dispatcher.py is a reasonable template.
Selective-op activation checkpointing tuning. The reference recomputes selectively over ["moe", "layernorm"]; the SELECTIVE_OP variant can express something similar, but the op list needs tuning for the SSD scan.
MTP and long-context. Multi-token prediction (mtp_hybrid_override_pattern="*E") and the 512k-sequence long-context phase are both absent; the latter needs context parallelism with Mamba state passing.
Convergence smoke test. Train the tiny variant to a known loss on lorem-ipsum, then a ~1B config on real data for a few thousand steps, watching expert-load entropy to confirm the bias balancing actually prevents collapse at scale. The unit tests prove the update rule is correct; only a real run proves it is sufficient.

Docs

See docs/components/nemotron.md for the full guide, and config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml for a ready-to-run 30B-A3B pretraining config.

🤖 Generated with Claude Code

…cture Adds support for hybrid Mamba-Transformer architectures with sparse mixture-of-experts feed-forward layers, targeting Nemotron-3 Nano 30B-A3B (arXiv:2512.20848). Reference behaviour taken from Megatron-Bridge's HybridModelProvider recipe and Megatron-LM's megatron/core/ssm. The architecture is a sequence of single-operator pre-norm residual layers described by a pattern string (M = Mamba-2, * = attention, E = MoE, - = dense MLP) rather than a stack of attention+MLP blocks. Nemotron-3 Nano uses 52 layers of which only 6 attend; the Mamba-2 layers carry the positional information, so the model has no positional embeddings. Network components are registered as layer *specs* (builders) rather than modules. The component factory memoises config nodes, so injecting instantiated layers would make all layers of a type share one weight tensor. The model calls build() once per layer position instead, which keeps every hyperparameter configurable from YAML while giving each layer independent parameters. This mirrors Megatron-LM's ModuleSpec. The module tree deliberately matches GPT2's (transformer.{wte,h,lm_head_norm,lm_head}, with h an nn.ModuleDict), so the existing activation-checkpointing, FSDP2, pipeline-splitting and chunked-loss components apply unchanged. New components: * model/nemotron and nemotron_layer_spec/{mamba2,attention,moe,mlp} * Mamba-2 mixer with a dependency-free pure-PyTorch state space scan (ssd_backend=native) plus an optional fused Triton backend (ssd_backend=fused, new "mamba" extra) * MoE block with sigmoid top-k router, grouped-matmul experts, shared experts * optimizer/moe_load_balanced: auxiliary-loss-free load balancing (arXiv:2408.15664) as an optimizer step pre-hook, which makes the expert bias update correct under gradient accumulation * loss/{moe_aux_loss,weighted_sum} for the sequence-level balancing penalty * stages_generator/nemotron_stages_generator, weighting pipeline stages by per-layer-type cost and guaranteeing complete module coverage * mfu_calculator/nemotron, using active parameters and charging the quadratic attention term only to the attention layers * model_initialization model_type "nemotron"; the state space parameters (A_log, D, dt_bias, conv1d) are excluded from the regex-driven initializer and initialized by Mamba2Mixer.reset_parameters instead Also fixes a pre-existing bug in NamedParameterwiseNormalInitialization: it only stripped torch.compile's _orig_mod. prefix from parameter names, so any config applying activation checkpointing before weight initialization silently matched no per-layer regex and kept the model's default initialization. This affected GPT2 configs as well. 258 tests. The chunked SSD scan is validated against a step-by-step transcription of the Mamba-2 recurrence, plus chunk-size invariance, causality and state passing. Two distributed FSDP2 tests cover expert sharding and a full training step with load balancing. The GPT2 implementation is untouched. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>

Modalities is MIT licensed; Megatron-LM, Megatron-Bridge and state-spaces/mamba are Apache 2.0 (NVIDIA / Tri Dao, Albert Gu) and TorchTitan is BSD 3-Clause. Preserving those notices in the derived files is a license obligation, so each source file that adapts upstream code now carries a header naming the upstream file, what was taken from it, the copyright holder and the license. Marked as adapted from Megatron-LM: * mamba2_mixer.py - packed [z, x, B, C, dt] projection layout, conv1d/A_log/D/ dt_bias shapes and initialization distributions * layer_pattern.py - the layer pattern symbols * nemotron_layers.py - single-operator pre-norm residual layer structure * nemotron_layer_specs.py - the declarative ModuleSpec/builder pattern * nemotron_model.py - hybrid model structure * router.py - sigmoid scoring, selection-only expert bias, renormalization * moe.py - the load-balancing loss formula and its sequence-level variant * load_balancing.py - the sign-based expert bias update rule * experts.py, nemotron_mlp.py - the squared ReLU activation Marked as adapted from state-spaces/mamba: * ssd.py - the chunk-parallel SSD block decomposition (ssd_minimal_discrete / segsum) and the GatedRMSNorm semantics Marked as adapted from TorchTitan (following the convention already used in model_factory.py and trainer.py): * experts.py - stacked per-expert weights and the torch._grouped_mm call * moe.py - dispatch/combine structure, expert-bias and token-count buffers * router.py - router interface shape * load_balancing.py - the optimizer step pre-hook approach * nemotron_stages_generator.py - pipeline split-point structure Marked as adapted from Meta's Llama: * nemotron_attention.py - grouped-query key/value head repetition The YAML configs note that their hyperparameter values are adopted from the Megatron-Bridge recipe; no code was taken from Megatron-Bridge. docs/components/nemotron.md gains an Attribution section tabulating all of the above in one place, and explicitly lists the files with no third-party derivation (norms.py, nemotron_model_factory.py, nemotron_mfu.py, and the partitioning algorithm in nemotron_stages_generator.py). Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>

…model Mirrors config_lorem_ipsum_long_fsdp2.yaml for the Nemotron hybrid Mamba- Transformer. Keeps the full *width* of Nemotron-3 Nano 30B-A3B - dimension 2688, 128 routed experts of dim 1856 with top-6 routing and 2 shared experts, 64 Mamba-2 heads of dim 64, 32 query / 2 KV attention heads of head dim 128 - and cuts only the depth, from 52 to 16 layers, so that it fits on four GPUs. The layer pattern is the first 16 characters of the published pattern, which preserves the Mamba / MoE / attention interleaving and ratio (7 : 7 : 2 here versus 23 : 23 : 6 in the full model). Verified end to end on 4x A100-SXM4-80GB: 162 steps, exit 0, no errors, loss 11.0 -> 0.03, one DCP checkpoint written and reloadable. Measured figures, both variants run: layers pattern parameters active/token peak alloc/GPU 16 MEMEM*EMEMEM*EME 9.67B 0.77B 46.1 GiB 9 MEMEM*EME 5.64B 0.53B 28.9 GiB The 9-layer variant is documented for 40GB cards. Throughput is 5.5 samples/s at 1-2% MFU, which is the expected cost of the native pure-PyTorch state space scan; use ssd_backend=fused for anything beyond a smoke test. Two things the run surfaced, both recorded in the config comments: * A DCP checkpoint of this model is ~109 GiB, so the config checkpoints once near the end. At the reference config's interval of 16 steps this would have written over 1 TB across a 162-step run. * DCPCheckpointSaving._delete_checkpoint uses Path.rmdir(), which only removes empty directories, so rotating a DCP checkpoint always raises OSError("Directory not empty"). That is a pre-existing bug unrelated to this model; the config therefore keeps every checkpoint (k: -1), as the reference config does. Left unfixed deliberately - it deletes directories and deserves its own change plus a test. Also makes NemotronMFUCalculator derive num_active_params from the model when it is omitted. Requiring it as a literal meant every config carried a magic number that silently went stale as soon as the layer pattern or expert count changed. numel() on a DTensor already reports the unsharded size, so the derivation is correct under FSDP2; a list of pipeline stages is rejected with an explicit error, since summing per-stage counts would not give the whole model's active parameters. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>

DCPCheckpointSaving._delete_checkpoint removed checkpoints with Path.rmdir(), which only removes *empty* directories. A distributed checkpoint is a directory of per-rank shard files (.metadata plus one __<rank>_0.distcp per rank), so deleting one always raised OSError("Directory not empty"). The effect was that every rotating checkpoint strategy was broken whenever it was combined with the dcp execution: save_k_most_recent_checkpoints_strategy with k >= 1 and keep_every_k_steps_and_m_most_recent_checkpointing_strategy both crash training at the first rotation, once enough checkpoints exist for one to be evicted. Only k = -1 (keep everything) worked, which is why the existing lorem-ipsum configs never hit it. Deletion is now recursive. Since that turns a no-op into a real recursive delete, it is guarded: the target must exist, must be a directory, and must resolve to a path located directly inside the configured checkpoint_path. The path itself is not user input - it is rebuilt from CHECKPOINT_FOLDER_STRUCTURE for a TrainingProgress this instance previously saved - but the guard makes that invariant explicit rather than implicit, and conservatively refuses a checkpoint folder that symlinks somewhere else. Adds tests/checkpointing/test_dcp_checkpoint_deletion.py, which builds folders shaped like real DCP checkpoints (shard files plus a nested directory). Five of its seven tests fail against the old rmdir implementation, including an end-to-end rotation through SaveKMostRecentCheckpointsStrategy. The FSDP1 execution is untouched: it deletes individual files with unlink(), which is correct for its single-file checkpoints. Also relaxes the Mamba fused-backend error assertion, which named both optional packages and so broke when only one of mamba-ssm / causal-conv1d was installed, and refreshes the checkpointing comment in the Nemotron lorem-ipsum config now that rotation is no longer a trap. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>

Copilot

Pull request overview

Adds first-class support for Nemotron-style hybrid Mamba–Transformer language models with sparse MoE layers, including YAML-configurable per-layer “spec builders”, load-balancing mechanisms, pipeline-stage generation, MFU estimation, and comprehensive unit/distributed tests. Also fixes a pre-existing correctness bug in DCP checkpoint deletion and hardens weight initialization name-matching under wrapper modules.

Changes:

Introduce Nemotron hybrid model stack driven by a layer-pattern string, plus layer-spec builders for Mamba2 / attention / MoE / dense MLP.
Add MoE routing, experts implementations, aux-loss plumbing, and an optimizer decorator for auxiliary-loss-free load balancing.
Add Nemotron pipeline stage generator, Nemotron MFU calculator, docs, training configs, and extensive test coverage (incl. FSDP2).

Reviewed changes

Copilot reviewed 45 out of 49 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/test_yaml_configs/nemotron_config_initialization.yaml	YAML fixture to build a small Nemotron model and initializer config via the registry.
tests/models/nemotron/test_ssd.py	Correctness tests for native SSD scan + conv + gated RMS norm primitives.
tests/models/nemotron/test_nemotron_layers.py	Unit tests for the four single-operator residual layer types and naming stability.
tests/models/nemotron/test_nemotron_initialization.py	Tests for regex-driven initialization coverage and wrapper-name normalization.
tests/models/nemotron/test_nemotron_config_build.py	End-to-end YAML→component-factory→model integration test for layer-spec builders.
tests/models/nemotron/test_nemotron_attention.py	Tests for decoupled head-dim GQA attention implementations and causality.
tests/models/nemotron/test_moe_load_balancing.py	Tests for expert-bias balancing rule, optimizer hook, and aux-loss plumbing.
tests/models/nemotron/test_mamba2_mixer.py	Tests for Mamba2Mixer shapes, causality, chunk invariance, init, and fused parity (skipped when unavailable).
tests/models/nemotron/test_layer_pattern.py	Tests for parsing and counting layer-pattern symbols, incl. Nemotron-3 Nano pattern.
tests/models/nemotron/init.py	Package marker for Nemotron tests.
tests/models/components/test_norms.py	Tests for NormWrapperConfig behavior (fresh instances per build, correct norm types).
tests/fsdp2_parallelization/test_nemotron_fsdp2.py	Multi-process CUDA tests for FSDP2 sharding/materialization/init and training-step execution.
tests/fsdp2_parallelization/nemotron_fsdp2_config.yaml	Minimal distributed config for the FSDP2 Nemotron tests (AC + MoE aux + load balancing hook).
tests/checkpointing/test_dcp_checkpoint_deletion.py	Regression tests proving DCP checkpoint deletion is recursive and path-safe.
src/modalities/utils/nemotron_mfu.py	Nemotron-specific MFU calculator using active parameters + attention-layer-only quadratic term.
src/modalities/registry/components.py	Registry wiring for Nemotron model, layer specs, MoE optimizer decorator, losses, stages generator, and MFU calculator.
src/modalities/nn/model_initialization/parameter_name_filters.py	Add Nemotron init regex groups and model type to init-filter registry.
src/modalities/nn/model_initialization/initialization_routines.py	Normalize wrapper-modified parameter FQNs before applying init regexes.
src/modalities/models/nemotron/nemotron_stages_generator.py	Cost-weighted, coverage-guaranteeing pipeline stage partitioner for hybrid stacks.
src/modalities/models/nemotron/nemotron_model.py	NemotronLLM model implementation mirroring GPT-2 module tree + aux-loss emission.
src/modalities/models/nemotron/nemotron_model_factory.py	Factory for constructing NemotronLLM (optionally on meta device).
src/modalities/models/nemotron/nemotron_mlp.py	Squared-ReLU (non-gated) MLP used for dense FFN and shared experts.
src/modalities/models/nemotron/nemotron_layers.py	Single-operator pre-norm residual layer wrappers (Mamba2/Attn/MoE/MLP).
src/modalities/models/nemotron/nemotron_layer_specs.py	Layer-spec builder implementations and configs for each layer symbol.
src/modalities/models/nemotron/nemotron_attention.py	Decoupled head-dim grouped-query causal self-attention module (manual/SDPA/dao).
src/modalities/models/nemotron/layer_pattern.py	Pattern parsing/counting utilities and LayerSymbol enum.
src/modalities/models/nemotron/init.py	Package marker for Nemotron model code.
src/modalities/models/components/norms.py	Shared normalization wrapper config for non-GPT2 models.
src/modalities/models/components/moe/router.py	Top-k router with sigmoid/softmax scoring and selection-only expert bias support.
src/modalities/models/components/moe/moe.py	MoE block with routed experts, optional shared experts, aux-loss, and token counters.
src/modalities/models/components/moe/moe_losses.py	Loss components to surface MoE aux loss and combine multiple loss terms.
src/modalities/models/components/moe/load_balancing.py	Auxiliary-loss-free expert-bias update + optimizer step pre-hook decorator.
src/modalities/models/components/moe/experts.py	GroupedExperts implementation with grouped_mm / looped backends and initialization.
src/modalities/models/components/moe/init.py	Package marker for MoE components.
src/modalities/models/components/mamba2/mamba2_mixer.py	Mamba2 mixer with native and optional fused SSD backend, plus distribution-safe init.
src/modalities/models/components/mamba2/init.py	Package marker for Mamba2 components.
src/modalities/config/pydantic_if_types.py	Add Pydantic type wrapper for NemotronLayerSpecIF for component factory validation.
src/modalities/config/config.py	Add MoELoadBalancedOptimizerConfig for the optimizer decorator component.
src/modalities/checkpointing/fsdp/fsdp_checkpoint_saving.py	Fix DCP checkpoint deletion to be recursive and path-guarded.
pyproject.toml	Add optional `mamba` extra for fused kernels (mamba-ssm + causal-conv1d).
docs/components/nemotron.md	Full Nemotron guide (architecture, config, kernels, balancing, init, parallelism, attribution).
docs/components/components.md	Register Nemotron model/spec/loss/optimizer variants in component documentation.
config_files/training/config_nemotron3_nano_30b_a3b_fsdp2.yaml	Ready-to-run Nemotron-3 Nano 30B-A3B FSDP2 training config (native SSD by default).
CHANGELOG_DEV.md	Changelog entry summarizing the new feature set and the two bug fixes.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+            for weight in (self.w1, self.w2):
+                # Match nn.Linear: kaiming_uniform_ with a=sqrt(5) over the fan-in of each expert.
+                fan_in = weight.shape[-1]
+                bound = math.sqrt(1.0 / fan_in) * math.sqrt(3.0)
+                nn.init.uniform_(weight, -bound, bound)


+        h = self.transformer.wte(inputs) if hasattr(self.transformer, "wte") else inputs
+
+        for layer_idx in self.transformer.h:
+            h = self.transformer.h[layer_idx](h)
+


le1nux and others added 4 commits July 28, 2026 18:41

BlueCrescent requested a review from Copilot July 29, 2026 12:52

Copilot started reviewing on behalf of BlueCrescent July 29, 2026 12:53 View session

Copilot AI reviewed Jul 29, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(model): add Nemotron-3 Nano hybrid Mamba-Transformer MoE architecture - #459

feat(model): add Nemotron-3 Nano hybrid Mamba-Transformer MoE architecture#459
le1nux wants to merge 4 commits into
mainfrom
nano_integration

le1nux commented Jul 28, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

le1nux commented Jul 28, 2026

What the architecture is, and why it needs new components

The central design decision: layer specs, not layer instances

Structural choice: mirror the GPT-2 module tree

Load balancing

Four real bugs found and fixed during the work

Testing

Caveats

Next steps for production at scale

Docs

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants