Integrating TiledMLP for a much smaller memory footprint
Author: stas00
Pull request: no
Labels: name: Feature request

Similar to `TiledFusedLogitsLoss` https://github.com/huggingface/transformers/issues/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:

<img width="1477" height="386" alt="Image" src="https://github.com/user-attachments/assets/ef8a66c4-a268-4ceb-aade-d1e0c0be5ade" />

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:

https://github.com/snowflakedb/ArcticTraining/blob/5959f72709ac40433e94d09d095c021e7466cf0d/arctic_training/model/tiled_compute.py

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 

Comments:
Hmm, the relevant people here are probably:

@sunmarc @mekkcyber for trainer/deepspeed/accelerate
@arthurzucker @cyrilvallez as core maintainers

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. 

https://github.com/linkedin/Liger-Kernel/pull/935

BTW sorry for the late late answer, this fits very well within our `kernels` approach @stas00 and is very welcome IMO! cc @MekkCyber !

At the moment TiledMLP is pure pytorch, no kernels involved - even in the LK version - though they plan to change it down the road.

I'd like to take this on.

Before I start implementing, I'd like to align on the integration shape:

1. @stas00 noted that TiledMLP is currently pure PyTorch, using a custom `autograd.Function` that recomputes the MLP forward per shard. Given @ArthurZucker's note about the `kernels` approach, which direction would you prefer for v1: a pure-PyTorch path in Transformers core, a Liger soft-dependency exposed through `use_liger_kernel` / `liger_kernel_config`, or adapting it to the `kernels` framework?
2. For the user surface, should this be Trainer-facing only, or model-level as well so non-`Trainer` training loops can enable it?
3. Is supporting common SwiGLU MLPs first, with GEGLU and MoE as follow-ups, an acceptable v1 scope?

For the first pass I'd plan on output/gradient parity tests against the regular MLP, `gradcheck` for the custom backward, gradient-checkpointing interaction tests, FSDP / `torch.compile` smoke tests where applicable, and a memory benchmark across sequence lengths.