[serve][LLM] Fix HF config loading for models with custom rope_scaling (ray-project#62464)

Signed-off-by: Kourosh Hakhamaneshi <kourosh@anyscale.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: 2 people

python/ray/llm/_internal/serve/core/configs/llm_config.py

@@ -256,50 +256,59 @@ def validate_server_cls(cls, value):
     _engine_config: EngineConfigType = PrivateAttr(None)
     _callback_instance: Optional[CallbackBase] = PrivateAttr(None)
 
-    def _infer_supports_vision(self, model_id_or_path: str) -> None:
-        """Called in llm node initializer together with other transformers calls. It
-        loads the model config from huggingface and sets the supports_vision
-        attribute based on whether the config has `vision_config`. All LVM models has
-        `vision_config` setup.
+    def _load_hf_config(self, model_id_or_path: str, trust_remote_code: bool = False):
+        """Load the HuggingFace config for a model.
+
+        Uses AutoConfig which loads the model-specific config class (e.g.
+        DeepseekV3Config) instead of the generic PretrainedConfig.  The generic
+        base class can fail for models whose config.json contains fields (like
+        ``rope_scaling``) that require model-specific post-init logic.
         """
         try:
-            hf_config = transformers.PretrainedConfig.from_pretrained(model_id_or_path)
-            self._supports_vision = hasattr(hf_config, "vision_config")
+            return transformers.AutoConfig.from_pretrained(
+                model_id_or_path, trust_remote_code=trust_remote_code
+            )
         except Exception as e:
             raise ValueError(
-                f"Failed to load Hugging Face config for model_id='{model_id_or_path}'.\
-                        Ensure `model_id` is a valid Hugging Face repo or a local path that \
-                        contains a valid `config.json` file. "
-                f"Original error: {repr(e)}"
+                f"Failed to load Hugging Face config for "
+                f"model_id='{model_id_or_path}'. Ensure `model_id` is a valid "
+                f"Hugging Face repo or a local path that contains a valid "
+                f"`config.json` file. Original error: {repr(e)}"
             ) from e
 
+    def _infer_supports_vision(
+        self, model_id_or_path: str, trust_remote_code: bool = False
+    ) -> None:
+        """Called in llm node initializer together with other transformers calls. It
+        loads the model config from huggingface and sets the supports_vision
+        attribute based on whether the config has `vision_config`. All LVM models has
+        `vision_config` setup.
+        """
+        hf_config = self._load_hf_config(
+            model_id_or_path, trust_remote_code=trust_remote_code
+        )
+        self._supports_vision = hasattr(hf_config, "vision_config")
+
     def _set_model_architecture(
         self,
         model_id_or_path: Optional[str] = None,
         model_architecture: Optional[str] = None,
+        trust_remote_code: bool = False,
     ) -> None:
         """Called in llm node initializer together with other transformers calls. It
         loads the model config from huggingface and sets the model_architecture
         attribute based on whether the config has `architectures`.
         """
         if model_id_or_path:
-            try:
-                hf_config = transformers.PretrainedConfig.from_pretrained(
-                    model_id_or_path
-                )
-                if (
-                    hf_config
-                    and hasattr(hf_config, "architectures")
-                    and hf_config.architectures
-                ):
-                    self._model_architecture = hf_config.architectures[0]
-            except Exception as e:
-                raise ValueError(
-                    f"Failed to load Hugging Face config for model_id='{model_id_or_path}'.\
-                        Ensure `model_id` is a valid Hugging Face repo or a local path that \
-                        contains a valid `config.json` file. "
-                    f"Original error: {repr(e)}"
-                ) from e
+            hf_config = self._load_hf_config(
+                model_id_or_path, trust_remote_code=trust_remote_code
+            )
+            if (
+                hf_config
+                and hasattr(hf_config, "architectures")
+                and hf_config.architectures
+            ):
+                self._model_architecture = hf_config.architectures[0]
 
         if model_architecture:
             self._model_architecture = model_architecture
@@ -308,8 +317,12 @@ def apply_checkpoint_info(
         self, model_id_or_path: str, trust_remote_code: bool = False
     ) -> None:
         """Apply the checkpoint info to the model config."""
-        self._infer_supports_vision(model_id_or_path)
-        self._set_model_architecture(model_id_or_path)
+        self._infer_supports_vision(
+            model_id_or_path, trust_remote_code=trust_remote_code
+        )
+        self._set_model_architecture(
+            model_id_or_path, trust_remote_code=trust_remote_code
+        )
 
     def get_or_create_callback(self) -> Optional[CallbackBase]:
         """Get or create the callback instance for this process.

python/ray/llm/tests/serve/cpu/configs/test_models.py

@@ -1,5 +1,6 @@
 import sys
 from pathlib import Path
+from unittest.mock import MagicMock, patch
 
 import pydantic
 import pytest
@@ -418,5 +419,28 @@ def test_requires_deferred_placement_group(self):
         assert tpu_accel_with_topo.requires_deferred_placement_group is True
 
 
+class TestCheckpointInfo:
+    def test_apply_checkpoint_info_uses_autoconfig_and_threads_trust_remote_code(self):
+        """apply_checkpoint_info uses AutoConfig (not PretrainedConfig) and forwards
+        trust_remote_code to every HF config load call."""
+        llm_config = LLMConfig(
+            model_loading_config=ModelLoadingConfig(model_id="test_model")
+        )
+        mock_hf_config = MagicMock(spec=["architectures", "vision_config"])
+        mock_hf_config.architectures = ["LlavaForCausalLM"]
+
+        with patch(
+            "transformers.AutoConfig.from_pretrained", return_value=mock_hf_config
+        ) as mock_auto:
+            llm_config.apply_checkpoint_info("vision/model", trust_remote_code=True)
+
+        assert all(
+            call.kwargs["trust_remote_code"] is True
+            for call in mock_auto.call_args_list
+        )
+        assert llm_config._supports_vision is True
+        assert llm_config._model_architecture == "LlavaForCausalLM"
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))