[kashif]

Related issue

Closes #1081

What does this PR do?

Makes the tiled MLP backward accumulate each weight gradient once per pass instead of once per shard, so DDP no longer fails with Expected to mark a variable ready only once.

Motivation

The backward drives each shard with its own torch.autograd.backward, so gate_proj.weight, up_proj.weight and down_proj.weight each reach their AccumulateGrad node num_shards times in a single pass. DDP registers one hook per parameter and marks it ready on every firing, which is one-shot per pass, so any model with more than one tiled MLP dies in backward.

ds_grad_is_ready (#1000) defers DeepSpeed's reduction to the last shard, but the accumulations themselves still happen once per shard, so it does not help DDP.

Tiling is a memory optimization and should not change how gradients reach the parameters.

Changes

• Pass the weights to LigerTiledMLPFunction.apply as explicit tensor inputs and return their gradients from backward, so autograd accumulates each one once, at the end of the pass.
• Use torch.autograd.grad instead of torch.autograd.backward on that path, summing the per-shard gradients locally. No AccumulateGrad node fires inside the loop, so DDP's per-parameter hook is untouched.
• Keep the existing per-shard path for ZeRO-3, selected on hasattr(param, "ds_id"). Partitioned weights are 0-sized placeholders outside their owning module's forward, so autograd would record (0,) as the expected gradient shape and the recompute would fail. Those still defer the reduction via ds_grad_is_ready.

Testing

Ran the reporter's repro from #1081, reduced to 2 gloo/CPU ranks and a 2-layer LlamaConfig so it needs no GPU. Before this change both ranks raise Expected to mark a variable ready only once; after it, backward completes.

Also checked against an untiled reference in float64, single process, max abs gradient error:

case	weight grads	input grads
baseline [2, 256, 64], 4 shards	3.4e-14	0.0
2d input (moe experts)	1.2e-14	0.0
seqlen not divisible by shards	1.2e-14	1.7e-16
auto num_shards	2.8e-14	0.0
single shard	0.0	0.0
out_features != in_features	2.3e-14	0.0
x.requires_grad=False	2.8e-14	n/a
ZeRO-3 (ds_id) path	2.5e-14	0.0
ZeRO-3 + non-divisible seqlen	1.4e-14	2.2e-16

And confirmed the accumulation count directly with a register_post_accumulate_grad_hook on each weight: one firing per weight per backward with 8 shards on the default path, num_shards firings on the ZeRO-3 path as intended.

liger-kernels has no test directory here, so the regression tests land upstream in linkedin/Liger-Kernel#1276 alongside the same fix.

Checklist

• This PR is linked to an issue that was discussed and approved
• I have tested these changes locally
• New/changed functionality has test coverage
• LLM disclosure:
  • I did not use an LLM to create this PR.
  • I used and LLM for assistance while creating this PR.
  • This PR was mostly or completely generated by an LLM.

[github-actions]

Hi @kashif, thanks for your interest in contributing!

This project requires that pull request authors are vouched, and you are not in the list of vouched users.

This PR will be closed automatically. See https://github.com/huggingface/kernels-community/blob/main/CONTRIBUTING.md for more details.

[danieldk]

/kernel-bot build liger-kernels

[github-actions]

Build request processed.

Command: /kernel-bot build liger-kernels
Mode: build only
Target branch: pr-1083
PR head SHA: 2f9b488750fa66947d0804061a82df8a96bdbe0a
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

• liger-kernels (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/31883217357

[sayakpaul]

/kernel-bot build-and-merge liger-kernels

[github-actions]

Invalid command. Use /kernel-bot <build|security|security-and-build|build-and-stage|merge-and-upload|release> <kernel1> [kernel2 ...] [--branch <target_branch>].
A kernel may be scoped to a subset of backends, e.g. flash-attn2[xpu,cpu].
The security command does not require kernel names.

[sayakpaul]

/kernel-bot merge-and-upload liger-kernels

[github-actions]

Build request processed.

Command: /kernel-bot merge-and-upload liger-kernels
Mode: merge, build and upload
Target branch: ``
PR head SHA: 2f9b488750fa66947d0804061a82df8a96bdbe0a
Workflows: `build.yaml, build-mac.yaml, build-windows.yaml`

Merge result: PR is already merged. Continuing with build/upload.

Dispatched (1):

• liger-kernels (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/32336985746

Hub uploads:

• Hub repo: https://huggingface.co/kernels/kernels-community/liger-kernels