Repository navigation
TiledMLP #893
Description
🚀 The feature, motivation and pitch
TiledMLP can save memory for large bs/long seq as well
Reference code: https://github.com/deepspeedai/DeepSpeed/blob/16c1bf429f10db3e681b49fac37d641909aad69d/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L781
Thanks to @upskyy, the base implementation has been done in #935
However, the current is not compatible with ddp and fsdp, see #935 (comment).
For ddp support, one can try registering hooks to temporily accumulate gradients and only return the accumulated gradients in the last iteration, based on https://github.com/axolotl-ai-cloud/axolotl/blob/dd78f2e0cc5cc6458daaad02cc29b649ff1046f5/src/axolotl/monkeypatch/tiled_mlp/base.py#L99
For fsdp, parameters are always resharded after the first backward pass, one need to find a way to unshard parameters in the latter iterations.
-
base implementation based on deepspeed: Add TiledMLP Implementation #935 @upskyy
-
ddp support (contribute welcome)
-
fsdp support (contribute welcome)
Alternatives
No response
Additional context
No response
Activity
0xtoward commentedon Nov 3, 2025
Hi, I'm interested in working on it!
Tcc0403 commentedon Nov 13, 2025
since this one is taken, @0xtoward would you like to work on #937 instead?update: This issue is not taken
0xtoward commentedon Nov 15, 2025
Ofc I'd work on it if no one else is.
Add TiledMLP Implementation (#935)
sfc-gh-sbekman commentedon Feb 3, 2026
@Tcc0403, so who is working on adding DDP/FSDP support for
LigerTiledMLPFunction?Is @upskyy working on it? It's been 2+ months but I haven't seen any related PRs.
Thanks.
Tcc0403 commentedon Feb 3, 2026
Sadly, it seems no one is working on it right now. Since TiledMLP is not yet intregrated into any model patches, DDP/FSDP support has lower priority in our roadmap.
For now, we are focusing on transformers v5 support in liger's next release. We will circle back to TiledMLP patch once the transformers modeling structure is settled, and DDP/FSDP support can follow on it.
sfc-gh-sbekman commentedon Feb 3, 2026
Thank you for the explanation, @Tcc0403. One of the key messages of my comment was to update this comment of yours #893 (comment)
perhaps edit it to update the reality so that contributors know this feature request doesn't have an active owner?
transformers v5 perturbations should have zero impact on this feature, so if a contributor is willing to work on it there should be no interference with your work on v5.
Tcc0403 commentedon Feb 3, 2026
Thanks for the suggestion. I've updated the comment and added a tag.
AlirezAkbary commentedon Feb 16, 2026
Hi, I would like to give it a shot! @Tcc0403
alektebel commentedon Mar 1, 2026
Hi @Tcc0403 , @AlirezAkbary ,
Happy to collaborate in elaborating a draft PR for these issue, if that's okey with you
Tcc0403 commentedon Mar 1, 2026
Yeah, feel free to open a PR.
alektebel commentedon Mar 3, 2026
Hi @Tcc0403 ,
I believe there might be some overlap with a recent PR: #1125 where DDP integration with TiledMLP is already underway. Would it make sense to narrow the scope of this draft PR to integrate TiledMLP with fsdp, in order to avoid duplication?
Let me know what you think.
Tcc0403 commentedon Mar 3, 2026
Yes, please work on fsdp support if it's possible.
Add FSDP support to TiledMLP by preventing premature resharding in ba…
alektebel commentedon Mar 4, 2026
Hi @Tcc0403 ,
Opened a draft PR with the fix #1128 .
As we talked earlier, I've adjusted the scope to just the FSDP, since DDP handling is already taking place in other PR. Basically I've added a context handler that activates when the wrapper is FSDP, in order to unshard the parameters when this happens. This seems to fix the error caused by the corruption of state, but needs to be verified in a multi gpu environment as well (I only verified it in one GPU)
Let me know what you think
fix(tiled_mlp): batch shard backward for FSDP compatibility
fix(tiled_mlp): FSDP-compatible backward via per-shard autograd.grad
fix (tiledmlp): FSDP-Compatible tiled backward via GradientAccumulato…
fix(tiled_mlp): use GradientAccumulator for FSDP-compatible backward
fix(tiled_mlp): FSDP-safe tiled backward via GradientAccumulator