Fix trust_remote_code and gradient checkpointing for custom models (#696)

* Fix: Add trust_remote_code=True for models with custom code

- Add trust_remote_code=True to all AutoConfig/AutoTokenizer.from_pretrained() calls
- Add torchrun path resolution (shutil.which with sys.executable fallback)
- Pass trust_remote_code=True to base_model_args and VLM helper functions

This fixes training failures for models like Nemotron that use remote code.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

* Fix: Handle models without gradient checkpointing support

Wrap gradient_checkpointing_enable() in try/except to handle models
like NemotronH that don't support gradient checkpointing.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

* Address reviewer feedback and fix ruff formatting

- Narrow exception handling in is_gpt_oss_model to catch specific
  exceptions (OSError, ValueError) instead of bare Exception, and
  log the failure details for debugging
- Add trust_remote_code=True to process_documents_for_pretraining()
  tokenizer loading for consistency with configure_tokenizer()
- Replace invalid fast_tokenizer kwarg with use_fast in
  tokenizer_utils.py setup_tokenizer()
- Create shared _enable_gradient_checkpointing_if_supported() helper
  on Model base class, catching ValueError, NotImplementedError, and
  AttributeError; use it in both LigerModel and CausalLMModel
- Improve torchrun fallback to use sys.executable -m
  torch.distributed.run instead of assuming a sibling script exists
- Fix ruff formatting for AutoConfig.from_pretrained call

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

* Fix unit tests to expect trust_remote_code=True

Update test assertions to expect trust_remote_code=True parameter
in AutoTokenizer.from_pretrained calls after adding this parameter
to process_documents_for_pretraining.

* Fix ruff formatting: break long assertion line

* Make trust_remote_code configurable via flag and environment variable

Instead of hardcoding trust_remote_code=True everywhere:

1. Add trust_remote_code field to TrainingArgs (default: False)
2. Add --trust_remote_code argparse flag to subprocess CLI
3. Support TRUST_REMOTE_CODE=1 environment variable
4. Thread the setting through Model, tokenizer, and config calls
5. Remove torchrun fallback — error if torchrun is not found
6. Remove unnecessary try/except in is_gpt_oss_model
7. Remove redundant use_fast=True from tokenizer_utils

The env var is exported by main() when the flag is set, so
downstream calls (data_process, tokenizer_utils, gpt_oss_utils)
automatically pick it up without needing explicit parameter threading.

* Document trust_remote_code in README

Add trust_remote_code to the TrainingArgs table and document the
TRUST_REMOTE_CODE environment variable in the environment variables
section.

* Enable local mamba kernel pre-population for NemotronH models

NemotronH has Mamba layers just like GraniteMoeHybrid and needs
the same _use_local_mamba_kernels() call to avoid causal_conv1d_cuda
import failures in torchrun subprocesses.

* Fix lint: revert torchrun shutil.which, remove unused imports, ruff format

The torchrun-not-found issue was caused by the venv not being
activated, not an installation problem. Revert to plain 'torchrun'
command. Remove now-unused shutil and sys imports. Run ruff format
on all modified files.

* Clarify trust_remote_code docs with security warning

* Add FP8 dequantization and requantization for Ministral VLM training

Ministral-3-3B ships with FP8 quantized weights that include scalar
parameters (weight_scale_inv, activation_scale) which FSDP rejects.
This change dequantizes FP8 weights to bf16 after VLM extraction for
training compatibility, preserves the original scales, and requantizes
back to FP8 at checkpoint save time so saved checkpoints match the
original FP8 format.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Ruff formatting fixes

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Code <claude@anthropic.com>
Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
Co-authored-by: Mustafa Eyceoz <meyceoz@redhat.com>

Author: 3 people

README.md

@@ -244,6 +244,7 @@ for training jobs. There are a number of options you can specify, such as settin
 | distributed_backend | Specifies which distributed training backend to use. Supported options are "fsdp" and "deepspeed". |
 | disable_flash_attn | Disables flash attention when set to true. This allows for training on older devices. |
 | keep_last_checkpoint_only | Determines whether we should only keep the last checkpoint directory - the previous checkpoint directory is always overwritten. The checkpoint directory is called `last_epoch`. |
+| trust_remote_code | Controls whether repository-provided Python code from HuggingFace Hub is executed when loading models and tokenizers. This is required for models that ship custom modeling code, such as Nemotron, Ministral, and Qwen3.5. Can also be enabled via the `TRUST_REMOTE_CODE=1` environment variable. Defaults to `False`. **Security note:** enabling this setting will execute remote code from the model repository — only enable it for sources you trust. |
 
 ### `DeepSpeedOptions`
 
@@ -507,6 +508,7 @@ run_training(
 Below is a list of custom environment variables users can set in the training library.
 
 1. `INSTRUCTLAB_NCCL_TIMEOUT_MS`, this environment variable controls the NCCL timeout in milliseconds. Consider increasing if seeing FSDP related NCCL errors.
+2. `TRUST_REMOTE_CODE`, when set to `1`, allows repository-provided Python code from HuggingFace Hub to be executed when loading models and tokenizers. This is required for models that ship custom modeling code (e.g. Nemotron, Ministral, Qwen3.5). Equivalent to setting `trust_remote_code=True` in `TrainingArgs`. Only enable for sources you trust.
 
 ## Developer Certificate of Origin
 

src/instructlab/training/config.py

@@ -302,6 +302,16 @@ class TrainingArgs(BaseModel):
         description="Whether to use Liger kernels for training.",
     )
 
+    trust_remote_code: bool = Field(
+        default=False,
+        description=(
+            "Whether to trust remote code when loading models and tokenizers "
+            "from HuggingFace Hub. Required for models with custom code such as "
+            "Nemotron, Ministral, and Qwen3.5. Can also be enabled via the "
+            "TRUST_REMOTE_CODE=1 environment variable."
+        ),
+    )
+
     log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = Field(
         default="INFO"
     )

src/instructlab/training/data_process.py

@@ -41,11 +41,18 @@
 logger = logging.getLogger(__name__)
 
 
+def _trust_remote_code() -> bool:
+    """Resolve trust_remote_code from the TRUST_REMOTE_CODE environment variable."""
+    return os.environ.get("TRUST_REMOTE_CODE", "").lower() in ("1", "true", "yes")
+
+
 @lru_cache()
 def is_gpt_oss_model(tokenizer: PreTrainedTokenizer) -> bool:
     """Check if this is a GPT-OSS model based on tokenizer."""
     model_name_or_path = tokenizer.name_or_path
-    config = AutoConfig.from_pretrained(model_name_or_path)
+    config = AutoConfig.from_pretrained(
+        model_name_or_path, trust_remote_code=_trust_remote_code()
+    )
     return config.model_type == "gpt_oss"
 
 
@@ -1182,7 +1189,9 @@ def process_documents_for_pretraining(
         )
 
     logger.info("Loading tokenizer from %s", model_path)
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, trust_remote_code=_trust_remote_code()
+    )
 
     if tokenizer.eos_token_id is None:
         raise ValueError("Tokenizer must have an EOS token defined for pretraining")
@@ -1294,7 +1303,9 @@ def load_and_validate_dataset(data_path: str, num_procs: int) -> Dataset:
 
 def configure_tokenizer(model_path: str) -> PreTrainedTokenizer:
     """Configure the tokenizer with necessary special tokens."""
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, trust_remote_code=_trust_remote_code()
+    )
 
     if not tokenizer.chat_template:
         raise ValueError(

src/instructlab/training/gpt_oss_utils_correct.py

@@ -8,6 +8,7 @@
 # Standard
 from typing import Dict
 import logging
+import os
 import re
 
 # Third Party
@@ -415,7 +416,14 @@ def is_known_model(
     # convert to config
     model_config = model_path_or_config
     if isinstance(model_path_or_config, str):
-        model_config = AutoConfig.from_pretrained(model_path_or_config)
+        _trust_remote = os.environ.get("TRUST_REMOTE_CODE", "").lower() in (
+            "1",
+            "true",
+            "yes",
+        )
+        model_config = AutoConfig.from_pretrained(
+            model_path_or_config, trust_remote_code=_trust_remote
+        )
 
     known_model_types = (
         [known_model_type] if isinstance(known_model_type, str) else known_model_type

src/instructlab/training/main_ds.py

@@ -394,7 +394,17 @@ def main(args):
     tokenizer = setup_tokenizer(args.model_name_or_path, args.chat_tmpl_path)
     # device = torch.device("cuda", args.local_rank)
 
-    model_conf = AutoConfig.from_pretrained(args.model_name_or_path)
+    # Resolve trust_remote_code from CLI flag or environment variable
+    trust_remote_code = getattr(args, "trust_remote_code", False) or os.environ.get(
+        "TRUST_REMOTE_CODE", ""
+    ).lower() in ("1", "true", "yes")
+    if trust_remote_code:
+        # Export so downstream calls (data_process, tokenizer_utils, etc.) pick it up
+        os.environ["TRUST_REMOTE_CODE"] = "1"
+
+    model_conf = AutoConfig.from_pretrained(
+        args.model_name_or_path, trust_remote_code=trust_remote_code
+    )
     args.model_type = model_conf.model_type
 
     #### distributed init #####
@@ -449,6 +459,7 @@ def main(args):
         flash_enabled=flash_enabled,
         noise_alpha=args.NEFTune_alpha,
         lora_quant_bits=args.lora_quant_bits,
+        trust_remote_code=trust_remote_code,
     )
 
     args.base_model_args = m.base_model_args
@@ -710,6 +721,13 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None:
     if train_args.use_liger:
         command.append("--use_liger")
 
+    # Resolve trust_remote_code from flag or environment variable
+    trust_remote_code = train_args.trust_remote_code or os.environ.get(
+        "TRUST_REMOTE_CODE", ""
+    ).lower() in ("1", "true", "yes")
+    if trust_remote_code:
+        command.append("--trust_remote_code")
+
     if train_args.keep_last_checkpoint_only:
         command.append("--keep_last_checkpoint_only")
 
@@ -1036,6 +1054,15 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None:
         help="Path to the chat template to set on the model for training. If none is provided, the chat template used in the model will be used.",
     )
     parser.add_argument("--disable_flash_attn", action="store_true")
+    parser.add_argument(
+        "--trust_remote_code",
+        action="store_true",
+        help=(
+            "Trust remote code when loading models/tokenizers from HuggingFace Hub. "
+            "Required for models with custom code (e.g. Nemotron, Ministral, Qwen3.5). "
+            "Can also be set via the TRUST_REMOTE_CODE=1 environment variable."
+        ),
+    )
     parser.add_argument(
         "--keep_last_checkpoint_only",
         action="store_true",

src/instructlab/training/model.py

@@ -65,6 +65,7 @@ def __init__(
         flash_enabled: bool = False,
         lora_config: Optional[LoraConfig] = None,
         lora_quant_bits: int = 0,
+        trust_remote_code: bool = False,
     ):
         self.lora_config = lora_config
         self.noise_alpha = noise_alpha
@@ -74,12 +75,13 @@ def __init__(
 
         # check model type & set on the mclasss
         self.is_granitemoehybrid = is_known_model(model_path, "granitemoehybrid")
+        self.is_nemotronh = is_known_model(model_path, "nemotron_h")
         self.is_gpt_oss = is_gpt_oss(model_path)
 
         # Pre-populate the Hub kernel cache with locally installed mamba_ssm
         # and causal_conv1d to avoid PyTorch/CUDA ABI mismatches with the
         # Hub-provided kernel builds.
-        if self.is_granitemoehybrid:
+        if self.is_granitemoehybrid or self.is_nemotronh:
             self._use_local_mamba_kernels()
 
         if self.is_gpt_oss:
@@ -100,6 +102,7 @@ def __init__(
         self.base_model_args = {
             "pretrained_model_name_or_path": model_path,
             "quantization_config": quant_config,
+            "trust_remote_code": trust_remote_code,
         }
 
         # load GPT-OSS in bfloat16 because it's a massive model, but otherwise
@@ -114,7 +117,7 @@ def __init__(
             # - M-RoPE models produce 3D position_ids that FA2 misinterprets
             # - Models with timm vision towers (TimmWrapperModel rejects FA2)
             # Detect these and fall back to SDPA.
-            use_sdpa = needs_sdpa(model_path)
+            use_sdpa = needs_sdpa(model_path, trust_remote_code=trust_remote_code)
             if use_sdpa:
                 logger.warning(
                     "Disabling flash_attention_2 — model is incompatible "
@@ -145,7 +148,7 @@ def __init__(
             # For models with timm vision towers: set vision config to eager
             # while keeping the text model's attention implementation.
             # timm's TimmWrapperModel rejects both FA2 and SDPA.
-            if has_timm_vision_tower(model_path):
+            if has_timm_vision_tower(model_path, trust_remote_code=trust_remote_code):
                 attn_impl = self.base_model_args.get(
                     "attn_implementation", "flash_attention_2"
                 )
@@ -196,6 +199,18 @@ def _use_local_mamba_kernels():
                 e,
             )
 
+    def _enable_gradient_checkpointing_if_supported(self) -> None:
+        """Enable gradient checkpointing if the model supports it.
+
+        Some models (e.g. NemotronH with hybrid Mamba/MoE/Attention architecture)
+        do not support gradient checkpointing and will raise an error. This helper
+        catches known exception types and logs a warning instead of crashing.
+        """
+        try:
+            self.model.gradient_checkpointing_enable()
+        except (ValueError, NotImplementedError, AttributeError) as e:
+            logger.warning("Gradient checkpointing not supported: %s", e)
+
     def _post_model_init(self):
         """Common initialization steps that should happen after model initialization."""
         self.reconcile_tokenizer()
@@ -544,6 +559,7 @@ def __init__(
         flash_enabled: bool = False,
         lora_config: Optional[LoraConfig] = None,
         lora_quant_bits: int = 0,
+        trust_remote_code: bool = False,
     ):
         super().__init__(
             model_path=model_path,
@@ -553,6 +569,7 @@ def __init__(
             flash_enabled=flash_enabled,
             lora_config=lora_config,
             lora_quant_bits=lora_quant_bits,
+            trust_remote_code=trust_remote_code,
         )
         try:
             # Third Party
@@ -570,7 +587,7 @@ def __init__(
             cross_entropy=True,
             fused_linear_cross_entropy=False,
         )
-        self.model.gradient_checkpointing_enable()
+        self._enable_gradient_checkpointing_if_supported()
         self._post_model_init()
 
 
@@ -586,6 +603,7 @@ def __init__(
         flash_enabled: bool = False,
         lora_config: Optional[LoraConfig] = None,
         lora_quant_bits: int = 0,
+        trust_remote_code: bool = False,
     ):
         super().__init__(
             model_path=model_path,
@@ -595,15 +613,16 @@ def __init__(
             flash_enabled=flash_enabled,
             lora_config=lora_config,
             lora_quant_bits=lora_quant_bits,
+            trust_remote_code=trust_remote_code,
         )
-        if is_vlm_with_causal_lm(model_path):
+        if is_vlm_with_causal_lm(model_path, trust_remote_code=trust_remote_code):
             self.model = extract_causal_lm_from_vlm(model_path, self.base_model_args)
-        elif is_vlm_for_direct_loading(model_path):
+        elif is_vlm_for_direct_loading(model_path, trust_remote_code=trust_remote_code):
             self.model = load_vlm_for_text_training(model_path, self.base_model_args)
         else:
             self.model = AutoModelForCausalLM.from_pretrained(**self.base_model_args)
         self._post_model_init()
-        self.model.gradient_checkpointing_enable()
+        self._enable_gradient_checkpointing_if_supported()
 
 
 def setup_optimizer(

src/instructlab/training/tokenizer_utils.py

@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
+# Standard
+import os
+
 # Third Party
 from transformers import AutoTokenizer, PreTrainedTokenizer
 import transformers
@@ -95,7 +98,14 @@ def setup_tokenizer(
     model_name_or_path,
     chat_tmpl_path: str | None = None,
 ) -> PreTrainedTokenizer:
-    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, fast_tokenizer=True)
+    trust_remote_code = os.environ.get("TRUST_REMOTE_CODE", "").lower() in (
+        "1",
+        "true",
+        "yes",
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name_or_path, trust_remote_code=trust_remote_code
+    )
     if not tokenizer.chat_template and chat_tmpl_path is None:
         raise ValueError(
             "Tokenizer does not have a chat template. Please provide a path to a chat template."

src/instructlab/training/utils.py

@@ -713,8 +713,31 @@ def _get_state_dict_patched(model, unwrap=False):
         if is_gpt_oss(model.module.config):
             add_gpt_oss_quantization_config(model.module.config)
 
+        # For FP8 models (e.g. Ministral), restore the original quantization
+        # config so the saved checkpoint matches the original FP8 format.
+        # We must also temporarily remove _fp8_* attrs since they contain
+        # tensors that are not JSON-serializable.
+        _fp8_quant_cfg = getattr(model.module.config, "_fp8_quantization_config", None)
+        _had_fp8_scales = hasattr(model.module.config, "_fp8_scales")
+        if _fp8_quant_cfg is not None:
+            model.module.config.quantization_config = _fp8_quant_cfg
+        if _had_fp8_scales:
+            _saved_scales = model.module.config._fp8_scales
+            del model.module.config._fp8_scales
+            if hasattr(model.module.config, "_fp8_quantization_config"):
+                del model.module.config._fp8_quantization_config
+
         model.module.config.to_json_file(output_config_file)
 
+        # Restore internal attrs and clear quantization_config so the live
+        # model doesn't look quantized during training
+        if _had_fp8_scales:
+            model.module.config._fp8_scales = _saved_scales
+            if _fp8_quant_cfg is not None:
+                model.module.config._fp8_quantization_config = _fp8_quant_cfg
+        if _fp8_quant_cfg is not None:
+            model.module.config.quantization_config = None
+
         tokenizer.save_pretrained(output_dir)
 
         if is_lora:
@@ -759,7 +782,25 @@ def _get_state_dict_patched(model, unwrap=False):
                     max_shard_size="5GB",
                     safe_serialization=True,
                 )
-        elif not is_gpt_oss(model.module.config):
+        elif getattr(model.module.config, "_fp8_scales", None):
+            # FP8 model (e.g. Ministral): re-quantize state dict before saving
+            if accelerator.is_main_process:
+                # First Party
+                from instructlab.training.vlm_utils import requantize_fp8_state_dict
+
+                log_rank_0("Re-quantizing FP8 parameters for checkpoint compatibility")
+                model_state = model.module.state_dict()
+                model_state = requantize_fp8_state_dict(
+                    model_state, model.module.config._fp8_scales
+                )
+                save_dict_accelerate(
+                    accelerator,
+                    model_state,
+                    save_directory=output_dir,
+                    max_shard_size="5GB",
+                    safe_serialization=True,
+                )
+        else:
             # Standard model saving
             accelerator.save_model(
                 model,