Count each token once in num_tokens under tensor parallelism

[qgallouedec]

Fixes #7100.

Under tensor parallelism every TP rank is handed the same batch, but num_tokens was gathered with gather_for_metrics and summed over all processes, so each token was counted tp_size times. With tp_size=32 the logged tokens/s was 32x the truth.

The gathered count is now divided by tp_size in the trainers that maintain _total_train_tokens themselves: SFT, DPO, KTO, Reward and (experimental) TPO.

CP and SP need no correction: context_parallel splits the buffers in place and the Ulysses dataloader adapter shards the batch, both before compute_loss runs, so each rank already holds a distinct slice. That is why the divisor is tp_size and not non_data_parallel_size, which is what Trainer uses for num_items_in_batch (counted before sharding).

GRPO, RLOO and Distillation report num_tokens from state.num_input_tokens_seen, which Trainer maintains and does not divide either. That one is upstream.

Verification

2xH100, tp_size=2, one 5-token sample per step:

before: num_tokens = 10
after:  num_tokens = 5

---

Note

Low Risk
Metrics-only change for logged token counts; no loss or optimization logic is modified. Wrong tp_size would only misreport throughput when parallelism_config is absent or misconfigured (defaults to 1).

Overview
Fixes inflated num_tokens (and derived tokens/s) when training with tensor parallelism: TP ranks each see the same batch, but gathered attention-mask sums were added across every process, so each token was counted once per TP rank.

SFT, DPO (standard and Liger paths), KTO, Reward, and experimental TPO now read tp_size from Accelerate’s parallelism_config (Accelerate ≥ 1.12.0; otherwise _tp_size = 1) and divide the gathered per-step token count by that factor before updating _total_train_tokens. Context and sequence parallelism are unchanged—those paths shard the batch before compute_loss, so no TP-style deduplication is applied.

^{Reviewed by Cursor Bugbot for commit f56614e. Bugbot is set up for automated code reviews on this repo. Configure here.}

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-07T19:05:42.343709Z	6f76553	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit f56614e. Configure here.}

[cursor]

TP version gate is too high

Low Severity

The new tp_size correction is gated on Accelerate 1.12.0, but ParallelismConfig.tp_size has existed since 1.10.0. With tensor parallelism on 1.10.x–1.11.x, _tp_size stays 1 and gathered num_tokens is still multiplied by the TP group size. The comment that parallelism_config requires 1.12.0 is also incorrect; the nearby context-parallel path already reads it from 1.10.1.

Additional Locations (2)

• trl/trainer/dpo_trainer.py#L938-L942
• trl/trainer/kto_trainer.py#L950-L954

 

^{Reviewed by Cursor Bugbot for commit f56614e. Configure here.}