fix: Use single method to determine trust_remote_code value throughout (#229)

## Summary
- Some `AutoConfig`, `AutoModelForCausalLM`, and `AutoTokenizer`
`from_pretrained` calls and `vLLM` init were missing
`trust_remote_code=True` for `nvidia/` models (e.g. Nemotron), causing
`ValueError` when loading models with custom code
- Consolidates the check into a single `trust_remote_code_for_model()`
in `llm/utils.py`, called by all 8 `ModelMetadata` subclasses,
`populate_derived_fields`, `LLMPromptConfig.from_tokenizer`,
`HuggingFaceBackend`, and `VllmBackend`
- Removes the redundant `TrainingBackend._trust_remote_code_for_model()`
method

## Test plan
- [x] Verified `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` loads
successfully through the SDK `SafeSynthesizer` pipeline (requires extra
deps to load, but the trust_remote_code won't cause a problem now)
- [x] Added unit tests for `trust_remote_code_from_model` 
- [X] Unit tests pass (`make test`)

## Other notes

Created #231 to
followup and make modifying this behavior use configurable.

Made with [Cursor](https://cursor.com)

---------

Signed-off-by: Kendrick Boyd <kendrickb@nvidia.com>
Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
Co-authored-by: Aaron Gonzales <aagonzales@nvidia.com>

Author: kendrickb-nvidia

src/nemo_safe_synthesizer/generation/vllm_backend.py

@@ -31,7 +31,7 @@
 from ..generation.regex_manager import build_json_based_regex
 from ..generation.results import GenerateJobResults, GenerationBatches, GenerationStatus
 from ..llm.metadata import ModelMetadata
-from ..llm.utils import cleanup_memory, get_max_vram
+from ..llm.utils import cleanup_memory, get_max_vram, trust_remote_code_for_model
 from ..observability import get_logger, heartbeat
 from ..utils import all_equal_type, load_json
 
@@ -212,6 +212,7 @@ def initialize(self, **kwargs) -> None:
                 max_lora_rank=self.config.training.lora_r,
                 structured_outputs_config=structured_outputs_config,
                 attention_config=attention_config,
+                trust_remote_code=trust_remote_code_for_model(self.config.training.pretrained_model),
             )
 
         # vLLM's get_tokenizer() returns a wider union than HF's PreTrainedTokenizerBase;

src/nemo_safe_synthesizer/llm/metadata.py

@@ -16,7 +16,7 @@
     field_validator,
     model_validator,
 )
-from transformers import AutoConfig, AutoTokenizer, PretrainedConfig
+from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedTokenizerBase
 
 from ..cli.artifact_structure import Workdir
 from ..config.parameters import SafeSynthesizerParameters
@@ -27,6 +27,7 @@
 )
 from ..observability import get_logger
 from ..utils import load_json, write_json
+from .utils import trust_remote_code_for_model
 
 logger = get_logger(__name__)
 
@@ -77,7 +78,7 @@ class LLMPromptConfig(BaseModel):
     """Integer id for the EOS token."""
 
     @classmethod
-    def from_tokenizer(cls, name: str, tokenizer: AutoTokenizer | None = None, **kwargs) -> LLMPromptConfig:
+    def from_tokenizer(cls, name: str, tokenizer: PreTrainedTokenizerBase | None = None, **kwargs) -> LLMPromptConfig:
         """Create a prompt config by reading from settings of a tokenizer.
 
         If no ``tokenizer`` is supplied one is loaded from ``name``
@@ -94,7 +95,9 @@ def from_tokenizer(cls, name: str, tokenizer: AutoTokenizer | None = None, **kwa
         Returns:
             A new ``LLMPromptConfig`` populated from the tokenizer.
         """
-        tokenizer = tokenizer or AutoTokenizer.from_pretrained(name)
+        tokenizer = tokenizer or AutoTokenizer.from_pretrained(
+            name, trust_remote_code=trust_remote_code_for_model(name)
+        )
         bos_token = kwargs.get("bos_token", getattr(tokenizer, "bos_token", None))
         bos_token_id = kwargs.get("bos_token_id", getattr(tokenizer, "bos_token_id", None))
         eos_token = kwargs.get("eos_token", getattr(tokenizer, "eos_token", None))
@@ -339,7 +342,11 @@ def populate_derived_fields(cls, data: dict) -> dict:
             The mutated ``data`` dict with derived fields populated.
         """
         if data.get("autoconfig") is None:
-            data["autoconfig"] = AutoConfig.from_pretrained(data["model_name_or_path"])
+            model_name_or_path = data["model_name_or_path"]
+            data["autoconfig"] = AutoConfig.from_pretrained(
+                model_name_or_path,
+                trust_remote_code=trust_remote_code_for_model(model_name_or_path),
+            )
 
         if data.get("base_max_seq_length") is None:
             data["base_max_seq_length"] = get_base_max_seq_length(data["autoconfig"])
@@ -447,6 +454,32 @@ def save_metadata(self) -> None:
             indent=4,
         )
 
+    @staticmethod
+    def _load_config_and_tokenizer(
+        model_name_or_path: str,
+        tokenizer: PreTrainedTokenizerBase | None = None,
+    ) -> tuple[PretrainedConfig, PreTrainedTokenizerBase]:
+        """Load ``PretrainedConfig`` and (optionally) ``AutoTokenizer`` for a model.
+
+        Centralises the repeated boilerplate present in every subclass
+        ``__init__``: loading the HuggingFace config and, when no
+        pre-loaded tokenizer is supplied, fetching one via
+        ``AutoTokenizer.from_pretrained``.
+
+        Args:
+            model_name_or_path: HuggingFace model identifier or local path.
+            tokenizer: Pre-loaded tokenizer to reuse.  When ``None`` a new
+                one is loaded from ``model_name_or_path``.
+
+        Returns:
+            A ``(config, tokenizer)`` tuple ready to pass to ``super().__init__``.
+        """
+        trust = trust_remote_code_for_model(model_name_or_path)
+        config: PretrainedConfig = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=trust)
+        if tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=trust)
+        return config, tokenizer
+
     @classmethod
     def _resolve_model_class(cls: type["ModelMetadata"], model_name_or_path: Path | str) -> type["ModelMetadata"]:
         """Resolve model name or path to the matching metadata subclass.
@@ -588,8 +621,7 @@ def __init__(
         rope_scaling_factor: float | None = None,
         **kwargs,
     ) -> None:
-        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) if tokenizer is None else tokenizer
-        config: PretrainedConfig = AutoConfig.from_pretrained(model_name_or_path)
+        config, tokenizer = ModelMetadata._load_config_and_tokenizer(model_name_or_path, tokenizer)
 
         super().__init__(
             autoconfig=config,
@@ -628,8 +660,7 @@ def __init__(
         rope_scaling_factor: float | None = None,
         **kwargs,
     ) -> None:
-        config: PretrainedConfig = AutoConfig.from_pretrained(model_name_or_path)
-        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) if tokenizer is None else tokenizer
+        config, tokenizer = ModelMetadata._load_config_and_tokenizer(model_name_or_path, tokenizer)
 
         super().__init__(
             autoconfig=config,
@@ -668,12 +699,11 @@ class Mistral(ModelMetadata):
     def __init__(
         self,
         model_name_or_path: str,
-        tokenizer: AutoTokenizer | None = None,
+        tokenizer: PreTrainedTokenizerBase | None = None,
         rope_scaling_factor: float | None = None,
         **kwargs,
     ) -> None:
-        tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(model_name_or_path) if tokenizer is None else tokenizer
-        config: PretrainedConfig = AutoConfig.from_pretrained(model_name_or_path)
+        config, tokenizer = ModelMetadata._load_config_and_tokenizer(model_name_or_path, tokenizer)
         if rope_scaling_factor:
             logger.warning(
                 f"Rope scaling factor {rope_scaling_factor} is not supported for Mistral due to longer default context lengths. Ignoring."
@@ -714,8 +744,7 @@ def __init__(
         rope_scaling_factor: float | None = None,
         **kwargs,
     ) -> None:
-        tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(model_name_or_path) if tokenizer is None else tokenizer
-        config: PretrainedConfig = AutoConfig.from_pretrained(model_name_or_path)
+        config, tokenizer = ModelMetadata._load_config_and_tokenizer(model_name_or_path, tokenizer)
 
         super().__init__(
             autoconfig=config,
@@ -751,8 +780,7 @@ def __init__(
         rope_scaling_factor: float | None = None,
         **kwargs,
     ) -> None:
-        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) if tokenizer is None else tokenizer
-        config = AutoConfig.from_pretrained(model_name_or_path)
+        config, tokenizer = ModelMetadata._load_config_and_tokenizer(model_name_or_path, tokenizer)
 
         super().__init__(
             autoconfig=config,
@@ -792,14 +820,13 @@ def __init__(
         rope_scaling_factor: float | None = None,
         **kwargs,
     ) -> None:
-        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) if tokenizer is None else tokenizer
-        config = AutoConfig.from_pretrained(model_name_or_path)
+        config, tokenizer = ModelMetadata._load_config_and_tokenizer(model_name_or_path, tokenizer)
         if rope_scaling_factor:
             logger.warning(
                 f"Rope scaling factor {rope_scaling_factor} is not supported for SmolLM2 due to longer default context lengths. Ignoring."
             )
 
-        im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
+        im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>")  # ty: ignore[unresolved-attribute] -- third-party stub
         super().__init__(
             autoconfig=config,
             instruction=DEFAULT_INSTRUCTION,
@@ -840,8 +867,7 @@ def __init__(
         rope_scaling_factor: float | None = None,
         **kwargs,
     ) -> None:
-        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) if tokenizer is None else tokenizer
-        config = AutoConfig.from_pretrained(model_name_or_path)
+        config, tokenizer = ModelMetadata._load_config_and_tokenizer(model_name_or_path, tokenizer)
 
         # we use the bos token here explicitly for support during group-by SFT.
         # the groupby assumes there is a bos token at the start of the prompt.
@@ -890,8 +916,7 @@ def __init__(
         rope_scaling_factor: float | None = None,
         **kwargs,
     ) -> None:
-        tokenizer = tokenizer or AutoTokenizer.from_pretrained(model_name_or_path)
-        config = AutoConfig.from_pretrained(model_name_or_path)
+        config, tokenizer = ModelMetadata._load_config_and_tokenizer(model_name_or_path, tokenizer)
 
         super().__init__(
             autoconfig=config,

src/nemo_safe_synthesizer/llm/utils.py

@@ -1,50 +1,60 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""GPU memory management, quantization, device mapping, and tokenizer helpers for LLM loading."""
+"""GPU memory management, quantization, device mapping, and tokenizer helpers for LLM loading.
+
+Optional LLM dependencies are imported inside the helpers that need them so
+lightweight utilities such as ``trust_remote_code_for_model`` remain usable
+without installing the full training or inference stack.
+"""
 
 from __future__ import annotations
 
 import gc
 from pathlib import Path
-from typing import Any, Literal
-
-import torch
-from accelerate import infer_auto_device_map, init_empty_weights
-from peft import (
-    PeftModel,
-)
-from transformers import (
-    AutoConfig,
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    BitsAndBytesConfig,
-    PreTrainedTokenizer,
-)
+from typing import TYPE_CHECKING, Any, Literal
 
 from ..observability import get_logger
 
+if TYPE_CHECKING:
+    from peft import PeftModel
+    from transformers import AutoConfig, BitsAndBytesConfig, PreTrainedTokenizer
+
 logger = get_logger(__name__)
 
 
 def trust_remote_code_for_model(model_name: str | Path) -> bool:
     """Determine whether to trust remote code when loading a model.
 
-    Returns ``True`` only for models whose name starts with
-    ``"nvidia/"``.
+    Returns ``True`` for NVIDIA-owned Hub model identifiers and for paths
+    inside Hugging Face's encoded cache directory for NVIDIA models.
 
     Args:
         model_name: HuggingFace model identifier or local path.
 
     Returns:
         Whether to set ``trust_remote_code=True`` when loading the model.
     """
-    mn = str(model_name)
-    return mn.startswith("nvidia/")
+    model_ref = str(model_name).casefold()
+    if model_ref.startswith("nvidia/"):
+        return True
+
+    path_parts = Path(model_ref).parts
+    while path_parts:
+        match path_parts:
+            case ("huggingface", "hub", *cache_parts):
+                # Brittle by design: this mirrors Hugging Face's current cache path layout.
+                return any(part.startswith("models--nvidia--") for part in cache_parts)
+            case (_, *remaining):
+                path_parts = remaining
+
+    return False
 
 
 def cleanup_memory() -> None:
     """Run garbage collection and empty the CUDA cache."""
+    import torch
+
     gc.collect()
     with torch.no_grad():
         torch.cuda.empty_cache()
@@ -56,6 +66,7 @@ def gpu_stats() -> None:
     Queries CUDA device 0 and logs the peak reserved memory and total
     available memory in GiB.
     """
+    import torch
 
     def round_gb(value: float) -> float:
         return round(value / 1024 / 1024 / 1024, 3)
@@ -80,6 +91,8 @@ def get_max_vram(max_vram_fraction: float | None = None) -> dict[int, float]:
     Returns:
         Mapping of CUDA device index to the usable memory fraction.
     """
+    import torch
+
     if max_vram_fraction is None:
         max_vram_fraction = 0.8
     max_memory = {}
@@ -148,6 +161,8 @@ def get_param_from_config(
     Raises:
         ValueError: If neither ``model_name`` nor ``config`` is provided.
     """
+    from transformers import AutoConfig
+
     if config is None:
         if model_name is None:
             raise ValueError("model_name is required if config is not provided")
@@ -170,6 +185,8 @@ def _get_auto_tokenizer(
     Returns:
         Configured ``PreTrainedTokenizer`` with BOS/EOS tokens enabled.
     """
+    from transformers import AutoTokenizer
+
     tokenizer = AutoTokenizer.from_pretrained(
         model_name,
         model_max_length=max_position_embeddings,
@@ -204,6 +221,9 @@ def get_device_map(
     Returns:
         Ordered dictionary mapping layer names to device identifiers.
     """
+    from accelerate import infer_auto_device_map, init_empty_weights
+    from transformers import AutoConfig, AutoModelForCausalLM
+
     config = autoconfig or AutoConfig.from_pretrained(
         model_target,
         revision=revision,
@@ -253,6 +273,9 @@ def get_quantization_config(quantization_bits: Literal[4, 8]) -> BitsAndBytesCon
     Raises:
         ValueError: If ``quantization_bits`` is not 4 or 8.
     """
+    import torch
+    from transformers import BitsAndBytesConfig
+
     if quantization_bits == 4:
         return BitsAndBytesConfig(
             load_in_4bit=True,

src/nemo_safe_synthesizer/training/backend.py

@@ -296,13 +296,3 @@ def teardown(self) -> None:
         this runs even when training raises.
         """
         pass
-
-    def _trust_remote_code_for_model(self) -> bool:
-        """Determine whether the model should be loaded with ``trust_remote_code=True``.
-
-        Currently returns ``True`` only for NVIDIA models on HuggingFace Hub.
-
-        Returns:
-            Whether to trust remote code when loading the model.
-        """
-        return str(self.params.training.pretrained_model).startswith("nvidia/")

src/nemo_safe_synthesizer/training/huggingface_backend.py

@@ -58,6 +58,7 @@
     get_device_map,
     get_max_vram,
     get_quantization_config,
+    trust_remote_code_for_model,
 )
 from ..observability import get_logger, traced_runtime, traced_user
 from ..privacy.dp_transformers.dp_utils import (
@@ -120,7 +121,8 @@ def __init__(self, *args, **kwargs):
         self.model_loader_type = AutoModelForCausalLM
         self.training_output_dir = Path(self.workdir.train.cache)
         self.autoconfig = AutoConfig.from_pretrained(
-            self.params.training.pretrained_model, trust_remote_code=self._trust_remote_code_for_model()
+            self.params.training.pretrained_model,
+            trust_remote_code=trust_remote_code_for_model(self.params.training.pretrained_model),
         )
 
     def _load_pretrained_model(self, **model_args: Any) -> None:
@@ -132,7 +134,9 @@ def _load_pretrained_model(self, **model_args: Any) -> None:
 
         self.tokenizer: PreTrainedTokenizer = add_bos_eos_tokens_to_tokenizer(
             AutoTokenizer.from_pretrained(
-                self.params.training.pretrained_model, model_max_length=model_args.get("max_seq_length", None)
+                self.params.training.pretrained_model,
+                trust_remote_code=trust_remote_code_for_model(self.params.training.pretrained_model),
+                model_max_length=model_args.get("max_seq_length", None),
             )
         )
 
@@ -202,10 +206,17 @@ def _build_base_framework_params(self, model_kwargs: dict) -> dict:
         Returns:
             Dictionary of parameters for ``from_pretrained``.
         """
+        trust_remote_code = trust_remote_code_for_model(self.params.training.pretrained_model)
         return dict(
             pretrained_model_name_or_path=self.params.training.pretrained_model,
+            trust_remote_code=trust_remote_code,
             device_map=model_kwargs.pop(
-                "device_map", get_device_map(self.params.training.pretrained_model, autoconfig=self.autoconfig)
+                "device_map",
+                get_device_map(
+                    self.params.training.pretrained_model,
+                    autoconfig=self.autoconfig,
+                    trust_remote_code=trust_remote_code,
+                ),
             ),
             attn_implementation=model_kwargs.pop(
                 "attn_implementation", self._resolve_attn_implementation(self.params.training.attn_implementation)