MsT: chunking the LM-head and MLP to extend sequence length and save memory
Author: wdlctc
Pull request: yes
Labels: name: Feature request

### Feature request
I would like to contribute an LM-head and MLP implementation that optimizes intermediate memory(logits and MLP intermediate). Chunking the LM-head and MLP over sequence dimension and accumulating the gradient during training and fine-tuning. This would not reduce throughput and mathematical equivalence to standard training. Combining with gradient checkpoint, we can extend 12x-24x sequence length over vanilla huggingface/transformers, and 4-7x sequence length over gradient checkpoint. This is all implemented in [this repo](https://github.com/wdlctc/transformers/)

### Usage:
from transformers.integrations import replace_with_minis
replace_with_minis()

### Motivation
State-of-the-art transformer models introduced larger tokenizers with a vocabulary of 128K tokens (Llama3) or 256 tokens (Gemma2). Training/Finetune these models can easily run out of GPU memory, and we found a vast majority of the memory is consumed by the intermediate memory(logits and MLP intermediate). The intermediate memory is not necessary to be stored on GPU memory, therefore, using chunking technologies for these intermediate memory can significant reduce the peak memory usage, like gradient accumulation. The memory save can be significantly enlarged with gradient checkpoint, where intermediate memory takes most activation memory.

### Your contribution
I provided an initial implementation on this https://github.com/wdlctc/transformers/. Not sure what is the right way to integrate it. Could be a modification of modeling_llama.py (and other model implementation), or an new class, not sure. I'm also not very familiar with the PR process, since this is my first issue, so maybe it would be better if someone from the HF team can shepherd this through.

### Additional context
Blog: https://wdlctc.github.io/mst.html
Model Finetune Guidence with our tech: [LLAMA3](https://github.com/wdlctc/mini-s/blob/main/doc/llama3.md), [Qwen2](https://github.com/wdlctc/mini-s/blob/main/doc/qwen.md), [Memba](https://github.com/wdlctc/mini-s/blob/main/doc/falcon-mamba.md), [Mistral](https://github.com/wdlctc/mini-s/blob/main/doc/mistral.md), [Gemma2](https://github.com/wdlctc/mini-s/blob/main/doc/gemma.md)
Arxiv: https://arxiv.org/abs/2407.15892

Comments:
> Hey! I like the idea, the issue is that such a change really does not scale! What would could do maybe to have a general class that monkey patches this!
> 
> Also this is more of a feature request, but chunking is indeed important for TTFT!
> 
> If we do this, we want something that does not involve any model changes!

We got it, we can do a monkey patching for LlamaForCausalLM.forward to achieve our goal. 

any more suggestions are welcome here. Then we would close this PR and propose a new one with monkey patching and feature request.

The main idea is to put this probably in `integration` or somewhere like this. 
We would also rather wait to see if the community is eager to have this merged! (via reaction on the proposed issue!)

> The main idea is to put this probably in `integration` or somewhere like this. We would also rather wait to see if the community is eager to have this merged! (via reaction on the proposed issue!)

We have made it in `integration` with monkey patches, and we highly welcome the community to try it!