From d40d07f1496c7edeadb12ebc0a59777be4106825 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Piotr=20=C5=BBelasko?= Date: Mon, 24 Aug 2026 09:47:35 -0400 Subject: [PATCH] Fix SpeechLM2 remote code opt-in MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Piotr Żelasko --- docs/source/speechlm2/models.rst | 15 ++++ .../speechlm2/models/duplex_ear_tts.py | 4 +- nemo/collections/speechlm2/parts/hf_hub.py | 14 ++++ .../speechlm2/test_duplex_eartts.py | 22 ++++++ tests/collections/speechlm2/test_hf_hub.py | 70 ++++++++++++++++++- 5 files changed, 123 insertions(+), 2 deletions(-) diff --git a/docs/source/speechlm2/models.rst b/docs/source/speechlm2/models.rst index a381fefcaa9f..3ff0b9888ea9 100644 --- a/docs/source/speechlm2/models.rst +++ b/docs/source/speechlm2/models.rst @@ -314,6 +314,21 @@ All models in the speechlm2 collection can be instantiated from pretrained check # Load NemotronVoiceChat (Inference Only) voicechat_model = slm.models.NemotronVoiceChat.from_pretrained("path/to/checkpoint") +Remote HuggingFace code is disabled by default. If a trusted checkpoint requires +custom code, opt in at runtime and pin the repository to a reviewed revision: + +.. code-block:: python + + model = slm.models.SALM.from_pretrained( + "trusted/model", + revision="reviewed-commit-sha", + trust_remote_code=True, + ) + +The ``trust_remote_code`` setting stored in a checkpoint configuration is ignored. +This prevents a model repository from opting itself into executing downloaded +Python code. + Fine-Tuning from a Checkpoint ------------------------------ diff --git a/nemo/collections/speechlm2/models/duplex_ear_tts.py b/nemo/collections/speechlm2/models/duplex_ear_tts.py index 45e49f41d1b3..02e89080ae3a 100644 --- a/nemo/collections/speechlm2/models/duplex_ear_tts.py +++ b/nemo/collections/speechlm2/models/duplex_ear_tts.py @@ -182,7 +182,9 @@ def _load_language_model(self, cfg): """Load language model for RVQ-EAR-TTS.""" if cfg.pretrained_lm_name: language_model = load_pretrained_hf( - self.cfg.pretrained_lm_name, pretrained_weights=True, trust_remote_code=True + self.cfg.pretrained_lm_name, + pretrained_weights=True, + trust_remote_code=self.cfg.get("trust_remote_code", False), ).eval() else: language_model = None diff --git a/nemo/collections/speechlm2/parts/hf_hub.py b/nemo/collections/speechlm2/parts/hf_hub.py index 006c8e019e4e..7bb66841a1ed 100644 --- a/nemo/collections/speechlm2/parts/hf_hub.py +++ b/nemo/collections/speechlm2/parts/hf_hub.py @@ -41,6 +41,7 @@ def _from_pretrained( token: Union[str, bool, None], map_location: str = "cpu", strict: bool = False, + trust_remote_code: bool = False, **model_kwargs, ): """ @@ -53,7 +54,15 @@ def _from_pretrained( >>> model = SALM.from_pretrained( ... "nvidia/salm-model", distributed_setup=strategy.distributed_setup ... ) + + ``trust_remote_code`` is a runtime security decision. It is deliberately + taken from the caller and overwrites any value stored in the downloaded + checkpoint config so that a model repository cannot opt itself into + executing remote code. """ + if not isinstance(trust_remote_code, bool): + raise TypeError(f"trust_remote_code must be a bool, got {type(trust_remote_code).__name__}") + distributed_setup = model_kwargs.pop("distributed_setup", None) device_mesh = distributed_setup.mesh_context.device_mesh if distributed_setup is not None else None torch_dtype = model_kwargs.pop("torch_dtype", None) @@ -73,6 +82,7 @@ def _from_pretrained( if resolved_config_file is None: raise RuntimeError(f"Missing {CONFIG_NAME} file for {model_id=}") model_kwargs['cfg'] = OmegaConf.to_container(OmegaConf.load(resolved_config_file)) + model_kwargs['cfg']['trust_remote_code'] = trust_remote_code _inject_local_artifact_paths(model_kwargs['cfg'], model_id, _cached_file_kwargs) # The setting below tells the model's __init__ not to load the original pretrained weights # for individual children modules. @@ -154,6 +164,10 @@ def save_pretrained( config = OmegaConf.to_container(self.cfg) # Ensure HF-compatible fields are present so vLLM / transformers can identify the model. if isinstance(config, dict): + config = dict(config) + # Remote-code trust is a runtime choice and must never be persisted + # in a checkpoint that can be loaded by another user. + config.pop("trust_remote_code", None) config.setdefault("model_type", "nemo_speechlm") config.setdefault("architectures", ["NeMoSpeechLMForConditionalGeneration"]) return super().save_pretrained( diff --git a/tests/collections/speechlm2/test_duplex_eartts.py b/tests/collections/speechlm2/test_duplex_eartts.py index 4468bde3e16d..2b1d6c54599a 100644 --- a/tests/collections/speechlm2/test_duplex_eartts.py +++ b/tests/collections/speechlm2/test_duplex_eartts.py @@ -13,10 +13,13 @@ # limitations under the License. import os +from types import SimpleNamespace + import pytest import torch from lhotse import CutSet, SupervisionSegment from lhotse.testing.dummies import dummy_cut, dummy_recording +from omegaconf import DictConfig from nemo.collections.common.data.utils import move_data_to_device from nemo.collections.speechlm2.data.duplex_ear_tts_dataset import ( @@ -31,6 +34,25 @@ torch.set_default_device('cuda') +def test_load_language_model_uses_configured_remote_code_policy(monkeypatch): + captured = {} + + class DummyLanguageModel: + def eval(self): + return self + + def fake_load_pretrained_hf(*args, **kwargs): + captured.update(kwargs) + return DummyLanguageModel() + + monkeypatch.setattr("nemo.collections.speechlm2.models.duplex_ear_tts.load_pretrained_hf", fake_load_pretrained_hf) + cfg = DictConfig({"pretrained_lm_name": "untrusted/repository", "trust_remote_code": False}) + + DuplexEARTTS._load_language_model(SimpleNamespace(cfg=cfg), cfg) + + assert captured["trust_remote_code"] is False + + test_eartts_config = { "model": { "pretrained_lm_name": "nvidia/NVIDIA-Nemotron-Nano-9B-v2", diff --git a/tests/collections/speechlm2/test_hf_hub.py b/tests/collections/speechlm2/test_hf_hub.py index 85f512292daf..972f2b57b0b2 100644 --- a/tests/collections/speechlm2/test_hf_hub.py +++ b/tests/collections/speechlm2/test_hf_hub.py @@ -12,7 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -from nemo.collections.speechlm2.parts.hf_hub import _inject_local_artifact_paths +import pytest +from huggingface_hub import PyTorchModelHubMixin + +import nemo.collections.speechlm2.parts.hf_hub as hf_hub +from nemo.collections.speechlm2.parts.hf_hub import HFHubMixin, _inject_local_artifact_paths + + +class _DummyHubModel(HFHubMixin): + pass def _cached_file_kwargs(): @@ -34,6 +42,66 @@ def _write_local_export_artifacts(tmp_path): (tmp_path / "llm_backbone" / "config.json").write_text("{}") +def _capture_pretrained_config(tmp_path, monkeypatch, repo_trust_remote_code, **model_kwargs): + config_path = tmp_path / "config.json" + config_path.write_text(f"trust_remote_code: {str(repo_trust_remote_code).lower()}\n") + + def fake_cached_file(_model_id, filename, **_kwargs): + return str(config_path) if filename == hf_hub.CONFIG_NAME else None + + captured = {} + + def fake_from_pretrained(_cls, **kwargs): + captured.update(kwargs) + return object() + + monkeypatch.setattr(hf_hub, "cached_file", fake_cached_file) + monkeypatch.setattr(PyTorchModelHubMixin, "_from_pretrained", classmethod(fake_from_pretrained)) + + _DummyHubModel._from_pretrained( + model_id="untrusted/repository", + revision=None, + cache_dir=None, + force_download=False, + local_files_only=True, + token=None, + **model_kwargs, + ) + return captured["cfg"] + + +@pytest.mark.parametrize( + ("repo_trust_remote_code", "model_kwargs", "expected"), + [ + pytest.param(True, {}, False, id="repository-cannot-opt-in"), + pytest.param(True, {"trust_remote_code": False}, False, id="explicit-opt-out-wins"), + pytest.param(False, {"trust_remote_code": True}, True, id="explicit-opt-in-wins"), + ], +) +def test_from_pretrained_remote_code_requires_explicit_opt_in( + tmp_path, monkeypatch, repo_trust_remote_code, model_kwargs, expected +): + cfg = _capture_pretrained_config(tmp_path, monkeypatch, repo_trust_remote_code, **model_kwargs) + + assert cfg["trust_remote_code"] is expected + + +def test_save_pretrained_does_not_persist_remote_code_trust(tmp_path, monkeypatch): + captured = {} + + def fake_save_pretrained(_self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(PyTorchModelHubMixin, "save_pretrained", fake_save_pretrained) + model = object.__new__(_DummyHubModel) + model.cfg = {"trust_remote_code": True} + + model.save_pretrained(tmp_path) + + assert "trust_remote_code" not in captured["config"] + assert model.cfg["trust_remote_code"] is True + + def test_inject_local_artifact_paths_salm_config(tmp_path): _write_local_export_artifacts(tmp_path) cfg = {