Repository navigation
Integrating TiledMLP for a much smaller memory footprint #41826
Description
Similar to TiledFusedLogitsLoss #41306 it's of a great benefit to perform TiledMLP as well, which can save a lot of HBM memory and allow longer seqlen/larger batch size or even enable a bs=1 training which wasn't possible before.
This technology comes from the Arctic Long Sequence Training https://arxiv.org/abs/2506.13996 which allowed us to train in bf16 with 500K tokens on a single H100 GPU, 3.7M on a single node, and 15M on Llama-8B using just four nodes. You will want to refer to section 3.1.1 TiledMLP for the full details.
Tiled MLP processes hidden_states in small slices. If we extract a single LlamaMLP layer from Llama-8B and run a bf16 hidden_states tensor of shape [1, 256_000, 4096] through its forward-backward, without and with sequence dimension tiling, we get about 10x memory saved as can be seen from the memory profiler plot:
There is the additional cost of recomputing forward path. So we trade time for memory.
Currently, we hack it in with monkey patching:
module_path = f"transformers.models.{model_type}.modeling_{model_type}"
model_cls_prefix, _ = get_causal_lm_model_cls_prefix(model_type)
module = __import__(module_path, fromlist=[f"{model_cls_prefix}MLP"])
mlp_cls = getattr(module, f"{model_cls_prefix}MLP")
setattr(mlp_cls, "forward", tiled_mlp_forward_common)
You can see the full code here:
This could be a built-in feature and enabled with a simple flag to from_pretrained(enable_tiled_mlp=True, ...)
Implementation-wise - most LLMs archs use the exact same MLP module, so the common case is very easy. Some are slightly different and would require some special cases. But perhaps those use cases could be added on demand if users request those, making the initial integration very simple (i.e. start with supporting just the models that use the common MLP.
TiledMLP lives here https://github.com/deepspeedai/DeepSpeed/blob/3631712bd796593bf6e2e70d4f1c352937a44cf8/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L807 but you can copy it if you don't want to make a dependency on deepspeed. It's a small independent self-contained autograd function.
cc: @Rocketknight1
Activity
Rocketknight1 commentedon Oct 24, 2025
Hmm, the relevant people here are probably:
@SunMarc @MekkCyber for trainer/deepspeed/accelerate
@ArthurZucker @Cyrilvallez as core maintainers
stas00 commentedon Nov 20, 2025
TiledMLP has just been integrated into Liger-Kernel, so it should be possible to just rely on it via a flag and just a soft-dependency on LK - or you could probably copy it - it's a very small class.
linkedin/Liger-Kernel#935
ArthurZucker commentedon Nov 21, 2025
BTW sorry for the late late answer, this fits very well within our
kernelsapproach @stas00 and is very welcome IMO! cc @MekkCyber !stas00 commentedon Nov 21, 2025
At the moment TiledMLP is pure pytorch, no kernels involved - even in the LK version - though they plan to change it down the road.
Incheonkirin commentedon Jun 17, 2026
I'd like to take this on.
Before I start implementing, I'd like to align on the integration shape:
autograd.Functionthat recomputes the MLP forward per shard. Given @ArthurZucker's note about thekernelsapproach, which direction would you prefer for v1: a pure-PyTorch path in Transformers core, a Liger soft-dependency exposed throughuse_liger_kernel/liger_kernel_config, or adapting it to thekernelsframework?Trainertraining loops can enable it?For the first pass I'd plan on output/gradient parity tests against the regular MLP,
gradcheckfor the custom backward, gradient-checkpointing interaction tests, FSDP /torch.compilesmoke tests where applicable, and a memory benchmark across sequence lengths.qgallouedec commentedon Aug 14, 2026
@Incheonkirin sorry for the long silence. Much of this has since been answered by code. TiledMLP is being done through
kernelsrather than as a pure-PyTorch path in core:From those PRs the shape looks like:
@use_kernel_forward_from_hubon the MLP class (so model-level, not Trainer-only), SwiGLU/GeGLU first. @vasqu can confirm the intended scope and whether anything is left for a contributor.Either way, probably best not to start a from-scratch implementation. Two things do look open:
torch.autograd.backwardonce per shard, so each weight'sAccumulateGradfires N times per backward. On a 2-layerLlamaForCausalLMunder 2-rank DDP that raisesExpected to mark a variable ready only once; it doesn't reproduce on a single-MLP model, so it's easy to miss. Passing the weights asautograd.Functioninputs gives one accumulation each and fixed it in a separate implementation (tested), but not unconditionally: under ZeRO-3 the params are 0-sized placeholders outside their leaf module's forward, and that path then fails withexpected shape compatible with [0]. Branching on whether the params carry ads_idworked there, though I haven't tried it against this kernel.nn.Linearmay clash and tile twice, to be confirmed.For these reasons, I'll keep this one open for now
vasqu commentedon Aug 24, 2026
Yea I think the biggest blocker was the double kernelize on linear modules when we have supersets like Swiglu MLP, see vasqu#2
But the idea is as shared by @qgallouedec: Use kernels and kernelize the nn linear instead of manual (monkey) patches. Re 1, I think it was fixed by Kashif iirc, but yea best to confirm things work as expected!