linkedin / linkedin/Liger-Kernel
Gradient checkpointing for `grad_weight` in LFCE
- Dominant language
- Python
- Stars
- 6.6k
- Forks
- 603
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 47
Description
### 🚀 The feature, motivation and pitch
The LFCE kernel allocates a `grad_weight` tensor:
https://github.com/linkedin/Liger-Kernel/blob/a8fa3bb37850e89500261024ff47da0c626ab75f/src/liger_kernel/ops/fused_linear_cross_entropy.py#L47
This tensor then gets updated throughout the chunked loss calculation and finally used in the backward as a custom grad operation:
https://github.com/linkedin/Liger-Kernel/blob/a8fa3bb37850e89500261024ff47da0c626ab75f/src/liger_kernel/ops/fused_linear_cross_entropy.py#L127-L136
This has shape `[vocab_size, hidden_size]`, which in situations where you have big models with big vocabularies, this becomes very large, which makes it impossible to do large pipeline parallel microbatches at long sequence lengths, as I have to keep this tensor in memory until the backward. It would be great to have gradient check-pointing here or even full recomputation would work.
### Alternatives
_No response_
### Additional context
_No response_
Contributor guide
Research direction
Start in src/liger_kernel/ops/fused_linear_cross_entropy.py at the grad_weight allocation around line 47 and its custom backward use around lines 127-136. Trace how the tensor is updated during chunked loss calculation and determine whether checkpointing or recomputation can avoid retaining it until backward. Done means reducing the retained grad_weight memory while preserving the backward result.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100