diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d44dd02..7bc8b831 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/forgelm/model.py b/forgelm/model.py index de6c7577..4d96ea14 100644 --- a/forgelm/model.py +++ b/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 diff --git a/tests/test_model.py b/tests/test_model.py index 771a671e..bc9f0530 100644 --- a/tests/test_model.py +++ b/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: