mimi_0.6b

Mimi-0.6B

感谢

• Qwen/Qwen3-4B-Instruct-2507
• llamafactory/alpaca_gpt4_zh
• Google Gemini 3 Pro

📣 项目宣言 (Manifesto)

模型的激动在于,自己从0开始做,做出一个模型,至少还给你说人话,虽然问题很多,就像小孩子叫那声"爸爸",那个心情...

本项目旨在完整记录一次从零开始的独立模型制作之旅。虽然最终模型问题很多,数据集覆盖面也不好,但关键在于我们实现了 从 0 到 0.1 的跃进。从训练数据到踩坑脚本,我全部开源,期待您去感受、改进,并创造出属于您自己的模型。

📖 模型介绍 (Introduction)

Mimi-0.6B 是一个由 Qwen3-4B 蒸馏而来的超轻量级语言模型。该模型在 NVIDIA RTX 5090 上进行全参数蒸馏 (Full-Parameter Distillation)，旨在探索消费级显卡下的高质量小模型训练极限。

It is a distilled version of Qwen3-4B, trained from scratch on a single RTX 5090. All training scripts and data are open-sourced.

🛠️ 训练配方与哲学 (Training Recipe & Philosophy)

开箱即用的训练套件

本仓库 train 目录中包括了一切所需文件，旨在实现一键复现：

• mimi_0.6b_distilled-train.py: 核心训练脚本
• mimi_0.6b_distilled-chat-test.py: 训练完成后的测试脚本
• test_token.py: 用于分析数据集 Token 长度
• alpaca_gpt4_zh/: 训练数据集
• Qwen3-4B-Instruct-2507/: 教师模型

只需运行 python mimi_0.6b_distilled-train.py，你就能得到一个属于自己的 0.6B 模型。

💡 核心理念与经验分享

AI 时代最好的制作模型工具,就是用最好的 AI。 任何代码不清楚的地方,直接上传到 Google Gemini 1.5 Pro,你会得到一切,包括如何根据你的硬件修改。能碰到的坑我都踩了,剩下就是您如何去做您自己心目中那个模型“儿子”或“女儿”了 XP。

经验：用任何技术完善自己的模型，不要觉得蒸馏等不好。 模型做好了,没人关心如何做的,关键是你的模型能用吗?

问题解释：我的模型只是展示一个可蒸馏的成果。 这个模型具体是什么,能做什么,考题能做多少分? 不在于我的展示,而在于看到模型的你,你要做出什么模型?

硬件环境 (Hardware)

• GPU: NVIDIA GeForce RTX 5090 (24G VRAM used)
• Power Consumption: ~580W Peak
• Strategy: Batch Size=1, Gradient Accumulation=32 (Effective BS=32)

核心坑点与解决方案 (Key Challenges & Solutions)

本项目在 Windows 原生环境 下完成，以下是关键避坑指南：

1. 报错: [WinError 1224] 或 DataLoader worker exited

   • 原因: Windows 多进程锁死。
   • 解决: TrainingArguments 里必须设 dataloader_num_workers=0，且所有执行代码放在 if __name__ == "__main__": 内。

2. 报错: CUDA OOM (爆显存)

   • 原因 1: Teacher 模型计算 Loss 瞬间申请大显存。
   • 解决 1: 在 Trainer 的 compute_loss 里，将 labels 从 inputs 剔除。
   • 原因 2: 显存碎片化。
   • 解决 2: 脚本开头加 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"。

3. 训练策略 (5090 专用)

   • per_device_train_batch_size = 1 (设为 2 必爆)
   • gradient_accumulation_steps = 32 (用时间换空间)
   • gradient_checkpointing = True (不开必爆)

🚀 如何使用 (How to Use)

不需要 trust_remote_code，直接使用 transformers 库即可。

from transformers import pipeline

# 你的模型保存路径
model_path = "aifeifei798/mimi_0.6b" # 已上传至 Hugging Face

print(f"正在加载训练好的模型: {model_path} ...")
generator = pipeline('text-generation', model=model_path, tokenizer=model_path, device="cuda:0")

# 测试集
prompts = [
    "User: 你是谁？\nAssistant:",
    "User: 为什么天空是蓝色的？\nAssistant:",
    "User: 写一首关于春天的诗。\nAssistant:",
    "User: 保持健康的建议。\nAssistant:"
]

print("="*50)
for prompt in prompts:
    outputs = generator(
        prompt, 
        max_new_tokens=256,
        do_sample=True, 
        temperature=0.7,
        top_k=50,
        repetition_penalty=1.2,
        top_p=0.9,
    )
    print(f"输入: {prompt.strip()}")
    # 精确提取 Assistant 的回答
    full_text = outputs['generated_text']
    assistant_response = full_text.split("Assistant:").strip()
    print(f"输出: {assistant_response}")
    print("-" * 50)