linkedin / linkedin/Liger-Kernel

Gradient checkpointing for `grad_weight` in LFCE

Open
#533 4 comments 3 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.