Enable Gradient Accumulation fix across all models + trainer fully in forward() #34283
Conversation
|
@ArthurZucker |
| @@ -3610,8 +3612,11 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N | |||
| labels = inputs.pop("labels") | |||
| else: | |||
| labels = None | |||
| # if num_items_in_batch is not None: | |||
| # inputs["num_items_in_batch"] = num_items_in_batch | |||
| if self.model_accepts_loss_kwargs: | |||
There was a problem hiding this comment.
This if condition doesn't seem to work for PeftModel class (it only has kwargs not loss_kwargs 🫠 )
I tried just changing that condition to if True and ran some tests, and the loss calculation worked perfectly for a LORA on a Llama 3 1B.
I'm wondering if there's a safe/non-breaking way to support peft models here as well?
There was a problem hiding this comment.
Thanks for the flag, that may be why I couldn't recreate #34263 (comment)
There was a problem hiding this comment.
@man-shar can you try giving peft another go? Should have fixed it
There was a problem hiding this comment.
Can confirm self.model_accepts_loss_kwargs is the correct value now for a peft model!
But seems like the compute_loss function inside trainer.py isn't getting the num_items_in_batch argument passed to it from training_step.
I notice that argument was removed in 4f3f86d and the commit still says "Experimental". So I assume it will be reverted once you're done experimenting! It should work after that!
|
Awesome work by all of you on this. Insane dev speed over the past few days 🙏 🔥 |
|
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. |
|
Rebasing should probably fix quality! |
muellerzr
force-pushed
the
fixup-loss_fn_issues
branch
from
c2a8453
to
fc6d674
Compare
| tr_loss_step = self.training_step(model, inputs, num_items_in_batch) | ||
| tr_loss_step = self.training_step(model, inputs) |
There was a problem hiding this comment.
Why num_items_in_batch is removed here?
| @@ -1114,6 +1114,7 @@ def forward( | |||
| return_dict: Optional[bool] = None, | |||
| cache_position: Optional[torch.LongTensor] = None, | |||
| num_logits_to_keep: int = 0, | |||
| **loss_kwargs, | |||
There was a problem hiding this comment.
I think making it an argument loss_args rather than kwargs would be better?
What does this PR do?
Since most users still want OOTB, this trickles the loss kwargs to the rest of the models so that causal loss can be calculated properly
Also relies on huggingface/accelerate#3185 for DDP bits
Fixes # (issue)
Fully fixes #34263 / finishes #34191 & #34198 & fixes #34242
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
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.
@ArthurZucker