Skip to content

🪂 Don't gather logits in SFT to avoid hanging #2890

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 3 commits into from
Feb 18, 2025
Merged

🪂 Don't gather logits in SFT to avoid hanging #2890

merged 3 commits into from
Feb 18, 2025

Conversation

qgallouedec
Copy link
Member

@qgallouedec qgallouedec commented Feb 18, 2025 •
edited
Loading

What does this PR do?

Fixes #2879

I'm not sure why it does always hang, but it seems that at some point in the training, always the same, it can hang while trying to gather the logits. The fix consist in gathering only the number of correct tokens and the total number of tokens.

The only way was to copy the content of compute_token_accuracy in the SFTTrainer, which I think should be ok as it is only used once in the codebase.

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a GitHub issue? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines.
  • Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@qgallouedec qgallouedec mentioned this pull request Feb 18, 2025
Closed
5 tasks
@qgallouedec qgallouedec requested a review from kashif February 18, 2025 11:21
@HuggingFaceDocBuilderDev
Copy link

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.

@qgallouedec qgallouedec marked this pull request as ready for review February 18, 2025 13:53
@qgallouedec qgallouedec merged commit 6c54f02 into main Feb 18, 2025
14 checks passed
@qgallouedec qgallouedec deleted the fix-hanging-sft branch February 18, 2025 14:31
qgallouedec added a commit that referenced this pull request Feb 18, 2025
@qgallouedec
🪂 Don't gather logits in SFT to avoid hanging (#2890)
8aa13d8 
* Don't gather logits

* Remove unused function and test
@coding-famer
Copy link

I guess the reason is the sizes of gathered logits(seq length) do not match. So gathering only the number of correct tokens and the total number of tokens will be a good fix.

@lewtun lewtun mentioned this pull request Feb 23, 2025
Merged
5 tasks
@dszhengyu
Copy link

dszhengyu commented Mar 28, 2025 •
edited
Loading

with trl==0.16.0, and SFTTrainer for "Qwen/Qwen2.5-1.5B-Instruct" and base_job_name = "mt5-large-full-peft", I still encounter similar issue when my training set is larger (~300M samples) on P4d.24xlarge instance (8 A100 GPU), while the error now suggesting mismatching shape:

    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/transformers/trainer.py", line 2245, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/transformers/trainer.py", line 2556, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/transformers/trainer.py", line 2556, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/transformers/trainer.py", line 2556, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/transformers/trainer.py", line 3718, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/transformers/trainer.py", line 3718, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/transformers/trainer.py", line 3718, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/trl/trainer/sft_trainer.py", line 531, in compute_loss
    total_tokens = self.accelerator.gather_for_metrics(total_tokens)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/trl/trainer/sft_trainer.py", line 531, in compute_loss
    total_tokens = self.accelerator.gather_for_metrics(total_tokens)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/trl/trainer/sft_trainer.py", line 531, in compute_loss
    total_tokens = self.accelerator.gather_for_metrics(total_tokens)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/site-packages/accelerate/accelerator.py", line 2613, in gather_for_metrics
    data = self.gather(input_data)
           ^^^^^^^^^^^^^^^^^^^^^^^
accelerate.utils.operations.DistributedOperationException: Cannot apply desired operation due to shape mismatches. All shapes across devices must be valid.
Operation: `accelerate.utils.operations.gather`
Input shapes:
  - Process 0: []
  - Process 1: []
  - Process 2: []
  - Process 3: [8]
  - Process 4: []
  - Process 5: []
  - Process 6: []
  - Process 7: [8]
  File "/opt/ml/code/sft_qwen.py", line 412, in <module>
    main(training_args)

Is it possible that, GPU might get different batches (for example, 10 batches for 8 GPUs), so the loss/mean_token_accuracy will run into issue?

@kashif
Copy link
Collaborator

kashif commented Mar 28, 2025

do you still have this issue @dszhengyu with 0.16.0 version of TRL?

@dszhengyu
Copy link

do you still have this issue @dszhengyu with 0.16.0 version of TRL?

Yes, it is 0.16.0, sorry I made a typo of 0.15.0, fixed.

@kashif
Copy link
Collaborator

kashif commented Mar 28, 2025

ok so yes its an edge case where there are things are all masked and thus the sum is none and should be zero...

@kashif
Copy link
Collaborator

kashif commented Mar 28, 2025

@dszhengyu l'll fix it and make a patch release

yxliu-TAMU pushed a commit to mincheolseong/ECEN743-GRPO-Project-Proposal that referenced this pull request Apr 20, 2025
@qgallouedec
🪂 Don't gather logits in SFT to avoid hanging (huggingface#2890)
492e33b 
* Don't gather logits

* Remove unused function and test
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@kashif kashif kashif approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

in SFT script, distributed training got stuck if set packing=false
5 participants
@qgallouedec @HuggingFaceDocBuilderDev @coding-famer @dszhengyu @kashif