Add support for MiniMax's MiniMax-Text-01
Author: geetu040
Pull request: yes
Labels: name: New model

# What does this PR do?

Fixes https://github.com/huggingface/transformers/issues/35710

This PR adds [MiniMaxAI](https://www.minimaxi.com/en)'s MiniMax-Text-01 model to Hugging Face Transformers.

- MiniMax-Text-01 is a powerful language model with 456 billion total parameters, of which 45.9 billion are activated per token.
- To better unlock the long context capabilities of the model, MiniMax-Text-01 adopts a hybrid architecture that combines Lightning Attention, Softmax Attention and Mixture-of-Experts (MoE).
- MiniMax-Text-01's training context length is extended to 1 million tokens, and it can handle a context of up to 4 million tokens during the inference.
- On various academic benchmarks, MiniMax-Text-01 also demonstrates the performance of a top-tier model.

Relevant Links
- Research Paper: [MiniMax-01: Scaling Foundation Models with Lightning Attention](https://arxiv.org/abs/2501.08313)
- Authors: [MiniMax](https://arxiv.org/search/cs?searchtype=author&query=MiniMax), [Aonian Li](https://arxiv.org/search/cs?searchtype=author&query=Li,+A), [Bangwei Gong](https://arxiv.org/search/cs?searchtype=author&query=Gong,+B), et al.
- Implementation: [MiniMaxAI/MiniMax-Text-01](https://huggingface.co/MiniMaxAI/MiniMax-Text-01/tree/main)
- Models Weights: [MiniMaxAI/MiniMax-Text-01](https://huggingface.co/MiniMaxAI/MiniMax-Text-01)
- Github Page: [MiniMax-AI/MiniMax-01](https://github.com/MiniMax-AI/MiniMax-01)

CC: @MiniMax-AI-Dev

## Before submitting
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?

## Who can review?

@ArthurZucker, @Rocketknight1

## Change Log

- **Tokenizer**: It uses the existing `GPT2Tokenizer`
- **Config**: Matches the `MixtralConfig` with a few additional parameters:
  - residual_post_norm, attn_type_list
  - layernorm_attention_alpha, layernorm_lightning_attention_alpha, layernorm_mlp_alpha
  - layernorm_attention_beta, layernorm_lightning_attention_beta, layernorm_mlp_beta
- **Weight Conversion Script**: No script needed, original weights can be loaded directly into the new architecutre
- **Model**: `MiniMax-Text-01` architecture matches and uses most of the `Mixtral` architecture, with a few changes in
  - **DecoderLayer**
    - `hidden_states` can be used as residual connections, before or after `layernorm` is applied
    - weighted sum is used in residual connection
    - selection between softmax and lightning attention based on the `layer_idx`
  - **LightningAttention**
    - intially used in [TransNormerLLM](https://arxiv.org/abs/2307.14995), upgraded in [Lightning Attention-2](https://arxiv.org/abs/2401.04658) and adopted in [MiniMax-01](https://arxiv.org/abs/2501.08313)
    - every 8th decoder layer uses a softmax attention, rest of the layers use lightning attention, which is not previously implemented in `transformers`

**To summarize above, the main area of review is the `LightningAttention` implementation.**

## TODOs

- [x] Update Documentation
- [x] Update Tests
- [x] Import Statements and Auto Modeling
- [x] Implement Model
  - [x] Implement End-to-End Architecture
  - [x] Implement `LightningAttention`
    - [x] Work with avialable code
    - [x] Refactor, Clean and Optimize
    - [x] Implement Decays
    - [x] Implement Caching
    - [x] Implement attention_mask
    - [x] Support `.generate()` method
- [x] Fix CI/CD tests
- [x] [MiniMax-Tiny](https://huggingface.co/geetu040/MiniMax-tiny) for slow integration tests
- [ ] Model Card (only config and README file is to be updated)

Comments:
Hi @geetu040, this looks quite good! You can ping us whenever it's ready for review. Also, code quality issues in the CI can be fixed with `pip install -e .[quality]` in the `transformers` directory, followed by `make fixup`.

Eager to see this! 🔥 Feel free to ping @Cyrilvallez  and me for a review! 

cc @Cyrilvallez can you have a look!

@ArthurZucker, @Rocketknight1, @Cyrilvallez

This PR is ready for review.

Please see `MiniMaxText01Cache`, the lightning attention uses a cache that doesnot grow with the sequence length. Though I have implemented this in the best of my knowledge, still I am not sure if this cache can be traced back for beam search or low memory contrasitive search where we need to crop back the cache given the sequence

@Cyrilvallez this is up for review again!

@Cyrilvallez, I have addressed the review comments.

How do you suggest we go with the checkpoint at [MiniMaxAI/MiniMax-Text-01](https://huggingface.co/MiniMaxAI/MiniMax-Text-01)?
The original safetensors can be used without any changes in the keys of state_dict, only the config is slightly different. We removed the `postnorm` and renamed attributes like `layernorm_full_attention_alpha` to `full_attn_alpha_factor` and the model architecture is `MiniMaxForCausalLM` instead of `MiniMaxText01ForCausalLM`.
Do we need a new repo like `MiniMaxAI/MiniMax-hf` or can the config be updated in the original repo?

Hi, @geetu040 @Shakib-IO . Thank you so much for all your hard work! Just wondering how things are going with the progress. Is there anything I can do to help move the PR forward?

> Hi, @geetu040 @Shakib-IO . Thank you so much for all your hard work! Just wondering how things are going with the progress. Is there anything I can do to help move the PR forward?

Hi @sriting, thanks for checking in! The PR is complete and currently waiting on a review from @Cyrilvallez. In the meantime, I’ll update it with the latest from the main branch.

Since you're on the MiniMax team, I’ll give you write access in case you’d like to make any changes directly. Otherwise, we’d really appreciate your review!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_35831). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Do you want me to commit directly to help get this merged? 

I tested the MiniMax model by using this code.

Prompt is "Hello! How are you?"

```
quantized_model = MiniMaxForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype="bfloat16",
    device_map="auto",
    quantization_config=quantization_config,
    attn_implementation="flash_attention_2",
    trust_remote_code=True,
)
```
the output is garbled, for example:

`
この広告はメリカ就很.pankajoudhia成这样brakk a Поједина and aプライバerjakan회견 grandissimumb。这是因为 a,カウンセ a`

However, if I switch to:

```
quantized_model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype="bfloat16",
    device_map="auto",
    quantization_config=quantization_config,
    attn_implementation="flash_attention_2",
    trust_remote_code=True,
)
```
then the output is correct, such as:

`I'm doing well, thank you for asking! How about you?`

It seems like there's something wrong when using MiniMaxForCausalLM.from_pretrained.

@qscqesze Thanks for testing this out!

This could be happening because of minor difference in naming the config attributes

can you trying doing this with the config
```py
config = AutoConfig.from_pretrained(
    MODEL_PATH,
    trust_remote_code=True,
)

config = config.to_dict()
config['full_attn_alpha_factor'] = config['layernorm_full_attention_alpha']
config['full_attn_beta_factor'] = config['layernorm_full_attention_beta']
config['linear_attn_alpha_factor'] = config['layernorm_linear_attention_alpha']
config['linear_attn_beta_factor'] = config['layernorm_linear_attention_beta']
config['mlp_alpha_factor'] = config['layernorm_mlp_alpha']
config['mlp_beta_factor'] = config['layernorm_mlp_beta']

config = MiniMaxConfig(**config)

quantized_model = MiniMaxForCausalLM.from_pretrained(
    MODEL_PATH,
    config=config,
    torch_dtype="bfloat16",
    device_map="auto",
    quantization_config=quantization_config,
    attn_implementation="flash_attention_2",
    trust_remote_code=True,
)
```

Also, don't use the code from latest commit of this branch, I am fixing some bugs here, please use the previous commit https://github.com/huggingface/transformers/pull/35831/commits/9c5397d3220bfeb7abfbceb40e3a96f2aefa71de code

> @qscqesze Thanks for testing this out!
> 
> This could be happening because of minor difference in naming the config attributes
> 
> can you trying doing this with the config
> 
> ```python
> config = AutoConfig.from_pretrained(
>     MODEL_PATH,
>     trust_remote_code=True,
> )
> 
> config = config.to_dict()
> config['full_attn_alpha_factor'] = config['layernorm_full_attention_alpha']
> config['full_attn_beta_factor'] = config['layernorm_full_attention_beta']
> config['linear_attn_alpha_factor'] = config['layernorm_linear_attention_alpha']
> config['linear_attn_beta_factor'] = config['layernorm_linear_attention_beta']
> config['mlp_alpha_factor'] = config['layernorm_mlp_alpha']
> config['mlp_beta_factor'] = config['layernorm_mlp_beta']
> 
> config = MiniMaxConfig(**config)
> 
> quantized_model = MiniMaxForCausalLM.from_pretrained(
>     MODEL_PATH,
>     config=config,
>     torch_dtype="bfloat16",
>     device_map="auto",
>     quantization_config=quantization_config,
>     attn_implementation="flash_attention_2",
>     trust_remote_code=True,
> )
> ```
> 
> Also, don't use the code from latest commit of this branch, I am fixing some bugs here, please use the previous commit [9c5397d](https://github.com/huggingface/transformers/commit/9c5397d3220bfeb7abfbceb40e3a96f2aefa71de) code

Thanks for your code—everything works fine on my side now.

@ArthurZucker please review the PR, if everything looks fine, we can move to handling the checkpoints on huggingface.

Hello. I'm curious—why do we need to change attn_type_list to layer_type_list? I checked the config file at [config.json](https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/config.json), and it actually doesn't contain a layer_type_list configuration.

When I use the latest code, I get the following error:


```
You are using a model of type minimax_text_01 to instantiate a model of type MiniMaxText01. This is not supported for all configurations of models and can yield errors.
You are using a model of type mixtral to instantiate a model of type MiniMaxText01. This is not supported for all configurations of models and can yield errors.
Traceback (most recent call last):
  File "/home/qingjun/vllm_test/huggingface_test.py", line 62, in <module>
    quantized_model = MiniMaxForCausalLM.from_pretrained(
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/qingjun/transformers/src/transformers/modeling_utils.py", line 314, in _wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/qingjun/transformers/src/transformers/modeling_utils.py", line 4622, in from_pretrained
    model = cls(config, *model_args, **model_kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/qingjun/transformers/src/transformers/models/minimax/modeling_minimax.py", line 883, in __init__
    self.model = MiniMaxModel(config)
                 ^^^^^^^^^^^^^^^^^^^^
  File "/home/qingjun/transformers/src/transformers/models/minimax/modeling_minimax.py", line 657, in __init__
    [MiniMaxDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/qingjun/transformers/src/transformers/models/minimax/modeling_minimax.py", line 500, in __init__
    self.layer_type = config.layer_type_list[layer_idx]
                      ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^
IndexError: list index out of range
```

However, if I modify it like this:

`config['layer_type_list'] = config['attn_type_list']`

It works fine.

@qscqesze yes, some of the config attributes are named differently here in transformers as suggested at https://github.com/huggingface/transformers/pull/35831#discussion_r2102783512 and https://github.com/huggingface/transformers/pull/35831#discussion_r1960692259.

We renamed these for better clarify and to avoid confusion by names.

Are you fine with the changes, or would prefer to have the original names or suggest a better one?

CC: @ArthurZucker @Cyrilvallez (could we rename them back?)

I understand the reasoning behind renaming these attributes for better clarity, and I appreciate the effort to make things more consistent. However, one drawback is that users who download the config files directly from Hugging Face might encounter issues if they’re unaware of the renaming—since the current setup requires them to manually remap the keys to make the model work. This could easily lead to confusion, especially for users who aren’t closely following the codebase changes.

I’d love to hear what other reviewers think about this as well. 

@qscqesze We actually handle this by creating another repository on huggingface, in this case something like `MiniMaxAI/MiniMax-Text-01-hf` which works with the implementation in transformers

@qscqesze when you use `trust_remote_code=True,` it is not expected to work with the code of this PR. 
We will open PRs to the repo to update the remote code! 

The reason we need to rename is to adhere to the standard: `layer_type` is now shared across models! https://github.com/huggingface/transformers/blob/afe78bdb892f1a502706c34a482fc5cd8336df9a/src/transformers/models/gemma2/modular_gemma2.py#L197 this means we can enable better support for different layer type! 

@ArthurZucker @geetu040 Thank you both!
﻿
That makes perfect sense regarding the need for a separate Hugging Face repo and the renaming to align with the new `layer_type` standards.

> My last comment then let's merge! cc @LysandreJik to merge once this is fixed

Hi @LysandreJik, This PR is ready for merge!

These checkpoints for integration tests should also be moved from [`geetu040/MiniMax-tiny`](https://huggingface.co/geetu040/MiniMax-tiny) to [`hf-internal-testing/MiniMax-tiny`](https://huggingface.co/hf-internal-testing/MiniMax-tiny).

Hey @geetu040, I mirrored your tiny checkpoint here: https://huggingface.co/hf-internal-testing/MiniMax-tiny/tree/main

I'm merging the PR! 

Thanks for your contribution 🤗