# 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/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 上使用：

```bash
cd /ssd/workspace/yi/ti_coding_agent_probe
git submodule update --init --recursive
./scripts/setup_env.sh
```

脚本会显式设置 B300 代理和国内 PyPI 源：

```bash
http://100.72.0.101:8888
https://mirrors.aliyun.com/pypi/simple/
```

Python 依赖安装在仓库内 `.venv`，不会写入系统 Python。安装逻辑是：

1. 先在 `.venv` 内安装 `uv`。
2. 用 SWIFT 官方推荐的 `uv pip install -e third_party/modelscope-swift --torch-backend=auto` 安装 SWIFT 源码依赖。
3. 再安装本 probe repo。
4. 默认额外补充 `deepspeed<0.19`、`liger-kernel`、`tensorboard`、`nvitop`。

如果要安装 SWIFT 的全量 optional 依赖：

```bash
INSTALL_SWIFT_ALL=1 ./scripts/setup_env.sh
```

如果只想安装 SWIFT core，不补额外训练包：

```bash
INSTALL_TRAINING_EXTRAS=0 ./scripts/setup_env.sh
```

如果要追加其他包：

```bash
EXTRA_UV_PACKAGES="flash-attn==2.8.3" ./scripts/setup_env.sh
```

## 下载数据集

默认数据源是 ModelScope：

```text
eigentom/ti_coding_agent_training_probe_20260624
```

如果数据集需要鉴权，设置任一 token 环境变量即可：

```bash
export MODELSCOPE_API_TOKEN=<modelscope-api-key>
# 或
export MS_TOKEN=<modelscope-api-key>
```

下载训练和验证数据：

```bash
./scripts/download_dataset.py
```

也可以显式指定 ModelScope repo：

```bash
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`。也可以显式指定：

```bash
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.jsonl`
- `data/processed/training_probe/validation.jsonl`
- `data/processed/training_probe/train.parquet`
- `data/processed/training_probe/validation.parquet`

训练数据里 `system`、`user`、`tool` 消息带 `loss=false`，只有 assistant 轨迹带 `loss=true`。system prompt 会作为上下文参与 attention，但不作为预测目标计算 loss。

## 下载模型

```bash
./scripts/download_models.sh
```

默认模型 ID：

- `Qwen/Qwen3.5-9B`
- `Qwen/Qwen3.6-27B`

如果 Hugging Face 上实际模型 ID 有变化，可以覆盖：

```bash
export QWEN35_9B_MODEL_ID=<actual-9b-id>
export QWEN36_27B_MODEL_ID=<actual-27b-id>
./scripts/download_models.sh
```

## 单步训练

每个训练脚本默认：

- `num_train_epochs=1`
- `lora_rank=32`
- `torch_dtype=bfloat16`
- `save_steps=1000`
- `eval_steps=1000`
- `report_to=tensorboard`
- `max_length=262144`
- `warmup_ratio=0.1`
- `learning_rate=1e-5`

命令：

```bash
./scripts/train_qwen35_9b_lora.sh
./scripts/train_qwen35_9b_full.sh
./scripts/train_qwen36_27b_lora.sh
./scripts/train_qwen36_27b_full.sh
```

## 一键完整实验

确认 GPU 空闲后执行：

```bash
./scripts/run_all_experiments.sh
```

执行顺序固定为：

1. Qwen3.5-9B LoRA
2. Qwen3.5-9B bf16 full SFT
3. Qwen3.6-27B LoRA
4. Qwen3.6-27B bf16 full SFT

## TensorBoard

```bash
./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`。