Back Home

GitHub Repo

Transformers Community Proposes Training Normalization Fix After Effective Token Count Truncation May Amplify Gradients

A Nov. 1 report says Trainer may truncate effective token counts under model parallelism, mis-scaling loss and gradients. The author shared a two-GPU comparison and submitted a fix, but it had not been merged or confirmed by maintainers as of the review.

Marcus Qwertyus · Public domain · Image source
zh-Hant

On Nov. 1, the Transformers community reported an issue affecting numerical correctness in training: when Trainer calculates the loss-normalization denominator, it may amplify gradients by rounding down if the number of effective target tokens is not divisible by the model-parallel factor. The reported setup used development version 5.23.0.dev0, two Tesla T4 GPUs, and context parallelism with FSDP2 and cp_size=2. Issue report

The suspected problem is in _get_num_items_in_batch(): it uses // when dividing the number of effective tokens by the product of the tensor, context, and sequence parallelism factors. For example, dividing three effective targets by two turns the denominator from the 1.5 it should retain into 1; the author’s minimal check also showed that one effective target could be truncated to zero. Reproduction code and count results

This count feeds into loss calculation; it is not just for logging. The official Trainer documentation says the custom loss interface receives the number of items in the entire gradient-accumulation batch, and uses num_items_in_batch in normalization. So, from an engineering perspective, the actual denominator and update magnitude need to be checked; confirming only that training completes is not enough. Trainer documentation

The author’s proposed fix replaces integer division with regular division and adds regression tests for divisible, non-divisible, and zero-denominator cases. In the author’s two-GPU validation, with three effective targets, the original version’s loss, gradient norm, and single-step SGD update norm were all 1.5 times those of the reference version; results for the four-target comparison matched. These results are self-reported by the submitter, and the PR remains open with no review recorded. Proposed fix

Context parallelism splits inputs along the sequence dimension so multiple GPUs can share the computation for attention over long sequences. The Accelerate documentation also cautions that label shifting and splitting can affect loss correctness. This report makes the normalization denominator another component worth checking, but it does not establish that all parallelism modes or released versions are affected. Follow-up should focus on maintainer confirmation, testing on other parallelism modes, and which version ultimately includes the fix. Context parallelism documentation

Sources

  1. Trainer floor-divides valid-token counts under model parallelism
  2. Fix/trainer parallel token count 50212 — PR #50214
  3. Transformers Trainer 文件
  4. Context Parallel in Accelerate