From 5e9f6a6690e6d0ba3959de97469dd99f1938e5a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=99=88=E7=AB=A5?= <31961076+tongchen126@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:33:05 +0800 Subject: [PATCH] feat: support epoch-based checkpoint intervals Co-authored-by: andrewwchen --- .../Instruction/Command-line-parameters.md | 1 + .../Megatron-SWIFT/Command-line-parameters.md | 1 + .../Instruction/Command-line-parameters.md | 2 + .../Megatron-SWIFT/Command-line-parameters.md | 2 + swift/arguments/sft_args.py | 3 + swift/megatron/arguments/megatron_args.py | 7 +- swift/ray/megatron/driver_utils.py | 12 +- swift/trainers/arguments.py | 7 ++ swift/trainers/patcher.py | 7 ++ swift/ui/llm_train/hyper.py | 11 ++ tests/train/test_save_epochs.py | 110 ++++++++++++++++++ 11 files changed, 158 insertions(+), 5 deletions(-) create mode 100644 tests/train/test_save_epochs.py diff --git a/docs/source/Instruction/Command-line-parameters.md b/docs/source/Instruction/Command-line-parameters.md index cba43833e4..51dcb4d425 100644 --- a/docs/source/Instruction/Command-line-parameters.md +++ b/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`的策略。 diff --git a/docs/source/Megatron-SWIFT/Command-line-parameters.md b/docs/source/Megatron-SWIFT/Command-line-parameters.md index 19a4642105..009cbf3896 100644 --- a/docs/source/Megatron-SWIFT/Command-line-parameters.md +++ b/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，默认为`[]`。 ## 训练参数 diff --git a/docs/source_en/Instruction/Command-line-parameters.md b/docs/source_en/Instruction/Command-line-parameters.md index 9302c55f7b..bf70832f58 100644 --- a/docs/source_en/Instruction/Command-line-parameters.md +++ b/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`. diff --git a/docs/source_en/Megatron-SWIFT/Command-line-parameters.md b/docs/source_en/Megatron-SWIFT/Command-line-parameters.md index 3ab71573a6..6fc6e77867 100644 --- a/docs/source_en/Megatron-SWIFT/Command-line-parameters.md +++ b/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 `[]`. diff --git a/swift/arguments/sft_args.py b/swift/arguments/sft_args.py index cb7abc98f3..401c49c6ce 100644 --- a/swift/arguments/sft_args.py +++ b/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 diff --git a/swift/megatron/arguments/megatron_args.py b/swift/megatron/arguments/megatron_args.py index f7f0115bd3..84db69774a 100644 --- a/swift/megatron/arguments/megatron_args.py +++ b/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`.') diff --git a/swift/ray/megatron/driver_utils.py b/swift/ray/megatron/driver_utils.py index 642b3cc60d..2ba676263a 100644 --- a/swift/ray/megatron/driver_utils.py +++ b/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: diff --git a/swift/trainers/arguments.py b/swift/trainers/arguments.py index 034595ce7a..27472da322 100644 --- a/swift/trainers/arguments.py +++ b/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: diff --git a/swift/trainers/patcher.py b/swift/trainers/patcher.py index a18071647e..0e81cfa3a3 100644 --- a/swift/trainers/patcher.py +++ b/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: diff --git a/swift/ui/llm_train/hyper.py b/swift/ui/llm_train/hyper.py index 3f27528af9..9ff6261060 100644 --- a/swift/ui/llm_train/hyper.py +++ b/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', diff --git a/tests/train/test_save_epochs.py b/tests/train/test_save_epochs.py new file mode 100644 index 0000000000..35d99e12b4 --- /dev/null +++ b/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()