[robertlangdonn]

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

Summary

Twin of #1552's sibling PR for internlm3 (#1545 / this repo's #1553): while fixing internlm3's DynamicNTKScalingRoPE, I found internlm2.py has an identical copy-pasted class with the same four bugs. No separate issue was filed for internlm2 specifically, so this PR describes and fixes it directly, scoped to that one file, rather than bundling it into #1553.

Four related bugs in DynamicNTKScalingRoPE and its construction in Attention.__init__ (config key is "type" here, vs "rope_type" in internlm3 — otherwise identical):

1. Position scale defaulted to 2.0 whenever scaling wasn't "linear" — including no rope_scaling at all — doubling every position index fed into mx.fast.rope. Only "linear" should scale positions (1/factor); 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.
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

Same shape as #1553: split scale (position scale) from a new ntk_factor (config factor, used only in the base-growth formula), thread rope_type through to gate the recompute, and read seq_len from x.shape[2].

Total change: +22/−9 in mlx_lm/models/internlm2.py.

Verification

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

• No rope_scaling → scale was 2.0, now 1.0.
• {"factor": 6.0, "type": "dynamic"} → scale was 2.0 (factor discarded), now scale=1.0 / ntk_factor=6.0.
• A tensor shaped (1, 4, 100, 8) (4 heads, real seq_len=100, max_position_embeddings=64) — pre-fix, seq_len read as 4 and the recompute never fired; post-fix it correctly fires on 100.
• "linear" mode past max_position_embeddings — pre-fix, base was rewritten anyway; post-fix, output is bit-identical to a run where max_position_embeddings is never exceeded.

Test plan

• New regression test test_internlm2_dynamic_ntk_rope, mirroring test_internlm3_dynamic_ntk_rope from Fix internlm3 dynamic-NTK RoPE: wrong seq axis, scale/factor conflation #1553. 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).
• pre-commit run --files mlx_lm/models/internlm2.py tests/test_models.py — clean.

[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.