From 8df52f0eaeb7ac0273f3724ba799dce7620d655a Mon Sep 17 00:00:00 2001 From: tongruiliu <952814125@qq.com> Date: Tue, 12 May 2026 16:40:38 +0800 Subject: [PATCH 1/3] support num_train_epoches while keep dynamic-step training --- examples/test/test_mix.yaml | 2 +- examples/test/test_select.yaml | 2 +- examples/test/test_static_mix.yaml | 2 +- examples/test/test_weight.yaml | 2 +- .../doremi_step1_static_qwen_pt_full.yaml | 2 +- .../doremi_step2_dynamic_qwen_pt_full.yaml | 2 +- .../doremi_step3_static_qwen_pt_full.yaml | 2 +- .../mixers/odm_dynamic_qwen_pt_full.yaml | 4 +- .../doremi_step2_dynamic_qwen_pt_lora.yaml | 2 +- examples/train_lora/mixers/random.yaml | 2 +- examples/train_lora/selectors/custom.yaml | 4 +- examples/train_lora/selectors/delta_loss.yaml | 4 +- examples/train_lora/selectors/less.yaml | 4 +- examples/train_lora/selectors/loss.yaml | 4 +- examples/train_lora/selectors/near.yaml | 4 +- examples/train_lora/selectors/nice.yaml | 4 +- examples/train_lora/selectors/random.yaml | 4 +- examples/train_lora/selectors/tsds.yaml | 2 +- examples/train_lora/weighters/custom.yaml | 2 +- examples/train_lora/weighters/loss.yaml | 2 +- src/dataflex/train/hparams/dynamic_params.py | 6 +- src/dataflex/train/trainer/mix_trainer.py | 94 +++++++++++++------ src/dataflex/train/trainer/select_trainer.py | 69 +++++++++++--- src/dataflex/train/trainer/weight_trainer.py | 27 ++++-- 24 files changed, 170 insertions(+), 82 deletions(-) diff --git a/examples/test/test_mix.yaml b/examples/test/test_mix.yaml index acbbee3..f4e849e 100644 --- a/examples/test/test_mix.yaml +++ b/examples/test/test_mix.yaml @@ -53,4 +53,4 @@ mixture_sample_rule: mixture # 初始采样规则，mixture为根据init_mix init_mixture_proportions: [0.7, 0.3] # 对应初始的比例 warmup_step: 4 update_step: 3 -update_times: 2 \ No newline at end of file +update_times: 2 # Updates per Flex epoch; repeats by num_train_epochs diff --git a/examples/test/test_select.yaml b/examples/test/test_select.yaml index 2fdf4c8..59a318a 100644 --- a/examples/test/test_select.yaml +++ b/examples/test/test_select.yaml @@ -60,4 +60,4 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: delta_loss # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 4 update_step: 3 -update_times: 4 \ No newline at end of file +update_times: 4 # Updates per Flex epoch; repeats by num_train_epochs diff --git a/examples/test/test_static_mix.yaml b/examples/test/test_static_mix.yaml index 4a59301..6821cc3 100644 --- a/examples/test/test_static_mix.yaml +++ b/examples/test/test_static_mix.yaml @@ -50,4 +50,4 @@ train_type: dynamic_mix # 选择训练器类型。可选值包括： static_mix: true # 是否固定初始静态混合比例（仅在dynamic_mix训练器中生效） mixture_sample_rule: mixture # 初始采样规则，mixture为根据init_mixture_proportions比例混合（可动态调整），stratified为固定按源数据集大小比例分层，uniform为固定均匀分布 init_mixture_proportions: [0.7, 0.3] # 对应初始的比例，可通过额外算法自行调整 -train_step: 3 # 总训练步数（仅在dynamic_mix训练器中生效），不考虑warmup和update steps \ No newline at end of file +train_step: 3 # Total steps; overrides num_train_epochs diff --git a/examples/test/test_weight.yaml b/examples/test/test_weight.yaml index 49c28bc..199d435 100644 --- a/examples/test/test_weight.yaml +++ b/examples/test/test_weight.yaml @@ -59,4 +59,4 @@ train_type: dynamic_weight # 选择训练器类型。可选值包括： components_cfg_file: src/dataflex/configs/components.yaml component_name: custom # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 1 -train_step: 3 # 总训练步数（包括warm_up） \ No newline at end of file +train_step: 3 # Total steps; overrides num_train_epochs diff --git a/examples/train_full/mixers/doremi_step1_static_qwen_pt_full.yaml b/examples/train_full/mixers/doremi_step1_static_qwen_pt_full.yaml index 2783999..b58f4b0 100644 --- a/examples/train_full/mixers/doremi_step1_static_qwen_pt_full.yaml +++ b/examples/train_full/mixers/doremi_step1_static_qwen_pt_full.yaml @@ -59,7 +59,7 @@ init_mixture_proportions: [0.5, 0.5] # corresponding initial proportions, here u static_mix: true warmup_step: 100 update_step: 200 -update_times: 3 +update_times: 3 # Ignored when static_mix=true # eval_dataset: mmlu_eval # TODO: Undefined dataset mmlu_eval in dataset_info.json. # eval_dataset: alpaca_zh_demo # SFT format, not suitable for PT diff --git a/examples/train_full/mixers/doremi_step2_dynamic_qwen_pt_full.yaml b/examples/train_full/mixers/doremi_step2_dynamic_qwen_pt_full.yaml index 7a7e836..fae829c 100644 --- a/examples/train_full/mixers/doremi_step2_dynamic_qwen_pt_full.yaml +++ b/examples/train_full/mixers/doremi_step2_dynamic_qwen_pt_full.yaml @@ -58,7 +58,7 @@ mixture_sample_rule: mixture # initial sampling rule, mixture is according to th init_mixture_proportions: [0.5, 0.5] # reference weight warmup_step: 100 update_step: 200 -update_times: 3 +update_times: 3 # Updates per Flex epoch; repeats by num_train_epochs # eval_dataset: mmlu_eval # TODO: Undefined dataset mmlu_eval in dataset_info.json. # eval_dataset: alpaca_zh_demo # SFT format, not suitable for PT diff --git a/examples/train_full/mixers/doremi_step3_static_qwen_pt_full.yaml b/examples/train_full/mixers/doremi_step3_static_qwen_pt_full.yaml index af3b517..67f7a89 100644 --- a/examples/train_full/mixers/doremi_step3_static_qwen_pt_full.yaml +++ b/examples/train_full/mixers/doremi_step3_static_qwen_pt_full.yaml @@ -59,7 +59,7 @@ init_mixture_proportions: [0.3, 0.7] # using the weights optimized from DoReMi S static_mix: true warmup_step: 100 update_step: 200 -update_times: 3 +update_times: 3 # Ignored when static_mix=true # eval_dataset: mmlu_eval # TODO: Undefined dataset mmlu_eval in dataset_info.json. # eval_dataset: alpaca_zh_demo # SFT format, not suitable for PT diff --git a/examples/train_full/mixers/odm_dynamic_qwen_pt_full.yaml b/examples/train_full/mixers/odm_dynamic_qwen_pt_full.yaml index ecea294..6f761a2 100644 --- a/examples/train_full/mixers/odm_dynamic_qwen_pt_full.yaml +++ b/examples/train_full/mixers/odm_dynamic_qwen_pt_full.yaml @@ -61,8 +61,8 @@ mixture_sample_rule: mixture # 初始采样规则，mixture为根据init_mixture init_mixture_proportions: [0.5, 0.5] # initial weights warmup_step: 10 update_step: 10 -update_times: -1 # -1 means to continuously update until the training ends (calculated based on num_train_epochs) -# train_step: 1000 # optional: explicitly specify the total training steps (priority over num_train_epochs) +update_times: -1 # No fixed updates; train by num_train_epochs/train_step +# train_step: 1000 # Total steps; overrides num_train_epochs # eval_dataset: mmlu_eval # TODO: Undefined dataset mmlu_eval in dataset_info.json. # eval_dataset: alpaca_zh_demo # SFT format, not suitable for PT diff --git a/examples/train_lora/mixers/doremi_step2_dynamic_qwen_pt_lora.yaml b/examples/train_lora/mixers/doremi_step2_dynamic_qwen_pt_lora.yaml index b8e72d1..58633af 100644 --- a/examples/train_lora/mixers/doremi_step2_dynamic_qwen_pt_lora.yaml +++ b/examples/train_lora/mixers/doremi_step2_dynamic_qwen_pt_lora.yaml @@ -61,7 +61,7 @@ mixture_sample_rule: mixture # 初始采样规则，mixture为根据init_mixture init_mixture_proportions: [0.5, 0.5] # 对应初始的比例，可通过额外算法自行调整 warmup_step: 100 update_step: 200 -update_times: 3 +update_times: 3 # Updates per Flex epoch; repeats by num_train_epochs # eval_dataset: mmlu_eval # TODO: Undefined dataset mmlu_eval in dataset_info.json. # eval_dataset: alpaca_zh_demo # SFT format, not suitable for PT diff --git a/examples/train_lora/mixers/random.yaml b/examples/train_lora/mixers/random.yaml index 30d068d..7450475 100644 --- a/examples/train_lora/mixers/random.yaml +++ b/examples/train_lora/mixers/random.yaml @@ -56,7 +56,7 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: random warmup_step: 100 update_step: 200 -update_times: 3 +update_times: 3 # Updates per Flex epoch; repeats by num_train_epochs eval_dataset: mmlu_eval diff --git a/examples/train_lora/selectors/custom.yaml b/examples/train_lora/selectors/custom.yaml index 01d4e07..8e4fd73 100644 --- a/examples/train_lora/selectors/custom.yaml +++ b/examples/train_lora/selectors/custom.yaml @@ -57,7 +57,7 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: custom # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 10 update_step: 10 -update_times: 2 +update_times: 2 # Updates per Flex epoch; repeats by num_train_epochs ## eval # val_size: 0.001 @@ -71,4 +71,4 @@ eval_steps: 10 # early_stopping_steps: 3 # early_stopping_min_delta: 0.01 -# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/custom.yaml \ No newline at end of file +# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/custom.yaml diff --git a/examples/train_lora/selectors/delta_loss.yaml b/examples/train_lora/selectors/delta_loss.yaml index 3f7e8a4..9ccff40 100644 --- a/examples/train_lora/selectors/delta_loss.yaml +++ b/examples/train_lora/selectors/delta_loss.yaml @@ -57,7 +57,7 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: delta_loss # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 10 update_step: 10 -update_times: 2 +update_times: 2 # Updates per Flex epoch; repeats by num_train_epochs ## eval # val_size: 0.001 @@ -71,4 +71,4 @@ eval_steps: 10 # early_stopping_steps: 3 # early_stopping_min_delta: 0.01 -# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/delta_loss.yaml \ No newline at end of file +# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/delta_loss.yaml diff --git a/examples/train_lora/selectors/less.yaml b/examples/train_lora/selectors/less.yaml index 4cbf06b..71802dd 100644 --- a/examples/train_lora/selectors/less.yaml +++ b/examples/train_lora/selectors/less.yaml @@ -57,7 +57,7 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: less # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 10 update_step: 10 -update_times: 2 +update_times: 2 # Updates per Flex epoch; repeats by num_train_epochs ## eval # val_size: 0.001 @@ -71,4 +71,4 @@ eval_steps: 10 # early_stopping_steps: 3 # early_stopping_min_delta: 0.01 -# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/less.yaml \ No newline at end of file +# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/less.yaml diff --git a/examples/train_lora/selectors/loss.yaml b/examples/train_lora/selectors/loss.yaml index f2ab0b2..b055c25 100644 --- a/examples/train_lora/selectors/loss.yaml +++ b/examples/train_lora/selectors/loss.yaml @@ -57,7 +57,7 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: loss # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 10 update_step: 10 -update_times: 2 +update_times: 2 # Updates per Flex epoch; repeats by num_train_epochs ## eval # val_size: 0.001 @@ -71,4 +71,4 @@ eval_steps: 10 # early_stopping_steps: 3 # early_stopping_min_delta: 0.01 -# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/loss.yaml \ No newline at end of file +# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/loss.yaml diff --git a/examples/train_lora/selectors/near.yaml b/examples/train_lora/selectors/near.yaml index f520eb4..0be9989 100644 --- a/examples/train_lora/selectors/near.yaml +++ b/examples/train_lora/selectors/near.yaml @@ -57,7 +57,7 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: near # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 10 update_step: 10 -update_times: 2 +update_times: 2 # Updates per Flex epoch; repeats by num_train_epochs ## eval # val_size: 0.001 @@ -71,4 +71,4 @@ eval_steps: 10 # early_stopping_steps: 3 # early_stopping_min_delta: 0.01 -# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/near.yaml \ No newline at end of file +# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/near.yaml diff --git a/examples/train_lora/selectors/nice.yaml b/examples/train_lora/selectors/nice.yaml index ed9a5f5..3c90460 100644 --- a/examples/train_lora/selectors/nice.yaml +++ b/examples/train_lora/selectors/nice.yaml @@ -57,7 +57,7 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: nice # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 10 update_step: 10 -update_times: 2 +update_times: 2 # Updates per Flex epoch; repeats by num_train_epochs ## eval # val_size: 0.001 @@ -71,4 +71,4 @@ eval_steps: 10 # early_stopping_steps: 3 # early_stopping_min_delta: 0.01 -# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/nice.yaml \ No newline at end of file +# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/nice.yaml diff --git a/examples/train_lora/selectors/random.yaml b/examples/train_lora/selectors/random.yaml index b26c12b..64c2846 100644 --- a/examples/train_lora/selectors/random.yaml +++ b/examples/train_lora/selectors/random.yaml @@ -57,7 +57,7 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: random # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 10 update_step: 10 -update_times: 2 +update_times: 2 # Updates per Flex epoch; repeats by num_train_epochs ## eval # val_size: 0.001 @@ -71,4 +71,4 @@ eval_steps: 10 # early_stopping_steps: 3 # early_stopping_min_delta: 0.01 -# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/random.yaml \ No newline at end of file +# FORCE_TORCHRUN=1 DISABLE_VERSION_CHECK=1 dataflex-cli train examples/train_lora/selectors/random.yaml diff --git a/examples/train_lora/selectors/tsds.yaml b/examples/train_lora/selectors/tsds.yaml index 5c756c1..e763706 100644 --- a/examples/train_lora/selectors/tsds.yaml +++ b/examples/train_lora/selectors/tsds.yaml @@ -57,7 +57,7 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: tsds # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 10 update_step: 10 -update_times: 2 +update_times: 2 # Updates per Flex epoch; repeats by num_train_epochs ## eval # val_size: 0.001 diff --git a/examples/train_lora/weighters/custom.yaml b/examples/train_lora/weighters/custom.yaml index 80c9d37..7a51ce7 100644 --- a/examples/train_lora/weighters/custom.yaml +++ b/examples/train_lora/weighters/custom.yaml @@ -57,4 +57,4 @@ train_type: dynamic_weight # 选择训练器类型。可选值包括： components_cfg_file: src/dataflex/configs/components.yaml component_name: custom # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 100 -train_step: 500 # 总训练步数（包括warm_up） \ No newline at end of file +train_step: 500 # Total steps; overrides num_train_epochs diff --git a/examples/train_lora/weighters/loss.yaml b/examples/train_lora/weighters/loss.yaml index ead1f7b..11953f2 100644 --- a/examples/train_lora/weighters/loss.yaml +++ b/examples/train_lora/weighters/loss.yaml @@ -57,4 +57,4 @@ train_type: dynamic_weight # 选择训练器类型。可选值包括： components_cfg_file: src/dataflex/configs/components.yaml component_name: loss # 选择组件名称，对应 components_cfg_file 中定义的组件 warmup_step: 100 -train_step: 500 # 总训练步数（包括warm_up） \ No newline at end of file +train_step: 500 # Total steps; overrides num_train_epochs diff --git a/src/dataflex/train/hparams/dynamic_params.py b/src/dataflex/train/hparams/dynamic_params.py index a4b8722..7eac532 100644 --- a/src/dataflex/train/hparams/dynamic_params.py +++ b/src/dataflex/train/hparams/dynamic_params.py @@ -547,7 +547,7 @@ class DynamicFinetuningArguments( ) update_times: int = field( default=1, - metadata={"help": "Total update times during the whole dynamic training progress for dynamic select or mix training"}, + metadata={"help": "Update times per Flex epoch for dynamic select or mix training. Use <= 0 for no fixed update count."}, ) static_mix: bool = field( default=False, @@ -555,7 +555,7 @@ class DynamicFinetuningArguments( ) train_step: int = field( default=0, - metadata={"help": "Only used in dynamic weight trainer and mix trainer (static_mix=True). Total training steps (including warmup)."}, + metadata={"help": "Optional total training steps. If set, overrides num_train_epochs."}, ) freeze_gate: bool = field( default=False, @@ -646,4 +646,4 @@ def __init__(self, early_stopping_patience: int): early_stopping_threshold=min_delta, ) - tuner.EarlyStoppingCallback = DataFlexEarlyStoppingCallback \ No newline at end of file + tuner.EarlyStoppingCallback = DataFlexEarlyStoppingCallback diff --git a/src/dataflex/train/trainer/mix_trainer.py b/src/dataflex/train/trainer/mix_trainer.py index 9783141..353547a 100644 --- a/src/dataflex/train/trainer/mix_trainer.py +++ b/src/dataflex/train/trainer/mix_trainer.py @@ -523,8 +523,8 @@ def _inner_training_loop( # Calculate steps from num_train_epochs and available data # First, get the original dataset size to estimate steps per epoch original_dataset_size = sum(len(dataset) for dataset in self.mixture_manager.sources.values()) - steps_per_epoch = max(1, original_dataset_size // total_train_batch_size) - total_steps = int(steps_per_epoch * args.num_train_epochs) + steps_per_epoch = max(1, int(np.ceil(original_dataset_size / total_train_batch_size))) + total_steps = int(np.ceil(steps_per_epoch * args.num_train_epochs)) logger.info(f"[Dataflex] Calculated train_step={total_steps} from num_train_epochs={args.num_train_epochs} " f"(dataset_size={original_dataset_size}, steps_per_epoch={steps_per_epoch})") @@ -547,31 +547,39 @@ def _inner_training_loop( len_dataloader, # 等于数据集长度/worldsize/micro_batchsize max_steps, ) = self.set_initial_training_values(args, train_dataloader, total_train_batch_size) + # Issue #49: support num_train_epochs while keeping dynamic-step training + epoch_update_steps = None if self.finetuning_args.static_mix: if self.finetuning_args.train_step > 0: max_steps = self.finetuning_args.train_step else: # Use the calculated total_steps from above original_dataset_size = sum(len(dataset) for dataset in self.mixture_manager.sources.values()) - steps_per_epoch = max(1, original_dataset_size // total_train_batch_size) - max_steps = int(steps_per_epoch * args.num_train_epochs) + steps_per_epoch = max(1, int(np.ceil(original_dataset_size / total_train_batch_size))) + max_steps = int(np.ceil(steps_per_epoch * args.num_train_epochs)) + epoch_update_steps = steps_per_epoch else: # Dynamic mixing: calculate max_steps based on update_times - if self.finetuning_args.update_times < 0: - # update_times=-1 means update until training finishes - # Calculate total steps from num_train_epochs - if self.finetuning_args.train_step > 0: - max_steps = self.finetuning_args.train_step - logger.info(f"[Dataflex] Dynamic mix: using train_step={max_steps}") - else: - original_dataset_size = sum(len(dataset) for dataset in self.mixture_manager.sources.values()) - steps_per_epoch = max(1, original_dataset_size // total_train_batch_size) - max_steps = int(steps_per_epoch * args.num_train_epochs) - logger.info(f"[Dataflex] Dynamic mix: calculated max_steps={max_steps} from num_train_epochs={args.num_train_epochs}") + dynamic_epoch_update_steps = None + if self.finetuning_args.update_times > 0: + dynamic_epoch_update_steps = max( + 1, + self.finetuning_args.warmup_step + self.finetuning_args.update_step * self.finetuning_args.update_times, + ) + if self.finetuning_args.train_step > 0: + max_steps = self.finetuning_args.train_step + epoch_update_steps = dynamic_epoch_update_steps + logger.info(f"[Dataflex] Dynamic mix: using train_step={max_steps}") + elif dynamic_epoch_update_steps is not None: + max_steps = int(np.ceil(args.num_train_epochs * dynamic_epoch_update_steps)) + epoch_update_steps = dynamic_epoch_update_steps + logger.info(f"[Dataflex] Dynamic mix: calculated max_steps={max_steps} from num_train_epochs={args.num_train_epochs}") else: - # Fixed number of updates - max_steps = (self.finetuning_args.warmup_step + self.finetuning_args.update_step * self.finetuning_args.update_times) - logger.info(f"[Dataflex] Dynamic mix: max_steps={max_steps} (warmup={self.finetuning_args.warmup_step} + update_step={self.finetuning_args.update_step} * update_times={self.finetuning_args.update_times})") + original_dataset_size = sum(len(dataset) for dataset in self.mixture_manager.sources.values()) + steps_per_epoch = max(1, int(np.ceil(original_dataset_size / total_train_batch_size))) + max_steps = int(np.ceil(steps_per_epoch * args.num_train_epochs)) + epoch_update_steps = steps_per_epoch + logger.info(f"[Dataflex] Dynamic mix: calculated max_steps={max_steps} from num_train_epochs={args.num_train_epochs}") epoch_based = False logger.info(f"[Dataflex]Set max train steps to {max_steps}") logger.info(f"[Dataflex]Set epoch_based = False") @@ -758,7 +766,6 @@ def _inner_training_loop( if args.eval_on_start: self._evaluate(trial, ignore_keys_for_eval, skip_scheduler=True) - # 放弃epoch逻辑，相当于只训练一个epoch，通过step来训练 current_dataloader = train_dataloader epoch = 0 if self.finetuning_args.static_mix == False and self.state.global_step < self.finetuning_args.warmup_step: @@ -770,7 +777,11 @@ def _inner_training_loop( if args.past_index >= 0: self._past = None - steps_in_epoch = max_steps * args.gradient_accumulation_steps + total_training_batches = max_steps * args.gradient_accumulation_steps + if epoch_update_steps is not None: + steps_in_epoch = epoch_update_steps * args.gradient_accumulation_steps + else: + steps_in_epoch = len_dataloader if len_dataloader is not None else total_training_batches self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control) if resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0: @@ -792,7 +803,7 @@ def _inner_training_loop( remainder = args.gradient_accumulation_steps update_step = -1 # 一个epoch中的模型总更新次数 - total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 + total_updates = total_training_batches // args.gradient_accumulation_steps + 1 if args.gradient_accumulation_steps == 1: total_updates -= 1 for _ in range(total_updates): @@ -801,12 +812,22 @@ def _inner_training_loop( num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder batch_samples, num_items_in_batch = self.get_batch_samples(current_iterator, num_batches, args.device) + if len(batch_samples) == 0: + epoch += 1 + if hasattr(current_dataloader, "set_epoch"): + current_dataloader.set_epoch(epoch) + current_iterator = iter(current_dataloader) + batch_samples, num_items_in_batch = self.get_batch_samples(current_iterator, num_batches, args.device) + if len(batch_samples) == 0: + self.control.should_training_stop = True + break # 遍历当前批次的样本 for i, inputs in enumerate(batch_samples): step += 1 # 每次迭代时增加全局步数 # 判断是否达到同步步数，或者是当前epoch的最后一个步数 - do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch + is_epoch_end = (step + 1 + steps_skipped) % steps_in_epoch == 0 + do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or is_epoch_end # 由于我们使用了预取（prefetching），我们需要手动设置同步梯度 self.accelerator.gradient_state._set_sync_gradients(do_sync_step) @@ -954,20 +975,33 @@ def _inner_training_loop( # learning_rate=learning_rate, ) - # 动态训练更新 + step_in_epoch = ( + self.state.global_step % epoch_update_steps + if epoch_update_steps is not None + else self.state.global_step + ) if ( + self.finetuning_args.static_mix == False + and step_in_epoch == 0 + and epoch_update_steps is not None + and self.state.global_step < max_steps + ): + self.train_dataset = self.mixture_manager.rebuild(num_samples = total_train_batch_size * self.finetuning_args.warmup_step) + current_dataloader = self.get_train_dataloader() + current_iterator = iter(current_dataloader) + elif ( self.finetuning_args.static_mix == False and( self.state.global_step < max_steps and ( - self.state.global_step == self.finetuning_args.warmup_step or - (self.state.global_step > self.finetuning_args.warmup_step and - (self.state.global_step - self.finetuning_args.warmup_step) % self.finetuning_args.update_step == 0))) + step_in_epoch == self.finetuning_args.warmup_step or + (step_in_epoch > self.finetuning_args.warmup_step and + (step_in_epoch - self.finetuning_args.warmup_step) % self.finetuning_args.update_step == 0))) ): self.accelerator.wait_for_everyone() torch.cuda.empty_cache() if dist.is_initialized(): dist.barrier() - update_times = (self.state.global_step - self.finetuning_args.warmup_step) // self.finetuning_args.update_step + 1 + update_times = (step_in_epoch - self.finetuning_args.warmup_step) // self.finetuning_args.update_step + 1 logger.info(f"[Dataflex] Model training paused, starting the {update_times}th dynamic data mixture...") # Pass current training batch to mixer for λ_t computation @@ -987,8 +1021,8 @@ def _inner_training_loop( self.mixture_manager.set_proportions(probs) self.train_dataset = self.mixture_manager.rebuild(num_samples = total_train_batch_size * self.finetuning_args.update_step) - train_loader = self.get_train_dataloader() - current_iterator = iter(train_loader) + current_dataloader = self.get_train_dataloader() + current_iterator = iter(current_dataloader) if self.accelerator.is_main_process: logger.info(f"[Dataflex] Updated dataloader at step {self.state.global_step}, new mixture proportions generated.") @@ -1099,4 +1133,4 @@ def _inner_training_loop( if self.neftune_noise_alpha is not None: self._deactivate_neftune(self.model) - return TrainOutput(self.state.global_step, train_loss, metrics) \ No newline at end of file + return TrainOutput(self.state.global_step, train_loss, metrics) diff --git a/src/dataflex/train/trainer/select_trainer.py b/src/dataflex/train/trainer/select_trainer.py index 7090046..b260407 100644 --- a/src/dataflex/train/trainer/select_trainer.py +++ b/src/dataflex/train/trainer/select_trainer.py @@ -382,7 +382,21 @@ def _inner_training_loop( len_dataloader, # 等于数据集长度/worldsize/micro_batchsize max_steps, ) = self.set_initial_training_values(args, train_dataloader, total_train_batch_size) - max_steps = (self.finetuning_args.warmup_step + self.finetuning_args.update_step * self.finetuning_args.update_times) + # Issue #49: support num_train_epochs while keeping dynamic-step training + epoch_update_steps = None + if self.finetuning_args.update_times > 0: + epoch_update_steps = max( + 1, + self.finetuning_args.warmup_step + self.finetuning_args.update_step * self.finetuning_args.update_times, + ) + if self.finetuning_args.train_step > 0: + max_steps = self.finetuning_args.train_step + elif epoch_update_steps is not None: + max_steps = int(np.ceil(args.num_train_epochs * epoch_update_steps)) + else: + steps_per_epoch = max(1, int(np.ceil(len(self.train_dataset) / total_train_batch_size))) + max_steps = int(np.ceil(args.num_train_epochs * steps_per_epoch)) + epoch_update_steps = steps_per_epoch epoch_based = False logger.info(f"[Dataflex]Set max train steps to {max_steps}") logger.info(f"[Dataflex]Set epoch_based = False") @@ -569,7 +583,6 @@ def _inner_training_loop( if args.eval_on_start: self._evaluate(trial, ignore_keys_for_eval, skip_scheduler=True) - # 放弃epoch逻辑，相当于只训练一个epoch，通过step来训练 current_dataloader = train_dataloader epoch = 0 if self.state.global_step < self.finetuning_args.warmup_step: @@ -581,7 +594,11 @@ def _inner_training_loop( if args.past_index >= 0: self._past = None - steps_in_epoch = max_steps * args.gradient_accumulation_steps + total_training_batches = max_steps * args.gradient_accumulation_steps + if epoch_update_steps is not None: + steps_in_epoch = epoch_update_steps * args.gradient_accumulation_steps + else: + steps_in_epoch = len_dataloader if len_dataloader is not None else total_training_batches self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control) if resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0: @@ -603,7 +620,7 @@ def _inner_training_loop( remainder = args.gradient_accumulation_steps update_step = -1 # 一个epoch中的模型总更新次数 - total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 + total_updates = total_training_batches // args.gradient_accumulation_steps + 1 if args.gradient_accumulation_steps == 1: total_updates -= 1 for _ in range(total_updates): @@ -612,12 +629,22 @@ def _inner_training_loop( num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder batch_samples, num_items_in_batch = self.get_batch_samples(current_iterator, num_batches, args.device) + if len(batch_samples) == 0: + epoch += 1 + if hasattr(current_dataloader, "set_epoch"): + current_dataloader.set_epoch(epoch) + current_iterator = iter(current_dataloader) + batch_samples, num_items_in_batch = self.get_batch_samples(current_iterator, num_batches, args.device) + if len(batch_samples) == 0: + self.control.should_training_stop = True + break # 遍历当前批次的样本 for i, inputs in enumerate(batch_samples): step += 1 # 每次迭代时增加全局步数 # 判断是否达到同步步数，或者是当前epoch的最后一个步数 - do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch + is_epoch_end = (step + 1 + steps_skipped) % steps_in_epoch == 0 + do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or is_epoch_end # 由于我们使用了预取（prefetching），我们需要手动设置同步梯度 self.accelerator.gradient_state._set_sync_gradients(do_sync_step) @@ -750,26 +777,40 @@ def _inner_training_loop( # learning_rate=learning_rate, ) - # 动态训练更新 - if ( + step_in_epoch = ( + self.state.global_step % epoch_update_steps + if epoch_update_steps is not None + else self.state.global_step + ) + if step_in_epoch == 0 and epoch_update_steps is not None and self.state.global_step < max_steps: + warmup_indices = self.selector.warmup(total_warmup_samples, replacement=True) + current_dataloader = self.get_train_dataloader(warmup_indices) + current_iterator = iter(current_dataloader) + elif ( self.state.global_step < max_steps and ( - self.state.global_step == self.finetuning_args.warmup_step or - (self.state.global_step > self.finetuning_args.warmup_step and - (self.state.global_step - self.finetuning_args.warmup_step) % self.finetuning_args.update_step == 0)) + step_in_epoch == self.finetuning_args.warmup_step or + (step_in_epoch > self.finetuning_args.warmup_step and + (step_in_epoch - self.finetuning_args.warmup_step) % self.finetuning_args.update_step == 0)) ): self.accelerator.wait_for_everyone() torch.cuda.empty_cache() if dist.is_initialized(): dist.barrier() - current_update_times = (self.state.global_step - self.finetuning_args.warmup_step) // self.finetuning_args.update_step + 1 + current_update_times = (step_in_epoch - self.finetuning_args.warmup_step) // self.finetuning_args.update_step + 1 + effective_update_times = self.finetuning_args.update_times + if effective_update_times <= 0 and epoch_update_steps is not None: + effective_update_times = max( + 1, + int(np.ceil(max(epoch_update_steps - self.finetuning_args.warmup_step, 0) / self.finetuning_args.update_step)), + ) logger.info(f"[Dataflex] Model training paused, starting the {current_update_times}th dynamic data selection...") # 这里传一些特定的selector参数 extra_args = dict( optimizer_state=self.optimizer.state, scheduler_state=self.lr_scheduler.state_dict(), current_update_times=current_update_times, - update_times=self.finetuning_args.update_times, + update_times=effective_update_times, tokenizer=self.tokenizer, ) new_indices = self.selector.select( @@ -780,8 +821,8 @@ def _inner_training_loop( ) # 每个进程根据 local_indices 构造 dataloader - train_loader = self.get_train_dataloader(indices=new_indices) - current_iterator = iter(train_loader) + current_dataloader = self.get_train_dataloader(indices=new_indices) + current_iterator = iter(current_dataloader) if self.accelerator.is_main_process: logger.info(f"[Dataflex] Updated dataloader at step {self.state.global_step}, {len(new_indices)} samples in total.") diff --git a/src/dataflex/train/trainer/weight_trainer.py b/src/dataflex/train/trainer/weight_trainer.py index e5a66b3..61476f1 100644 --- a/src/dataflex/train/trainer/weight_trainer.py +++ b/src/dataflex/train/trainer/weight_trainer.py @@ -306,7 +306,9 @@ def _inner_training_loop( len_dataloader, # 等于数据集长度/worldsize/micro_batchsize max_steps, ) = self.set_initial_training_values(args, train_dataloader, total_train_batch_size) - max_steps = self.finetuning_args.train_step + # Issue #49: support num_train_epochs while keeping dynamic-step training + if self.finetuning_args.train_step > 0: + max_steps = self.finetuning_args.train_step epoch_based = False logger.info(f"[Dataflex]Set max train steps to {max_steps}") logger.info(f"[Dataflex]Set epoch_based = False") @@ -493,7 +495,7 @@ def _inner_training_loop( if args.eval_on_start: self._evaluate(trial, ignore_keys_for_eval, skip_scheduler=True) - # 放弃epoch逻辑，相当于只训练一个epoch，通过step来训练 + current_dataloader = train_dataloader epoch = 0 if self.state.global_step < self.finetuning_args.warmup_step: logger.info("[Dataflex] Model warmup in progress...") @@ -503,7 +505,8 @@ def _inner_training_loop( if args.past_index >= 0: self._past = None - steps_in_epoch = max_steps * args.gradient_accumulation_steps + total_training_batches = max_steps * args.gradient_accumulation_steps + steps_in_epoch = len_dataloader if len_dataloader is not None else total_training_batches self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control) if resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0: @@ -512,20 +515,20 @@ def _inner_training_loop( rng_to_sync = False steps_skipped = 0 if steps_trained_in_current_epoch > 0: - train_dataloader = skip_first_batches(train_dataloader, steps_trained_in_current_epoch) + current_dataloader = skip_first_batches(current_dataloader, steps_trained_in_current_epoch) steps_skipped = steps_trained_in_current_epoch steps_trained_in_current_epoch = 0 rng_to_sync = True step = -1 - current_iterator = iter(train_dataloader) + current_iterator = iter(current_dataloader) # We chunkify the epoch iterator into gradient accumulation steps `n` batches remainder = num_examples % args.gradient_accumulation_steps if remainder == 0: remainder = args.gradient_accumulation_steps update_step = -1 # 一个epoch中的模型总更新次数 - total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 + total_updates = total_training_batches // args.gradient_accumulation_steps + 1 if args.gradient_accumulation_steps == 1: total_updates -= 1 for _ in range(total_updates): @@ -534,12 +537,22 @@ def _inner_training_loop( num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder batch_samples, num_items_in_batch = self.get_batch_samples(current_iterator, num_batches, args.device) + if len(batch_samples) == 0: + epoch += 1 + if hasattr(current_dataloader, "set_epoch"): + current_dataloader.set_epoch(epoch) + current_iterator = iter(current_dataloader) + batch_samples, num_items_in_batch = self.get_batch_samples(current_iterator, num_batches, args.device) + if len(batch_samples) == 0: + self.control.should_training_stop = True + break # 遍历当前批次的样本 for i, inputs in enumerate(batch_samples): step += 1 # 每次迭代时增加全局步数 # 判断是否达到同步步数，或者是当前epoch的最后一个步数 - do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch + is_epoch_end = (step + 1 + steps_skipped) % steps_in_epoch == 0 + do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or is_epoch_end # 由于我们使用了预取（prefetching），我们需要手动设置同步梯度 self.accelerator.gradient_state._set_sync_gradients(do_sync_step) From d4830abd7dbd90b8d761e77025213d8972cd03a1 Mon Sep 17 00:00:00 2001 From: tongruiliu <952814125@qq.com> Date: Tue, 12 May 2026 18:46:19 +0800 Subject: [PATCH 2/3] support num_train_epoches while keep dynamic-step training --- examples/test/test_mix.yaml | 2 +- examples/test/test_static_mix.yaml | 2 +- .../doremi_step1_static_qwen_pt_full.yaml | 2 +- .../doremi_step2_dynamic_qwen_pt_full.yaml | 2 +- .../doremi_step3_static_qwen_pt_full.yaml | 2 +- .../mixers/odm_dynamic_qwen_pt_full.yaml | 4 +- .../doremi_step2_dynamic_qwen_pt_lora.yaml | 2 +- examples/train_lora/mixers/random.yaml | 2 +- skills/how_to_use.md | 40 +++++++- src/dataflex/train/hparams/dynamic_params.py | 2 +- src/dataflex/train/trainer/mix_trainer.py | 94 ++++++------------- src/dataflex/train/trainer/select_trainer.py | 14 ++- src/dataflex/train/trainer/weight_trainer.py | 2 + 13 files changed, 91 insertions(+), 79 deletions(-) diff --git a/examples/test/test_mix.yaml b/examples/test/test_mix.yaml index f4e849e..acbbee3 100644 --- a/examples/test/test_mix.yaml +++ b/examples/test/test_mix.yaml @@ -53,4 +53,4 @@ mixture_sample_rule: mixture # 初始采样规则，mixture为根据init_mix init_mixture_proportions: [0.7, 0.3] # 对应初始的比例 warmup_step: 4 update_step: 3 -update_times: 2 # Updates per Flex epoch; repeats by num_train_epochs +update_times: 2 \ No newline at end of file diff --git a/examples/test/test_static_mix.yaml b/examples/test/test_static_mix.yaml index 6821cc3..4a59301 100644 --- a/examples/test/test_static_mix.yaml +++ b/examples/test/test_static_mix.yaml @@ -50,4 +50,4 @@ train_type: dynamic_mix # 选择训练器类型。可选值包括： static_mix: true # 是否固定初始静态混合比例（仅在dynamic_mix训练器中生效） mixture_sample_rule: mixture # 初始采样规则，mixture为根据init_mixture_proportions比例混合（可动态调整），stratified为固定按源数据集大小比例分层，uniform为固定均匀分布 init_mixture_proportions: [0.7, 0.3] # 对应初始的比例，可通过额外算法自行调整 -train_step: 3 # Total steps; overrides num_train_epochs +train_step: 3 # 总训练步数（仅在dynamic_mix训练器中生效），不考虑warmup和update steps \ No newline at end of file diff --git a/examples/train_full/mixers/doremi_step1_static_qwen_pt_full.yaml b/examples/train_full/mixers/doremi_step1_static_qwen_pt_full.yaml index b58f4b0..2783999 100644 --- a/examples/train_full/mixers/doremi_step1_static_qwen_pt_full.yaml +++ b/examples/train_full/mixers/doremi_step1_static_qwen_pt_full.yaml @@ -59,7 +59,7 @@ init_mixture_proportions: [0.5, 0.5] # corresponding initial proportions, here u static_mix: true warmup_step: 100 update_step: 200 -update_times: 3 # Ignored when static_mix=true +update_times: 3 # eval_dataset: mmlu_eval # TODO: Undefined dataset mmlu_eval in dataset_info.json. # eval_dataset: alpaca_zh_demo # SFT format, not suitable for PT diff --git a/examples/train_full/mixers/doremi_step2_dynamic_qwen_pt_full.yaml b/examples/train_full/mixers/doremi_step2_dynamic_qwen_pt_full.yaml index fae829c..7a7e836 100644 --- a/examples/train_full/mixers/doremi_step2_dynamic_qwen_pt_full.yaml +++ b/examples/train_full/mixers/doremi_step2_dynamic_qwen_pt_full.yaml @@ -58,7 +58,7 @@ mixture_sample_rule: mixture # initial sampling rule, mixture is according to th init_mixture_proportions: [0.5, 0.5] # reference weight warmup_step: 100 update_step: 200 -update_times: 3 # Updates per Flex epoch; repeats by num_train_epochs +update_times: 3 # eval_dataset: mmlu_eval # TODO: Undefined dataset mmlu_eval in dataset_info.json. # eval_dataset: alpaca_zh_demo # SFT format, not suitable for PT diff --git a/examples/train_full/mixers/doremi_step3_static_qwen_pt_full.yaml b/examples/train_full/mixers/doremi_step3_static_qwen_pt_full.yaml index 67f7a89..af3b517 100644 --- a/examples/train_full/mixers/doremi_step3_static_qwen_pt_full.yaml +++ b/examples/train_full/mixers/doremi_step3_static_qwen_pt_full.yaml @@ -59,7 +59,7 @@ init_mixture_proportions: [0.3, 0.7] # using the weights optimized from DoReMi S static_mix: true warmup_step: 100 update_step: 200 -update_times: 3 # Ignored when static_mix=true +update_times: 3 # eval_dataset: mmlu_eval # TODO: Undefined dataset mmlu_eval in dataset_info.json. # eval_dataset: alpaca_zh_demo # SFT format, not suitable for PT diff --git a/examples/train_full/mixers/odm_dynamic_qwen_pt_full.yaml b/examples/train_full/mixers/odm_dynamic_qwen_pt_full.yaml index 6f761a2..ecea294 100644 --- a/examples/train_full/mixers/odm_dynamic_qwen_pt_full.yaml +++ b/examples/train_full/mixers/odm_dynamic_qwen_pt_full.yaml @@ -61,8 +61,8 @@ mixture_sample_rule: mixture # 初始采样规则，mixture为根据init_mixture init_mixture_proportions: [0.5, 0.5] # initial weights warmup_step: 10 update_step: 10 -update_times: -1 # No fixed updates; train by num_train_epochs/train_step -# train_step: 1000 # Total steps; overrides num_train_epochs +update_times: -1 # -1 means to continuously update until the training ends (calculated based on num_train_epochs) +# train_step: 1000 # optional: explicitly specify the total training steps (priority over num_train_epochs) # eval_dataset: mmlu_eval # TODO: Undefined dataset mmlu_eval in dataset_info.json. # eval_dataset: alpaca_zh_demo # SFT format, not suitable for PT diff --git a/examples/train_lora/mixers/doremi_step2_dynamic_qwen_pt_lora.yaml b/examples/train_lora/mixers/doremi_step2_dynamic_qwen_pt_lora.yaml index 58633af..b8e72d1 100644 --- a/examples/train_lora/mixers/doremi_step2_dynamic_qwen_pt_lora.yaml +++ b/examples/train_lora/mixers/doremi_step2_dynamic_qwen_pt_lora.yaml @@ -61,7 +61,7 @@ mixture_sample_rule: mixture # 初始采样规则，mixture为根据init_mixture init_mixture_proportions: [0.5, 0.5] # 对应初始的比例，可通过额外算法自行调整 warmup_step: 100 update_step: 200 -update_times: 3 # Updates per Flex epoch; repeats by num_train_epochs +update_times: 3 # eval_dataset: mmlu_eval # TODO: Undefined dataset mmlu_eval in dataset_info.json. # eval_dataset: alpaca_zh_demo # SFT format, not suitable for PT diff --git a/examples/train_lora/mixers/random.yaml b/examples/train_lora/mixers/random.yaml index 7450475..30d068d 100644 --- a/examples/train_lora/mixers/random.yaml +++ b/examples/train_lora/mixers/random.yaml @@ -56,7 +56,7 @@ components_cfg_file: src/dataflex/configs/components.yaml component_name: random warmup_step: 100 update_step: 200 -update_times: 3 # Updates per Flex epoch; repeats by num_train_epochs +update_times: 3 eval_dataset: mmlu_eval diff --git a/skills/how_to_use.md b/skills/how_to_use.md index 92275c7..ada403b 100644 --- a/skills/how_to_use.md +++ b/skills/how_to_use.md @@ -59,9 +59,38 @@ A DataFlex training config is a standard LlamaFactory YAML with additional DataF | `component_name` | str | Which algorithm to use, matching a key in `components_cfg_file` | | `warmup_step` | int | Number of warmup steps before dynamic behavior kicks in | | `update_step` | int | Interval (in steps) between dynamic updates | -| `update_times` | int | Total number of dynamic updates. Use `-1` for continuous updates until training ends | +| `update_times` | int | For `dynamic_select`, number of dynamic updates per Flex epoch. Use `-1` for dataset-sized epochs with continuous updates | | `static_mix` | bool | If `true` with `dynamic_mix`, use fixed proportions (no dynamic updates). Used in DoReMi Step 1 & 3 | -| `train_step` | int | Total training steps. Required for `dynamic_weight` and `dynamic_mix` with `static_mix: true` | +| `train_step` | int | Optional total training steps. If set to a positive value, it overrides `num_train_epochs`-derived steps | + +### Step and Epoch Semantics + +DataFlex dynamic trainers run a step-based training loop internally. Prefer `eval_strategy: "steps"` / `save_strategy: "steps"` or disable them with `"no"` when using dynamic training. Epoch-based evaluation or saving depends on the internal step-to-epoch bookkeeping and may not align with Flex epoch boundaries. + +For `dynamic_select`, `num_train_epochs` repeats Flex epochs. One Flex epoch contains: + +```text +warmup_step + update_step * update_times +``` + +For example: + +```yaml +warmup_step: 10 +update_step: 10 +update_times: 2 +num_train_epochs: 3 +``` + +This runs `3 * (10 + 10 * 2) = 90` optimization steps. To keep the old single-Flex-epoch behavior, set: + +```yaml +num_train_epochs: 1.0 +``` + +If `train_step > 0`, DataFlex uses `train_step` as the exact total number of optimization steps and does not derive total steps from `num_train_epochs`. + +For `dynamic_weight`, `warmup_step` is a global step threshold. Reweighting starts when `global_step >= warmup_step` and does not reset at epoch boundaries. ### Data Mixture Fields (for `dynamic_mix` only) @@ -83,12 +112,14 @@ component_name: less # choices: less, nice, loss, delta_loss, tsds, near, warmup_step: 10 update_step: 10 update_times: 2 +num_train_epochs: 1.0 ``` **How it works:** 1. Warmup phase: train on randomly sampled data for `warmup_step` steps. 2. At `warmup_step` and every `update_step` steps: pause training, run the selector to pick new samples, rebuild the dataloader. -3. Total steps = `warmup_step + update_step * update_times`. +3. One Flex epoch has `warmup_step + update_step * update_times` steps. +4. Total steps are derived from `num_train_epochs` unless `train_step > 0`. **Example:** ```bash @@ -144,7 +175,8 @@ train_step: 500 **How it works:** 1. Standard training for `warmup_step` steps (no reweighting). 2. After warmup: each training step computes per-sample losses and applies the weighting strategy. -3. Total steps = `train_step`. +3. `warmup_step` is measured in global optimization steps and does not reset per epoch. +4. If `train_step > 0`, total steps = `train_step`; otherwise total steps follow the standard `num_train_epochs` calculation. **Example:** ```bash diff --git a/src/dataflex/train/hparams/dynamic_params.py b/src/dataflex/train/hparams/dynamic_params.py index 7eac532..e6bd388 100644 --- a/src/dataflex/train/hparams/dynamic_params.py +++ b/src/dataflex/train/hparams/dynamic_params.py @@ -547,7 +547,7 @@ class DynamicFinetuningArguments( ) update_times: int = field( default=1, - metadata={"help": "Update times per Flex epoch for dynamic select or mix training. Use <= 0 for no fixed update count."}, + metadata={"help": "Update times per Flex epoch for dynamic selection. Use <= 0 for no fixed update count."}, ) static_mix: bool = field( default=False, diff --git a/src/dataflex/train/trainer/mix_trainer.py b/src/dataflex/train/trainer/mix_trainer.py index 353547a..9783141 100644 --- a/src/dataflex/train/trainer/mix_trainer.py +++ b/src/dataflex/train/trainer/mix_trainer.py @@ -523,8 +523,8 @@ def _inner_training_loop( # Calculate steps from num_train_epochs and available data # First, get the original dataset size to estimate steps per epoch original_dataset_size = sum(len(dataset) for dataset in self.mixture_manager.sources.values()) - steps_per_epoch = max(1, int(np.ceil(original_dataset_size / total_train_batch_size))) - total_steps = int(np.ceil(steps_per_epoch * args.num_train_epochs)) + steps_per_epoch = max(1, original_dataset_size // total_train_batch_size) + total_steps = int(steps_per_epoch * args.num_train_epochs) logger.info(f"[Dataflex] Calculated train_step={total_steps} from num_train_epochs={args.num_train_epochs} " f"(dataset_size={original_dataset_size}, steps_per_epoch={steps_per_epoch})") @@ -547,39 +547,31 @@ def _inner_training_loop( len_dataloader, # 等于数据集长度/worldsize/micro_batchsize max_steps, ) = self.set_initial_training_values(args, train_dataloader, total_train_batch_size) - # Issue #49: support num_train_epochs while keeping dynamic-step training - epoch_update_steps = None if self.finetuning_args.static_mix: if self.finetuning_args.train_step > 0: max_steps = self.finetuning_args.train_step else: # Use the calculated total_steps from above original_dataset_size = sum(len(dataset) for dataset in self.mixture_manager.sources.values()) - steps_per_epoch = max(1, int(np.ceil(original_dataset_size / total_train_batch_size))) - max_steps = int(np.ceil(steps_per_epoch * args.num_train_epochs)) - epoch_update_steps = steps_per_epoch + steps_per_epoch = max(1, original_dataset_size // total_train_batch_size) + max_steps = int(steps_per_epoch * args.num_train_epochs) else: # Dynamic mixing: calculate max_steps based on update_times - dynamic_epoch_update_steps = None - if self.finetuning_args.update_times > 0: - dynamic_epoch_update_steps = max( - 1, - self.finetuning_args.warmup_step + self.finetuning_args.update_step * self.finetuning_args.update_times, - ) - if self.finetuning_args.train_step > 0: - max_steps = self.finetuning_args.train_step - epoch_update_steps = dynamic_epoch_update_steps - logger.info(f"[Dataflex] Dynamic mix: using train_step={max_steps}") - elif dynamic_epoch_update_steps is not None: - max_steps = int(np.ceil(args.num_train_epochs * dynamic_epoch_update_steps)) - epoch_update_steps = dynamic_epoch_update_steps - logger.info(f"[Dataflex] Dynamic mix: calculated max_steps={max_steps} from num_train_epochs={args.num_train_epochs}") + if self.finetuning_args.update_times < 0: + # update_times=-1 means update until training finishes + # Calculate total steps from num_train_epochs + if self.finetuning_args.train_step > 0: + max_steps = self.finetuning_args.train_step + logger.info(f"[Dataflex] Dynamic mix: using train_step={max_steps}") + else: + original_dataset_size = sum(len(dataset) for dataset in self.mixture_manager.sources.values()) + steps_per_epoch = max(1, original_dataset_size // total_train_batch_size) + max_steps = int(steps_per_epoch * args.num_train_epochs) + logger.info(f"[Dataflex] Dynamic mix: calculated max_steps={max_steps} from num_train_epochs={args.num_train_epochs}") else: - original_dataset_size = sum(len(dataset) for dataset in self.mixture_manager.sources.values()) - steps_per_epoch = max(1, int(np.ceil(original_dataset_size / total_train_batch_size))) - max_steps = int(np.ceil(steps_per_epoch * args.num_train_epochs)) - epoch_update_steps = steps_per_epoch - logger.info(f"[Dataflex] Dynamic mix: calculated max_steps={max_steps} from num_train_epochs={args.num_train_epochs}") + # Fixed number of updates + max_steps = (self.finetuning_args.warmup_step + self.finetuning_args.update_step * self.finetuning_args.update_times) + logger.info(f"[Dataflex] Dynamic mix: max_steps={max_steps} (warmup={self.finetuning_args.warmup_step} + update_step={self.finetuning_args.update_step} * update_times={self.finetuning_args.update_times})") epoch_based = False logger.info(f"[Dataflex]Set max train steps to {max_steps}") logger.info(f"[Dataflex]Set epoch_based = False") @@ -766,6 +758,7 @@ def _inner_training_loop( if args.eval_on_start: self._evaluate(trial, ignore_keys_for_eval, skip_scheduler=True) + # 放弃epoch逻辑，相当于只训练一个epoch，通过step来训练 current_dataloader = train_dataloader epoch = 0 if self.finetuning_args.static_mix == False and self.state.global_step < self.finetuning_args.warmup_step: @@ -777,11 +770,7 @@ def _inner_training_loop( if args.past_index >= 0: self._past = None - total_training_batches = max_steps * args.gradient_accumulation_steps - if epoch_update_steps is not None: - steps_in_epoch = epoch_update_steps * args.gradient_accumulation_steps - else: - steps_in_epoch = len_dataloader if len_dataloader is not None else total_training_batches + steps_in_epoch = max_steps * args.gradient_accumulation_steps self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control) if resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0: @@ -803,7 +792,7 @@ def _inner_training_loop( remainder = args.gradient_accumulation_steps update_step = -1 # 一个epoch中的模型总更新次数 - total_updates = total_training_batches // args.gradient_accumulation_steps + 1 + total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 if args.gradient_accumulation_steps == 1: total_updates -= 1 for _ in range(total_updates): @@ -812,22 +801,12 @@ def _inner_training_loop( num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder batch_samples, num_items_in_batch = self.get_batch_samples(current_iterator, num_batches, args.device) - if len(batch_samples) == 0: - epoch += 1 - if hasattr(current_dataloader, "set_epoch"): - current_dataloader.set_epoch(epoch) - current_iterator = iter(current_dataloader) - batch_samples, num_items_in_batch = self.get_batch_samples(current_iterator, num_batches, args.device) - if len(batch_samples) == 0: - self.control.should_training_stop = True - break # 遍历当前批次的样本 for i, inputs in enumerate(batch_samples): step += 1 # 每次迭代时增加全局步数 # 判断是否达到同步步数，或者是当前epoch的最后一个步数 - is_epoch_end = (step + 1 + steps_skipped) % steps_in_epoch == 0 - do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or is_epoch_end + do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch # 由于我们使用了预取（prefetching），我们需要手动设置同步梯度 self.accelerator.gradient_state._set_sync_gradients(do_sync_step) @@ -975,33 +954,20 @@ def _inner_training_loop( # learning_rate=learning_rate, ) - step_in_epoch = ( - self.state.global_step % epoch_update_steps - if epoch_update_steps is not None - else self.state.global_step - ) + # 动态训练更新 if ( - self.finetuning_args.static_mix == False - and step_in_epoch == 0 - and epoch_update_steps is not None - and self.state.global_step < max_steps - ): - self.train_dataset = self.mixture_manager.rebuild(num_samples = total_train_batch_size * self.finetuning_args.warmup_step) - current_dataloader = self.get_train_dataloader() - current_iterator = iter(current_dataloader) - elif ( self.finetuning_args.static_mix == False and( self.state.global_step < max_steps and ( - step_in_epoch == self.finetuning_args.warmup_step or - (step_in_epoch > self.finetuning_args.warmup_step and - (step_in_epoch - self.finetuning_args.warmup_step) % self.finetuning_args.update_step == 0))) + self.state.global_step == self.finetuning_args.warmup_step or + (self.state.global_step > self.finetuning_args.warmup_step and + (self.state.global_step - self.finetuning_args.warmup_step) % self.finetuning_args.update_step == 0))) ): self.accelerator.wait_for_everyone() torch.cuda.empty_cache() if dist.is_initialized(): dist.barrier() - update_times = (step_in_epoch - self.finetuning_args.warmup_step) // self.finetuning_args.update_step + 1 + update_times = (self.state.global_step - self.finetuning_args.warmup_step) // self.finetuning_args.update_step + 1 logger.info(f"[Dataflex] Model training paused, starting the {update_times}th dynamic data mixture...") # Pass current training batch to mixer for λ_t computation @@ -1021,8 +987,8 @@ def _inner_training_loop( self.mixture_manager.set_proportions(probs) self.train_dataset = self.mixture_manager.rebuild(num_samples = total_train_batch_size * self.finetuning_args.update_step) - current_dataloader = self.get_train_dataloader() - current_iterator = iter(current_dataloader) + train_loader = self.get_train_dataloader() + current_iterator = iter(train_loader) if self.accelerator.is_main_process: logger.info(f"[Dataflex] Updated dataloader at step {self.state.global_step}, new mixture proportions generated.") @@ -1133,4 +1099,4 @@ def _inner_training_loop( if self.neftune_noise_alpha is not None: self._deactivate_neftune(self.model) - return TrainOutput(self.state.global_step, train_loss, metrics) + return TrainOutput(self.state.global_step, train_loss, metrics) \ No newline at end of file diff --git a/src/dataflex/train/trainer/select_trainer.py b/src/dataflex/train/trainer/select_trainer.py index b260407..be6b5ab 100644 --- a/src/dataflex/train/trainer/select_trainer.py +++ b/src/dataflex/train/trainer/select_trainer.py @@ -397,6 +397,8 @@ def _inner_training_loop( steps_per_epoch = max(1, int(np.ceil(len(self.train_dataset) / total_train_batch_size))) max_steps = int(np.ceil(args.num_train_epochs * steps_per_epoch)) epoch_update_steps = steps_per_epoch + if self.finetuning_args.train_step > 0 and args.num_train_epochs != 1: + logger.warning("[Dataflex] train_step is set; num_train_epochs will be ignored.") epoch_based = False logger.info(f"[Dataflex]Set max train steps to {max_steps}") logger.info(f"[Dataflex]Set epoch_based = False") @@ -782,7 +784,17 @@ def _inner_training_loop( if epoch_update_steps is not None else self.state.global_step ) - if step_in_epoch == 0 and epoch_update_steps is not None and self.state.global_step < max_steps: + if ( + step_in_epoch == 0 + and epoch_update_steps is not None + and self.state.global_step > 0 + and self.state.global_step < max_steps + ): + self.accelerator.wait_for_everyone() + torch.cuda.empty_cache() + if dist.is_initialized(): + dist.barrier() + warmup_indices = self.selector.warmup(total_warmup_samples, replacement=True) current_dataloader = self.get_train_dataloader(warmup_indices) current_iterator = iter(current_dataloader) diff --git a/src/dataflex/train/trainer/weight_trainer.py b/src/dataflex/train/trainer/weight_trainer.py index 61476f1..114b07b 100644 --- a/src/dataflex/train/trainer/weight_trainer.py +++ b/src/dataflex/train/trainer/weight_trainer.py @@ -309,6 +309,8 @@ def _inner_training_loop( # Issue #49: support num_train_epochs while keeping dynamic-step training if self.finetuning_args.train_step > 0: max_steps = self.finetuning_args.train_step + if args.num_train_epochs != 1: + logger.warning("[Dataflex] train_step is set; num_train_epochs will be ignored.") epoch_based = False logger.info(f"[Dataflex]Set max train steps to {max_steps}") logger.info(f"[Dataflex]Set epoch_based = False") From 0153ecc957ec0a68de510c80053ca784dc921ccd Mon Sep 17 00:00:00 2001 From: tongruiliu <952814125@qq.com> Date: Wed, 13 May 2026 15:44:59 +0800 Subject: [PATCH 3/3] add how to use.md for weighter --- skills/how_to_use.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/skills/how_to_use.md b/skills/how_to_use.md index ada403b..2101ae5 100644 --- a/skills/how_to_use.md +++ b/skills/how_to_use.md @@ -89,6 +89,7 @@ num_train_epochs: 1.0 ``` If `train_step > 0`, DataFlex uses `train_step` as the exact total number of optimization steps and does not derive total steps from `num_train_epochs`. +For multi-epoch tests, make sure example configs do not leave a positive `train_step`; pass `train_step=0` on the CLI if you want `num_train_epochs` to control training length. For `dynamic_weight`, `warmup_step` is a global step threshold. Reweighting starts when `global_step >= warmup_step` and does not reset at epoch boundaries. @@ -169,7 +170,7 @@ train_type: dynamic_weight components_cfg_file: src/dataflex/configs/components.yaml component_name: loss # choices: loss, custom warmup_step: 100 -train_step: 500 +train_step: 500 # fixed-step example; set to 0 for num_train_epochs-based multi-epoch runs ``` **How it works:**