Skip to content

Commit 39e5d08

Browse files
dimapihtardimapihtar
authored
guard trust_remote_code (NVIDIA-NeMo#15065)
* guard trust_remote_code Signed-off-by: dimapihtar <dpihtar@gmail.com> * Apply isort and black reformatting Signed-off-by: dimapihtar <dimapihtar@users.noreply.github.com> * fix import Signed-off-by: dimapihtar <dpihtar@gmail.com> * Apply isort and black reformatting Signed-off-by: dimapihtar <dimapihtar@users.noreply.github.com> * fix trust_remote_code logic Signed-off-by: dimapihtar <dpihtar@gmail.com> * Apply isort and black reformatting Signed-off-by: dimapihtar <dimapihtar@users.noreply.github.com> * fix unit tests Signed-off-by: dimapihtar <dpihtar@gmail.com> * fix tests Signed-off-by: dimapihtar <dpihtar@gmail.com> * fix multiple values issue Signed-off-by: dimapihtar <dpihtar@gmail.com> * fix kwargs Signed-off-by: dimapihtar <dpihtar@gmail.com> * revert changes Signed-off-by: dimapihtar <dpihtar@gmail.com> * fix tests Signed-off-by: dimapihtar <dpihtar@gmail.com> * revert changes Signed-off-by: dimapihtar <dpihtar@gmail.com> * Apply isort and black reformatting Signed-off-by: dimapihtar <dimapihtar@users.noreply.github.com> * fix tests Signed-off-by: dimapihtar <dpihtar@gmail.com> --------- Signed-off-by: dimapihtar <dpihtar@gmail.com> Signed-off-by: dimapihtar <dimapihtar@users.noreply.github.com> Co-authored-by: dimapihtar <dimapihtar@users.noreply.github.com>
1 parent d14d05b commit 39e5d08

26 files changed

+484
-112
lines changed

‎nemo/collections/llm/gpt/model/baichuan.py‎

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
from torch import nn
2222

2323
from nemo.collections.llm.gpt.model.base import GPTConfig, GPTModel, torch_dtype_from_mcore_config
24-
from nemo.collections.llm.utils import Config
24+
from nemo.collections.llm.utils import Config, is_safe_repo
2525
from nemo.lightning import OptimizerModule, io, teardown
2626
from nemo.lightning.io.state import TransformFns
2727
from nemo.lightning.pytorch.utils import dtype_from_hf
@@ -104,19 +104,28 @@ def init(self) -> Baichuan2Model:
104104
"""
105105
return Baichuan2Model(self.config, tokenizer=self.tokenizer)
106106

107-
def apply(self, output_path: Path) -> Path:
107+
def apply(self, output_path: Path, trust_remote_code: bool | None = None) -> Path:
108108
"""
109109
Apply the conversion from HF to NeMo format.
110110
111111
Args:
112112
output_path: Path where the converted model will be saved
113+
trust_remote_code: Whether remote code execution should be trusted for a given HF path
113114
114115
Returns:
115116
Path: Path to the saved NeMo model
116117
"""
117118
from transformers import AutoModelForCausalLM
118119

119-
source = AutoModelForCausalLM.from_pretrained(str(self), trust_remote_code=True, torch_dtype='auto')
120+
self.trust_remote_code = trust_remote_code
121+
source = AutoModelForCausalLM.from_pretrained(
122+
str(self),
123+
trust_remote_code=is_safe_repo(
124+
trust_remote_code=self.trust_remote_code,
125+
hf_path=str(self),
126+
),
127+
torch_dtype='auto',
128+
)
120129
target = self.init()
121130
trainer = self.nemo_setup(target)
122131
self.convert_state(source, target)
@@ -173,7 +182,13 @@ def tokenizer(self) -> "AutoTokenizer":
173182
"""
174183
from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
175184

176-
return AutoTokenizer(self.save_hf_tokenizer_assets(str(self)), trust_remote_code=True)
185+
return AutoTokenizer(
186+
self.save_hf_tokenizer_assets(str(self)),
187+
trust_remote_code=is_safe_repo(
188+
trust_remote_code=self.trust_remote_code,
189+
hf_path=str(self),
190+
),
191+
)
177192

178193
@property
179194
def config(self) -> Baichuan2Config:
@@ -188,7 +203,13 @@ def config(self) -> Baichuan2Config:
188203
"""
189204
from transformers import AutoConfig as HFAutoConfig
190205

191-
source = HFAutoConfig.from_pretrained(str(self), trust_remote_code=True)
206+
source = HFAutoConfig.from_pretrained(
207+
str(self),
208+
trust_remote_code=is_safe_repo(
209+
trust_remote_code=self.trust_remote_code,
210+
hf_path=str(self),
211+
),
212+
)
192213

193214
def make_vocab_size_divisible_by(vocab_size):
194215
base = 128
@@ -243,17 +264,27 @@ def init(self, dtype=torch.bfloat16, model_name=None) -> "AutoModelForCausalLM":
243264
# Since Baichuan2 is not importable from transformers, we can only initialize the HF model
244265
# from a known checkpoint folder containing the config file and modeling files.
245266
# The model_name will need to be passed in.
246-
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
267+
config = AutoConfig.from_pretrained(
268+
model_name,
269+
trust_remote_code=is_safe_repo(
270+
trust_remote_code=self.trust_remote_code,
271+
hf_path=model_name,
272+
),
273+
)
247274
hf_model = AutoModelForCausalLM.from_config(
248275
config,
249-
trust_remote_code=True,
276+
trust_remote_code=is_safe_repo(
277+
trust_remote_code=self.trust_remote_code,
278+
hf_path=model_name,
279+
),
250280
torch_dtype=dtype,
251281
)
252282
# Register the AutoModel Hook so that the custom modeling files are saved during save_pretrained()
253283
type(hf_model).register_for_auto_class("AutoModelForCausalLM")
254284
return hf_model
255285

256-
def apply(self, output_path: Path, target_model_name=None) -> Path:
286+
def apply(self, output_path: Path, target_model_name=None, trust_remote_code: bool | None = None) -> Path:
287+
self.trust_remote_code = trust_remote_code
257288
source, _ = self.nemo_load(str(self))
258289
target = self.init(torch_dtype_from_mcore_config(source.config), model_name=target_model_name)
259290
target = self.convert_state(source, target)

‎nemo/collections/llm/gpt/model/chatglm.py‎

Lines changed: 39 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
from torch import nn
2222

2323
from nemo.collections.llm.gpt.model.base import GPTConfig, GPTModel, torch_dtype_from_mcore_config
24-
from nemo.collections.llm.utils import Config
24+
from nemo.collections.llm.utils import Config, is_safe_repo
2525
from nemo.lightning import OptimizerModule, io, teardown
2626
from nemo.lightning.io.state import TransformFns
2727
from nemo.lightning.pytorch.utils import dtype_from_hf
@@ -114,19 +114,28 @@ def init(self) -> ChatGLMModel:
114114
"""
115115
return ChatGLMModel(self.config, tokenizer=self.tokenizer)
116116

117-
def apply(self, output_path: Path) -> Path:
117+
def apply(self, output_path: Path, trust_remote_code: bool | None = None) -> Path:
118118
"""
119119
Apply the conversion from HF to NeMo format.
120120
121121
Args:
122122
output_path: Path where the converted model will be saved
123+
trust_remote_code: Whether remote code execution should be trusted for a given HF path
123124
124125
Returns:
125126
Path: Path to the saved NeMo model
126127
"""
127128
from transformers import AutoModelForCausalLM
128129

129-
source = AutoModelForCausalLM.from_pretrained(str(self), trust_remote_code=True, torch_dtype='auto')
130+
self.trust_remote_code = trust_remote_code
131+
source = AutoModelForCausalLM.from_pretrained(
132+
str(self),
133+
trust_remote_code=is_safe_repo(
134+
trust_remote_code=self.trust_remote_code,
135+
hf_path=str(self),
136+
),
137+
torch_dtype='auto',
138+
)
130139
target = self.init()
131140
trainer = self.nemo_setup(target)
132141
self.convert_state(source, target)
@@ -177,7 +186,13 @@ def tokenizer(self) -> "AutoTokenizer":
177186
"""
178187
from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
179188

180-
return AutoTokenizer(self.save_hf_tokenizer_assets(str(self)), trust_remote_code=True)
189+
return AutoTokenizer(
190+
self.save_hf_tokenizer_assets(str(self)),
191+
trust_remote_code=is_safe_repo(
192+
trust_remote_code=self.trust_remote_code,
193+
hf_path=str(self),
194+
),
195+
)
181196

182197
@property
183198
def config(self) -> ChatGLMConfig:
@@ -192,7 +207,13 @@ def config(self) -> ChatGLMConfig:
192207
"""
193208
from transformers import AutoConfig as HFAutoConfig
194209

195-
source = HFAutoConfig.from_pretrained(str(self), trust_remote_code=True)
210+
source = HFAutoConfig.from_pretrained(
211+
str(self),
212+
trust_remote_code=is_safe_repo(
213+
trust_remote_code=self.trust_remote_code,
214+
hf_path=str(self),
215+
),
216+
)
196217
output = ChatGLMConfig(
197218
num_layers=source.num_layers,
198219
hidden_size=source.hidden_size,
@@ -228,17 +249,27 @@ def init(self, dtype=torch.bfloat16, model_name=None) -> "AutoModelForCausalLM":
228249
# Since ChatGLM is not importable from transformers, we can only initialize the HF model
229250
# from a known checkpoint folder containing the config file and modeling files.
230251
# The model_name will need to be passed in.
231-
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
252+
config = AutoConfig.from_pretrained(
253+
model_name,
254+
trust_remote_code=is_safe_repo(
255+
trust_remote_code=self.trust_remote_code,
256+
hf_path=model_name,
257+
),
258+
)
232259
hf_model = AutoModelForCausalLM.from_config(
233260
config,
234-
trust_remote_code=True,
261+
trust_remote_code=is_safe_repo(
262+
trust_remote_code=self.trust_remote_code,
263+
hf_path=model_name,
264+
),
235265
torch_dtype=dtype,
236266
)
237267
# Register the AutoModel Hook so that the custom modeling files are saved during save_pretrained()
238268
type(hf_model).register_for_auto_class("AutoModelForCausalLM")
239269
return hf_model
240270

241-
def apply(self, output_path: Path, target_model_name=None) -> Path:
271+
def apply(self, output_path: Path, target_model_name=None, trust_remote_code: bool | None = None) -> Path:
272+
self.trust_remote_code = trust_remote_code
242273
source, _ = self.nemo_load(str(self))
243274
target = self.init(torch_dtype_from_mcore_config(source.config), model_name=target_model_name)
244275
target = self.convert_state(source, target)

‎nemo/collections/llm/gpt/model/deepseek.py‎

Lines changed: 38 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
gpt_data_step,
3636
torch_dtype_from_dict_config,
3737
)
38+
from nemo.collections.llm.utils import is_safe_repo
3839
from nemo.export.trt_llm.nemo_ckpt_loader.nemo_file import load_distributed_model_weights
3940
from nemo.lightning import io, teardown
4041
from nemo.lightning.io.state import TransformFns, _ModelState
@@ -225,12 +226,20 @@ class HFDeepSeekImporter(io.ModelConnector["AutoModelForCausalLM", DeepSeekModel
225226
def init(self) -> DeepSeekModel:
226227
return DeepSeekModel(self.config, tokenizer=self.tokenizer)
227228

228-
def apply(self, output_path: Path, convert_mtp: bool = False) -> Path:
229+
def apply(self, output_path: Path, convert_mtp: bool = False, trust_remote_code: bool | None = None) -> Path:
229230
from transformers import AutoModelForCausalLM
230231

232+
self.trust_remote_code = trust_remote_code
231233
self.convert_mtp = convert_mtp
232234
self._verify_source()
233-
source = AutoModelForCausalLM.from_pretrained(str(self), trust_remote_code=True, torch_dtype='auto')
235+
source = AutoModelForCausalLM.from_pretrained(
236+
str(self),
237+
trust_remote_code=is_safe_repo(
238+
trust_remote_code=self.trust_remote_code,
239+
hf_path=str(self),
240+
),
241+
torch_dtype='auto',
242+
)
234243
target = self.init()
235244
trainer = self.nemo_setup(target)
236245
self.convert_state(source, target)
@@ -244,7 +253,13 @@ def apply(self, output_path: Path, convert_mtp: bool = False) -> Path:
244253
return output_path
245254

246255
def _verify_source(self):
247-
source_config = AutoConfig.from_pretrained(str(self), trust_remote_code=True)
256+
source_config = AutoConfig.from_pretrained(
257+
str(self),
258+
trust_remote_code=is_safe_repo(
259+
trust_remote_code=self.trust_remote_code,
260+
hf_path=str(self),
261+
),
262+
)
248263
assert 'quantization_config' not in source_config, (
249264
"HuggingFace cannot load DeepSeek V3's FP8 checkpoint directly. You must convert the checkpoint "
250265
"to BF16. See NeMo documentation for more details: "
@@ -407,7 +422,13 @@ def config(self) -> DeepSeekConfig:
407422
from transformers import AutoConfig as HFAutoConfig
408423
from transformers import GenerationConfig
409424

410-
source = HFAutoConfig.from_pretrained(str(self), trust_remote_code=True)
425+
source = HFAutoConfig.from_pretrained(
426+
str(self),
427+
trust_remote_code=is_safe_repo(
428+
trust_remote_code=self.trust_remote_code,
429+
hf_path=str(self),
430+
),
431+
)
411432
try:
412433
generation_config = GenerationConfig.from_pretrained(str(self))
413434
except OSError:
@@ -463,10 +484,19 @@ def init(self, dtype=torch.bfloat16, model_name="deepseek-ai/DeepSeek-V3") -> "A
463484
# Since DeepSeek is not importable from transformers, we can only initialize the HF model
464485
# from a known checkpoint folder containing the config file and modeling files.
465486
# The model_name will need to be passed in.
466-
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
487+
config = AutoConfig.from_pretrained(
488+
model_name,
489+
trust_remote_code=is_safe_repo(
490+
trust_remote_code=self.trust_remote_code,
491+
hf_path=model_name,
492+
),
493+
)
467494
hf_model = AutoModelForCausalLM.from_config(
468495
config,
469-
trust_remote_code=True,
496+
trust_remote_code=is_safe_repo(
497+
trust_remote_code=self.trust_remote_code,
498+
hf_path=model_name,
499+
),
470500
torch_dtype=dtype,
471501
)
472502
# Register the AutoModel Hook so that the custom modeling files are saved during save_pretrained()
@@ -528,8 +558,9 @@ def ckpt_load(self, path: Path) -> Tuple[Dict, Dict]:
528558
state_dict[new_k] = v
529559
return state_dict, config['config']
530560

531-
def apply(self, output_path: Path, target_model_name=None) -> Path:
561+
def apply(self, output_path: Path, target_model_name=None, trust_remote_code: bool | None = None) -> Path:
532562
logging.info("Loading DeepSeek NeMo checkpoint. This may take a while...")
563+
self.trust_remote_code = trust_remote_code
533564
source, source_config = self.ckpt_load(self)
534565
logging.info("DeepSeek NeMo checkpoint loaded.")
535566
if target_model_name is None:

‎nemo/collections/llm/gpt/model/gpt_oss.py‎

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
from nemo.collections.common.tokenizers import AutoTokenizer
2828
from nemo.collections.common.tokenizers.tiktoken_tokenizer import TiktokenTokenizer
2929
from nemo.collections.llm.gpt.model.base import GPTConfig, GPTModel, torch_dtype_from_mcore_config
30-
from nemo.collections.llm.utils import Config
30+
from nemo.collections.llm.utils import Config, is_safe_repo
3131
from nemo.lightning import OptimizerModule, io, teardown
3232
from nemo.lightning.io.state import TransformFns, _ModelState
3333
from nemo.utils import logging
@@ -229,8 +229,9 @@ class HFGPTOSSImporter(_BaseGPTOSSImporter):
229229
"""Importer for GPT-OSS models from Hugging Face"""
230230

231231
# pylint: disable=C0115,C0116
232-
def apply(self, output_path: Path) -> Path:
232+
def apply(self, output_path: Path, trust_remote_code: bool | None = None) -> Path:
233233
logging.setLevel(logging.DEBUG)
234+
self.trust_remote_code = trust_remote_code
234235
source_state = self.hf_ckpt_load()
235236
source = _ModelState(source_state)
236237
target = self.init()
@@ -307,14 +308,26 @@ def convert_state(self, source, target):
307308
def tokenizer(self) -> "AutoTokenizer":
308309
from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
309310

310-
return AutoTokenizer(self.save_hf_tokenizer_assets(str(self)), trust_remote_code=True)
311+
return AutoTokenizer(
312+
self.save_hf_tokenizer_assets(str(self)),
313+
trust_remote_code=is_safe_repo(
314+
trust_remote_code=self.trust_remote_code,
315+
hf_path=str(self),
316+
),
317+
)
311318

312319
@cached_property
313320
def config(self) -> GPTOSSConfig:
314321
from transformers import AutoConfig as HFAutoConfig
315322
from transformers import GenerationConfig
316323

317-
source = HFAutoConfig.from_pretrained(str(self), trust_remote_code=True)
324+
source = HFAutoConfig.from_pretrained(
325+
str(self),
326+
trust_remote_code=is_safe_repo(
327+
trust_remote_code=self.trust_remote_code,
328+
hf_path=str(self),
329+
),
330+
)
318331
generation_config = GenerationConfig.from_pretrained(str(self))
319332
return GPTOSSConfig(
320333
num_layers=source.num_hidden_layers,
@@ -426,7 +439,14 @@ def init(self, dtype=torch.bfloat16) -> "AutoModelForCausalLM":
426439
from transformers.modeling_utils import no_init_weights
427440

428441
with no_init_weights():
429-
return AutoModelForCausalLM.from_config(self.config, trust_remote_code=True, torch_dtype=dtype)
442+
return AutoModelForCausalLM.from_config(
443+
self.config,
444+
trust_remote_code=is_safe_repo(
445+
trust_remote_code=self.trust_remote_code,
446+
hf_path=str(self),
447+
),
448+
torch_dtype=dtype,
449+
)
430450

431451
def apply(self, output_path: Path) -> Path:
432452
source, _ = self.nemo_load(str(self))
@@ -549,7 +569,7 @@ def init(self, dtype=torch.bfloat16) -> "AutoPeftModelForCausalLM":
549569
model.name_or_path = os.path.join(model_ckpt_path.split("/")[-2:])
550570
return get_peft_model(model, self.peft_config, autocast_adapter_dtype=False)
551571

552-
def apply(self, output_path: Path) -> Path:
572+
def apply(self, output_path: Path, trust_remote_code: bool | None = None) -> Path:
553573
"""Apply the conversion from NeMo PEFT model to HF format.
554574
555575
Args:
@@ -560,6 +580,7 @@ def apply(self, output_path: Path) -> Path:
560580
"""
561581
from nemo.collections.llm.peft import CanonicalLoRA, DoRA, LoRA
562582

583+
self.trust_remote_code = trust_remote_code
563584
self.peft_obj: Union[LoRA, DoRA, CanonicalLoRA] = io.load_context(str(self), subpath="model.model_transform")
564585

565586
source, _ = self.nemo_load(str(self))

0 commit comments

Comments
 (0)