TI Coding Agent Training Probe
这个仓库用于复现一组 coding-agent SFT probing 实验:从 ModelScope 下载已经构造好的 Open-SWE-Traces probe 训练/验证数据集,下载 Qwen3.5-9B 和 Qwen3.6-27B,然后用 ModelScope-SWIFT 依次跑四个 1 epoch 训练任务。Hugging Face 数据源仍保留为 fallback。
目录
third_party/modelscope-swift/: ModelScope-SWIFT submodule。scripts/setup_env.sh: 一键创建 repo 内.venv并安装本项目和 SWIFT。scripts/download_dataset.py: 下载 ModelScope 或 Hugging Face 数据集并解压train.jsonl、validation.jsonl。scripts/download_models.sh: 下载 Qwen3.5-9B 和 Qwen3.6-27B 到models/。scripts/train_qwen35_9b_lora.sh: Qwen3.5-9B rank=32 LoRA。scripts/train_qwen35_9b_full.sh: Qwen3.5-9B bf16 full SFT。scripts/train_qwen36_27b_lora.sh: Qwen3.6-27B rank=32 LoRA。scripts/train_qwen36_27b_full.sh: Qwen3.6-27B bf16 full SFT。scripts/train_qwen36_27b_megatron_full.sh: 基于 Megatron-SWIFT/MCore-Bridge 的 Qwen3.6-27B full SFT 入口。scripts/run_all_experiments.sh: 按 LoRA 9B -> full 9B -> LoRA 27B -> full 27B 的顺序执行完整实验。runs/: TensorBoard 日志目录。logs/: 训练 stdout/stderr 和实际命令记录。outputs/: checkpoint 和最终模型权重输出目录。data/: 下载后的训练和验证数据,默认不进 git。models/: 下载后的 base model,默认不进 git。
环境部署
在 B300 上使用:
cd /ssd/workspace/yi/ti_coding_agent_probe
git submodule update --init --recursive
./scripts/setup_env.sh
脚本默认使用国内源和 B300 机器特定代理,避免直接访问外网导致下载失败:
http_proxy=http://100.72.0.101:8888
https_proxy=http://100.72.0.101:8888
HTTP_PROXY=http://100.72.0.101:8888
HTTPS_PROXY=http://100.72.0.101:8888
PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/
UV_DEFAULT_INDEX=https://mirrors.aliyun.com/pypi/simple/
UV_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/
HF_ENDPOINT=https://hf-mirror.com
其中 proxy 是当前 B300/g0049 环境的 machine-specific 配置。换到其他机器、其他集群或本地环境时,需要先确认代理地址是否仍然可用;如果不可用,应覆盖 http_proxy、https_proxy、HTTP_PROXY、HTTPS_PROXY,或在无墙环境下显式清空这些变量。PyPI/uv/Hugging Face 镜像也可以通过上面的环境变量覆盖。
Python 依赖安装在仓库内 .venv,不会写入系统 Python。安装逻辑是:
- 先在
.venv内安装uv。 - 用 SWIFT 官方推荐的
uv pip install -e third_party/modelscope-swift --torch-backend=auto安装 SWIFT 源码依赖。 - 再安装本 probe repo。
- 默认额外补充
deepspeed<0.19、liger-kernel、tensorboard、nvitop。
如果要安装 SWIFT 的全量 optional 依赖:
INSTALL_SWIFT_ALL=1 ./scripts/setup_env.sh
如果只想安装 SWIFT core,不补额外训练包:
INSTALL_TRAINING_EXTRAS=0 ./scripts/setup_env.sh
如果要追加其他包:
EXTRA_UV_PACKAGES="flash-attn==2.8.3" ./scripts/setup_env.sh
下载数据集
默认数据源是 ModelScope:
eigentom/ti_coding_agent_training_probe_20260624
如果数据集需要鉴权,设置任一 token 环境变量即可:
export MODELSCOPE_API_TOKEN=<modelscope-api-key>
# 或
export MS_TOKEN=<modelscope-api-key>
下载训练和验证数据:
./scripts/download_dataset.py
也可以显式指定 ModelScope repo:
export MS_DATASET_REPO_ID=eigentom/ti_coding_agent_training_probe_20260624
./scripts/download_dataset.py --backend modelscope
Hugging Face 后端仍保留为 fallback。使用 HF 时,如果环境里设置了 HF_TOKEN,脚本会用 token owner 自动拼成 owner/ti_coding_agent_training_probe_20260624。也可以显式指定:
export HF_ENDPOINT=https://hf-mirror.com
export HF_DATASET_REPO_ID=<owner>/ti_coding_agent_training_probe_20260624
./scripts/download_dataset.py --backend huggingface
输出:
data/raw/training_probe/: 原始下载文件。data/processed/training_probe/train.jsonldata/processed/training_probe/validation.jsonldata/processed/training_probe/train.parquetdata/processed/training_probe/validation.parquet
训练数据里 system、user、tool 消息带 loss=false,只有 assistant 轨迹带 loss=true。system prompt 会作为上下文参与 attention,但不作为预测目标计算 loss。
下载模型
./scripts/download_models.sh
默认模型 ID:
Qwen/Qwen3.5-9BQwen/Qwen3.6-27B
如果 Hugging Face 上实际模型 ID 有变化,可以覆盖:
export QWEN35_9B_MODEL_ID=<actual-9b-id>
export QWEN36_27B_MODEL_ID=<actual-27b-id>
./scripts/download_models.sh
模型下载脚本使用新版 Hugging Face CLI:
hf download <model-id> --local-dir <dir> --max-workers <n>
旧版 huggingface-cli download 和 --local-dir-use-symlinks 参数已经不再使用。默认并发是 HF_DOWNLOAD_MAX_WORKERS=8,如果网络或镜像不稳定可以调小:
export HF_DOWNLOAD_MAX_WORKERS=4
export HF_TOKEN=<optional-token>
./scripts/download_models.sh
单步训练
每个训练脚本默认:
num_train_epochs=1lora_rank=32torch_dtype=bfloat16save_steps=1000eval_steps=1000report_to=tensorboardmax_length=262144warmup_ratio=0.1lr_scheduler_type=cosine- full SFT:
learning_rate=1e-5 - LoRA:
learning_rate=5e-5 - 默认每卡训练 batch size 为
1 - 默认
NPROC_PER_NODE=8 - 默认
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 - 默认
deepspeed=zero2
B300/g0049 当前是 8 张 B300,每卡约 275GB 显存。由于本实验默认上下文长度是 262144,预设 batch size 保守取 1,防止 27B/full 或长样本直接 OOM。吞吐测试时可以从脚本外部调大,建议先从 LoRA/9B 开始试。
常用覆盖方式:
# 全部训练统一覆盖
export PER_DEVICE_BATCH_SIZE=2
export GRAD_ACCUM_STEPS=2
# 只覆盖 LoRA 或 full
export LORA_PER_DEVICE_BATCH_SIZE=2
export FULL_PER_DEVICE_BATCH_SIZE=1
export LORA_LEARNING_RATE=5e-5
export FULL_LEARNING_RATE=1e-5
# 只覆盖某一个 run;优先级最高
export QWEN35_9B_LORA_R32_PER_DEVICE_BATCH_SIZE=2
export QWEN35_9B_LORA_R32_GRAD_ACCUM_STEPS=2
export QWEN36_27B_FULL_BF16_PER_DEVICE_BATCH_SIZE=1
# scheduler/warmup
export WARMUP_RATIO=0.1
export LR_SCHEDULER_TYPE=cosine
# 短跑 throughput/debug,不跑完整 epoch
export MAX_STEPS=10
命令:
./scripts/train_qwen35_9b_lora.sh
./scripts/train_qwen35_9b_full.sh
./scripts/train_qwen36_27b_lora.sh
./scripts/train_qwen36_27b_full.sh
Megatron-SWIFT
仓库也提供了基于 Megatron-SWIFT/MCore-Bridge 的 full SFT 入口,参考官方 quick-start 的 megatron sft 写法:
./scripts/train_qwen36_27b_megatron_full.sh
默认参数:
NPROC_PER_NODE=8CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7MEGATRON_MODEL=Qwen/Qwen3.6-27BTENSOR_MODEL_PARALLEL_SIZE=4PIPELINE_MODEL_PARALLEL_SIZE=1CONTEXT_PARALLEL_SIZE=1MICRO_BATCH_SIZE=1GLOBAL_BATCH_SIZE=8MAX_LENGTH=262144LR=1e-5MIN_LR=1e-6LR_WARMUP_FRACTION=0.1
Megatron 多机/共享盘训练要求 dataset cache 一致。脚本默认把 MODELSCOPE_CACHE、Megatron 输出和日志放在共享路径:
/mnt/beegfs/workspace/ti_coding_agent_probe/
如果换机器或换共享盘,需要覆盖:
export BEEGFS_ROOT=/mnt/beegfs/workspace/ti_coding_agent_probe
export MODELSCOPE_CACHE=$BEEGFS_ROOT/modelscope_cache
export MEGATRON_OUTPUT_ROOT=$BEEGFS_ROOT/megatron_outputs
export MEGATRON_LOG_ROOT=$BEEGFS_ROOT/megatron_logs
一键完整实验
确认 GPU 空闲后执行:
./scripts/run_all_experiments.sh
执行顺序固定为:
- Qwen3.5-9B LoRA
- Qwen3.5-9B bf16 full SFT
- Qwen3.6-27B LoRA
- Qwen3.6-27B bf16 full SFT
TensorBoard
./scripts/tensorboard.sh
训练日志会写到 runs/<run_name>/。SWIFT/Transformers 的 TensorBoard 标量通常包括 loss、learning rate、eval loss、runtime、samples/sec、steps/sec 等 throughput 指标;同时 stdout 会保存在 logs/<run_name>.log。