Add NemotronH hybrid Mamba2-Attention model (Nemotron-3-Nano-4B) - #3464
Add NemotronH hybrid Mamba2-Attention model (Nemotron-3-Nano-4B)#3464OmarAzizi wants to merge 3 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the NemotronH hybrid Mamba2-Attention-MoE architecture to MLC-LLM. The changes include the core model implementation, a weight loader for HuggingFace models, a 4B model preset, and unit tests. Feedback identifies an unused variable in the Mamba2 mixer, suggests consolidating redundant dtype casts for better clarity, and recommends using identity checks for booleans in the test suite.
| proj = self.in_proj(hidden_states) | ||
| i0 = self.intermediate_size | ||
| i1 = i0 + self.conv_dim | ||
| i2 = i1 + self.num_heads |
There was a problem hiding this comment.
The variable i2 is defined but never used. It can be removed to improve code clarity.
| x = op.reshape(x.astype(dtype), (b, s, self.num_heads, self.head_dim)) | ||
| B = op.reshape(B, (b, s, self.n_groups, self.ssm_state_size)) | ||
| C = op.reshape(C, (b, s, self.n_groups, self.ssm_state_size)) |
There was a problem hiding this comment.
There are several inconsistencies and redundancies in dtype casting within this function. To improve clarity and correctness, I suggest consolidating the dtype conversions.
BandCare used asfloat32and then cast todtypemultiple times later. It would be cleaner to cast them here along withx.- There are several redundant
.astype(dtype)calls on tensors that are already ofdtype(e.g., onLat line 324,dx_tat line 346, andnew_statesat line 384).
By casting B and C here, you can then remove the subsequent casts on G (line 333), B_t (line 342), and C_t (line 383), as well as the redundant casts mentioned above.
| x = op.reshape(x.astype(dtype), (b, s, self.num_heads, self.head_dim)) | |
| B = op.reshape(B, (b, s, self.n_groups, self.ssm_state_size)) | |
| C = op.reshape(C, (b, s, self.n_groups, self.ssm_state_size)) | |
| x = op.reshape(x.astype(dtype), (b, s, self.num_heads, self.head_dim)) | |
| B = op.reshape(B.astype(dtype), (b, s, self.n_groups, self.ssm_state_size)) | |
| C = op.reshape(C.astype(dtype), (b, s, self.n_groups, self.ssm_state_size)) |
| assert config.mamba_d_conv == 4 # mapped from conv_kernel | ||
| assert config.mamba_chunk_size == 256 # mapped from chunk_size | ||
| assert config.mamba_n_groups == 8 # mapped from n_groups | ||
| assert config.mamba_conv_bias == True # mapped from use_conv_bias |
There was a problem hiding this comment.
For better adherence to Python style conventions, it's recommended to use is for boolean identity checks instead of ==.
| assert config.mamba_conv_bias == True # mapped from use_conv_bias | |
| assert config.mamba_conv_bias is True # mapped from use_conv_bias |
|
Hey @OmarAzizi thanks for the PR! I'll try and get some tests run / code review done on this soon. Just note that it sometimes takes us a second to get model PRs merged (our team is small), so just keep an eye on this branch, and make sure to rebase it off of main in case anything relevant ends up landing before it. This way, it'll keep it easy to merge in. Thanks again! |
|
Also out of curiosity, when you were developing this PR, did you feel like there were any particular sources of significant boilerplate that you had to contend with? Just collecting some feedback from the community! |
|
Hi @babusid, thanks for the heads up, I have been busy lately, but I’ll make sure to keep the branch rebased on main 👍 Regarding testing, I actually tried running the model on a more powerful machine (both CPU and CUDA), but I ran into quite a few errors during execution. I’m still digging into them and would really appreciate your help. |
Regarding this, to be honest, I’m still quite new to these types of projects, so I’m not entirely sure. |
no worries! I'll try to take a look at the branch sometime this week. |
Resolves mlc-ai#3458 Adds support for the NemotronH (nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16) hybrid architecture to MLC-LLM. Architecture: - 42 layers: 21 Mamba2 SSM + 17 MLP + 4 GQA Attention - Layer pattern from hybrid_override_pattern in HF config - Only 4 attention layers use KV cache slots (not all 42) - Pure Relax/TVM implementation of chunked SSD algorithm - No custom CUDA kernels required Changes: - python/mlc_llm/model/nemotron_h/: new model package - nemotron_h_model.py: NemotronHConfig, NemotronHMamba2Mixer, NemotronHAttention, NemotronHMLP, NemotronHMoE (for future 30B), NemotronHDecoderLayer, NemotronHForCausalLM - nemotron_h_loader.py: HF weight name mapping (backbone.* prefix) - python/mlc_llm/model/model.py: register nemotron_h model type - python/mlc_llm/model/model_preset.py: add nemotron_h_4b preset - tests/python/model/test_nemotron_h.py: 5 pytest tests Implementation notes: - HF config uses hybrid_override_pattern (M/*/- chars) instead of layers_block_type list; __post_init__ handles both formats - HF field names differ from our internal names (conv_kernel -> mamba_d_conv, n_groups -> mamba_n_groups, etc); mapped in __post_init__ - HF weight prefix is backbone.* not model.*; handled in loader - MoE layer type included for forward compatibility with 30B Nano model which uses Mamba2+MoE+Attention; not used by 4B model Tested: - All 5 pytest tests pass - mlc_llm convert_weight: 263/263 params, 3.97B total, 2.08GB q4f16_1 - mlc_llm gen_config: succeeds - mlc_llm compile --device cpu: succeeds, lib.so generated Signed-off-by: OmarAzizi <oalazizi75@gmail.com>
…cal fixes After mlc-ai#3449 introduced kHybrid KVStateKind, NemotronH crashes at runtime with an unrecognized AttnKind. This PR migrates to the new runtime path and fixes several numerical bugs in the chunked SSD math that were causing all inputs to map to a single special token. Runtime / compile: - Route nemotron_h through "hybrid" in _infer_kv_state_kind. - Add create_rnn_state and thread rnn_state through batch_*; update spec. - Drop layer_partition arg in create_paged_kv_cache (was [0, 42], should default to [0, num_attention_layers] = [0, 4]; was the actual cause of the AttnKind crash). - Migrate tir -> tirx (catches up with mlc-ai#3462). GPU compile (Mamba2-specific): - Replace concat+split-based dc construction with a direct te.compute (avoids buffer aliasing during lowering). - Hand-schedule inter-chunk state propagation as a T.prim_func with tirx.is_scheduled=1; mirrors qwen35's GatedDeltaNet pattern. Dlight's GPU reduction rule fails to bind threads when the reduction dim is small for short prefills. - Fix dtype mismatch in causal_conv1d's te.if_then_else. SSD numerical fixes: - Add causal mask on intra-chunk L. Without it, exp(cumsum[i] - cumsum[j]) for j > i blows up to inf and propagates NaN through Y_diag. - Compute dt/A/dA, Y_diag, Y_off, D skip in fp32 (HF runs SSM in fp32). NemotronH's dt_bias up to 33.5 + A_log up to 8.6 makes |dA| ~2e5 which overflows fp16. - Use cumulative LC for inter-chunk decay; raw per-chunk lc didn't correspond to any physical decay. - Fix off-by-one in dc and add causal mask in the chunk dimension. Validated: 4 unit tests pass; convert_weight, gen_config, compile (CPU + CUDA) all succeed; GPU inference runs at ~20 tok/s on RTX 3050. Known follow-up: output is varied real text (no longer a special-token loop) but still not coherent. Suspect q4f16_1 quantization aggression on Mamba2's sensitive matrices or one more subtle bug; needs HF transformers reference comparison to localize.
7e7286a to
788964c
Compare
|
Hi @babusid! Quick update. I rebased onto main and pushed the fixes (force push since the rebase rewrote history). Sorry for the rebase noise. I dug into the runtime errors I mentioned. Main one was the kHybrid KVStateKind transition from #3449: NemotronH was still using Other fixes in the same commit:
Pipeline runs end-to-end on CUDA at ~20 tok/s on my RTX 3050. Output varies per input, but still isn't coherent, likely q4f16_1 being aggressive on Mamba2's sensitive matrices, or one more subtle bug. I would appreciate your review whenever you have time! |
Hey sorry for the slow review, it's been a bit hectic. I will try to get to this soon! I will try some known tricks and see if I can spot where a bug might be happening. Question for you - is the output sometimes coherent and sometimes not? Do you observe any crashes? |
No problem at all, I have been really busy myself lately. To answer your question, the model output always seems to be incoherent, but I experienced no crashes on my end. |
okay. This could point to numerical issues / errors that still manages to emit token ids within the valid vocabulary range. I will take a look and run some tests, but for your reference we have seen slight errors cause this type of behavior before. A potential source of error (not sure if it's core to this model, just something I've seen before) is incorrect positional encoding leading to incorrect attention outputs. Dumping and numerically comparing tensors to see where divergence happens is a potential strategy to fixing this. |
2 participants
Summary
Adds support for the NemotronH architecture family to MLC-LLM, resolving #3458.
The first supported model is nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16.
Architecture
NemotronH uses a hybrid layer pattern defined by
hybrid_override_patternin the HF config:Only the 4 attention layers allocate KV cache slots, keeping memory usage low.
Implementation
Pure Relax/TVM - no custom CUDA kernels required.
The Mamba2 chunked SSD algorithm is implemented entirely in TVM Relax + TIR ops, translated from the reference
torch_forwardpath in HuggingFace transformers.Key implementation decisions:
hybrid_override_patternstring (M/*/-chars) auto-converted tolayers_block_typelist in__post_init__.backbone.*weight prefix instead ofmodel.*; handled in loader.conv_kernel,n_groups,chunk_size, etc); mapped in__post_init__.kHybridKVStateKind: provides bothcreate_paged_kv_cache(for the 4 attention layers) andcreate_rnn_state(for the Mamba2 layers, currently a stub matching this PR's prefill-only design — proper recurrent caching is a follow-up).dt_bias/A_logvalues, producing NaN indA_cumsumand breaking inference.Changes
python/mlc_llm/model/nemotron_h/: new model packagepython/mlc_llm/model/model.py: registernemotron_hmodel typepython/mlc_llm/model/model_preset.py: addnemotron_h_4bpresetpython/mlc_llm/interface/compile.py: routenemotron_hthrough thekHybridruntime pathtests/python/model/test_nemotron_h.py: 4 pytest testsTest Plan
Unit tests (no download required):
Full pipeline test (requires ~8GB download, GPU recommended):