fix: respect trust_remote_code when building AutoConfig (#1007)

Author: ooooo-create

nemo_automodel/recipes/llm/train_ft.py

@@ -405,13 +405,24 @@ def build_loss_fn(cfg_loss):
     return cfg_loss.instantiate()
 
 
-def _build_tokenizer(cfg_model, cfg_ds):
-    def compute_trust_remote_code():
-        if hasattr(cfg_model, "trust_remote_code"):
-            return getattr(cfg_model, "trust_remote_code")
-        return resolve_trust_remote_code(_get_model_name(cfg_model))
+def compute_trust_remote_code_from_model(cfg_model):
+    """Compute the value of trust_remote_code based on the model configuration.
+
+    Args:
+        cfg_model (ConfigNode): Model configuration.
 
-    trust_remote_code = compute_trust_remote_code()
+    Returns:
+        bool: Whether to trust remote code.
+    """
+    if hasattr(cfg_model, "trust_remote_code"):
+        return getattr(cfg_model, "trust_remote_code")
+    elif hasattr(cfg_model, "config") and hasattr(cfg_model.config, "trust_remote_code"):
+        return getattr(cfg_model.config, "trust_remote_code")
+    return resolve_trust_remote_code(_get_model_name(cfg_model))
+
+
+def _build_tokenizer(cfg_model, cfg_ds):
+    trust_remote_code = compute_trust_remote_code_from_model(cfg_model)
     # if tokenizer is not provided, use the model config to instantiate it
     if "tokenizer" not in cfg_ds and _get_model_name(cfg_model) is not None:
         logging.info("Using model config to instantiate tokenizer")
@@ -592,7 +603,9 @@ def build_dataloader(
         if pp_enabled:
             from nemo_automodel.components.datasets.utils import add_causal_masks_to_batch
 
-            hf_model_config = AutoConfig.from_pretrained(_get_model_name(cfg_model))
+            hf_model_config = AutoConfig.from_pretrained(
+                _get_model_name(cfg_model), trust_remote_code=compute_trust_remote_code_from_model(cfg_model)
+            )
 
             if "collate_fn" in dl_kwargs:
                 # Case 1: PP enabled + collate_fn exists -> chain them

tests/unit_tests/recipes/test_train_ft.py

@@ -28,6 +28,7 @@
     build_dataloader,
     build_model_and_optimizer,
     build_validation_dataloader,
+    compute_trust_remote_code_from_model,
 )
 from torch.utils.data import IterableDataset
 
@@ -632,3 +633,36 @@ def _build_model_and_optimizer_stub(*args, **kwargs):
         (parts[0], "PipelineStage_0"),
         (parts[1], "PipelineStage_1"),
     ]
+
+
+def test_compute_trust_remote_code_prefers_cfg_flag():
+    cfg_model = ConfigNode({"trust_remote_code": False, "pretrained_model_name_or_path": "ignored"})
+
+    with patch("nemo_automodel.recipes.llm.train_ft.resolve_trust_remote_code") as mock_resolve:
+        result = compute_trust_remote_code_from_model(cfg_model)
+
+    assert result is False
+    mock_resolve.assert_not_called()
+
+
+def test_compute_trust_remote_code_prefers_nested_config():
+    cfg_model = ConfigNode({"config": {"trust_remote_code": True}})
+
+    with patch("nemo_automodel.recipes.llm.train_ft.resolve_trust_remote_code") as mock_resolve:
+        result = compute_trust_remote_code_from_model(cfg_model)
+
+    assert result is True
+    mock_resolve.assert_not_called()
+
+
+def test_compute_trust_remote_code_falls_back_to_resolve():
+    cfg_model = ConfigNode({"pretrained_model_name_or_path": "nvidia/foo"})
+
+    with patch(
+        "nemo_automodel.recipes.llm.train_ft.resolve_trust_remote_code",
+        return_value=True,
+    ) as mock_resolve:
+        result = compute_trust_remote_code_from_model(cfg_model)
+
+    assert result is True
+    mock_resolve.assert_called_once_with("nvidia/foo")