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

mlx_lm/models/internlm3.py

@@ -53,22 +53,32 @@ def __init__(
         traditional: bool = False,
         base: float = 10000,
         scale: float = 1.0,
+        rope_type: str = "default",
+        ntk_factor: float = 1.0,
     ):
         super().__init__()
         self.max_position_embeddings = max_position_embeddings
         self.original_base = base
         self.dims = dims
         self.traditional = traditional
         self.scale = scale
+        self.rope_type = rope_type
+        self.ntk_factor = ntk_factor
 
     def extra_repr(self):
-        return f"{self.dims}, traditional={self.traditional}, max_position_embeddings={self.max_position_embeddings}, scaling_factor={self.scaling_factor}"
+        return f"{self.dims}, traditional={self.traditional}, max_position_embeddings={self.max_position_embeddings}, rope_type={self.rope_type}"
 
     def __call__(self, x, offset: int = 0):
-        seq_len = x.shape[1] + offset
-        if seq_len > self.max_position_embeddings:
+        # x is [B, n_heads, L, head_dim] by the time it reaches rope() -- the
+        # sequence axis is 2, not 1.
+        seq_len = x.shape[2] + offset
+        # Only "dynamic" NTK recomputes the base, and only off the config
+        # factor -- self.scale is the *position* scale (only nonzero for
+        # "linear"), a different knob from the NTK base-growth factor.
+        if self.rope_type == "dynamic" and seq_len > self.max_position_embeddings:
             base = self.original_base * (
-                (self.scale * seq_len / self.max_position_embeddings) - (self.scale - 1)
+                (self.ntk_factor * seq_len / self.max_position_embeddings)
+                - (self.ntk_factor - 1)
             ) ** (self.dims / (self.dims - 2))
         else:
             base = self.original_base
@@ -101,19 +111,23 @@ def __init__(self, args: ModelArgs):
         self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=qkv_bias)
         self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=qkv_bias)
 
-        rope_scale = (
-            1 / args.rope_scaling["factor"]
-            if args.rope_scaling is not None
-            and args.rope_scaling["rope_type"] == "linear"
-            else 2.0
+        rope_type = (
+            args.rope_scaling["rope_type"] if args.rope_scaling is not None else None
         )
+        # Position scale: only "linear" mode scales positions (1/factor).
+        # "dynamic" mode leaves positions alone and instead grows the base
+        # (see DynamicNTKScalingRoPE.__call__).
+        rope_scale = 1 / args.rope_scaling["factor"] if rope_type == "linear" else 1.0
+        ntk_factor = args.rope_scaling["factor"] if rope_type == "dynamic" else 1.0
 
         self.rope = DynamicNTKScalingRoPE(
             head_dim,
             max_position_embeddings=args.max_position_embeddings,
             traditional=args.rope_traditional,
             base=args.rope_theta,
             scale=rope_scale,
+            rope_type=rope_type or "default",
+            ntk_factor=ntk_factor,
         )
 
     def __call__(

tests/test_models.py

@@ -1930,6 +1930,81 @@ def test_internlm3(self):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_internlm3_dynamic_ntk_rope(self):
+        # Regression test for #1545: dynamic-NTK RoPE conflated the position
+        # scale with the NTK base-growth factor, defaulted the position scale
+        # to 2.0 whenever scaling wasn't "linear", and read the sequence
+        # length off the wrong axis (post-transpose, seq is axis 2, not 1).
+        from mlx_lm.models import internlm3
+
+        def make_attn(rope_scaling):
+            args = internlm3.ModelArgs(
+                model_type="internlm3",
+                hidden_size=64,
+                num_hidden_layers=1,
+                intermediate_size=64,
+                num_attention_heads=4,
+                rms_norm_eps=1e-5,
+                vocab_size=100,
+                max_position_embeddings=32768,
+                rope_scaling=rope_scaling,
+            )
+            return internlm3.Attention(args)
+
+        # No scaling config at all: positions must be left unscaled (was 2.0).
+        self.assertEqual(make_attn(None).rope.scale, 1.0)
+
+        # Dynamic: positions unscaled, and the config factor reaches the NTK
+        # base formula as `ntk_factor` (was silently discarded / hard-coded).
+        dyn = make_attn({"factor": 6.0, "rope_type": "dynamic"})
+        self.assertEqual(dyn.rope.scale, 1.0)
+        self.assertEqual(dyn.rope.ntk_factor, 6.0)
+
+        # Linear: only the position scale (1/factor) applies.
+        lin = make_attn({"factor": 4.0, "rope_type": "linear"})
+        self.assertEqual(lin.rope.scale, 0.25)
+
+        # Sequence length must come from the post-transpose axis (2), so a
+        # long prompt actually triggers the dynamic NTK recompute even though
+        # the head-count axis (1) is small. Detect the recompute by its
+        # effect: a rotated non-zero vector changes when the base changes.
+        mx.random.seed(0)
+        x = mx.random.normal((1, 4, 100, 8))  # n_heads=4 < 64, real seq_len=100 > 64
+
+        def make_rope(max_position_embeddings, rope_type, ntk_factor):
+            return internlm3.DynamicNTKScalingRoPE(
+                dims=8,
+                max_position_embeddings=max_position_embeddings,
+                base=10000.0,
+                scale=1.0,
+                rope_type=rope_type,
+                ntk_factor=ntk_factor,
+            )
+
+        recomputed = make_rope(64, "dynamic", 6.0)(x)  # 100 > 64 -> base grows
+        static = make_rope(1_000_000, "dynamic", 6.0)(x)  # 100 < 1e6 -> base unchanged
+        self.assertFalse(mx.array_equal(recomputed, static))
+
+        # Linear mode must never rewrite the base, even past
+        # max_position_embeddings -- only the position scale applies there.
+        lin_short = internlm3.DynamicNTKScalingRoPE(
+            dims=8,
+            max_position_embeddings=1_000_000,
+            base=10000.0,
+            scale=0.25,
+            rope_type="linear",
+            ntk_factor=1.0,
+        )(x)
+        lin_long = internlm3.DynamicNTKScalingRoPE(
+            dims=8,
+            max_position_embeddings=8,  # real seq_len=100 > 8
+            base=10000.0,
+            scale=0.25,
+            rope_type="linear",
+            ntk_factor=1.0,
+        )(x)
+        self.assertTrue(mx.array_equal(lin_short, lin_long))
+
     def test_iquestloopcoder(self):
         from mlx_lm.models import iquestloopcoder