[kashif]

The tiled MLP kernel passes compute_params through but never uses it, and computes num_shards locally on each rank. Both bite under sharded-parameter backends:

• DeepSpeed ZeRO-3 crashes: the per-shard recompute fires ZeRO's gradient hook multiple times for the same parameter, so the same ds_id lands in the reduction bucket twice and it asserts.
• ZeRO-3 / FSDP can deadlock: if ranks see different sequence lengths they pick different num_shards, so they issue a different number of weight all-gathers and hang.

This ports the fixes that DeepSpeed's own TiledMLP uses:

• In the backward, defer ZeRO-3 reduction by setting param.ds_grad_is_ready = False for every shard except the last (gated on the ds_id marker, so non-ZeRO backends are untouched).
• In apply_tiled_mlp, all-reduce num_shards to the per-rank max when distributed, so every rank runs the same number of shards.

Same fix is going into linkedin/Liger-Kernel. Verified there on 2x H100: tiled gradients match a non-tiled reference under FSDP2 and DeepSpeed ZeRO-3 (with and without gradient checkpointing). No change for single-GPU or non-distributed runs.

[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.

[vasqu]

/kernel-bot build liger-kernels

[github-actions]

Build request processed.

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

Dispatched (1):

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

[vasqu]

Thanks 🤗

[vasqu]

/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: e156b7a70cf24a5e29ce92d738d5f79503bb9096
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/28387310314

[danieldk]

And congrats on making PR #1000 @kashif :)

[kashif]

thanks @danieldk happy to help!