[bugfix] fix max_shard_size transformers 5.x (#8209)

Author: Jintao-Huang

docs/source/Instruction/Command-line-parameters.md

@@ -219,8 +219,10 @@ gradient_checkpointing: true
 - router_aux_loss_coef: 用于moe模型训练时，设置 aux_loss 的权重，默认为`0.`。
 - enable_dft_loss: 是否在SFT训练中使用[DFT](https://arxiv.org/abs/2508.05629) (Dynamic Fine-Tuning) loss，默认为False。
 - enable_channel_loss: 启用channel loss，默认为`False`。你需要在数据集中准备"channel"字段，ms-swift会根据该字段分组统计loss（若未准备"channel"字段，则归为默认`None` channel）。数据集格式参考[channel loss](../Customization/Custom-dataset.md#channel-loss)。channel loss兼容packing/padding_free/loss_scale等技术。
+- safe_serialization: 是否存储为safetensors，默认为True。
+- max_shard_size: 单存储文件最大大小，默认'5GB'。
 - logging_dir: tensorboard日志保存路径。默认为None，即设置为`f'{self.output_dir}/runs'`。
-- 🔥predict_with_generate: 验证时使用生成式的方式，默认为False。
+- predict_with_generate: 验证时使用生成式的方式，默认为False。
 - metric_for_best_model: 默认为None，即当`predict_with_generate`设置为False时，设置为'loss'，否则设置为'rouge-l'（在PPO训练时，不进行默认值设置；GRPO训练设置为'reward'）。
 - greater_is_better: 默认为None，即当`metric_for_best_model`含'loss'时，设置为False，否则设置为True。
 - max_epochs: 训练到`max_epochs`时强制退出训练，并对权重进行验证和保存。该参数在使用流式数据集时很有用。默认为None。

docs/source_en/Instruction/Command-line-parameters.md

@@ -222,6 +222,8 @@ This list inherits from the Transformers `Seq2SeqTrainingArguments`, with ms-swi
 - router_aux_loss_coef: Used in MoE model training to set the weight of auxiliary loss. Default is `0.`.
 - enable_dft_loss: Whether to use [DFT](https://arxiv.org/abs/2508.05629) (Dynamic Fine-Tuning) loss during SFT training. Default is `False`.
 - enable_channel_loss: Enable channel-based loss. Default is `False`. Requires a `"channel"` field in the dataset. ms-swift groups and computes loss by this field (samples without `"channel"` are grouped into the default `None` channel). Dataset format reference: [channel loss](../Customization/Custom-dataset.md#channel-loss).  Channel loss is compatible with packing, padding_free, and loss_scale techniques.
+- safe_serialization: Whether to save the model in safetensors format. Default is True.
+- max_shard_size: Maximum size of a single storage file, default is '5GB'.
 - logging_dir: Directory for TensorBoard logs. Default is `None`, automatically set to `f'{self.output_dir}/runs'`.
 - predict_with_generate: Use generation during evaluation. Default is `False`.
 - metric_for_best_model: Default is `None`. If `predict_with_generate=False`, it's set to `'loss'`; otherwise `'rouge-l'` (in PPO training, no default; in GRPO, set to `'reward'`).

swift/trainers/arguments.py

@@ -135,6 +135,8 @@ class TrainArgumentsMixin:
     router_aux_loss_coef: float = 0.
     enable_dft_loss: bool = False  # https://arxiv.org/abs/2508.05629
     enable_channel_loss: bool = False
+    safe_serialization: bool = True
+    max_shard_size: str = '5GB'
 
     weight_decay: float = 0.1
     adam_beta2: float = 0.95

swift/trainers/mixin.py

@@ -54,11 +54,6 @@
 from .utils import (can_return_loss, dynamic_gradient_checkpointing, find_labels, get_function, get_resume_dir,
                     is_instance_of_ms_model, patch_modelscope_hub_timeout, replace_index_file)
 
-try:
-    from trl import AutoModelForCausalLMWithValueHead
-except (ImportError, RuntimeError):
-    AutoModelForCausalLMWithValueHead = None
-
 logger = get_logger()
 
 
@@ -275,9 +270,7 @@ def _save_model(self, output_dir: Optional[str] = None, state_dict=None):
         # model
         supported_classes = (SwiftModel, PreTrainedModel, PeftModel)
         supported_names = ('SentenceTransformer', )
-        if AutoModelForCausalLMWithValueHead is not None:
-            supported_classes = supported_classes + (AutoModelForCausalLMWithValueHead, )
-        save_safetensors = getattr(self.args, 'save_safetensors', True)
+        safe_serialization = self.args.safe_serialization
         use_flash_ckpt = self.args.use_flash_ckpt
 
         if not isinstance(self.model, supported_classes) and self.model.__class__.__name__ not in supported_names:
@@ -286,7 +279,7 @@ def _save_model(self, output_dir: Optional[str] = None, state_dict=None):
 
             _unwrap_model = unwrap_model(self.model)
             if isinstance(_unwrap_model, supported_classes):
-                save_kwargs = {'state_dict': state_dict}
+                save_kwargs = {'state_dict': state_dict, 'max_shard_size': self.args.max_shard_size}
                 if isinstance(_unwrap_model, PeftModel):
                     save_kwargs['selected_adapters'] = ['default']
                 if use_flash_ckpt:
@@ -296,33 +289,16 @@ def _save_model(self, output_dir: Optional[str] = None, state_dict=None):
                         save_function=self.flash_checkpointer.ckpt_agent.save,
                         **save_kwargs)
                 else:
-                    _unwrap_model.save_pretrained(output_dir, safe_serialization=save_safetensors, **save_kwargs)
+                    _unwrap_model.save_pretrained(output_dir, safe_serialization=safe_serialization, **save_kwargs)
             else:
                 logger.info('Trainer.model is not a `PreTrainedModel`, only saving its state dict.')
                 if use_flash_ckpt:
                     self.flash_checkpointer.ckpt_agent.save(state_dict, os.path.join(output_dir, 'pytorch_model.bin'))
                 else:
-                    if save_safetensors:
+                    if safe_serialization:
                         safetensors.torch.save_file(state_dict, os.path.join(output_dir, 'model.safetensors'))
                     else:
                         torch.save(state_dict, os.path.join(output_dir, 'pytorch_model.bin'))
-        elif AutoModelForCausalLMWithValueHead and isinstance(self.model, AutoModelForCausalLMWithValueHead):
-            # save reward model
-            state_dict = self.model.state_dict()
-            decoder_state_dict, v_head_state_dict = {}, {}
-            for name, param in state_dict.items():
-                if name.startswith('v_head.'):
-                    v_head_state_dict[name] = param
-                else:
-                    decoder_state_dict[name.replace('pretrained_model.', '', 1)] = param
-            self.model.pretrained_model.save_pretrained(
-                output_dir, state_dict=decoder_state_dict or None, safe_serialization=save_safetensors)
-            if save_safetensors:
-                from safetensors.torch import save_file
-                save_file(
-                    v_head_state_dict, os.path.join(output_dir, 'value_head.safetensors'), metadata={'format': 'pt'})
-            else:
-                torch.save(v_head_state_dict, os.path.join(output_dir, 'value_head.bin'))
         elif is_instance_of_ms_model(self.model):
             if use_flash_ckpt:
                 PreTrainedModel.save_pretrained(
@@ -334,13 +310,13 @@ def _save_model(self, output_dir: Optional[str] = None, state_dict=None):
             else:
                 # modelscope save_pretrained does not support safe_serialization
                 PreTrainedModel.save_pretrained(
-                    self.model, output_dir, state_dict=state_dict, safe_serialization=save_safetensors)
+                    self.model, output_dir, state_dict=state_dict, safe_serialization=safe_serialization)
         elif self.args.tuner_type in tuners_map:
             tuners_map[self.args.tuner_type].save_pretrained(
-                self.model, output_dir, state_dict=state_dict, safe_serialization=save_safetensors)
+                self.model, output_dir, state_dict=state_dict, safe_serialization=safe_serialization)
         else:
             if self.model.__class__.__name__ != 'SentenceTransformer':
-                save_kwargs = {'state_dict': state_dict}
+                save_kwargs = {'state_dict': state_dict, 'max_shard_size': self.args.max_shard_size}
                 if isinstance(self.model, PeftModel):
                     save_kwargs['selected_adapters'] = ['default']
                 if use_flash_ckpt:
@@ -350,7 +326,7 @@ def _save_model(self, output_dir: Optional[str] = None, state_dict=None):
                         save_function=self.flash_checkpointer.ckpt_agent.save,
                         **save_kwargs)
                 else:
-                    self.model.save_pretrained(output_dir, safe_serialization=save_safetensors, **save_kwargs)
+                    self.model.save_pretrained(output_dir, safe_serialization=safe_serialization, **save_kwargs)
             else:
 
                 @contextmanager
@@ -373,7 +349,7 @@ def save_context():
                             safe_serialization=False,
                             save_function=self.flash_checkpointer.ckpt_agent.save)
                     else:
-                        self.model.save_pretrained(output_dir, safe_serialization=save_safetensors)
+                        self.model.save_pretrained(output_dir, safe_serialization=safe_serialization)
                         # copy sentencetransformers files
                     copy_files_by_pattern(
                         self.model.model_dir, output_dir, '*.py', exclude_patterns=['model.safetensors.index.json'])
@@ -636,7 +612,7 @@ def _save_flash_checkpoint(self, model, trial, metrics=None):
                 rng_states,
                 os.path.join(output_dir, f'rng_state_{self.args.process_index}.pth'),
             )
-        if self.args.save_safetensors:
+        if self.args.safe_serialization:
             torch.save({'safe_serialization': True}, 'safe_serialization')
             replace_index_file(output_dir)