Files
ars-opd-rebuild/scripts/train_whitebox.sh
T
iomgaa b7f24d635e 层2/U5: white-box OPD 自包含训练脚本(train_whitebox.py + .sh)
对应 docs/03 §5 U5,对齐层 1 train_sft 骨架,换成层 2 装配:
- 双模型:student Qwen3-0.6B(fp32+bf16混训) + teacher Qwen3-4B(bf16 推理)
- prompt_only collator(无 teacher 缓存,现场 on-policy 生成)
- DistillTrainer 装配:teacher_model/teacher_tokenizer + beta/温度/生成参数
- FULL = §5 默认(beta=1、lr=1e-6、B=4×GA=4×4卡=全局64、max_new_tokens=1024)
- 首 prompt 自检(末尾须为生成引导符);sanity=50步冒烟
- .sh 前置清单强调白盒扛两份全词表 logits、OOM 阶梯、首跑下载 4B teacher

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-19 05:23:42 -04:00

32 lines
1.9 KiB
Bash
Executable File
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env bash
# 层 2white-box OPD 训练(远程 gpu-a800-060 专用;本地不跑训练)。
#
# 用法(tmux 内执行,日志实时可查):
# bash scripts/train_whitebox.sh sanity # 50 步冒烟:看首 prompt 自检 + KL loss +
# # 生成 token 数;预期见 loss 毛刺(梯度爆炸实况)
# bash scripts/train_whitebox.sh # 正式:1k 子集 1 epoch
#
# 前置检查清单:
# 1. nvidia-smi 确认下方 GPUS 四张卡空闲(只许用 8 卡中的 4 张,严禁自动选卡)。
# ⚠️ 白盒显存比层 1 紧:student 训练全套 + teacher(4B) 推理副本 + **两份**全词表
# logits(student/teacher),§5 估算 B=4/T=2048 起步安全,但首跑必须盯 nvidia-smi;
# 若 OOM,降 per_device_train_batch_size 到 2,仍不够再开 gradient_checkpointing
# (改 DistillConfig,注意 checkpointing 与 generate 的 use_cache 交互)。
# 2. data/dapo-math-17k-unique.parquet 已在(层 2 无需 teacher 缓存,纯 prompt-only):
# scp data/dapo-math-17k-unique.parquet <远程>:/data/zym/ars-opd-rebuild/data/
# 3. 代码最新:git -C /data/zym/ars-opd-rebuild pull
# 4. 首跑会下载 teacher Qwen3-4BGB 级)到 HF_HOME,确保 /data 有空间
set -euo pipefail
cd "$(dirname "$0")/.." # 锚定仓库根
GPUS=0,1,2,3 # ⚠️ 改这里前先 nvidia-smi
MODE=${1:-full}
export CUDA_VISIBLE_DEVICES=$GPUS
export PYTHONUNBUFFERED=1 # 禁止日志缓存(CLAUDE.md §5
export PYTORCH_ALLOC_CONF=expandable_segments:True # 变长生成序列易碎片化,按需扩段
export HF_ENDPOINT=${HF_ENDPOINT:-https://hf-mirror.com}
export HF_HOME=${HF_HOME:-/data/zym/hf_cache} # 模型缓存落 /data,根分区已满
torchrun --nproc_per_node=4 --master_port=29572 scripts/train_whitebox.py "$MODE"