[fsdp] fix: Support trust_remote_code during FSDP HugginFace checkpoint save

.pre-commit-config.yaml

@@ -40,6 +40,6 @@ repos:
     hooks:
       - id: compileall
         name: Compile all python files
-        entry: sh -c 'PYTHONWARNINGS=error python3 -m compileall -q .'
+        entry: sh -c 'PYTHONWARNINGS=error python3 -m compileall -q . -x "\.venv|venv|\.git"'
         language: python
         pass_filenames: false

verl/experimental/vla/fsdp_workers.py

@@ -305,6 +305,7 @@ def init_model(self):
                 lr_scheduler=self.actor_lr_scheduler,
                 processing_class=self.processor if self.processor is not None else self.tokenizer,
                 checkpoint_config=self.config.actor.checkpoint,
+                trust_remote_code=self.config.model.trust_remote_code,
             )
 
         torch.distributed.barrier()

verl/trainer/fsdp_sft_trainer.py

@@ -609,6 +609,7 @@ def _init_checkpoint_manager(self):
             lr_scheduler=self.lr_scheduler,
             processing_class=self.tokenizer,
             checkpoint_config=checkpoint_config_dict,
+            trust_remote_code=self.config.model.trust_remote_code,
         )
 
     def load_checkpoint(self):

verl/utils/checkpoint/fsdp_checkpoint_manager.py

@@ -70,6 +70,7 @@ class FSDPCheckpointManager(BaseCheckpointManager):
         checkpoint_contents DictConfig: Configuration for checkpoint contents.
             - 'load': Components to load; must contain 'model'. Defaults to ['model', 'optimizer', 'extra'].
             - 'save': Components to save; must contain 'model'. Defaults to ['model', 'optimizer', 'extra'].
+        trust_remote_code: Whether to trust_remote_code when loading the model configuration
     """
 
     def __init__(
@@ -79,6 +80,7 @@ def __init__(
         lr_scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,
         processing_class: PreTrainedTokenizer | ProcessorMixin = None,
         checkpoint_config: DictConfig = None,
+        trust_remote_code: bool = False,
         **kwargs,
     ):
         if processing_class is None and "tokenizer" in kwargs:
@@ -94,6 +96,7 @@ def __init__(
             processing_class=processing_class,
             checkpoint_config=checkpoint_config,
         )
+        self.trust_remote_code = trust_remote_code
 
     def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_after_load=False):
         """
@@ -333,7 +336,10 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i
                     raise NotImplementedError(f"Unknown architecture {model_config['architectures']}")
 
                 with init_empty_weights():
-                    save_model = auto_model_cls.from_config(model_config, torch_dtype=torch.bfloat16)
+                    save_model = auto_model_cls.from_config(
+                        model_config, torch_dtype=torch.bfloat16, trust_remote_code=self.trust_remote_code
+                    )
+
                 save_model.to_empty(device="cpu")
 
                 if save_model.can_generate():

verl/workers/engine/fsdp/transformer_impl.py

@@ -164,6 +164,7 @@ def initialize(self):
             lr_scheduler=self.lr_scheduler,
             processing_class=self.model_config.get_processor(),
             checkpoint_config=self.checkpoint_config,
+            trust_remote_code=self.model_config.trust_remote_code,
         )
 
         self.to(

verl/workers/engine/veomni/transformer_impl.py

@@ -134,6 +134,7 @@ def initialize(self):
             lr_scheduler=self.lr_scheduler,
             processing_class=self.model_config.get_processor(),
             checkpoint_config=self.checkpoint_config,
+            trust_remote_code=self.model_config.trust_remote_code,
         )
 
         self.to(

verl/workers/fsdp_workers.py

@@ -873,6 +873,7 @@ def init_model(self):
                 lr_scheduler=self.actor_lr_scheduler,
                 processing_class=self.processor if self.processor is not None else self.tokenizer,
                 checkpoint_config=self.config.actor.checkpoint,
+                trust_remote_code=self.config.model.get("trust_remote_code", False),
             )
 
         if not self._is_actor and self._is_rollout:
@@ -1255,7 +1256,7 @@ def __init__(self, config: FSDPCriticConfig):
         )
         self.use_orig_params = self.config.model.fsdp_config.get("use_orig_params", False)
 
-    def _build_critic_model_optimizer(self, config):
+    def _build_critic_model_optimizer(self, config: FSDPCriticConfig):
         # the following line is necessary
         from torch.distributed.fsdp import MixedPrecision
 
@@ -1533,6 +1534,7 @@ def init_model(self):
             lr_scheduler=self.critic_lr_scheduler,
             processing_class=self.processor if self.processor is not None else self.tokenizer,
             checkpoint_config=self.config.checkpoint,
+            trust_remote_code=self.config.model.get("trust_remote_code", False),
         )
 
     @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name="critic"))