Skip to content

fix(trainer): normalize partial gradient accumulation windows - #10084

Open
MrCapricornLiu wants to merge 1 commit into
modelscope:mainfrom
MrCapricornLiu:lch/fix-partial-gradient-accumulation
Open

fix(trainer): normalize partial gradient accumulation windows#10084
MrCapricornLiu wants to merge 1 commit into
modelscope:mainfrom
MrCapricornLiu:lch/fix-partial-gradient-accumulation

Conversation

@MrCapricornLiu

Copy link
Copy Markdown
Contributor

PR type

  • Bug Fix
  • New Feature
  • Document Updates
  • More Models or Datasets Support

PR information

When the last accumulation window of an epoch has fewer micro-batches than
gradient_accumulation_steps, Trainer.compute_loss and
RerankerTrainer.compute_loss still divide their mean loss by the configured
window size. For models accepting loss kwargs, HF's training_step does not
divide that loss again. With a configured window of four and an actual window
of two, this produces half the full-batch reference gradient.

Use HF Trainer's current_gradient_accumulation_steps for that division, with
the configured value as a fallback for versions/callers without this attribute.
Full windows and the existing normalization branch conditions are unchanged.

The regression runs the actual HF training_step/Accelerate backward path
with Swift's Trainer and RerankerTrainer loss implementations, then compares
parameter gradients and accumulated loss with an independent concatenated-batch
reference. It uses a tiny CPU model, no model downloads. The test initializes
HF's training machinery directly and supplies Swift's required loss context;
it does not run the complete Swift initialization or epoch loop.

Experiment results

CUDA_VISIBLE_DEVICES='' PYTHONPATH=. python -m pytest -q \
  tests/train/test_trainer_partial_accumulation.py
  • Original code: 4 failed, 5 passed. Failures are partial windows of one/two
    micro-batches in Trainer and RerankerTrainer.
  • Fixed code: 9 passed with Transformers 5.3.0 / Accelerate 1.14.0;
    9 passed with Transformers 4.57.1 / Accelerate 1.11.0.
  • Controls cover full windows and Trainer's model-without-loss-kwargs path.
  • Changed-file pre-commit, flake8/isort/YAPF and git diff --check pass.

These are CPU backward tests with equal-size micro-batches, not distributed,
AMP, full fine-tuning or variable-batch weighting tests. The new pytest module
must be invoked explicitly; the repository's existing unittest CI does not
automatically collect it.

Related: #9510 refactors the training loop but retains the configured-size
loss division in these implementations; this fix is independent of that work.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant