Model Support : MiMo-Audio

[pengchengneo]

Model description

MiMo-Audio support

Open source examples / papers / references

• Paper: {https://github.com/XiaomiMiMo/MiMo-Audio/blob/main/MiMo-Audio-Technical-Report.pdf}
• Website: {https://github.com/XiaomiMiMo/MiMo-Audio}
• Model code: {https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/src}
• Model weights: {https://huggingface.co/collections/XiaomiMiMo/mimo-audio}

Why this model?

• {MiMo-Audio-7B-Instruct achieves open-source SOTA on audio understanding benchmarks, spoken dialogue benchmarks and instruct-TTS evaluations, approaching or surpassing closed-source models.}

[pengchengneo]

@Mozoltov821

[Mozoltov821]

MiMoAudio PyTorch to JAX Migration Plan.md

[pengchengneo]

MiMoAudio PyTorch to JAX Migration Plan

1. Project Background

Migrating the MiMoAudio MultiModal Machine Learning audio language model from the PyTorch ecosystem to the bonsai JAX framework, maintaining functional equivalence with the original version, and verifying the correctness of the model results are the prerequisite tasks for integrating MiMoAudio into sglang-jax.

2. Technical Architecture Design

2.1 Overall Architecture Adaptation

2.1.1 Follow Bonsai Design Principles

Since the llm backbone, patch encoder and patch decoder of mimo-audio are all based on the qwen2 architecture, it is necessary to implement the qwen2 model in bonsai first.

bonsai/models/qwen2/
├── README.md
├── __init__.py
├── modeling.py
├── params.py
└── tests/
    ├── __init__.py
    ├── run_model.py
    ├── test_outputs.py
    └── integration_tests.py

bonsai/models/mimoaudio/
├── README.md
├── __init__.py
├── modeling.py
├── params.py
├── audio_tokenizer.py
└── tests/
    ├── __init__.py
    ├── run_model.py
    ├── test_outputs.py
    └── integration_tests.py

2.1.2 Key Architecture Decisions

• Separate implementation: Main model (modeling.py) + Audio tokenizer (audio_tokenizer.py)
• Single File Principle: Each major component is fully implemented in a single file
• JAX First: Fully based on JAX + Flax NNX

2.2 Component Mapping Strategy

2.2.1 Core Model Component Mapping

PyTorch Components	Bonsai JAX Component
MiMoAudioForCausalLM	MiMoAudioModel
MiMoAudioTokenizer	AudioTokenizer
Qwen2Model	qwen2
ResidualVectorQuantizer	VectorQuantizer
MelSpectrogram	MelTransform

3. Development Phase Breakdown

Phase 1: Infrastructure

• Implement Qwen2
• Test and verify the correctness of the Qwen2 model

Deliverables: Qwen2 Model

Phase 2: Audio Processing Component

• Implement the JAX version of the Mel spectrogram processor
• Implement audio encoder and decoder (Corresponding to Audio_Encoder & Audio_Decoder in paper)
• Implementing Residual Vector Quantizer
• Integrate the above components into a complete audio tokenizer

Deliverables: audio_tokenizer.py, numerical validation test

Phase 3: Inference Model Component

• Implement the main class of MiMoAudioModel
• Completed input processing system
• LLM backbone(Mimo-7B-Base，based on Qwen2)
• Implement Patch Encoder and Decoder (Based on Qwen2 Model)

Deliverables: modeling.py, end-to-end inference functionality

Phase 4: Integration and Optimization

• Complete system integration and testing
• Documentation and examples improvement

Deliverables: Complete bonsai MiMoAudio implementation

4. Core Component Design

MiMoAudioModel Architecture

class MiMoAudioModel(nnx.Module):
    def __init__(self, config: MiMoAudioConfig, rngs: nnx.Rngs):
        # backbone LM
        self.model = Qwen2Model(config.qwen_config, rngs)

        # Corresponding to patch_decoder in the paper
        self.local_transformer = Qwen2Model(config.local_config, rngs)

        # Corresponding to patch_encoder in the paper
        self.input_local_transformer = Qwen2Model(config.input_local_config, rngs)

        self.speech_embeddings = [
            nnx.Embed(vocab_size, config.input_local_dim, rngs=rngs)
            for vocab_size in config.speech_vocab_sizes
        ]

        self.local_transformer_lm_heads = [
            nnx.Linear(config.local_dim, vocab_size, rngs=rngs)
            for vocab_size in config.speech_vocab_sizes
        ]

        self.speech_group_downcast = nnx.Linear(
            config.input_local_dim * config.group_size,
            config.hidden_size,
            rngs=rngs
        )

AudioTokenizer Architecture

class AudioTokenizer(nnx.Module):
    def __init__(self, config: AudioTokenizerConfig, rngs: nnx.Rngs):
        self.mel_transform = MelTransform(config)
        self.encoder = AudioEncoder(config, rngs)
        self.quantizer = ResidualVectorQuantizer(config, rngs)
        self.decoder = AudioDecoder(config, rngs)

[Mozoltov821]

MiMoAudio PyTorch to JAX Migration Plan

1. Project Background

Migrating the MiMoAudio MultiModal Machine Learning audio language model from the PyTorch ecosystem to the bonsai JAX framework, maintaining functional equivalence with the original version, and verifying the correctness of the model results are the prerequisite tasks for integrating MiMoAudio into sglang-jax.

2. Technical Architecture Design

2.1 Overall Architecture Adaptation

2.1.1 Follow Bonsai Design Principles

Since the llm backbone, patch encoder and patch decoder of mimo-audio are all based on the qwen2 architecture, it is necessary to implement the qwen2 model in bonsai first.

bonsai/models/qwen2/
├── README.md
├── __init__.py
├── modeling.py
├── params.py
└── tests/
    ├── __init__.py
    ├── run_model.py
    ├── test_outputs.py
    └── integration_tests.py

bonsai/models/mimoaudio/
├── README.md
├── __init__.py
├── modeling.py
├── params.py
├── audio_tokenizer.py
└── tests/
    ├── __init__.py
    ├── run_model.py
    ├── test_outputs.py
    └── integration_tests.py

2.1.2 Key Architecture Decisions

• Separate implementation: Main model (modeling.py) + Audio tokenizer (audio_tokenizer.py)
• Single File Principle: Each major component is fully implemented in a single file
• JAX First: Fully based on JAX + Flax NNX

2.2 Component Mapping Strategy

2.2.1 Core Model Component Mapping

PyTorch Components	Bonsai JAX Component
MiMoAudioForCausalLM	MiMoAudioModel
MiMoAudioTokenizer	AudioTokenizer
Qwen2Model	qwen2
ResidualVectorQuantizer	VectorQuantizer
MelSpectrogram	MelTransform

3. Development Phase Breakdown

Phase 1: Infrastructure

• Implement Qwen2
• Test and verify the correctness of the Qwen2 model

Deliverables: Qwen2 Model

Phase 2: Audio Processing Component

• Implement the JAX version of the Mel spectrogram processor
• Implement audio encoder and decoder (Corresponding to Audio_Encoder & Audio_Decoder in paper)
• Implementing Residual Vector Quantizer
• Integrate the above components into a complete audio tokenizer

Deliverables: audio_tokenizer.py, numerical validation test

Phase 3: Inference Model Component

• Implement the main class of MiMoAudioModel
• Completed input processing system
• LLM backbone(Mimo-7B-Base，based on Qwen2)
• Implement Patch Encoder and Decoder (Based on Qwen2 Model)

Deliverables: modeling.py, end-to-end inference functionality

Phase 4: Integration and Optimization

• Complete system integration and testing
• Documentation and examples improvement

Deliverables: Complete bonsai MiMoAudio implementation

4. Core Component Design

MiMoAudioModel Architecture

class MiMoAudioModel(nnx.Module):
    def __init__(self, config: MiMoAudioConfig, rngs: nnx.Rngs):
        # backbone LM
        self.model = Qwen2Model(config.qwen_config, rngs)

        # Corresponding to patch_decoder in the paper
        self.local_transformer = Qwen2Model(config.local_config, rngs)

        # Corresponding to patch_encoder in the paper
        self.input_local_transformer = Qwen2Model(config.input_local_config, rngs)

        self.speech_embeddings = [
            nnx.Embed(vocab_size, config.input_local_dim, rngs=rngs)
            for vocab_size in config.speech_vocab_sizes
        ]

        self.local_transformer_lm_heads = [
            nnx.Linear(config.local_dim, vocab_size, rngs=rngs)
            for vocab_size in config.speech_vocab_sizes
        ]

        self.speech_group_downcast = nnx.Linear(
            config.input_local_dim * config.group_size,
            config.hidden_size,
            rngs=rngs
        )

AudioTokenizer Architecture

class AudioTokenizer(nnx.Module):
    def __init__(self, config: AudioTokenizerConfig, rngs: nnx.Rngs):
        self.mel_transform = MelTransform(config)
        self.encoder = AudioEncoder(config, rngs)
        self.quantizer = ResidualVectorQuantizer(config, rngs)
        self.decoder = AudioDecoder(config, rngs)

[jenriver]

This would be a great addition! Feel free to open a PR once it's ready for review.

[pengchengneo]

@jenriver @chapman20j hello, we already finish this model , we will open PR soon, please help us review code, thanks !

[chapman20j]

Hi @SiqiLi-Fighting. Thanks for making the PR! We start taking a look today. Thanks for your hard work on this!