feat: support epoch-based checkpoint intervals

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

@@ -275,6 +275,7 @@ ENV:
 - 🔥num_train_epochs: 训练的epoch数，默认为3。
 - 🔥save_strategy: 保存模型的策略，可选为'no'、'steps'、'epoch'，默认为'steps'。
 - 🔥save_steps: 默认为500。
+- save_epochs: 每N个epoch保存一次checkpoint，并自动使用`save_strategy='epoch'`。当启用epoch评估时，评估跟随相同的间隔。默认为None。
 - 🔥eval_strategy: 评估策略。默认为None，跟随`save_strategy`的策略。
   - 若不使用`val_dataset`和`eval_dataset`且`split_dataset_ratio`为0，则默认为'no'。
 - 🔥eval_steps: 默认为None，如果存在评估数据集，则跟随`save_steps`的策略。

docs/source/Megatron-SWIFT/Command-line-parameters.md

@@ -305,6 +305,7 @@ lora训练：
 - mrl_dims: Embedding训练的[Matryoshka表征学习（MRL）](https://arxiv.org/abs/2205.13147)维度配置，默认为None。格式为`Dict[int, float]`或Json字符串，key为截断的embedding维度，value为该维度对应的loss权重，例如`'{"32": 1.0, "64": 1.0, "128": 1.0}'`。开启后，trainer会对`last_hidden_state`分别截断到每个维度并做L2归一化，再调用`loss_type`对应的loss加权累加。仅在`task_type='embedding'`下生效。
   - 注意：可支持的最大embedding维度由模型`config.json`中的`hidden_size`决定，key大于`hidden_size`的KV对将被自动忽略。
 - 🔥save_strategy: 保存策略，可选项为'steps'和'epoch'。默认为'steps'。当设置为'epoch'时，会根据数据集大小自动计算`save_steps`和`eval_steps`以实现每个epoch保存一次，用户传入的`save_steps`和`eval_steps`参数值将被忽略。
+- save_epochs: 每N个epoch保存一次checkpoint，并自动使用`save_strategy='epoch'`，评估跟随相同的间隔。默认为None。
 - callbacks: 自定义trainer callback，默认为`[]`。
 
 ## 训练参数

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

@@ -280,6 +280,8 @@ Other important parameters:
 - 🔥num_train_epochs: Number of training epochs. Default is 3.
 - 🔥save_strategy: Strategy for saving checkpoints. Options: `'no'`, `'steps'`, `'epoch'`. Default is `'steps'`.
 - 🔥save_steps: Default is 500.
+- save_epochs: Save a checkpoint every N epochs and automatically use `save_strategy='epoch'`. When epoch-based
+  evaluation is enabled, evaluation follows the same interval. Defaults to None.
 - 🔥eval_strategy: Evaluation strategy. Default is `None`, following `save_strategy`.
   - If neither `val_dataset` nor `eval_dataset` is used and `split_dataset_ratio=0`, defaults to `'no'`.
 - 🔥eval_steps: Default is `None`. If evaluation dataset exists, follows `save_steps`.

docs/source_en/Megatron-SWIFT/Command-line-parameters.md

@@ -320,6 +320,8 @@ LoRA Training:
 - mrl_dims: Dimension configuration for [Matryoshka Representation Learning (MRL)](https://arxiv.org/abs/2205.13147) on embedding training. Default is None. Format is `Dict[int, float]` or a JSON string, where the key is the truncated embedding dimension and the value is the corresponding loss weight, e.g. `'{"32": 1.0, "64": 1.0, "128": 1.0}'`. When enabled, the trainer slices `last_hidden_state` to each dimension, applies L2 normalization, and aggregates the per-dimension `loss_type` losses with the configured weights. Only effective when `task_type='embedding'`.
   - Note: The maximum supported embedding dimension is determined by `hidden_size` in the model's `config.json`. Any key-value pair whose key is greater than `hidden_size` will be silently ignored.
 - 🔥save_strategy: Saving strategy, options are 'steps' and 'epoch'. Defaults to 'steps'. When set to 'epoch', `save_steps` and `eval_steps` are automatically calculated to save at each epoch, so any user-provided values for these arguments are ignored.
+- save_epochs: Save a checkpoint every N epochs and automatically use `save_strategy='epoch'`. Evaluation follows the
+  same interval. Defaults to None.
 - callbacks: Custom trainer callbacks. Defaults to `[]`.
 
 

swift/arguments/sft_args.py

@@ -241,6 +241,9 @@ def _init_override(self):
         self._init_output_dir()
         self._init_metric()
 
+        if self.save_epochs is not None:
+            self.save_strategy = 'epoch'
+
         if self.learning_rate is None:
             if self.tuner_type == 'full':
                 self.learning_rate = 1e-5

swift/megatron/arguments/megatron_args.py

@@ -561,6 +561,7 @@ class MegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin):
     # checkpoint
     output_dir: Optional[str] = None
     save_steps: int = 500
+    save_epochs: Optional[int] = None
     no_save_optim: bool = False
     no_save_rng: bool = False
     mcore_model: Optional[str] = None
@@ -788,6 +789,10 @@ def _check_bridge_backend(self):
             self._check_mcore_bridge()
 
     def __post_init__(self):
+        if self.save_epochs is not None:
+            if self.save_epochs < 1:
+                raise ValueError('`save_epochs` must be greater than or equal to 1.')
+            self.save_strategy = 'epoch'
         if self.tuner_type != 'full':
             require_version('peft>=0.15', 'Please install peft>=0.15 to use LoRA in Megatron-SWIFT.')
         RLHFMegatronArgumentsMixin.__post_init__(self)
@@ -1043,7 +1048,7 @@ def init_iters(self, train_dataset, val_dataset):
         if self.save_strategy == 'epoch':
             if hasattr(train_dataset, '__len__'):
                 dataset_sample = len(train_dataset) // step_batch_size * step_batch_size * num_generations
-                self.save_steps = dataset_sample // self.global_batch_size
+                self.save_steps = dataset_sample // self.global_batch_size * (self.save_epochs or 1)
                 self.eval_steps = self.save_steps
             else:
                 raise ValueError('streaming dataset is not supported with `--save_strategy epoch`.')

swift/ray/megatron/driver_utils.py

@@ -209,6 +209,7 @@ def build_dataset_from_dict(cfg: Dict[str, Any]):
         'num_train_epochs': args.num_train_epochs,
         'train_iters': args.train_iters,
         'save_strategy': args.save_strategy,
+        'save_epochs': args.save_epochs,
         'eval_iters': args.eval_iters,
         'num_generations': args.num_generations,
         'template': template,
@@ -254,10 +255,13 @@ def compute_iter_params(data_info: Dict[str, Any], dp_size: int) -> Dict[str, An
 
     result: Dict[str, Any] = {}
 
-    if data_info.get('save_strategy') == 'epoch' and train_len > 0:
-        ds_sample = train_len // step_batch_size * step_batch_size * num_gen
-        result['save_steps'] = ds_sample // gbs
-        result['eval_steps'] = result['save_steps']
+    if data_info.get('save_strategy') == 'epoch':
+        if data_info.get('save_epochs') is not None and train_ds is not None and not hasattr(train_ds, '__len__'):
+            raise ValueError('streaming dataset is not supported with `--save_strategy epoch`.')
+        if train_len > 0:
+            ds_sample = train_len // step_batch_size * step_batch_size * num_gen
+            result['save_steps'] = ds_sample // gbs * (data_info.get('save_epochs') or 1)
+            result['eval_steps'] = result['save_steps']
 
     train_iters = data_info.get('train_iters')
     if data_info.get('num_train_epochs') is not None and train_len > 0:

swift/trainers/arguments.py

@@ -61,6 +61,8 @@ class TrainArgumentsMixin:
             training dataset (with a random factor).
         max_epochs (Optional[int]): The total number of training epochs to perform. Overrides `num_train_epochs`.
             Defaults to None.
+        save_epochs (Optional[int]): Save a checkpoint every `save_epochs` epochs. When set, `save_strategy` is
+            automatically set to 'epoch'. Defaults to None.
         aligner_lr (Optional[float]): A specific learning rate for the aligner part of the model. Defaults to None.
         vit_lr (Optional[float]): A specific learning rate for the Vision Transformer part of the model. Defaults to
             None.
@@ -154,6 +156,7 @@ class TrainArgumentsMixin:
     train_dataloader_shuffle: bool = True
     group_by_length: bool = False
     max_epochs: Optional[int] = None
+    save_epochs: Optional[int] = None
     aligner_lr: Optional[float] = None
     vit_lr: Optional[float] = None
     use_logits_to_keep: Optional[bool] = None
@@ -240,6 +243,10 @@ def _init_callbacks(self):
             self.callbacks.append('activation_cpu_offload')
 
     def __post_init__(self):
+        if self.save_epochs is not None:
+            if self.save_epochs < 1:
+                raise ValueError('`save_epochs` must be greater than or equal to 1.')
+            self.save_strategy = 'epoch'
         if hasattr(self, 'output_dir'):
             self.output_dir = os.path.abspath(os.path.expanduser(self.output_dir))
         if is_mp() and self.use_liger_kernel:

swift/trainers/patcher.py

@@ -74,6 +74,13 @@ def on_step_end(self, args: TrainingArguments, state: TrainerState, control: Tra
     def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
         control = super().on_epoch_end(args, state, control, **kwargs)
         evaluation_strategy = args.eval_strategy if hasattr(args, 'eval_strategy') else args.evaluation_strategy
+        save_epochs = getattr(args, 'save_epochs', None)
+        if save_epochs is not None:
+            epoch = math.ceil(state.epoch)
+            should_save = epoch > 0 and epoch % save_epochs == 0
+            control.should_save = should_save
+            if evaluation_strategy == IntervalStrategy.EPOCH:
+                control.should_evaluate = should_save
         if args.max_epochs is not None and args.max_epochs <= math.ceil(state.epoch):
             logger.info('Training has reached `max_epochs`. The model will be saved and the training will be exited.')
             if evaluation_strategy != IntervalStrategy.NO:

swift/ui/llm_train/hyper.py

@@ -102,6 +102,16 @@ class Hyper(BaseUI):
                 'en': 'Set the save steps',
             }
         },
+        'save_epochs': {
+            'label': {
+                'zh': '存储轮次',
+                'en': 'Save epochs',
+            },
+            'info': {
+                'zh': '设置每隔多少轮进行一次存储',
+                'en': 'Set the epoch interval to save',
+            }
+        },
         'output_dir': {
             'label': {
                 'zh': '存储目录',
@@ -133,6 +143,7 @@ def do_build_ui(cls, base_tab: Type['BaseUI']):
                 with gr.Row():
                     gr.Textbox(elem_id='eval_steps', lines=1, value='500', scale=20)
                     gr.Textbox(elem_id='save_steps', value='500', lines=1, scale=20)
+                    gr.Textbox(elem_id='save_epochs', lines=1, scale=20)
                     gr.Textbox(elem_id='output_dir', scale=20)
                     gr.Dropdown(
                         elem_id='attn_impl',

tests/train/test_save_epochs.py

@@ -0,0 +1,110 @@
+import unittest
+from types import SimpleNamespace
+
+from transformers.trainer_callback import TrainerControl, TrainerState
+from transformers.trainer_utils import IntervalStrategy
+
+
+class TestSaveEpochs(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        try:
+            from swift.ray.megatron.driver_utils import compute_iter_params
+            from swift.trainers.patcher import DefaultFlowCallbackNew
+            cls.compute_iter_params = compute_iter_params
+            cls.DefaultFlowCallbackNew = DefaultFlowCallbackNew
+        except (ImportError, RuntimeError) as e:
+            raise unittest.SkipTest(f'ms-swift dependencies are not available: {e}')
+
+    @staticmethod
+    def _training_args(**kwargs):
+        args = dict(
+            eval_strategy=IntervalStrategy.EPOCH,
+            logging_strategy=IntervalStrategy.NO,
+            eval_delay=0,
+            save_strategy=IntervalStrategy.EPOCH,
+            save_epochs=2,
+            max_epochs=None,
+        )
+        args.update(kwargs)
+        return SimpleNamespace(**args)
+
+    def test_hf_callback_aligns_epoch_evaluation_with_saves(self):
+        callback = self.DefaultFlowCallbackNew()
+        args = self._training_args()
+
+        first_epoch = callback.on_epoch_end(args, TrainerState(epoch=1), TrainerControl())
+        self.assertFalse(first_epoch.should_save)
+        self.assertFalse(first_epoch.should_evaluate)
+
+        second_epoch = callback.on_epoch_end(args, TrainerState(epoch=2), TrainerControl())
+        self.assertTrue(second_epoch.should_save)
+        self.assertTrue(second_epoch.should_evaluate)
+
+    def test_ray_compute_iter_params_uses_save_epochs(self):
+        data_info = {
+            'micro_batch_size': 2,
+            'global_batch_size': 8,
+            'num_generations': 1,
+            'train_dataset': list(range(100)),
+            'val_dataset': None,
+            'save_strategy': 'epoch',
+            'save_epochs': 3,
+            'train_iters': None,
+            'num_train_epochs': 2,
+            'eval_iters': 0,
+        }
+        result = self.compute_iter_params(data_info, dp_size=1)
+        self.assertEqual(result['save_steps'], 36)
+        self.assertEqual(result['eval_steps'], 36)
+        self.assertEqual(result['train_iters'], 25)
+
+    def test_ray_streaming_save_epochs_is_rejected(self):
+
+        class StreamingDataset:
+
+            def __iter__(self):
+                return iter(())
+
+        data_info = {
+            'micro_batch_size': 2,
+            'global_batch_size': 8,
+            'num_generations': 1,
+            'train_dataset': StreamingDataset(),
+            'val_dataset': None,
+            'save_strategy': 'epoch',
+            'save_epochs': 2,
+            'train_iters': 10,
+            'num_train_epochs': None,
+            'eval_iters': 0,
+        }
+        with self.assertRaisesRegex(ValueError, 'streaming dataset'):
+            self.compute_iter_params(data_info, dp_size=1)
+
+    def test_ray_existing_streaming_epoch_strategy_remains_compatible(self):
+
+        class StreamingDataset:
+
+            def __iter__(self):
+                return iter(())
+
+        data_info = {
+            'micro_batch_size': 2,
+            'global_batch_size': 8,
+            'num_generations': 1,
+            'train_dataset': StreamingDataset(),
+            'val_dataset': None,
+            'save_strategy': 'epoch',
+            'save_epochs': None,
+            'train_iters': 10,
+            'num_train_epochs': None,
+            'eval_iters': 0,
+        }
+        result = self.compute_iter_params(data_info, dp_size=1)
+        self.assertEqual(result['train_iters'], 10)
+        self.assertNotIn('save_steps', result)
+
+
+if __name__ == '__main__':
+    unittest.main()