cleanup: remove deprecated doc files, update src training utils

.gitignore

@@ -7,11 +7,11 @@ wheels/
 *.egg-info
 
 # dataset
-dataset/pretrain_t2t_mini.jsonl
-dataset/sft_t2t_mini.jsonl
-dataset/dpo.jsonl
-dataset/rlaif.jsonl
-dataset/agent_rl.jsonl
-dataset/agent_rl_math.jsonl
+# dataset/pretrain_t2t_mini.jsonl
+# dataset/sft_t2t_mini.jsonl
+# dataset/dpo.jsonl
+# dataset/rlaif.jsonl
+# dataset/agent_rl.jsonl
+# dataset/agent_rl_math.jsonl
 # Virtual environments
 .venv

docs/dataset/index.md

@@ -1,33 +0,0 @@
-# data/ · 数据集
-
-`dataset/` 下**每类数据集一个文件**（已去掉 `_dataset` 后缀），公共辅助函数在 `common.py`。
-
-| 文件 | 数据集 | 用途 |
-| --- | --- | --- |
-| `pretrain.py` | `PretrainDataset` | 预训练（纯文本） |
-| `sft.py` | `SFTDataset` | 全量 SFT（chat 模板 + loss mask） |
-| `dpo.py` | `DPODataset` | DPO 偏好数据 |
-| `rlaif.py` | `RLAIFDataset` | RLHF/RLAIF prompt 采样 |
-| `agent_rl.py` | `AgentRLDataset` | Agent 强化学习轨迹 |
-| `vlm.py` | `VLMDataset` | 图文对（parquet） |
-| `vam.py` | `VAMDataset` | 全模态（图文 + 音频，parquet） |
-
-统一从 `dataset`（本仓库包）导入，例如 `from dataset import SFTDataset`。
-
-## 共同模式
-
-- 都继承 `torch.utils.data.Dataset`，实现 `__len__` / `__getitem__`。
-- 文本类通过 `tokenizer.apply_chat_template` 渲染对话；用 `bos_id/eos_id` 在 `generate_labels` 里**只对 assistant 回复计算损失**（prompt 部分 label=-100）。
-- 多模态（`vlm`/`vam`）额外加载图像/音频，并产出 `(input_ids, labels, 视觉/音频特征)` 元组。
-
-## 训练时拼 batch
-
-- `trainers/lm/*` 用 `SkipBatchSampler` + `DataLoader`；
-- `vlm` 用 `vlm_collate_fn`（把变长视觉特征 pad/stack）；
-- `vam` 的 `VAMDataset.__getitem__` 直接返回定长张量（文本+音频各层 label + 音频特征 + spk_emb）。
-
-## 要点（面试）
-
-- **loss mask 为什么只标 assistant？** 让模型只学习生成回复，不拟合用户输入/系统提示，避免「学用户说话」。
-- 多模态数据常用 parquet + `HFDataset.from_parquet`，便于大批量流式读取。
-- `pre_processing_chat` 随机插入 system prompt，`post_processing_chat` 随机去掉空 think 段，做数据增广。

docs/encoders/index.md

@@ -1,14 +0,0 @@
-# encoders/ · 多模态编码器
-
-把「外部模态」的原始信号编码为模态特征，冻结预训练权重、不参与 LLM 主干训练。
-
-| 文件 | 类 | 作用 |
-| --- | --- | --- |
-| `encoders/vision/siglip.py` | `SiglipVisionEncoder` | 图像 → patch 特征（SigLIP） |
-| `encoders/audio/sensevoice.py` | `SenseVoiceAudioEncoder` | 语音 → 语义/声学特征（SenseVoice） |
-
-## 设计要点（面试）
-
-- **encoder 冻结 + 只训 projector**：大模型主干易灾难性遗忘，冻结预训练视觉/音频 encoder、只训轻量投影层，是高效多模态对齐的常用做法。
-- **`encoders/` 与 `projectors/` 分离**：encoder 负责「理解模态」，projector 负责「对齐到 LLM 空间」，二者职责清晰、可独立替换。
-- **可扩展性**：加新模态 = 加一对 encoder+projector + 占位 token，主干不动（符合本仓库分层初衷）。

docs/overview.md

@@ -1,108 +0,0 @@
-# Omni 开发文档
-
-> Omni — 把 MiniMind / MiniMind-V / MiniMind-O 三套代码融合成一个按模态分层（文本 / 视觉 / 全模态）的训练 + 推理框架。面向**学习与面试**的模块化开发文档。
-
-文档目录与 `src/` 代码模块一一对应：
-
-```
-docs/                 ↔  src/
-├── architecture/     #   概念总览（入口，非代码模块）
-├── core/             ↔  src/core/        纯 Transformer 组件
-├── models/           ↔  src/models/      lm / vlm / vam 模型拼装
-├── dataset/          ↔  src/dataset/     数据集
-├── encoders/         ↔  src/encoders/    视觉/音频编码器
-├── projectors/       ↔  src/projectors/  模态桥接层
-├── serve/            ↔  src/serve/       实时语音会话
-├── trainers/         ↔  src/trainers/    训练脚本
-├── utils/            ↔  src/utils/       训练/分布式/checkpoint 工具
-└── interview/        #   面试速查（学习辅助）
-```
-
-## 目录
-
-### 架构（概念入口）
-| 文档 | 说明 |
-| --- | --- |
-| [架构总览](architecture/overview.md) | 能力分层设计、模型能力矩阵、前向数据流、训练入口、关键设计取舍 |
-
-### core/（基础组件）
-| 文档 | 说明 |
-| --- | --- |
-| [组件索引](core/index.md) | 与模态无关的纯 Transformer 组件总览 |
-| [RMSNorm](core/norm.md) | RMSNorm 实现与数值技巧 |
-| [RoPE](core/rope.md) | 旋转位置编码（含 YaRN）、`repeat_kv` |
-| [Attention](core/attention.md) | 带 QK-Norm + GQA 的注意力 |
-| [MLP / MoE](core/mlp.md) | SwiGLU FFN 与 MoE 前馈 |
-| [Block](core/block.md) | Pre-Norm Transformer 块、MoE 可插拔 |
-
-### models/（模型拼装）
-| 文档 | 说明 |
-| --- | --- |
-| [模型索引](models/index.md) | `lm` / `vlm` / `vam` 三子包拼装、继承链、共享主干 |
-| [LM](models/lm.md) | 纯文本主干 `LM` + `LMForCausalLM` |
-| [VLM](models/vlm.md) | 文本 + 图像（SigLIP 编码器 + 视觉投影） |
-| [VAM](models/vam.md) | 文本 + 图像 + 语音，`TalkerModule` 双 head |
-
-### dataset/（数据集）
-| 文档 | 说明 |
-| --- | --- |
-| [数据集索引](dataset/index.md) | `dataset/` 每类一个文件的数据集、loss mask、批次拼接 |
-
-### encoders/（多模态编码器）
-| 文档 | 说明 |
-| --- | --- |
-| [编码器索引](encoders/index.md) | `SiglipVisionEncoder` / `SenseVoiceAudioEncoder` |
-
-### projectors/（模态桥接层）
-| 文档 | 说明 |
-| --- | --- |
-| [桥接层索引](projectors/index.md) | `MMVisionProjector` / `MMAudioProjector`、注入范式 |
-
-### serve/（实时语音会话）
-| 文档 | 说明 |
-| --- | --- |
-| [服务索引](serve/index.md) | `SileroVAD` / `RealtimeSession` 端到端语音链路 |
-
-### trainers/（训练脚本）
-| 文档 | 说明 |
-| --- | --- |
-| [训练索引](trainers/index.md) | 按模态组织的 trainer 模块分布 |
-| [Trainers](trainers/trainers.md) | 各 trainer 模块概览与通用训练循环 |
-
-### utils/（训练工具）
-| 文档 | 说明 |
-| --- | --- |
-| [工具索引](utils/index.md) | `training` / `checkpoint` / `distributed` / `multimodal` 工具分布 |
-| [配置与命令行](utils/config-and-cli.md) | YAML 配置驱动训练、`apply_config` 机制、tokenizer 训练、启动示例 |
-
-### 面试速查
-| 文档 | 说明 |
-| --- | --- |
-| [面试准备](interview/面试速查.md) | 高频问题 + 一句话答法，按主题组织 |
-
-## 命名约定（本仓库）
-
-类名已在融合过程中统一重命名，阅读代码/面试时对照：
-
-| 旧名（上游） | 本仓库 |
-| --- | --- |
-| `MiniMindModel` | `LM` |
-| `MiniMindForCausalLM` | `LMForCausalLM` |
-| `MiniMindConfig` | `LMConfig` |
-| `MiniMindVLM` | `VLM` |
-| `MiniMindOmni` | `VAM` |
-| `OmniConfig` | `VAMConfig` |
-
-包结构（已扁平化，无 `omni` 中间层）：
-
-```
-src/
-├── core/        # 纯组件：norm / rope / attention / mlp / block
-├── models/      # lm / vlm / vam 三个子包（config + model）
-├── dataset/     # 每类一个文件的数据集
-├── encoders/    # vision / audio 编码器
-├── projectors/  # 视觉/音频桥接层
-├── serve/       # 实时语音会话
-├── trainers/    # lm / vlm / vam 训练脚本
-└── utils/       # training / distributed / checkpoint / multimodal
-```

docs/projectors/index.md

@@ -1,21 +0,0 @@
-# projectors/ · 模态桥接层
-
-把 encoder 输出的模态特征投影到 LLM 隐藏维度（`hidden_size`），随后替换 `input_ids` 中的占位 token 位置。
-
-| 文件 | 类 | 作用 |
-| --- | --- | --- |
-| `projectors/vision.py` | `MMVisionProjector` | 视觉特征 → LLM `hidden_size` |
-| `projectors/audio.py` | `MMAudioProjector` | 音频特征 → LLM `hidden_size` |
-
-## 注入范式
-
-```
-原始输入(图/音)
-   └─ encoder      → 模态特征
-        └─ projector → 投影到 hidden_size
-             └─ 替换 input_ids 中的占位 token 位置的 embedding
-                  └─ 进入共享 LM 主干
-```
-
-- 文本侧用特殊占位 token（`<|image_pad|>` / `<|audio_pad|>`）预留位置；
-- `forward` 把投影特征写到这些位置，LLM 对文本/视觉/音频 token 一视同仁。

docs/trainers/index.md

@@ -1,13 +0,0 @@
-# trainers/ · 训练脚本
-
-按模态组织的训练入口，每个脚本暴露 `main(default_config=None)`（可由 `python -m trainers.<mod>` 或 `runs/*.sh` 调用）。
-
-- `trainers/lm/`：pretrain / full_sft / lora / dpo / distillation / ppo / grpo / agent / rollout_engine / train_tokenizer
-- `trainers/vlm/`：pretrain / full_sft
-- `trainers/vam/`：full_sft
-
-相关：配置与 CLI 机制见 [../utils/config-and-cli.md](../utils/config-and-cli.md)（基于 `utils/training.apply_config`）。
-
-详见：
-
-- [trainers.md](trainers.md) — 各 trainer 模块概览与通用训练循环

docs/trainers/trainers.md

@@ -1,50 +0,0 @@
-# trainers/trainers.md · 训练脚本
-
-`trainers/` 按模态分子包，每个脚本暴露 `main(default_config=None)`（可由 `python -m trainers.<mod>` 或 `runs/*.sh` 调用）。
-
-## 文本（lm/）
-
-| 脚本 | 任务 | 关键损失/算法 |
-| --- | --- | --- |
-| `pretrain.py` | 预训练 | 下一 token CE |
-| `full_sft.py` | 全量 SFT | CE（loss mask 仅 assistant） |
-| `lora.py` | LoRA 微调 | 低秩适配，仅训 A/B |
-| `dpo.py` | DPO | 偏好对齐（参考比损失） |
-| `distillation.py` | 知识蒸馏 | 师生 KL |
-| `ppo.py` | PPO | Actor-Critic + 奖励 |
-| `grpo.py` | GRPO | 分组相对策略优化 |
-| `agent.py` | Agent RL | 工具调用强化学习 |
-| `rollout_engine.py` | — | torch / sglang 推理引擎（被 ppo/grpo/agent 复用） |
-| `train_tokenizer.py` | — | tokenizer 训练（学习用），结果保存到 `checkpoint/tokenizer/` |
-
-## 视觉（vlm/）
-
-- `pretrain.py`：视觉预训练
-- `full_sft.py`：视觉 SFT（含 `vlm_collate_fn`）
-
-## 全模态（vam/）
-
-- `full_sft.py`：全模态 SFT（文本 + 视觉 + 音频，双 head 损失）
-
-## 通用训练循环（以 full_sft 为例）
-
-```python
-for epoch in range(epochs):
-    loader = DataLoader(ds, batch_sampler=SkipBatchSampler(...))
-    for step, (input_ids, labels) in enumerate(loader):
-        loss = model(input_ids, labels=labels).loss + res.aux_loss
-        loss = loss / accumulation_steps
-        scaler.scale(loss).backward()
-        if step % accumulation_steps == 0:
-            clip_grad_norm_; scaler.step(optimizer); zero_grad()
-        # 定期保存权重到 save_dir + 保存 optimizer/ckpt 到 checkpoint/
-```
-
-支持：分布式（`init_distributed_mode` + `DistributedDataParallel`）、混合精度（`autocast` + `GradScaler`）、梯度累积、断点续训（`from_resume`）、可选 wandb/swanlab。
-
-## 要点（面试）
-
-- **`SkipBatchSampler`**：分布式下跳过已训 step，配合 `from_resume` 实现精确续训。
-- **`aux_loss`**：MoE 路由均衡损失，只在 `use_moe` 时非 0，需显式加到总损失。
-- **RL trainer 复用 `rollout_engine`**：生成样本与训练解耦，可换 torch / sglang 后端。
-- 保存分两份：`save_dir`（最终权重 `.pth`）+ `checkpoint/`（optimizer/scheduler 状态用于续训）。`train_tokenizer.py` 额外把训练出的 tokenizer 写到 `checkpoint/tokenizer/`。

docs/utils/config-and-cli.md

@@ -1,83 +0,0 @@
-# utils/config-and-cli.md · 配置与命令行
-
-## 机制
-
-`utils/training.apply_config(parser, default_config=None)` 让 **YAML 成为 argparse 默认值，CLI 显式参数覆盖 YAML**：
-
-```python
-def apply_config(parser, default_config=None):
-    pre, _ = parser.parse_known_args()          # 先读出 --config
-    config_path = getattr(pre, 'config', None) or default_config
-    if config_path and os.path.exists(config_path):
-        defaults = _load_yaml_config(config_path)
-        parser.set_defaults(**defaults)          # 注入默认值
-    return parser.parse_args()                   # 最终解析（CLI 覆盖）
-```
-
-`_load_yaml_config` 把 YAML 的 `model / train / paths` 三段**扁平化**为顶层键：
-`model.hidden_size → hidden_size`，`train.epochs → epochs`，`paths.data_path → data_path` ……
-这些键名与 trainer 的 `--hidden_size` / `--epochs` / `--data_path` 等 argparse 参数一一对应。
-
-## 配置文件位置
-
-```
-configs/
-├── model/      # lm / lm_moe / vlm / vlm_moe / vam / vam_moe .yaml
-└── tokenizer/  # tokenizer.json / tokenizer_config.json
-```
-
-每个 YAML 例：
-
-```yaml
-model:
-  hidden_size: 768
-  num_hidden_layers: 8
-  use_moe: 0
-  vocab_size: 6400
-train:
-  epochs: 2
-  batch_size: 16
-  learning_rate: 1.0e-5
-  from_weight: pretrain
-paths:
-  save_dir: checkpoint/lm
-  data_path: dataset/sft.jsonl
-```
-
-## 配置如何真正生效
-
-trainer 用 `LMConfig(**vars(args))`（或 `VLMConfig` / `VAMConfig`）构造模型配置，
-因此 YAML 里**所有** `model` 字段（含 `vocab_size`、head 数、MoE 专家数、talker 层数等）
-都会驱动模型结构，而非仅少数硬编码字段。
-
-## 启动
-
-```bash
-bash runs/lm.sh                                   # 默认 configs/model/lm.yaml
-bash runs/vlm.sh --config configs/model/vlm_moe.yaml
-bash runs/vam.sh --epochs 10                       # 覆盖单字段
-bash runs/train_tokenizer.sh                      # 训练 tokenizer（学习用），保存到 checkpoint/tokenizer/
-```
-
-## Tokenizer 训练
-
-`trainers/lm/train_tokenizer.py` 仅供学习参考（MiniMind 已自带 tokenizer，重复训练会导致词表不统一）。
-训练得到的 tokenizer 直接保存到 **`checkpoint/tokenizer/`**（与模型权重同目录），包含：
-
-- `tokenizer.json` / `vocab.json` / `merges.txt`：BPE 词表
-- `tokenizer_config.json`：special token、chat template 等配置
-
-常用参数：
-
-```bash
-bash runs/train_tokenizer.sh --data_path dataset/sft_t2t_mini.jsonl \
-                             --vocab_size 6400 \
-                             --checkpoint_dir ../checkpoint \
-                             --no_eval
-```
-
-## 要点（面试）
-
-- 设计亮点：**单一事实源（YAML）+ CLI 覆盖**，实验可复现、参数可微调。
-- `set_defaults` 注入后，argparse 仍允许命令行覆盖 → 优先级 `CLI > YAML > 代码默认`。
-- 扁平化约定要求 YAML 键名与 `--arg` 名一致，故 `trainers` 的参数命名需与 YAML 段对齐。

docs/utils/index.md

@@ -1,14 +0,0 @@
-# utils/ · 训练工具
-
-训练与分布式通用工具，被 `trainers/` 复用。
-
-| 文件 | 内容 | 作用 |
-| --- | --- | --- |
-| `training.py` | `apply_config` / `lm_checkpoint` / `init_model` / `SkipBatchSampler` / `get_lr` | YAML 配置驱动、checkpoint 读写、采样器 |
-| `checkpoint.py` | `save_checkpoint` / `load_checkpoint` / `iter_module_state_dict` | 权重序列化辅助 |
-| `distributed.py` | 分布式初始化辅助 | 多卡训练 |
-| `multimodal.py` | `init_vlm_model` / `vlm_checkpoint` / `init_omni_model` / `omni_checkpoint` | VLM / VAM 模型装配与 checkpoint |
-
-详见：
-
-- [配置与命令行](config-and-cli.md) — YAML 配置如何驱动训练、`apply_config` 机制、tokenizer 训练、启动示例

src/trainers/lm/agent.py

@@ -373,6 +373,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Agent RL")
     parser.add_argument('--config', type=str, default=None, help='YAML 配置路径，其字段作为 argparse 默认值，CLI 显式传参可覆盖')
     parser.add_argument("--save_dir", type=str, default="../checkpoint", help="模型保存目录")
+    parser.add_argument("--tokenizer_dir", type=str, default="checkpoint/tokenizer", help="tokenizer 目录路径")
     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="批次大小")
@@ -432,9 +433,9 @@ if __name__ == "__main__":
         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)
+    model, tokenizer = init_model(lm_config, args.from_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, device=args.device)
 
-    ref_model, _ = init_model(lm_config, args.from_weight, device=args.device)
+    ref_model, _ = init_model(lm_config, args.from_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, device=args.device)
     ref_model = ref_model.eval().requires_grad_(False)
 
     reward_model = LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16)

src/trainers/lm/distillation.py

@@ -146,6 +146,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Knowledge Distillation")
     parser.add_argument('--config', type=str, default=None, help='YAML 配置路径，其字段作为 argparse 默认值，CLI 显式传参可覆盖')
     parser.add_argument("--save_dir", type=str, default="../checkpoint", help="模型保存目录")
+    parser.add_argument("--tokenizer_dir", type=str, default="checkpoint/tokenizer", help="tokenizer 目录路径")
     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")
@@ -202,9 +203,9 @@ if __name__ == "__main__":
         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)
+    model, tokenizer = init_model(lm_config_student, args.from_student_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, 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, _ = init_model(lm_config_teacher, args.from_teacher_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, 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')

src/trainers/lm/dpo.py

@@ -130,6 +130,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind DPO (Direct Preference Optimization)")
     parser.add_argument('--config', type=str, default=None, help='YAML 配置路径，其字段作为 argparse 默认值，CLI 显式传参可覆盖')
     parser.add_argument("--save_dir", type=str, default="../checkpoint", help="模型保存目录")
+    parser.add_argument("--tokenizer_dir", type=str, default="checkpoint/tokenizer", help="tokenizer 目录路径")
     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")
@@ -180,10 +181,10 @@ if __name__ == "__main__":
         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)
+    model, tokenizer = init_model(lm_config, args.from_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, 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, _ = init_model(lm_config, args.from_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, 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')

src/trainers/lm/full_sft.py

@@ -83,6 +83,7 @@ def main(default_config=None):
     parser = argparse.ArgumentParser(description="MiniMind Full SFT")
     parser.add_argument('--config', type=str, default=None, help='YAML 配置路径，其字段作为 argparse 默认值，CLI 显式传参可覆盖')
     parser.add_argument("--save_dir", type=str, default="../checkpoint", help="模型保存目录")
+    parser.add_argument("--tokenizer_dir", type=str, default="checkpoint/tokenizer", help="tokenizer 目录路径")
     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")
@@ -132,7 +133,7 @@ def main(default_config=None):
         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)
+    model, tokenizer = init_model(lm_config, args.from_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, 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'))

src/trainers/lm/grpo.py

@@ -205,6 +205,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind GRPO (Group Relative Policy Optimization)")
     parser.add_argument('--config', type=str, default=None, help='YAML 配置路径，其字段作为 argparse 默认值，CLI 显式传参可覆盖')
     parser.add_argument("--save_dir", type=str, default="../checkpoint", help="模型保存目录")
+    parser.add_argument("--tokenizer_dir", type=str, default="checkpoint/tokenizer", help="tokenizer 目录路径")
     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")
@@ -271,9 +272,9 @@ if __name__ == "__main__":
     # ========== 5. 初始化模型和数据 ==========
     base_weight = args.from_weight
     # Policy模型
-    model, tokenizer = init_model(lm_config, base_weight, device=args.device)
+    model, tokenizer = init_model(lm_config, base_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, device=args.device)
     # Reference模型
-    ref_model, _ = init_model(lm_config, base_weight, device=args.device)
+    ref_model, _ = init_model(lm_config, base_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, 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)

src/trainers/lm/lora.py

@@ -77,6 +77,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind LoRA Fine-tuning")
     parser.add_argument('--config', type=str, default=None, help='YAML 配置路径，其字段作为 argparse 默认值，CLI 显式传参可覆盖')
     parser.add_argument("--save_dir", type=str, default="../checkpoint", help="模型保存目录")
+    parser.add_argument("--tokenizer_dir", type=str, default="checkpoint/tokenizer", help="tokenizer 目录路径")
     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")
@@ -126,7 +127,7 @@ if __name__ == "__main__":
         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)
+    model, tokenizer = init_model(lm_config, args.from_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, device=args.device)
     apply_lora(model)
     
     # 统计参数

src/trainers/lm/ppo.py

@@ -295,6 +295,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind PPO (Proximal Policy Optimization)")
     parser.add_argument('--config', type=str, default=None, help='YAML 配置路径，其字段作为 argparse 默认值，CLI 显式传参可覆盖')
     parser.add_argument("--save_dir", type=str, default="../checkpoint", help="模型保存目录")
+    parser.add_argument("--tokenizer_dir", type=str, default="checkpoint/tokenizer", help="tokenizer 目录路径")
     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")
@@ -365,8 +366,8 @@ if __name__ == "__main__":
     # ========== 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)
+    actor_model, tokenizer = init_model(lm_config, base_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, device=args.device)
+    ref_model, _ = init_model(lm_config, base_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, 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'

src/trainers/lm/pretrain.py

@@ -82,6 +82,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Pretraining")
     parser.add_argument('--config', type=str, default=None, help='YAML 配置路径，其字段作为 argparse 默认值，CLI 显式传参可覆盖')
     parser.add_argument("--save_dir", type=str, default="../checkpoint", help="模型保存目录")
+    parser.add_argument("--tokenizer_dir", type=str, default="checkpoint/tokenizer", help="tokenizer 目录路径")
     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")
@@ -131,7 +132,7 @@ if __name__ == "__main__":
         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)
+    model, tokenizer = init_model(lm_config, args.from_weight, save_dir=args.save_dir, tokenizer_dir=args.tokenizer_dir, 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'))

src/utils/training.py

@@ -158,8 +158,9 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc
         return None
 
 
-def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../checkpoint', device='cuda'):
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+def init_model(lm_config, from_weight='pretrain', save_dir='../checkpoint', tokenizer_dir=None, device='cuda'):
+    tokenizer_dir = tokenizer_dir or os.path.join(save_dir, 'tokenizer')
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
     model = LMForCausalLM(lm_config)
 
     if from_weight != 'none':