[robertlangdonn]

Fix internlm3 dynamic-NTK RoPE: wrong seq axis, scale/factor conflation

Summary

Fixes #1545. Credit to @magicnight for the full root-cause diagnosis in the issue (checked against InternLM's reference modeling_internlm3.py) — this PR independently verifies each of the four defects against the real code and adds regression tests.

Four related bugs in internlm3.py's DynamicNTKScalingRoPE and its construction in Attention.__init__:

1. Position scale defaulted to 2.0 whenever scaling wasn't "linear" — including the common case of no rope_scaling at all — doubling every position index fed into mx.fast.rope. Only "linear" should scale positions (1/factor); the default should be 1.0.
2. The dynamic-NTK base formula reused that same wrong 2.0 instead of the config's factor, so a real config (e.g. factor=6.0) never reached the base recomputation.
3. seq_len was read from x.shape[1] — the head-count axis post-transpose ([B, n_heads, L, head_dim]) — not the sequence axis (shape[2]). The dynamic recompute could never fire at prefill regardless of prompt length (head count is typically far below max_position_embeddings).
4. The base recompute wasn't gated on rope_type — "linear" scaling also rewrote the base once past max_position_embeddings, when it should stay static and only the position scale should apply.

Fix

• Split the single conflated scale into scale (position scale, mx.fast.rope's scale=) and a new ntk_factor (the config factor, used only in the base-growth formula).
• rope_type is now threaded into DynamicNTKScalingRoPE and gates the base recompute (if self.rope_type == "dynamic" and seq_len > self.max_position_embeddings).
• seq_len = x.shape[2] + offset instead of x.shape[1].

Total change: +23/−9 in mlx_lm/models/internlm3.py.

Verification

Falsified all four defects against the pre-fix code directly (no checkpoint needed — these are pure position-math bugs reproducible with tiny synthetic tensors), then verified each is fixed:

• No rope_scaling → self.rope.scale was 2.0, now 1.0.
• {"factor": 6.0, "rope_type": "dynamic"} → self.rope.scale was 2.0 (factor discarded), now scale=1.0 / ntk_factor=6.0.
• A tensor shaped (1, 4, 100, 8) (4 heads, 100 real sequence positions, max_position_embeddings=64) — pre-fix, seq_len read as 4 and the dynamic recompute never fired; post-fix, seq_len=100 correctly triggers it.
• "linear" mode with real seq_len=50 > max_position_embeddings=8 — pre-fix, the base was rewritten anyway; post-fix, output is identical to a run where max_position_embeddings is never exceeded (base provably stayed static).
• End-to-end: full internlm3.Model (tiny synthetic config, rope_scaling={"factor": 4.0, "rope_type": "dynamic"}), prefill past max_position_embeddings then several decode steps — runs cleanly.

Test plan

• New regression test test_internlm3_dynamic_ntk_rope — covers all four defects (scale defaults, factor threading, correct seq axis via an actual-effect comparison rather than just checking it doesn't crash, and linear-mode base staying static). Confirmed it fails on pre-fix code (2.0 != 1.0) and passes post-fix.
• Full suite: python -m unittest discover tests/ — pre-existing unrelated import errors only (datasets/lm_eval/requests, optional deps not installed locally).
• pre-commit run --files mlx_lm/models/internlm3.py tests/test_models.py — clean.

Note

internlm2.py has an identical copy-pasted DynamicNTKScalingRoPE with the same four bugs. Filing that as a separate PR against that file rather than bundling it here, to keep this one scoped to the reported issue.

[zcbenz]

Thanks for the PR but at the moment the number of PRs is way beyond our capacity to review so I'm closing the non-essential ones so we can actually work on this repo.