fix(model): apply training.rope_scaling through the config on transformers 5

CHANGELOG.md

@@ -68,6 +68,13 @@ All notable changes to ForgeLM are documented here.
 
 ### Fixed
 
+- `training.rope_scaling` works on transformers 5 again. The setting was passed to `from_pretrained(rope_scaling=...)`,
+  which no longer works: the documented `{type: linear, factor: 4}` raised `KeyError: 'rope_type'` while the model was
+  built, and `{rope_type: ..., factor: ...}` replaced the whole parameter dict so `rope_theta` was lost (`TypeError`).
+  The requested keys are now merged over the model's own `rope_parameters` (the legacy key `type` is read as
+  `rope_type`), applied to the text config of a multimodal checkpoint, and the finished config is loaded. The
+  `sliding_window_attention` override travels in the same config object. Verified on transformers 5.13.0 for
+  `linear`, `dynamic` and `yarn` on a tiny local Llama; not exercised on the 5.3-5.9 range the floor still admits (#103).
 - Fix the KTO bullet on the Alignment overview user-manual page (English and Turkish): it named the
   fields `(prompt, response, good?)`, but the KTO loader reads `prompt`, `completion` and a boolean
   `label` (#965).

forgelm/model.py

@@ -420,6 +420,35 @@ def _device_map_for(config: Any, is_distributed: bool):
     return None  # CPU-only
 
 
+def _config_with_rope_scaling(
+    name_or_path: str, rope_scaling: dict, *, trust_remote_code: bool, revision: Optional[str] = None
+) -> Any:
+    """Load the model's config and apply ``training.rope_scaling`` to it.
+
+    transformers 5 keeps the RoPE settings in ``config.rope_parameters`` (which also carries ``rope_theta``) and
+    expects the key ``rope_type``.  Passing ``rope_scaling=`` to ``from_pretrained`` no longer works: the documented
+    ``{type: linear, factor: 4}`` raises ``KeyError: 'rope_type'`` while the model is built, and
+    ``{rope_type: ..., factor: ...}`` replaces the whole dict so ``rope_theta`` is lost (``TypeError`` in the
+    frequency computation).  So the requested keys are merged *over* the model's own ``rope_parameters``, the legacy
+    key ``type`` is read as ``rope_type``, and the finished config object is what gets loaded.
+
+    For a multimodal checkpoint the RoPE settings live in the text sub-config, which is where they are applied.
+    A transformers without ``rope_parameters`` gets the legacy ``rope_scaling`` attribute instead.
+    """
+    from transformers import AutoConfig
+
+    model_config = AutoConfig.from_pretrained(name_or_path, trust_remote_code=trust_remote_code, revision=revision)
+    target = model_config.get_text_config() if hasattr(model_config, "get_text_config") else model_config
+    requested = dict(rope_scaling)
+    if "type" in requested and "rope_type" not in requested:
+        requested["rope_type"] = requested.pop("type")
+    if hasattr(target, "rope_parameters"):
+        target.rope_parameters = {**(getattr(target, "rope_parameters", None) or {}), **requested}
+    else:
+        target.rope_scaling = requested
+    return model_config
+
+
 def _build_model_kwargs(config: Any, trust_remote_code: bool, revision: Optional[str] = None) -> dict:
     """Assemble from_pretrained kwargs (device_map, BnB, RoPE, sliding window, revision).
 
@@ -461,12 +490,19 @@ def _build_model_kwargs(config: Any, trust_remote_code: bool, revision: Optional
         )
 
     rope_scaling = config.training.rope_scaling
+    sliding_window = config.training.sliding_window_attention
     if rope_scaling:
         logger.info("RoPE scaling enabled: %s", rope_scaling)
-        kwargs["rope_scaling"] = rope_scaling
-
-    sliding_window = config.training.sliding_window_attention
-    if sliding_window:
+        model_config = _config_with_rope_scaling(
+            config.model.name_or_path, rope_scaling, trust_remote_code=trust_remote_code, revision=revision
+        )
+        if sliding_window:
+            # Once ``config=`` is a finished object, from_pretrained hands unrecognised keyword arguments to the
+            # model constructor instead of the config, so the override is applied to the config itself.
+            logger.info("Sliding window attention override: %d tokens", sliding_window)
+            model_config.sliding_window = sliding_window
+        kwargs["config"] = model_config
+    elif sliding_window:
         logger.info("Sliding window attention override: %d tokens", sliding_window)
         kwargs["sliding_window"] = sliding_window
 

tests/test_model.py

@@ -115,15 +115,105 @@ def test_quantization_config_added_with_cuda(self):
         assert "quantization_config" in kwargs
         fake_bnb.assert_called_once()
 
-    def test_rope_and_sliding_window_passed_through(self):
+    def test_sliding_window_alone_is_a_keyword_argument(self):
         cfg = _cfg(
             model={"name_or_path": "org/model", "load_in_4bit": False},
-            training={"rope_scaling": {"type": "linear", "factor": 2.0}, "sliding_window_attention": 4096},
+            training={"sliding_window_attention": 4096},
         )
         with patch("torch.cuda.is_available", return_value=False):
             kwargs = model_mod._build_model_kwargs(cfg, trust_remote_code=False)
-        assert kwargs["rope_scaling"] == {"type": "linear", "factor": 2.0}
-        assert kwargs["sliding_window"] == 4096
+        assert kwargs["sliding_window"] == 4096 and "config" not in kwargs
+
+    def test_rope_scaling_travels_in_a_config_object_not_a_keyword(self):
+        # transformers 5 no longer accepts `rope_scaling=` on from_pretrained (#103): the settings go into the
+        # config, together with the sliding-window override, which would otherwise reach the model constructor.
+        model_config = SimpleNamespace(rope_parameters={"rope_theta": 10000.0, "rope_type": "default"})
+        fake_transformers = MagicMock()
+        fake_transformers.AutoConfig.from_pretrained.return_value = model_config
+        cfg = _cfg(
+            model={"name_or_path": "org/model", "load_in_4bit": False},
+            training={"rope_scaling": {"type": "linear", "factor": 2.0}, "sliding_window_attention": 4096},
+        )
+        with (
+            patch("torch.cuda.is_available", return_value=False),
+            patch.dict("sys.modules", {"transformers": fake_transformers}),
+        ):
+            kwargs = model_mod._build_model_kwargs(cfg, trust_remote_code=False, revision="abc123")
+        assert "rope_scaling" not in kwargs and "sliding_window" not in kwargs
+        assert kwargs["config"] is model_config
+        assert model_config.rope_parameters == {"rope_theta": 10000.0, "rope_type": "linear", "factor": 2.0}
+        assert model_config.sliding_window == 4096
+        fake_transformers.AutoConfig.from_pretrained.assert_called_once_with(
+            "org/model", trust_remote_code=False, revision="abc123"
+        )
+
+
+class TestConfigWithRopeScaling:
+    """Against a real (tiny, local, CPU) Llama config: the merged settings must build and run a model (#103)."""
+
+    @staticmethod
+    def _tiny_model_dir(tmp_path):
+        from transformers import LlamaConfig
+
+        LlamaConfig(
+            vocab_size=64,
+            hidden_size=32,
+            intermediate_size=64,
+            num_hidden_layers=1,
+            num_attention_heads=2,
+            num_key_value_heads=2,
+            max_position_embeddings=128,
+        ).save_pretrained(tmp_path)
+        return str(tmp_path)
+
+    @pytest.mark.parametrize(
+        "requested, expected_type",
+        [
+            ({"type": "linear", "factor": 4.0}, "linear"),  # the key every documented example and the wizard use
+            ({"rope_type": "linear", "factor": 4.0}, "linear"),
+            ({"type": "dynamic", "factor": 2.0}, "dynamic"),
+            ({"rope_type": "yarn", "factor": 2.0, "original_max_position_embeddings": 128}, "yarn"),
+        ],
+    )
+    def test_the_merged_config_builds_a_model_that_runs(self, tmp_path, requested, expected_type):
+        import torch
+        from transformers import AutoModelForCausalLM
+
+        path = self._tiny_model_dir(tmp_path)
+        model_config = model_mod._config_with_rope_scaling(path, requested, trust_remote_code=False)
+        assert model_config.rope_parameters["rope_type"] == expected_type
+        assert model_config.rope_parameters["rope_theta"] == 10000.0  # kept from the model, not lost
+        assert "type" not in model_config.rope_parameters
+        logits = AutoModelForCausalLM.from_config(model_config)(input_ids=torch.tensor([[1, 2, 3, 4]])).logits
+        assert tuple(logits.shape) == (1, 4, 64)
+
+    def test_requested_rope_theta_overrides_the_models(self, tmp_path):
+        path = self._tiny_model_dir(tmp_path)
+        model_config = model_mod._config_with_rope_scaling(
+            path, {"type": "linear", "factor": 2.0, "rope_theta": 500000.0}, trust_remote_code=False
+        )
+        assert model_config.rope_parameters["rope_theta"] == 500000.0
+
+    def test_a_transformers_without_rope_parameters_gets_the_legacy_attribute(self):
+        legacy = SimpleNamespace()
+        fake_transformers = MagicMock()
+        fake_transformers.AutoConfig.from_pretrained.return_value = legacy
+        with patch.dict("sys.modules", {"transformers": fake_transformers}):
+            out = model_mod._config_with_rope_scaling(
+                "org/model", {"type": "linear", "factor": 2.0}, trust_remote_code=False
+            )
+        assert out is legacy and legacy.rope_scaling == {"rope_type": "linear", "factor": 2.0}
+
+    def test_multimodal_settings_land_in_the_text_config(self):
+        text = SimpleNamespace(rope_parameters={"rope_theta": 1.0, "rope_type": "default"})
+        outer = SimpleNamespace(get_text_config=lambda: text)
+        fake_transformers = MagicMock()
+        fake_transformers.AutoConfig.from_pretrained.return_value = outer
+        with patch.dict("sys.modules", {"transformers": fake_transformers}):
+            out = model_mod._config_with_rope_scaling(
+                "org/vlm", {"type": "linear", "factor": 2.0}, trust_remote_code=False
+            )
+        assert out is outer and text.rope_parameters["rope_type"] == "linear"
 
 
 class TestLoadTokenizer: