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_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_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/skills/how_to_use.md b/skills/how_to_use.md index 92275c7..2101ae5 100644 --- a/skills/how_to_use.md +++ b/skills/how_to_use.md @@ -59,9 +59,39 @@ 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 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. ### Data Mixture Fields (for `dynamic_mix` only) @@ -83,12 +113,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 @@ -138,13 +170,14 @@ 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:** 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 a4b8722..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": "Total update times during the whole dynamic training progress for dynamic select or mix training"}, + metadata={"help": "Update times per Flex epoch for dynamic selection. 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/select_trainer.py b/src/dataflex/train/trainer/select_trainer.py index 7090046..be6b5ab 100644 --- a/src/dataflex/train/trainer/select_trainer.py +++ b/src/dataflex/train/trainer/select_trainer.py @@ -382,7 +382,23 @@ 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 + 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") @@ -569,7 +585,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 +596,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 +622,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 +631,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 +779,50 @@ 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 ( + 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) + 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 +833,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..114b07b 100644 --- a/src/dataflex/train/trainer/weight_trainer.py +++ b/src/dataflex/train/trainer/weight_trainer.py @@ -306,7 +306,11 @@ 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 + 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") @@ -493,7 +497,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 +507,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 +517,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 +539,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)