Fuse minimind/minimind-v/minimind-o into omni with layered architecture

Migrate the three MiniMind repos into the omni package using a clean
layered design:

- core/: split the monolithic Transformer kernel into norm/rope/attention/
mlp/block/model modules (replacing the old components.py)
- models/: assemble core components into MiniMindForCausalLM (text),
MiniMindVLM (vision), MiniMindOmni (speech/full-modal, with TalkerModule)
- encoders/: SiglipVisionEncoder (vision/), SenseVoiceAudioEncoder (audio/)
- projectors/: MMVisionProjector / MMAudioProjector bridging encoders to LLM
- serve/: SileroVAD / RealtimeSession realtime speech session layer
- trainers/ + datasets/ + utils/: ported train_sft/pretrain/dpo/etc plus
VLM/Omni trainers, VLMDataset/OmniDataset, init_vlm/omni_model helpers
- examples/: eval_llm/vlm/omni, omni_web_demo, serve_openai_api, convert_model
- README + pyproject updated; all modules import-compile-smoke verified

Co-Authored-By: opencode <noreply@opencode.ai>

README.md

@@ -0,0 +1,124 @@
+# Omni
+
+Omni 是一个以 **多模态 (omni)** 为目标的 LLM 训练 / 推理框架，已完整集成
+[MiniMind](https://github.com/jingyaogong/minimind)（纯文本）、
+[miniMind-V](https://github.com/jingyaogong/minimind-v)（视觉多模态）与
+[miniMind-O](https://github.com/jingyaogong/minimind-o)（语音 / 全模态）三套代码。
+
+项目采用标准 `src/` 布局（`pip install -e .` 即可安装为 `omni` 包），
+按 **core（组件）/ models（拼装）/ encoders（模态编码器）/ projectors（桥接层）** 分层。
+
+## 设计分层
+
+- `core/`：可复用模型**底层组件**，按层级细分为独立模块——
+  `norm.py`（`RMSNorm`）、`rope.py`（`precompute_freqs_cis` / `apply_rotary_pos_emb` / `repeat_kv`）、
+  `attention.py`（`Attention`）、`mlp.py`（`FeedForward` / `MOEFeedForward`）、
+  `block.py`（`MiniMindBlock`）、`model.py`（`MiniMindModel` Transformer 主干）。
+- `models/`：把 `core` 组件**拼装**成成品模型——
+  `MiniMindForCausalLM`（文本）、`MiniMindVLM`（视觉）、`MiniMindOmni`（语音/全模态）。
+- `encoders/`：外部模态编码器，按模态分目录——`vision/`（SigLIP）、`audio/`（SenseVoice）。
+- `projectors/`：把 encoder 输出**桥接**到 LLM 隐藏维度的拼接层（`MMVisionProjector`、`MMAudioProjector`）。
+- `serve/`：实时语音会话工程层（`SileroVAD`、`RealtimeSession`）。
+
+## 目录结构
+
+```
+src/omni/
+├── core/               # 模型底层组件（按层级拆分）
+│   ├── norm.py         #   RMSNorm
+│   ├── rope.py         #   precompute_freqs_cis / apply_rotary_pos_emb / repeat_kv
+│   ├── attention.py    #   Attention
+│   ├── mlp.py          #   FeedForward / MOEFeedForward
+│   ├── block.py        #   MiniMindBlock
+│   └── model.py        #   MiniMindModel（Transformer 主干）
+├── models/             # 模型拼装
+│   ├── minimind.py     # MiniMindConfig + MiniMindForCausalLM
+│   ├── vlm.py          # VLMConfig + MiniMindVLM（视觉多模态）
+│   ├── omni.py         # OmniConfig + MiniMindOmni + TalkerModule（语音/全模态）
+│   └── lora.py         # LoRA 注入 / 保存 / 合并
+├── encoders/           # 多模态编码器（按模态分目录）
+│   ├── vision/        #   SiglipVisionEncoder
+│   └── audio/         #   SenseVoiceAudioEncoder
+├── projectors/         # 多模态桥接层
+│   ├── vision.py      #   MMVisionProjector
+│   └── audio.py       #   MMAudioProjector
+├── trainers/           # 训练脚本（可直接 python -m 运行）
+│   ├── pretrain.py     # 文本预训练
+│   ├── full_sft.py     # 文本全量 SFT
+│   ├── lora.py / dpo.py / distillation.py / ppo.py / grpo.py / agent.py
+│   ├── rollout_engine.py  # torch / sglang 推理引擎
+│   ├── train_tokenizer.py # tokenizer 训练（学习用）
+│   ├── pretrain_vlm.py # 视觉预训练
+│   ├── full_sft_vlm.py # 视觉 SFT
+│   └── full_sft_omni.py  # 全模态 SFT
+├── datasets/           # 数据集（Pretrain/SFT/DPO/RLAIF/Agent/VLM/Omni）
+│   └── lm_dataset.py
+├── utils/             # 工具
+│   ├── training.py     # get_lr / init_model / lm_checkpoint / SkipBatchSampler / LMForRewardModel
+│   ├── multimodal.py   # init_vlm_model / vlm_checkpoint / init_omni_model / omni_checkpoint
+│   ├── distributed.py  # 分布式初始化
+│   └── checkpoint.py   # checkpoint 读写辅助
+├── serve/             # 实时语音会话（SileroVAD / RealtimeSession）
+└── __init__.py
+examples/               # 推理 / 服务 / 转换脚本
+├── eval_llm.py        # 命令行推理与对话
+├── eval_vlm.py        # 视觉多模态推理
+├── eval_omni.py       # 全模态推理
+├── serve_openai_api.py # OpenAI 兼容 API 服务
+├── omni_web_demo.py   # 网页演示（含实时语音）
+├── eval_toolcall.py   # 工具调用评测
+└── convert_model.py   # torch <-> transformers 权重互转
+weights/MiniMind2/      # tokenizer 与模型配置（从 MiniMind 迁移）
+configs/                # 训练配置（按需补充）
+```
+
+## 安装
+
+```bash
+pip install -e .
+# 可选依赖：RL 训练 / API 服务 / 演示
+pip install -e ".[rl,serve,demo]"
+```
+
+## 快速开始
+
+### 推理 / 对话
+
+```bash
+python -m examples.eval_llm --load_from weights/MiniMind2 --weight full_sft
+```
+
+### 训练
+
+每个训练脚本都是 `omni.trainers` 下的一个模块，直接运行即可：
+
+```bash
+# 预训练
+python -m omni.trainers.pretrain --data_path dataset/pretrain.jsonl
+# 全量 SFT
+python -m omni.trainers.full_sft --data_path dataset/sft.jsonl
+# LoRA 微调
+python -m omni.trainers.lora --data_path dataset/lora.jsonl
+# DPO / 蒸馏 / PPO / GRPO / Agent RL
+python -m omni.trainers.dpo    --data_path dataset/dpo.jsonl
+python -m omni.trainers.distillation --data_path dataset/sft.jsonl
+python -m omni.trainers.grpo    --data_path dataset/rlaif.jsonl
+python -m omni.trainers.agent  --data_path dataset/agent_rl.jsonl
+```
+
+所有脚本参数与原 MiniMind 保持一致（hidden_size / num_hidden_layers / use_moe / data_path 等）。
+
+## 与原 MiniMind 的差异
+
+- 去除所有 `sys.path` 注入 hack，统一使用 `omni.*` 包导入；
+- `trainer_utils.py` 拆分为 `utils/training.py`（训练工具）、`utils/distributed.py`、`utils/checkpoint.py`；
+- 训练脚本从「`if __name__ == '__main__'` 内联」改为可被 `python -m omni.trainers.<name>` 调用的模块；
+- 预留 `encoders/`（vision/audio）、`projectors/`、`core/` 供多模态扩展。
+
+## 多模态扩展方向
+
+在 `models/minimind.py` 的 `MiniMindModel` 之上接入：
+
+1. `encoders/vision` / `encoders/audio` —— 各自的模态编码器；
+2. `projectors` —— 将编码器输出投影到 LLM 隐藏维度；
+3. 在 `MiniMindModel.forward` 中把投影特征拼接到 `embed_tokens` 之后。

configs/example_train.yaml

@@ -0,0 +1,26 @@
+# Omni / MiniMind 训练配置示例
+# 用法：将此处参数直接作为命令行参数传给对应训练模块，例如
+#   python -m omni.trainers.full_sft --hidden_size 768 --num_hidden_layers 8 --use_moe 0 --data_path dataset/sft.jsonl
+
+model:
+  hidden_size: 768
+  num_hidden_layers: 8
+  use_moe: 0            # 1 启用 MoE
+  vocab_size: 6400
+  max_seq_len: 768
+
+train:
+  epochs: 2
+  batch_size: 16
+  learning_rate: 1.0e-5
+  accumulation_steps: 1
+  grad_clip: 1.0
+  dtype: bfloat16
+  save_interval: 1000
+  log_interval: 100
+  from_weight: pretrain   # none / pretrain / full_sft / ...
+  from_resume: 0
+
+paths:
+  save_dir: out
+  data_path: dataset/sft.jsonl

dataset/dataset.md

@@ -0,0 +1,5 @@
+# MiniMind Datasets
+
+将所有下载的数据集文件放置到当前目录.
+
+Place the downloaded dataset file in the current directory.

examples/convert_model.py

@@ -0,0 +1,141 @@
+import os
+import json
+
+import torch
+import transformers
+import warnings
+from transformers import AutoTokenizer, AutoModelForCausalLM, Qwen3Config, Qwen3ForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM
+from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models.lora import apply_lora, merge_lora
+
+warnings.filterwarnings('ignore', category=UserWarning)
+
+def convert_torch2transformers_minimind(torch_path, transformers_path, dtype=torch.float16):
+    MiniMindConfig.register_for_auto_class()
+    MiniMindForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+    lm_model = MiniMindForCausalLM(lm_config)
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    state_dict = torch.load(torch_path, map_location=device)
+    lm_model.load_state_dict(state_dict, strict=False)
+    lm_model = lm_model.to(dtype)  # 转换模型权重精度
+    model_params = sum(p.numel() for p in lm_model.parameters() if p.requires_grad)
+    print(f'模型参数: {model_params / 1e6} 百万 = {model_params / 1e9} B (Billion)')
+    lm_model.save_pretrained(transformers_path, safe_serialization=False)
+    tokenizer = AutoTokenizer.from_pretrained('../model/')
+    tokenizer.save_pretrained(transformers_path)
+    # ======= transformers-5.0的兼容低版本写法 =======
+    if int(transformers.__version__.split('.')[0]) >= 5:
+        tokenizer_config_path, config_path = os.path.join(transformers_path, "tokenizer_config.json"), os.path.join(transformers_path, "config.json")
+        json.dump({**json.load(open(tokenizer_config_path, 'r', encoding='utf-8')), "tokenizer_class": "PreTrainedTokenizerFast", "extra_special_tokens": {}}, open(tokenizer_config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
+        config = json.load(open(config_path, 'r', encoding='utf-8'))
+        config['rope_theta'] = lm_config.rope_theta; config['rope_scaling'] = None; del config['rope_parameters']
+        json.dump(config, open(config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
+    print(f"模型已保存为 Transformers-MiniMind 格式: {transformers_path}")
+
+
+# QwenForCausalLM/LlamaForCausalLM结构兼容生态
+def convert_torch2transformers(torch_path, transformers_path, dtype=torch.float16):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    state_dict = torch.load(torch_path, map_location=device)
+    common_config = {
+        "vocab_size": lm_config.vocab_size,
+        "hidden_size": lm_config.hidden_size,
+        "intermediate_size": lm_config.intermediate_size,
+        "num_hidden_layers": lm_config.num_hidden_layers,
+        "num_attention_heads": lm_config.num_attention_heads,
+        "num_key_value_heads": lm_config.num_key_value_heads,
+        "head_dim": lm_config.hidden_size // lm_config.num_attention_heads,
+        "max_position_embeddings": lm_config.max_position_embeddings,
+        "rms_norm_eps": lm_config.rms_norm_eps,
+        "rope_theta": lm_config.rope_theta,
+        "tie_word_embeddings": lm_config.tie_word_embeddings
+    }
+    if not lm_config.use_moe:
+        qwen_config = Qwen3Config(
+            **common_config, 
+            use_sliding_window=False, 
+            sliding_window=None
+        )
+        qwen_model = Qwen3ForCausalLM(qwen_config)
+    else:
+        qwen_config = Qwen3MoeConfig(
+            **common_config,
+            num_experts=lm_config.num_experts,
+            num_experts_per_tok=lm_config.num_experts_per_tok,
+            moe_intermediate_size=lm_config.moe_intermediate_size,
+            norm_topk_prob=lm_config.norm_topk_prob
+        )
+        qwen_model = Qwen3MoeForCausalLM(qwen_config)
+        # ======= transformers-5.0的兼容低版本写法 =======
+        if int(transformers.__version__.split('.')[0]) >= 5:
+            new_sd = {k: v for k, v in state_dict.items() if 'experts.' not in k or 'gate.weight' in k}
+            for l in range(lm_config.num_hidden_layers):
+                p = f'model.layers.{l}.mlp.experts'
+                new_sd[f'{p}.gate_up_proj'] = torch.cat([torch.stack([state_dict[f'{p}.{e}.gate_proj.weight'] for e in range(lm_config.num_experts)]), torch.stack([state_dict[f'{p}.{e}.up_proj.weight'] for e in range(lm_config.num_experts)])], dim=1)
+                new_sd[f'{p}.down_proj'] = torch.stack([state_dict[f'{p}.{e}.down_proj.weight'] for e in range(lm_config.num_experts)])
+            state_dict = new_sd
+
+    qwen_model.load_state_dict(state_dict, strict=True)
+    qwen_model = qwen_model.to(dtype)  # 转换模型权重精度
+    qwen_model.save_pretrained(transformers_path)
+    model_params = sum(p.numel() for p in qwen_model.parameters() if p.requires_grad)
+    print(f'模型参数: {model_params / 1e6} 百万 = {model_params / 1e9} B (Billion)')
+    tokenizer = AutoTokenizer.from_pretrained('../model/')
+    tokenizer.save_pretrained(transformers_path)
+
+    # ======= transformers-5.0的兼容低版本写法 =======
+    if int(transformers.__version__.split('.')[0]) >= 5:
+        tokenizer_config_path, config_path = os.path.join(transformers_path, "tokenizer_config.json"), os.path.join(transformers_path, "config.json")
+        json.dump({**json.load(open(tokenizer_config_path, 'r', encoding='utf-8')), "tokenizer_class": "PreTrainedTokenizerFast", "extra_special_tokens": {}}, open(tokenizer_config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
+        config = json.load(open(config_path, 'r', encoding='utf-8'))
+        config['rope_theta'] = lm_config.rope_theta; config['rope_scaling'] = None; del config['rope_parameters']
+        json.dump(config, open(config_path, 'w', encoding='utf-8'), indent=2, ensure_ascii=False)
+    print(f"模型已保存为 Transformers 格式: {transformers_path}")
+
+
+def convert_transformers2torch(transformers_path, torch_path):
+    model = AutoModelForCausalLM.from_pretrained(transformers_path, trust_remote_code=True)
+    torch.save({k: v.cpu().half() for k, v in model.state_dict().items()}, torch_path)
+    print(f"模型已保存为 PyTorch 格式: {torch_path}")
+
+
+def convert_merge_base_lora(base_torch_path, lora_path, merged_torch_path):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    lm_model = MiniMindForCausalLM(lm_config).to(device)
+    state_dict = torch.load(base_torch_path, map_location=device)
+    lm_model.load_state_dict(state_dict, strict=False)
+    apply_lora(lm_model)
+    merge_lora(lm_model, lora_path, merged_torch_path)
+    print(f"LoRA 已合并并保存为基模结构 PyTorch 格式: {merged_torch_path}")
+
+
+def convert_jinja_to_json(jinja_path):
+    with open(jinja_path, 'r') as f: template = f.read()
+    escaped = json.dumps(template)
+    print(f'"chat_template": {escaped}')
+
+
+def convert_json_to_jinja(json_file_path, output_path):
+    with open(json_file_path, 'r') as f: config = json.load(f)
+    template = config['chat_template']
+    with open(output_path, 'w') as f: f.write(template)
+    print(f"模板已保存为 jinja 文件: {output_path}")
+
+
+if __name__ == '__main__':
+    lm_config = MiniMindConfig(hidden_size=768, num_hidden_layers=8, max_seq_len=8192, use_moe=False)
+
+    # convert torch to transformers
+    torch_path = f"../out/full_sft_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+    transformers_path = '../minimind-3'
+    convert_torch2transformers(torch_path, transformers_path)
+
+    # # merge lora
+    # base_torch_path = f"../out/full_sft_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+    # lora_path = f"../out/lora_identity_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+    # merged_torch_path = f"../out/merge_identity_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+    # convert_merge_base_lora(base_torch_path, lora_path, merged_torch_path)
+
+    # convert_transformers2torch(transformers_path, torch_path)
+    # convert_json_to_jinja('../model/tokenizer_config.json', '../model/chat_template.jinja')
+    # convert_jinja_to_json('../model/chat_template.jinja')

examples/eval_llm.py

@@ -0,0 +1,94 @@
+import time
+import argparse
+import random
+import warnings
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models.lora import *  # noqa: F401,F403
+from omni.utils.training import setup_seed, get_model_params
+warnings.filterwarnings('ignore')
+
+def init_model(args):
+    tokenizer = AutoTokenizer.from_pretrained(args.load_from)
+    if 'model' in args.load_from:
+        model = MiniMindForCausalLM(MiniMindConfig(
+            hidden_size=args.hidden_size,
+            num_hidden_layers=args.num_hidden_layers,
+            use_moe=bool(args.use_moe),
+            inference_rope_scaling=args.inference_rope_scaling
+        ))
+        moe_suffix = '_moe' if args.use_moe else ''
+        ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
+        model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
+        if args.lora_weight != 'None':
+            apply_lora(model)
+            load_lora(model, f'./{args.save_dir}/{args.lora_weight}_{args.hidden_size}.pth')
+    else:
+        model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
+    get_model_params(model, model.config)
+    return model.half().eval().to(args.device), tokenizer
+
+def main():
+    parser = argparse.ArgumentParser(description="MiniMind模型推理与对话")
+    parser.add_argument('--load_from', default='model', type=str, help="模型加载路径（model=原生torch权重，其他路径=transformers格式）")
+    parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
+    parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀（pretrain, full_sft, rlhf, reason, ppo_actor, grpo, spo）")
+    parser.add_argument('--lora_weight', default='None', type=str, help="LoRA权重名称（None表示不使用，可选：lora_identity, lora_medical）")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--inference_rope_scaling', default=False, action='store_true', help="启用RoPE位置编码外推（4倍，仅解决位置编码问题）")
+    parser.add_argument('--max_new_tokens', default=8192, type=int, help="最大生成长度（注意：并非模型实际长文本能力）")
+    parser.add_argument('--temperature', default=0.85, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
+    parser.add_argument('--top_p', default=0.95, type=float, help="nucleus采样阈值（0-1）")
+    parser.add_argument('--open_thinking', default=0, type=int, help="是否开启自适应思考（0=否，1=是）")
+    parser.add_argument('--historys', default=0, type=int, help="携带历史对话轮数（需为偶数，0表示不携带历史）")
+    parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）")
+    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
+    args = parser.parse_args()
+    
+    prompts = [
+        '你有什么特长？',
+        '为什么天空是蓝色的',
+        '请用Python写一个计算斐波那契数列的函数',
+        '解释一下"光合作用"的基本过程',
+        '如果明天下雨，我应该如何出门',
+        '比较一下猫和狗作为宠物的优缺点',
+        '解释什么是机器学习',
+        '推荐一些中国的美食'
+    ]
+    
+    conversation = []
+    model, tokenizer = init_model(args)
+    input_mode = int(input('[0] 自动测试\n[1] 手动输入\n'))
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
+    prompt_iter = prompts if input_mode == 0 else iter(lambda: input('💬: '), '')
+    for prompt in prompt_iter:
+        setup_seed(random.randint(0, 31415926))
+        if input_mode == 0: print(f'💬: {prompt}')
+        conversation = conversation[-args.historys:] if args.historys else []
+        conversation.append({"role": "user", "content": prompt})
+        if 'pretrain' in args.weight:
+            inputs = tokenizer.bos_token + prompt
+        else:
+            inputs = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True, open_thinking=bool(args.open_thinking))
+        
+        inputs = tokenizer(inputs, return_tensors="pt", truncation=True).to(args.device)
+
+        print('🧠: ', end='')
+        st = time.time()
+        generated_ids = model.generate(
+            inputs=inputs["input_ids"], attention_mask=inputs["attention_mask"],
+            max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
+            pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
+            top_p=args.top_p, temperature=args.temperature, repetition_penalty=1
+        )
+        response = tokenizer.decode(generated_ids[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
+        conversation.append({"role": "assistant", "content": response})
+        gen_tokens = len(generated_ids[0]) - len(inputs["input_ids"][0])
+        print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n')
+
+if __name__ == "__main__":
+    main()

examples/eval_omni.py

@@ -0,0 +1,244 @@
+import argparse
+import os
+import random
+import time
+import warnings
+import torch
+from PIL import Image
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from omni.models import MiniMindOmni, OmniConfig
+from omni.datasets.lm_dataset import OmniDataset
+from omni.utils import setup_seed, log_model_params
+warnings.filterwarnings('ignore')
+
+
+def init_model(args):
+    tokenizer = AutoTokenizer.from_pretrained(args.load_from)
+    if 'model' in args.load_from:
+        moe_suffix = '_moe' if args.use_moe else ''
+        ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
+        model = MiniMindOmni(
+            OmniConfig(
+                hidden_size=args.hidden_size,
+                num_hidden_layers=args.num_hidden_layers,
+                use_moe=bool(args.use_moe)
+            ),
+            audio_encoder_path="./model/SenseVoiceSmall",
+            vision_model_path="./model/siglip2-base-p32-256-ve"
+        )
+        model.load_state_dict(torch.load(ckp, map_location=args.device), strict=False)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
+        model.audio_encoder, model.audio_processor = MiniMindOmni.load_sensevoice("./model/SenseVoiceSmall")
+        model.vision_encoder, model.vision_processor = MiniMindOmni.load_vision("./model/siglip2-base-p32-256-ve")
+    log_model_params(model)
+    if model.audio_encoder is not None: model.audio_encoder.to(args.device)
+    if model.vision_encoder is not None: model.vision_encoder.to(args.device)
+    from transformers import MimiModel
+    model.mimi_model = MimiModel.from_pretrained("./model/mimi").eval()
+    return model.half().eval().to(args.device), tokenizer
+
+
+def eval_sample(model, tokenizer, args, idx, prompt, audio_inputs, output_name, pixel_values=None, history=None, audio_lens=None, ref_codes=None, spk_emb=None):
+    import soundfile as sf
+    from pydub import AudioSegment
+    messages = (history or []) + [{"role": "user", "content": prompt}]
+    inputs_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, open_thinking=bool(args.open_thinking))
+    x = torch.tensor(tokenizer(inputs_text).data['input_ids'], dtype=torch.long, device=args.device)[None, ...]
+
+    audio_frames = []
+    with torch.no_grad():
+        res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=args.max_new_tokens,
+                               temperature=args.temperature, top_p=args.top_p, stream=True,
+                               return_audio_codes=True, open_thinking=bool(args.open_thinking),
+                               audio_inputs=audio_inputs, audio_lens=audio_lens, pixel_values=pixel_values,
+                               ref_codes=ref_codes, spk_emb=spk_emb)
+        print('📒 [Thinker]: ', end='', flush=True)
+        history_idx = 0
+        for y, audio_frame in res_y:
+            if y is not None:
+                answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
+                if answer and answer[-1] != '�':
+                    print(answer[history_idx:], end='', flush=True)
+                    history_idx = len(answer)
+            if audio_frame:
+                audio_frames.append(audio_frame)
+        print()
+
+        if audio_frames:
+            print(f'🎹 [Talker]: {len(audio_frames)} frames', end=" ")
+            if args.decode_audio:
+                try:
+                    codes = [f for f in audio_frames if f and len(f) == 8]
+                    if not codes:
+                        print('⚠️  生成的Mimi codes为空，跳过保存。')
+                        return
+                    mimi_codes = torch.tensor(codes, dtype=torch.long).T.unsqueeze(0).to(args.device)
+                    filtered = torch.where(mimi_codes >= 2049, torch.zeros_like(mimi_codes), mimi_codes)
+                    audio = model.mimi_model.decode(filtered).audio_values
+                    output_path = os.path.join(args.output_dir, output_name)
+                    wav_path = output_path.rsplit('.', 1)[0] + '.wav'
+                    sf.write(wav_path, audio.squeeze().float().cpu().numpy(), 24000)
+                    AudioSegment.from_wav(wav_path).export(output_path, format='mp3', bitrate='64k')
+                    os.remove(wav_path)
+                    print(f'| Audio decoded to: {output_path}')
+                except Exception as e:
+                    print(f'⚠️  保存音频失败: {str(e)}')
+            else:
+                print("(decode_audio=off)\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MiniMind-O Chat")
+    parser.add_argument('--load_from', default='model', type=str, help="模型加载路径（model=原生torch权重）")
+    parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
+    parser.add_argument('--weight', default='sft_omni', type=str, help="权重名称前缀")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构")
+    parser.add_argument('--max_new_tokens', default=512, type=int, help="最大生成长度")
+    parser.add_argument('--temperature', default=0.7, type=float, help="Thinker生成温度")
+    parser.add_argument('--top_p', default=0.85, type=float, help="nucleus采样阈值")
+    parser.add_argument('--output_dir', default='./output_audio/', type=str, help="输出音频保存目录")
+    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
+    parser.add_argument('--audio_dir', default='./dataset/eval_omni/', type=str, help="测试音频目录")
+    parser.add_argument('--image_dir', default='./dataset/eval_omni/', type=str, help="测试图像目录")
+    parser.add_argument('--open_thinking', default=0, type=int, help="是否开启思考模式（0=否，1=是）（思考模式下禁用audio输出）")
+    parser.add_argument('--decode_audio', default=1, type=int, help="是否解码音频输出（0=否，1=是）")
+    parser.add_argument('--mode', default='0', type=str, help="评估模式：-1=all 0=text 1=multi 2=audio 3=clone 4=image 5=mix（逗号组合，如 2,5）")
+    parser.add_argument('--prompt_lang', default=0, type=int, choices=[0, 1, 2], help="问题语言：0=英文 1=中文 2=英文+中文")
+    args = parser.parse_args()
+    modes = set(args.mode.replace(',', '').replace('-1', '012345'))
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    model, tokenizer = init_model(args)
+    setup_seed(int(time.time()) % 31415926)
+
+    if '0' in modes:
+        print('\n\n==================== text -> {text, audio} ====================')
+        test_prompts_en = [
+            "Tell me an interesting fact about space.", "How do I make a cup of coffee?", "What's the weather like today?",
+            "Will it rain tomorrow?", "Tell me a joke.", "Can you sing a song for me?", "Please introduce yourself."
+        ]
+        test_prompts_zh = [
+            "告诉我一个关于太空的有趣事实。", "如何制作一杯咖啡？", "今天的天气怎么样？",
+            "明天会下雨吗？", "给我讲个笑话吧", "你能为我唱首歌吗？", "介绍一下你自己"
+        ]
+        test_prompts = [test_prompts_en, test_prompts_zh, test_prompts_en + test_prompts_zh][args.prompt_lang]
+        for idx, prompt in enumerate(test_prompts):
+            print(f'\n📝 [text-{idx+1}]: {prompt}')
+            eval_sample(model, tokenizer, args, idx, prompt, None, f"text-{idx:02d}.mp3")
+
+    if '1' in modes:
+        print('\n\n==================== multi-turn -> {text, audio} ====================')
+        multi_turn_tests_zh = [
+            {
+                "history": [
+                    {"role": "user", "content": "你好"},
+                    {"role": "assistant", "content": "你好！有什么可以帮你的吗？"}
+                ],
+                "prompt": "我想找点事做，你有什么建议吗？"
+            },
+            {
+                "history": [
+                    {"role": "user", "content": "你好"},
+                    {"role": "assistant", "content": "你好！有什么可以帮你的吗？"},
+                    {"role": "user", "content": "我想找点事做，你有什么建议吗？"},
+                    {"role": "assistant", "content": "可以听听音乐或者看看书，放松一下心情。"}
+                ],
+                "prompt": "好的，那我去照做了，谢谢你"
+            }
+        ]
+        multi_turn_tests_en = [
+            {
+                "history": [
+                    {"role": "user", "content": "Hello"},
+                    {"role": "assistant", "content": "Hello! How can I help you?"}
+                ],
+                "prompt": "I want to find something to do. Do you have any suggestions?"
+            },
+            {
+                "history": [
+                    {"role": "user", "content": "Hello"},
+                    {"role": "assistant", "content": "Hello! How can I help you?"},
+                    {"role": "user", "content": "I want to find something to do. Do you have any suggestions?"},
+                    {"role": "assistant", "content": "You can listen to music or read a book to relax a little."}
+                ],
+                "prompt": "Okay, I will try that. Thank you."
+            }
+        ]
+        multi_turn_tests = [multi_turn_tests_en, multi_turn_tests_zh, multi_turn_tests_en + multi_turn_tests_zh][args.prompt_lang]
+        for idx, test in enumerate(multi_turn_tests):
+            print(f'\n💬 [multi-{idx+1}]')
+            for msg in test["history"]: print(f'   {msg["role"]}: {msg["content"]}')
+            print(f'   user: {test["prompt"]}')
+            eval_sample(model, tokenizer, args, idx, test["prompt"], None, f"multi-{idx:02d}.mp3", history=test["history"])
+
+    if '2' in modes:
+        print('\n\n==================== audio -> {text, audio} ====================')
+        audio_files_en = sorted([f for f in os.listdir(args.audio_dir) if f.startswith('audio-en-') and f.lower().endswith(('.mp3', '.wav'))])
+        audio_files_zh = sorted([f for f in os.listdir(args.audio_dir) if f.startswith('audio-zh-') and f.lower().endswith(('.mp3', '.wav'))])
+        audio_files = [audio_files_en, audio_files_zh, audio_files_en + audio_files_zh][args.prompt_lang]
+        for idx, audio_file in enumerate(audio_files):
+            print(f'\n🎤 [audio-{idx+1}]: {audio_file}')
+            mel, valid_len = OmniDataset.process_audio(os.path.join(args.audio_dir, audio_file), model.audio_processor)
+            audio_inputs = mel.unsqueeze(0).to(args.device)
+            audio_lens = torch.tensor([valid_len], device=args.device)
+            audio_token_len = valid_len or 1
+            prompt = model.config.audio_special_token * audio_token_len
+            eval_sample(model, tokenizer, args, idx, prompt, audio_inputs, f"audio-{idx:02d}-{os.path.splitext(audio_file)[0]}.mp3", audio_lens=audio_lens)
+
+    if '3' in modes:
+        print('\n\n==================== clone voice -> {text, audio} ====================')
+        clone_prompts_en = ["Hello, please introduce yourself.", "What's the weather like today?", "Tell me a joke."]
+        clone_prompts_zh = ["你好，请介绍一下你自己。", "今天天气怎么样？", "给我讲个笑话吧"]
+        clone_prompts = [clone_prompts_en, clone_prompts_zh, clone_prompts_en + clone_prompts_zh][args.prompt_lang]
+        voices_pt = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'model', 'speaker', 'voices_unseen.pt')
+        voices = [('default', None, None)]
+        if os.path.exists(voices_pt):
+            voice_data = torch.load(voices_pt, map_location=args.device)
+            for speaker, v in sorted(voice_data.items()):
+                rc = v['ref_codes'].unsqueeze(0).to(args.device)
+                se = v['spk_emb'].half().unsqueeze(0).to(args.device) if 'spk_emb' in v else None
+                voices.append((speaker, rc, se))
+        for speaker, rc, se in voices:
+            info = f'ref_codes: {rc.shape[2]} frames, spk_emb: {"+" if se is not None else "-"}' if rc is not None else ('spk_emb only' if se is not None else 'default')
+            print(f'\n🎵 [clone: {speaker}] {info}')
+            for idx, prompt in enumerate(clone_prompts):
+                print(f'  📝 [text-{idx+1}]: {prompt}')
+                history = [{"role": "system", "content": "你是一个专业的语音助手，请用给定的音色风格来回答用户的问题。请尽量详细地回答，给出有价值的信息。"}]
+                eval_sample(model, tokenizer, args, idx, prompt, None, f"clone-{speaker}-{idx:02d}.mp3", ref_codes=rc, history=history, spk_emb=se)
+
+    if '4' in modes:
+        print('\n\n==================== image -> {text, audio} ====================')
+        image_files = sorted([f for f in os.listdir(args.image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
+        for idx, image_file in enumerate(image_files):
+            print(f'\n🖼️ [image-{idx+1}]: {image_file}')
+            image = Image.open(os.path.join(args.image_dir, image_file)).convert('RGB')
+            pixel_values = {k: v.to(args.device) for k, v in model.vision_processor(images=image, return_tensors="pt").items()}
+            prompts = [["Please describe this image."], ["请描述这张图片"], ["Please describe this image.", "请描述这张图片"]][args.prompt_lang]
+            for lang_idx, prompt_text in enumerate(prompts):
+                prompt = prompt_text + "\n\n" + model.config.image_special_token * model.config.image_token_len
+                eval_sample(model, tokenizer, args, idx, prompt, None, f"image-{idx:02d}-{lang_idx}-{os.path.splitext(image_file)[0]}.mp3", pixel_values=pixel_values)
+
+    if '5' in modes:
+        print('\n\n==================== text+audio+image -> {text, audio} ====================')
+        img_audio_files = sorted([f for f in os.listdir(args.audio_dir) if f.startswith('img-') and f.lower().endswith(('.mp3', '.wav'))])
+        image_files = sorted([f for f in os.listdir(args.image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
+        text_hints = [["Please answer me: "], ["请回答我："], ["Please answer me: ", "请回答我："]][args.prompt_lang]
+        for idx, image_file in enumerate(image_files):
+            audio_file = random.choice(img_audio_files)
+            image = Image.open(os.path.join(args.image_dir, image_file)).convert('RGB')
+            pixel_values = {k: v.to(args.device) for k, v in model.vision_processor(images=image, return_tensors="pt").items()}
+            for lang_idx, text_hint in enumerate(text_hints):
+                print(f'\n🌀 [mix-{idx+1}-{lang_idx}]: {text_hint} | {audio_file} | {image_file}')
+                mel, valid_len = OmniDataset.process_audio(os.path.join(args.audio_dir, audio_file), model.audio_processor)
+                audio_inputs = mel.unsqueeze(0).to(args.device)
+                audio_lens = torch.tensor([valid_len], device=args.device)
+                audio_token_len = valid_len or 1
+                prompt = text_hint + model.config.audio_special_token * audio_token_len + "\n\n" + model.config.image_special_token * model.config.image_token_len
+                eval_sample(model, tokenizer, args, idx, prompt, audio_inputs, f"mix-{idx:02d}-{lang_idx}-{os.path.splitext(image_file)[0]}.mp3", pixel_values=pixel_values, audio_lens=audio_lens)
+
+
+if __name__ == "__main__":
+    main()

examples/eval_toolcall.py

@@ -0,0 +1,238 @@
+import os
+import re
+import json
+import time
+import random
+import argparse
+import warnings
+import torch
+from datetime import datetime
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from openai import OpenAI
+from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.utils.training import setup_seed, get_model_params
+warnings.filterwarnings('ignore')
+
+TOOLS = [
+    {"type": "function", "function": {"name": "calculate_math", "description": "计算数学表达式的结果，支持加减乘除、幂运算、开方等", "parameters": {"type": "object", "properties": {"expression": {"type": "string", "description": "数学表达式，如123+456、2**10、sqrt(144)"}}, "required": ["expression"]}}},
+    {"type": "function", "function": {"name": "get_current_time", "description": "获取当前日期和时间，支持指定时区", "parameters": {"type": "object", "properties": {"timezone": {"type": "string", "description": "时区名称，如Asia/Shanghai、America/New_York", "default": "Asia/Shanghai"}}, "required": []}}},
+    {"type": "function", "function": {"name": "random_number", "description": "生成指定范围内的随机数", "parameters": {"type": "object", "properties": {"min": {"type": "integer", "description": "最小值", "default": 0}, "max": {"type": "integer", "description": "最大值", "default": 100}}, "required": []}}},
+    {"type": "function", "function": {"name": "text_length", "description": "计算文本的字符数和单词数", "parameters": {"type": "object", "properties": {"text": {"type": "string", "description": "要统计的文本"}}, "required": ["text"]}}},
+    {"type": "function", "function": {"name": "unit_converter", "description": "进行单位换算，支持长度、重量、温度等", "parameters": {"type": "object", "properties": {"value": {"type": "number", "description": "要转换的数值"}, "from_unit": {"type": "string", "description": "源单位，如km、miles、kg、pounds、celsius、fahrenheit"}, "to_unit": {"type": "string", "description": "目标单位"}}, "required": ["value", "from_unit", "to_unit"]}}},
+    {"type": "function", "function": {"name": "get_current_weather", "description": "获取指定城市的当前天气信息，包括温度、湿度和天气状况", "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "城市名称，如北京、上海、New York"}, "unit": {"type": "string", "description": "温度单位，celsius或fahrenheit", "enum": ["celsius", "fahrenheit"], "default": "celsius"}}, "required": ["location"]}}},
+    {"type": "function", "function": {"name": "get_exchange_rate", "description": "查询两种货币之间的实时汇率", "parameters": {"type": "object", "properties": {"from_currency": {"type": "string", "description": "源货币代码，如USD、CNY、EUR"}, "to_currency": {"type": "string", "description": "目标货币代码，如USD、CNY、EUR"}}, "required": ["from_currency", "to_currency"]}}},
+    {"type": "function", "function": {"name": "translate_text", "description": "将文本翻译成目标语言", "parameters": {"type": "object", "properties": {"text": {"type": "string", "description": "要翻译的文本"}, "target_language": {"type": "string", "description": "目标语言，如english、chinese、japanese、french"}}, "required": ["text", "target_language"]}}},
+]
+
+MOCK_RESULTS = {
+    "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("²", "**2").replace("³", "**3").replace("（", "(").replace("）", ")")))},
+    "get_current_time": lambda args: {"datetime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "timezone": args.get("timezone", "Asia/Shanghai")},
+    "random_number": lambda args: {"result": random.randint(int(args.get("min", 0)), int(args.get("max", 100)))},
+    "text_length": lambda args: {"characters": len(args.get("text", "")), "words": len(args.get("text", "").split())},
+    "unit_converter": lambda args: {"result": round(float(args.get("value", 0)) * 0.621371, 2), "from": f"{args.get('value', 0)} {args.get('from_unit', '')}", "to": args.get("to_unit", "")},
+    "get_current_weather": lambda args: {"city": args.get("location"), "temperature": "22°C", "humidity": "65%", "condition": "晴"},
+    "get_exchange_rate": lambda args: {"from": args.get("from_currency", ""), "to": args.get("to_currency", ""), "rate": 7.15},
+    "translate_text": lambda args: {"translated": "hello world"},
+}
+
+TOOL_MAP = {t["function"]["name"]: t for t in TOOLS}
+
+def get_tools(names):
+    return [TOOL_MAP[n] for n in names]
+
+TEST_CASES = [
+    {"prompt": "帮我算一下 256 乘以 37 等于多少", "tools": ["calculate_math", "get_current_time"]},
+    {"prompt": "现在几点了？", "tools": ["get_current_time", "random_number"]},
+    {"prompt": "帮我把100公里换算成英里", "tools": ["unit_converter", "calculate_math"]},
+    {"prompt": "帮我生成一个1到1000的随机数，然后计算它的平方", "tools": ["random_number", "calculate_math", "text_length"]},
+    {"prompt": "北京今天天气怎么样？", "tools": ["get_current_weather", "get_current_time"]},
+    {"prompt": "查一下美元兑人民币汇率", "tools": ["get_exchange_rate", "get_current_time"]},
+    {"prompt": "把'你好世界'翻译成英文", "tools": ["translate_text", "text_length"]},
+    {"prompt": "What is the weather in Tokyo? Also convert 30 celsius to fahrenheit.", "tools": ["get_current_weather", "unit_converter", "get_current_time"]},
+]
+
+
+def init_model(args):
+    tokenizer = AutoTokenizer.from_pretrained(args.load_from)
+    if 'model' in args.load_from:
+        model = MiniMindForCausalLM(MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe)))
+        moe_suffix = '_moe' if args.use_moe else ''
+        ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
+        model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
+    get_model_params(model, model.config)
+    return model.half().eval().to(args.device), tokenizer
+
+
+def parse_tool_calls(text):
+    matches = re.findall(r'<tool_call>(.*?)</tool_call>', text, re.DOTALL)
+    calls = []
+    for m in matches:
+        try:
+            calls.append(json.loads(m.strip()))
+        except Exception:
+            pass
+    return calls
+
+
+def parse_tool_call_from_text(content):
+    pattern = r'<tool_call>\s*(\{.*?\})\s*</tool_call>'
+    matches = re.findall(pattern, content, re.DOTALL)
+    if not matches:
+        return None
+    tool_calls = []
+    for i, match in enumerate(matches):
+        try:
+            data = json.loads(match)
+            tool_calls.append({
+                "id": f"call_{i}",
+                "function": {"name": data.get("name", ""), "arguments": json.dumps(data.get("arguments", {}), ensure_ascii=False)}
+            })
+        except Exception:
+            pass
+    return tool_calls if tool_calls else None
+
+
+def execute_tool(call, arguments=None):
+    name = call.get("name", "") if isinstance(call, dict) else call
+    try:
+        raw_args = call.get("arguments", {}) if isinstance(call, dict) else arguments
+        args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
+    except Exception:
+        args = {}
+    fn = MOCK_RESULTS.get(name)
+    if not fn:
+        return {"error": f"未知工具: {name}"}
+    try:
+        return fn(args)
+    except Exception as e:
+        return {"error": f"工具执行失败: {str(e)[:80]}"}
+
+
+def generate(model, tokenizer, messages, tools, args):
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools, open_thinking=False)
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True).to(args.device)
+    st = time.time()
+    print('🧠: ', end='')
+    generated_ids = model.generate(
+        inputs["input_ids"], attention_mask=inputs["attention_mask"],
+        max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
+        pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
+        top_p=args.top_p, temperature=args.temperature
+    )
+    response = tokenizer.decode(generated_ids[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
+    gen_tokens = len(generated_ids[0]) - len(inputs["input_ids"][0])
+    print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s') if args.show_speed else print()
+    return response
+
+
+def chat_api(client, messages, tools, args, stream=True):
+    response = client.chat.completions.create(
+        model=args.api_model, messages=messages, tools=tools,
+        stream=stream, temperature=args.temperature,
+        max_tokens=8192, top_p=args.top_p
+    )
+    if not stream:
+        choice = response.choices[0]
+        content = choice.message.content or ""
+        tool_calls = choice.message.tool_calls
+        if not tool_calls:
+            tool_calls = parse_tool_call_from_text(content)
+        print(f'🧠: {content}')
+        return content, tool_calls
+    print('🧠: ', end='', flush=True)
+    content, tool_calls = "", None
+    for chunk in response:
+        delta = chunk.choices[0].delta
+        if delta.content:
+            print(delta.content, end="", flush=True)
+            content += delta.content
+        if delta.tool_calls:
+            if tool_calls is None:
+                tool_calls = []
+            for tc_chunk in delta.tool_calls:
+                idx = tc_chunk.index if tc_chunk.index is not None else len(tool_calls)
+                while len(tool_calls) <= idx:
+                    tool_calls.append({
+                        "id": "",
+                        "function": {"name": "", "arguments": ""}
+                    })
+                if tc_chunk.id:
+                    tool_calls[idx]["id"] += tc_chunk.id
+                if tc_chunk.function:
+                    if tc_chunk.function.name:
+                        tool_calls[idx]["function"]["name"] += tc_chunk.function.name
+                    if tc_chunk.function.arguments:
+                        tool_calls[idx]["function"]["arguments"] += tc_chunk.function.arguments
+    print()
+    if not tool_calls:
+        tool_calls = parse_tool_call_from_text(content)
+    return content, tool_calls
+
+
+def run_case(prompt, tools, args, model=None, tokenizer=None, client=None):
+    messages = [{"role": "user", "content": prompt}]
+    while True:
+        if args.backend == 'local':
+            content = generate(model, tokenizer, messages, tools, args)
+            tool_calls = parse_tool_calls(content)
+        else:
+            content, tool_calls = chat_api(client, messages, tools, args, stream=bool(args.stream))
+        if not tool_calls:
+            break
+        tool_calls = [{
+            "id": tc.id if hasattr(tc, 'id') else tc.get("id", ""),
+            "name": tc.function.name if hasattr(tc, 'function') else tc["function"]["name"],
+            "arguments": tc.function.arguments if hasattr(tc, 'function') else tc["function"]["arguments"]
+        } for tc in tool_calls] if args.backend == 'api' else tool_calls
+        messages.append({"role": "assistant", "content": content} if args.backend == 'local' else {"role": "assistant", "content": content, "tool_calls": [{"id": tc["id"], "type": "function", "function": {"name": tc["name"], "arguments": tc["arguments"]}} for tc in tool_calls]})
+        for tc in tool_calls:
+            name = tc["name"]
+            arguments = tc["arguments"]
+            print(f'📞 [Tool Calling]: {name} | args={arguments}')
+            result = execute_tool(tc if args.backend == 'local' else name, arguments)
+            print(f'✅ [Tool Called]: {json.dumps(result, ensure_ascii=False)}')
+            messages.append({"role": "tool", "content": json.dumps(result, ensure_ascii=False)} if args.backend == 'local' else {"role": "tool", "content": json.dumps(result, ensure_ascii=False), "tool_call_id": tc["id"]})
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MiniMind ToolCall评测")
+    parser.add_argument('--backend', default='local', choices=['local', 'api'], type=str, help="推理后端（local=本地模型，api=OpenAI兼容接口）")
+    parser.add_argument('--load_from', default='../model', type=str, help="模型加载路径（model=原生torch权重，其他路径=transformers格式）")
+    parser.add_argument('--save_dir', default='../out', type=str, help="模型权重目录")
+    parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀（pretrain, full_sft, rlhf, reason, ppo_actor, grpo, spo）")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--max_new_tokens', default=512, type=int, help="最大生成长度")
+    parser.add_argument('--temperature', default=0.9, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
+    parser.add_argument('--top_p', default=0.9, type=float, help="nucleus采样阈值（0-1）")
+    parser.add_argument('--show_speed', default=0, type=int, help="显示decode速度（tokens/s）")
+    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
+    parser.add_argument('--api_base_url', default="http://localhost:11434/v1", type=str, help="OpenAI兼容接口的base_url")
+    parser.add_argument('--api_key', default='sk-123', type=str, help="OpenAI兼容接口的api_key")
+    parser.add_argument('--api_model', default='jingyaogong/minimind-3:latest', type=str, help="API请求时使用的模型名称")
+    parser.add_argument('--stream', default=1, type=int, help="API模式下是否流式输出（0=否，1=是）")
+    args = parser.parse_args()
+
+    model = tokenizer = client = None
+    if args.backend == 'local': model, tokenizer = init_model(args)
+    else: client = OpenAI(api_key=args.api_key, base_url=args.api_base_url)
+
+    input_mode = int(input('[0] 自动测试\n[1] 手动输入\n'))
+
+    cases = [{"prompt": case["prompt"], "tools": get_tools(case["tools"]), "tool_names": case["tools"]} for case in TEST_CASES] if input_mode == 0 else iter(lambda: {"prompt": input('💬: '), "tools": TOOLS, "tool_names": [t["function"]["name"] for t in TOOLS]}, {"prompt": "", "tools": TOOLS, "tool_names": []})
+    for case in cases:
+        if not case["prompt"]: break
+        setup_seed(random.randint(0, 31415926))
+        if input_mode == 0:
+            print(f'📦 可用工具: {case["tool_names"]}\n')
+            print(f'💬: {case["prompt"]}')
+        run_case(case["prompt"], case["tools"], args, model=model, tokenizer=tokenizer, client=client)
+        print('\n' + '-' * 50 + '\n')
+
+
+if __name__ == "__main__":
+    main()

examples/eval_vlm.py

@@ -0,0 +1,79 @@
+import time
+import argparse
+import os
+import warnings
+import torch
+import random
+from PIL import Image
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from omni.models import MiniMindVLM, VLMConfig
+from omni.utils import setup_seed, get_vlm_model_params
+warnings.filterwarnings('ignore')
+
+def init_model(args):
+    tokenizer = AutoTokenizer.from_pretrained(args.load_from, trust_remote_code=True)
+    if 'model' in args.load_from:
+        moe_suffix = '_moe' if args.use_moe else ''
+        ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
+        model = MiniMindVLM(
+            VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe)),
+            vision_model_path="./model/siglip2-base-p32-256-ve"
+        )
+        state_dict = torch.load(ckp, map_location=args.device)
+        model.load_state_dict({k: v for k, v in state_dict.items() if 'mask' not in k}, strict=False)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
+        model.vision_encoder, model.processor = MiniMindVLM.get_vision_model("./model/siglip2-base-p32-256-ve")
+    get_vlm_model_params(model, model.config)
+    model = model.eval()
+    if "cuda" in args.device: model = model.half()
+    return model.to(args.device), tokenizer, model.processor
+
+
+def main():
+    parser = argparse.ArgumentParser(description="MiniMind-V Chat")
+    parser.add_argument('--load_from', default='model', type=str, help="模型加载路径（model=原生torch权重，其他路径=transformers格式）")
+    parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
+    parser.add_argument('--weight', default='sft_vlm', type=str, help="权重名称前缀（pretrain_vlm, sft_vlm）")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--max_new_tokens', default=512, type=int, help="最大生成长度")
+    parser.add_argument('--temperature', default=0.7, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
+    parser.add_argument('--top_p', default=0.85, type=float, help="nucleus采样阈值（0-1）")
+    parser.add_argument('--image_dir', default='./dataset/eval_images/', type=str, help="测试图像目录")
+    parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）")
+    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
+    parser.add_argument('--open_thinking', default=0, type=int, help="是否开启自适应思考（0=否，1=是）")
+    args = parser.parse_args()
+
+    model, tokenizer, preprocess = init_model(args)
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    # 自动测试image_dir中的所有图像
+    prompt = "<image>\n请描述这张图中的主要物体和场景。"
+    for image_file in sorted(os.listdir(args.image_dir)):
+        if image_file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
+            setup_seed(random.randint(1, 31415926))
+            image_path = os.path.join(args.image_dir, image_file)
+            image = Image.open(image_path).convert('RGB')
+            pixel_values = {k: v.to(args.device) for k, v in MiniMindVLM.image2tensor(image, preprocess).items()}
+
+            messages = [{"role": "user", "content": prompt.replace('<image>', model.config.image_special_token * model.config.image_token_len)}]
+            inputs_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, open_thinking=bool(args.open_thinking))
+            inputs = tokenizer(inputs_text, return_tensors="pt", truncation=True).to(args.device)
+
+            print(f'[图像]: {image_file}')
+            print(f"💬: {repr(prompt)}")
+            print('🤖: ', end='')
+            st = time.time()
+            generated_ids = model.generate(
+                inputs=inputs["input_ids"], attention_mask=inputs["attention_mask"],
+                max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer,
+                pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id,
+                top_p=args.top_p, temperature=args.temperature, pixel_values=pixel_values
+            )
+            gen_tokens = len(generated_ids[0]) - len(inputs["input_ids"][0])
+            print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n')
+
+if __name__ == "__main__":
+    main()

examples/omni_web_demo.py

@@ -0,0 +1,512 @@
+import argparse, os, sys, json, time, math, torch, threading, queue, base64, io, logging, contextlib
+import numpy as np
+from PIL import Image
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from omni.models import MiniMindOmni
+from omni.serve.realtime import RealtimeSession
+from omni.utils import log_model_params
+logging.getLogger().setLevel(logging.ERROR)
+
+M = {}  # model / tokenizer / device / mimi / asr / cfg
+V = {}  # voice_name -> {ref_codes, spk_emb}
+V_builtin, V_unseen, V_manual = [], [], []
+MODEL_LOCK = threading.Lock()
+SAMPLES_PER_FRAME = 1920
+REF_FRAMES = 300
+CLONE_VOICE = 'voice_clone'
+CLONE_FILE = 'voice_clone.pt'
+
+# -------- helpers --------
+def sse(d): return f"data: {json.dumps(d)}\n\n"
+
+def scan_hf_models(base_dir):
+    models = {}
+    base_dir = os.path.abspath(base_dir)
+    for d in sorted(os.listdir(base_dir), reverse=True):
+        full_path = os.path.join(base_dir, d)
+        if not os.path.isdir(full_path) or d.startswith('.') or d.startswith('_'):
+            continue
+        files = set(os.listdir(full_path))
+        has_model = bool(files & {'pytorch_model.bin', 'model.safetensors', 'pytorch_model.bin.index.json', 'model.safetensors.index.json'})
+        if has_model:
+            models[d] = full_path
+    return models
+
+def asr_run(samples):
+    from funasr.utils.postprocess_utils import rich_transcription_postprocess
+    r = M['asr'].generate(input=samples, cache={}, language='auto', use_itn=True)
+    return rich_transcription_postprocess(r[0]['text']).strip() if r else ''
+
+def prep_audio(samples):
+    m = M['model']
+    proc = m.audio_processor(samples, sampling_rate=16000, return_tensors="pt", return_attention_mask=True)
+    mel = proc.input_features.squeeze(0).unsqueeze(0).to(M['device'])
+    vlen = proc.attention_mask.sum().item()
+    prompt = m.config.audio_special_token * (vlen or 1)
+    return mel, torch.tensor([vlen], device=M['device']), prompt
+
+def prep_image(b64):
+    img = Image.open(io.BytesIO(base64.b64decode(b64))).convert('RGB')
+    return {k: v.to(M['device']) for k, v in M['model'].vision_processor(images=img, return_tensors="pt").items()}
+
+def build_ids(prompt, history):
+    tok, dev, n = M['tokenizer'], M['device'], M['cfg'].max_history_turns
+    hist = history[-n:] if n > 0 else []
+    msgs = hist + [{"role": "user", "content": prompt}]
+    t = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+    return torch.tensor(tok(t).data['input_ids'], dtype=torch.long, device=dev)[None, ...]
+
+def _mimi_decode(frames):
+    codes = [f for f in frames if f and len(f) == 8]
+    if not codes or not M['mimi']: return None
+    mc = torch.tensor(codes, dtype=torch.long, device=M['device']).T.unsqueeze(0)
+    mc = torch.where(mc >= 2049, torch.zeros_like(mc), mc)
+    with torch.no_grad():
+        au = M['mimi'].decode(mc).audio_values.squeeze().cpu().numpy()
+    return au, mc.shape[-1]
+
+def pcm_bytes(frames, ov):
+    r = _mimi_decode(frames)
+    if r is None: return None
+    au, T = r
+    if ov > 0: au = au[int(ov * len(au) / T):]
+    return (au * 32767).astype('int16').tobytes()
+
+def stream_pcm(frames, flush=False):
+    """yield (pcm_bytes,) on chunk boundaries or on final flush."""
+    if not M['mimi']: return
+    cf, ov_max, n = M['cfg'].audio_chunk_frames, M['cfg'].audio_overlap, len(frames)
+    if not flush and n >= cf and n % cf == 0:
+        ov = min(ov_max, n - cf)
+        p = pcm_bytes(frames[-(cf + ov):], ov)
+        if p: yield p
+    elif flush:
+        rem = n % cf
+        if rem:
+            ov = min(ov_max, n - rem)
+            p = pcm_bytes(frames[-(rem + ov):], ov)
+            if p: yield p
+
+def voice_args(name):
+    if name and name != 'default' and name in V:
+        v = V[name]
+        dev = M['device']
+        rc = v['ref_codes'].unsqueeze(0).to(dev)
+        se = v['spk_emb'].half().unsqueeze(0).to(dev) if 'spk_emb' in v else None
+        return {'ref_codes': rc, 'spk_emb': se}
+    return {}
+
+def register_voice(name, value, group='manual'):
+    V[name] = value
+    groups = {'builtin': V_builtin, 'unseen': V_unseen, 'manual': V_manual}
+    dst = groups[group]
+    if name not in dst:
+        dst.append(name)
+    for k, lst in groups.items():
+        if k != group and name in lst:
+            lst.remove(name)
+
+def clone_voice_path():
+    return os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'model', 'speaker', CLONE_FILE)
+
+def delete_manual_voice(name):
+    if name not in V_manual:
+        raise RuntimeError('只能删除手动克隆的音色')
+    out_path = clone_voice_path()
+    saved = torch.load(out_path, map_location='cpu') if os.path.exists(out_path) else {}
+    if name in saved:
+        saved.pop(name)
+        torch.save(saved, out_path)
+    V.pop(name, None)
+    if name in V_manual:
+        V_manual.remove(name)
+
+def normalize_voice_name(name):
+    name = ' '.join(str(name or '').split())
+    if not name:
+        name = CLONE_VOICE
+    if len(name) > 24:
+        raise RuntimeError('音色名太长，建议控制在 24 个字以内')
+    if name.lower() == 'default':
+        raise RuntimeError('default 是保留名称，请换一个')
+    if name in V_builtin or name in V_unseen:
+        raise RuntimeError('该名称已被现有音色占用，请换一个')
+    return name
+
+def validate_clone_audio(w16):
+    if w16.numel() < int(16000 * 1.8):
+        raise RuntimeError('录音太短，请把整句话读完')
+    peak = w16.abs().max().item()
+    frame, hop = 800, 400
+    if w16.numel() >= frame:
+        rms = w16.unfold(0, frame, hop).pow(2).mean(dim=1).sqrt().cpu().numpy()
+    else:
+        rms = np.array([w16.pow(2).mean().sqrt().item()])
+    hi = float(np.quantile(rms, 0.95))
+    lo = float(np.quantile(rms, 0.2))
+    if hi < 0.008:
+        raise RuntimeError('录音太轻，请靠近麦克风一点')
+    if hi > 0 and lo / hi > 0.45:
+        raise RuntimeError('环境噪声太大，请换安静一点的环境')
+    if peak > 0.995:
+        raise RuntimeError('录音有爆音，请离麦克风远一点')
+
+def build_clone_voice(audio_b64):
+    if M.get('mimi') is None or M.get('campplus') is None or M.get('mel_fn') is None:
+        raise RuntimeError('Mimi 或 CAM++ 未加载')
+    from pydub import AudioSegment
+    seg = AudioSegment.from_file(io.BytesIO(base64.b64decode(audio_b64))).set_channels(1).set_sample_width(2)
+    if len(seg) < 1000:
+        raise RuntimeError('录音太短，至少读 1 秒')
+    try:
+        seg = seg.speedup(playback_speed=1.5, chunk_size=150, crossfade=25)
+    except Exception:
+        seg = seg.speedup(playback_speed=1.5)
+    seg24 = seg.set_frame_rate(24000)
+    seg16 = seg.set_frame_rate(16000)
+    w24 = torch.tensor(np.frombuffer(seg24.raw_data, dtype=np.int16).astype(np.float32) / 32768.0)
+    w16 = torch.tensor(np.frombuffer(seg16.raw_data, dtype=np.int16).astype(np.float32) / 32768.0)
+    validate_clone_audio(w16)
+    mimi_dev = next(M['mimi'].parameters()).device
+    mimi_dtype = torch.float16 if mimi_dev.type != 'cpu' else torch.float32
+    with torch.inference_mode():
+        t = w24.unsqueeze(0).unsqueeze(0).to(device=mimi_dev, dtype=mimi_dtype)
+        codes = M['mimi'].encode(t).audio_codes
+        nf = math.ceil(w24.shape[-1] / SAMPLES_PER_FRAME)
+        ref_codes = codes[0, :8, :nf].cpu()[:, :min(nf, REF_FRAMES)]
+    with torch.no_grad():
+        mel = M['mel_fn'](w16.unsqueeze(0).to(M['device']))
+        feat = mel.clamp(min=1e-10).log().transpose(1, 2)
+        feat = feat - feat.mean(dim=1, keepdim=True)
+        spk_emb = M['campplus'](feat).squeeze(0).cpu()
+    return {'ref_codes': ref_codes, 'spk_emb': spk_emb}
+
+def run_generate(x, audio_inputs, audio_lens, pixel_values, **kw):
+    with MODEL_LOCK, torch.no_grad():
+        yield from M['model'].generate(
+            x, M['tokenizer'].eos_token_id, stream=True, return_audio_codes=True,
+            audio_inputs=audio_inputs, audio_lens=audio_lens, pixel_values=pixel_values, **kw)
+
+def load_main_model(model_path, model_name):
+    with MODEL_LOCK:
+        [sys.modules.pop(k) for k in list(sys.modules) if 'transformers_modules' in k]
+        M.pop('model', None); M.pop('tokenizer', None)
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
+        tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        m = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
+        vision_encoder, vision_processor = MiniMindOmni.load_vision('../model/siglip2-base-p32-256-ve')
+        audio_encoder, audio_processor = MiniMindOmni.load_sensevoice('../model/SenseVoiceSmall')
+        object.__setattr__(m, 'vision_encoder', vision_encoder)
+        object.__setattr__(m, 'vision_processor', vision_processor)
+        object.__setattr__(m, 'audio_encoder', audio_encoder)
+        object.__setattr__(m, 'audio_processor', audio_processor)
+        m = m.half().eval().to(M['device'])
+        if m.audio_encoder: m.audio_encoder.to(M['device'])
+        if m.vision_encoder: m.vision_encoder.to(M['device'])
+        M['tokenizer'], M['model'], M['model_name'] = tok, m, model_name
+        params = sum(p.numel() for p in m.parameters()) / 1e6
+        print(f'Loaded model: {model_name} ({params:.2f}M)')
+        return round(params, 2)
+
+def prepare_turn(text, samples, image_b64, do_asr_for_image):
+    """返回 (audio_inputs, audio_lens, pixel_values, prompt_for_model, user_text_for_history, asr_thread, asr_result)"""
+    audio_inputs = audio_lens = pixel_values = None
+    prompt = text or ''
+    user_text = text or ''
+    asr_thread, asr_result = None, [None]
+    if samples is not None:
+        from pydub import AudioSegment
+        if image_b64 and do_asr_for_image:
+            user_text = asr_run(samples)
+            prompt = user_text
+        else:
+            audio_inputs, audio_lens, prompt = prep_audio(samples)
+            if M['cfg'].max_history_turns > 0:
+                sa = samples.copy()
+                def _a(): asr_result[0] = asr_run(sa)
+                asr_thread = threading.Thread(target=_a); asr_thread.start()
+    if image_b64:
+        pixel_values = prep_image(image_b64)
+        m = M['model']
+        prompt = (prompt + "\n\n" if prompt else "") + "请描述这张图片\n\n" + m.config.image_special_token * m.config.image_token_len
+    return audio_inputs, audio_lens, pixel_values, prompt, user_text, asr_thread, asr_result
+
+
+# -------- web app (lazy: requires flask/flask_sock at runtime) --------
+def init_web_app():
+    from flask import Flask, request, Response, send_from_directory
+    from flask_cors import CORS
+    from flask_sock import Sock
+
+    app = Flask(__name__, static_folder='.')
+    CORS(app)
+    sock = Sock(app)
+
+    @app.route('/')
+    def index(): return send_from_directory('.', 'web_demo.html')
+    @app.route('/call')
+    def call_page(): return send_from_directory('.', 'web_demo.html')
+
+    @app.route('/voices')
+    def get_voices():
+        return json.dumps({'builtin': sorted(V_builtin), 'unseen': sorted(V_unseen), 'manual': sorted(V_manual)})
+
+    @app.route('/models')
+    def get_models():
+        return json.dumps({'models': list(M.get('models', {}).keys()), 'current': M.get('model_name')})
+
+    @app.route('/switch_model', methods=['POST'])
+    def switch_model():
+        name = (request.json or {}).get('name')
+        if name not in M.get('models', {}):
+            return Response(json.dumps({'ok': False, 'error': 'unknown model'}), status=400, mimetype='application/json')
+        try:
+            params = load_main_model(M['models'][name], name)
+            return Response(json.dumps({'ok': True, 'model': name, 'params': params}), mimetype='application/json')
+        except Exception as e:
+            return Response(json.dumps({'ok': False, 'error': str(e)}), status=500, mimetype='application/json')
+
+    @app.route('/clone_voice', methods=['POST'])
+    def clone_voice():
+        d = request.json or {}
+        if not d.get('audio'):
+            return Response(json.dumps({'ok': False, 'error': 'missing audio'}), status=400, mimetype='application/json')
+        try:
+            name = normalize_voice_name(d.get('name'))
+            value = build_clone_voice(d['audio'])
+            out_path = clone_voice_path()
+            saved = torch.load(out_path, map_location='cpu') if os.path.exists(out_path) else {}
+            saved[name] = value
+            torch.save(saved, out_path)
+            register_voice(name, value, group='manual')
+            return Response(json.dumps({'ok': True, 'voice': name, 'path': './model/speaker/' + CLONE_FILE}), mimetype='application/json')
+        except Exception as e:
+            return Response(json.dumps({'ok': False, 'error': str(e)}), status=500, mimetype='application/json')
+
+    @app.route('/delete_voice', methods=['POST'])
+    def delete_voice():
+        d = request.json or {}
+        name = ' '.join(str(d.get('name') or '').split())
+        if not name:
+            return Response(json.dumps({'ok': False, 'error': 'missing name'}), status=400, mimetype='application/json')
+        try:
+            delete_manual_voice(name)
+            return Response(json.dumps({'ok': True, 'voice': name}), mimetype='application/json')
+        except Exception as e:
+            return Response(json.dumps({'ok': False, 'error': str(e)}), status=500, mimetype='application/json')
+
+    @app.route('/chat', methods=['POST'])
+    def chat():
+        d = request.json
+        history = d.get('history', [])
+        samples = None
+        if d.get('audio'):
+            seg = AudioSegment.from_file(io.BytesIO(base64.b64decode(d['audio']))).set_frame_rate(16000).set_channels(1).set_sample_width(2)
+            samples = np.frombuffer(seg.raw_data, dtype=np.int16).astype(np.float32) / 32768.0
+        va = voice_args(d.get('voice', 'default'))
+
+        def gen():
+            audio_inputs, audio_lens, pixel_values, prompt, user_text, asr_th, asr_res = prepare_turn(
+                d.get('text', ''), samples, d.get('image'), do_asr_for_image=True)
+            x = build_ids(prompt, history)
+            asr_sent = False
+            if user_text and samples is not None and d.get('image'):
+                yield sse({'type': 'user_prompt', 'content': user_text}); asr_sent = True
+            frames, text_ttft, audio_ttft = [], None, None
+            t0 = time.time(); hi = 0
+            for y, af in run_generate(x, audio_inputs, audio_lens, pixel_values,
+                                       max_new_tokens=d.get('max_tokens', 512),
+                                       temperature=d.get('temperature', 1), top_p=0.85, **va):
+                if not asr_sent and asr_th and not asr_th.is_alive():
+                    asr_th.join()
+                    if asr_res[0]: yield sse({'type': 'user_prompt', 'content': asr_res[0]})
+                    asr_sent = True
+                if y is not None:
+                    if text_ttft is None:
+                        text_ttft = (time.time() - t0) * 1000
+                        yield sse({'type': 'ttft', 'text_ttft': round(text_ttft, 1)})
+                    ans = M['tokenizer'].decode(y[0].tolist(), skip_special_tokens=True)
+                    if ans and ans[-1] != '\ufffd' and len(ans) > hi:
+                        yield sse({'type': 'text', 'content': ans[hi:]}); hi = len(ans)
+                if af:
+                    if audio_ttft is None:
+                        audio_ttft = (time.time() - t0) * 1000
+                        yield sse({'type': 'ttft', 'audio_ttft': round(audio_ttft, 1)})
+                    frames.append(af)
+                    for pcm in stream_pcm(frames):
+                        b64 = base64.b64encode(pcm).decode()
+                        for i in range(0, len(b64), 2000):
+                            yield sse({'type': 'pcm', 'c': b64[i:i+2000], 'd': i+2000 >= len(b64)})
+            for pcm in stream_pcm(frames, flush=True):
+                b64 = base64.b64encode(pcm).decode()
+                for i in range(0, len(b64), 2000):
+                    yield sse({'type': 'pcm', 'c': b64[i:i+2000], 'd': i+2000 >= len(b64)})
+            if not asr_sent:
+                if asr_th:
+                    asr_th.join()
+                    if asr_res[0]: yield sse({'type': 'user_prompt', 'content': asr_res[0]})
+                else:
+                    yield sse({'type': 'user_prompt', 'content': prompt})
+            yield sse({'type': 'done'})
+
+        return Response(gen(), mimetype='text/event-stream', headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'})
+
+    @sock.route('/ws/realtime')
+    def realtime(ws):
+        session = RealtimeSession(M['vad_path'])
+        q = queue.Queue(); alive = [True]; state = {'history': [], 'image': None}
+        n_hist = M['cfg'].max_history_turns
+
+        def push_audio(data):
+            return session.push_chunk(np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0)
+
+        def set_ctx(msg):
+            h = msg.get('history') or []
+            state['history'] = h[-n_hist:] if n_hist > 0 else []
+            if 'image' in msg: state['image'] = msg.get('image')
+            if 'voice' in msg: state['voice'] = msg.get('voice', 'default')
+
+        def poll_interrupt():
+            while True:
+                try: data = q.get_nowait()
+                except queue.Empty: return False
+                if isinstance(data, bytes):
+                    if push_audio(data) == 'interrupt': return True
+                    ws.send(json.dumps({'type': 'vad', 'speaking': session.speaking}))
+                else:
+                    m = json.loads(data)
+                    if m.get('type') == 'context': set_ctx(m)
+                    elif m.get('type') in ('stop', 'end'):
+                        if m['type'] == 'end': alive[0] = False
+                        session.interrupt = True; return True
+
+        def recv_loop():
+            while alive[0]:
+                try:
+                    data = ws.receive(timeout=1)
+                    if data is None: alive[0] = False; break
+                    q.put(data)
+                except: alive[0] = False; break
+
+        threading.Thread(target=recv_loop, daemon=True).start()
+        try:
+            while alive[0]:
+                try: data = q.get(timeout=0.05)
+                except queue.Empty: continue
+                if isinstance(data, str):
+                    m = json.loads(data)
+                    if m.get('type') == 'context': set_ctx(m)
+                    elif m.get('type') == 'stop': session.interrupt = True
+                    elif m.get('type') == 'end': break
+                    continue
+                if session.generating:
+                    push_audio(data); ws.send(json.dumps({'type': 'vad', 'speaking': session.speaking})); continue
+                status = push_audio(data)
+                ws.send(json.dumps({'type': 'vad', 'speaking': session.speaking}))
+                if status != 'speech_end': continue
+
+                session.generating = True
+                audio = session.get_audio()
+                ws.send(json.dumps({'type': 'generating'}))
+                audio_inputs, audio_lens, pixel_values, prompt, user_text, asr_th, asr_res = prepare_turn(
+                    '', audio, state['image'], do_asr_for_image=True)
+                if state['image']: state['image'] = None
+                x = build_ids(prompt, state['history'])
+                va_rt = voice_args(state.get('voice', 'default'))
+
+                frames, full_text, interrupted = [], '', False
+                for y, af in run_generate(x, audio_inputs, audio_lens, pixel_values,
+                                           max_new_tokens=512, temperature=0.7, **va_rt):
+                    if poll_interrupt() or session.interrupt: interrupted = True; break
+                    if y is not None:
+                        ans = M['tokenizer'].decode(y[0].tolist(), skip_special_tokens=True)
+                        if ans and ans[-1] != '\ufffd' and len(ans) > len(full_text):
+                            ws.send(json.dumps({'type': 'text', 'content': ans[len(full_text):]})); full_text = ans
+                    if af:
+                        frames.append(af)
+                        for pcm in stream_pcm(frames):
+                            ws.send(json.dumps({'type': 'pcm', 'data': base64.b64encode(pcm).decode()}))
+                if not interrupted:
+                    for pcm in stream_pcm(frames, flush=True):
+                        ws.send(json.dumps({'type': 'pcm', 'data': base64.b64encode(pcm).decode()}))
+                if asr_th:
+                    asr_th.join(); user_text = asr_res[0] or user_text
+                if n_hist > 0:
+                    if user_text: state['history'].append({'role': 'user', 'content': user_text})
+                    if full_text: state['history'].append({'role': 'assistant', 'content': full_text})
+                    state['history'] = state['history'][-n_hist:]
+                ws.send(json.dumps({'type': 'done', 'interrupted': interrupted or session.interrupt}))
+                session.generating = False; session.interrupt = False
+        finally:
+            alive[0] = False
+
+    return app
+
+
+def init_model(args):
+    M['cfg'] = args; M['device'] = args.device
+    with contextlib.redirect_stdout(io.StringIO()):
+        from funasr import AutoModel
+        from funasr.utils.postprocess_utils import rich_transcription_postprocess
+        M['asr'] = AutoModel(model='../model/SenseVoiceSmall', trust_remote_code=True, device=args.device, disable_update=True)
+    M['models'] = scan_hf_models(args.load_from)
+    if not M['models']:
+        raise RuntimeError(f"未在 {os.path.abspath(args.load_from)} 找到 transformers 模型")
+    model_name = next(iter(M['models']))
+    load_main_model(M['models'][model_name], model_name)
+    try:
+        import torchaudio
+        from transformers import MimiModel
+        M['mimi'] = MimiModel.from_pretrained('../model/mimi').eval().to(args.device)
+        if args.device != 'cpu': M['mimi'] = M['mimi'].half()
+        print('Mimi model loaded')
+    except Exception:
+        M['mimi'] = None
+    try:
+        from modelscope.models.audio.sv.DTDNN import CAMPPlus
+        M['campplus'] = CAMPPlus(feat_dim=80, embedding_size=192, growth_rate=32, bn_size=4,
+                                 init_channels=128, config_str='batchnorm-relu', memory_efficient=True)
+        sd = torch.load('../model/campplus/campplus_cn_common.pt', map_location='cpu')
+        M['campplus'].load_state_dict({k: v.float() for k, v in sd.items()})
+        M['campplus'] = M['campplus'].eval().to(args.device)
+        import torchaudio
+        M['mel_fn'] = torchaudio.transforms.MelSpectrogram(
+            sample_rate=16000, n_fft=512, win_length=400, hop_length=160,
+            n_mels=80, f_min=20, f_max=7600, norm='slaney', mel_scale='slaney',
+        ).to(args.device)
+        print('CAM++ loaded')
+    except Exception as e:
+        M['campplus'], M['mel_fn'] = None, None
+        print(f'CAM++ load failed: {e}')
+    M['vad_path'] = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'model', 'vad', 'silero_vad.onnx')
+    spk_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'model', 'speaker')
+    for fn, group in [('voices.pt', 'builtin'), ('voices_unseen.pt', 'unseen'), (CLONE_FILE, 'manual')]:
+        fp = os.path.join(spk_dir, fn)
+        if os.path.exists(fp):
+            for speaker, v in torch.load(fp, map_location=args.device).items():
+                if speaker not in V or fn == CLONE_FILE:
+                    register_voice(speaker, v, group=group)
+    if V: print(f'Loaded {len(V)} voices: builtin={sorted(V_builtin)}, unseen={sorted(V_unseen)}, manual={sorted(V_manual)}')
+    log_model_params(M['model'])
+    print('Warmup...')
+    with torch.no_grad():
+        ids = torch.tensor([[1, 2, 3]], device=args.device)
+        au = torch.full((1, 8, 3), 2049, dtype=torch.long, device=args.device)
+        M['model'].forward(torch.cat((au, ids.unsqueeze(1)), dim=1))
+        if M['model'].audio_encoder: M['model'].audio_encoder(torch.zeros(1, 100, 560, device=args.device), torch.tensor([100], device=args.device))
+        if M['mimi']: M['mimi'].decode(torch.zeros(1, 8, 1, dtype=torch.long, device=args.device))
+    print('Warmup done!')
+
+
+if __name__ == '__main__':
+    p = argparse.ArgumentParser()
+    p.add_argument('--load_from', default='./', help='模型权重搜索目录；目录下可放多个 HF 格式模型，WebUI 会自动扫描并允许切换。')
+    p.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', help='推理设备；CUDA 可用时默认 cuda。显存不足或排查环境问题时可改为 cpu。')
+    p.add_argument('--port', default=7860, type=int, help='WebUI 服务端口；端口被占用或需要同时启动多个实例时调整。')
+    p.add_argument('--audio_chunk_frames', default=4, type=int, help='流式播放每次解码的 Mimi frame 数；默认 4 约 320ms。WebUI 播放卡顿时可调大到 8/12，低延迟优先时保持 4。')
+    p.add_argument('--audio_overlap', default=2, type=int, help='分块 Mimi 解码的重叠帧数；默认 2 用于缓解块边界断裂。一般不需要调整，边界杂音明显时可适当增大。')
+    p.add_argument('--max_history_turns', default=0, type=int, help='对话历史轮数；默认 0 不带历史以降低延迟和显存。需要多轮上下文时调大，但会增加 prefill 成本。')
+    args = p.parse_args()
+    init_model(args)
+    app = init_web_app()
+    app.run(host='0.0.0.0', port=args.port, threaded=True)

examples/serve_openai_api.py

@@ -0,0 +1,249 @@
+import argparse
+import json
+import re
+import os
+
+import time
+import torch
+import warnings
+import uvicorn
+
+from threading import Thread
+from queue import Queue
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models.lora import apply_lora, load_lora
+
+warnings.filterwarnings('ignore')
+
+app = FastAPI()
+
+
+def init_model(args):
+    tokenizer = AutoTokenizer.from_pretrained(args.load_from)
+    if 'model' in args.load_from:
+        moe_suffix = '_moe' if args.use_moe else ''
+        ckp = f'../{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
+        model = MiniMindForCausalLM(MiniMindConfig(
+            hidden_size=args.hidden_size,
+            num_hidden_layers=args.num_hidden_layers,
+            max_seq_len=args.max_seq_len,
+            use_moe=bool(args.use_moe),
+            inference_rope_scaling=args.inference_rope_scaling
+        ))
+        model.load_state_dict(torch.load(ckp, map_location=device), strict=True)
+        if args.lora_weight != 'None':
+            apply_lora(model)
+            load_lora(model, f'../{args.save_dir}/lora/{args.lora_weight}_{args.hidden_size}.pth')
+    else:
+        model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True)
+    print(f'MiniMind模型参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f} M(illion)')
+    return model.half().eval().to(device), tokenizer
+
+
+class ChatRequest(BaseModel):
+    model: str
+    messages: list
+    temperature: float = 0.7
+    top_p: float = 0.92
+    max_tokens: int = 8192
+    stream: bool = True
+    tools: list = Field(default_factory=list)
+    open_thinking: bool = False
+    chat_template_kwargs: dict = None
+    
+    def get_open_thinking(self) -> bool:
+        """兼容多种方式开启 thinking"""
+        if self.open_thinking:
+            return True
+        if self.chat_template_kwargs:
+            return self.chat_template_kwargs.get('open_thinking', False) or \
+                   self.chat_template_kwargs.get('enable_thinking', False)
+        return False
+
+
+class CustomStreamer(TextStreamer):
+    def __init__(self, tokenizer, queue):
+        super().__init__(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        self.queue = queue
+        self.tokenizer = tokenizer
+
+    def on_finalized_text(self, text: str, stream_end: bool = False):
+        self.queue.put(text)
+        if stream_end:
+            self.queue.put(None)
+
+
+def parse_response(text):
+    reasoning_content = None
+    think_match = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
+    if think_match:
+        reasoning_content = think_match.group(1).strip()
+        text = re.sub(r'<think>.*?</think>\s*', '', text, flags=re.DOTALL)
+    elif '</think>' in text:
+        parts = text.split('</think>', 1)
+        reasoning_content = parts[0].strip()
+        text = parts[1].strip() if len(parts) > 1 else ''
+    tool_calls = []
+    for i, m in enumerate(re.findall(r'<tool_call>(.*?)</tool_call>', text, re.DOTALL)):
+        try:
+            call = json.loads(m.strip())
+            tool_calls.append({"id": f"call_{int(time.time())}_{i}", "type": "function", "function": {"name": call.get("name", ""), "arguments": json.dumps(call.get("arguments", {}), ensure_ascii=False)}})
+        except Exception:
+            pass
+    if tool_calls:
+        text = re.sub(r'<tool_call>.*?</tool_call>', '', text, flags=re.DOTALL)
+    return text.strip(), reasoning_content, tool_calls or None
+
+
+def generate_stream_response(messages, temperature, top_p, max_tokens, tools=None, open_thinking=False):
+    try:
+        new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools or None, open_thinking=open_thinking)
+        inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
+
+        queue = Queue()
+        streamer = CustomStreamer(tokenizer, queue)
+
+        def _generate():
+            try:
+                model.generate(
+                    inputs.input_ids,
+                    max_new_tokens=max_tokens,
+                    do_sample=True,
+                    temperature=temperature,
+                    top_p=top_p,
+                    attention_mask=inputs.attention_mask,
+                    pad_token_id=tokenizer.pad_token_id,
+                    eos_token_id=tokenizer.eos_token_id,
+                    streamer=streamer
+                )
+            except Exception as e:
+                queue.put({"error": str(e)})
+                queue.put(None)
+
+        Thread(target=_generate).start()
+
+        full_text = ""
+        emitted = 0
+        thinking_ended = not bool(open_thinking)
+
+        while True:
+            text = queue.get()
+            if text is None:
+                break
+            if isinstance(text, dict):
+                yield json.dumps(text, ensure_ascii=False)
+                continue
+            full_text += text
+
+            if not thinking_ended:
+                pos = full_text.find('</think>')
+                if pos >= 0:
+                    thinking_ended = True
+                    new_r = full_text[emitted:pos]
+                    if new_r:
+                        yield json.dumps({"choices": [{"delta": {"reasoning_content": new_r}}]}, ensure_ascii=False)
+                    emitted = pos + len('</think>')
+                    after = full_text[emitted:].lstrip('\n')
+                    emitted = len(full_text) - len(after)
+                    if after:
+                        yield json.dumps({"choices": [{"delta": {"content": after}}]}, ensure_ascii=False)
+                        emitted = len(full_text)
+                else:
+                    new_r = full_text[emitted:]
+                    if new_r:
+                        yield json.dumps({"choices": [{"delta": {"reasoning_content": new_r}}]}, ensure_ascii=False)
+                        emitted = len(full_text)
+            else:
+                new_c = full_text[emitted:]
+                if new_c:
+                    yield json.dumps({"choices": [{"delta": {"content": new_c}}]}, ensure_ascii=False)
+                    emitted = len(full_text)
+
+        _, _, tool_calls = parse_response(full_text)
+        if tool_calls:
+            yield json.dumps({"choices": [{"delta": {"tool_calls": tool_calls}}]}, ensure_ascii=False)
+        yield json.dumps({"choices": [{"delta": {}, "finish_reason": "tool_calls" if tool_calls else "stop"}]}, ensure_ascii=False)
+
+    except Exception as e:
+        yield json.dumps({"error": str(e)})
+
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: ChatRequest):
+    try:
+        if request.stream:
+            return StreamingResponse(
+                (f"data: {chunk}\n\n" for chunk in generate_stream_response(
+                    messages=request.messages,
+                    temperature=request.temperature,
+                    top_p=request.top_p,
+                    max_tokens=request.max_tokens,
+                    tools=request.tools,
+                    open_thinking=request.get_open_thinking()
+                )),
+                media_type="text/event-stream"
+            )
+        else:
+            new_prompt = tokenizer.apply_chat_template(
+                request.messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                tools=request.tools or None,
+                open_thinking=request.get_open_thinking()
+            )
+            inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
+            with torch.no_grad():
+                generated_ids = model.generate(
+                    inputs["input_ids"],
+                    max_length=inputs["input_ids"].shape[1] + request.max_tokens,
+                    do_sample=True,
+                    attention_mask=inputs["attention_mask"],
+                    pad_token_id=tokenizer.pad_token_id,
+                    eos_token_id=tokenizer.eos_token_id,
+                    top_p=request.top_p,
+                    temperature=request.temperature
+                )
+                answer = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+            content, reasoning_content, tool_calls = parse_response(answer)
+            message = {"role": "assistant", "content": content}
+            if reasoning_content:
+                message["reasoning_content"] = reasoning_content
+            if tool_calls:
+                message["tool_calls"] = tool_calls
+            return {
+                "id": f"chatcmpl-{int(time.time())}",
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": "minimind",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": message,
+                        "finish_reason": "tool_calls" if tool_calls else "stop"
+                    }
+                ]
+            }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Server for MiniMind")
+    parser.add_argument('--load_from', default='../model', type=str, help="模型加载路径（model=原生torch权重，其他路径=transformers格式）")
+    parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
+    parser.add_argument('--weight', default='full_sft', type=str, help="权重名称前缀（pretrain, full_sft, dpo, reason, ppo_actor, grpo, spo）")
+    parser.add_argument('--lora_weight', default='None', type=str, help="LoRA权重名称（None表示不使用，可选：lora_identity, lora_medical）")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--max_seq_len', default=8192, type=int, help="最大序列长度")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--inference_rope_scaling', default=False, action='store_true', help="启用RoPE位置编码外推（4倍，仅解决位置编码问题）")
+    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
+    args = parser.parse_args()
+    device = args.device
+    model, tokenizer = init_model(args)
+    uvicorn.run(app, host="0.0.0.0", port=8998)

examples/web_demo.html

@@ -0,0 +1,748 @@
+<!DOCTYPE html>
+<html lang="zh">
+<head>
+<meta charset="UTF-8">
+<meta name="viewport" content="width=device-width, initial-scale=1.0">
+<title>MiniMind-O</title>
+<link rel="icon" type="image/png" href="https://www.modelscope.cn/studio/gongjy/MiniMind/resolve/master/images/logo2.png">
+<style>
+:root {
+  --bg:#1e1e20; --p1:#2a2a2e; --p2:#35353b; --p3:#40404a;
+  --bd:#55555e; --bd2:#7a7a85;
+  --fg:#ececef; --mu:#9a9aa0; --dim:#8a8a92;
+  --u:#7fc5e0; --a:#e5b57a; --danger:#d84545; --cam:#7ad88a;
+  --serif:'Georgia','Times New Roman',serif; --mono:'SF Mono',Consolas,monospace;
+}
+*{margin:0;padding:0;box-sizing:border-box}
+html,body{height:100%}
+body{height:100vh;background:var(--bg);color:var(--fg);font-family:var(--serif);display:flex;align-items:center;justify-content:center;overflow:hidden}
+.app{max-width:860px;width:100%;height:calc(100vh - 40px);display:flex;flex-direction:column;padding:20px 24px;background:var(--p1);border:1px solid var(--bd);overflow:hidden}
+.topbar{display:flex;align-items:center;justify-content:space-between;padding:4px 0 12px;flex-shrink:0;gap:10px}
+.brand{display:flex;align-items:center;gap:10px;min-width:0;flex:1}
+.brand>div{min-width:0}
+.brand img{height:34px}
+.brand h1{font-size:1.15rem;font-weight:700;font-style:italic;color:#fff}
+.brand .sub{font-size:.72rem;color:var(--mu);font-style:italic}
+.topright{display:flex;gap:8px;align-items:center}
+.btn{padding:5px 10px;font-size:.75rem;cursor:pointer;background:var(--p2);border:1px solid var(--bd);color:var(--fg);font-family:var(--serif);font-style:italic;display:inline-flex;align-items:center;gap:5px}
+.btn:hover:not(:disabled){background:var(--p3);border-color:var(--bd2)}
+.btn:disabled{opacity:.4;cursor:not-allowed}
+.btn.danger{color:var(--danger);border-color:#6a3a3a}
+.btn.danger:hover:not(:disabled){background:var(--danger);border-color:var(--danger);color:#fff}
+.btn svg{width:12px;height:12px;fill:currentColor}
+.model-sel{max-width:150px;padding:5px 8px;font-size:.72rem;background:var(--p2);border:1px solid var(--bd);color:var(--fg);font-family:var(--serif);font-style:italic;outline:none}
+.brand .model-sel{max-width:260px;padding:0 34px 0 0;background:transparent;border:none;font-size:1.15rem;font-weight:700;color:#fff}
+.model-sel:focus{border-color:var(--bd2)}
+.mini-sel{max-width:130px;padding:5px 8px;font-size:.72rem;background:var(--p2);border:1px solid var(--bd);color:var(--fg);font-family:var(--serif);font-style:italic;outline:none}
+.mini-sel:focus{border-color:var(--bd2)}
+.voice-sel{display:none}
+.more{position:relative}
+.more-pop{top:100%;right:0;left:auto;margin-top:6px;min-width:130px}
+.vgroup{padding:4px 10px;font-size:.6rem;color:var(--mu);font-style:italic}
+.vitem{padding:5px 10px;display:flex;justify-content:space-between;align-items:center;cursor:pointer;font-size:.72rem;font-family:var(--serif);font-style:italic;color:var(--fg);gap:8px}
+.vitem:hover{background:var(--p3)}
+.vitem.on{color:var(--u)}
+.vx{color:var(--mu);cursor:pointer;font-style:normal;font-size:.8rem;padding:0 2px;line-height:1}
+.vx:hover{color:var(--danger)}
+.vadd{padding:5px 10px;font-size:.72rem;color:var(--mu);cursor:pointer;font-family:var(--serif);font-style:italic;border-top:1px solid var(--bd)}
+.vadd:hover{background:var(--p3);color:var(--fg)}
+.status{display:inline-flex;align-items:center;gap:10px;padding:5px 10px;font-size:.65rem;font-family:var(--mono);background:var(--p2);border:1px solid var(--bd);transition:opacity .3s}
+.status.hidden{opacity:0;pointer-events:none}
+.status .row{display:inline-flex;gap:5px;white-space:nowrap}
+.status .lab{color:var(--mu);font-style:italic}
+
+.chat{flex:1;overflow-y:auto;padding:12px 0;min-height:0;display:flex;flex-direction:column;gap:18px;scrollbar-width:none;-ms-overflow-style:none}
+.chat::-webkit-scrollbar{display:none}
+.msg{display:flex;flex-direction:column;gap:6px;animation:fi .25s ease}
+.msg.user{align-items:flex-end}
+.msg.user .body{align-items:flex-end}
+@keyframes fi{from{opacity:0;transform:translateY(3px)}to{opacity:1;transform:translateY(0)}}
+.meta{font-size:.72rem;font-style:italic}
+.msg.user .meta{color:var(--u)}
+.msg.assistant .meta{color:var(--a)}
+.body{display:flex;flex-direction:column;gap:8px}
+.text{line-height:1.65;white-space:pre-wrap;word-break:break-word}
+.img{display:block;max-width:160px;max-height:160px}
+.msg.user .text:not(.voice){color:var(--u)}
+.voice{font-family:var(--mono);font-size:1.05rem;line-height:1.2;color:var(--bd)}
+.msg.user .voice{color:rgba(127,197,224,.5)}
+.prog{text-decoration:underline dashed;text-decoration-color:var(--bd);text-underline-offset:4px;text-decoration-thickness:1px}
+.prog .tk{cursor:pointer;padding:1px 0;transition:color .08s,background .08s}
+.msg.assistant .prog .tk.played{color:var(--a)}
+.msg.assistant .prog .tk:hover{background:rgba(229,181,122,.2)}
+.msg.assistant .prog.on .tk.cur{background:rgba(229,181,122,.45)}
+.msg.user .prog .tk.played{color:var(--u)}
+.msg.user .prog .tk:hover{background:rgba(127,197,224,.18)}
+.msg.user .prog.on .tk.cur{background:rgba(127,197,224,.35)}
+.player{display:inline-flex;align-items:center;gap:10px;padding:6px 12px;width:320px;max-width:100%;background:var(--p2);border:1px solid var(--bd)}
+.pp{background:transparent;border:none;color:inherit;cursor:pointer;font-family:var(--mono);font-size:.9rem;width:20px;text-align:left}
+.pp:hover{opacity:.75}
+.bar-wrap{flex:1;height:3px;cursor:pointer;background:var(--p1)}
+.bar{height:100%;width:0;transition:width .08s linear}
+.msg.user .bar,.msg.user .pp{background:var(--u);color:var(--u)}
+.msg.assistant .bar,.msg.assistant .pp{background:var(--a);color:var(--a)}
+.msg.user .bar-wrap{background:var(--p1)}
+.time{font-size:.7rem;font-family:var(--mono);min-width:76px;text-align:right;color:var(--mu)}
+
+.input-area{flex-shrink:0;padding:10px 0 4px}
+.compose{display:flex;gap:6px}
+.compose button{padding:0 12px;min-width:42px;cursor:pointer;font-size:.85rem;background:var(--p2);border:1px solid var(--bd);color:var(--fg);display:flex;align-items:center;justify-content:center;font-family:inherit}
+.compose button svg{width:16px;height:16px;fill:currentColor}
+.compose button:hover:not(:disabled){background:var(--p3);border-color:var(--bd2)}
+.compose button:disabled{opacity:.4;cursor:not-allowed}
+.compose button.primary{background:var(--fg);color:#1a1a1e;border-color:var(--fg)}
+.compose button.primary:hover:not(:disabled){background:#fff}
+.compose button.rec{animation:pulse 1s infinite;background:var(--danger);color:#fff;border-color:var(--danger)}
+@keyframes pulse{0%,100%{opacity:1}50%{opacity:.5}}
+.text-wrap{flex:1;display:flex;align-items:center;gap:6px;padding:0 6px 0 10px;min-width:0;background:var(--p2);border:1px solid var(--bd)}
+.text-wrap:focus-within{border-color:var(--fg)}
+.text-wrap input{flex:1;min-width:0;padding:10px 4px;font-size:.9rem;outline:none;background:transparent;border:none;color:inherit;font-family:inherit}
+.text-wrap input::placeholder{color:var(--dim)}
+.attach{display:flex;gap:4px;flex-wrap:nowrap;flex-shrink:0;max-width:45%;overflow-x:auto;scrollbar-width:none}
+.attach::-webkit-scrollbar{display:none}
+.attach:empty{display:none}
+.chip{display:inline-flex;align-items:center;gap:6px;padding:3px 7px;font-size:.72rem;background:var(--p1);border:1px solid var(--bd);flex-shrink:0}
+.chip .x{cursor:pointer;opacity:.55;font-size:.85rem}
+.chip .x:hover{opacity:1}
+.chip .play{cursor:pointer}
+.chip img{width:22px;height:22px;object-fit:cover}
+.pop{position:absolute;background:var(--p2);border:1px solid var(--bd);display:none;flex-direction:column;z-index:50;min-width:110px;font-family:var(--serif);font-style:italic;white-space:nowrap}
+.pop.on{display:flex}
+.pop button{padding:5px 10px;background:transparent;border:none;color:var(--fg);cursor:pointer;text-align:left;font-size:.72rem;font-family:inherit;font-style:inherit;line-height:1.2}
+.pop button:hover{background:var(--p3)}
+.snap{position:fixed;inset:0;z-index:200;background:rgba(0,0,0,.85);display:none;align-items:center;justify-content:center;flex-direction:column;gap:14px}
+.snap.on{display:flex}
+.snap video{max-width:80vw;max-height:70vh;background:#000;border:1px solid var(--bd)}
+.snap .sr{display:flex;gap:10px}
+.snap .sr button{padding:8px 18px;background:var(--p2);border:1px solid var(--bd);color:var(--fg);cursor:pointer;font-family:var(--serif);font-style:italic}
+.snap .sr button.primary{background:var(--fg);color:#1a1a1e;border-color:var(--fg)}
+.clone-card{width:min(92vw,520px);padding:22px;background:var(--p1);border:1px solid var(--bd);display:flex;flex-direction:column;gap:14px}
+.clone-title{font-size:1rem;font-style:italic;color:#fff}
+.clone-hint{font-size:.76rem;color:var(--mu);font-style:italic;line-height:1.6}
+.clone-name{width:100%;padding:10px 12px;background:var(--p2);border:1px solid var(--bd);color:var(--fg);font-family:inherit;font-size:.85rem;outline:none}
+.clone-name:focus{border-color:var(--fg)}
+.clone-script{padding:12px 14px;background:var(--p2);border:1px solid var(--bd);font-size:1rem;line-height:1.85}
+.clone-stat{min-height:18px;font-size:.72rem;color:var(--mu);font-family:var(--mono)}
+
+/* ===== Call mode ===== */
+body>.app,#call{will-change:transform,opacity,filter;transform-origin:center center}
+body>.app{transition:opacity .55s cubic-bezier(.22,1,.36,1),transform .7s cubic-bezier(.22,1,.36,1),filter .55s cubic-bezier(.22,1,.36,1)}
+body.call>.app{opacity:0;transform:scale(.92);filter:blur(8px);pointer-events:none}
+#call{position:fixed;inset:0;z-index:100;background:var(--bg);overflow:hidden;opacity:0;transform:scale(1.08);filter:blur(10px);pointer-events:none;visibility:hidden;transition:opacity .55s cubic-bezier(.22,1,.36,1),transform .7s cubic-bezier(.22,1,.36,1),filter .55s cubic-bezier(.22,1,.36,1),visibility 0s .7s}
+body.call #call{opacity:1;transform:scale(1);filter:blur(0);pointer-events:auto;visibility:visible;transition:opacity .55s cubic-bezier(.22,1,.36,1),transform .7s cubic-bezier(.22,1,.36,1),filter .55s cubic-bezier(.22,1,.36,1),visibility 0s 0s}
+#call .app{max-width:560px;margin:20px auto;padding:18px 22px;position:relative}
+#call .topbar{border-bottom:1px dashed var(--bd)}
+#call .brand h1{font-size:1.05rem}
+#call .brand .model-sel{max-width:230px;padding-right:38px;font-size:1.05rem}
+#call .brand .sub{font-size:.68rem}
+.pill{display:inline-flex;align-items:center;gap:6px;padding:5px 10px;font-size:.7rem;font-family:var(--mono);background:var(--p2);border:1px solid var(--bd);color:var(--mu);font-style:italic;line-height:1.2}
+.pill .dot{width:6px;height:6px;background:var(--mu);border-radius:50%}
+.pill.live{color:var(--u);border-color:#3c5966}
+.pill.live .dot{background:var(--u);box-shadow:0 0 6px var(--u);animation:cp 1.4s ease-in-out infinite}
+@keyframes cp{0%,100%{opacity:1}50%{opacity:.4}}
+#call .btn.active{background:var(--u);color:#1a1a1e;border-color:var(--u)}
+
+.hero{flex:1;display:flex;align-items:center;justify-content:center;min-height:0;padding:14px 0;position:relative}
+.orb{position:relative;width:280px;height:280px;max-width:60vh;max-height:60vh;display:flex;align-items:center;justify-content:center;will-change:transform}
+.orb svg{width:100%;height:100%;overflow:visible}
+.blob{fill:none;stroke-linejoin:round;transition:stroke .35s,opacity .45s}
+.blob.l0{stroke:#5a5a64;stroke-width:1.2;opacity:.55}
+.blob.l1{stroke:var(--bd2);stroke-width:1.4;opacity:.85}
+.blob.l2{stroke:var(--fg);stroke-width:1.6}
+.orb.camon .blob.l0,.orb.camon .blob.l1{opacity:0}
+.orb.camon .blob.l2{stroke:var(--cam);filter:drop-shadow(0 0 4px rgba(122,216,138,.45))}
+svg.listen .blob.l2{stroke:var(--u);filter:drop-shadow(0 0 4px rgba(127,197,224,.45))}
+svg.listen .blob.l1{stroke:rgba(127,197,224,.6)}
+svg.speak .blob.l2{stroke:var(--a);filter:drop-shadow(0 0 5px rgba(229,181,122,.5))}
+svg.speak .blob.l1{stroke:rgba(229,181,122,.6)}
+svg.gen .blob.l2{stroke:var(--fg);filter:drop-shadow(0 0 4px rgba(236,236,239,.35))}
+.core{display:none}
+svg.listen .core{fill:var(--u)}
+svg.speak .core{fill:var(--a)}
+.ripple{fill:none;stroke:var(--mu);stroke-width:1;transform-origin:center;animation:rip 1.6s ease-out forwards}
+.ripple.user{stroke:var(--u)}
+.ripple.assistant{stroke:var(--a)}
+@keyframes rip{0%{transform:scale(.55);opacity:.55}100%{transform:scale(1.35);opacity:0}}
+.lbl{position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);text-align:center;pointer-events:none;display:flex;flex-direction:column;align-items:center;gap:6px}
+.ic{width:34px;height:34px;color:var(--mu);transition:color .3s;display:flex;align-items:center;justify-content:center}
+.ic svg{width:100%;height:100%;stroke:currentColor;fill:none;stroke-width:2;stroke-linecap:round;stroke-linejoin:round}
+.lbl.listen .ic{color:var(--u)}
+.lbl.speak .ic{color:var(--a)}
+.lbl.gen .ic{color:var(--fg)}
+.lbl .st{display:none}
+.ic svg *{transform-box:view-box;transform-origin:12px 12px}
+.lbl.idle .ic .dot{animation:breathe 2.2s ease-in-out infinite;transform-origin:center;transform-box:view-box}
+@keyframes breathe{0%,100%{opacity:.4;transform:scale(.8)}50%{opacity:1;transform:scale(1.15)}}
+.lbl.listen .ic .w{animation:sigin 1.6s ease-out infinite}
+.lbl.listen .ic .w2{animation-delay:.3s}
+.lbl.listen .ic .w3{animation-delay:.6s}
+@keyframes sigin{0%{opacity:0;transform:scale(1.2)}60%{opacity:1;transform:scale(1)}100%{opacity:0;transform:scale(.85)}}
+.lbl.gen .ic svg{animation:spin 1.2s cubic-bezier(.5,.1,.5,.9) infinite}
+@keyframes spin{to{transform:rotate(360deg)}}
+.lbl.speak .ic .a1{animation:radout 1.3s ease-out infinite}
+.lbl.speak .ic .a2{animation:radout 1.3s ease-out .25s infinite}
+.lbl.speak .ic .a3{animation:radout 1.3s ease-out .5s infinite}
+@keyframes radout{0%{opacity:0;transform:scale(.7)}40%{opacity:1}100%{opacity:0;transform:scale(1.15)}}
+.up{display:none}
+.cam{position:absolute;top:50%;left:50%;width:180px;height:180px;object-fit:cover;background:var(--bg);border:1px solid var(--bd);transform:translate(-50%,-50%) scale(.3);opacity:0;transition:transform .5s cubic-bezier(.4,0,.2,1),opacity .3s,border-color .35s,box-shadow .35s;pointer-events:none;will-change:transform,opacity}
+.orb.camon .cam{transform:translate(-50%,-50%) scale(1);opacity:1;border-color:var(--cam);box-shadow:0 0 14px rgba(122,216,138,.35)}
+.orb.camon .core{opacity:0}
+.orb.camon .lbl{top:auto;bottom:-8px;transform:translate(-50%,100%)}
+.orb.camon .lbl .st{font-size:.9rem}
+svg.listen ~ .cam{border-color:var(--u);box-shadow:0 0 14px rgba(127,197,224,.45)}
+svg.speak ~ .cam{border-color:var(--a);box-shadow:0 0 14px rgba(229,181,122,.45)}
+
+.latest{width:100%;padding:10px 14px;background:var(--p2);border:1px solid var(--bd);display:flex;flex-direction:column;gap:5px;flex-shrink:0;max-height:120px;overflow-y:auto;scrollbar-width:none}
+.latest::-webkit-scrollbar{display:none}
+.latest .lm{font-family:var(--mono);font-size:.6rem;color:var(--mu);font-style:italic}
+.latest .lt{font-size:.82rem;line-height:1.55;font-style:italic;opacity:.9;white-space:pre-wrap;word-break:break-word}
+.latest.user .lm{color:var(--u)}
+.latest.assistant .lm{color:var(--a)}
+.latest.empty .lt{color:var(--bd)}
+.ctrls{display:flex;justify-content:center;gap:10px;padding:14px 0 2px;margin-top:12px;border-top:1px dashed var(--bd);flex-shrink:0}
+.cbtn{width:72px;height:58px;background:var(--p2);border:1px solid var(--bd);color:var(--fg);cursor:pointer;display:flex;flex-direction:column;align-items:center;justify-content:center;gap:3px;font-family:var(--mono);font-size:.58rem;letter-spacing:1px;text-transform:uppercase;transition:all .18s}
+.cbtn svg{width:18px;height:18px;fill:currentColor}
+.cbtn:hover{background:var(--p3);border-color:var(--bd2)}
+.cbtn.active{background:var(--u);color:#1a1a1e;border-color:var(--u)}
+.cbtn.danger{color:var(--danger);border-color:#6a3a3a}
+.cbtn.danger:hover{background:var(--danger);color:#fff}
+
+.trans{position:absolute;inset:0;background:var(--p1);z-index:20;display:flex;flex-direction:column;padding:18px 22px;transform:translateY(100%);transition:transform .3s ease;will-change:transform;contain:layout paint}
+.trans.on{transform:translateY(0)}
+.trans .th{display:flex;justify-content:space-between;padding-bottom:10px;border-bottom:1px dashed var(--bd)}
+.trans .tt{font-style:italic;font-size:.95rem;display:flex;gap:10px;align-items:center}
+.trans .cnt{font-family:var(--mono);font-size:.68rem;color:var(--mu)}
+#call .chat:empty::before{content:'no turns yet · say something';display:block;text-align:center;color:var(--bd);font-style:italic;padding:40px 0;font-size:.82rem}
+@media(max-width:540px){
+  body{display:block}
+  .app{height:100vh;height:100dvh;max-width:none;border:none;padding:14px 16px}
+  .brand img{height:28px}
+  .brand h1{font-size:1rem}
+  .brand .sub{display:none}
+  .topbar{gap:6px}
+  .status{display:none!important}
+  .player{width:100%}
+  .vitem,.vadd{padding:8px 12px}
+  .compose button{min-width:38px;padding:0 10px}
+  .text-wrap input{padding:8px 4px;font-size:.85rem}
+  .clone-card{padding:16px}
+  .clone-hint{font-size:.7rem}
+  #call .app{margin:0;height:100vh;height:100dvh;border:none;padding:14px 16px}
+  .orb{width:240px;height:240px}
+  .cbtn{width:58px;height:52px}
+}
+</style>
+</head>
+<body>
+<div class="app">
+  <div class="topbar">
+    <div class="brand">
+      <img src="https://www.modelscope.cn/studio/gongjy/MiniMind/resolve/master/images/logo2.png" alt="logo">
+      <div><select class="model-sel" id="modelSel"></select><div class="sub">multimodal inference · text / image / audio → text + audio</div></div>
+    </div>
+    <div class="topright">
+      <div class="status hidden" id="stat"><div class="row"><span class="lab">text_ttft</span><span id="tt">—</span></div><div class="row"><span class="lab">audio_ttft</span><span id="ta">—</span></div><div class="row"><span class="lab">status</span><span id="ts">idle</span></div></div>
+      <select class="voice-sel" id="voiceSel"><option value="default">default</option></select>
+      <div class="more">
+        <button class="btn" id="moreBtn">default</button>
+        <div class="pop more-pop" id="morePop"></div>
+      </div>
+      <button class="btn" id="callBtn"><svg viewBox="0 0 24 24"><path d="M6.62 10.79a15.05 15.05 0 0 0 6.59 6.59l2.2-2.2a1 1 0 0 1 1.02-.24c1.12.37 2.32.57 3.57.57a1 1 0 0 1 1 1V20a1 1 0 0 1-1 1A17 17 0 0 1 3 4a1 1 0 0 1 1-1h3.5a1 1 0 0 1 1 1c0 1.25.2 2.45.57 3.57.1.35 0 .74-.25 1.02l-2.2 2.2z"/></svg>Call</button>
+    </div>
+  </div>
+  <div class="chat" id="chat"></div>
+  <div class="input-area"><div class="compose">
+    <div style="position:relative;display:flex">
+      <button id="imgBtn" title="image"><svg viewBox="0 0 24 24"><path d="M21 19V5a2 2 0 0 0-2-2H5a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h14a2 2 0 0 0 2-2zm-2 0H5l3.5-4.5 2.5 3.01L14.5 13l4.5 6zM8.5 8A1.5 1.5 0 1 1 7 9.5 1.5 1.5 0 0 1 8.5 8z"/></svg></button>
+      <div class="pop" id="imgPop" style="bottom:100%;left:0;margin-bottom:6px"><button id="fileBtn">from file</button><button id="shotBtn">take photo</button></div>
+    </div>
+    <input type="file" id="imgIn" accept="image/*" hidden>
+    <button id="recBtn" title="录音"><svg viewBox="0 0 24 24"><path d="M12 14c1.66 0 3-1.34 3-3V5c0-1.66-1.34-3-3-3S9 3.34 9 5v6c0 1.66 1.34 3 3 3zm5.91-3c-.49 0-.9.36-.98.85C16.52 14.2 14.47 16 12 16s-4.52-1.8-4.93-4.15c-.08-.49-.49-.85-.98-.85-.61 0-1.09.54-1 1.14.49 3 2.89 5.35 5.91 5.78V20c0 .55.45 1 1 1s1-.45 1-1v-2.08c3.02-.43 5.42-2.78 5.91-5.78.1-.6-.39-1.14-1-1.14z"/></svg></button>
+    <div class="text-wrap"><div class="attach" id="att"></div><input type="text" id="txt" placeholder="Message MiniMind-O..." autocomplete="off"></div>
+    <button id="sendBtn" class="primary"><svg viewBox="0 0 24 24"><path d="M2.01 21L23 12 2.01 3 2 10l15 2-15 2z"/></svg></button>
+  </div></div>
+</div>
+
+<div class="snap" id="snap"><video id="snapV" autoplay playsinline muted></video><div class="sr"><button id="snapCancel">cancel</button><button id="snapOk" class="primary">capture</button></div></div>
+<div class="snap" id="cloneModal"><div class="clone-card"><div class="clone-title">Clone Voice</div><div class="clone-hint">请自然朗读下面这句话。建议 3 到 6 秒，环境尽量安静。录完后会自动做 1.5 倍速、提取 `CAM++` speaker embedding 与 `Mimi` ref codes，然后立即加入音色列表。</div><input class="clone-name" id="cloneName" placeholder="音色名，例如：小王" maxlength="24"><div class="clone-script" id="cloneText">今天阳光很好，欢迎来到我的语音克隆小实验。</div><div class="clone-stat" id="cloneStat">点击 start 开始录音</div><div class="sr"><button id="cloneCancel">cancel</button><button id="cloneRec" class="primary">start</button></div></div></div>
+
+<div id="call"><div class="app">
+  <div class="topbar">
+    <div class="brand"><img src="https://www.modelscope.cn/studio/gongjy/MiniMind/resolve/master/images/logo2.png" alt=""><div><select class="model-sel" id="callModelSel"></select><div class="sub">live voice · bidirectional streaming</div></div></div>
+    <div class="topright">
+      <select class="mini-sel" id="callVoiceSel"><option value="default">default</option></select>
+      <button class="btn" id="logBtn"><svg viewBox="0 0 24 24"><path d="M14 2H6c-1.1 0-1.99.9-1.99 2L4 20c0 1.1.89 2 1.99 2H18c1.1 0 2-.9 2-2V8l-6-6zM6 20V4h7v5h5v11H6zm2-3h8v-2H8v2zm0-4h8v-2H8v2z"/></svg>log</button>
+      <span class="pill" id="pill"><span class="dot"></span><span id="connT">connecting</span></span>
+    </div>
+  </div>
+  <div class="hero"><div class="orb" id="orb">
+    <svg viewBox="-160 -160 320 320" id="osvg"><g id="rip"></g><path class="blob l0" id="b0"/><path class="blob l1" id="b1"/><path class="blob l2" id="b2"/><circle class="core" id="core" cx="0" cy="0" r="4"/></svg>
+    <video class="cam" id="camV" autoplay playsinline muted></video>
+    <div class="lbl" id="lbl"><div class="ic" id="stIc"></div><div class="st" id="stT">init</div><div class="up" id="upT">00:00:00</div></div>
+  </div></div>
+  <div class="latest empty" id="lat"><span class="lm" id="lmE">latest · awaiting first turn</span><span class="lt" id="ltE">speak naturally — model responds with voice and text</span></div>
+  <div class="ctrls">
+    <button class="cbtn" id="camBtn"><svg viewBox="0 0 24 24"><path d="M17 10.5V7c0-.55-.45-1-1-1H4c-.55 0-1 .45-1 1v10c0 .55.45 1 1 1h12c.55 0 1-.45 1-1v-3.5l4 4v-11l-4 4z"/></svg><span>cam</span></button>
+    <button class="cbtn danger" id="endBtn"><svg viewBox="0 0 24 24"><path d="M12 9c-1.6 0-3.15.25-4.6.72v3.1c0 .39-.23.74-.56.9-.98.49-1.87 1.12-2.66 1.85-.18.18-.43.28-.7.28-.28 0-.53-.11-.71-.29L.29 13.08c-.18-.17-.29-.42-.29-.7 0-.28.11-.53.29-.71C3.34 8.78 7.46 7 12 7s8.66 1.78 11.71 4.67c.18.18.29.43.29.71 0 .28-.11.53-.29.71l-2.48 2.48c-.18.18-.43.29-.71.29-.27 0-.52-.11-.7-.28-.79-.74-1.69-1.36-2.67-1.85-.33-.16-.56-.5-.56-.9v-3.1C15.15 9.25 13.6 9 12 9z"/></svg><span>end</span></button>
+  </div>
+  <div class="trans" id="trans"><div class="th"><div class="tt">session transcript <span class="cnt" id="tcnt">0 turns</span></div><button class="btn" id="closeBtn">× close</button></div><div class="chat" id="cchat"></div></div>
+</div></div>
+
+<script>
+const $=id=>document.getElementById(id);
+const chat=$('chat'),txt=$('txt'),sendBtn=$('sendBtn'),recBtn=$('recBtn'),imgBtn=$('imgBtn'),imgIn=$('imgIn'),callBtn=$('callBtn'),att=$('att');
+const stat=$('stat'),tt=$('tt'),ta=$('ta'),ts=$('ts'),voiceSel=$('voiceSel'),modelSel=$('modelSel'),callModelSel=$('callModelSel'),callVoiceSel=$('callVoiceSel');
+const moreBtn=$('moreBtn'),morePop=$('morePop');
+const CLONE_PROMPT='今天阳光很好，欢迎来到我的语音克隆小实验。';
+let voiceGroups={builtin:[],unseen:[],manual:[]};
+const closeMore=()=>morePop.classList.remove('on');
+const modelLabel=m=>m;
+async function loadVoices(sel){
+  const keep=sel||voiceSel.value||'default';
+  const keepCall=callVoiceSel.value||'default';
+  voiceSel.innerHTML='<option value="default">default</option>';
+  const d=await fetch('/voices').then(r=>r.json());
+  voiceGroups={builtin:d.builtin||[],unseen:d.unseen||[],manual:d.manual||[]};
+  for(const [label,names] of [['builtin',voiceGroups.builtin],['unseen',voiceGroups.unseen],['manual',voiceGroups.manual]]){
+    if(!names.length)continue;
+    const g=document.createElement('optgroup');
+    g.label=label;
+    names.forEach(v=>{const o=document.createElement('option');o.value=v;o.textContent=v;g.appendChild(o)});
+    voiceSel.appendChild(g);
+  }
+  if([...voiceSel.options].some(o=>o.value===keep))voiceSel.value=keep;
+  callVoiceSel.innerHTML=voiceSel.innerHTML;
+  if([...callVoiceSel.options].some(o=>o.value===keepCall))callVoiceSel.value=keepCall;
+  renderVoiceList();
+}
+function renderVoiceList(){
+  let h='';const sel=voiceSel.value;
+  moreBtn.textContent=sel;
+  h+=`<div class="vitem${sel==='default'?' on':''}" data-voice="default">default</div>`;
+  for(const [g,ns] of [['builtin',voiceGroups.builtin],['unseen',voiceGroups.unseen],['manual',voiceGroups.manual]]){
+    if(!ns.length)continue;
+    h+=`<div class="vgroup">${g}</div>`;
+    for(const v of ns)h+=`<div class="vitem${v===sel?' on':''}" data-voice="${v}"><span>${v}</span>${g==='manual'?`<span class="vx" data-del="${v}">×</span>`:''}</div>`;
+  }
+  h+='<div class="vadd">+ clone</div>';
+  morePop.innerHTML=h;
+}
+loadVoices().catch(()=>{});
+async function loadModels(){
+  const d=await fetch('/models').then(r=>r.json());
+  modelSel.innerHTML='';callModelSel.innerHTML='';
+  for(const m of d.models||[]){const o=document.createElement('option');o.value=m;o.textContent=modelLabel(m);modelSel.appendChild(o);callModelSel.appendChild(o.cloneNode(true))}
+  if(d.current){modelSel.value=d.current;callModelSel.value=d.current}
+}
+async function switchModelSelect(sel, oldSel){
+  const old=sel.dataset.current||sel.value,name=sel.value;
+  const wasCall=document.body.classList.contains('call');
+  stopChat();if(callStop)callStop();stat.classList.remove('hidden');ts.textContent='switching';sel.disabled=true;
+  try{
+    const r=await fetch('/switch_model',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({name})});
+    const d=await r.json();
+    if(!r.ok||!d.ok)throw new Error(d.error||'switch failed');
+    modelSel.value=name;callModelSel.value=name;modelSel.dataset.current=name;callModelSel.dataset.current=name;ts.textContent=`loaded ${name}`;if(wasCall)openCall();
+  }catch(e){sel.value=old;if(oldSel)oldSel.value=old;ts.textContent='switch failed';alert(e.message)}
+  sel.disabled=false;
+}
+modelSel.onchange=()=>switchModelSelect(modelSel,callModelSel);
+callModelSel.onchange=()=>switchModelSelect(callModelSel,modelSel);
+callVoiceSel.onchange=()=>{if(window.callSendCtx)window.callSendCtx()};
+loadModels().then(()=>{modelSel.dataset.current=modelSel.value;callModelSel.dataset.current=callModelSel.value}).catch(()=>{});
+
+/* ------------- shared helpers ------------- */
+function decPcm(b64){
+  const bin=atob(b64),i16=new Int16Array(bin.length/2);
+  for(let i=0;i<i16.length;i++)i16[i]=bin.charCodeAt(i*2)|(bin.charCodeAt(i*2+1)<<8);
+  const f32=new Float32Array(i16.length),fade=240;
+  for(let i=0;i<i16.length;i++)f32[i]=i16[i]/32768;
+  for(let i=0;i<fade&&i<f32.length;i++){f32[i]*=i/fade;f32[f32.length-1-i]*=i/fade;}
+  return {i16,f32};
+}
+function buildWav(chunks,sr){
+  let n=0;for(const c of chunks)n+=c.length;
+  const buf=new ArrayBuffer(44+n*2),v=new DataView(buf),w=(o,s)=>{for(let i=0;i<s.length;i++)v.setUint8(o+i,s.charCodeAt(i))};
+  w(0,'RIFF');v.setUint32(4,36+n*2,true);w(8,'WAVE');w(12,'fmt ');v.setUint32(16,16,true);v.setUint16(20,1,true);v.setUint16(22,1,true);
+  v.setUint32(24,sr,true);v.setUint32(28,sr*2,true);v.setUint16(32,2,true);v.setUint16(34,16,true);w(36,'data');v.setUint32(40,n*2,true);
+  let o=44;for(const c of chunks)for(let i=0;i<c.length;i++){v.setInt16(o,c[i],true);o+=2;}
+  return new Blob([buf],{type:'audio/wav'});
+}
+const blobToB64=blob=>new Promise(r=>{const fr=new FileReader();fr.onloadend=()=>r(fr.result.split(',')[1]);fr.readAsDataURL(blob)});
+const fmtT=s=>{if(!isFinite(s)||s<0)s=0;s=Math.floor(s);return `${Math.floor(s/60)}:${String(s%60).padStart(2,'0')}`};
+const fmtDur=s=>{s=Math.max(0,Math.floor(s));const h=s/3600|0,m=(s%3600/60)|0,x=s%60;return [h,m,x].map(n=>String(n).padStart(2,'0')).join(':')};
+const waveChars=n=>{const c='▁▂▃▄▅▆▇';let s='';for(let i=0;i<n;i++){const v=(Math.sin(i*.7)+Math.sin(i*.31+1.3)+Math.sin(i*1.9+.4))/3;s+=c[Math.max(0,Math.min(c.length-1,Math.floor((v+1)/2*c.length)))];}return s};
+
+function attachProg(audio,el,durHint,initPlayed){
+  const raw=el.textContent;if(!raw)return;
+  el.innerHTML='';el.classList.add('prog');
+  const chars=[...raw],wts=chars.map(c=>/\s/.test(c)?.2:/[,、;；:：，]/.test(c)?.5:/[.。!！?？…—]/.test(c)?1:1);
+  const cum=[0];for(const w of wts)cum.push(cum[cum.length-1]+w);
+  const tot=cum[cum.length-1]||1;
+  const spans=chars.map(c=>{const s=document.createElement('span');s.className='tk';s.textContent=c;el.appendChild(s);return s});
+  if(initPlayed)spans.forEach(s=>s.classList.add('played'));
+  const dur=()=>isFinite(audio.duration)&&audio.duration>0?audio.duration:(durHint||0);
+  const tOf=i=>dur()*cum[i]/tot;
+  const findCur=t=>{const d=dur();if(d<=0)return -1;const tg=t/d*tot;let lo=0,hi=spans.length-1;while(lo<hi){const m=(lo+hi+1)>>1;cum[m]<=tg?lo=m:hi=m-1}return lo};
+  spans.forEach((s,i)=>s.onclick=()=>{if(!audio.paused&&s.classList.contains('cur')){audio.pause();return}audio.currentTime=tOf(i);if(audio.paused)audio.play()});
+  const upd=()=>{const t=audio.currentTime,ci=!audio.paused?findCur(t):-1;spans.forEach((s,i)=>{s.classList.toggle('played',tOf(i)<=t);s.classList.toggle('cur',i===ci)});el.classList.toggle('on',!audio.paused)};
+  ['timeupdate','seeked','play','pause'].forEach(e=>audio.addEventListener(e,upd));
+}
+function mkPlayer(src){
+  const w=document.createElement('div');w.className='player';
+  w.innerHTML='<button class="pp">▶</button><div class="bar-wrap"><div class="bar"></div></div><span class="time">0:00 / 0:00</span>';
+  const a=new Audio(src),pp=w.querySelector('.pp'),bw=w.querySelector('.bar-wrap'),b=w.querySelector('.bar'),tm=w.querySelector('.time');
+  pp.onclick=()=>a.paused?a.play():a.pause();
+  a.onplay=()=>pp.textContent='❚❚';a.onpause=()=>pp.textContent='▶';
+  a.ontimeupdate=()=>{const d=a.duration||0;b.style.width=d?(a.currentTime/d*100)+'%':'0';tm.textContent=`${fmtT(a.currentTime)} / ${fmtT(d)}`};
+  a.onended=()=>b.style.width='0';a.onloadedmetadata=()=>tm.textContent=`0:00 / ${fmtT(a.duration)}`;
+  bw.onclick=e=>{const r=bw.getBoundingClientRect();if(a.duration)a.currentTime=(e.clientX-r.left)/r.width*a.duration};
+  return w;
+}
+
+let msgN=0;
+function addMsg(root,role,opts={}){
+  const d=document.createElement('div');d.className='msg '+role;msgN++;
+  const ts=new Date().toLocaleTimeString('en-US',{hour12:false});
+  d.innerHTML=`<div class="meta">[${msgN}] ${role} @ ${ts}</div><div class="body"></div>`;
+  const body=d.querySelector('.body');
+  if(opts.image){const i=document.createElement('img');i.className='img';i.src=opts.image;body.appendChild(i)}
+  if(opts.audioUrl){
+    if(role==='user'&&opts.audioDuration){
+      const n=Math.max(14,Math.min(80,Math.round(opts.audioDuration/.15)));
+      const v=document.createElement('div');v.className='text voice';v.textContent=waveChars(n);body.appendChild(v);
+      attachProg(new Audio(opts.audioUrl),v,opts.audioDuration);
+    } else body.appendChild(mkPlayer(opts.audioUrl));
+  }
+  if(opts.text!==undefined){const t=document.createElement('div');t.className='text';t.textContent=opts.text;body.appendChild(t)}
+  root.appendChild(d);root.scrollTop=root.scrollHeight;
+  return d;
+}
+
+/* ------------- Chat mode ------------- */
+let chatHist=[],pendImg=null,pendAud=null,rec=null,recChunks=[],recording=false,gen=false,actx=null,nextT=0,aborter=null,chatPlaying=[],chatEpoch=0;
+function stopChatAudio(){for(const s of chatPlaying)try{s.stop()}catch{}chatPlaying=[];if(actx){actx.close().catch(()=>{});actx=null;nextT=0}}
+function stopChat(){chatEpoch++;if(aborter)try{aborter.abort()}catch{}stopChatAudio();gen=false;sendBtn.disabled=false;ts.textContent='idle'}
+
+function renderChips(){
+  att.innerHTML='';txt.disabled=!!pendAud;txt.placeholder=pendAud?'(voice attached — text disabled)':'Message MiniMind-O...';
+  if(pendImg){const c=document.createElement('span');c.className='chip';c.innerHTML=`<img src="${pendImg.url}"><span>image</span><span class="x">×</span>`;c.querySelector('.x').onclick=()=>{URL.revokeObjectURL(pendImg.url);pendImg=null;renderChips()};att.appendChild(c)}
+  if(pendAud){const c=document.createElement('span');c.className='chip';c.innerHTML=`<span class="play">▶ voice · ${pendAud.duration.toFixed(1)}s</span><span class="x">×</span>`;c.querySelector('.play').onclick=()=>new Audio(pendAud.url).play();c.querySelector('.x').onclick=()=>{URL.revokeObjectURL(pendAud.url);pendAud=null;renderChips()};att.appendChild(c)}
+}
+
+async function send(){
+  if(gen)return;
+  stopChatAudio();
+  const t=txt.value.trim();if(!t&&!pendImg&&!pendAud)return;
+  const o={};if(pendImg)o.image=pendImg.url;if(pendAud){o.audioUrl=pendAud.url;o.audioDuration=pendAud.duration}if(t)o.text=t;
+  addMsg(chat,'user',o);
+  const payload={text:t||'',audio:pendAud?.b64||null,image:pendImg?.b64||null,temperature:.7,max_tokens:512,history:chatHist,voice:voiceSel.value};
+  txt.value='';pendImg=null;pendAud=null;renderChips();
+  sendBtn.disabled=true;gen=true;stat.classList.remove('hidden');tt.textContent='—';ta.textContent='—';ts.textContent='sending';
+  const ad=addMsg(chat,'assistant',{text:'…'}),ab=ad.querySelector('.body'),at=ab.querySelector('.text');
+  let full='',segs=[],acc='';
+  try{
+    const myEpoch=++chatEpoch;
+    aborter=new AbortController();
+    const res=await fetch('/chat',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(payload),signal:aborter.signal});
+    const rd=res.body.getReader(),dec=new TextDecoder();ts.textContent='streaming';
+    while(true){
+      if(myEpoch!==chatEpoch)throw new DOMException('chat stopped','AbortError');
+      const {done,value}=await rd.read();if(done)break;
+      for(const ln of dec.decode(value).split('\n')){
+        if(myEpoch!==chatEpoch)throw new DOMException('chat stopped','AbortError');
+        if(!ln.startsWith('data: '))continue;
+        let d;try{d=JSON.parse(ln.slice(6))}catch{continue}
+        if(d.type==='user_prompt')chatHist.push({role:'user',content:d.content});
+        else if(d.type==='ttft'){if(d.text_ttft!==undefined)tt.textContent=d.text_ttft+' ms';if(d.audio_ttft!==undefined)ta.textContent=d.audio_ttft+' ms'}
+        else if(d.type==='text'){full+=d.content;at.textContent=full;chat.scrollTop=chat.scrollHeight}
+        else if(d.type==='pcm'){
+          acc+=d.c;
+          if(d.d){
+            if(!actx){actx=new AudioContext({sampleRate:24000});nextT=actx.currentTime+.1}
+            const {i16,f32}=decPcm(acc);
+            const buf=actx.createBuffer(1,f32.length,24000);buf.getChannelData(0).set(f32);
+            const s=actx.createBufferSource();s.buffer=buf;s.connect(actx.destination);
+            s.onended=()=>{chatPlaying=chatPlaying.filter(x=>x!==s)};
+            chatPlaying.push(s);
+            if(nextT<actx.currentTime)nextT=actx.currentTime;
+            s.start(nextT);nextT+=buf.duration-.01;segs.push(i16);acc='';
+          }
+        } else if(d.type==='error')at.textContent='⚠ '+d.content;
+      }
+    }
+    if(segs.length){const url=URL.createObjectURL(buildWav(segs,24000));attachProg(new Audio(url),at,null,true)}
+    if(full)chatHist.push({role:'assistant',content:full});
+    ts.textContent='done';
+  } catch(e){if(e.name!=='AbortError'){at.textContent='⚠ '+e.message;ts.textContent='error'}}
+  gen=false;sendBtn.disabled=false;aborter=null;
+}
+
+sendBtn.onclick=send;
+txt.addEventListener('keydown',e=>{if(e.key==='Enter'&&!e.shiftKey){e.preventDefault();send()}});
+const imgPop=$('imgPop'),fileBtn=$('fileBtn'),shotBtn=$('shotBtn'),snap=$('snap'),snapV=$('snapV'),snapOk=$('snapOk'),snapCancel=$('snapCancel');
+const cloneModal=$('cloneModal'),cloneText=$('cloneText'),cloneStat=$('cloneStat'),cloneRecBtn=$('cloneRec'),cloneCancelBtn=$('cloneCancel'),cloneName=$('cloneName');
+cloneText.textContent=CLONE_PROMPT;
+moreBtn.onclick=e=>{e.stopPropagation();morePop.classList.toggle('on')};
+morePop.onclick=e=>{
+  e.stopPropagation();
+  const del=e.target.closest('.vx');
+  if(del){const name=del.dataset.del;if(!confirm('确认删除 '+name+'？'))return;fetch('/delete_voice',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({name})}).then(r=>r.json()).then(d=>{if(!d.ok)throw new Error(d.error);loadVoices(voiceSel.value===name?'default':voiceSel.value)}).catch(e=>alert(e.message));return}
+  const add=e.target.closest('.vadd');
+  if(add){closeMore();stopChat();cloneModal.classList.add('on');if(!cloneName.value.trim())cloneName.value='voice_clone';cloneStat.textContent='请自然读完，不用刻意加速';return}
+  const item=e.target.closest('.vitem');
+  if(item&&item.dataset.voice){voiceSel.value=item.dataset.voice;renderVoiceList();closeMore()}
+};
+imgBtn.onclick=e=>{e.stopPropagation();imgPop.classList.toggle('on')};
+document.addEventListener('click',e=>{if(!imgPop.contains(e.target)&&e.target!==imgBtn)imgPop.classList.remove('on');if(!morePop.contains(e.target)&&e.target!==moreBtn)closeMore()});
+fileBtn.onclick=()=>{imgPop.classList.remove('on');imgIn.click()};
+const setImg=(b64,url)=>{if(pendImg)URL.revokeObjectURL(pendImg.url);pendImg={b64,url};renderChips()};
+imgIn.onchange=()=>{const f=imgIn.files?.[0];if(!f)return;const r=new FileReader();r.onloadend=()=>setImg(r.result.split(',')[1],URL.createObjectURL(f));r.readAsDataURL(f);imgIn.value=''};
+let snapStream=null;
+const closeSnap=()=>{snap.classList.remove('on');if(snapStream){snapStream.getTracks().forEach(t=>t.stop());snapStream=null}snapV.srcObject=null};
+shotBtn.onclick=async()=>{
+  imgPop.classList.remove('on');
+  try{snapStream=await navigator.mediaDevices.getUserMedia({video:{facingMode:'user',width:1280,height:720}});snapV.srcObject=snapStream;snap.classList.add('on')}
+  catch(e){alert('camera error: '+e.message)}
+};
+snapCancel.onclick=closeSnap;
+snapOk.onclick=()=>{
+  if(!snapV.videoWidth)return;
+  const cv=document.createElement('canvas');cv.width=snapV.videoWidth;cv.height=snapV.videoHeight;cv.getContext('2d').drawImage(snapV,0,0);
+  const dataUrl=cv.toDataURL('image/jpeg',.9),b64=dataUrl.split(',')[1];
+  cv.toBlob(blob=>setImg(b64,URL.createObjectURL(blob)),'image/jpeg',.9);
+  closeSnap();
+};
+const CLONE_MAX_SECS=6;
+let cloneRec=null,cloneChunks=[],cloneRecording=false,cloneBusy=false,cloneCanceled=false,cloneStream=null,cloneDeadline=0,cloneTimer=null;
+const stopCloneTimer=()=>{if(cloneTimer){clearInterval(cloneTimer);cloneTimer=null}};
+const resetClone=()=>{stopCloneTimer();if(cloneStream){cloneStream.getTracks().forEach(t=>t.stop());cloneStream=null}cloneRec=null;cloneChunks=[];cloneRecording=false;cloneBusy=false;cloneCanceled=false;cloneDeadline=0;cloneRecBtn.classList.remove('rec');cloneRecBtn.textContent='start';cloneRecBtn.disabled=false;cloneCancelBtn.disabled=false};
+const hideClone=()=>{cloneModal.classList.remove('on');resetClone();cloneStat.textContent='点击 start 开始录音'};
+cloneCancelBtn.onclick=()=>{if(cloneBusy)return;if(cloneRecording&&cloneRec){cloneCanceled=true;cloneStat.textContent='已取消';cloneRec.stop()}else hideClone()};
+cloneRecBtn.onclick=async()=>{
+  if(cloneBusy)return;
+  if(!cloneRecording){
+    try{
+      cloneStream=await navigator.mediaDevices.getUserMedia({audio:true});
+      cloneRec=new MediaRecorder(cloneStream);cloneChunks=[];cloneCanceled=false;
+      cloneRec.ondataavailable=e=>cloneChunks.push(e.data);
+      cloneRec.onstop=async()=>{
+        const canceled=cloneCanceled,blob=new Blob(cloneChunks,{type:(cloneRec&&cloneRec.mimeType)||'audio/webm'});
+        resetClone();
+        if(canceled){cloneModal.classList.remove('on');cloneStat.textContent='点击 start 开始录音';return}
+        cloneBusy=true;cloneRecBtn.disabled=true;cloneCancelBtn.disabled=true;cloneStat.textContent='正在提取音色…';
+        try{
+          const audio=await blobToB64(blob);
+          const res=await fetch('/clone_voice',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({audio,name:cloneName.value.trim()})});
+          const d=await res.json();
+          if(!res.ok||!d.ok)throw new Error(d.error||'clone failed');
+          await loadVoices(d.voice);
+          voiceSel.value=d.voice;
+          cloneStat.textContent='音色已生成，已自动切换到 '+d.voice;
+          setTimeout(()=>hideClone(),600);
+        }catch(e){cloneBusy=false;cloneRecBtn.disabled=false;cloneCancelBtn.disabled=false;cloneStat.textContent='失败：'+e.message}
+      };
+      cloneRec.start();cloneRecording=true;cloneRecBtn.classList.add('rec');cloneRecBtn.textContent='stop';cloneDeadline=Date.now()+CLONE_MAX_SECS*1000;
+      const tick=()=>{
+        const left=Math.max(0,Math.ceil((cloneDeadline-Date.now())/1000));
+        cloneStat.textContent=`正在录音，请读完这句话 · ${left}s`;
+        if(left<=0&&cloneRecording&&cloneRec){cloneBusy=true;cloneRecording=false;cloneRecBtn.disabled=true;cloneRecBtn.classList.remove('rec');cloneRecBtn.textContent='start';cloneStat.textContent='录音结束，正在上传…';cloneRec.stop()}
+      };
+      tick();
+      cloneTimer=setInterval(tick,200);
+    }catch(e){alert('mic error: '+e.message)}
+  }else{
+    stopCloneTimer();cloneBusy=true;cloneRecording=false;cloneRecBtn.disabled=true;cloneRecBtn.classList.remove('rec');cloneRecBtn.textContent='start';cloneStat.textContent='录音结束，正在上传…';cloneRec.stop();
+  }
+};
+recBtn.onclick=async()=>{
+  if(!recording){
+    try{
+      const st=await navigator.mediaDevices.getUserMedia({audio:true});
+      rec=new MediaRecorder(st);recChunks=[];const t0=Date.now();
+      rec.ondataavailable=e=>recChunks.push(e.data);
+      rec.onstop=()=>{
+        st.getTracks().forEach(t=>t.stop());
+        const blob=new Blob(recChunks,{type:'audio/webm'}),dur=(Date.now()-t0)/1000,r=new FileReader();
+        r.onloadend=()=>{const b64=r.result.split(',')[1],url=URL.createObjectURL(blob);if(pendAud)URL.revokeObjectURL(pendAud.url);pendAud={b64,url,duration:dur};renderChips()};
+        r.readAsDataURL(blob);
+      };
+      rec.start();recording=true;recBtn.classList.add('rec');
+    } catch(e){alert('mic error: '+e.message)}
+  } else {rec.stop();recording=false;recBtn.classList.remove('rec')}
+};
+
+/* ------------- Call mode ------------- */
+let callInit=false,callStop=null;
+callBtn.onclick=openCall;
+function openCall(){stopChat();document.body.classList.add('call');history.pushState({call:1},'','/call');if(!callInit){callInit=true;initCall()}}
+function closeCall(){document.body.classList.remove('call');if(callStop)callStop();if(location.pathname==='/call')history.pushState({},'','/')}
+window.addEventListener('popstate',()=>location.pathname==='/call'?openCall():closeCall());
+if(location.pathname==='/call')openCall();
+
+function initCall(){
+  const stT=$('stT'),upT=$('upT'),stIc=$('stIc'),lbl=$('lbl'),osvg=$('osvg'),core=$('core'),ripG=$('rip');
+  const ICONS={
+    idle:'<svg viewBox="0 0 24 24"><circle class="dot" cx="12" cy="12" r="3" fill="currentColor" stroke="none"/></svg>',
+    listen:'<svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="2.2" fill="currentColor" stroke="none"/><path class="w w3" d="M7 8.5a6 6 0 0 0 0 7"/><path class="w w3" d="M17 8.5a6 6 0 0 1 0 7"/><path class="w w2" d="M4 6a9.5 9.5 0 0 0 0 12"/><path class="w w2" d="M20 6a9.5 9.5 0 0 1 0 12"/></svg>',
+    gen:'<svg viewBox="0 0 24 24"><circle cx="12" cy="4" r="1.8" fill="currentColor" stroke="none"/><circle cx="18.93" cy="15.5" r="1.4" fill="currentColor" stroke="none" opacity=".65"/><circle cx="5.07" cy="15.5" r="1" fill="currentColor" stroke="none" opacity=".35"/></svg>',
+    speak:'<svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="2.2" fill="currentColor" stroke="none"/><path class="a1" d="M8 9a5 5 0 0 1 0 6"/><path class="a1" d="M16 9a5 5 0 0 0 0 6"/><path class="a2" d="M5 6.5a9 9 0 0 1 0 11"/><path class="a2" d="M19 6.5a9 9 0 0 0 0 11"/><path class="a3" d="M2 4a13 13 0 0 1 0 16"/><path class="a3" d="M22 4a13 13 0 0 0 0 16"/></svg>',
+    err:'<svg viewBox="0 0 24 24"><circle cx="12" cy="12" r="9"/><line x1="8.5" y1="8.5" x2="15.5" y2="15.5"/><line x1="15.5" y1="8.5" x2="8.5" y2="15.5"/></svg>',
+  };
+  const cchat=$('cchat'),lat=$('lat'),lmE=$('lmE'),ltE=$('ltE');
+  const camBtn=$('camBtn'),logBtn=$('logBtn'),endBtn=$('endBtn'),orb=$('orb'),camV=$('camV');
+  const pill=$('pill'),connT=$('connT'),tcnt=$('tcnt'),trans=$('trans'),closeBtn=$('closeBtn');
+  const blobs=[$('b0'),$('b1'),$('b2')];
+
+  let ac,ana,mic,camStream=null,camOn=false,hist=[],nextT=0;
+  let ws=null,wsActive=false,proc=null,rtc=null,rtSrc=null,playing=[],turns=0,mn=0;
+  let curMsg=null,curText='',segs=[],start=0,ripT=0,sVol=0,asstPlay=false,raf=0;
+  let spk=false,userBuf=[];
+  const RT_SAMPLE_RATE=16000,RT_BLOCK=4096;
+  const resamplePcm=(input,inRate,outRate)=>{
+    if(!input.length||inRate===outRate)return input;
+    const outLen=Math.max(1,Math.round(input.length*outRate/inRate));
+    const out=new Float32Array(outLen);
+    const step=inRate/outRate;
+    for(let i=0;i<outLen;i++){
+      const pos=i*step,idx=Math.floor(pos),frac=pos-idx;
+      const a=input[idx]||0;
+      const b=input[Math.min(idx+1,input.length-1)]||a;
+      out[i]=a+(b-a)*frac;
+    }
+    return out;
+  };
+  const floatToI16=input=>{const out=new Int16Array(input.length);for(let i=0;i<input.length;i++)out[i]=Math.max(-32768,Math.min(32767,input[i]*32768));return out};
+  const flushUser=()=>{
+    if(!userBuf.length)return;
+    const dur=userBuf.reduce((a,c)=>a+c.length,0)/16000;
+    const url=URL.createObjectURL(buildWav(userBuf,16000));
+    userBuf=[];
+    const d=addCM('user'),body=d.querySelector('.body'),n=Math.max(14,Math.min(80,Math.round(dur/.15)));
+    const v=document.createElement('div');v.className='text voice';v.textContent=waveChars(n);body.appendChild(v);
+    attachProg(new Audio(url),v,dur);
+    setLatest('user',`[voice · ${dur.toFixed(1)}s]`);
+  };
+
+  const N=32,NC=2,NR=5;let shT=0;
+  const LAY=[{c:56,r:112,a:1.8,sp:.0005,ph:.6,vc:34,vr:18},{c:80,r:124,a:2.6,sp:.0008,ph:1.3,vc:52,vr:26},{c:104,r:138,a:3.6,sp:.0011,ph:2.1,vc:74,vr:34}];
+
+  const catmull=pts=>{let d=`M ${pts[0][0].toFixed(2)} ${pts[0][1].toFixed(2)} `;const n=pts.length;for(let i=0;i<n;i++){const p0=pts[(i-1+n)%n],p1=pts[i],p2=pts[(i+1)%n],p3=pts[(i+2)%n];const c1x=p1[0]+(p2[0]-p0[0])/6,c1y=p1[1]+(p2[1]-p0[1])/6,c2x=p2[0]-(p3[0]-p1[0])/6,c2y=p2[1]-(p3[1]-p1[1])/6;d+=`C ${c1x.toFixed(2)} ${c1y.toFixed(2)}, ${c2x.toFixed(2)} ${c2y.toFixed(2)}, ${p2[0].toFixed(2)} ${p2[1].toFixed(2)} `}return d+'Z'};
+  const setSt=(cls,t)=>{osvg.setAttribute('class',cls||'');lbl.className='lbl '+(cls||'');stT.textContent=t;const key=cls||(/error/i.test(t)?'err':'idle');stIc.innerHTML=ICONS[key]||ICONS.idle};
+  const setConn=(live,t)=>{pill.className='pill'+(live?' live':'');connT.textContent=t};
+  const ripple=k=>{const now=performance.now();if(now-ripT<550)return;ripT=now;const c=document.createElementNS('http://www.w3.org/2000/svg','circle');c.setAttribute('cx',0);c.setAttribute('cy',0);c.setAttribute('r',85);c.setAttribute('class','ripple '+k);ripG.appendChild(c);c.addEventListener('animationend',()=>c.remove())};
+  const upTimer=setInterval(()=>{if(start)upT.textContent=fmtDur((Date.now()-start)/1000)},1000);
+
+  const setLatest=(role,t)=>{if(!t)return;lat.classList.remove('empty','user','assistant');lat.classList.add(role);const ts=new Date().toLocaleTimeString('en-US',{hour12:false});lmE.textContent=`latest · ${role} @ ${ts}`;ltE.textContent=t;lat.scrollTop=0};
+  const addCM=role=>{const d=document.createElement('div');d.className='msg '+role;mn++;const ts=new Date().toLocaleTimeString('en-US',{hour12:false});d.innerHTML=`<div class="meta">[${mn}] ${role} @ ${ts}</div><div class="body"></div>`;cchat.appendChild(d);cchat.scrollTop=cchat.scrollHeight;return d};
+  const playPcm=b64=>{const {i16,f32}=decPcm(b64);segs.push(i16);const buf=ac.createBuffer(1,f32.length,24000);buf.getChannelData(0).set(f32);const s=ac.createBufferSource();s.buffer=buf;s.connect(ac.destination);s.onended=()=>{playing=playing.filter(x=>x!==s);if(!playing.length){asstPlay=false;if(!curMsg)setSt('','idle')}};playing.push(s);asstPlay=true;if(nextT<ac.currentTime)nextT=ac.currentTime;s.start(nextT);nextT+=buf.duration-.01};
+  const stopPcm=()=>{for(const s of playing)try{s.stop()}catch{}playing=[];asstPlay=false;if(ac)nextT=ac.currentTime};
+  const capCam=()=>{if(!camOn||!camV.videoWidth)return null;const cv=document.createElement('canvas');cv.width=camV.videoWidth;cv.height=camV.videoHeight;cv.getContext('2d').drawImage(camV,0,0);return cv.toDataURL('image/jpeg',.8).split(',')[1]};
+  const sendCtx=()=>{if(!ws||ws.readyState!==1)return;ws.send(JSON.stringify({type:'context',history:hist,image:camOn?capCam():null,voice:callVoiceSel.value}))};
+  window.callSendCtx=sendCtx;
+
+  function render(){
+    const now=performance.now();let vol=0,fq=null;
+    if(ana){const d=new Uint8Array(ana.frequencyBinCount);ana.getByteFrequencyData(d);fq=d;let s=0;for(let i=0;i<d.length;i++)s+=d[i];vol=s/d.length/255}
+    if(asstPlay)vol=Math.max(vol,.6+.25*Math.sin(now*.013));
+    sVol=sVol*.7+vol*.3;
+    const tgt=camOn?1:0;shT+=(tgt-shT)*.08;
+    const shN=NC+(NR-NC)*shT,inv=-1/shN;
+    for(let li=0;li<LAY.length;li++){
+      const L=LAY[li],hw=L.c+(L.r-L.c)*shT,hh=hw,vs=L.vc+(L.vr-L.vc)*shT;
+      const pts=new Array(N);
+      for(let i=0;i<N;i++){
+        const ang=(i/N)*Math.PI*2,cA=Math.cos(ang),sA=Math.sin(ang);
+        const n=Math.sin(now*L.sp+i*L.ph)+.55*Math.sin(now*L.sp*1.7+i*L.ph*.7+li);
+        let fv=0;if(fq){const idx=Math.floor((i/N)*Math.min(64,fq.length))+li*4;fv=(fq[idx%fq.length]||0)/255}
+        if(asstPlay)fv=Math.max(fv,.5+.45*Math.abs(Math.sin(now*.011+i*.55+li*.9)));
+        const g=Math.min(1,fv*.65+sVol*.55);
+        const tx=Math.pow(Math.abs(cA)/hw,shN),ty=Math.pow(Math.abs(sA)/hh,shN);
+        const r=Math.pow(tx+ty,inv)+n*L.a+g*vs;
+        pts[i]=[cA*r,sA*r];
+      }
+      blobs[li].setAttribute('d',catmull(pts));
+    }
+    core.setAttribute('r',(3+sVol*10).toFixed(1));
+    raf=requestAnimationFrame(render);
+  }
+  raf=requestAnimationFrame(render);
+
+  async function initMic(){
+    try{mic=await navigator.mediaDevices.getUserMedia({audio:true});ac=new AudioContext();const s=ac.createMediaStreamSource(mic);ana=ac.createAnalyser();ana.fftSize=256;s.connect(ana);setSt('','connecting');setConn(false,'connecting')}
+    catch(e){setSt('','mic error');setConn(false,'mic error')}
+  }
+
+  camBtn.onclick=async()=>{
+    if(camOn){if(camStream){camStream.getTracks().forEach(t=>t.stop());camStream=null}camV.srcObject=null;orb.classList.remove('camon');camBtn.classList.remove('active');camOn=false}
+    else{try{camStream=await navigator.mediaDevices.getUserMedia({video:{facingMode:'user',width:640,height:480}});camV.srcObject=camStream;orb.classList.add('camon');camBtn.classList.add('active');camOn=true}catch{setSt('','cam error')}}
+    if(wsActive)sendCtx();
+  };
+  logBtn.onclick=()=>{trans.classList.toggle('on');logBtn.classList.toggle('active',trans.classList.contains('on'))};
+  closeBtn.onclick=()=>{trans.classList.remove('on');logBtn.classList.remove('active')};
+  endBtn.onclick=closeCall;
+
+  async function startRT(){
+    wsActive=true;if(!mic)await initMic();
+    rtc=new AudioContext();rtSrc=rtc.createMediaStreamSource(mic);proc=rtc.createScriptProcessor(4096,1,1);
+    let rtAcc=new Float32Array(0);
+    const pr=location.protocol==='https:'?'wss':'ws';
+    ws=new WebSocket(`${pr}://${location.host}/ws/realtime`);ws.binaryType='arraybuffer';
+    ws.onmessage=e=>{
+      const m=JSON.parse(e.data);
+      if(m.type==='vad'){
+        if(m.speaking&&!spk){sendCtx();stopPcm();ripple('user');userBuf=[]}
+        if(!m.speaking&&spk)flushUser();
+        spk=m.speaking;
+        if(m.speaking)setSt('listen','listening');else if(!asstPlay&&!curMsg)setSt('','idle');
+      } else if(m.type==='generating'){
+        setSt('gen','thinking');if(spk){flushUser();spk=false}
+        turns++;tcnt.textContent=turns+(turns===1?' turn':' turns');
+        curText='';segs=[];
+        const d=addCM('assistant'),body=d.querySelector('.body'),t=document.createElement('div');t.className='text';body.appendChild(t);
+        curMsg={div:d,body,textEl:t};stopPcm();
+      } else if(m.type==='text'){
+        curText+=m.content;if(curMsg)curMsg.textEl.textContent=curText;
+        setLatest('assistant',curText);cchat.scrollTop=cchat.scrollHeight;
+      } else if(m.type==='pcm'){setSt('speak','speaking');ripple('assistant');playPcm(m.data)}
+      else if(m.type==='done'){
+        if(curText)hist.push({role:'assistant',content:curText});
+        if(m.interrupted)stopPcm();
+        if(curMsg&&segs.length){const url=URL.createObjectURL(buildWav(segs,24000));attachProg(new Audio(url),curMsg.textEl,null,true)}
+        curMsg=null;if(!asstPlay)setSt('','idle');
+      }
+    };
+    ws.onclose=()=>{if(wsActive)stopRT()};
+    ws.onopen=()=>{
+      start=Date.now();setConn(true,'live');setSt('','idle');sendCtx();
+      proc.onaudioprocess=e=>{
+        if(ws.readyState!==1)return;
+        const rs=resamplePcm(e.inputBuffer.getChannelData(0),e.inputBuffer.sampleRate,RT_SAMPLE_RATE);
+        const merged=new Float32Array(rtAcc.length+rs.length);merged.set(rtAcc);merged.set(rs,rtAcc.length);
+        let off=0;
+        for(;merged.length-off>=RT_BLOCK;off+=RT_BLOCK){
+          const i16=floatToI16(merged.subarray(off,off+RT_BLOCK));
+          ws.send(i16.buffer);
+          if(spk)userBuf.push(new Int16Array(i16));
+        }
+        rtAcc=merged.slice(off);
+      };
+      rtSrc.connect(proc);proc.connect(rtc.destination);
+    };
+  }
+  function stopRT(){wsActive=false;if(ws){ws.close();ws=null}if(proc){proc.disconnect();proc.onaudioprocess=null;proc=null}if(rtSrc){rtSrc.disconnect();rtSrc=null}if(rtc){rtc.close();rtc=null}stopPcm();setSt('','disconnected');setConn(false,'offline')}
+
+  callStop=()=>{
+    stopRT();
+    window.callSendCtx=null;
+    if(camStream){camStream.getTracks().forEach(t=>t.stop());camStream=null;camOn=false;orb.classList.remove('camon');camV.srcObject=null;camBtn.classList.remove('active')}
+    if(mic){mic.getTracks().forEach(t=>t.stop());mic=null}
+    if(ac){try{ac.close()}catch{}ac=null}
+    if(raf)cancelAnimationFrame(raf);
+    clearInterval(upTimer);callInit=false;callStop=null;
+  };
+  initMic().then(startRT);
+}
+</script>
+</body>
+</html>

pyproject.toml

@@ -1,10 +1,35 @@
 [project]
 name = "omni"
 version = "0.1.0"
-description = "Add your description here"
+description = "MiniMind LLM (and future multimodal) training & inference framework, structured as an installable package."
 readme = "README.md"
 requires-python = ">=3.12"
-dependencies = []
+dependencies = [
+    "torch>=2.6.0",
+    "transformers>=4.57.6",
+    "datasets>=3.6.0",
+    "tokenizers>=0.20.0",
+    "tiktoken>=0.10.0",
+    "jinja2>=3.1.2",
+    "numpy>=1.26.4",
+    "safetensors>=0.5.0",
+    "accelerate>=1.0.0",
+    "sentencepiece>=0.2.0",
+    "jieba>=0.42.1",
+    "pydantic>=2.11.5",
+    "rich>=13.7.1",
+    "regex",
+]
+
+[project.optional-dependencies]
+rl = ["swanlab>=0.7.11", "wandb>=0.18.3"]
+serve = ["fastapi>=3.0.3", "uvicorn", "pydantic>=2.11.5", "openai>=1.59.6", "flask>=3.0.3", "flask-cors>=4.0.0"]
+demo = ["streamlit>=1.50.0"]
+all = ["swanlab>=0.7.11", "wandb>=0.18.3", "fastapi>=3.0.3", "uvicorn", "openai>=1.59.6", "flask>=3.0.3", "flask-cors>=4.0.0", "streamlit>=1.50.0", "sglang"]
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/omni/core/__init__.py

@@ -0,0 +1,18 @@
+from omni.core.norm import RMSNorm
+from omni.core.rope import precompute_freqs_cis, apply_rotary_pos_emb, repeat_kv
+from omni.core.attention import Attention
+from omni.core.mlp import FeedForward, MOEFeedForward
+from omni.core.block import MiniMindBlock
+from omni.core.model import MiniMindModel
+
+__all__ = [
+    "RMSNorm",
+    "precompute_freqs_cis",
+    "apply_rotary_pos_emb",
+    "repeat_kv",
+    "Attention",
+    "FeedForward",
+    "MOEFeedForward",
+    "MiniMindBlock",
+    "MiniMindModel",
+]

src/omni/core/attention.py

@@ -0,0 +1,55 @@
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from omni.core.norm import RMSNorm
+from omni.core.rope import apply_rotary_pos_emb, repeat_kv
+
+
+class Attention(nn.Module):
+    def __init__(self, config: "MiniMindConfig"):
+        super().__init__()
+        self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
+        self.n_local_heads = config.num_attention_heads
+        self.n_local_kv_heads = self.num_key_value_heads
+        self.n_rep = self.n_local_heads // self.n_local_kv_heads
+        self.head_dim = config.head_dim
+        self.is_causal = True
+        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.dropout = config.dropout
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn
+
+    def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
+        bsz, seq_len, _ = x.shape
+        xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+        xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
+        xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
+        xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
+        xq, xk = self.q_norm(xq), self.k_norm(xk)
+        cos, sin = position_embeddings
+        xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
+        if past_key_value is not None:
+            xk = torch.cat([past_key_value[0], xk], dim=1)
+            xv = torch.cat([past_key_value[1], xv], dim=1)
+        past_kv = (xk, xv) if use_cache else None
+        xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
+        if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
+            output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)
+        else:
+            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            if self.is_causal:
+                scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
+            if attention_mask is not None:
+                scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
+            output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
+        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
+        output = self.resid_dropout(self.o_proj(output))
+        return output, past_kv

src/omni/core/block.py

@@ -0,0 +1,24 @@
+from torch import nn
+
+from omni.core.norm import RMSNorm
+from omni.core.attention import Attention
+from omni.core.mlp import FeedForward, MOEFeedForward
+
+
+class MiniMindBlock(nn.Module):
+    def __init__(self, layer_id: int, config: "MiniMindConfig"):
+        super().__init__()
+        self.self_attn = Attention(config)
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
+
+    def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
+        residual = hidden_states
+        hidden_states, present_key_value = self.self_attn(
+            self.input_layernorm(hidden_states), position_embeddings,
+            past_key_value, use_cache, attention_mask
+        )
+        hidden_states += residual
+        hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
+        return hidden_states, present_key_value

src/omni/core/mlp.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.activations import ACT2FN
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config: "MiniMindConfig", intermediate_size: int = None):
+        super().__init__()
+        intermediate_size = intermediate_size or config.intermediate_size
+        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class MOEFeedForward(nn.Module):
+    def __init__(self, config: "MiniMindConfig"):
+        super().__init__()
+        self.config = config
+        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
+        self.experts = nn.ModuleList([FeedForward(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)])
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        batch_size, seq_len, hidden_dim = x.shape
+        x_flat = x.view(-1, hidden_dim)
+        scores = F.softmax(self.gate(x_flat), dim=-1)
+        topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
+        if self.config.norm_topk_prob:
+            topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
+        y = torch.zeros_like(x_flat)
+        for i, expert in enumerate(self.experts):
+            mask = (topk_idx == i)
+            if mask.any():
+                token_idx = mask.any(dim=-1).nonzero().flatten()
+                weight = topk_weight[mask].view(-1, 1)
+                y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
+            elif self.training:
+                y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
+        if self.training and self.config.router_aux_loss_coef > 0:
+            load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)
+            self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
+        else:
+            self.aux_loss = scores.new_zeros(1).squeeze()
+        return y.view(batch_size, seq_len, hidden_dim)

src/omni/core/model.py

@@ -0,0 +1,45 @@
+from torch import nn
+
+from omni.core.norm import RMSNorm
+from omni.core.rope import precompute_freqs_cis
+from omni.core.block import MiniMindBlock
+from omni.core.mlp import MOEFeedForward
+
+
+class MiniMindModel(nn.Module):
+    def __init__(self, config: "MiniMindConfig"):
+        super().__init__()
+        self.config = config
+        self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.dropout)
+        self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling)
+        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
+        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+
+    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
+        batch_size, seq_length = input_ids.shape
+        if hasattr(past_key_values, 'layers'):
+            past_key_values = None
+        past_key_values = past_key_values or [None] * len(self.layers)
+        start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
+        hidden_states = self.dropout(self.embed_tokens(input_ids))
+        if self.freqs_cos[0, 0] == 0:
+            freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
+            self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
+        position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length])
+        presents = []
+        for layer, past_key_value in zip(self.layers, past_key_values):
+            hidden_states, present = layer(
+                hidden_states,
+                position_embeddings,
+                past_key_value=past_key_value,
+                use_cache=use_cache,
+                attention_mask=attention_mask
+            )
+            presents.append(present)
+        hidden_states = self.norm(hidden_states)
+        aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
+        return hidden_states, presents, aux_loss

src/omni/core/norm.py

@@ -0,0 +1,15 @@
+import torch
+from torch import nn
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        return (self.weight * self.norm(x.float())).type_as(x)

src/omni/core/rope.py

@@ -0,0 +1,37 @@
+import math
+import torch
+
+
+def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: dict = None):
+    freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0
+    if rope_scaling is not None:  # YaRN
+        orig_max, factor, beta_fast, beta_slow, attn_factor = (
+            rope_scaling.get("original_max_position_embeddings", 2048), rope_scaling.get("factor", 16),
+            rope_scaling.get("beta_fast", 32.0), rope_scaling.get("beta_slow", 1.0), rope_scaling.get("attention_factor", 1.0)
+        )
+        if end / orig_max > 1.0:
+            inv_dim = lambda b: (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base))
+            low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim // 2 - 1)
+            ramp = torch.clamp((torch.arange(dim // 2, device=freqs.device).float() - low) / max(high - low, 0.001), 0, 1)
+            freqs = freqs * (1 - ramp + ramp / factor)
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor
+    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor
+    return freqs_cos, freqs_sin
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    def rotate_half(x):
+        return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
+
+    q_embed = ((q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))).to(q.dtype)
+    k_embed = ((k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))).to(k.dtype)
+    return q_embed, k_embed
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    bs, slen, num_key_value_heads, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    return (x[:, :, :, None, :].expand(bs, slen, num_key_value_heads, n_rep, head_dim).reshape(bs, slen, num_key_value_heads * n_rep, head_dim))

src/omni/datasets/__init__.py

@@ -0,0 +1 @@
+from omni.datasets.lm_dataset import *  # noqa: F401,F403

src/omni/datasets/lm_dataset.py

@@ -0,0 +1,606 @@
+from torch.utils.data import Dataset, DataLoader
+import torch
+import io
+import json
+import os
+import random
+from datasets import load_dataset, Features, Sequence, Value
+from PIL import Image
+from datasets import Dataset as HFDataset
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+from omni.models import MiniMindVLM
+
+
+def pre_processing_chat(conversations, add_system_ratio=0.2):
+    if any(conv.get('tools') for conv in conversations):
+        return conversations
+
+    SYSTEM_PROMPTS = [
+        "你是一个知识丰富的AI，尽力为用户提供准确的信息。",
+        "你是minimind，一个小巧但有用的语言模型。",
+        "你是一个专业的AI助手，请提供有价值的回答。",
+        "你是minimind，请尽力帮助用户解决问题。",
+        "你是一个可靠的AI，请给出准确的回答。",
+        "You are a helpful AI assistant.",
+        "You are minimind, a lightweight intelligent assistant.",
+        "You are a friendly chatbot. Please answer the user's questions carefully.",
+        "You are a knowledgeable AI. Try your best to provide accurate information.",
+        "You are minimind, a small but useful language model."
+    ]
+    if conversations[0].get('role') != 'system':
+        if random.random() < add_system_ratio:
+            return [{'role': 'system', 'content': random.choice(SYSTEM_PROMPTS)}] + conversations
+    return conversations
+
+
+def post_processing_chat(prompt_content, empty_think_ratio=0.2):
+    if '<think>\n\n</think>\n\n' in prompt_content and random.random() > empty_think_ratio:
+        prompt_content = prompt_content.replace('<think>\n\n</think>\n\n', '')
+    return prompt_content
+
+
+class PretrainDataset(Dataset):
+    def __init__(self, data_path, tokenizer, max_length=512):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.samples = load_dataset('json', data_files=data_path, split='train')
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        tokens = self.tokenizer(str(sample['text']), add_special_tokens=False, max_length=self.max_length - 2, truncation=True).input_ids
+        tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]
+        input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens))
+        input_ids = torch.tensor(input_ids, dtype=torch.long)
+        labels = input_ids.clone()
+        labels[input_ids == self.tokenizer.pad_token_id] = -100
+        return input_ids, labels
+
+
+class SFTDataset(Dataset):
+    def __init__(self, jsonl_path, tokenizer, max_length=1024):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        features = Features({'conversations': [{'role': Value('string'), 'content': Value('string'), 'reasoning_content': Value('string'), 'tools': Value('string'), 'tool_calls': Value('string')}]})
+        self.samples = load_dataset('json', data_files=jsonl_path, split='train', features=features)
+        self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids
+        self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
+
+    def __len__(self):
+        return len(self.samples)
+
+    def create_chat_prompt(self, conversations):
+        messages = []
+        tools = None
+        for message in conversations:
+            message = dict(message)
+            if message.get("role") == "system" and message.get("tools"):
+                tools = json.loads(message["tools"]) if isinstance(message["tools"], str) else message["tools"]
+            if message.get("tool_calls") and isinstance(message["tool_calls"], str):
+                message["tool_calls"] = json.loads(message["tool_calls"])
+            messages.append(message)
+        return self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+            tools=tools
+        )
+
+    def generate_labels(self, input_ids):
+        labels = [-100] * len(input_ids)
+        i = 0
+        while i < len(input_ids):
+            if input_ids[i:i + len(self.bos_id)] == self.bos_id:
+                start = i + len(self.bos_id)
+                end = start
+                while end < len(input_ids):
+                    if input_ids[end:end + len(self.eos_id)] == self.eos_id:
+                        break
+                    end += 1
+                for j in range(start, min(end + len(self.eos_id), self.max_length)):
+                    labels[j] = input_ids[j]
+                i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
+            else:
+                i += 1
+        return labels
+
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        conversations = pre_processing_chat(sample['conversations'])
+        prompt = self.create_chat_prompt(conversations)
+        prompt = post_processing_chat(prompt)
+        input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
+        input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
+        labels = self.generate_labels(input_ids)
+        return torch.tensor(input_ids, dtype=torch.long), torch.tensor(labels, dtype=torch.long)
+
+
+class DPODataset(Dataset):
+    def __init__(self, file_path, tokenizer, max_length=4096):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+        self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids
+        self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
+        self.samples = load_dataset('json', data_files=file_path, split='train')
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        chosen = sample['chosen']
+        rejected = sample['rejected']
+        chosen_prompt = self.tokenizer.apply_chat_template(
+            chosen, tokenize=False, add_generation_prompt=False
+        )
+        chosen_prompt = post_processing_chat(chosen_prompt)
+
+        rejected_prompt = self.tokenizer.apply_chat_template(
+            rejected, tokenize=False, add_generation_prompt=False
+        )
+        rejected_prompt = post_processing_chat(rejected_prompt)
+        chosen_encoding = self.tokenizer(
+            chosen_prompt, truncation=True, max_length=self.max_length, padding='max_length'
+        )
+        rejected_encoding = self.tokenizer(
+            rejected_prompt, truncation=True, max_length=self.max_length, padding='max_length'
+        )
+
+        chosen_input_ids = chosen_encoding['input_ids']
+        chosen_loss_mask = self.generate_loss_mask(chosen_input_ids)
+
+        rejected_input_ids = rejected_encoding['input_ids']
+        rejected_loss_mask = self.generate_loss_mask(rejected_input_ids)
+        x_chosen = torch.tensor(chosen_input_ids[:-1], dtype=torch.long)
+        y_chosen = torch.tensor(chosen_input_ids[1:], dtype=torch.long)
+        mask_chosen = torch.tensor(chosen_loss_mask[1:], dtype=torch.long)
+        x_rejected = torch.tensor(rejected_input_ids[:-1], dtype=torch.long)
+        y_rejected = torch.tensor(rejected_input_ids[1:], dtype=torch.long)
+        mask_rejected = torch.tensor(rejected_loss_mask[1:], dtype=torch.long)
+
+        return {
+            'x_chosen': x_chosen,
+            'y_chosen': y_chosen,
+            'mask_chosen': mask_chosen,
+            'x_rejected': x_rejected,
+            'y_rejected': y_rejected,
+            'mask_rejected': mask_rejected
+        }
+
+    def generate_loss_mask(self, input_ids):
+        loss_mask = [0] * len(input_ids)
+        i = 0
+        while i < len(input_ids):
+            if input_ids[i:i + len(self.bos_id)] == self.bos_id:
+                start = i + len(self.bos_id)
+                end = start
+                while end < len(input_ids):
+                    if input_ids[end:end + len(self.eos_id)] == self.eos_id:
+                        break
+                    end += 1
+                for j in range(start, min(end + len(self.eos_id), self.max_length)):
+                    loss_mask[j] = 1
+                i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
+            else:
+                i += 1
+        return loss_mask
+
+
+class RLAIFDataset(Dataset):
+    def __init__(self, jsonl_path, tokenizer, max_length=1024, thinking_ratio=0.5):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.thinking_ratio = thinking_ratio
+        self.samples = load_dataset('json', data_files=jsonl_path, split='train')
+        self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant', add_special_tokens=False).input_ids
+        self.eos_id = tokenizer(f'{tokenizer.eos_token}', add_special_tokens=False).input_ids
+
+    def __len__(self):
+        return len(self.samples)
+
+    def create_chat_prompt(self, conversations):
+        conversations = pre_processing_chat(conversations)
+        use_thinking = random.random() < self.thinking_ratio
+        return self.tokenizer.apply_chat_template(
+            conversations[:-1],
+            tokenize=False,
+            open_thinking=use_thinking,
+            add_generation_prompt=True
+        )
+
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        prompt = self.create_chat_prompt(sample['conversations'])
+
+        return {
+            'prompt': prompt,
+            'answer': ""
+        }
+
+
+class AgentRLDataset(Dataset):
+    def __init__(self, jsonl_path, tokenizer, max_length=1024):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.samples = []
+        with open(jsonl_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                self.samples.append(json.loads(line.strip()))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def parse_conversations(self, conversations):
+        messages = []
+        tools = None
+        for message in conversations:
+            message = dict(message)
+            if message.get("role") == "system" and message.get("tools"):
+                tools = json.loads(message["tools"]) if isinstance(message["tools"], str) else message["tools"]
+            messages.append(message)
+        return messages[:-1], tools
+
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        messages, tools = self.parse_conversations(sample['conversations'])
+        return {'messages': messages, 'tools': tools, 'gt': sample['gt']}
+
+
+class VLMDataset(Dataset):
+    def __init__(self, parquet_path, tokenizer, preprocess=None, max_length=512, image_special_token='<|image_pad|>', image_token_len=64):
+        super().__init__()
+        self.dataset = HFDataset.from_parquet(parquet_path)
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.preprocess = preprocess
+        self.image_special_token = image_special_token * image_token_len
+        self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids
+        self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def create_chat_prompt(self, conversations):
+        messages = []
+        for turn in conversations:
+            content = turn['content'].replace('<image>', self.image_special_token) if turn.get('role') != 'system' else turn['content']
+            messages.append({"role": turn['role'], "content": content})
+        tools = conversations[0]["functions"] if (conversations and conversations[0]["role"] == "system" and conversations[0].get("functions")) else None
+        return self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+            tools=tools
+        )
+
+    def generate_labels(self, input_ids):
+        labels = [-100] * len(input_ids)
+        i = 0
+        while i < len(input_ids):
+            if input_ids[i:i + len(self.bos_id)] == self.bos_id:
+                start = i + len(self.bos_id)
+                end = start
+                while end < len(input_ids):
+                    if input_ids[end:end + len(self.eos_id)] == self.eos_id:
+                        break
+                    end += 1
+                for j in range(start, min(end + len(self.eos_id), self.max_length)):
+                    labels[j] = input_ids[j]
+                i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
+            else:
+                i += 1
+        return labels
+
+    def __getitem__(self, index: int):
+        row = self.dataset[index]
+        conversations = json.loads(row['conversations']) if isinstance(row['conversations'], str) else row['conversations']
+        image_bytes = row['image_bytes']
+        if not isinstance(image_bytes, list): image_bytes = [image_bytes]
+
+        conversations = pre_processing_chat(conversations)
+        prompt = self.create_chat_prompt(conversations)
+        prompt = post_processing_chat(prompt)
+        input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
+        input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
+        labels = self.generate_labels(input_ids)
+
+        image_inputs_list = [MiniMindVLM.image2tensor(Image.open(io.BytesIO(img)), self.preprocess) for img in image_bytes]
+        if hasattr(image_inputs_list[0], 'keys'):
+            image_data = {k: torch.cat([inp[k] for inp in image_inputs_list], dim=0) for k in image_inputs_list[0].keys()}
+        else:
+            image_data = torch.stack(image_inputs_list)
+
+        return torch.tensor(input_ids, dtype=torch.long), torch.tensor(labels, dtype=torch.long), image_data
+
+
+class OmniDataset(Dataset):
+    def __init__(self, data_path, tokenizer, audio_processor=None, vision_processor=None,
+                 max_length=1200, audio_special_token='<|audio_pad|>', image_special_token='<|image_pad|>',
+                 audio_stop_token=2050,  # <|audio_stop|>
+                 audio_pad_token=2049,  # <|audio_pad|>
+                 audio_spk_token=2051,  # <|audio_spk|>
+                 audio_vocab_size=2112,  # 2048 mimi codes + 64 special tokens
+                 scheduled_sampling=0.05,
+                 image_token_len=64):
+        super().__init__()
+        import pyarrow as pa
+        import pyarrow.parquet as pq
+        tables = [pa.Table.from_batches(pq.ParquetFile(p.strip()).iter_batches()) for p in data_path.split(',')]
+        tables = [t.cast(pa.schema([f.with_type(pa.large_string()) if pa.types.is_string(f.type) else f for f in t.schema])) for t in tables]
+        self.table = pa.concat_tables(tables, promote_options='default')
+        self.tokenizer = tokenizer
+        self.audio_processor = audio_processor
+        self.vision_processor = vision_processor
+        self.max_length = max_length
+        self.audio_token = audio_special_token
+        self.image_token_len = image_token_len
+        self.image_token = image_special_token * image_token_len
+        self.audio_stop_token = audio_stop_token
+        self.audio_pad_token = audio_pad_token
+        self.audio_spk_token = audio_spk_token
+        self.audio_vocab_size = audio_vocab_size
+        self.scheduled_sampling_prob = scheduled_sampling
+        self.text_vocab_size = len(tokenizer)
+        self.image_token_id = tokenizer.encode(image_special_token, add_special_tokens=False)[0]
+        self.audio_token_id = tokenizer.encode(audio_special_token, add_special_tokens=False)[0]
+        self.think_end_ids = tokenizer.encode('</think>\n\n', add_special_tokens=False)
+        self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids
+        self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
+
+    def __len__(self):
+        return len(self.table)
+
+    @staticmethod
+    def process_audio(audio_path, audio_processor):
+        import soundfile as sf
+        import numpy as np
+        wav, sr = sf.read(audio_path)
+        if wav.ndim > 1: wav = wav.mean(axis=1)
+        if sr != 16000:
+            import librosa
+            wav = librosa.resample(wav.astype(float), orig_sr=sr, target_sr=16000)
+        inputs = audio_processor(wav.astype(np.float32), sampling_rate=16000, return_tensors="pt", return_attention_mask=True)
+        valid_len = inputs.attention_mask.sum().item()
+        return inputs.input_features.squeeze(0), valid_len
+
+    def augment_wav(self, wav, sr=16000):
+        import numpy as np
+        from scipy.signal import resample
+        if random.random() < 0.5:
+            speed = random.uniform(0.7, 1.6)
+            wav = resample(wav, int(len(wav) / speed)).astype(np.float32)
+        if random.random() < 0.3:
+            noise = np.random.randn(len(wav)).astype(np.float32) * random.uniform(0.001, 0.01)
+            wav = wav + noise
+        if random.random() < 0.3:
+            wav = wav * random.uniform(0.8, 1.2)
+        if random.random() < 0.2 and len(wav) > sr:
+            start = random.randint(0, len(wav) - sr // 4)
+            wav[start:start + sr // 4] = 0
+        if random.random() < 0.2:
+            k = random.choice([3, 5, 7])
+            wav = np.convolve(wav, np.ones(k) / k, mode='same').astype(np.float32)
+        if random.random() < 0.3:
+            ir_len = int(sr * random.uniform(0.05, 0.2))
+            ir = np.random.randn(ir_len).astype(np.float32) * np.exp(-np.linspace(0, 10, ir_len))
+            ir[0] = 1.0
+            ir /= np.sqrt(np.sum(ir ** 2) + 1e-6)
+            wav = np.convolve(wav, ir, mode='same').astype(np.float32)
+        if random.random() < 0.2:
+            pink = np.cumsum(np.random.randn(len(wav))).astype(np.float32)
+            pink /= np.max(np.abs(pink)) + 1e-6
+            wav = wav + pink * random.uniform(0.003, 0.015)
+        return np.clip(wav, -1.0, 1.0).astype(np.float32)
+
+    def augment_mel(self, fbank):
+        import numpy as np
+        T, D = fbank.shape
+        if random.random() < 0.5:
+            f = random.randint(1, 64)
+            f0 = random.randint(0, D - f)
+            fbank[:, f0:f0 + f] = 0
+        if random.random() < 0.5 and T > 1:
+            t = random.randint(1, min(10, T))
+            t0 = random.randint(0, T - t)
+            fbank[t0:t0 + t, :] = 0
+        return fbank
+
+    def load_audio_inputs(self, audio_bytes):
+        import soundfile as sf
+        import numpy as np
+        import io
+        import librosa
+        if not audio_bytes: return None, 0
+        wav, sr = sf.read(io.BytesIO(audio_bytes))
+        if wav.ndim > 1: wav = wav.mean(axis=1)
+        if sr != 16000: wav = librosa.resample(wav.astype(float), orig_sr=sr, target_sr=16000)
+        wav = self.augment_wav(wav.astype(np.float32))
+        inputs = self.audio_processor(wav, sampling_rate=16000, return_tensors="pt", return_attention_mask=True)
+        valid_len = inputs.attention_mask.sum().item()
+        return self.augment_mel(inputs.input_features.squeeze(0)), valid_len
+
+    def load_image_inputs(self, image_bytes):
+        import io
+        from PIL import Image
+        if not image_bytes or self.vision_processor is None: return None
+        image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
+        inputs = self.vision_processor(images=image, return_tensors="pt")
+        if hasattr(inputs, 'keys'): return {k: v for k, v in inputs.items()}
+        return inputs.pixel_values
+
+    def create_chat_prompt(self, conversations, audio_features_length=0):
+        conversations = pre_processing_chat(conversations)
+        messages = []
+        is_last_user = lambda i: i == max(j for j, t in enumerate(conversations) if t['role'] == 'user')
+        for idx, turn in enumerate(conversations):
+            role, content = turn['role'], turn['content']
+            if role == 'user' and is_last_user(idx) and audio_features_length > 0:
+                ap = self.audio_token * audio_features_length
+                r = random.random()
+                if r < 0.4: content = ap
+                elif r < 0.6: content = content
+                elif r < 0.8: content = ap + '\n\n' + content
+                else: content = content + '\n\n' + ap
+            if '<image>' in content:
+                r = random.random()
+                if r < 0.2: content = '<image>\n' + content.replace('<image>', '').strip()
+                elif r < 0.4: content = '<image>\n\n' + content.replace('<image>', '').strip()
+                elif r < 0.6: content = content.replace('<image>', '').strip() + '\n' + '<image>'
+                else: content = content.replace('<image>', '').strip() + '\n\n' + '<image>'
+            messages.append({"role": role, "content": content})
+        prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
+        return post_processing_chat(prompt)
+
+    def generate_text_labels(self, input_ids):
+        labels = [-100] * len(input_ids)
+        ranges = []
+        i = 0
+        while i < len(input_ids):
+            if input_ids[i:i + len(self.bos_id)] == self.bos_id:
+                start = i + len(self.bos_id)
+                end = start
+                while end < len(input_ids):
+                    if input_ids[end:end + len(self.eos_id)] == self.eos_id:
+                        break
+                    end += 1
+                ranges.append((start, end))
+                for j in range(start, min(end + len(self.eos_id), self.max_length)):
+                    labels[j] = input_ids[j]
+                i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
+            else:
+                i += 1
+        return labels, ranges
+
+    def apply_scheduled_sampling(self, input_ids, audio_labels, text_labels):
+        if self.scheduled_sampling_prob <= 0:
+            return input_ids
+        audio_mask = (audio_labels != -100).any(dim=0) & (torch.rand(input_ids.size(1)) < self.scheduled_sampling_prob)
+        for i in range(8):
+            input_ids[i] = torch.where(audio_mask, torch.randint(0, self.audio_vocab_size, input_ids[i].shape), input_ids[i])
+        text_mask = (text_labels != -100) & (input_ids[8] != self.image_token_id) & (torch.rand(input_ids.size(1)) < self.scheduled_sampling_prob)
+        input_ids[8] = torch.where(text_mask, torch.randint(0, self.text_vocab_size, input_ids[8].shape), input_ids[8])
+        return input_ids
+
+    def __getitem__(self, index: int):
+        import numpy as np
+        conversations = json.loads(self.table['conversations'][index].as_py())
+        question_audios = self.table['question_audios'][index].as_py() if 'question_audios' in self.table.column_names else []
+        answer_audios = self.table['answer_audios'][index].as_py() if 'answer_audios' in self.table.column_names else []
+        image_bytes = self.table['image_bytes'][index].as_py() if 'image_bytes' in self.table.column_names else []
+        if image_bytes and not isinstance(image_bytes, list): image_bytes = [image_bytes]
+        ref_audios = self.table['ref_audios'][index].as_py() if 'ref_audios' in self.table.column_names else []
+        spk_emb_raw = self.table['spk_emb'][index].as_py() if 'spk_emb' in self.table.column_names else []
+
+        asst_indices = [i for i, t in enumerate(conversations) if t['role'] == 'assistant']
+        if len(asst_indices) > 1:
+            rand_idx = random.randint(0, len(asst_indices) - 1)
+            for i in range(rand_idx, -1, -1):
+                conversations = conversations[:asst_indices[i] + 1]
+                test_prompt = self.create_chat_prompt(conversations, 0)
+                if len(self.tokenizer(test_prompt).input_ids) + 100 < self.max_length:
+                    break
+
+        pixel_values = None
+        user_count = sum(1 for t in conversations if t['role'] == 'user')
+        if image_bytes and len(image_bytes) > 0 and self.vision_processor:
+            pixel_values = self.load_image_inputs(image_bytes[0])
+
+        audio_inputs, audio_len, audio_features_length = None, 0, 0
+        user_count = sum(1 for t in conversations if t['role'] == 'user')
+        if question_audios and user_count > 0 and user_count <= len(question_audios) and self.audio_processor:
+            audio_bytes = question_audios[user_count - 1]
+            if audio_bytes:
+                mel, valid_len = self.load_audio_inputs(audio_bytes)
+                if mel is not None:
+                    audio_inputs = mel.unsqueeze(0)
+                    audio_len = valid_len
+                    audio_features_length = valid_len or 1
+
+        if audio_inputs is None and self.audio_processor:
+            audio_inputs = torch.zeros(1, 1, 560)
+            audio_len = 0
+        if pixel_values is None and self.vision_processor:
+            pixel_values = {'pixel_values': torch.zeros(1, 3, 256, 256)}
+
+        last_audio_codes = None
+        asst_count = sum(1 for t in conversations if t['role'] == 'assistant')
+        if answer_audios and asst_count > 0 and asst_count <= len(answer_audios):
+            tokens = answer_audios[asst_count - 1]
+            if tokens:
+                audio_codes_8layers = [[] for _ in range(8)]
+                for i in range(0, len(tokens) - 7, 8):
+                    for j in range(8): audio_codes_8layers[j].append(tokens[i + j])
+                for layer in audio_codes_8layers: layer.append(self.audio_stop_token)
+                last_audio_codes = audio_codes_8layers
+
+        prompt = self.create_chat_prompt(conversations, audio_features_length)
+        if pixel_values is not None: prompt = prompt.replace('<image>', self.image_token)
+        input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
+
+        input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
+
+        text_labels, assistant_ranges = self.generate_text_labels(input_ids)
+        for start, end in assistant_ranges[:-1]:
+            mask_end = min(end + len(self.eos_id), self.max_length)
+            text_labels[start:mask_end] = [-100] * (mask_end - start)
+
+        Y_audio_layers = [[self.audio_pad_token] * self.max_length for _ in range(8)]
+        audio_labels = [[-100] * self.max_length for _ in range(8)]
+        if assistant_ranges and last_audio_codes:
+            assistant_start, assistant_end = assistant_ranges[-1]
+            for pos in range(assistant_start, min(assistant_end, assistant_start + 50)):
+                if input_ids[pos:pos + len(self.think_end_ids)] == self.think_end_ids:
+                    assistant_start = pos + len(self.think_end_ids)
+                    break
+            has_spk = bool(spk_emb_raw)
+            has_ref = bool(ref_audios) and random.random() > 0.5
+            spk_reserve = 1 if has_spk else 0
+            if has_ref:
+                ref_codes = [[] for _ in range(8)]
+                for i in range(0, len(ref_audios) - 7, 8):
+                    for j in range(8): ref_codes[j].append(ref_audios[i + j])
+                ref_len = len(ref_codes[0])
+                ref_start = max(spk_reserve, assistant_start - ref_len)
+                for layer_idx in range(8):
+                    codes = ref_codes[layer_idx][-(assistant_start - ref_start):] if ref_len > (assistant_start - ref_start) else ref_codes[layer_idx]
+                    for i, code in enumerate(codes):
+                        Y_audio_layers[layer_idx][ref_start + i] = code
+            else:
+                ref_start = assistant_start
+            if has_spk and ref_start > 0:
+                spk_pos = ref_start - 1
+                for layer_idx in range(8):
+                    Y_audio_layers[layer_idx][spk_pos] = self.audio_spk_token
+            for layer_idx in range(8):
+                codes = last_audio_codes[layer_idx]
+                start_pos = assistant_start + layer_idx + 1
+                for i, code in enumerate(codes):
+                    if start_pos + i < self.max_length:
+                        Y_audio_layers[layer_idx][start_pos + i] = code
+                        audio_labels[layer_idx][start_pos + i] = code
+
+        X_audio = torch.tensor([layer[:-1] for layer in Y_audio_layers], dtype=torch.long)  # (8, T-1)
+        X_text = torch.tensor(input_ids[:-1], dtype=torch.long)  # (T-1,)
+        input_ids = torch.cat((X_audio, X_text.unsqueeze(0)), dim=0)  # (9, T-1)
+        text_labels = torch.tensor(text_labels[1:], dtype=torch.long)  # (T-1,)
+        audio_labels = torch.tensor([layer[1:] for layer in audio_labels], dtype=torch.long)  # (8, T-1)
+
+        input_ids = self.apply_scheduled_sampling(input_ids, audio_labels, text_labels)
+        spk_emb = torch.tensor(spk_emb_raw, dtype=torch.float32) if spk_emb_raw else torch.zeros(192)
+        return input_ids, text_labels, audio_labels, audio_inputs, audio_len, pixel_values, spk_emb
+
+
+if __name__ == "__main__":
+    pass

src/omni/encoders/__init__.py

@@ -0,0 +1,3 @@
+from omni.encoders.vision import SiglipVisionEncoder
+
+__all__ = ["SiglipVisionEncoder"]

src/omni/encoders/audio/__init__.py

@@ -0,0 +1,3 @@
+from omni.encoders.audio.sensevoice import SenseVoiceAudioEncoder, SenseVoiceAudioProcessor
+
+__all__ = ["SenseVoiceAudioEncoder", "SenseVoiceAudioProcessor"]

src/omni/encoders/audio/sensevoice.py

@@ -0,0 +1,70 @@
+import os
+import logging
+import contextlib
+import io
+import warnings
+import torch
+import numpy as np
+from torch import nn
+from types import SimpleNamespace
+
+from transformers import logging as hf_logging
+
+
+class SenseVoiceAudioProcessor:
+    def __init__(self, frontend):
+        self.frontend = frontend
+
+    def __call__(self, wav, sampling_rate=16000, return_tensors="pt", return_attention_mask=True, **kwargs):
+        if isinstance(wav, np.ndarray):
+            wav = torch.from_numpy(wav).float()
+        if wav.dim() == 1:
+            wav = wav.unsqueeze(0)
+        with torch.no_grad():
+            fbank, flen = self.frontend(wav, torch.tensor([wav.size(1)]))
+        return SimpleNamespace(
+            input_features=fbank,
+            attention_mask=(torch.arange(fbank.size(1)) < flen[0]).long().unsqueeze(0),
+        )
+
+
+class SenseVoiceAudioEncoder(nn.Module):
+    """Frozen SenseVoice audio encoder (funasr AutoModel). Returns frame embeddings."""
+
+    def __init__(self, model_path: str = None):
+        super().__init__()
+        self.model = None
+        self.processor = None
+        if model_path is not None and os.path.exists(model_path):
+            self.load(model_path)
+
+    @torch.no_grad()
+    def load(self, model_path: str):
+        if not os.path.exists(model_path):
+            warnings.warn(f"[SenseVoiceAudioEncoder] path not found: {model_path}")
+            return
+        logging.getLogger().setLevel(logging.ERROR)
+        hf_logging.set_verbosity_error()
+        with contextlib.redirect_stdout(io.StringIO()):
+            from funasr import AutoModel
+            m = AutoModel(model=model_path, trust_remote_code=True, disable_update=True, device="cpu")
+        encoder, frontend = m.model.encoder, m.kwargs["frontend"]
+        for p in encoder.parameters():
+            p.requires_grad = False
+        self.model = encoder.eval().float()
+        self.processor = SenseVoiceAudioProcessor(frontend.eval())
+
+    @property
+    def hidden_size(self) -> int:
+        if self.model is None:
+            return 0
+        return self.model.config.d_model if hasattr(self.model.config, "d_model") else self.model.config.hidden_size
+
+    @torch.no_grad()
+    def encode(self, fbank, audio_lens=None):
+        if self.model is None:
+            return None
+        if audio_lens is None:
+            audio_lens = torch.tensor([fbank.size(1)] * fbank.size(0), device=fbank.device)
+        emb, _ = self.model(fbank, audio_lens)
+        return emb

src/omni/encoders/vision/__init__.py

@@ -0,0 +1,3 @@
+from omni.encoders.vision.siglip import SiglipVisionEncoder
+
+__all__ = ["SiglipVisionEncoder"]

src/omni/encoders/vision/siglip.py

@@ -0,0 +1,54 @@
+import os
+import torch
+import torch.nn.functional as F
+import warnings
+from torch import nn
+from typing import Optional
+
+from transformers import SiglipVisionModel, SiglipImageProcessor
+from transformers import logging as hf_logging
+
+
+class SiglipVisionEncoder(nn.Module):
+    """Frozen SigLIP vision encoder. Returns last-hidden-state embeddings (B, T, D)."""
+
+    def __init__(self, model_path: Optional[str] = None):
+        super().__init__()
+        self.model = None
+        self.processor = None
+        if model_path is not None and os.path.exists(model_path):
+            self.load(model_path)
+
+    @torch.no_grad()
+    def load(self, model_path: str):
+        hf_logging.set_verbosity_error()
+        try:
+            model = SiglipVisionModel.from_pretrained(model_path)
+        except (RuntimeError, ValueError):
+            warnings.warn(f"[SiglipVisionEncoder] failed to load from {model_path}")
+            return
+        processor = SiglipImageProcessor.from_pretrained(model_path)
+        for param in model.parameters():
+            param.requires_grad = False
+        self.model = model.eval()
+        self.processor = processor
+
+    @property
+    def hidden_size(self) -> int:
+        if self.model is None:
+            return 0
+        return self.model.config.hidden_size
+
+    def preprocess(self, image):
+        if image.mode in ('RGBA', 'LA'):
+            image = image.convert('RGB')
+        return self.processor(images=image, return_tensors='pt')
+
+    @torch.no_grad()
+    def encode(self, pixel_values):
+        if self.model is None:
+            return None
+        if hasattr(pixel_values, 'keys'):
+            pixel_values = {k: (v.squeeze(1) if v.ndim > 2 and v.shape[1] == 1 else v) for k, v in pixel_values.items()}
+        outputs = self.model(**pixel_values)
+        return outputs.last_hidden_state

src/omni/models/__init__.py

@@ -0,0 +1,53 @@
+from omni.core import (
+    RMSNorm,
+    Attention,
+    FeedForward,
+    MOEFeedForward,
+    MiniMindBlock,
+    MiniMindModel,
+    precompute_freqs_cis,
+    apply_rotary_pos_emb,
+)
+from omni.models.minimind import (
+    MiniMindConfig,
+    MiniMindForCausalLM,
+)
+from omni.models.vlm import (
+    VLMConfig,
+    MiniMindVLM,
+)
+from omni.models.omni import (
+    OmniConfig,
+    MiniMindOmni,
+    TalkerModule,
+)
+from omni.models.lora import (
+    LoRA,
+    apply_lora,
+    load_lora,
+    save_lora,
+    merge_lora,
+)
+
+__all__ = [
+    "MiniMindConfig",
+    "MiniMindModel",
+    "MiniMindForCausalLM",
+    "VLMConfig",
+    "MiniMindVLM",
+    "OmniConfig",
+    "MiniMindOmni",
+    "TalkerModule",
+    "RMSNorm",
+    "Attention",
+    "FeedForward",
+    "MOEFeedForward",
+    "MiniMindBlock",
+    "precompute_freqs_cis",
+    "apply_rotary_pos_emb",
+    "LoRA",
+    "apply_lora",
+    "load_lora",
+    "save_lora",
+    "merge_lora",
+]

src/omni/models/lora.py

@@ -0,0 +1,61 @@
+import torch
+from torch import optim, nn
+
+
+class LoRA(nn.Module):
+    def __init__(self, in_features, out_features, rank):
+        super().__init__()
+        self.rank = rank
+        self.A = nn.Linear(in_features, rank, bias=False)
+        self.B = nn.Linear(rank, out_features, bias=False)
+        self.A.weight.data.normal_(mean=0.0, std=0.02)
+        self.B.weight.data.zero_()
+
+    def forward(self, x):
+        return self.B(self.A(x))
+
+
+def apply_lora(model, rank=16):
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear) and module.in_features == module.out_features:
+            lora = LoRA(module.in_features, module.out_features, rank=rank).to(model.device)
+            setattr(module, "lora", lora)
+            original_forward = module.forward
+
+            def forward_with_lora(x, layer1=original_forward, layer2=lora):
+                return layer1(x) + layer2(x)
+
+            module.forward = forward_with_lora
+
+
+def load_lora(model, path):
+    state_dict = torch.load(path, map_location=model.device)
+    state_dict = {(k[7:] if k.startswith('module.') else k): v for k, v in state_dict.items()}
+
+    for name, module in model.named_modules():
+        if hasattr(module, 'lora'):
+            lora_state = {k.replace(f'{name}.lora.', ''): v for k, v in state_dict.items() if f'{name}.lora.' in k}
+            module.lora.load_state_dict(lora_state)
+
+
+def save_lora(model, path):
+    raw_model = getattr(model, '_orig_mod', model)
+    state_dict = {}
+    for name, module in raw_model.named_modules():
+        if hasattr(module, 'lora'):
+            clean_name = name[7:] if name.startswith("module.") else name
+            lora_state = {f'{clean_name}.lora.{k}': v.cpu().half() for k, v in module.lora.state_dict().items()}
+            state_dict.update(lora_state)
+    torch.save(state_dict, path)
+
+
+def merge_lora(model, lora_path, save_path):
+    load_lora(model, lora_path)
+    raw_model = getattr(model, '_orig_mod', model)
+    state_dict = {k: v.cpu().half() for k, v in raw_model.state_dict().items() if '.lora.' not in k}
+    for name, module in raw_model.named_modules():
+        if isinstance(module, nn.Linear) and '.lora.' not in name:
+            state_dict[f'{name}.weight'] = module.weight.data.clone().cpu().half()
+            if hasattr(module, 'lora'):
+                state_dict[f'{name}.weight'] += (module.lora.B.weight.data @ module.lora.A.weight.data).cpu().half()
+    torch.save(state_dict, save_path)

src/omni/models/minimind.py

@@ -0,0 +1,113 @@
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+
+from omni.core import MiniMindModel
+
+
+class MiniMindConfig(PretrainedConfig):
+    model_type = "minimind"
+
+    def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.use_moe = use_moe
+        self.dropout = kwargs.get("dropout", 0.0)
+        self.vocab_size = kwargs.get("vocab_size", 6400)
+        self.bos_token_id = kwargs.get("bos_token_id", 1)
+        self.eos_token_id = kwargs.get("eos_token_id", 2)
+        self.flash_attn = kwargs.get("flash_attn", True)
+        self.num_attention_heads = kwargs.get("num_attention_heads", 8)
+        self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
+        self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)
+        self.hidden_act = kwargs.get("hidden_act", 'silu')
+        self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
+        self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768)
+        self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)
+        self.rope_theta = kwargs.get("rope_theta", 1e6)
+        self.tie_word_embeddings = kwargs.get("tie_word_embeddings", True)
+        self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)
+        self.rope_scaling = {
+            "beta_fast": 32,
+            "beta_slow": 1,
+            "factor": 16,
+            "original_max_position_embeddings": 2048,
+            "attention_factor": 1.0,
+            "type": "yarn"
+        } if self.inference_rope_scaling else None
+        # MoE specific configs (ignored if use_moe = False)
+        self.num_experts = kwargs.get("num_experts", 4)
+        self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1)
+        self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size)
+        self.norm_topk_prob = kwargs.get("norm_topk_prob", True)
+        self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4)
+
+
+class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MiniMindConfig
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: MiniMindConfig = None):
+        self.config = config or MiniMindConfig()
+        super().__init__(self.config)
+        self.model = MiniMindModel(self.config)
+        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        if self.config.tie_word_embeddings:
+            self.model.embed_tokens.weight = self.lm_head.weight
+        self.post_init()
+
+    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs):
+        hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        loss = None
+        if labels is not None:
+            x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
+            loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
+        return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states)
+
+    @torch.inference_mode()
+    def generate(self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=2, streamer=None, use_cache=True, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs):
+        input_ids = kwargs.pop("input_ids", inputs).repeat(num_return_sequences, 1)
+        attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
+        past_key_values = kwargs.pop("past_key_values", None)
+        finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device)
+        if streamer:
+            streamer.put(input_ids.cpu())
+        for _ in range(max_new_tokens):
+            past_len = past_key_values[0][0].shape[1] if past_key_values else 0
+            outputs = self.forward(input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, **kwargs)
+            attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
+            logits = outputs.logits[:, -1, :] / temperature
+            if repetition_penalty != 1.0:
+                for i in range(input_ids.shape[0]):
+                    seen = torch.unique(input_ids[i])
+                    score = logits[i, seen]
+                    logits[i, seen] = torch.where(score > 0, score / repetition_penalty, score * repetition_penalty)
+            if top_k > 0:
+                logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
+                mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
+                logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')
+            next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
+            if eos_token_id is not None:
+                next_token = torch.where(finished.unsqueeze(-1), next_token.new_full((next_token.shape[0], 1), eos_token_id), next_token)
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+            past_key_values = outputs.past_key_values if use_cache else None
+            if streamer:
+                streamer.put(next_token.cpu())
+            if eos_token_id is not None:
+                finished |= next_token.squeeze(-1).eq(eos_token_id)
+                if finished.all():
+                    break
+        if streamer:
+            streamer.end()
+        if kwargs.get("return_kv"):
+            return {'generated_ids': input_ids, 'past_kv': past_key_values}
+        return input_ids

src/omni/models/omni.py

@@ -0,0 +1,409 @@
+import os
+import math
+import torch
+import warnings
+import logging
+import contextlib
+import io
+from torch import nn
+from torch.nn import functional as F
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+from transformers import SiglipVisionModel, SiglipImageProcessor, logging as hf_logging
+
+from omni.core import RMSNorm, precompute_freqs_cis, apply_rotary_pos_emb, MiniMindBlock, MOEFeedForward
+from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.encoders.audio import SenseVoiceAudioEncoder
+from omni.encoders.vision import SiglipVisionEncoder
+from omni.projectors import MMVisionProjector, MMAudioProjector
+
+
+class OmniConfig(MiniMindConfig):
+    model_type = "minimind-o"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.num_talker_hidden_layers = kwargs.get("num_talker_hidden_layers", 4)
+        self.talker_hidden_size = kwargs.get("talker_hidden_size", 768)
+        self.audio_ids = kwargs.get("audio_ids", [16])
+        self.audio_special_token = kwargs.get("audio_special_token", "<|audio_pad|>")
+        self.audio_hidden_size = kwargs.get("audio_hidden_size", 512)
+        self.audio_vocab_size = kwargs.get("audio_vocab_size", 2112)
+        self.audio_pad_token = kwargs.get("audio_pad_token", 2049)
+        self.audio_stop_token = kwargs.get("audio_stop_token", 2050)
+        self.audio_spk_token = kwargs.get("audio_spk_token", 2051)
+        self.spk_emb_size = kwargs.get("spk_emb_size", 192)
+        self.think_end_ids = kwargs.get("think_end_ids", [26, 234, 234])
+        self.image_ids = kwargs.get("image_ids", [12])
+        self.image_special_token = kwargs.get("image_special_token", "<|image_pad|>")
+        self.image_hidden_size = kwargs.get("image_hidden_size", 768)
+        self.image_token_len = kwargs.get("image_token_len", 64)
+        self.bridge_layer = kwargs.get("bridge_layer", self.num_hidden_layers // 2 - 1)
+
+
+class TalkerHead(nn.Module):
+    def __init__(self, in_features, out_features, num_layers=8, rank=256):
+        super().__init__()
+        self.num_layers = num_layers
+        self.base = nn.Linear(in_features, out_features, bias=False)
+        self.adapters = nn.ModuleList([
+            nn.Sequential(nn.Linear(in_features, rank, bias=False), nn.GELU(), nn.Linear(rank, out_features, bias=False))
+            for _ in range(num_layers)
+        ])
+
+    def forward(self, x):
+        base_out = self.base(x)
+        return [base_out + adapter(x) for adapter in self.adapters]
+
+
+class TalkerEmbedding(nn.Module):
+    def __init__(self, num_embeddings, embedding_dim, num_layers=8, rank=256):
+        super().__init__()
+        self.num_layers = num_layers
+        self.base = nn.Embedding(num_embeddings, embedding_dim)
+        self.adapters = nn.ModuleList([
+            nn.Sequential(nn.Embedding(num_embeddings, rank), nn.GELU(), nn.Linear(rank, embedding_dim, bias=False))
+            for _ in range(num_layers)
+        ])
+
+    def forward(self, x):
+        base_out = self.base(x)
+        return sum(base_out[:, i, :] + self.adapters[i](x[:, i, :]) for i in range(len(self.adapters))) / self.num_layers
+
+
+class TalkerModule(nn.Module):
+    def __init__(self, config: OmniConfig):
+        super().__init__()
+        self.talker_config = MiniMindConfig(hidden_size=config.talker_hidden_size, use_moe=config.use_moe)
+        self.layers = nn.ModuleList([MiniMindBlock(l, self.talker_config) for l in range(config.num_talker_hidden_layers)])
+        self.norm = RMSNorm(config.talker_hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = TalkerHead(config.talker_hidden_size, config.audio_vocab_size)
+        self.embed_tokens = TalkerEmbedding(config.audio_vocab_size, config.talker_hidden_size)
+        self.codec_proj = nn.Sequential(
+            nn.Linear(config.talker_hidden_size, config.talker_hidden_size),
+            nn.GELU(),
+            nn.Linear(config.talker_hidden_size, config.talker_hidden_size),
+            RMSNorm(config.talker_hidden_size, eps=config.rms_norm_eps),
+        )
+        self.embed_proj = nn.Sequential(
+            nn.Linear(config.hidden_size, config.hidden_size),
+            nn.GELU(),
+            nn.Linear(config.hidden_size, config.talker_hidden_size),
+            RMSNorm(config.talker_hidden_size, eps=config.rms_norm_eps),
+        )
+        self.text_scale, self.audio_scale = nn.Parameter(torch.tensor(3.0)), nn.Parameter(torch.tensor(1.0))
+        self.spk_proj = nn.Linear(config.spk_emb_size, config.talker_hidden_size, bias=False)
+        freqs_cos, freqs_sin = precompute_freqs_cis(
+            dim=self.talker_config.head_dim, end=config.max_position_embeddings,
+            rope_base=config.rope_theta, rope_scaling=config.rope_scaling
+        )
+        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
+        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+
+
+class MiniMindOmni(MiniMindForCausalLM):
+    config_class = OmniConfig
+
+    def __init__(self, config: OmniConfig = None, audio_encoder_path: str = None, vision_model_path: str = None):
+        config = config or OmniConfig()
+        super().__init__(config)
+        object.__setattr__(self, 'thinker', self.model)
+        object.__setattr__(self.model, 'lm_head', self.lm_head)
+        self.talker = TalkerModule(config)
+        self.audio_proj = MMAudioProjector(config.audio_hidden_size, config.hidden_size)
+        self.vision_proj = MMVisionProjector(config.image_hidden_size, config.hidden_size, target_tokens=config.image_token_len)
+        self.audio_pad_token, self.audio_stop_token, self.audio_spk_token = config.audio_pad_token, config.audio_stop_token, config.audio_spk_token
+        audio_encoder = SenseVoiceAudioEncoder(audio_encoder_path) if audio_encoder_path else SenseVoiceAudioEncoder()
+        object.__setattr__(self, 'audio_encoder', audio_encoder)
+        object.__setattr__(self, 'audio_processor', audio_encoder.processor)
+        vision_encoder = SiglipVisionEncoder(vision_model_path) if vision_model_path else SiglipVisionEncoder()
+        object.__setattr__(self, 'vision_encoder', vision_encoder)
+        object.__setattr__(self, 'vision_processor', vision_encoder.processor)
+
+    @staticmethod
+    def load_sensevoice(path):
+        if not os.path.exists(path):
+            warnings.warn(f"[MiniMindOmni] SenseVoice path not found: {path}")
+            return None, None
+        logging.getLogger().setLevel(logging.ERROR)
+        hf_logging.set_verbosity_error()
+        with contextlib.redirect_stdout(io.StringIO()):
+            from funasr import AutoModel
+            m = AutoModel(model=path, trust_remote_code=True, disable_update=True, device="cpu")
+        encoder, frontend = m.model.encoder, m.kwargs["frontend"]
+        for p in encoder.parameters():
+            p.requires_grad = False
+        return encoder.eval().float(), SenseVoiceAudioProcessor(frontend.eval())
+
+    @staticmethod
+    def load_vision(path):
+        if path is None or not os.path.exists(path):
+            warnings.warn(f"[MiniMindOmni] Vision model path not found: {path}. vision_encoder will be None!")
+            return None, None
+        hf_logging.set_verbosity_error()
+        try:
+            model = SiglipVisionModel.from_pretrained(path)
+        except (RuntimeError, ValueError):
+            return None, None
+        processor = SiglipImageProcessor.from_pretrained(path)
+        for p in model.parameters():
+            p.requires_grad = False
+        return model.eval(), processor
+
+    @torch.compiler.disable
+    def encode_audio_inputs(self, audio_inputs, audio_lens=None):
+        if (audio_inputs is None) or (self.audio_encoder is None) or (not audio_inputs.any()):
+            return None
+        batch_mask = audio_inputs.flatten(1).any(1)
+        enc_dtype = next(self.audio_encoder.parameters()).dtype
+        valid_fbank = audio_inputs[batch_mask].to(dtype=enc_dtype)
+        if audio_lens is not None:
+            valid_lens = audio_lens[batch_mask].to(valid_fbank.device)
+        else:
+            valid_lens = torch.tensor([valid_fbank.size(1)] * valid_fbank.size(0), device=valid_fbank.device)
+        with torch.no_grad():
+            emb, _ = self.audio_encoder.model(valid_fbank, valid_lens)
+        proj_dtype = next(self.audio_proj.parameters()).dtype
+        emb_list = [self.audio_proj(emb[i, :max(1, min(valid_lens[i].item(), emb.size(1)))].unsqueeze(0).to(proj_dtype)).squeeze(0) for i in range(emb.size(0))]
+        if batch_mask.all():
+            return emb_list
+        out = [None] * audio_inputs.size(0)
+        j = 0
+        for i in range(audio_inputs.size(0)):
+            if batch_mask[i]:
+                out[i] = emb_list[j]
+                j += 1
+        return out
+
+    @torch.compiler.disable
+    def inject_audio_features(self, tokens, h, audio_feats, seqlen):
+        if audio_feats is None or not self.config.audio_ids:
+            return h
+        marker = self.config.audio_ids[0]
+        out = []
+        for b in range(h.size(0)):
+            hb, seq, i = h[b], tokens[b].tolist(), 0
+            af = audio_feats[b] if audio_feats[b] is not None else None
+            while i < len(seq):
+                if seq[i] == marker:
+                    start = i
+                    while i < len(seq) and seq[i] == marker:
+                        i += 1
+                    if af is not None:
+                        inject_len = min(af.size(0), i - start)
+                        hb = torch.cat((hb[:start], af[:inject_len], hb[start + inject_len:]), dim=0)
+                        af = None
+                else:
+                    i += 1
+            out.append(hb)
+        return torch.stack(out)
+
+    @torch.compiler.disable
+    def get_image_embeddings(self, image_inputs):
+        if hasattr(image_inputs, 'keys'):
+            image_inputs = {k: (v.squeeze(1) if v.ndim > 2 and v.shape[1] == 1 else v) for k, v in image_inputs.items()}
+            pixel_attention_mask = image_inputs.get('pixel_attention_mask')
+            if pixel_attention_mask is not None and not pixel_attention_mask.any():
+                pv = image_inputs['pixel_values']
+                return pv.new_zeros(pv.size(0), pv.size(1), self.config.image_hidden_size)
+        with torch.no_grad():
+            outputs = self.vision_encoder.model(**image_inputs)
+        return outputs.last_hidden_state
+
+    @torch.compiler.disable
+    def encode_image_inputs(self, pixel_values):
+        if pixel_values is None or self.vision_encoder is None:
+            return None
+        mask = pixel_values.flatten(1).any(1)
+        if not mask.any():
+            return pixel_values.new_zeros(pixel_values.size(0), self.config.image_token_len, self.config.hidden_size)
+        with torch.no_grad():
+            emb = self.vision_encoder.model(pixel_values=pixel_values[mask]).last_hidden_state
+        if emb.dim() == 2:
+            emb = emb.unsqueeze(0)
+        emb = self.vision_proj(emb)
+        if mask.all():
+            return emb
+        idx = mask.nonzero().view(-1, 1, 1).expand_as(emb)
+        return emb.new_zeros(pixel_values.size(0), *emb.shape[1:]).scatter(0, idx, emb)
+
+    @torch.compiler.disable
+    def count_vision_proj(self, tokens, h, vision_tensors=None, seqlen=512):
+        if vision_tensors is None or not self.config.image_ids:
+            return h
+        marker, vf = self.config.image_ids[0], vision_tensors
+        if vf.dim() == 3:
+            vf = vf.unsqueeze(1)
+        out = []
+        for b in range(h.size(0)):
+            hb, seq, k, i = h[b], tokens[b].tolist(), 0, 0
+            while i < len(seq):
+                if seq[i] == marker:
+                    start = i
+                    while i < len(seq) and seq[i] == marker:
+                        i += 1
+                    if k < vf.size(1):
+                        hb = torch.cat((hb[:start], vf[b][k][:i - start], hb[i:]), dim=0)[:seqlen]
+                        k += 1
+                else:
+                    i += 1
+            out.append(hb)
+        return torch.stack(out)
+
+    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0,
+                audio_inputs=None, audio_lens=None, pixel_values=None, **args):
+        if len(input_ids.shape) == 2:
+            batch_size, seq_length = input_ids.shape
+            text_ids = input_ids
+            audio_ids = torch.full((batch_size, 8, seq_length), self.audio_pad_token, dtype=torch.long, device=input_ids.device)
+        else:
+            batch_size, _, seq_length = input_ids.shape
+            text_ids, audio_ids = input_ids[:, 8, :], input_ids[:, :8, :]
+        if hasattr(past_key_values, 'layers'):
+            past_key_values = None
+        n_thinker, n_talker = len(self.thinker.layers), len(self.talker.layers)
+        past_key_values = past_key_values or ([None] * (n_thinker + n_talker))
+        start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
+        if self.thinker.freqs_cos[0, 0] == 0:
+            freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
+            self.thinker.freqs_cos, self.thinker.freqs_sin = freqs_cos.to(input_ids.device), freqs_sin.to(input_ids.device)
+        if self.talker.freqs_cos[0, 0] == 0:
+            freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.talker.talker_config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
+            self.talker.freqs_cos, self.talker.freqs_sin = freqs_cos.to(input_ids.device), freqs_sin.to(input_ids.device)
+        presents = []
+
+        hidden_states = self.thinker.dropout(self.thinker.embed_tokens(text_ids))
+        position_embeddings = (self.thinker.freqs_cos[start_pos:start_pos + seq_length], self.thinker.freqs_sin[start_pos:start_pos + seq_length])
+        if audio_inputs is not None and start_pos == 0:
+            audio_features = self.encode_audio_inputs(audio_inputs, audio_lens)
+            hidden_states = self.inject_audio_features(text_ids, hidden_states, audio_features, seq_length)
+        if pixel_values is not None and start_pos == 0:
+            if hasattr(pixel_values, 'keys'):
+                img_emb = self.get_image_embeddings(pixel_values).to(hidden_states.dtype)
+                vision_tensors = self.vision_proj(img_emb)
+            else:
+                if len(pixel_values.shape) == 6:
+                    pixel_values = pixel_values.squeeze(2)
+                if len(pixel_values.shape) == 4:
+                    pixel_values = pixel_values.unsqueeze(1)
+                bs, num, c, im_h, im_w = pixel_values.shape
+                stack_dim = 1 if bs > 1 else 0
+                vision_tensors = torch.stack([self.encode_image_inputs(pixel_values[:, i, :, :, :]) for i in range(num)], dim=stack_dim)
+            hidden_states = self.count_vision_proj(tokens=text_ids, h=hidden_states, vision_tensors=vision_tensors, seqlen=seq_length)
+        bridge_states = hidden_states
+        for i, (layer, past_key_value) in enumerate(zip(self.thinker.layers, past_key_values[:n_thinker])):
+            hidden_states, present = layer(hidden_states, position_embeddings, past_key_value=past_key_value, use_cache=use_cache, attention_mask=attention_mask)
+            presents.append(present)
+            if i == self.config.bridge_layer:
+                bridge_states = hidden_states
+        h_thinker = self.thinker.norm(hidden_states)
+
+        talker_emb = self.talker.embed_tokens(audio_ids)
+        spk_emb = args.get('spk_emb', None)
+        if spk_emb is not None:
+            spk_mask = (audio_ids[:, 0, :] == self.audio_spk_token).unsqueeze(-1)
+            talker_emb = torch.where(spk_mask, self.talker.spk_proj(spk_emb).unsqueeze(1), talker_emb)
+        hidden_states = self.talker.embed_proj(bridge_states) * self.talker.text_scale + self.talker.codec_proj(talker_emb) * self.talker.audio_scale
+        talker_pos_emb = (self.talker.freqs_cos[start_pos:start_pos + seq_length], self.talker.freqs_sin[start_pos:start_pos + seq_length])
+        for layer, past_key_value in zip(self.talker.layers, past_key_values[n_thinker:]):
+            hidden_states, present = layer(hidden_states, talker_pos_emb, past_key_value=past_key_value, use_cache=use_cache, attention_mask=attention_mask)
+            presents.append(present)
+        h_talker = self.talker.norm(hidden_states)
+
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        aux_loss = sum(l.mlp.aux_loss for l in list(self.thinker.layers) + list(self.talker.layers) if isinstance(l.mlp, MOEFeedForward))
+        aux_loss += sum(p.sum() for p in self.audio_proj.parameters()) * 0 + sum(p.sum() for p in self.vision_proj.parameters()) * 0 + sum(p.sum() for p in self.talker.lm_head.adapters.parameters()) * 0 + sum(p.sum() for p in self.talker.spk_proj.parameters()) * 0
+        text_logits = self.thinker.lm_head(h_thinker[:, slice_indices, :])
+        audio_logits = self.talker.lm_head(h_talker[:, slice_indices, :])
+
+        out = MoeCausalLMOutputWithPast(aux_loss=aux_loss, logits=text_logits, past_key_values=presents)
+        out.audio_logits = audio_logits
+        return out
+
+    @torch.inference_mode()
+    def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,
+                 stream=False, rp=1., use_cache=True, return_audio_codes=False, **args):
+        if stream:
+            return self.stream_generate(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, return_audio_codes, **args)
+        tokens = list(self.stream_generate(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, return_audio_codes, **args))
+        return tokens[-1] if tokens else input_ids
+
+    def stream_generate(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, return_audio_codes=False, **args):
+        start_pos, past_kvs, text_finished, first_finished = input_ids.shape[1], None, False, True
+        audio_codes = [[] for _ in range(8)]
+        audio_stop_pos = [None] * 8
+        audio_buffer = torch.full((1, 8, start_pos), self.audio_pad_token, dtype=torch.long, device=input_ids.device)
+        spk_emb = args.get('spk_emb', None)
+        ref_codes = args.get('ref_codes', None)
+        ref_len = ref_codes.shape[2] if ref_codes is not None else 0
+        spk_reserve = 1 if spk_emb is not None else 0
+        fill_end = start_pos
+        fill_start = max(spk_reserve, start_pos - ref_len)
+        if ref_codes is not None and fill_start < fill_end:
+            audio_buffer[:, :, fill_start:fill_end] = ref_codes[:, :, -(fill_end - fill_start):]
+        if spk_emb is not None and fill_start > 0:
+            audio_buffer[:, :, fill_start - 1] = self.audio_spk_token
+        think_end_step, generated_tokens = None, ([] if args.get('open_thinking', False) else None)
+        while input_ids.shape[1] < start_pos + max_new_tokens:
+            if past_kvs is None or not use_cache:
+                out = self.forward(torch.cat((audio_buffer, input_ids.unsqueeze(1)), dim=1), past_key_values=past_kvs, use_cache=use_cache, **args)
+            else:
+                out = self.forward(torch.cat((audio_buffer[:, :, -1:], input_ids[:, -1:].unsqueeze(1)), dim=1), past_key_values=past_kvs, use_cache=use_cache, **args)
+            past_kvs = out.past_key_values
+
+            logits = out.logits[0, -1, :].clone() / (temperature + 1e-9)
+            if rp != 1.0:
+                seen = list(set(input_ids[0].tolist()))
+                score = logits[seen]
+                logits[seen] = torch.where(score > 0, score / rp, score * rp)
+            if top_p and top_p < 1.0:
+                sorted_l, sorted_i = torch.sort(logits, descending=True)
+                mask = torch.cumsum(F.softmax(sorted_l, dim=-1), dim=-1) > top_p
+                mask[1:], mask[0] = mask[:-1].clone(), False
+                logits[sorted_i[mask]] = -float('Inf')
+            text_token = torch.multinomial(F.softmax(logits, dim=-1), 1).item()
+
+            if text_finished:
+                text_token = args.get('enter_token_id', 201) if first_finished else args.get('pad_token_id', 0)
+                first_finished = False
+
+            step = input_ids.shape[1] - start_pos
+            audio_step = step - 1
+            if generated_tokens is not None:
+                generated_tokens.append(text_token)
+                if not think_end_step and generated_tokens[-len(self.config.think_end_ids):] == list(self.config.think_end_ids):
+                    think_end_step = step + 2
+                audio_step = (step - think_end_step) if think_end_step else -1
+            for i, al in enumerate(out.audio_logits):
+                if audio_step < i:
+                    audio_codes[i].append(self.audio_pad_token)
+                else:
+                    logits_i = al[0, -1, :].clone() / 0.2
+                    for prev_code in audio_codes[i][-3:]:
+                        score = logits_i[prev_code]
+                        logits_i[prev_code] = torch.where(score > 0, score / 1.05, score * 1.05)
+                    top_val, top_idx = logits_i.topk(50)
+                    code = top_idx[torch.multinomial(F.softmax(top_val, dim=-1), 1)].item()
+                    audio_codes[i].append(code)
+                    if audio_stop_pos[i] is None and code >= 2048:
+                        audio_stop_pos[i] = len(audio_codes[i]) - 1
+
+            if text_finished and all(audio_stop_pos[i] is not None for i in range(8)):
+                break
+
+            input_ids = torch.cat((input_ids, torch.tensor([[text_token]], device=input_ids.device)), dim=1)
+            audio_buffer = torch.cat((audio_buffer, torch.full((1, 8, 1), self.audio_pad_token, dtype=torch.long, device=input_ids.device)), dim=2)
+            for i in range(min(audio_step + 1, 8)):
+                audio_buffer[0, i, -1] = audio_codes[i][-1]
+
+            audio_frame = None
+            if return_audio_codes and audio_step >= 7:
+                frame = [audio_codes[i][step - 7 + i] for i in range(8)]
+                active_layers = sum(1 for i in range(8) if audio_stop_pos[i] is None or step - 7 + i < audio_stop_pos[i])
+                if active_layers >= 8:
+                    audio_frame = frame
+            if not text_finished:
+                yield input_ids[:, start_pos:], audio_frame
+                if text_token == eos_token_id:
+                    text_finished = True
+            else:
+                yield None, audio_frame

src/omni/models/vlm.py

@@ -0,0 +1,155 @@
+import os
+import torch
+import torch.nn.functional as F
+import warnings
+from typing import Optional, Tuple, List, Union
+from torch import nn
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+
+from omni.core import precompute_freqs_cis, MOEFeedForward
+from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.encoders.vision import SiglipVisionEncoder
+from omni.projectors import MMVisionProjector
+
+warnings.filterwarnings('ignore')
+
+
+class VLMConfig(MiniMindConfig):
+    model_type = "minimind-v"
+
+    def __init__(self, image_special_token='<|image_pad|>', image_ids=[12], **kwargs):
+        self.image_special_token = image_special_token
+        self.image_ids = image_ids
+        self.image_hidden_size = kwargs.get("image_hidden_size", 768)
+        self.image_token_len = kwargs.get("image_token_len", 64)
+        super().__init__(**kwargs)
+
+
+class MiniMindVLM(MiniMindForCausalLM):
+    config_class = VLMConfig
+
+    def __init__(self, config: VLMConfig = None, vision_model_path: Optional[str] = None):
+        self.config = config or VLMConfig()
+        super().__init__(self.config)
+        self.vision_encoder = SiglipVisionEncoder(vision_model_path) if vision_model_path else SiglipVisionEncoder()
+        self.vision_proj = MMVisionProjector(
+            self.config.image_hidden_size, self.config.hidden_size, target_tokens=self.config.image_token_len
+        )
+
+    @staticmethod
+    def get_image_embeddings(image_inputs, vision_model):
+        if vision_model is None:
+            return None
+        return vision_model.encode(image_inputs)
+
+    @torch.compiler.disable
+    def count_vision_proj(self, tokens, h, vision_tensors=None, seqlen=512):
+        if vision_tensors is None or not self.config.image_ids:
+            return h
+        marker, vf = self.config.image_ids[0], vision_tensors
+        if vf.dim() == 3:
+            vf = vf.unsqueeze(1)
+        out = []
+        for b in range(h.size(0)):
+            hb, seq, k, i = h[b], tokens[b].tolist(), 0, 0
+            while i < len(seq):
+                if seq[i] == marker:
+                    start = i
+                    while i < len(seq) and seq[i] == marker:
+                        i += 1
+                    if k < vf.size(1):
+                        hb = torch.cat((hb[:start], vf[b][k][:i - start], hb[i:]), dim=0)[:seqlen]
+                        k += 1
+                else:
+                    i += 1
+            out.append(hb)
+        return torch.stack(out)
+
+    def forward(self,
+                input_ids: Optional[torch.Tensor] = None,
+                attention_mask: Optional[torch.Tensor] = None,
+                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+                use_cache: bool = False,
+                logits_to_keep: Union[int, torch.Tensor] = 0,
+                labels: Optional[torch.Tensor] = None,
+                pixel_values: Optional[torch.FloatTensor] = None,
+                **args):
+        batch_size, seq_length = input_ids.shape
+        if hasattr(past_key_values, 'layers'):
+            past_key_values = None
+        past_key_values = past_key_values or [None] * len(self.model.layers)
+        start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
+
+        hidden_states = self.model.dropout(self.model.embed_tokens(input_ids))
+
+        if pixel_values is not None and start_pos == 0:
+            if hasattr(pixel_values, 'keys'):
+                sample_val = next(iter(pixel_values.values()))
+                if sample_val.ndim == 5:
+                    bs, num = sample_val.shape[:2]
+                    vision_tensors = self.vision_proj(
+                        MiniMindVLM.get_image_embeddings(
+                            {k: v.flatten(0, 1) for k, v in pixel_values.items()}, self.vision_encoder
+                        )
+                    ).view(bs, num, self.config.image_token_len, -1)
+                else:
+                    vision_tensors = self.vision_proj(
+                        MiniMindVLM.get_image_embeddings(pixel_values, self.vision_encoder)
+                    )
+            else:
+                if len(pixel_values.shape) == 6:
+                    pixel_values = pixel_values.squeeze(2)
+                bs, num, c, im_h, im_w = pixel_values.shape
+                vision_tensors = torch.stack(
+                    [self.vision_proj(MiniMindVLM.get_image_embeddings(pixel_values[:, i, :, :, :], self.vision_encoder))
+                     for i in range(num)], dim=1
+                )
+            hidden_states = self.count_vision_proj(
+                tokens=input_ids, h=hidden_states, vision_tensors=vision_tensors, seqlen=input_ids.shape[1]
+            )
+
+        if self.model.freqs_cos[0, 0] == 0:
+            freqs_cos, freqs_sin = precompute_freqs_cis(
+                dim=self.config.head_dim, end=self.config.max_position_embeddings,
+                rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling
+            )
+            self.model.freqs_cos, self.model.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
+        position_embeddings = (
+            self.model.freqs_cos[start_pos:start_pos + seq_length],
+            self.model.freqs_sin[start_pos:start_pos + seq_length]
+        )
+
+        presents = []
+        for layer_idx, (layer, past_key_value) in enumerate(zip(self.model.layers, past_key_values)):
+            hidden_states, present = layer(
+                hidden_states, position_embeddings,
+                past_key_value=past_key_value, use_cache=use_cache, attention_mask=attention_mask
+            )
+            presents.append(present)
+
+        hidden_states = self.model.norm(hidden_states)
+
+        aux_loss = sum([l.mlp.aux_loss for l in self.model.layers if isinstance(l.mlp, MOEFeedForward)],
+                       hidden_states.new_zeros(1).squeeze())
+        aux_loss = aux_loss + sum(p.sum() for p in self.vision_proj.parameters()) * 0
+
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100)
+
+        output = MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=presents, hidden_states=hidden_states)
+        return output
+
+    def generate(self, *args, num_return_sequences=1, **kwargs):
+        if num_return_sequences > 1 and 'pixel_values' in kwargs:
+            pv = kwargs['pixel_values']
+            if hasattr(pv, 'keys'):
+                kwargs['pixel_values'] = {k: v.repeat(num_return_sequences, *([1] * (v.ndim - 1))) for k, v in pv.items()}
+            else:
+                kwargs['pixel_values'] = pv.repeat(num_return_sequences, *([1] * (pv.ndim - 1)))
+        return super().generate(*args, num_return_sequences=num_return_sequences, **kwargs)

src/omni/projectors/__init__.py

@@ -0,0 +1,4 @@
+from omni.projectors.vision import MMVisionProjector
+from omni.projectors.audio import MMAudioProjector
+
+__all__ = ["MMVisionProjector", "MMAudioProjector"]

src/omni/projectors/audio.py

@@ -0,0 +1,16 @@
+import torch
+from torch import nn
+
+
+class MMAudioProjector(nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.LayerNorm(in_dim),
+            nn.Linear(in_dim, out_dim),
+            nn.GELU(),
+            nn.Linear(out_dim, out_dim),
+        )
+
+    def forward(self, x):
+        return self.mlp(x)

src/omni/projectors/vision.py

@@ -0,0 +1,17 @@
+import math
+import torch
+from torch import nn
+
+
+class MMVisionProjector(nn.Module):
+    def __init__(self, in_dim, out_dim, source_tokens=64, target_tokens=64):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.LayerNorm(in_dim),
+            nn.Linear(in_dim, out_dim),
+            nn.GELU(),
+            nn.Linear(out_dim, out_dim),
+        )
+
+    def forward(self, x):
+        return self.mlp(x)

src/omni/serve/__init__.py

@@ -0,0 +1 @@
+import omni.serve

src/omni/serve/realtime.py

@@ -0,0 +1,68 @@
+import numpy as np
+
+
+class SileroVAD:
+    def __init__(self, path):
+        import onnxruntime as ort
+        opts = ort.SessionOptions()
+        opts.inter_op_num_threads = opts.intra_op_num_threads = 1
+        opts.log_severity_level = 4
+        self.session = ort.InferenceSession(path, providers=["CPUExecutionProvider"], sess_options=opts)
+        self.h, self.c = np.zeros((2, 1, 64), dtype=np.float32), np.zeros((2, 1, 64), dtype=np.float32)
+
+    def reset(self):
+        self.h[:], self.c[:] = 0, 0
+
+    def __call__(self, chunk, sr=16000):
+        out, self.h, self.c = self.session.run(None, {"input": chunk.reshape(1, -1).astype(np.float32), "h": self.h, "c": self.c, "sr": np.array(sr, dtype="int64")})
+        return float(out[0][0])
+
+
+class RealtimeSession:
+    def __init__(self, vad_path, sr=16000, threshold=0.8, min_speech_ms=128, min_silence_ms=800):
+        self.vad, self.sr, self.threshold = SileroVAD(vad_path), sr, threshold
+        self.min_speech, self.min_silence = int(sr * min_speech_ms / 1000), int(sr * min_silence_ms / 1000)
+        self.reset()
+
+    def reset(self):
+        self.vad.reset()
+        self.buffer, self.ring, self.speaking, self.generating, self.interrupt = [], [], False, False, False
+        self.speech_samples = self.silence_samples = self.tail_silence = 0
+
+    def push_chunk(self, chunk, W=1024):
+        for i in range(0, max(len(chunk), 1), W):
+            w = chunk[i:i + W]
+            if len(w) < W:
+                w = np.pad(w, (0, W - len(w)))
+            prob = self.vad(w, self.sr)
+            if prob > self.threshold:
+                self.silence_samples = self.tail_silence = 0
+                self.speech_samples += len(w)
+                self.buffer.append(w)
+                if self.speech_samples >= self.min_speech and not self.speaking:
+                    self.speaking = True
+                    self.buffer = self.ring + self.buffer
+                    self.ring = []
+                if self.generating and self.speaking:
+                    self.interrupt = True
+                    return 'interrupt'
+            elif self.speaking:
+                self.silence_samples += len(w)
+                self.tail_silence += 1
+                self.buffer.append(w)
+                if self.silence_samples >= self.min_silence:
+                    if self.tail_silence > 1:
+                        del self.buffer[-(self.tail_silence - 1):]
+                    self.speaking, self.speech_samples, self.silence_samples, self.tail_silence = False, 0, 0, 0
+                    return 'speech_end'
+            else:
+                if self.speech_samples > 0:
+                    self.buffer.clear()
+                self.speech_samples = 0
+                self.ring = [w]
+        return 'listening'
+
+    def get_audio(self):
+        audio = np.concatenate(self.buffer) if self.buffer else np.array([], dtype=np.float32)
+        self.buffer.clear()
+        return audio

src/omni/trainers/__init__.py

@@ -0,0 +1,17 @@
+from omni.trainers import pretrain, full_sft, lora, dpo, distillation, ppo, grpo, agent, rollout_engine, train_tokenizer, pretrain_vlm, full_sft_vlm, full_sft_omni
+
+__all__ = [
+    "pretrain",
+    "full_sft",
+    "lora",
+    "dpo",
+    "distillation",
+    "ppo",
+    "grpo",
+    "agent",
+    "rollout_engine",
+    "train_tokenizer",
+    "pretrain_vlm",
+    "full_sft_vlm",
+    "full_sft_omni",
+]

src/omni/trainers/agent.py

@@ -0,0 +1,489 @@
+import os
+
+
+import datasets  # noqa: F401  # Windows pyarrow/torch DLL conflict workaround (issue #771)
+import re
+import gc
+import json
+import math
+import random
+import signal
+import argparse
+import warnings
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+from contextlib import nullcontext
+from torch import optim
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from transformers import AutoTokenizer
+from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.datasets.lm_dataset import AgentRLDataset
+from omni.utils.training import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from omni.trainers.rollout_engine import create_rollout_engine, compute_per_token_logps
+
+warnings.filterwarnings('ignore')
+
+# ================================ 工具与 Reward = Start ================================
+
+def rep_penalty(text, n=3, cap=0.5):
+    toks = re.findall(r"\w+|[^\w\s]", text.lower())
+    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
+    return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
+
+# ======== 工具定义 ========
+TOOLS = [
+    {"type": "function", "function": {"name": "calculate_math", "description": "计算数学表达式", "parameters": {"type": "object", "properties": {"expression": {"type": "string"}}, "required": ["expression"]}}},
+    {"type": "function", "function": {"name": "unit_converter", "description": "单位换算", "parameters": {"type": "object", "properties": {"value": {"type": "number"}, "from_unit": {"type": "string"}, "to_unit": {"type": "string"}}, "required": ["value", "from_unit", "to_unit"]}}},
+    {"type": "function", "function": {"name": "get_current_weather", "description": "获取天气", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}}},
+    {"type": "function", "function": {"name": "get_current_time", "description": "获取时间", "parameters": {"type": "object", "properties": {"timezone": {"type": "string", "default": "Asia/Shanghai"}}, "required": []}}},
+    {"type": "function", "function": {"name": "get_exchange_rate", "description": "查询汇率", "parameters": {"type": "object", "properties": {"from_currency": {"type": "string"}, "to_currency": {"type": "string"}}, "required": ["from_currency", "to_currency"]}}},
+    {"type": "function", "function": {"name": "translate_text", "description": "翻译文本", "parameters": {"type": "object", "properties": {"text": {"type": "string"}, "target_language": {"type": "string"}}, "required": ["text", "target_language"]}}},
+]
+
+# ======== 模拟数据 ========
+WEATHER_DATA = {"北京": ("28°C", "晴"), "上海": ("15°C", "多云"), "广州": ("32°C", "闷热"), "深圳": ("30°C", "晴"), "杭州": ("22°C", "阴"), "成都": ("18°C", "小雨"), "武汉": ("25°C", "多云"), "南京": ("20°C", "晴"), "西安": ("16°C", "大风"), "重庆": ("26°C", "阴"), "Tokyo": ("12°C", "晴"), "New York": ("8°C", "多云"), "London": ("5°C", "小雨"), "Paris": ("10°C", "阴"), "Sydney": ("25°C", "晴朗")}
+TIME_DATA = {"Asia/Shanghai": "2025-03-07 14:30:00", "America/New_York": "2025-03-07 01:30:00", "Europe/London": "2025-03-07 06:30:00", "Asia/Tokyo": "2025-03-07 15:30:00", "Europe/Paris": "2025-03-07 07:30:00", "Australia/Sydney": "2025-03-07 17:30:00"}
+EXCHANGE_DATA = {("USD", "CNY"): 7.21, ("EUR", "CNY"): 7.85, ("GBP", "CNY"): 9.12, ("JPY", "CNY"): 0.048, ("USD", "EUR"): 0.92, ("USD", "GBP"): 0.79, ("CNY", "JPY"): 20.83, ("AUD", "CNY"): 4.72}
+TRANSLATE_DATA = {("你好世界", "english"): "Hello World", ("Good morning", "chinese"): "早上好", ("今天天气真好", "english"): "The weather is nice today", ("I love programming", "chinese"): "我喜欢编程", ("机器学习很有趣", "english"): "Machine learning is interesting", ("Happy birthday", "chinese"): "生日快乐"}
+UNIT_DATA = {"km_miles": 0.621371, "miles_km": 1.60934, "kg_pounds": 2.20462, "pounds_kg": 0.453592, "meters_feet": 3.28084, "feet_meters": 0.3048, "celsius_fahrenheit": 1.8, "fahrenheit_celsius": 0.5556}
+
+# ======== 模拟执行 ========
+MOCK_RESULTS = {
+    "calculate_math": lambda args: {"result": str(eval(str(args.get("expression", "0")).replace("^", "**").replace("×", "*").replace("÷", "/").replace("−", "-").replace("（", "(").replace("）", ")"), {"__builtins__": {}, "math": math}))},
+    "unit_converter": lambda args: {"result": round(float(args.get("value", 0)) * UNIT_DATA.get(f"{args.get('from_unit', '').lower()}_{args.get('to_unit', '').lower()}", 1), 4)},
+    "get_current_weather": lambda args: (lambda w: {"city": args.get("location"), "temperature": w[0], "humidity": "65%", "condition": w[1]})(WEATHER_DATA.get(args.get("location"), ("22°C", "晴"))),
+    "get_current_time": lambda args: {"datetime": TIME_DATA.get(args.get("timezone", "Asia/Shanghai"), "2025-03-07 14:30:00"), "timezone": args.get("timezone", "Asia/Shanghai")},
+    "get_exchange_rate": lambda args: {"from": args.get("from_currency"), "to": args.get("to_currency"), "rate": EXCHANGE_DATA.get((args.get("from_currency"), args.get("to_currency")), 1.0)},
+    "translate_text": lambda args: {"translated_text": TRANSLATE_DATA.get((args.get("text"), args.get("target_language")), args.get("text", ""))},
+}
+
+# ======== 参数校验 ========
+CHECK_ARGS = {
+    "calculate_math": lambda a: bool(a.get("expression")),
+    "unit_converter": lambda a: a.get("value") is not None and a.get("from_unit") and a.get("to_unit"),
+    "get_current_weather": lambda a: bool(a.get("location")),
+    "get_current_time": lambda a: True,
+    "get_exchange_rate": lambda a: bool(a.get("from_currency")) and bool(a.get("to_currency")),
+    "translate_text": lambda a: bool(a.get("text")) and bool(a.get("target_language")),
+}
+
+# ======== 工具调用解析与执行 ========
+def parse_tool_calls(text):
+    calls = []
+    for m in re.findall(r'<tool_call>(.*?)</tool_call>', text, re.DOTALL):
+        try: calls.append(json.loads(m.strip()))
+        except: pass
+    return calls
+
+def execute_tool(name, args):
+    fn = MOCK_RESULTS.get(name)
+    if not fn: return None
+    try:
+        signal.signal(signal.SIGALRM, lambda *_: (_ for _ in ()).throw(TimeoutError()))
+        signal.alarm(1)
+        return fn(args)
+    except:
+        return None
+    finally:
+        try: signal.alarm(0)
+        except: pass
+
+# ======== 多轮 Rollout ========
+def rollout_single(rollout_engine, tokenizer, messages, tools, max_turns=3, max_new_tokens=256, thinking_ratio=0.5, device="cuda"):
+    all_outputs = []
+    prompt_ids = None
+    response_ids = []
+    response_mask = []
+    response_old_logps = []
+    final_context = ""
+    unfinished = False
+    open_thinking = random.random() < thinking_ratio
+    for turn in range(max_turns):
+        context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools, open_thinking=open_thinking)
+        inputs = tokenizer(context, return_tensors="pt", add_special_tokens=False).to(device)
+        context_ids = inputs["input_ids"][0].tolist()
+        if prompt_ids is None:
+            prompt_ids = context_ids
+        rollout_result = rollout_engine.rollout(
+            prompt_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            num_generations=1,
+            max_new_tokens=max_new_tokens,
+            temperature=0.8,
+        )
+        new_ids = rollout_result.completion_ids[0].tolist()
+        new_logps = rollout_result.per_token_logps[0].tolist()
+        if len(new_ids) != len(new_logps): Logger(f"rollout token/logprob length mismatch: {len(new_ids)} vs {len(new_logps)}")
+        pairs = [(t, lp) for t, lp in zip(new_ids, new_logps) if t != tokenizer.pad_token_id and t != tokenizer.eos_token_id]
+        new_ids = [t for t, _ in pairs]
+        new_logps = [lp for _, lp in pairs]
+        new_text = rollout_result.completions[0]
+        all_outputs.append(new_text)
+        response_ids.extend(new_ids)
+        response_mask.extend([1] * len(new_ids))
+        response_old_logps.extend(new_logps)
+        final_context = context + new_text
+        calls = parse_tool_calls(new_text)
+        if not calls:
+            break
+        unfinished = turn == max_turns - 1
+        messages.append({"role": "assistant", "content": new_text})
+        for call in calls:
+            name, raw = call.get("name", ""), call.get("arguments", {})
+            if isinstance(raw, str):
+                try: raw = json.loads(raw)
+                except: raw = {}
+            result = execute_tool(name, raw)
+            result_str = (json.dumps(result, ensure_ascii=False) if result else '{"error": "tool not found"}')[:2048]  # 防止天文数字撑爆tokenizer
+            messages.append({"role": "tool", "content": result_str})
+
+        observe_context = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=not unfinished, tools=tools, open_thinking=open_thinking)
+        observe_ids = tokenizer(observe_context, return_tensors="pt", add_special_tokens=False)["input_ids"][0].tolist()
+        current_len = len(prompt_ids) + len(response_ids)
+        obs_delta = observe_ids[current_len:]
+        response_ids.extend(obs_delta)
+        response_mask.extend([0] * len(obs_delta))
+        response_old_logps.extend([0.0] * len(obs_delta))
+        final_context = observe_context
+
+    final_output = all_outputs[-1] if all_outputs else ""
+    prompt_ids = prompt_ids or []
+    return final_output, final_context, prompt_ids, response_ids, response_mask, response_old_logps, list(all_outputs), unfinished
+
+def rollout_batch(rollout_engine, tokenizer, messages_batch, tools_batch, num_gen, max_turns=3, max_new_tokens=256, thinking_ratio=0.5, device="cuda"):
+    all_completions = []
+    all_contexts = []
+    all_prompt_ids = []
+    all_response_ids = []
+    all_response_masks = []
+    all_response_old_logps = []
+    all_turn_outputs = []
+    all_unfinished = []
+    for messages, tools in zip(messages_batch, tools_batch):
+        for _ in range(num_gen):
+            msgs_copy = [dict(m) for m in messages]
+            completion, context, prompt_ids, response_ids, response_mask, response_old_logps, turn_outputs, unfinished = rollout_single(rollout_engine, tokenizer, msgs_copy, tools, max_turns, max_new_tokens, thinking_ratio, device)
+            all_completions.append(completion)
+            all_contexts.append(context)
+            all_prompt_ids.append(prompt_ids)
+            all_response_ids.append(response_ids)
+            all_response_masks.append(response_mask)
+            all_response_old_logps.append(response_old_logps)
+            all_turn_outputs.append(turn_outputs)
+            all_unfinished.append(unfinished)
+    return all_completions, all_contexts, all_prompt_ids, all_response_ids, all_response_masks, all_response_old_logps, all_turn_outputs, all_unfinished
+
+# ======== Reward 计算 ========
+def validate_gt_in_text(text, gt_list):
+    text, text_num = str(text), str(text).replace(',', '')
+    nums = [float(x) for x in re.findall(r'(?<![\w.])[-+]?\d+(?:\.\d+)?(?![\w.])', text_num)]
+    return {g for g in gt_list if ((s := str(g).strip()) and s.lower() in text.lower()) or (re.fullmatch(r'[-+]?\d+(?:\.\d+)?', str(g).strip().replace(',', '')) and any(abs(float(str(g).strip().replace(',', '')) - n) < 1e-6 for n in nums))}
+
+def calculate_rewards(prompts, completions, gt_batch, tools_batch, num_gen, reward_model=None, device="cuda", turn_outputs_batch=None, unfinished_batch=None):
+    rewards = torch.zeros(len(completions), device=device)
+    for idx, response in enumerate(completions):
+        reward, answer = 0.0, response
+        sample_idx = idx // num_gen
+        tools = tools_batch[sample_idx]
+        turn_outputs = turn_outputs_batch[idx] if turn_outputs_batch is not None else [response]
+        unfinished = unfinished_batch[idx] if unfinished_batch is not None else False
+        turn_answers = [turn.split('</think>', 1)[-1].strip() if '</think>' in turn else turn.strip() for turn in turn_outputs]
+        answer = turn_answers[-1] if turn_answers else response.strip()
+        valid_names = {t['function']['name'] for t in tools} if tools else set()
+        tool_calls = []
+        for turn_answer in turn_answers: tool_calls.extend(parse_tool_calls(turn_answer))  # 解析tool调用
+        reward -= 0.5 * sum(abs(turn.count('<tool_call>') - turn.count('</tool_call>')) for turn in turn_answers)  # 标签扣分
+        # -------- 无工具调用：格式+reward奖励 --------
+        if not tool_calls:
+            reward += 0.5 if 5 <= len(response.strip()) <= 800 else -0.5  # 长度分
+            if '</think>' in response:
+                think, answer = response.split('</think>', 1)
+                reward += 1.0 if 20 <= len(think.strip()) <= 300 else -0.5  # 思考长度分
+                reward += 0.25 if response.count('</think>') == 1 else -0.25  # 思考闭合分
+                answer = answer.strip()
+            if reward_model is not None:
+                prompt = prompts[sample_idx]
+                pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
+                matches = re.findall(pattern, prompt, re.DOTALL)
+                messages = [{"role": role, "content": content.strip()} for role, content in matches]
+                score = reward_model.get_score(messages, answer)
+                reward += score  # RM分
+            reward -= rep_penalty(answer)
+            rewards[idx] = max(min(reward, 3.0), -3.0)  # 总分Clip
+        # -------- 有工具调用：执行结果奖励 --------
+        else:
+            gt = gt_batch[sample_idx]
+            valid_call_count = 0
+            for tool_call in tool_calls:
+                name, raw = tool_call.get("name", ""), tool_call.get("arguments", {})
+                if isinstance(raw, str):
+                    try: raw = json.loads(raw)
+                    except: raw = {}
+                check = CHECK_ARGS.get(name)
+                valid_call_count += int(bool(name in valid_names and check and check(raw)))
+            tool_gap = abs(valid_call_count - len(gt)) + max(0, len(tool_calls) - valid_call_count)  # tool数差值
+            reward += 0.5 if tool_gap == 0 else -0.5 * tool_gap  # tool对齐分
+            
+            final_text = "" if unfinished else (answer.split('</tool_call>')[-1] if '</tool_call>' in answer else answer)
+            verified = validate_gt_in_text(final_text, gt) if gt else set()
+            if gt: reward += 2.5 * len(verified) / len(gt)  # GT分
+            if unfinished: reward -= 0.5  # 未完成扣分
+            reward -= rep_penalty(final_text if final_text else answer)
+            rewards[idx] = max(min(reward, 3.0), -3.0)  # 总分Clip
+    return rewards
+
+# ================================ 工具与 Reward = End ================================
+def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model=None, start_step=0, wandb=None, use_sglang=False):
+    last_step = start_step
+    for step, batch in enumerate(loader, start=start_step + 1):
+        messages_batch = batch['messages']
+        tools_batch = batch['tools']
+        gt_batch = batch['gt']
+        last_step = step
+
+        with torch.no_grad():
+            completions, contexts, prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch, turn_outputs_batch, unfinished_batch = rollout_batch(rollout_engine, tokenizer, messages_batch, tools_batch, args.num_generations, max_turns=3, max_new_tokens=args.max_gen_len, thinking_ratio=args.thinking_ratio, device=args.device)
+
+        prompts = [tokenizer.apply_chat_template(m, tokenize=False, add_generation_prompt=True, tools=t) for m, t in zip(messages_batch, tools_batch)]
+        packed_samples = []
+        for p, r, m, old_lp in zip(prompt_ids_batch, response_ids_batch, response_masks_batch, response_old_logps_batch):
+            ids = p + r
+            mask = [0] * len(p) + m
+            old_logps = [0.0] * max(len(p) - 1, 0) + old_lp
+            if len(ids) > args.max_total_len:
+                ids = ids[-args.max_total_len:]
+                mask = mask[-args.max_total_len:]
+                old_logps = old_logps[-(len(ids) - 1):]
+            prompt_len = next((i for i, v in enumerate(mask) if v == 1), len(mask))
+            packed_samples.append((ids, mask, prompt_len, old_logps))
+        seq_lens = torch.tensor([len(ids) for ids, _, _, _ in packed_samples], device=args.device)
+        max_len = seq_lens.max().item()
+        input_ids = torch.tensor([ids + [tokenizer.pad_token_id] * (max_len - len(ids)) for ids, _, _, _ in packed_samples], device=args.device)
+        prompt_lens = torch.tensor([prompt_len for _, _, prompt_len, _ in packed_samples], device=args.device)
+        full_response_masks = torch.tensor([mask + [0] * (max_len - len(mask)) for _, mask, _, _ in packed_samples], device=args.device, dtype=torch.float32)
+        old_per_token_logps = torch.tensor([old_logps + [0.0] * ((max_len - 1) - len(old_logps)) for _, _, _, old_logps in packed_samples], device=args.device, dtype=torch.float32)
+        full_mask = (input_ids != tokenizer.pad_token_id).long()
+
+        rewards = calculate_rewards(prompts, completions, gt_batch, tools_batch, args.num_generations, reward_model, device=args.device, turn_outputs_batch=turn_outputs_batch, unfinished_batch=unfinished_batch)
+
+        model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
+        with autocast_ctx:
+            res = model_unwrapped(input_ids, attention_mask=full_mask)
+            aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
+            logits = res.logits[:, :-1, :]
+            per_token_logps = F.log_softmax(logits, dim=-1).gather(2, input_ids[:, 1:].unsqueeze(-1)).squeeze(-1)
+
+        with torch.no_grad():
+            ref_per_token_logps = compute_per_token_logps(ref_model, input_ids, input_ids.size(1) - 1, attention_mask=full_mask)
+
+        completion_mask = full_response_masks[:, 1:]
+        is_eos = (input_ids[:, 1:] == tokenizer.eos_token_id) & completion_mask.bool()
+        eos_idx = torch.full((completion_mask.size(0),), completion_mask.size(1) - 1, device=args.device, dtype=torch.long)
+        has_eos = is_eos.any(dim=1)
+        eos_idx[has_eos] = is_eos.int().argmax(dim=1)[has_eos]
+        pos = torch.arange(completion_mask.size(1), device=args.device).unsqueeze(0)
+        completion_mask = completion_mask * (pos <= eos_idx.unsqueeze(1)).float()
+        token_counts = completion_mask.sum(dim=1)
+        valid_rows = token_counts > 0
+
+        if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
+            for i in range(len(messages_batch)):
+                Logger(f"[DEBUG] step={step}, gt[{i}]: {repr(gt_batch[i])}")
+                Logger('-'*100)
+                for j in range(args.num_generations):
+                    idx = i * args.num_generations + j
+                    plen, slen = prompt_lens[idx].item(), seq_lens[idx].item()
+                    Logger(f"{'=' * 30} [DEBUG] gen[{i}][{j}] CONTEXT_BEGIN {'=' * 30}")
+                    Logger(contexts[idx])
+                    Logger(f"{'=' * 31} [DEBUG] gen[{i}][{j}] CONTEXT_END {'=' * 31}")
+                    Logger(f"[DEBUG] gen[{i}][{j}] prompt_len={plen}, seq_len={slen}")
+                    tokens = input_ids[idx, plen:slen].tolist()
+                    text = tokenizer.decode(tokens, skip_special_tokens=False)
+                    Logger(f"{'=' * 28} [DEBUG] gen[{i}][{j}] COMPLETION_BEGIN [{plen}:{slen}] {'=' * 28}")
+                    Logger(text)
+                    Logger(f"{'=' * 29} [DEBUG] gen[{i}][{j}] COMPLETION_END {'=' * 29}")
+                    Logger(f"[DEBUG] gen[{i}][{j}] reward={rewards[idx].item():.4f}")
+                    Logger('='*100)
+
+        grouped_rewards = rewards.view(-1, args.num_generations)
+        mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations)
+        std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(args.num_generations)
+        advantages = (rewards - mean_r) / (std_r + 1e-4)
+
+        kl_div = ref_per_token_logps - per_token_logps
+        per_token_kl = torch.exp(kl_div) - kl_div - 1
+        ratio = torch.exp(per_token_logps - old_per_token_logps)
+        if args.loss_type == "cispo":
+            clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()
+            per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * per_token_logps - args.beta * per_token_kl)
+        else:
+            clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
+            per_token_loss1 = ratio * advantages.unsqueeze(1)
+            per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
+            per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
+        policy_loss = (((per_token_loss * completion_mask).sum(dim=1)[valid_rows] / token_counts[valid_rows].clamp(min=1)).mean()
+                       if valid_rows.any() else per_token_loss.sum() * 0.0)
+        loss = (policy_loss + aux_loss) / args.accumulation_steps
+        loss.backward()
+
+        if step % args.accumulation_steps == 0:
+            if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            optimizer.step(); scheduler.step(); optimizer.zero_grad()
+
+        if step % args.log_interval == 0 or step == iters:
+            pl = loss.item() * args.accumulation_steps
+            ar = rewards.mean().item()
+            al = token_counts.float().mean().item()
+            kl = ((ref_per_token_logps - per_token_logps) * completion_mask).sum().item() / max(token_counts.sum().item(), 1)
+            gs = grouped_rewards.std(dim=1, unbiased=False).mean().item()
+            am, ast = advantages.mean().item(), advantages.std().item()
+            lr = optimizer.param_groups[0]['lr']
+            Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}), Reward:{ar:.4f}, KL:{kl:.4f}, GrpStd:{gs:.4f}, AdvStd:{ast:.4f}, Loss:{pl:.4f}, AvgLen:{al:.2f}, AdvMean:{am:.4f}, LR:{lr:.8f}')
+            if wandb and is_main_process():
+                wandb.log({"reward":ar,"kl_ref":kl,"group_reward_std":gs,"advantages_std":ast,"policy_loss":pl,"avg_response_len":al,"advantages_mean":am,"learning_rate":lr})
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if lm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            state_dict = raw_model.state_dict()
+            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
+                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler)
+            model.train()
+            del state_dict
+
+        if step % args.save_interval == 0 or step == iters: rollout_engine.update_policy(model)
+
+        del per_token_logps, ref_per_token_logps
+        del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask
+
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        optimizer.step(); scheduler.step(); optimizer.zero_grad()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind Agent RL")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='agent', type=str, help="保存权重名称")
+    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=2, help="批次大小")
+    parser.add_argument("--learning_rate", type=float, default=3e-7, help="学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型 bfloat16/float16")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=10, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="模型隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="模型层数")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE")
+    parser.add_argument('--max_seq_len', default=1024, type=int, help="最大序列长度")
+    parser.add_argument("--max_gen_len", type=int, default=768, help="单次最大生成长度")
+    parser.add_argument("--max_total_len", type=int, default=2500, help="训练侧最终总长度上界")
+    parser.add_argument("--data_path", type=str, default="../dataset/agent_rl.jsonl", help="训练数据路径")
+    parser.add_argument("--num_generations", type=int, default=4, help="每个prompt生成数量")
+    parser.add_argument("--beta", type=float, default=0.1, help="KL散度惩罚系数")
+    parser.add_argument("--loss_type", type=str, default="cispo", choices=["grpo", "cispo"], help="loss类型")
+    parser.add_argument("--epsilon", type=float, default=0.2, help="GRPO的PPO clip epsilon")
+    parser.add_argument("--epsilon_high", type=float, default=5.0, help="epsilon上界")
+    parser.add_argument('--from_weight', default='full_sft', type=str, help="加载预训练权重名称")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否从checkpoint恢复")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb记录")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-Agent-RL", help="wandb项目名称")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile")
+    parser.add_argument("--debug_mode", action="store_true", help="调试模式")
+    parser.add_argument("--debug_interval", type=int, default=20, help="调试日志间隔")
+    parser.add_argument("--thinking_ratio", type=float, default=0.1, help="按概率开启thinking（0.0~1.0）")
+    parser.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward", help="Reward模型路径")
+    parser.add_argument("--rollout_engine", type=str, default="torch", choices=["torch", "sglang"], help="rollout引擎类型")
+    parser.add_argument("--sglang_base_url", type=str, default="http://localhost:8998", help="SGLang服务器URL")
+    parser.add_argument("--sglang_model_path", type=str, default="../model", help="SGLang tokenizer路径")
+    parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_agent", help="SGLang共享存储路径")
+    args = parser.parse_args()
+
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+
+    os.makedirs(args.save_dir, exist_ok=True)
+    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+                               max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
+    ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None
+
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb.init(project=args.wandb_project, name=f"Agent-RL-E{args.epochs}-B{args.batch_size}-LR{args.learning_rate}", id=wandb_id, resume=resume)
+
+    model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
+
+    ref_model, _ = init_model(lm_config, args.from_weight, device=args.device)
+    ref_model = ref_model.eval().requires_grad_(False)
+
+    reward_model = LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16)
+    Logger(f'Loaded reward model from {args.reward_model_path}')
+    # Rollout引擎
+    rollout_engine = create_rollout_engine(
+        engine_type=args.rollout_engine,
+        policy_model=model,
+        tokenizer=tokenizer,
+        device=args.device,
+        autocast_ctx=autocast_ctx,
+        sglang_base_url=args.sglang_base_url,
+        sglang_model_path=args.sglang_model_path,
+        sglang_shared_path=args.sglang_shared_path,
+    )
+    train_ds = AgentRLDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    def collate_fn(batch): return {'messages': [b['messages'] for b in batch], 'tools': [b['tools'] for b in batch], 'gt': [b['gt'] for b in batch]}
+    loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler, collate_fn=collate_fn)
+    iters = len(loader_for_count)
+    total_optimizer_steps = math.ceil(iters / args.accumulation_steps) * args.epochs
+    scheduler = CosineAnnealingLR(optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
+
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'])
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scheduler.load_state_dict(ckp_data['scheduler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
+        rollout_engine.update_policy(model)
+    if dist.is_initialized():
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+    rollout_engine.update_policy(model)
+
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate_fn)
+        if skip > 0:
+            Logger(f'Epoch [{epoch+1}/{args.epochs}]: skip {start_step} steps')
+            rl_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
+        else:
+            rl_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
+
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

src/omni/trainers/distillation.py

@@ -0,0 +1,245 @@
+import os
+
+
+import datasets  # noqa: F401  # Windows pyarrow/torch DLL conflict workaround (issue #771)
+import argparse
+import time
+import warnings
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+from contextlib import nullcontext
+from torch import optim
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from omni.models.minimind import MiniMindConfig
+from omni.datasets.lm_dataset import SFTDataset
+from omni.utils.training import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
+
+warnings.filterwarnings('ignore')
+
+
+def distillation_loss(student_logits, teacher_logits, temperature=1.0, reduction='batchmean'):
+    with torch.no_grad():
+        teacher_probs = F.softmax(teacher_logits / temperature, dim=-1).detach()
+
+    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)
+
+    kl = F.kl_div(
+        student_log_probs,
+        teacher_probs,
+        reduction=reduction
+    )
+    return (temperature ** 2) * kl
+
+
+def train_epoch(epoch, loader, iters, teacher_model, lm_config_student, start_step=0, wandb=None, alpha=0.0, temperature=1.0):
+    start_time = time.time()
+    last_step = start_step
+    
+    if teacher_model is not None:
+        teacher_model.eval()
+        teacher_model.requires_grad_(False)
+
+    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
+        last_step = step
+        input_ids = input_ids.to(args.device)
+        labels = labels.to(args.device)
+        loss_mask = (labels[..., 1:] != -100).float()
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        # 前向传播（学生模型）
+        with autocast_ctx:
+            res = model(input_ids)
+            student_logits = res.logits[..., :-1, :].contiguous()
+
+        # 教师模型前向传播（只在eval & no_grad）
+        if teacher_model is not None:
+            with torch.no_grad():
+                teacher_logits = teacher_model(input_ids).logits[..., :-1, :].contiguous()
+                vocab_size_student = student_logits.size(-1)
+                teacher_logits = teacher_logits[..., :vocab_size_student]
+
+        # ========== 计算损失 ==========
+        # 1) Ground-Truth CE Loss
+        shift_labels = labels[..., 1:].contiguous()
+        loss_mask_flat = loss_mask.view(-1)
+        ce_loss = F.cross_entropy(
+            student_logits.view(-1, student_logits.size(-1)),
+            shift_labels.view(-1),
+            ignore_index=-100,
+            reduction='none'
+        )
+        ce_loss_raw = torch.sum(ce_loss * loss_mask_flat) / (loss_mask_flat.sum() + 1e-8)
+        if lm_config_student.use_moe: ce_loss = ce_loss_raw + res.aux_loss
+        else: ce_loss = ce_loss_raw
+
+        # 2) Distillation Loss
+        if teacher_model is not None:
+            distill_loss = distillation_loss(
+                student_logits.view(-1, student_logits.size(-1))[loss_mask_flat == 1],
+                teacher_logits.view(-1, teacher_logits.size(-1))[loss_mask_flat == 1],
+                temperature=temperature
+            )
+        else:
+            distill_loss = torch.tensor(0.0, device=args.device)
+
+        # 3) 总损失 = alpha * CE + (1-alpha) * Distill
+        loss = (alpha * ce_loss + (1 - alpha) * distill_loss) / args.accumulation_steps
+
+        scaler.scale(loss).backward()
+
+        if step % args.accumulation_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad(set_to_none=True)
+
+        if step % args.log_interval == 0 or step == iters:
+            spend_time = time.time() - start_time
+            current_loss = loss.item() * args.accumulation_steps
+            current_ce_loss = ce_loss_raw.item()
+            current_aux_loss = res.aux_loss.item() if lm_config_student.use_moe else 0.0
+            current_lr = optimizer.param_groups[-1]['lr']
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
+            
+            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, ce: {current_ce_loss:.4f}, aux_loss: {current_aux_loss:.4f}, distill: {distill_loss.item():.4f}, learning_rate: {current_lr:.8f}, epoch_time: {eta_min:.3f}min')
+            
+            if wandb:
+                wandb.log({
+                    "loss": current_loss,
+                    "ce_loss": current_ce_loss,
+                    "aux_loss": current_aux_loss,
+                    "distill_loss": distill_loss.item() if teacher_model is not None else 0.0,
+                    "learning_rate": current_lr,
+                    "epoch_time": eta_min
+                })
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if lm_config_student.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config_student.hidden_size}{moe_suffix}.pth'
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            state_dict = raw_model.state_dict()
+            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+            lm_checkpoint(lm_config_student, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
+            model.train()
+            del state_dict
+
+        del input_ids, labels, loss_mask, res, student_logits, ce_loss, distill_loss, loss
+
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
+
+if __name__ == "__main__":
+    # 模拟用moe模型蒸馏dense模型，也可以用更大teacher_hidden_size模型蒸馏更小student_hidden_size的
+    parser = argparse.ArgumentParser(description="MiniMind Knowledge Distillation")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='full_dist', type=str, help="保存权重的前缀名")
+    parser.add_argument("--epochs", type=int, default=6, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=32, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=5e-6, help="初始学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
+    parser.add_argument("--max_seq_len", type=int, default=340, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
+    parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl", help="训练数据路径")
+    parser.add_argument('--student_hidden_size', default=768, type=int, help="学生模型隐藏层维度")
+    parser.add_argument('--student_num_layers', default=8, type=int, help="学生模型隐藏层数量")
+    parser.add_argument('--teacher_hidden_size', default=768, type=int, help="教师模型隐藏层维度")
+    parser.add_argument('--teacher_num_layers', default=8, type=int, help="教师模型隐藏层数量")
+    parser.add_argument('--student_use_moe', default=0, type=int, choices=[0, 1], help="学生模型是否使用MoE（0=否，1=是）")
+    parser.add_argument('--teacher_use_moe', default=1, type=int, choices=[0, 1], help="教师模型是否使用MoE（0=否，1=是）")
+    parser.add_argument('--from_student_weight', default='full_sft', type=str, help="学生模型基于哪个权重")
+    parser.add_argument('--from_teacher_weight', default='full_sft', type=str, help="教师模型基于哪个权重")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
+    parser.add_argument('--alpha', default=0.5, type=float, help="CE损失权重，总损失=alpha*CE+(1-alpha)*KL")
+    parser.add_argument('--temperature', default=1.5, type=float, help="蒸馏温度（推荐范围1.0-2.0）")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-Distillation", help="wandb项目名")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+    
+    # ========== 2. 配置目录、模型参数、检查ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    lm_config_student = MiniMindConfig(hidden_size=args.student_hidden_size, num_hidden_layers=args.student_num_layers, use_moe=bool(args.student_use_moe))
+    lm_config_teacher = MiniMindConfig(hidden_size=args.teacher_hidden_size, num_hidden_layers=args.teacher_num_layers, use_moe=bool(args.teacher_use_moe))
+    ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
+    
+    # ========== 3. 设置混合精度 ==========
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+    
+    # ========== 4. 配wandb ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb_run_name = f"MiniMind-Distill-S{args.student_hidden_size}T{args.teacher_hidden_size}-Epoch-{args.epochs}-BS-{args.batch_size}-LR-{args.learning_rate}"
+        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
+    
+    # ========== 5. 定义学生和教师模型 ==========
+    model, tokenizer = init_model(lm_config_student, args.from_student_weight, device=args.device)
+    Logger(f'学生模型总参数量：{sum(p.numel() for p in model.parameters()) / 1e6:.3f} M')
+    teacher_model, _ = init_model(lm_config_teacher, args.from_teacher_weight, device=args.device)
+    teacher_model.eval()
+    teacher_model.requires_grad_(False)
+    Logger(f'教师模型总参数量：{sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f} M')
+    train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    
+    # ========== 6. 从ckp恢复状态 ==========
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'])
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scaler.load_state_dict(ckp_data['scaler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+    
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
+    if dist.is_initialized():
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+    
+    # ========== 8. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
+        if skip > 0: 
+            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
+            train_epoch(epoch, loader, len(loader) + skip, teacher_model, lm_config_student, start_step, wandb, args.alpha, args.temperature)
+        else:
+            train_epoch(epoch, loader, len(loader), teacher_model, lm_config_student, 0, wandb, args.alpha, args.temperature)
+    
+    # ========== 9. 清理分布进程 ==========
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

src/omni/trainers/dpo.py

@@ -0,0 +1,225 @@
+import os
+
+
+import datasets  # noqa: F401  # Windows pyarrow/torch DLL conflict workaround (issue #771)
+import argparse
+import time
+import warnings
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+from contextlib import nullcontext
+from torch import optim
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from omni.models.minimind import MiniMindConfig
+from omni.datasets.lm_dataset import DPODataset
+from omni.utils.training import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
+
+warnings.filterwarnings('ignore')
+
+
+def logits_to_log_probs(logits, labels):
+    # logits shape: (batch_size, seq_len, vocab_size)
+    # labels shape: (batch_size, seq_len)
+    # log_probs shape: (batch_size, seq_len)
+    log_probs = F.log_softmax(logits, dim=2)
+    log_probs_per_token = torch.gather(log_probs, dim=2, index=labels.unsqueeze(2)).squeeze(-1)
+    return log_probs_per_token
+
+
+def dpo_loss(ref_log_probs, policy_log_probs, mask, beta):
+    # ref_log_probs 和 policy_log_probs 都是 shape: (batch_size, seq_len)
+    ref_log_probs = (ref_log_probs * mask).sum(dim=1)
+    policy_log_probs = (policy_log_probs * mask).sum(dim=1)
+
+    # 将 chosen 和 rejected 数据分开
+    batch_size = ref_log_probs.shape[0]
+    chosen_ref_log_probs = ref_log_probs[:batch_size // 2]
+    reject_ref_log_probs = ref_log_probs[batch_size // 2:]
+    chosen_policy_log_probs = policy_log_probs[:batch_size // 2]
+    reject_policy_log_probs = policy_log_probs[batch_size // 2:]
+
+    pi_logratios = chosen_policy_log_probs - reject_policy_log_probs
+    ref_logratios = chosen_ref_log_probs - reject_ref_log_probs
+    logits = pi_logratios - ref_logratios
+    loss = -F.logsigmoid(beta * logits)
+    return loss.mean()
+
+
+def train_epoch(epoch, loader, iters, ref_model, lm_config, start_step=0, wandb=None, beta=0.1):
+    start_time = time.time()
+    last_step = start_step
+
+    for step, batch in enumerate(loader, start=start_step + 1):
+        last_step = step
+        x_chosen = batch['x_chosen'].to(args.device)
+        x_rejected = batch['x_rejected'].to(args.device)
+        y_chosen = batch['y_chosen'].to(args.device)
+        y_rejected = batch['y_rejected'].to(args.device)
+        mask_chosen = batch['mask_chosen'].to(args.device)
+        mask_rejected = batch['mask_rejected'].to(args.device)
+        x = torch.cat([x_chosen, x_rejected], dim=0)
+        y = torch.cat([y_chosen, y_rejected], dim=0)
+        mask = torch.cat([mask_chosen, mask_rejected], dim=0)
+
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        with autocast_ctx:
+            with torch.no_grad():
+                ref_outputs = ref_model(x)
+                ref_logits = ref_outputs.logits
+            ref_log_probs = logits_to_log_probs(ref_logits, y)
+            
+            outputs = model(x)
+            logits = outputs.logits
+            policy_log_probs = logits_to_log_probs(logits, y)
+            
+            dpo_loss_val = dpo_loss(ref_log_probs, policy_log_probs, mask, beta=beta)
+            loss = dpo_loss_val + outputs.aux_loss
+            loss = loss / args.accumulation_steps
+
+        scaler.scale(loss).backward()
+
+        if step % args.accumulation_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad(set_to_none=True)
+
+        if step % args.log_interval == 0 or step == iters:
+            spend_time = time.time() - start_time
+            current_loss = loss.item() * args.accumulation_steps
+            current_dpo_loss = dpo_loss_val.item()
+            current_aux_loss = outputs.aux_loss.item()
+            current_lr = optimizer.param_groups[-1]['lr']
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
+            
+            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, dpo_loss: {current_dpo_loss:.4f}, aux_loss: {current_aux_loss:.4f}, learning_rate: {current_lr:.8f}, epoch_time: {eta_min:.3f}min')
+            
+            if wandb: wandb.log({"loss": current_loss, "dpo_loss": current_dpo_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if lm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            state_dict = raw_model.state_dict()
+            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
+            model.train()
+            del state_dict
+
+        del x_chosen, x_rejected, y_chosen, y_rejected, mask_chosen, mask_rejected, x, y, mask
+        del ref_outputs, ref_logits, ref_log_probs, outputs, logits, policy_log_probs, loss
+
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind DPO (Direct Preference Optimization)")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='dpo', type=str, help="保存权重的前缀名")
+    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=4, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=4e-8, help="初始学习率（建议<=5e-8避免遗忘）")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--max_seq_len', default=1024, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument("--data_path", type=str, default="../dataset/dpo.jsonl", help="DPO训练数据路径")
+    parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
+    parser.add_argument('--beta', default=0.15, type=float, help="DPO中的beta参数")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-DPO", help="wandb项目名")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+    
+    # ========== 2. 配置目录、模型参数、检查ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
+    ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
+    
+    # ========== 3. 设置混合精度 ==========
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+    
+    # ========== 4. 配wandb ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb_run_name = f"MiniMind-DPO-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LR-{args.learning_rate}"
+        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
+    
+    # ========== 5. 定义模型和参考模型 ==========
+    model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
+    Logger(f'策略模型总参数量：{sum(p.numel() for p in model.parameters()) / 1e6:.3f} M')
+    # 初始化参考模型（ref_model冻结）
+    ref_model, _ = init_model(lm_config, args.from_weight, device=args.device)
+    ref_model.eval()
+    ref_model.requires_grad_(False)
+    Logger(f'参考模型总参数量：{sum(p.numel() for p in ref_model.parameters()) / 1e6:.3f} M')
+    
+    train_ds = DPODataset(args.data_path, tokenizer, max_length=args.max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    
+    # ========== 6. 从ckp恢复状态 ==========
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'])
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scaler.load_state_dict(ckp_data['scaler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+    
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
+    if dist.is_initialized():
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+    
+    # ========== 8. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
+        if skip > 0: 
+            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
+            train_epoch(epoch, loader, len(loader) + skip, ref_model, lm_config, start_step, wandb, args.beta)
+        else:
+            train_epoch(epoch, loader, len(loader), ref_model, lm_config, 0, wandb, args.beta)
+    
+    # ========== 9. 清理分布进程 ==========
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

src/omni/trainers/full_sft.py

@@ -0,0 +1,170 @@
+import os
+
+
+import datasets  # noqa: F401  # Windows pyarrow/torch DLL conflict workaround (issue #771)
+import argparse
+import time
+import warnings
+import torch
+import torch.distributed as dist
+from contextlib import nullcontext
+from torch import optim, nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from omni.models.minimind import MiniMindConfig
+from omni.datasets.lm_dataset import SFTDataset
+from omni.utils.training import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
+
+warnings.filterwarnings('ignore')
+
+
+def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
+    start_time = time.time()
+    last_step = start_step
+    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
+        input_ids = input_ids.to(args.device)
+        labels = labels.to(args.device)
+        last_step = step
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        with autocast_ctx:
+            res = model(input_ids, labels=labels)
+            loss = res.loss + res.aux_loss
+            loss = loss / args.accumulation_steps
+
+        scaler.scale(loss).backward()
+
+        if step % args.accumulation_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+
+            scaler.step(optimizer)
+            scaler.update()
+
+            optimizer.zero_grad(set_to_none=True)
+
+        if step % args.log_interval == 0 or step == iters:
+            spend_time = time.time() - start_time
+            current_loss = loss.item() * args.accumulation_steps
+            current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
+            current_logits_loss = current_loss - current_aux_loss
+            current_lr = optimizer.param_groups[-1]['lr']
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
+            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
+            if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if lm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            state_dict = raw_model.state_dict()
+            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, 
+                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
+            model.train()
+            del state_dict
+
+        del input_ids, labels, res, loss
+
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind Full SFT")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='full_sft', type=str, help="保存权重的前缀名")
+    parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=16, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=1e-5, help="初始学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--max_seq_len', default=768, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl", help="训练数据路径")
+    parser.add_argument('--from_weight', default='pretrain', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT", help="wandb项目名")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+    
+    # ========== 2. 配置目录、模型参数、检查ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
+    ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
+    
+    # ========== 3. 设置混合精度 ==========
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+    
+    # ========== 4. 配wandb ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb_run_name = f"MiniMind-Full-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
+    
+    # ========== 5. 定义模型、数据、优化器 ==========
+    model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
+    train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    
+    # ========== 6. 从ckp恢复状态 ==========
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'])
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scaler.load_state_dict(ckp_data['scaler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+    
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
+    if dist.is_initialized():
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+    
+    # ========== 8. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
+        if skip > 0: 
+            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
+            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
+        else:
+            train_epoch(epoch, loader, len(loader), 0, wandb)
+    
+    # ========== 9. 清理分布进程 ==========
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

src/omni/trainers/full_sft_omni.py

@@ -0,0 +1,257 @@
+import os
+import argparse
+import time
+import warnings
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+from contextlib import nullcontext
+from torch import optim
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from omni.models import OmniConfig, MiniMindOmni
+from omni.datasets.lm_dataset import OmniDataset
+from omni.utils import get_lr, Logger, is_main_process, init_distributed_mode, setup_seed, init_omni_model, omni_checkpoint, SkipBatchSampler, log_model_params
+
+warnings.filterwarnings('ignore')
+
+
+def omni_collate_fn(batch):
+    """自定义collate函数，处理变长audio_inputs和pixel_values"""
+    input_ids, labels, audio_labels, audio_inputs, audio_lens, pixel_values, spk_emb = zip(*batch)
+    input_ids = torch.stack(input_ids)
+    labels = torch.stack(labels)
+    audio_labels = torch.stack(audio_labels)
+    audio_lens = torch.tensor(audio_lens, dtype=torch.long)
+    valid_audios = [a for a in audio_inputs if a is not None]
+    if valid_audios:
+        max_t = max(a.size(1) for a in valid_audios)
+        padded = [a if a.size(1) == max_t else torch.nn.functional.pad(a, (0, 0, 0, max_t - a.size(1))) for a in valid_audios]
+        audio_inputs = torch.cat(padded, dim=0)
+    else:
+        audio_inputs = None
+    valid_images = [p for p in pixel_values if p is not None]
+    if valid_images:
+        if hasattr(valid_images[0], 'keys'):
+            keys = set.intersection(*[set(d.keys()) for d in valid_images])
+            pixel_values = {k: torch.cat([d[k] for d in valid_images], dim=0) for k in keys}
+        else:
+            pixel_values = torch.cat(valid_images, dim=0)
+    else:
+        pixel_values = None
+    spk_emb = torch.stack(spk_emb)
+    return input_ids, labels, audio_labels, audio_inputs, audio_lens, pixel_values, spk_emb
+
+
+def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
+    start_time = time.time()
+    last_step = start_step
+    for step, (input_ids, labels, audio_labels, audio_inputs, audio_lens, pixel_values, spk_emb) in enumerate(loader, start=start_step + 1):
+        input_ids = input_ids.to(args.device)
+        labels = labels.to(args.device)
+        audio_labels = audio_labels.to(args.device)
+        audio_lens = audio_lens.to(args.device)
+        if audio_inputs is not None:
+            audio_inputs = audio_inputs.to(args.device)
+        if pixel_values is not None:
+            if hasattr(pixel_values, 'keys'):
+                pixel_values = {k: v.to(args.device) for k, v in pixel_values.items()}
+            else:
+                pixel_values = pixel_values.to(args.device)
+        spk_emb = spk_emb.to(args.device)
+        last_step = step
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        with autocast_ctx:
+            res = model(input_ids, audio_inputs=audio_inputs, audio_lens=audio_lens, pixel_values=pixel_values, spk_emb=spk_emb)
+            loss_fct = nn.CrossEntropyLoss(reduction='none')
+
+            # Text loss
+            text_loss_raw = loss_fct(res.logits.view(-1, res.logits.size(-1)), labels.view(-1))
+            text_mask = (labels.view(-1) != -100).float()
+            text_loss = (text_loss_raw * text_mask).sum() / (text_mask.sum() + 1e-9)
+
+            # Audio loss
+            audio_loss = res.audio_logits[0].sum() * 0
+            for i, al in enumerate(res.audio_logits):
+                al_flat = al.view(-1, al.size(-1))
+                target_flat = audio_labels[:, i, :].reshape(-1)
+                layer_loss = loss_fct(al_flat, target_flat)
+                valid_mask = (target_flat != -100).float()
+                stop_mask = (target_flat == 2050).float()
+                weighted_loss = layer_loss * valid_mask * (1 + stop_mask * 9)
+                msum = valid_mask.sum()
+                if msum > 0:
+                    audio_loss = audio_loss + weighted_loss.sum() / (msum + 1e-9)
+            audio_loss = audio_loss / 8
+
+            loss = (text_loss + audio_loss + res.aux_loss) / args.accumulation_steps
+
+        scaler.scale(loss).backward()
+        if step % args.accumulation_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad(set_to_none=True)
+
+        if step % args.log_interval == 0 or step == iters:
+            spend_time = time.time() - start_time
+            current_loss = loss.item() * args.accumulation_steps
+            text_loss_val = text_loss.item() if isinstance(text_loss, torch.Tensor) else 0
+            audio_loss_val = audio_loss.item() if isinstance(audio_loss, torch.Tensor) else 0
+            current_lr = optimizer.param_groups[-1]['lr']
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
+            Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, text: {text_loss_val:.4f}, audio: {audio_loss_val:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
+            if wandb:
+                wandb.log({"loss": current_loss, "text_loss": text_loss_val,
+                          "audio_loss": audio_loss_val, "lr": current_lr, "epoch_time": eta_min})
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if omni_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{omni_config.hidden_size}{moe_suffix}.pth'
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            clean_state_dict = {k: v for k, v in raw_model.state_dict().items() if not k.startswith('audio_encoder.')}
+            torch.save({k: v.half().cpu() for k, v in clean_state_dict.items()}, ckp)
+            omni_checkpoint(omni_config, weight=args.save_weight, model=model, optimizer=optimizer,
+                          epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
+            model.train()
+
+        del input_ids, labels, audio_labels, audio_inputs, audio_lens, pixel_values, spk_emb, res, loss
+
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind-O SFT")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='sft_omni', type=str, help="保存权重的前缀名")
+    parser.add_argument("--epochs", type=int, default=15, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=32, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
+    parser.add_argument("--num_workers", type=int, default=4, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--max_seq_len', default=512, type=int, help="训练的最大截断长度")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构")
+    parser.add_argument("--data_path", type=str, default="../dataset/train_t2a_mini.parquet", help="训练数据路径（parquet格式）")
+    parser.add_argument("--audio_encoder_dir", type=str, default="../model/SenseVoiceSmall", help="音频encoder路径(SenseVoice)")
+    parser.add_argument("--vision_dir", type=str, default="../model/siglip2-base-p32-256-ve", help="CLIP视觉模型路径")
+    parser.add_argument('--from_weight', default='llm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
+    parser.add_argument('--freeze_backbone', default='none', type=str, choices=['none', 'all', 'last1'], help="冻结主干模型: none=全量训练, all=只训练audio层, last1=只训练最后1层+audio层")
+    parser.add_argument('--mode', default='all', type=str, choices=['all', 'audio_proj', 'vision_proj'], help="训练模式: all=全量训练, audio_proj=只训练audio_proj, vision_proj=只训练vision_proj")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-O-SFT", help="wandb项目名")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized():
+        args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+
+    # ========== 2. 配置目录、模型参数、检查ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    omni_config = OmniConfig(
+        hidden_size=args.hidden_size,
+        num_hidden_layers=args.num_hidden_layers,
+        use_moe=bool(args.use_moe)
+    )
+    ckp_data = omni_checkpoint(omni_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
+
+    # ========== 3. 设置混合精度 ==========
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+
+    # ========== 4. 配wandb ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb_run_name = f"MiniMind-O-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LR-{args.learning_rate}"
+        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
+
+    # ========== 5. 定义模型、数据、优化器 ==========
+    model, tokenizer = init_omni_model(omni_config, from_weight=args.from_weight,
+                                        audio_encoder_path=args.audio_encoder_dir,
+                                        vision_model_path=args.vision_dir,
+                                        save_dir=args.save_dir, device=args.device,
+                                        freeze_backbone=args.freeze_backbone, from_resume=args.from_resume)
+
+    if args.use_compile == 1:
+        model = torch.compile(model)
+
+    if model.audio_encoder is not None: model.audio_encoder.to(args.device)
+    if model.vision_encoder is not None: model.vision_encoder.to(args.device)
+
+    if args.mode == 'audio_proj':
+        for p in model.parameters(): p.requires_grad = False
+        for p in model.audio_proj.parameters(): p.requires_grad = True
+    elif args.mode == 'vision_proj':
+        for p in model.parameters(): p.requires_grad = False
+        for p in model.vision_proj.parameters(): p.requires_grad = True
+    log_model_params(model)
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6
+    Logger(f'Trainable: {trainable:.2f}M | Mode: {args.mode} | Freeze: {args.freeze_backbone} | Compile: {"on" if args.use_compile else "off"}')
+
+    # scheduled_sampling 现在会自动保护 image/audio token 的连续性
+    train_ds = OmniDataset(
+        args.data_path,
+        tokenizer,
+        audio_processor=model.audio_processor,
+        vision_processor=model.vision_processor,
+        max_length=args.max_seq_len,
+        image_token_len=model.config.image_token_len
+    )
+
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+
+    # ========== 6. 从ckp恢复状态 ==========
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'], strict=False)
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scaler.load_state_dict(ckp_data['scaler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+
+    # ========== 7. DDP包模型 ==========
+    if dist.is_initialized():
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+
+    # ========== 8. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, collate_fn=omni_collate_fn, num_workers=args.num_workers, pin_memory=True)
+        if skip > 0:
+            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
+            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
+        else:
+            train_epoch(epoch, loader, len(loader), 0, wandb)
+
+    # ========== 9. 清理分布进程 ==========
+    if dist.is_initialized(): dist.destroy_process_group()

src/omni/trainers/full_sft_vlm.py

@@ -0,0 +1,175 @@
+import argparse
+import os
+import time
+import warnings
+import datasets
+import torch
+import torch.distributed as dist
+from contextlib import nullcontext
+from torch import optim, nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from transformers import AutoTokenizer
+from omni.models import MiniMindVLM, VLMConfig
+from omni.datasets.lm_dataset import VLMDataset
+from omni.utils import get_lr, Logger, is_main_process, init_distributed_mode, setup_seed, init_vlm_model, vlm_checkpoint, SkipBatchSampler, vlm_collate_fn
+
+warnings.filterwarnings('ignore')
+
+
+def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
+    start_time = time.time()
+    last_step = start_step
+    for step, (input_ids, labels, pixel_values) in enumerate(loader, start=start_step + 1):
+        input_ids = input_ids.to(args.device)
+        labels = labels.to(args.device)
+        pixel_values = {k: v.to(args.device) for k, v in pixel_values.items()} if isinstance(pixel_values, dict) else pixel_values.to(args.device)
+        last_step = step
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        with autocast_ctx:
+            res = model(input_ids, labels=labels, pixel_values=pixel_values)
+            loss = res.loss + res.aux_loss
+            loss = loss / args.accumulation_steps
+
+        scaler.scale(loss).backward()
+
+        if step % args.accumulation_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+
+            scaler.step(optimizer)
+            scaler.update()
+
+            optimizer.zero_grad(set_to_none=True)
+
+        if step % args.log_interval == 0 or step == iters:
+            spend_time = time.time() - start_time
+            current_loss = loss.item() * args.accumulation_steps
+            current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
+            current_logits_loss = current_loss - current_aux_loss
+            current_lr = optimizer.param_groups[-1]['lr']
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
+            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
+            if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if vlm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth'
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            state_dict = raw_model.state_dict()
+            clean_state_dict = {
+                key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.')
+            }
+            clean_state_dict = {k: v.half().cpu() for k, v in clean_state_dict.items()}  # 半精度保存并移到CPU
+            torch.save(clean_state_dict, ckp)
+            vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer,
+                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
+            model.train()
+            del state_dict, clean_state_dict
+
+        del input_ids, labels, pixel_values, res, loss
+
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind-V SFT")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='sft_vlm', type=str, help="保存权重的前缀名")
+    parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=4, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=5e-6, help="初始学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
+    parser.add_argument("--num_workers", type=int, default=2, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--max_seq_len', default=768, type=int, help="训练的最大截断长度")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument("--data_path", type=str, default="../dataset/sft_i2t.parquet", help="训练数据路径")
+    parser.add_argument('--from_weight', default='pretrain_vlm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
+    parser.add_argument('--freeze_llm', default=1, type=int, choices=[0, 1, 2], help="冻结策略（0=完全可训练，1=冻结+解冻首尾层，2=完全冻结仅训练proj）")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-V-SFT", help="wandb项目名")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+
+    # ========== 2. 配置目录、模型参数、检查ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    vlm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, max_seq_len=args.max_seq_len, use_moe=bool(args.use_moe))
+    ckp_data = vlm_checkpoint(vlm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
+
+    # ========== 3. 设置混合精度 ==========
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+
+    # ========== 4. 配wandb ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb_run_name = f"MiniMind-V-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
+
+    # ========== 5. 定义模型、数据、优化器 ==========
+    model, tokenizer = init_vlm_model(vlm_config, from_weight=args.from_weight, device=args.device, freeze_llm=args.freeze_llm)
+    preprocess = model.vision_encoder.processor
+    train_ds = VLMDataset(args.data_path, tokenizer, preprocess=preprocess, image_special_token=vlm_config.image_special_token, image_token_len=vlm_config.image_token_len, max_length=vlm_config.max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+
+    # ========== 6. 从ckp恢复状态 ==========
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'], strict=False)
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scaler.load_state_dict(ckp_data['scaler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
+    if dist.is_initialized():
+        model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+
+    # ========== 8. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=vlm_collate_fn)
+        if skip > 0:
+            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
+            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
+        else:
+            train_epoch(epoch, loader, len(loader), 0, wandb)
+
+    # ========== 9. 清理分布进程 ==========
+    if dist.is_initialized(): dist.destroy_process_group()

src/omni/trainers/grpo.py

@@ -0,0 +1,331 @@
+import os
+
+
+import datasets  # noqa: F401  # Windows pyarrow/torch DLL conflict workaround (issue #771)
+import argparse
+import math
+import re
+import gc
+import warnings
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+from transformers import AutoTokenizer
+from contextlib import nullcontext
+from torch import optim
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from transformers import AutoModel
+from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.datasets.lm_dataset import RLAIFDataset
+from omni.utils.training import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from omni.trainers.rollout_engine import create_rollout_engine
+
+warnings.filterwarnings('ignore')
+
+
+def rep_penalty(text, n=3, cap=0.5):
+    toks = re.findall(r"\w+|[^\w\s]", text.lower())
+    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
+    return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
+
+
+def calculate_rewards(prompts, responses, reward_model):
+    rewards = torch.zeros(len(responses), device=args.device)
+
+    with torch.no_grad():
+        reward_model_scores = []
+        batch_size = len(prompts)
+
+        for i in range(batch_size):
+            for j in range(args.num_generations):
+                response_idx = i * args.num_generations + j
+                response = responses[response_idx]
+                prompt = prompts[i]
+
+                pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
+                matches = re.findall(pattern, prompt, re.DOTALL)
+                messages = [{"role": role, "content": content.strip()} for role, content in matches]
+                answer = response
+                rewards[response_idx] += 0.5 if 20 <= len(response.strip()) <= 800 else -0.5
+                if '</think>' in response:
+                    thinking_content, answer_content = response.split('</think>', 1)
+                    rewards[response_idx] += 1.0 if 20 <= len(thinking_content.strip()) <= 300 else -0.5
+                    rewards[response_idx] += 0.25 if response.count('</think>') == 1 else -0.25
+                    answer = answer_content.strip()
+                rewards[response_idx] -= rep_penalty(answer)
+
+                score = reward_model.get_score(messages, answer)
+                reward_model_scores.append(score)
+
+        reward_model_scores = torch.tensor(reward_model_scores, device=args.device)
+        rewards += reward_model_scores
+
+    return rewards
+
+
+def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model, start_step=0, wandb=None, use_sglang=False):
+    for step, batch in enumerate(loader, start=start_step + 1):
+        prompts = batch['prompt']  # list[str], length B
+        prompt_inputs = tokenizer(prompts, return_tensors="pt", padding=True, return_token_type_ids=False,
+                                  padding_side="left", add_special_tokens=False).to(args.device)
+        if args.max_seq_len:
+            prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -args.max_seq_len:]
+            prompt_inputs["attention_mask"] = prompt_inputs["attention_mask"][:, -args.max_seq_len:]
+
+        rollout_result = rollout_engine.rollout(
+            prompt_ids=prompt_inputs["input_ids"],
+            attention_mask=prompt_inputs["attention_mask"],
+            num_generations=args.num_generations,
+            max_new_tokens=args.max_gen_len,
+            temperature=0.8,
+        )
+        outputs = rollout_result.output_ids
+        completion_ids = rollout_result.completion_ids
+        completions = rollout_result.completions
+        old_per_token_logps = rollout_result.per_token_logps.to(args.device).detach()
+        prompt_lens = rollout_result.prompt_lens.to(args.device)
+        full_mask = (outputs != tokenizer.pad_token_id).long()
+        logp_pos = prompt_lens.unsqueeze(1) - 1 + torch.arange(completion_ids.size(1), device=args.device).unsqueeze(0)
+
+        rewards = calculate_rewards(prompts, completions, reward_model).to(args.device)  # [B*num_gen]
+
+        model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
+        with autocast_ctx:
+            res = model_unwrapped(outputs, attention_mask=full_mask)
+            aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
+            per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)
+
+        with torch.no_grad():
+            ref_per_token_logps = F.log_softmax(ref_model(outputs, attention_mask=full_mask).logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)
+
+        if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
+            for i in range(len(prompts)):
+                Logger(f"[DEBUG] step={step}, sample[{i}]")
+                Logger('-'*100)
+                Logger(f"{'=' * 30} [DEBUG] sample[{i}] CONTEXT_BEGIN {'=' * 30}")
+                Logger(prompts[i])
+                Logger(f"{'=' * 31} [DEBUG] sample[{i}] CONTEXT_END {'=' * 31}")
+                for j in range(args.num_generations):
+                    idx = i * args.num_generations + j
+                    Logger(f"{'=' * 28} [DEBUG] gen[{j}] RESPONSE_BEGIN {'=' * 28}")
+                    Logger(completions[idx])
+                    Logger(f"{'=' * 29} [DEBUG] gen[{j}] RESPONSE_END {'=' * 29}")
+                    Logger(f"[DEBUG] gen[{j}] reward={rewards[idx].item():.4f}")
+                Logger('='*100)
+
+        grouped_rewards = rewards.view(-1, args.num_generations)  # [B, num_gen]
+        mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations)  # [B*num_gen]
+        std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(args.num_generations)  # [B*num_gen]
+        advantages = (rewards - mean_r) / (std_r + 1e-4)  # [B*num_gen]
+
+        completion_pad_mask = rollout_result.completion_mask.to(args.device).bool()
+        is_eos = (completion_ids == tokenizer.eos_token_id) & completion_pad_mask  # [B*num_gen, R]
+        eos_idx = torch.full((is_eos.size(0),), is_eos.size(1) - 1, dtype=torch.long, device=args.device)
+        eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
+        completion_mask = ((torch.arange(is_eos.size(1), device=args.device).expand(is_eos.size(0), -1) <= eos_idx.unsqueeze(1)) & completion_pad_mask).int()  # [B*num_gen, R]
+
+        kl_div = ref_per_token_logps - per_token_logps
+        per_token_kl = torch.exp(kl_div) - kl_div - 1  # [B*num_gen, R]
+        ratio = torch.exp(per_token_logps - old_per_token_logps)  # [B*num_gen, R]
+        if args.loss_type == "cispo":
+            clamped_ratio = torch.clamp(ratio, max=args.epsilon_high).detach()
+            per_token_loss = -(clamped_ratio * advantages.unsqueeze(1) * per_token_logps - args.beta * per_token_kl)
+        else:
+            clipped_ratio = torch.clamp(ratio, 1 - args.epsilon, 1 + args.epsilon)
+            per_token_loss1 = ratio * advantages.unsqueeze(1)
+            per_token_loss2 = clipped_ratio * advantages.unsqueeze(1)
+            per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl)
+        policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp(min=1)).mean()
+        loss = (policy_loss + aux_loss) / args.accumulation_steps  # scalar
+        loss.backward()
+
+        if step % args.accumulation_steps == 0:
+            if args.grad_clip > 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+
+        if step % args.log_interval == 0 or step == iters:
+            policy_loss_val = loss.item() * args.accumulation_steps
+            current_aux_loss = aux_loss.item()
+            avg_reward_val = rewards.mean().item()
+            avg_len_val = completion_mask.sum(dim=1).float().mean().item()
+            kl_ref_val = ((ref_per_token_logps - per_token_logps) * completion_mask).sum().item() / max(completion_mask.sum().item(), 1)
+            advantages_mean_val = advantages.mean().item()
+            advantages_std_val = advantages.std().item()
+            current_lr = optimizer.param_groups[0]['lr']
+
+            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), '
+                   f'Reward: {avg_reward_val:.4f}, KL_ref: {kl_ref_val:.4f}, '
+                   f'Adv Std: {advantages_std_val:.4f}, Adv Mean: {advantages_mean_val:.4f}, '
+                   f'Actor Loss: {policy_loss_val:.4f}, Avg Response Len: {avg_len_val:.2f}, Learning Rate: {current_lr:.8f}')
+
+            if wandb and is_main_process():
+                wandb.log({
+                    "reward": avg_reward_val,
+                    "kl_ref": kl_ref_val,
+                    "advantages_std": advantages_std_val,
+                    "advantages_mean": advantages_mean_val,
+                    "policy_loss": policy_loss_val,
+                    "avg_response_len": avg_len_val,
+                    "learning_rate": current_lr
+                })
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if lm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            state_dict = raw_model.state_dict()
+            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, 
+                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler)
+            model.train()
+            del state_dict
+
+        if step % args.save_interval == 0 or step == iters: rollout_engine.update_policy(model)
+
+        del prompt_inputs, outputs, completion_ids, per_token_logps, ref_per_token_logps
+        del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask, completion_pad_mask, prompt_lens, logp_pos
+
+    if step > start_step and step % args.accumulation_steps != 0:
+        if args.grad_clip > 0:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        optimizer.step()
+        scheduler.step()
+        optimizer.zero_grad()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind GRPO (Group Relative Policy Optimization)")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='grpo', type=str, help="保存权重的前缀名")
+    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=2, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=3e-7, help="初始学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=10, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--max_seq_len', default=768, type=int, help="Prompt最大长度")
+    parser.add_argument("--max_gen_len", type=int, default=1024, help="生成的最大长度")
+    parser.add_argument("--data_path", type=str, default="../dataset/rlaif.jsonl", help="RLAIF数据路径")
+    parser.add_argument("--num_generations", type=int, default=6, help="每个prompt生成的样本数")
+    parser.add_argument("--beta", type=float, default=0.1, help="KL惩罚系数")
+    parser.add_argument("--loss_type", type=str, default="cispo", choices=["grpo", "cispo"], help="loss类型")
+    parser.add_argument("--epsilon", type=float, default=0.2, help="GRPO的PPO clip epsilon")
+    parser.add_argument("--epsilon_high", type=float, default=5.0, help="epsilon上界")
+    parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练")
+    parser.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward", help="Reward模型路径")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-GRPO", help="wandb项目名")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--debug_mode", action="store_true", help="是否打印训练调试采样")
+    parser.add_argument("--debug_interval", type=int, default=20, help="debug模式下每隔多少step打印一次采样")
+    parser.add_argument("--thinking_ratio", type=float, default=0.9, help="按概率开启thinking（0.0~1.0）")
+    parser.add_argument("--rollout_engine", type=str, default="torch", choices=["torch", "sglang"], help="rollout引擎类型")
+    parser.add_argument("--sglang_base_url", type=str, default="http://localhost:8998", help="SGLang服务器URL")
+    parser.add_argument("--sglang_model_path", type=str, default="../model", help="SGLang tokenizer路径")
+    parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_grpo", help="SGLang共享存储路径")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+    
+    # ========== 2. 配置目录、模型参数、检查ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+                               max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe))
+    ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
+    
+    # ========== 3. 设置混合精度 ==========
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+    
+    # ========== 4. 配wandb ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb_run_name = f"MiniMind-GRPO-Epoch-{args.epochs}-BS-{args.batch_size}-LR-{args.learning_rate}"
+        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
+    
+    # ========== 5. 初始化模型和数据 ==========
+    base_weight = args.from_weight
+    # Policy模型
+    model, tokenizer = init_model(lm_config, base_weight, device=args.device)
+    # Reference模型
+    ref_model, _ = init_model(lm_config, base_weight, device=args.device)
+    ref_model = ref_model.eval().requires_grad_(False)
+    # Reward模型
+    reward_model = LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16)
+    # Rollout引擎（可插拔替换，只负责 policy 推理）
+    rollout_engine = create_rollout_engine(
+        engine_type=args.rollout_engine,
+        policy_model=model,
+        tokenizer=tokenizer,
+        device=args.device,
+        autocast_ctx=autocast_ctx,
+        sglang_base_url=args.sglang_base_url,
+        sglang_model_path=args.sglang_model_path,
+        sglang_shared_path=args.sglang_shared_path,
+    )
+    # 数据和优化器
+    train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len, thinking_ratio=args.thinking_ratio)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler)
+    iters = len(loader_for_count)
+    total_optimizer_steps = math.ceil(iters / args.accumulation_steps) * args.epochs
+    scheduler = CosineAnnealingLR(optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
+    
+    # ========== 6. 从ckp恢复状态 ==========
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'])
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scheduler.load_state_dict(ckp_data['scheduler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+    
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
+        rollout_engine.update_policy(model)
+    if dist.is_initialized():
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+    rollout_engine.update_policy(model)
+    
+    # ========== 8. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
+        if skip > 0: 
+            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
+            grpo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
+        else:
+            grpo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
+    
+    # ========== 9. 清理分布进程 ==========
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

src/omni/trainers/lora.py

@@ -0,0 +1,183 @@
+import os
+
+
+import datasets  # noqa: F401  # Windows pyarrow/torch DLL conflict workaround (issue #771)
+import argparse
+import time
+import warnings
+import torch
+import torch.distributed as dist
+from contextlib import nullcontext
+from torch import optim, nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from omni.models.minimind import MiniMindConfig
+from omni.datasets.lm_dataset import SFTDataset
+from omni.models.lora import save_lora, apply_lora
+from omni.utils.training import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
+
+warnings.filterwarnings('ignore')
+
+
+def train_epoch(epoch, loader, iters, lora_params, start_step=0, wandb=None):
+    start_time = time.time()
+    last_step = start_step
+    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
+        input_ids = input_ids.to(args.device)
+        labels = labels.to(args.device)
+        last_step = step
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        with autocast_ctx:
+            res = model(input_ids, labels=labels)
+            loss = res.loss + res.aux_loss
+            loss = loss / args.accumulation_steps
+
+        scaler.scale(loss).backward()
+
+        if step % args.accumulation_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip)
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer.zero_grad(set_to_none=True)
+
+        if step % args.log_interval == 0 or step == iters:
+            spend_time = time.time() - start_time
+            current_loss = loss.item() * args.accumulation_steps
+            current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
+            current_logits_loss = current_loss - current_aux_loss
+            current_lr = optimizer.param_groups[-1]['lr']
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
+            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
+            if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if lm_config.use_moe else ''
+            lora_save_path = f'{args.save_dir}/{args.lora_name}_{lm_config.hidden_size}{moe_suffix}.pth'
+            # LoRA只保存LoRA权重
+            save_lora(model, lora_save_path)
+            lm_checkpoint(lm_config, weight=args.lora_name, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
+            model.train()
+
+        del input_ids, labels, res, loss
+
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(lora_params, args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind LoRA Fine-tuning")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument("--lora_name", type=str, default="lora_medical", help="LoRA权重名称(如lora_identity/lora_medical等)")
+    parser.add_argument("--epochs", type=int, default=10, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=32, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=10, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument("--data_path", type=str, default="../dataset/lora_medical.jsonl", help="LoRA训练数据路径")
+    parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练，默认full_sft")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-LoRA", help="wandb项目名")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+    
+    # ========== 2. 配置目录、模型参数、检查ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
+    ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None
+    
+    # ========== 3. 设置混合精度 ==========
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+    
+    # ========== 4. 配wandb ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb_run_name = f"MiniMind-LoRA-{args.lora_name}-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LR-{args.learning_rate}"
+        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
+    
+    # ========== 5. 定义模型、应用LoRA、冻结非LoRA参数 ==========
+    model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
+    apply_lora(model)
+    
+    # 统计参数
+    total_params = sum(p.numel() for p in model.parameters())
+    lora_params_count = sum(p.numel() for name, p in model.named_parameters() if 'lora' in name)
+    Logger(f"LLM 总参数量: {total_params / 1e6:.3f} M")
+    Logger(f"LoRA 参数量: {lora_params_count / 1e6:.3f} M")
+    Logger(f"LoRA 参数占比: {lora_params_count / total_params * 100:.2f}%")
+    
+    # 冻结非LoRA参数，收集LoRA参数
+    lora_params = []
+    for name, param in model.named_parameters():
+        if 'lora' in name:
+            param.requires_grad = True
+            lora_params.append(param)
+        else:
+            param.requires_grad = False
+    
+    # ========== 6. 定义数据和优化器 ==========
+    train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
+    optimizer = optim.AdamW(lora_params, lr=args.learning_rate)
+    
+    # ========== 7. 从ckp恢复状态 ==========
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'], strict=False)
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scaler.load_state_dict(ckp_data['scaler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+    
+    # ========== 8. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        args.use_compile = 0
+        Logger('[LoRA] monkey-patch forward 与 torch.compile 不兼容，use_compile 已自动关闭')
+    if dist.is_initialized():
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+    
+    # ========== 9. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
+        if skip > 0: 
+            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
+            train_epoch(epoch, loader, len(loader) + skip, lora_params, start_step, wandb)
+        else:
+            train_epoch(epoch, loader, len(loader), lora_params, 0, wandb)
+    
+    # ========== 10. 清理分布进程 ==========
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

src/omni/trainers/ppo.py

@@ -0,0 +1,434 @@
+import os
+
+
+import datasets  # noqa: F401  # Windows pyarrow/torch DLL conflict workaround (issue #771)
+import argparse
+import math
+import re
+import warnings
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+from contextlib import nullcontext
+from torch import optim, nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.nn.utils import clip_grad_norm_
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.datasets.lm_dataset import RLAIFDataset
+from omni.utils.training import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
+from omni.trainers.rollout_engine import create_rollout_engine
+
+warnings.filterwarnings('ignore')
+
+
+def rep_penalty(text, n=3, cap=0.5):
+    toks = re.findall(r"\w+|[^\w\s]", text.lower())
+    grams = [tuple(toks[i:i + n]) for i in range(len(toks) - n + 1)]
+    return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0
+
+
+# 自定义的Critic模型，继承自MiniMindLM
+class CriticModel(MiniMindForCausalLM):
+    def __init__(self, params):
+        super().__init__(params)
+        # 替换lm_head为输出单一价值的线性层
+        self.value_head = nn.Linear(params.hidden_size, 1)
+
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        # 使用基础模型获取隐藏状态
+        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **kwargs)
+        hidden_states = self.model.norm(outputs[0])
+        # 使用value_head获取价值估计
+        values = self.value_head(hidden_states).squeeze(-1)
+        return values
+
+
+def calculate_rewards(prompts, responses, reward_model):
+    rewards = torch.zeros(len(responses), device=args.device)
+
+    with torch.no_grad():
+        reward_model_scores = []
+        for i, (prompt, response) in enumerate(zip(prompts, responses)):
+            pattern = r"<\|im_start\|>(system|user|assistant)\s+(.*?)<\|im_end\|>"
+            matches = re.findall(pattern, prompt, re.DOTALL)
+            messages = [{"role": role, "content": content.strip()} for role, content in matches]
+            answer = response
+            rewards[i] += 0.5 if 20 <= len(response.strip()) <= 800 else -0.5
+            if '</think>' in response:
+                thinking_content, answer_content = response.split('</think>', 1)
+                rewards[i] += 1.0 if 20 <= len(thinking_content.strip()) <= 300 else -0.5
+                rewards[i] += 0.25 if response.count('</think>') == 1 else -0.25
+                answer = answer_content.strip()
+            rewards[i] -= rep_penalty(answer)
+
+            score = reward_model.get_score(messages, answer)
+            reward_model_scores.append(score)
+
+        reward_model_scores = torch.tensor(reward_model_scores, device=args.device)
+        rewards += reward_model_scores
+
+    return rewards
+
+
+def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step=0, wandb=None, use_sglang=False):
+    actor_model.train()
+    critic_model.train()
+    grad_accum_step = 0
+
+    for step, batch in enumerate(loader, start=start_step + 1):
+        prompts = batch["prompt"]  # list[str], length B
+        enc = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True, max_length=args.max_seq_len,
+                        padding_side="left").to(args.device)  # input_ids: [B, P], attention_mask: [B, P]
+
+        rollout_result = rollout_engine.rollout(
+            prompt_ids=enc.input_ids,
+            attention_mask=enc.attention_mask,
+            num_generations=1,
+            max_new_tokens=args.max_gen_len,
+            temperature=0.8,
+        )
+        gen_out = rollout_result.output_ids
+        completion_ids = rollout_result.completion_ids
+        prompt_lens = rollout_result.prompt_lens.to(args.device)
+        responses_text = rollout_result.completions
+        old_resp_logp = rollout_result.per_token_logps.to(args.device)
+        rewards = calculate_rewards(prompts, responses_text, reward_model)  # [B]
+
+        if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
+            for i in range(len(prompts)):
+                Logger(f"[DEBUG] step={step}, sample[{i}]")
+                Logger('-'*100)
+                Logger(f"{'=' * 30} [DEBUG] sample[{i}] CONTEXT_BEGIN {'=' * 30}")
+                Logger(prompts[i])
+                Logger(f"{'=' * 31} [DEBUG] sample[{i}] CONTEXT_END {'=' * 31}")
+                Logger(f"[DEBUG] prompt_len={prompt_lens[i].item()}, response_len={len(responses_text[i])}")
+                Logger(f"{'=' * 28} [DEBUG] sample[{i}] RESPONSE_BEGIN {'=' * 28}")
+                Logger(responses_text[i])
+                Logger(f"{'=' * 29} [DEBUG] sample[{i}] RESPONSE_END {'=' * 29}")
+                Logger(f"[DEBUG] reward={rewards[i].item():.4f}")
+                Logger('='*100)
+
+        full_mask = (gen_out != tokenizer.pad_token_id).long()  # [B, P+R]
+        labels = gen_out[:, 1:].clone()  # [B, P+R-1]
+        B = len(prompts)
+        resp_labels = completion_ids
+        resp_idx = torch.arange(resp_labels.size(1), device=gen_out.device).unsqueeze(0)
+        logp_pos = prompt_lens.unsqueeze(1) - 1 + resp_idx
+        resp_pad_mask = rollout_result.completion_mask.to(args.device).bool()
+        resp_lengths = resp_pad_mask.sum(dim=1); valid_resp = resp_lengths > 0; eos_mask = resp_labels.eq(tokenizer.eos_token_id) & resp_pad_mask
+        has_eos = eos_mask.any(dim=1); eos_pos = torch.argmax(eos_mask.int(), dim=1)
+        resp_lengths = torch.where(has_eos, eos_pos + 1, resp_lengths).long().clamp(min=1)
+        resp_policy_mask = ((resp_idx < resp_lengths.unsqueeze(1)) & resp_pad_mask).float()
+        resp_value_mask = resp_policy_mask.clone()
+
+        with torch.no_grad():  # Rollout阶段只需推理获取old_logp和old_values，切断梯度省显存
+            critic_for_rollout = critic_model.module if isinstance(critic_model, DistributedDataParallel) else critic_model
+            values_seq = critic_for_rollout(input_ids=gen_out, attention_mask=full_mask)
+            old_resp_values = values_seq.gather(1, logp_pos) * resp_value_mask
+            
+            ref_resp_logp = F.log_softmax(ref_model(input_ids=gen_out, attention_mask=full_mask).logits[:, :-1], dim=-1).gather(2, labels.unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)
+            token_rewards = torch.zeros_like(old_resp_logp)
+            last_idx = resp_lengths - 1  # [B]
+            token_rewards[torch.arange(B, device=args.device)[valid_resp], last_idx[valid_resp]] += rewards[valid_resp]  # 末尾加外部奖励
+
+            gen_len = old_resp_values.size(1); lastgaelam = torch.zeros(B, device=args.device); advs_rev = []
+            for t in reversed(range(gen_len)):
+                nv = old_resp_values[:, t + 1] if t < gen_len - 1 else 0.0
+                delta = token_rewards[:, t] + args.gamma * nv - old_resp_values[:, t]
+                lastgaelam = delta + args.gamma * args.lam * lastgaelam
+                advs_rev.append(lastgaelam)
+            advantages = torch.stack(advs_rev[::-1], dim=1)  # [B, R]
+            returns = advantages + old_resp_values  # [B, R]
+
+            adv_mean = (advantages * resp_policy_mask).sum() / resp_policy_mask.sum().clamp(min=1)
+            adv_var = ((advantages - adv_mean) ** 2 * resp_policy_mask).sum() / resp_policy_mask.sum().clamp(min=1)
+            advantages = (advantages - adv_mean) * torch.rsqrt(adv_var + 1e-8) * resp_policy_mask
+
+        mb_size = max(1, min(args.mini_batch_size, B))
+        stop_ppo = False
+        policy_loss_sum = 0.0
+        value_loss_sum = 0.0
+        kl_sum = 0.0
+        kl_ref_sum = 0.0
+        clipfrac_sum = 0.0
+        aux_loss_sum = 0.0
+        log_count = 0
+        actor_unwrapped = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
+        critic_unwrapped = critic_model.module if isinstance(critic_model, DistributedDataParallel) else critic_model
+        for ppo_epoch in range(args.ppo_update_iters):
+            if stop_ppo:
+                break
+            b_inds = torch.randperm(B, device=args.device)
+            for i in range(0, B, mb_size):
+                inds = b_inds[i:i + mb_size]
+                
+                mb_values_seq = critic_unwrapped(input_ids=gen_out[inds], attention_mask=full_mask[inds])
+                mb_resp_values = mb_values_seq.gather(1, logp_pos[inds])
+
+                with autocast_ctx:
+                    res = actor_unwrapped(input_ids=gen_out[inds], attention_mask=full_mask[inds])
+                    aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
+
+                mb_resp_logp = F.log_softmax(res.logits[:, :-1], dim=-1).gather(2, labels[inds].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos[inds])
+                
+                log_ratio = mb_resp_logp - old_resp_logp[inds]
+                approx_kl = (0.5 * (log_ratio ** 2) * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
+                
+                # 同步各卡的 approx_kl，防止某卡 break 而其它卡继续导致 DDP 死锁
+                approx_kl_val = approx_kl.detach().clone()
+                if dist.is_initialized():
+                    dist.all_reduce(approx_kl_val, op=dist.ReduceOp.AVG)
+                    
+                if approx_kl_val > args.early_stop_kl:
+                    stop_ppo = True
+                
+                ratio = torch.exp(log_ratio)
+                clipfrac = ((((ratio - 1.0).abs() > args.clip_epsilon).float() * resp_policy_mask[inds]).sum()
+                            / resp_policy_mask[inds].sum().clamp(min=1))
+                kl_ref_penalty = ((torch.exp(ref_resp_logp[inds] - mb_resp_logp) - (ref_resp_logp[inds] - mb_resp_logp) - 1.0)
+                                  * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
+                policy_loss = ((torch.max(-advantages[inds] * ratio,
+                                          -advantages[inds] * torch.clamp(ratio, 1.0 - args.clip_epsilon, 1.0 + args.clip_epsilon))
+                               * resp_policy_mask[inds]).sum() / resp_policy_mask[inds].sum().clamp(min=1)
+                               + args.kl_coef * kl_ref_penalty)
+                value_loss = 0.5 * (torch.max((mb_resp_values - returns[inds]) ** 2,
+                                              (torch.clamp(mb_resp_values, old_resp_values[inds] - args.cliprange_value,
+                                                           old_resp_values[inds] + args.cliprange_value) - returns[inds]) ** 2)
+                                    * resp_value_mask[inds]).sum() / resp_value_mask[inds].sum().clamp(min=1)
+
+                kl = approx_kl_val
+                kl_ref = kl_ref_penalty.detach()
+
+                # 早停时必须保证 forward-backward 闭环，故只截断 loss 不中断 DDP 通信
+                if stop_ppo:
+                    loss = (policy_loss + args.vf_coef * value_loss + aux_loss) * 0.0
+                else:
+                    loss = (policy_loss + args.vf_coef * value_loss + aux_loss) / args.accumulation_steps
+                
+                loss.backward()
+
+                policy_loss_sum += policy_loss.item()
+                value_loss_sum += value_loss.item()
+                kl_sum += kl.item()
+                kl_ref_sum += kl_ref.item()
+                clipfrac_sum += clipfrac.item()
+                aux_loss_sum += aux_loss.item()
+                log_count += 1
+
+                grad_accum_step += 1
+
+                if grad_accum_step % args.accumulation_steps == 0:
+                    clip_grad_norm_(actor_model.parameters(), args.grad_clip)
+                    clip_grad_norm_(critic_model.parameters(), args.grad_clip)
+                    actor_optimizer.step()
+                    critic_optimizer.step()
+                    actor_scheduler.step()
+                    critic_scheduler.step()
+                    actor_optimizer.zero_grad()
+                    critic_optimizer.zero_grad()
+
+        if grad_accum_step % args.accumulation_steps != 0:
+            clip_grad_norm_(actor_model.parameters(), args.grad_clip)
+            clip_grad_norm_(critic_model.parameters(), args.grad_clip)
+            actor_optimizer.step()
+            critic_optimizer.step()
+            actor_scheduler.step()
+            critic_scheduler.step()
+            actor_optimizer.zero_grad()
+            critic_optimizer.zero_grad()
+        
+        if step % args.save_interval == 0 or step == iters: rollout_engine.update_policy(actor_model)
+
+        if is_main_process():
+            critic_loss_val = value_loss_sum / max(log_count, 1)
+            reward_val = rewards.mean().item()
+            approx_kl_val = kl_sum / max(log_count, 1)
+            kl_ref_val = kl_ref_sum / max(log_count, 1)
+            clipfrac_val = clipfrac_sum / max(log_count, 1)
+            avg_len_val = resp_lengths.float().mean().item()
+            actor_lr, critic_lr = actor_optimizer.param_groups[0]['lr'], critic_optimizer.param_groups[0]['lr']
+
+            if wandb is not None:
+                wandb.log({
+                    "reward": reward_val,
+                    "kl_ref": kl_ref_val,
+                    "approx_kl": approx_kl_val,
+                    "clipfrac": clipfrac_val,
+                    "critic_loss": critic_loss_val,
+                    "avg_response_len": avg_len_val,
+                    "actor_lr": actor_lr,
+                    "critic_lr": critic_lr,
+                })
+
+            Logger(f"Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), "
+                   f"Reward: {reward_val:.4f}, KL_ref: {kl_ref_val:.4f}, Approx KL: {approx_kl_val:.4f}, "
+                   f"ClipFrac: {clipfrac_val:.4f}, Critic Loss: {critic_loss_val:.4f}, "
+                   f"Avg Response Len: {avg_len_val:.2f}, Actor LR: {actor_lr:.8f}, Critic LR: {critic_lr:.8f}")
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            actor_model.eval()
+            moe_suffix = '_moe' if lm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            raw_actor = actor_model.module if isinstance(actor_model, DistributedDataParallel) else actor_model
+            raw_actor = getattr(raw_actor, '_orig_mod', raw_actor)
+            actor_state = raw_actor.state_dict()
+            torch.save({k: v.half().cpu() for k, v in actor_state.items()}, ckp)
+            
+            # 使用 lm_checkpoint 保存完整状态（包括 critic）
+            lm_checkpoint(lm_config, weight=args.save_weight, model=actor_model, optimizer=actor_optimizer, 
+                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints',
+                         scheduler=actor_scheduler, critic_model=critic_model, 
+                         critic_optimizer=critic_optimizer, critic_scheduler=critic_scheduler)
+            actor_model.train()
+            del actor_state
+
+        del enc, gen_out, completion_ids, responses_text, rewards, full_mask, values_seq, advantages
+        del labels, resp_labels, resp_idx, resp_pad_mask, valid_resp, eos_mask, has_eos, eos_pos, resp_lengths, resp_policy_mask, resp_value_mask, old_resp_logp, ref_resp_logp
+        del kl, kl_ref, policy_loss, value_loss, loss, token_rewards, returns, old_resp_values, prompt_lens, logp_pos
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind PPO (Proximal Policy Optimization)")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='ppo_actor', type=str, help="保存权重的前缀名")
+    parser.add_argument("--epochs", type=int, default=1, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=2, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=3e-7, help="Actor学习率")
+    parser.add_argument("--critic_learning_rate", type=float, default=5e-7, help="Critic学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=1, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=10, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument('--max_seq_len', default=768, type=int, help="Prompt最大长度")
+    parser.add_argument("--max_gen_len", type=int, default=1024, help="生成的最大长度")
+    parser.add_argument("--data_path", type=str, default="../dataset/rlaif.jsonl", help="RLAIF数据路径")
+    parser.add_argument("--clip_epsilon", type=float, default=0.2, help="PPO裁剪参数")
+    parser.add_argument("--vf_coef", type=float, default=0.5, help="Value function系数")
+    parser.add_argument("--kl_coef", type=float, default=0.02, help="KL散度惩罚系数")
+    parser.add_argument("--gamma", type=float, default=1.0, help="GAE折扣因子")
+    parser.add_argument("--lam", type=float, default=0.95, help="GAE lambda参数")
+    parser.add_argument("--cliprange_value", type=float, default=0.2, help="Value function裁剪范围")
+    parser.add_argument("--ppo_update_iters", type=int, default=2, help="同一批rollout重复更新次数")
+    parser.add_argument("--early_stop_kl", type=float, default=0.25, help="PPO early stop 的 KL 阈值")
+    parser.add_argument("--mini_batch_size", type=int, default=2, help="PPO每次更新的minibatch大小")
+    parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练")
+    parser.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward", help="Reward模型路径")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-PPO", help="wandb项目名")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--debug_mode", action="store_true", help="是否打印训练调试采样")
+    parser.add_argument("--debug_interval", type=int, default=20, help="debug模式下每隔多少step打印一次采样")
+    parser.add_argument("--thinking_ratio", type=float, default=0.9, help="按概率开启thinking（0.0~1.0）")
+    parser.add_argument("--rollout_engine", type=str, default="torch", choices=["torch", "sglang"], help="rollout引擎类型")
+    parser.add_argument("--sglang_base_url", type=str, default="http://localhost:8998", help="SGLang服务器URL")
+    parser.add_argument("--sglang_model_path", type=str, default="../model", help="SGLang tokenizer路径")
+    parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_ppo", help="SGLang共享存储路径")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+    
+    # ========== 2. 配置目录、模型参数、检查ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
+    ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
+    
+    # ========== 3. 设置混合精度 ==========
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+    
+    # ========== 4. 配wandb ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb_run_name = f"MiniMind-PPO-Epoch-{args.epochs}-BS-{args.batch_size}-LR-{args.learning_rate}"
+        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
+    
+    # ========== 5. 初始化模型和数据 ==========
+    base_weight = args.from_weight
+    # Actor模型
+    actor_model, tokenizer = init_model(lm_config, base_weight, device=args.device)
+    ref_model, _ = init_model(lm_config, base_weight, device=args.device)
+    ref_model = ref_model.eval().requires_grad_(False)
+    moe_suffix = '_moe' if lm_config.use_moe else ''
+    ckp = f'{args.save_dir}/{base_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+    state_dict = torch.load(ckp, map_location=args.device)
+    critic_model = CriticModel(lm_config)
+    critic_model.load_state_dict(state_dict, strict=False)
+    critic_model = critic_model.to(args.device)
+    reward_model = LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16)
+    # Rollout引擎
+    rollout_engine = create_rollout_engine(
+        engine_type=args.rollout_engine,
+        policy_model=actor_model,
+        tokenizer=tokenizer,
+        device=args.device,
+        autocast_ctx=autocast_ctx,
+        sglang_base_url=args.sglang_base_url,
+        sglang_model_path=args.sglang_model_path,
+        sglang_shared_path=args.sglang_shared_path,
+    )
+    train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=(args.max_seq_len + args.max_gen_len), thinking_ratio=args.thinking_ratio)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    actor_optimizer = optim.AdamW(actor_model.parameters(), lr=args.learning_rate)
+    critic_optimizer = optim.AdamW(critic_model.parameters(), lr=args.critic_learning_rate)
+    loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler)
+    iters = len(loader_for_count)
+    mb_factor = max(1, math.ceil(args.batch_size / args.mini_batch_size))
+    total_optimizer_steps = math.ceil(iters * args.epochs * args.ppo_update_iters * mb_factor / args.accumulation_steps)
+    actor_scheduler = CosineAnnealingLR(actor_optimizer, T_max=total_optimizer_steps, eta_min=args.learning_rate / 10)
+    critic_scheduler = CosineAnnealingLR(critic_optimizer, T_max=total_optimizer_steps, eta_min=args.critic_learning_rate / 10)
+
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        actor_model.load_state_dict(ckp_data['model'])
+        critic_model.load_state_dict(ckp_data['critic_model'])
+        actor_optimizer.load_state_dict(ckp_data['optimizer'])
+        critic_optimizer.load_state_dict(ckp_data['critic_optimizer'])
+        actor_scheduler.load_state_dict(ckp_data['scheduler'])
+        critic_scheduler.load_state_dict(ckp_data['critic_scheduler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+    
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        actor_model = torch.compile(actor_model)
+        Logger('torch.compile enabled')
+        rollout_engine.update_policy(actor_model)
+    if dist.is_initialized():
+        actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank])
+        critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank])
+    rollout_engine.update_policy(actor_model)
+    
+    # ========== 8. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
+        if skip > 0: 
+            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
+            ppo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang"))
+        else:
+            ppo_train_epoch(epoch, loader, len(loader), rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, 0, wandb, use_sglang = (args.rollout_engine == "sglang"))
+    
+    # ========== 9. 清理分布进程 ==========
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

src/omni/trainers/pretrain.py

@@ -0,0 +1,169 @@
+import os
+
+
+import datasets  # noqa: F401  # Windows pyarrow/torch DLL conflict workaround (issue #771)
+import argparse
+import time
+import warnings
+import torch
+import torch.distributed as dist
+from contextlib import nullcontext
+from torch import optim, nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from omni.models.minimind import MiniMindConfig
+from omni.datasets.lm_dataset import PretrainDataset
+from omni.utils.training import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
+
+warnings.filterwarnings('ignore')
+
+
+def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
+    start_time = time.time()
+    last_step = start_step
+    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
+        input_ids = input_ids.to(args.device)
+        labels = labels.to(args.device)
+        last_step = step
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        with autocast_ctx:
+            res = model(input_ids, labels=labels)
+            loss = res.loss + res.aux_loss
+            loss = loss / args.accumulation_steps
+
+        scaler.scale(loss).backward()
+
+        if step % args.accumulation_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+
+            scaler.step(optimizer)
+            scaler.update()
+
+            optimizer.zero_grad(set_to_none=True)
+
+        if step % args.log_interval == 0 or step == iters:
+            spend_time = time.time() - start_time
+            current_loss = loss.item() * args.accumulation_steps
+            current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
+            current_logits_loss = current_loss - current_aux_loss
+            current_lr = optimizer.param_groups[-1]['lr']
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
+            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
+            if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if lm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            state_dict = raw_model.state_dict()
+            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
+            model.train()
+            del state_dict
+
+        del input_ids, labels, res, loss
+
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind Pretraining")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='pretrain', type=str, help="保存权重的前缀名")
+    parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=32, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
+    parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐���层数量")
+    parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument("--data_path", type=str, default="../dataset/pretrain_t2t_mini.jsonl", help="预训练数据路径")
+    parser.add_argument('--from_weight', default='none', type=str, help="基于哪个权重训练，为none则从头开始")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain", help="wandb项目名")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+    
+    # ========== 2. 配置目录、模型参数、检查ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
+    ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
+    
+    # ========== 3. 设置混合精度 ==========
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+    
+    # ========== 4. 配wandb ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
+    
+    # ========== 5. 定义模型、数据、优化器 ==========
+    model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
+    train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    
+    # ========== 6. 从ckp恢复状态 ==========
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'])
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scaler.load_state_dict(ckp_data['scaler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+    
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
+    if dist.is_initialized():
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+    
+    # ========== 8. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
+        if skip > 0: 
+            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
+            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
+        else:
+            train_epoch(epoch, loader, len(loader), 0, wandb)
+    
+    # ========== 9. 清理分布进程 ==========
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()

src/omni/trainers/pretrain_vlm.py

@@ -0,0 +1,175 @@
+import argparse
+import os
+import time
+import warnings
+import datasets
+import torch
+import torch.distributed as dist
+from contextlib import nullcontext
+from torch import optim, nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from transformers import AutoTokenizer
+from omni.models import MiniMindVLM, VLMConfig
+from omni.datasets.lm_dataset import VLMDataset
+from omni.utils import get_lr, Logger, is_main_process, init_distributed_mode, setup_seed, init_vlm_model, vlm_checkpoint, SkipBatchSampler, vlm_collate_fn
+
+warnings.filterwarnings('ignore')
+
+
+def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
+    start_time = time.time()
+    last_step = start_step
+    for step, (input_ids, labels, pixel_values) in enumerate(loader, start=start_step + 1):
+        input_ids = input_ids.to(args.device)
+        labels = labels.to(args.device)
+        pixel_values = {k: v.to(args.device) for k, v in pixel_values.items()} if isinstance(pixel_values, dict) else pixel_values.to(args.device)
+        last_step = step
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+
+        with autocast_ctx:
+            res = model(input_ids, labels=labels, pixel_values=pixel_values)
+            loss = res.loss + res.aux_loss
+            loss = loss / args.accumulation_steps
+
+        scaler.scale(loss).backward()
+
+        if step % args.accumulation_steps == 0:
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+
+            scaler.step(optimizer)
+            scaler.update()
+
+            optimizer.zero_grad(set_to_none=True)
+
+        if step % args.log_interval == 0 or step == iters:
+            spend_time = time.time() - start_time
+            current_loss = loss.item() * args.accumulation_steps
+            current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
+            current_logits_loss = current_loss - current_aux_loss
+            current_lr = optimizer.param_groups[-1]['lr']
+            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
+            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
+            if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
+
+        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+            model.eval()
+            moe_suffix = '_moe' if vlm_config.use_moe else ''
+            ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth'
+            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+            state_dict = raw_model.state_dict()
+            clean_state_dict = {
+                key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.')
+            }
+            clean_state_dict = {k: v.half().cpu() for k, v in clean_state_dict.items()}  # 半精度保存并移到CPU
+            torch.save(clean_state_dict, ckp)
+            vlm_checkpoint(vlm_config, weight=args.save_weight, model=model, optimizer=optimizer,
+                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
+            model.train()
+            del state_dict, clean_state_dict
+
+        del input_ids, labels, pixel_values, res, loss
+
+    if last_step > start_step and last_step % args.accumulation_steps != 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="MiniMind-V Pretrain")
+    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
+    parser.add_argument('--save_weight', default='pretrain_vlm', type=str, help="保存权重的前缀名")
+    parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
+    parser.add_argument("--batch_size", type=int, default=16, help="batch size")
+    parser.add_argument("--learning_rate", type=float, default=4e-4, help="初始学习率")
+    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
+    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
+    parser.add_argument("--num_workers", type=int, default=2, help="数据加载线程数")
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度��积步数")
+    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
+    parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
+    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
+    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
+    parser.add_argument('--max_seq_len', default=450, type=int, help="训练的最大截断长度")
+    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
+    parser.add_argument("--data_path", type=str, default="../dataset/pretrain_i2t.parquet", help="训练数据路径")
+    parser.add_argument('--from_weight', default='llm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
+    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
+    parser.add_argument('--freeze_llm', default=2, type=int, choices=[0, 1, 2], help="冻结策略（0=完全可训练，1=冻结+解冻首尾层，2=完全冻结仅训练proj）")
+    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
+    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
+    parser.add_argument("--wandb_project", type=str, default="MiniMind-V-Pretrain", help="wandb项目名")
+    args = parser.parse_args()
+
+    # ========== 1. 初始化环境和随机种子 ==========
+    local_rank = init_distributed_mode()
+    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
+    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
+
+    # ========== 2. 配置目录、模型参数、检查ckp ==========
+    os.makedirs(args.save_dir, exist_ok=True)
+    vlm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, max_seq_len=args.max_seq_len, use_moe=bool(args.use_moe))
+    ckp_data = vlm_checkpoint(vlm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None
+
+    # ========== 3. 设置混合精度 ==========
+    device_type = "cuda" if "cuda" in args.device else "cpu"
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
+
+    # ========== 4. 配wandb ==========
+    wandb = None
+    if args.use_wandb and is_main_process():
+        import swanlab as wandb
+        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
+        resume = 'must' if wandb_id else None
+        wandb_run_name = f"MiniMind-V-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
+        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
+
+    # ========== 5. 定义模型、数据、优化器 ==========
+    model, tokenizer = init_vlm_model(vlm_config, from_weight=args.from_weight, device=args.device, freeze_llm=args.freeze_llm)
+    preprocess = model.vision_encoder.processor
+    train_ds = VLMDataset(args.data_path, tokenizer, preprocess=preprocess, image_special_token=vlm_config.image_special_token, image_token_len=vlm_config.image_token_len, max_length=vlm_config.max_seq_len)
+    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
+    optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=args.learning_rate)
+
+    # ========== 6. 从ckp恢复状态 ==========
+    start_epoch, start_step = 0, 0
+    if ckp_data:
+        model.load_state_dict(ckp_data['model'], strict=False)
+        optimizer.load_state_dict(ckp_data['optimizer'])
+        scaler.load_state_dict(ckp_data['scaler'])
+        start_epoch = ckp_data['epoch']
+        start_step = ckp_data.get('step', 0)
+
+    # ========== 7. 编译和分布式包装 ==========
+    if args.use_compile == 1:
+        model = torch.compile(model)
+        Logger('torch.compile enabled')
+    if dist.is_initialized():
+        model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
+        model = DistributedDataParallel(model, device_ids=[local_rank])
+
+    # ========== 8. 开始训练 ==========
+    for epoch in range(start_epoch, args.epochs):
+        train_sampler and train_sampler.set_epoch(epoch)
+        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
+        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
+        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
+        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=vlm_collate_fn)
+        if skip > 0:
+            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
+            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
+        else:
+            train_epoch(epoch, loader, len(loader), 0, wandb)
+
+    # ========== 9. 清理分布进程 ==========
+    if dist.is_initialized(): dist.destroy_process_group()

src/omni/trainers/rollout_engine.py

@@ -0,0 +1,218 @@
+"""Rollout engines for RL training.
+
+To use the SGLang backend, first launch the server (with a transformers-format model):
+    python -m sglang.launch_server --model-path ./minimind-3 --attention-backend triton --host 0.0.0.0 --port 8998
+"""
+import os
+import requests
+import torch
+import torch.distributed as dist
+from abc import ABC, abstractmethod
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+from torch import Tensor
+from torch.nn.parallel import DistributedDataParallel
+from transformers import AutoTokenizer
+
+
+def compute_per_token_logps(model, input_ids: Tensor, n_keep: int, attention_mask: Optional[Tensor] = None) -> Tensor:
+    if n_keep <= 0:
+        return input_ids.new_empty((input_ids.size(0), 0), dtype=torch.float32)
+    unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
+    input_ids = input_ids.detach().clone() if input_ids.is_inference() else input_ids
+    logits = unwrapped(input_ids, attention_mask=attention_mask, logits_to_keep=n_keep + 1).logits[:, :-1, :]
+    per_token_logps = []
+    for logits_row, ids_row in zip(logits, input_ids[:, -n_keep:]):
+        ids_row = ids_row.detach().clone() if ids_row.is_inference() else ids_row
+        per_token_logps.append(
+            torch.gather(logits_row.log_softmax(dim=-1), 1, ids_row.unsqueeze(1)).squeeze(1)
+        )
+    return torch.stack(per_token_logps)
+
+
+@dataclass
+class RolloutResult:
+    output_ids: Tensor
+    completion_ids: Tensor
+    per_token_logps: Tensor
+    completions: List[str]
+    prompt_lens: Tensor
+    completion_mask: Tensor
+
+
+class RolloutEngine(ABC):
+    tokenizer = None
+
+    @abstractmethod
+    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
+        pass
+
+    @abstractmethod
+    def update_policy(self, model: torch.nn.Module):
+        pass
+
+
+class TorchRolloutEngine(RolloutEngine):
+    def __init__(self, policy_model: torch.nn.Module, tokenizer, device: str = "cuda", autocast_ctx=None):
+        self.policy_model = policy_model
+        self.tokenizer = tokenizer
+        self.device = device
+        self.autocast_ctx = autocast_ctx
+
+    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
+        model = self.policy_model.module if isinstance(self.policy_model, DistributedDataParallel) else self.policy_model
+        ctx = self.autocast_ctx if self.autocast_ctx else nullcontext()
+        with torch.no_grad(), ctx:
+            output_ids = model.generate(
+                input_ids=prompt_ids.repeat_interleave(num_generations, dim=0),
+                attention_mask=attention_mask.repeat_interleave(num_generations, dim=0),
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                temperature=temperature,
+                num_return_sequences=1,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+            ).clone()
+            prompt_len = prompt_ids.size(1)
+            completion_ids = output_ids[:, prompt_len:]
+            full_mask = (output_ids != self.tokenizer.pad_token_id).long()
+            per_token_logps = compute_per_token_logps(self.policy_model, output_ids, completion_ids.size(1), attention_mask=full_mask)
+        completions = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True)
+        return RolloutResult(output_ids, completion_ids, per_token_logps, completions,
+                            prompt_ids.new_full((output_ids.size(0),), prompt_len),
+                            attention_mask.new_ones(output_ids.size(0), completion_ids.size(1)))
+
+    def update_policy(self, model: torch.nn.Module):
+        self.policy_model = model
+
+
+class SGLangRolloutEngine(RolloutEngine):
+    def __init__(self, base_url: str, model_path: str, shared_ckpt_path: str = "./sglang_ckpt", timeout: int = 120):
+        self.base_url = base_url.rstrip('/')
+        self.shared_ckpt_path = shared_ckpt_path
+        self.timeout = timeout
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.http = requests
+
+    def rollout(self, prompt_ids: Tensor, attention_mask: Tensor, num_generations: int, max_new_tokens: int, temperature: float = 0.8) -> RolloutResult:
+        input_ids_list = []
+        for ids, mask in zip(prompt_ids, attention_mask):
+            valid_ids = ids[mask.bool()].tolist()
+            input_ids_list.append(valid_ids)
+        all_input_ids = [ids for ids in input_ids_list for _ in range(num_generations)]
+
+        payload = {
+            "input_ids": all_input_ids,
+            "sampling_params": {
+                "temperature": temperature,
+                "max_new_tokens": max_new_tokens,
+                "stop_token_ids": [self.tokenizer.eos_token_id] if self.tokenizer.eos_token_id else [],
+            },
+            "return_logprob": True,
+        }
+
+        resp = self.http.post(f"{self.base_url}/generate", json=payload, timeout=self.timeout)
+        resp.raise_for_status()
+
+        results = resp.json()
+        if not isinstance(results, list):
+            results = [results]
+
+        all_output_ids, all_completion_ids, all_logprobs = [], [], []
+        completions = []
+
+        for i, result in enumerate(results):
+            meta = result.get("meta_info", {})
+            completion_ids = meta.get("output_ids", result.get("output_ids", []))
+            raw_logprobs = meta.get("output_token_logprobs", [])
+
+            logprobs = []
+            for item in raw_logprobs:
+                if isinstance(item, (list, tuple)) and len(item) >= 1:
+                    logprobs.append(item[0])
+                elif isinstance(item, (int, float)):
+                    logprobs.append(item)
+
+            if len(logprobs) < len(completion_ids):
+                logprobs = [0.0] * (len(completion_ids) - len(logprobs)) + logprobs
+            elif len(logprobs) > len(completion_ids):
+                logprobs = logprobs[-len(completion_ids):] if completion_ids else []
+            prompt = all_input_ids[i]
+            full_output = prompt + completion_ids
+            all_output_ids.append(full_output)
+            all_completion_ids.append(completion_ids)
+            all_logprobs.append(logprobs)
+            completions.append(self.tokenizer.decode(completion_ids, skip_special_tokens=True))
+
+        device = prompt_ids.device
+        max_comp_len = max(1, max(len(ids) for ids in all_completion_ids))
+        max_out_len = max(len(ids) for ids in all_input_ids) + max_comp_len
+
+        def pad_to_tensor(seqs, max_len, pad_val=0):
+            return torch.tensor([s + [pad_val] * (max_len - len(s)) for s in seqs], device=device)
+
+        pad_id = self.tokenizer.pad_token_id
+        return RolloutResult(
+            output_ids=pad_to_tensor(all_output_ids, max_out_len, pad_val=pad_id),
+            completion_ids=pad_to_tensor(all_completion_ids, max_comp_len, pad_val=pad_id),
+            per_token_logps=pad_to_tensor(all_logprobs, max_comp_len, pad_val=0.0),
+            completions=completions,
+            prompt_lens=torch.tensor([len(ids) for ids in all_input_ids], device=device),
+            completion_mask=torch.tensor([[1] * len(ids) + [0] * (max_comp_len - len(ids)) for ids in all_completion_ids], device=device),
+        )
+
+    def update_policy(self, model: torch.nn.Module):
+        ok = True
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            try:
+                unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
+                unwrapped = getattr(unwrapped, '_orig_mod', unwrapped)
+                abs_path = os.path.abspath(self.shared_ckpt_path)
+                state_dict = {k: v.detach().half().cpu() for k, v in unwrapped.state_dict().items()}
+                unwrapped.save_pretrained(abs_path, state_dict=state_dict, safe_serialization=False)
+                self.tokenizer.save_pretrained(abs_path)
+                resp = self.http.post(f"{self.base_url}/update_weights_from_disk", json={"model_path": abs_path}, timeout=self.timeout)
+                if resp.status_code != 200:
+                    print(f"[SGLANG WARNING] update_weights 失败: {resp.status_code}, {resp.text}")
+                ok = resp.status_code == 200
+            except Exception as e:
+                print(f"[SGLANG WARNING] update_weights 异常: {e}")
+                ok = False
+        if dist.is_initialized():
+            ok_t = torch.tensor(int(ok), device=next(model.parameters()).device)
+            dist.broadcast(ok_t, src=0)
+            dist.barrier()
+            ok = bool(ok_t.item())
+        if not ok:
+            raise RuntimeError("SGLang update_policy failed")
+        return ok
+
+    def flush_cache(self) -> bool:
+        resp = self.http.post(f"{self.base_url}/flush_cache", timeout=30)
+        return resp.status_code == 200
+
+    def health(self) -> bool:
+        try:
+            resp = self.http.get(f"{self.base_url}/health", timeout=5)
+            return resp.status_code == 200
+        except Exception:
+            return False
+
+
+def create_rollout_engine(
+    engine_type: str = "torch",
+    policy_model: torch.nn.Module = None,
+    tokenizer=None,
+    device: str = "cuda",
+    autocast_ctx=None,
+    sglang_base_url: str = None,
+    sglang_model_path: str = None,
+    sglang_shared_path: str = None,
+) -> RolloutEngine:
+    if engine_type == "torch":
+        return TorchRolloutEngine(policy_model, tokenizer, device, autocast_ctx)
+    elif engine_type == "sglang":
+        return SGLangRolloutEngine(sglang_base_url, sglang_model_path, sglang_shared_path)
+    else:
+        raise ValueError(f"不支持的引擎类型: {engine_type}")

src/omni/trainers/train_tokenizer.py

@@ -0,0 +1,168 @@
+# 注：不建议再重复训练tokenizer（“词典”），MiniMind已自带，此脚本仅供学习和参考。基于不同词典训练的模型将导致输出完全不统一，降低社区的模型复用性
+# Note: It is not recommended to re-train the tokenizer. MiniMind already includes one. This script is for learning and reference only. Training models with different tokenizers will lead to inconsistent outputs and reduce model reusability in the community.
+import os
+import json
+from tokenizers import decoders, models, pre_tokenizers, trainers, Tokenizer
+
+DATA_PATH = '../dataset/sft_t2t_mini.jsonl'
+TOKENIZER_DIR = '../model_learn_tokenizer/'
+VOCAB_SIZE = 6400
+SPECIAL_TOKENS_NUM = 36
+
+def get_texts(data_path):
+    with open(data_path, 'r', encoding='utf-8', errors='ignore') as f:
+        for i, line in enumerate(f):
+            if i >= 10000: break # 选10000行测试
+            try:
+                data = json.loads(line)
+                contents = [item.get('content') for item in data.get('conversations', []) if item.get('content')]
+                if contents:
+                    yield "\n".join(contents)
+            except json.JSONDecodeError:
+                continue
+
+def train_tokenizer(data_path, tokenizer_dir, vocab_size, special_tokens_num=SPECIAL_TOKENS_NUM):
+    tokenizer = Tokenizer(models.BPE())
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    
+    special_tokens_list = [
+        "<|endoftext|>", "<|im_start|>", "<|im_end|>", 
+        "<|object_ref_start|>", "<|object_ref_end|>", "<|box_start|>", "<|box_end|>", "<|quad_start|>", "<|quad_end|>", 
+        "<|vision_start|>", "<|vision_end|>", "<|vision_pad|>", "<|image_pad|>", "<|video_pad|>", 
+        "<|audio_start|>", "<|audio_end|>", "<|audio_pad|>", "<tts_pad>", "<tts_text_bos>", "<tts_text_eod>", "<tts_text_bos_single>"
+    ]
+    
+    additional_tokens_list = [
+        "<tool_call>", "</tool_call>",
+        "<tool_response>", "</tool_response>",
+        "<think>", "</think>"
+    ]
+    num_buffer = special_tokens_num - len(special_tokens_list + additional_tokens_list)
+    buffer_tokens = [f"<|buffer{i}|>" for i in range(1, num_buffer + 1)] # 预留一定数量的token位置
+    all_special_tokens = special_tokens_list + additional_tokens_list + buffer_tokens
+    trainer = trainers.BpeTrainer(
+        vocab_size=vocab_size,
+        show_progress=True,
+        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+        special_tokens=all_special_tokens
+    )
+    texts = get_texts(data_path)
+    tokenizer.train_from_iterator(texts, trainer=trainer)
+    tokenizer.decoder = decoders.ByteLevel()
+    tokenizer.add_special_tokens(special_tokens_list)
+
+    os.makedirs(tokenizer_dir, exist_ok=True)
+    tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
+    tokenizer.model.save(tokenizer_dir)
+    tokenizer_json_path = os.path.join(tokenizer_dir, "tokenizer.json")
+    with open(tokenizer_json_path, 'r', encoding='utf-8') as f:
+        tokenizer_data = json.load(f)
+    for token_info in tokenizer_data.get('added_tokens', []):
+        if token_info['content'] not in special_tokens_list:
+            token_info['special'] = False
+    with open(tokenizer_json_path, 'w', encoding='utf-8') as f:
+        json.dump(tokenizer_data, f, ensure_ascii=False, indent=2)
+    
+    added_tokens_decoder = {}
+    for i, token in enumerate(all_special_tokens):
+        idx = tokenizer.token_to_id(token)
+        added_tokens_decoder[str(idx)] = {
+            "content": token,
+            "lstrip": False,
+            "normalized": False,
+            "rstrip": False,
+            "single_word": False,
+            "special": True if token in special_tokens_list else False
+        }
+
+    config = {
+        "add_bos_token": False,
+        "add_eos_token": False,
+        "add_prefix_space": False,
+        "added_tokens_decoder": added_tokens_decoder,
+        "additional_special_tokens": [t for t in special_tokens_list if t not in ["<|endoftext|>"]],
+        "bos_token": "<|im_start|>",
+        "clean_up_tokenization_spaces": False,
+        "eos_token": "<|im_end|>",
+        "legacy": True,
+        "model_max_length": 131072,
+        "pad_token": "<|endoftext|>",
+        "sp_model_kwargs": {},
+        "spaces_between_special_tokens": False,
+        "unk_token": "<|endoftext|>",
+        "image_token": "<|image_pad|>",
+        "audio_token": "<|audio_pad|>",
+        "video_token": "<|video_pad|>",
+        "vision_bos_token": "<|vision_start|>",
+        "vision_eos_token": "<|vision_end|>",
+        "audio_bos_token": "<|audio_start|>",
+        "audio_eos_token": "<|audio_end|>",
+        "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in content %}\n                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if true %}\n            {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if open_thinking is defined and open_thinking is true %}\n        {{- '<think>\\n' }}\n    {%- else %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+        "tokenizer_class": "PreTrainedTokenizerFast"
+    }
+
+    with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w", encoding="utf-8") as f:
+        json.dump(config, f, ensure_ascii=False, indent=4)
+    print("Tokenizer training completed.")
+
+def eval_tokenizer(tokenizer_dir):
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
+    messages = [
+        {"role": "system", "content": "你是一个优秀的聊天机器人，总是给我正确的回应！"},
+        {"role": "user", "content": '你来自哪里？'},
+        {"role": "assistant", "content": '我来自月球'},
+        {"role": "user", "content": '你到底来自哪里？'},
+        {"role": "assistant", "content": '我来自地球'}
+    ]
+    new_prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False
+    )
+    print('-'*100)
+    print(new_prompt)
+    print('-'*100)
+    print('tokenizer词表长度：', len(tokenizer))
+    model_inputs = tokenizer(new_prompt)
+    print('encoder长度：', len(model_inputs['input_ids']))
+    response = tokenizer.decode(model_inputs['input_ids'], skip_special_tokens=False)
+    print('decoder一致性：', response == new_prompt, "\n")
+    print('-'*100)
+    print('压缩率测试（Chars/Tokens）：')
+    test_texts = [
+        # 中文样本 (约200字)
+        "人工智能是计算机科学的一个分支，它企图了解智能的实质，并生产出一种新的能以人类智能相似的方式做出反应的智能机器，该领域的研究包括机器人、语言识别、图像识别、自然语言处理和专家系统等。人工智能从诞生以来，理论和技术日益成熟，应用领域也不断扩大，可以设想，未来人工智能带来的科技产品，将会是人类智慧的“容器”。人工智能可以对人的意识、思维的信息过程的模拟。人工智能不是人的智能，但能像人那样思考、也可能超过人的智能。",
+        "星际航行是指在星系内甚至星系间的空间中进行的航行。由于宇宙空间极其广阔，传统的化学火箭动力在恒星间航行时显得力不从心。科学家们提出了多种方案，包括离子推进器、核热火箭、甚至是利用反物质作为能源的设想。此外，曲率驱动和虫洞旅行等科幻概念也在理论物理研究中被反复探讨。尽管目前人类的足迹仅限于月球，但随着核聚变技术和材料科学的突破，前往火星乃至更遥远的太阳系边缘将成为可能。",
+        # 英文样本 (约200词/字符)
+        "Large language models (LLMs) are a type of artificial intelligence (AI) trained on vast amounts of text data to understand and generate human-like language. These models use deep learning techniques, specifically transformers, to process and predict the next word in a sequence. LLMs like GPT-4, Llama, and Claude have demonstrated remarkable capabilities in coding, translation, and creative writing. However, they also face challenges such as hallucinations, where the model generates factually incorrect information, and the need for significant computational resources.",
+        "The development of sustainable energy is crucial for the future of our planet. As climate change continues to impact global weather patterns, transitioning from fossil fuels to renewable sources like solar, wind, and hydroelectric power has become an urgent priority. Innovations in battery storage technology and smart grid management are essential to ensure a reliable energy supply. International cooperation and policy frameworks are also necessary to drive the global shift towards a greener economy and reduce carbon emissions.",
+        # 混合样本
+        "Python 是一种高级编程语言，以其简洁的语法和强大的生态系统而闻名。It is widely used in data science, machine learning, and web development. 开发者可以利用 NumPy, Pandas, and PyTorch 等库快速构建复杂的应用。学习 Python 的过程非常愉快，因为它的代码读起来就像英语一样。Whether you are a beginner or an expert, Python offers something for everyone.",
+    ]
+    
+    total_compression = 0
+    for i, text in enumerate(test_texts):
+        encoded = tokenizer.encode(text)
+        token_count = len(encoded)
+        char_count = len(text)
+        compression_ratio = char_count / token_count
+        total_compression += compression_ratio
+        print(f"样本 {i+1} | 字符数: {char_count:4} | Tokens: {token_count:3} | 压缩率: {compression_ratio:.2f}")
+    
+    print(f"平均压缩率: {total_compression / len(test_texts):.2f}")
+    print('-'*100)
+    print('流式解码（字节缓冲）测试：')
+    input_ids = model_inputs['input_ids']
+    token_cache = []
+    for tid in input_ids:
+        token_cache.append(tid)
+        current_decode = tokenizer.decode(token_cache)
+        if current_decode and '\ufffd' not in current_decode:
+            display_ids = token_cache[0] if len(token_cache) == 1 else token_cache
+            raw_tokens = [tokenizer.convert_ids_to_tokens(int(t)) for t in (token_cache if isinstance(token_cache, list) else [token_cache])]
+            print(f'Token ID: {str(display_ids):15} -> Raw: {str(raw_tokens):20} -> Decode Str: {current_decode}')
+            token_cache = []
+
+if __name__ == '__main__':
+    train_tokenizer(DATA_PATH, TOKENIZER_DIR, VOCAB_SIZE)
+    eval_tokenizer(TOKENIZER_DIR)