MiMo-7B-RL: Mô hình ngôn ngữ suy luận từ Xiaomi

Đội ngũ LLM-Core của Xiaomi giới thiệu MiMo-7B-RL, thách thức ý tưởng rằng khả năng suy luận hàng đầu trong AI đòi hỏi các mô hình khổng lồ. Mô hình 7 tỷ tham số này, được thiết kế đặc biệt cho các tác vụ toán học và lập trình, thể hiện hiệu suất cạnh tranh với các mô hình lớn hơn nhiều và các hệ thống chuyên biệt như o1-mini của OpenAI. Thành tựu này là kết quả của một chiến lược toàn diện tối ưu hóa toàn bộ vòng đời mô hình, chứng minh rằng khả năng suy luận mạnh mẽ có thể được khai phá trong các kiến trúc hiệu quả hơn.

MiMo-7B là gì

Sự phát triển của MiMo-7B dựa trên niềm tin rằng khả năng suy luận nền tảng của một mô hình được thiết lập trong giai đoạn tiền huấn luyện (pre-training). Mặc dù các giai đoạn tinh chỉnh (fine-tuning) sau này là cần thiết, nền tảng ban đầu vẫn đóng vai trò cực kỳ quan trọng. Xiaomi nhận thấy rằng nhiều mô hình nhỏ hơn gặp khó khăn với khả năng suy luận phức tạp vì quá trình huấn luyện cơ bản của chúng thiếu sự tiếp xúc đầy đủ với các mẫu logic.

Để khắc phục điều này, quá trình tiền huấn luyện của MiMo được thiết kế tỉ mỉ để tối đa hóa "mật độ mẫu suy luận". Điều này liên quan đến xử lý dữ liệu phức tạp: tăng cường trích xuất văn bản để nắm bắt các cấu trúc phức tạp trong tài liệu kỹ thuật và mã nguồn, áp dụng các bộ lọc đa chiều để tập trung các ví dụ suy luận, và tạo ra các bộ dữ liệu tổng hợp khổng lồ thể hiện các bước logic và giải quyết vấn đề. Một chiến lược kết hợp dữ liệu ba giai đoạn đã được sử dụng trong quá trình tiền huấn luyện, tận dụng khoảng 25 nghìn tỷ token để xây dựng mô hình MiMo-7B-Base.

Hơn nữa, Xiaomi đã tích hợp Dự đoán Đa Token (Multiple-Token Prediction - MTP) như một mục tiêu huấn luyện bổ trợ. Kỹ thuật này, trong đó mô hình dự đoán nhiều token phía trước, có khả năng tăng cường hiểu biết về các phụ thuộc phức tạp và có thể tăng tốc suy luận thông qua giải mã suy đoán (speculative decoding).

Học Tăng cường Nâng cao

Dựa trên mô hình MiMo-7B-SFT đã được tinh chỉnh, giai đoạn Học Tăng cường (Reinforcement Learning - RL) đặc biệt nhắm đến khả năng thành thạo toán học và lập trình. Một bộ dữ liệu chất lượng cao gồm 130.000 bài toán và bài lập trình được tuyển chọn cẩn thận, tất cả đều có thể kiểm chứng thông qua các kiểm tra dựa trên quy tắc (như unit test hoặc xác thực số học), đã tạo thành cơ sở cho quá trình huấn luyện.

Để đảm bảo cải thiện khả năng thực sự và tránh "gian lận phần thưởng" (reward hacking), chỉ sử dụng các phần thưởng chính xác dựa trên quy tắc, mang tính khách quan. Một hệ thống "phần thưởng mã nguồn dựa trên độ khó kiểm tra" mới đã được giới thiệu để giải quyết vấn đề phần thưởng thưa thớt (sparse reward) vốn có trong việc tạo mã phức tạp. Thay vì phần thưởng được ăn cả ngã về không, hệ thống này cấp điểm một phần khi vượt qua các trường hợp kiểm tra dễ hơn trong một bài toán, cung cấp tín hiệu gradient dày đặc hơn để mô hình học hỏi.

Hiệu quả cũng là yếu tố then chốt. Khi mô hình cải thiện, một chiến lược lấy mẫu lại dữ liệu đã giảm trọng số các bài toán dễ hơn, tập trung huấn luyện vào các ví dụ khó khăn hơn. Xiaomi cũng đã phát triển một "Công cụ Triển khai Liên tục" (Seamless Rollout Engine), một hạ tầng RL được tối ưu hóa tích hợp tạo liên tục, tính toán phần thưởng không đồng bộ và kết thúc sớm để giảm thiểu thời gian GPU nhàn rỗi, mang lại tốc độ huấn luyện đáng kể (2.29x) và xác thực (1.96x).

Họ MiMo-7B-RL: Tổng quan nhanh

Xiaomi đã phát hành một số mô hình thể hiện các giai đoạn phát triển:

Điểm chuẩn của MiMo-7B-RL

Kết quả đánh giá làm nổi bật thế mạnh của MiMo-7B-RL, đặc biệt khi so sánh với các mô hình hàng đầu sử dụng nhiệt độ tạo (generation temperature) là 0.6.

Hiệu suất so sánh:

(Hiển thị các điểm chuẩn toán học/mã nguồn được chọn lọc)

MiMo-7B-RL thể hiện hiệu suất vượt trội trong toán học và lập trình, thường vượt qua các mô hình lớn hơn đáng kể và các mô hình suy luận chuyên biệt như o1-mini trên các điểm chuẩn thách thức như MATH, AIME và các phiên bản LiveCodeBench gần đây. Mặc dù điểm số suy luận chung của nó mạnh mẽ so với kích thước, chúng tự nhiên vẫn kém hơn các mô hình tiên tiến lớn nhất, phản ánh trọng tâm huấn luyện chuyên biệt của nó.

Hiệu suất trong Dòng MiMo:

So sánh nội bộ này minh họa hiệu quả của từng giai đoạn huấn luyện. Mô hình cơ sở cho thấy khả năng suy luận ban đầu mạnh mẽ, được tăng cường đáng kể bởi SFT, và được tinh chỉnh thêm để đạt hiệu suất cao nhất bởi giai đoạn RL cuối cùng nhắm vào toán học và mã nguồn. Áp dụng RL trực tiếp cho mô hình cơ sở (RL-Zero) cũng hiệu quả, nhưng bước trung gian SFT dường như có lợi để đạt được điểm số cao nhất.

Chạy MiMo-7B-RL

Các mô hình có sẵn trên Hugging Face Hub.

Truy cập Mô hình:

Tìm MiMo-7B-RL và các mô hình khác trong dòng tại trang tổ chức XiaomiMiMo trên Hugging Face. Kích thước mô hình khoảng 7.83 tỷ tham số (độ chính xác BF16, định dạng Safetensors).

Chạy Suy luận với vLLM (Được khuyến nghị)

Xiaomi khuyến nghị sử dụng fork vLLM của họ (dựa trên v0.7.3) để chạy suy luận, vì nó hỗ trợ tính năng Dự đoán Đa Token (Multi-Token Prediction) để có khả năng tạo nhanh hơn.

• Sử dụng Xiaomi vLLM Fork (có MTP):

# Ensure Xiaomi's vLLM fork is installed
from vllm import LLM, SamplingParams

# --- FACTUAL CODE SNIPPET START ---
# Source: https://huggingface.co/XiaomiMiMo/MiMo-7B-RL Model Card
model_path = "/path/to/XiaomiMiMo/MiMo-7B-RL" # Replace with your download path

llm = LLM(
    model=model_path,
    trust_remote_code=True,  # Required for MiMo's custom code
    num_speculative_tokens=1, # Enables MTP speculative decoding
    disable_log_stats=False
)
# Recommended sampling temperature for benchmark replication
sampling_params = SamplingParams(temperature=0.6)

# Example conversation structure (empty system prompt recommended)
conversation = [
    {
        "role": "system",
        "content": "" # Use an empty system prompt
    },
    {
        "role": "user",
        "content": "Write a python function to compute the nth Fibonacci number.",
    },
]

# Generate the response
outputs = llm.chat(conversation,
                   sampling_params=sampling_params,
                   use_tqdm=False)

# Process and print output
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}")
    print("-" * 20)
    print(f"Generated text: {generated_text!r}")
# --- FACTUAL CODE SNIPPET END ---

print("=" * 80)

• Sử dụng Standard vLLM (không có MTP):
  Nếu không sử dụng tính năng MTP hoặc sử dụng bản dựng vLLM tiêu chuẩn, hãy đăng ký kiến trúc MiMo trước bằng cách sử dụng script register_mimo_in_vllm.py do Xiaomi cung cấp.

# --- FACTUAL CODE SNIPPET START ---
# Source: https://huggingface.co/XiaomiMiMo/MiMo-7B-RL Model Card
# Ensure register_mimo_in_vllm.py is accessible
import register_mimo_in_vllm

from vllm import LLM, SamplingParams

model_path = "/path/to/XiaomiMiMo/MiMo-7B-RL" # Replace with your download path
llm = LLM(
    model=model_path,
    trust_remote_code=True,
    # Do not set num_speculative_tokens if not using MTP
    disable_log_stats=False
)
sampling_params = SamplingParams(temperature=0.6)

# Conversation setup and generation call is the same as the MTP example...
conversation = [
    {"role": "system", "content": ""},
    {"role": "user", "content": "Write a python function to compute the nth Fibonacci number."},
]
outputs = llm.chat(conversation, sampling_params=sampling_params, use_tqdm=False)
# Processing output is the same...
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}\n{'-'*20}\nGenerated text: {generated_text!r}")
# --- FACTUAL CODE SNIPPET END ---

Sử dụng HuggingFace Transformers

Cũng có thể chạy suy luận bằng thư viện transformers tiêu chuẩn của HuggingFace. Lưu ý rằng trust_remote_code=True là cần thiết.

# --- FACTUAL CODE SNIPPET START ---
# Source: https://huggingface.co/XiaomiMiMo/MiMo-7B-RL Model Card
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "/path/to/XiaomiMiMo/MiMo-7B-RL" # Replace with your download path

# Load the model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    trust_remote_code=True, # Essential for loading MiMo
    device_map="auto"       # Use GPU if available
)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# Prepare the input prompt
prompt = "Write a python function to compute the nth Fibonacci number."
# Tokenize the input
inputs = tokenizer([prompt], return_tensors='pt').to(model.device)

# Generate the output sequence
output_sequences = model.generate(
    **inputs,
    max_new_tokens=256,      # Control output length
    temperature=0.6,         # Recommended temperature
    do_sample=True           # Use sampling for temperatures != 1.0
)

# Decode the output
generated_text = tokenizer.decode(output_sequences[0], skip_special_tokens=True)
print(generated_text)
# --- FACTUAL CODE SNIPPET END ---

Khuyến nghị Sử dụng

Để có kết quả tốt nhất, đặc biệt khi cố gắng tái tạo điểm chuẩn, hãy sử dụng thiết lập được khuyến nghị: fork vLLM của Xiaomi (dựa trên v0.7.3) và lời nhắc hệ thống (system prompt) rỗng.

Lời kết: Khả năng suy luận hiệu quả được Xiaomi hiện thực hóa?

MiMo-7B-RL của Xiaomi chứng minh rằng hiệu suất suy luận đặc biệt trong các lĩnh vực chuyên biệt như toán học và lập trình có thể đạt được mà không cần dùng đến các mô hình có kích thước khổng lồ. Thông qua quá trình tiền huấn luyện cẩn thận tập trung vào các mẫu suy luận và các kỹ thuật học tăng cường đổi mới, họ đã tạo ra một mô hình hiệu quả cạnh tranh hiệu quả với các đối thủ lớn hơn nhiều. Việc phát hành mở dòng MiMo cung cấp các công cụ và hiểu biết có giá trị, thúc đẩy sự phát triển của các khả năng suy luận AI mạnh mẽ, dễ tiếp cận.