fix(internlm2): honor rope_scaling (dynamic NTK, linear) and rope_traditional #1320

Summary

internlm2 ignores rope_scaling entirely and ignores rope_traditional after parsing it. The public checkpoint mlx-community/internlm2_5-7b-chat-4bit ships rope_scaling: {"type": "dynamic", "factor": 2.0} with max_position_embeddings: 32768, so the rotary base must grow once a request passes 32768 tokens; today it stays at rope_theta and long contexts run with the wrong rotation frequencies. A "linear" entry would be ignored too. This change wires the family to the shared DynamicNtkRope helper (src/models/dynamic_ntk_rope.rs) introduced by #1324.

Below 32768 tokens the current output is already correct (positions unscaled, base unchanged), which is why this is medium rather than high.

Current behavior

src/models/internlm2.rs:28-58: ModelArgs has no rope_scaling field, so the key is dropped at deserialization. rope_traditional is parsed at line 50-51 but never read.

src/models/internlm2.rs:157-158:

q = mlxcel_core::fast_rope(&q, self.rope_dims, false, self.rope_base, 1.0, offset);
k = mlxcel_core::fast_rope(&k, self.rope_dims, false, self.rope_base, 1.0, offset);

traditional is hard-coded false, scale is hard-coded 1.0, and rope_base is the constant args.rope_theta set at src/models/internlm2.rs:207.

Checkpoint config (models/internlm2-7b-4bit/config.json): rope_theta: 1000000, max_position_embeddings: 32768, rope_scaling: {"type": "dynamic", "factor": 2.0}, num_attention_heads: 32, hidden_size: 4096 (so head_dim = 128). Note the key is spelled type, not rope_type.

Expected behavior

Same schedule as internlm3 (see #1324), with dims = head_dim = 128, base = rope_theta, max_pos = max_position_embeddings:

• no rope_scaling, or type/rope_type in {"default", absent}: fast_rope(x, dims, rope_traditional, base, 1.0, offset).
• "linear", factor = f: fast_rope(x, dims, rope_traditional, base, 1.0 / f, offset).
• "dynamic", factor = f: seq_len = max(L + offset, max_pos); base_eff = base * ((f * seq_len / max_pos) - (f - 1)) ^ (dims / (dims - 2)); fast_rope(x, dims, rope_traditional, base_eff, 1.0, offset).
• any other type, or linear/dynamic without factor: load error.

Worked values for the checkpoint (base = 1e6, f = 2.0, dims = 128, max_pos = 32768):

rope_traditional selects the rotation layout: false rotates the two halves of the head (x[..., :d/2], x[..., d/2:]), true rotates adjacent pairs. The checkpoint's own modeling_internlm2.py uses the half-split layout, so false (the serde default) remains the value for every public checkpoint; the flag just has to reach the call.

Implementation plan

1. src/models/internlm2.rs ModelArgs: add

   #[serde(default)]
   pub rope_scaling: Option<crate::models::config::RopeScaling>,

   config::RopeScaling (src/models/config.rs:309) already accepts both type and rope_type.

2. Replace the rope_dims: i32, rope_base: f32 fields of Attention (src/models/internlm2.rs:103-104) with rope: crate::models::dynamic_ntk_rope::DynamicNtkRope. In Attention::from_weights build it with DynamicNtkRope::from_scaling(head_dim, args.rope_theta, args.rope_traditional, args.max_position_embeddings, args.rope_scaling.as_ref())?.

3. In Attention::forward replace the two hard-coded fast_rope calls with self.rope.apply(&q, offset, l + offset) and self.rope.apply(&k, offset, l + offset) (offset = cache.offset, l = shape[1]).

4. Extend the // Used by: comment on DynamicNtkRope to list InternLM2.

Validation

(a) Unit tests, src/models/internlm2_tests.rs (new, attached with #[cfg(test)] #[path = "internlm2_tests.rs"] mod tests;):

• internlm2_config_parses_type_spelled_rope_scaling: deserialize the checkpoint config excerpt above into ModelArgs; rope_scaling.scaling_type == Some("dynamic"), factor == Some(2.0).
• internlm2_rope_is_unit_scale_dynamic_for_checkpoint_config: the built DynamicNtkRope has position_scale() == 1.0, base_for(56) == 1_000_000.0, and base_for(65536) within 1e-3 relative of 3_052_773.67.
• internlm2_rope_linear_uses_inverse_factor: rope_scaling: {"type": "linear", "factor": 4.0} gives position_scale() == 0.25.
• internlm2_rope_traditional_reaches_the_helper: rope_traditional: true yields rope.traditional == true.
• internlm2_rejects_unknown_rope_type: {"type": "yarn", "factor": 2.0} makes Attention::from_weights (or a thin ModelArgs::rope() accessor) return Err.

(b) Real checkpoint: mlx-community/internlm2_5-7b-chat-4bit (local models/internlm2-7b-4bit).

./target/release/mlxcel generate -m models/internlm2-7b-4bit -p "Write a short paragraph explaining why the sky is blue." -n 64

Acceptance: token-exact with the output produced before this change on a prompt of at least 64 tokens with -n 64 (below 32768 tokens the schedule is unchanged, so any drift is a regression). Then a long-context smoke: a prompt of about 33000 tokens (for example a repeated paragraph) with -n 32 must produce finite, fluent text; log the base_eff once at the first forward past max_position_embeddings (debug level) and confirm it matches the table within 1e-3 relative.

Acceptance criteria

• rope_scaling with type or rope_type is parsed; dynamic rescales the base past 32768 tokens exactly as in the table; linear scales positions by 1 / factor
• rope_traditional reaches fast_rope
• Output on prompts shorter than max_position_embeddings is token-identical to the current binary
• cargo test --workspace --profile test-fast --features metal,accelerate passes
• cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass

Dependencies

Blocked by #1324 (provides src/models/dynamic_ntk_rope.rs).