RFC: refactor causal lm loss to handle lm_head in loss function#39047
RFC: refactor causal lm loss to handle lm_head in loss function#39047winglian wants to merge 1 commit into
Conversation
|
Maybe some input from liger-kernel and cce maintainers would be helpful too @shimizust @lancerts @vaibhavjindal @erikwijmans |
|
This is definitely a core change to the library layout, so cc @ArthurZucker @Cyrilvallez! |
|
I definitely like the spirit behind this. This may get a bit complicated as sometimes how the logits are compute is specific to the model implementation (softcapping, like in Gemma 2, is what comes to mind). I think given the memory saving that fusing the cross-entropy and linear layer enables, this is worth figuring out tho :) I do think the example given would work because you can still specify the logits from the caller (the model) if needed and any uncommon extra arguments could be put in **kwargs. The default implementation should probably support softcapping and a bias, since both aren't common it may make sense to have those be in kwargs. |
|
cc @ArthurZucker - could indeed allow nice training optimizations more easily, but we lose a bit of explicitness, and add a redundant codepaths if the loss is not overriden. WDYT? |
|
One other thought: Accessing the lm head weights via PyTorch's FSDP2 relies on the forward call to know when to unshard the weights. The default way transformers shards is for the lm_head to be part of the model-level sharding group, so this isn't a problem as it gets unsharded by I haven't seen this issue with ZeRO as it relies on a different mechanism. I don't have much experience with PyTorch's original FSDP, but I believe you to explicitly use summon_full_params to access the weights in this way. |
What does this PR do?
The ecosystem is has a lot of patches to optimize the loss for causal models. Due to frequent api changes, these patches are hard to maintain and keep up to date.
The RFC proposes changes to the ForCausalLMLoss to accept the lm_head weights and last hidden_states as part of the loss function. By moving the loss calculation here, it's much easier to add various loss variations that leverage fusing the linear + cross entropy calculations.
Fixes # (issue)
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.