Repository navigation
fix(nemotron_nas): greedy output is still degraded after the RoPE fix in b429947c #1688
Description
Problem / Background
Commit b429947c fixed one real defect in src/models/nemotron_nas.rs: NASAttention::forward reshaped to [B, L, H, D], applied fast_rope, and only then transposed to [B, H, L, D]. fast_rope reads token positions off the second-to-last axis, so the head index was serving as the position. That was found by auditing the tree after the same defect turned up in recurrent_gemma.rs under #1687, and the audit also fixed a third instance in the unused layers::Attention.
The fix moved the family from unusable to partially working, but greedy output is still degraded. Facts are retrieved and then the continuation falls apart, so something else is still wrong.
Current Behavior
Checkpoint: mlx-community/Llama-3_3-Nemotron-Super-49B-v1-mlx-4bit, 26 GB on disk, model_type: nemotron-nas, architectures: ["DeciLMForCausalLM"], 80 block_configs entries. mlxcel routes it to NemotronNAS and reports "80 layers (49 with attention, 31 no-op)". Fetch with ./target/release/mlxcel download mlx-community/Llama-3_3-Nemotron-Super-49B-v1-mlx-4bit, which lands in ~/.cache/mlxcel/models/mlx-community/Llama-3_3-Nemotron-Super-49B-v1-mlx-4bit.
What b429947c changed, at --temp 0 --no-chat-template. Before: "The capital of France is" continued "the the the the the the the the". After: "The capital of France is Paris, the most visited by millions of tourists every year by millions of people every year, but there are attracted to its beauty".
What is still wrong, same flags:
- "The first five prime numbers are" continues "1, 1, 2, 3, 3, 5, 7", which is wrong.
- "Paris is the capital of" continues "France and most visited city of France and one of France. It is known as the most visited city", which starts correct and then repeats.
- The chat-template path is worse: "List the first five prime numbers." at
--temp 0with the template produces "Here are the numbers from 1. \n\n1. 1 to 2. **1 to 2.3.3.3.5\n\n1.5.5.7.9.7.9.11.11.13".
Two prompts do come out right: "2 + 2 =" gives "2 = 4" and "Water boils at a temperature of" gives "100C". The model retrieves facts and then degrades rather than being uniformly broken.
Why it was not localized further
The mlx-lm reference cannot load this checkpoint in the repo's baseline venv. .venv-mlxlm/bin/python -m mlx_lm generate --model <path> --prompt "The capital of France is" --max-tokens 25 --temp 0 fails inside transformers with AttributeError: 'PreTrainedConfig' object has no attribute 'max_position_embeddings', raised from transformers/integrations/heterogeneity/configuration_utils.py. Adding --trust-remote-code does not change it. Without a working reference there is no per-layer activation to diff against, which is how #1687 was localized.
Already checked and matching the reference
Read against mlx_lm/models/nemotron-nas.py:
ffn_mult_to_intermediate_sizecomputesint(2 * ffn_mult * n_embd / 3)rounded up to a multiple of 256, matching_ffn_mult_to_intermediate_sizeand_find_multiple.- Config parsing covers every
block_configsfield the checkpoint carries:no_op,replace_with_linear,sparsify,n_heads_in_group,window_length,num_sink_tokens,unshifted_sink,use_prefill_window_in_sink_attention,ffn_mult. - Note that
src/models/nemotron_nas.rs:573computesffn_mult_to_intermediate_sizeinto a discarded_hidden_dim, so the MLP widths actually come from the loaded weight shapes rather than from that helper.
Candidates to investigate
These are candidates, not conclusions:
- The
no_opandreplace_with_linearbranches for both attention and FFN atsrc/models/nemotron_nas.rs:473-575, since 31 of 80 layers take a no-op attention path. - The sink-attention and
window_lengthfields, which the config parses but whose actual use in the forward path should be confirmed. - Getting a reference running at all, either by pinning a
transformersversion that loads this config or by using a smallernemotron-nascheckpoint.
Acceptance Criteria
- A working mlx-lm reference on some
nemotron-nascheckpoint, so outputs can be diffed per layer rather than eyeballed. - Greedy completions at
--temp 0 --no-chat-templatestay coherent past 25 tokens on the prompts listed above. - The chat-template path produces a sane answer to a simple arithmetic or list prompt.
- A unit test covering whatever the defect turns out to be.
Verification
cargo build --release --features metal,accelerate
./target/release/mlxcel generate -m ~/.cache/mlxcel/models/mlx-community/Llama-3_3-Nemotron-Super-49B-v1-mlx-4bit -p "The first five prime numbers are" -n 25 --temp 0 --no-chat-template
./target/release/mlxcel generate -m ~/.cache/mlxcel/models/mlx-community/Llama-3_3-Nemotron-Super-49B-v1-mlx-4bit -p "List the first five prime numbers." -n 40 --temp 0
cargo test --release nemotron_nas
cargo clippy --all-targets -- -D warningsTechnical Considerations
A benchmark row for this checkpoint is being added to benchmarks/metal_m5max_2026-09-06.csv so the family has a tracked number. Read that number as provisional while the output is still wrong: it measures the speed of a computation that is not yet producing correct text.
Environment: mlxcel 0.7.0-beta.1 on branch bench/0.7.0-refresh, MacBook Pro M5 Max 128GB, macOS 26.6.2, MLX pin 9a795735, built with cargo build --release --features metal,accelerate.
Related: #1687 (the recurrent_gemma instance of the same RoPE defect), commit b429947c.
Activity
bench: track nemotron-nas and recurrent_gemma, fix a bad figure
fix(nemotron_nas): resolve llama3 rope scaling through the shared reader
fix(nemotron_nas): read eos_token_id from the config
inureyes commentedon Sep 8, 2026
Fixed in
a5386990. The remaining defect was the rope scaling scheme, not the rotation order.Cause
NASAttentionread onlyfactorout ofrope_scalingand passed1 / factortofast_ropeas a uniform position scale. That is thelinearscheme. This checkpoint declaresrope_type: "llama3", whose rescaling is piecewise in the frequency: long wavelengths divide byfactor, short ones are left alone, and the band between interpolates. Compressing every frequency by 8x is what left the continuation falling apart after the facts were retrieved.llama3.rs:350already stated the distinction ("Position scale handed tofast_rope.1.0for every checkpoint except arope_scalingblock of typelinear, where it is1 / factor"), andrope_utils::RopeScalingKind::resolvealready built the llama3 table for every other decoder. This block was the one that had its own partial copy of the config struct and used it wrongly.Verification
Checkpoint
mlx-community/Llama-3_3-Nemotron-Super-49B-v1-mlx-4bit,--temp 0 --no-chat-template, atadc4f773with the binary confirmed in sync.The three cases this issue reported as wrong:
Five further short prompts (primary colors, capital of France, Korean translation, prime list, one sentence about the ocean) all answer correctly, and a needle placed at the head of a 1300-token context is recalled exactly, which is the positional behavior this issue asked about.
Two notes
b429947cwas a real fix and insufficient on its own. Before it the model emitted a single "I" and blank lines; after it a fluent but incoherent loop; only with the rope scheme corrected does it answer. Both defects had to go.This family had never run against a real checkpoint. No
nemotron-nasmodel existed in either machine's benchmark store: the only one present was reachable through a symlink into an excluded download cache, and itsmodel_typeis spellednemotron-naswhile the module isnemotron_nas, so two independent scans reported the family absent. The checkpoint is now in the standard store.Left open elsewhere
Two observations that are not this issue and are not regressions from it. With the chat template and no system message the model answers correctly and then stops after about 15 tokens; the raw path generates 90 without trouble, so this looks like the template emitting an empty system message to a reasoning model that expects "detailed thinking on/off" there, and
mlxcel generatehas no flag to supply one. Separately,eos_token_idsreturned a hardcoded2where the checkpoint declares[128001, 128008, 128009]; that is fixed in0a032422but changed no output here, because token 2 never appeared in these runs.mlx-lm still cannot serve as a reference for this checkpoint: it raises
AttributeError: 'PreTrainedConfig' object has no attribute 'max_position_embeddings'fromstandardize_rope_paramsbefore the model loads, which is a transformers regression on heterogeneous configs rather than anything in the checkpoint, whosemax_position_embeddingsis present and set to 131072.inureyes commentedon Sep 8, 2026
Follow-up: the verification in the previous comment did not reach the code it was verifying, and a longer run now confirms the fix.
rope_scalingon this checkpoint declaresoriginal_max_position_embeddings: 8192. The llama3 rescaling is piecewise in the frequency, and the low-frequency band it changes only carries positions past that length, so the 1300-token needle used above exercised the head layout and the GQA path and never touched the scheme the fix replaced. It passed for reasons unrelated to the change.Re-run at 14737 tokens, 1.8x the boundary, needle at position ~10,
--temp 0, binary confirmed in sync with the tree:VERMILION-8842.src/models/nemotron_nas.rsreverted toa5386990~1and rebuilt: answers "You've provided" and never recovers the code.So the test discriminates, which the short one did not. Both arms share a binary lineage and differ only in that file.
Worth stating plainly because it nearly went unrecorded: a passing check at a length that cannot reach the changed branch is not evidence, and it looks exactly like evidence. The boundary is in the config, so the length a rope-scaling change has to be tested at is knowable before the test is written rather than after.