Add SWIFT coding agent probe experiment scripts

This commit is contained in:
Codex
2026-06-24 21:55:23 +08:00
commit 73e9752331
19 changed files with 526 additions and 0 deletions

10
.gitignore vendored Normal file
View File

@@ -0,0 +1,10 @@
.venv/
__pycache__/
*.py[cod]
.pytest_cache/
data/
models/
outputs/
runs/
logs/
*.log

3
.gitmodules vendored Normal file
View File

@@ -0,0 +1,3 @@
[submodule "third_party/modelscope-swift"]
path = third_party/modelscope-swift
url = https://github.com/modelscope/ms-swift.git

121
README.md Normal file
View File

@@ -0,0 +1,121 @@
# TI Coding Agent Training Probe
这个仓库用于复现一组 coding-agent SFT probing 实验:从 Hugging Face 下载已经构造好的 Open-SWE-Traces probe 数据集,下载 Qwen3.5-9B 和 Qwen3.6-27B然后用 ModelScope-SWIFT 依次跑四个 1 epoch 训练任务。
## 目录
- `third_party/modelscope-swift/`: ModelScope-SWIFT submodule。
- `scripts/setup_env.sh`: 一键创建 repo 内 `.venv` 并安装本项目和 SWIFT。
- `scripts/download_dataset.py`: 下载 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 代理:
```bash
http://100.72.0.101:8888
```
Python 依赖安装在仓库内 `.venv`,不会写入系统 Python。
## 下载数据集
数据集默认名是 `ti_coding_agent_training_probe_20260624`。如果环境里设置了 `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
```
输出:
- `data/raw/training_probe/`: Hugging Face snapshot。
- `data/processed/training_probe/train.jsonl`
- `data/processed/training_probe/validation.jsonl`
训练数据里 `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`

81
SKILL.md Normal file
View File

@@ -0,0 +1,81 @@
# TI Coding Agent Probe Skill
Use this skill when the user wants to run or modify the TI coding-agent SFT probe experiments based on the Hugging Face dataset `ti_coding_agent_training_probe_20260624`.
## Repository Contract
- Work in `/ssd/workspace/yi/ti_coding_agent_probe` on B300 unless the user says otherwise.
- Use the B300 proxy for all network access:
- `http_proxy=http://100.72.0.101:8888`
- `https_proxy=http://100.72.0.101:8888`
- `HF_ENDPOINT=https://hf-mirror.com`
- Do not install Python packages globally. Use the repository-local `.venv` created by `scripts/setup_env.sh`.
- Do not start GPU training before checking GPU occupancy with `nvidia-smi`.
- Do not commit or upload `data/`, `models/`, `outputs/`, `runs/`, or `logs/`.
## Key Commands
Initialize the repo and environment:
```bash
git submodule update --init --recursive
./scripts/setup_env.sh
```
Download the training probe dataset:
```bash
export HF_ENDPOINT=https://hf-mirror.com
export HF_DATASET_REPO_ID=<owner>/ti_coding_agent_training_probe_20260624
./scripts/download_dataset.py
```
Download base models:
```bash
./scripts/download_models.sh
```
Run the full ordered experiment:
```bash
./scripts/run_all_experiments.sh
```
Run only one stage:
```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
```
Open TensorBoard:
```bash
./scripts/tensorboard.sh
```
## Training Semantics
The dataset uses SWIFT-style chat messages. `system`, `user`, and `tool` messages should remain masked with `loss=false`; only assistant trajectories should contribute to loss. This keeps scaffold prompts and tool outputs as conditioning context rather than targets to memorize.
The default experiment uses:
- 1 epoch
- LoRA rank 32 for LoRA runs
- bf16 full fine-tuning for full runs
- `max_length=262144`
- checkpoint save every 1000 steps
- validation every 1000 steps
- TensorBoard logging under `runs/`
Override model IDs or paths with:
```bash
export QWEN35_9B_MODEL_ID=<hf-id>
export QWEN36_27B_MODEL_ID=<hf-id>
export QWEN35_9B_MODEL_PATH=<local-path>
export QWEN36_27B_MODEL_PATH=<local-path>
```

18
pyproject.toml Normal file
View File

@@ -0,0 +1,18 @@
[build-system]
requires = ["setuptools>=68", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "ti-coding-agent-probe"
version = "0.1.0"
description = "Reproducible SWIFT SFT probe scripts for Open-SWE-Traces derived coding-agent data."
requires-python = ">=3.10"
dependencies = [
"huggingface_hub>=0.23",
"datasets>=2.20",
"pyarrow>=15",
"tensorboard>=2.15",
]
[tool.setuptools.packages.find]
where = ["src"]

84
scripts/download_dataset.py Executable file
View File

@@ -0,0 +1,84 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import gzip
import os
import shutil
from pathlib import Path
from huggingface_hub import HfApi, snapshot_download
DEFAULT_REPO_NAME = "ti_coding_agent_training_probe_20260624"
def resolve_dataset_id(raw: str, token: str | None, endpoint: str | None) -> str:
if "/" in raw:
return raw
if not token:
raise SystemExit(
"HF_DATASET_REPO_ID must be owner/name when HF_TOKEN is not set. "
f"Got unqualified repo name: {raw}"
)
owner = HfApi(token=token, endpoint=endpoint).whoami()["name"]
return f"{owner}/{raw}"
def gunzip_if_needed(src: Path, dst: Path) -> None:
if dst.exists() and dst.stat().st_size > 0:
return
dst.parent.mkdir(parents=True, exist_ok=True)
with gzip.open(src, "rb") as fin, dst.open("wb") as fout:
shutil.copyfileobj(fin, fout, length=1024 * 1024)
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--dataset-id", default=os.environ.get("HF_DATASET_REPO_ID", DEFAULT_REPO_NAME))
parser.add_argument("--raw-dir", default="data/raw/training_probe")
parser.add_argument("--out-dir", default="data/processed/training_probe")
args = parser.parse_args()
token = os.environ.get("HF_TOKEN")
endpoint = os.environ.get("HF_ENDPOINT")
dataset_id = resolve_dataset_id(args.dataset_id, token, endpoint)
raw_dir = Path(args.raw_dir)
out_dir = Path(args.out_dir)
raw_dir.mkdir(parents=True, exist_ok=True)
out_dir.mkdir(parents=True, exist_ok=True)
local_path = snapshot_download(
repo_id=dataset_id,
repo_type="dataset",
local_dir=raw_dir,
token=token,
endpoint=endpoint,
allow_patterns=[
"README.md",
"metadata.json",
"train.parquet",
"validation.parquet",
"train.jsonl.gz",
"validation.jsonl.gz",
],
)
local = Path(local_path)
for name in ("train", "validation"):
gz = local / f"{name}.jsonl.gz"
if gz.exists():
gunzip_if_needed(gz, out_dir / f"{name}.jsonl")
parquet = local / f"{name}.parquet"
if parquet.exists():
target = out_dir / f"{name}.parquet"
if not target.exists():
target.symlink_to(parquet.resolve())
print(f"DATASET_ID={dataset_id}")
print(f"TRAIN_JSONL={out_dir / 'train.jsonl'}")
print(f"VALIDATION_JSONL={out_dir / 'validation.jsonl'}")
return 0
if __name__ == "__main__":
raise SystemExit(main())

20
scripts/download_models.sh Executable file
View File

@@ -0,0 +1,20 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT_DIR}"
export http_proxy="${http_proxy:-http://100.72.0.101:8888}"
export https_proxy="${https_proxy:-http://100.72.0.101:8888}"
export HTTP_PROXY="${HTTP_PROXY:-${http_proxy}}"
export HTTPS_PROXY="${HTTPS_PROXY:-${https_proxy}}"
export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
source .venv/bin/activate
mkdir -p models
QWEN35_9B_MODEL_ID="${QWEN35_9B_MODEL_ID:-Qwen/Qwen3.5-9B}"
QWEN36_27B_MODEL_ID="${QWEN36_27B_MODEL_ID:-Qwen/Qwen3.6-27B}"
huggingface-cli download "${QWEN35_9B_MODEL_ID}" --local-dir "models/qwen3.5-9b" --local-dir-use-symlinks False
huggingface-cli download "${QWEN36_27B_MODEL_ID}" --local-dir "models/qwen3.6-27b" --local-dir-use-symlinks False

16
scripts/run_all_experiments.sh Executable file
View File

@@ -0,0 +1,16 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT_DIR}"
./scripts/download_dataset.py
./scripts/download_models.sh
./scripts/train_qwen35_9b_lora.sh
./scripts/train_qwen35_9b_full.sh
./scripts/train_qwen36_27b_lora.sh
./scripts/train_qwen36_27b_full.sh
echo "All experiments finished."
echo "TensorBoard: tensorboard --logdir runs --host 0.0.0.0 --port 6006"

24
scripts/setup_env.sh Executable file
View File

@@ -0,0 +1,24 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT_DIR}"
export http_proxy="${http_proxy:-http://100.72.0.101:8888}"
export https_proxy="${https_proxy:-http://100.72.0.101:8888}"
export HTTP_PROXY="${HTTP_PROXY:-${http_proxy}}"
export HTTPS_PROXY="${HTTPS_PROXY:-${https_proxy}}"
export PIP_INDEX_URL="${PIP_INDEX_URL:-https://mirrors.aliyun.com/pypi/simple/}"
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -U pip setuptools wheel
python -m pip install -e .
python -m pip install -e third_party/modelscope-swift
mkdir -p data/raw data/processed models outputs runs logs
python - <<'PY'
import swift, sys
print("python", sys.version)
print("swift", getattr(swift, "__version__", "unknown"))
PY

88
scripts/swift_train_common.sh Executable file
View File

@@ -0,0 +1,88 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT_DIR}"
export http_proxy="${http_proxy:-http://100.72.0.101:8888}"
export https_proxy="${https_proxy:-http://100.72.0.101:8888}"
export HTTP_PROXY="${HTTP_PROXY:-${http_proxy}}"
export HTTPS_PROXY="${HTTPS_PROXY:-${https_proxy}}"
export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}"
if [[ -f .venv/bin/activate ]]; then
source .venv/bin/activate
elif [[ "${DRY_RUN:-0}" != "1" ]]; then
echo "Missing .venv. Run ./scripts/setup_env.sh first." >&2
exit 2
fi
mkdir -p outputs runs logs
TRAIN_JSONL="${TRAIN_JSONL:-data/processed/training_probe/train.jsonl}"
VAL_JSONL="${VAL_JSONL:-data/processed/training_probe/validation.jsonl}"
MAX_LENGTH="${MAX_LENGTH:-262144}"
SAVE_STEPS="${SAVE_STEPS:-1000}"
EVAL_STEPS="${EVAL_STEPS:-1000}"
LOGGING_STEPS="${LOGGING_STEPS:-1}"
GRAD_ACCUM_STEPS="${GRAD_ACCUM_STEPS:-1}"
PER_DEVICE_BATCH_SIZE="${PER_DEVICE_BATCH_SIZE:-1}"
NUM_EPOCHS="${NUM_EPOCHS:-1}"
LEARNING_RATE="${LEARNING_RATE:-1e-5}"
WARMUP_RATIO="${WARMUP_RATIO:-0.1}"
LORA_RANK="${LORA_RANK:-32}"
require_file() {
if [[ ! -f "$1" ]]; then
echo "Missing required file: $1" >&2
exit 2
fi
}
run_swift_train() {
local model_path="$1"
local train_type="$2"
local run_name="$3"
local output_dir="outputs/${run_name}"
local tb_dir="runs/${run_name}"
local log_file="logs/${run_name}.log"
require_file "${TRAIN_JSONL}"
require_file "${VAL_JSONL}"
mkdir -p "${output_dir}" "${tb_dir}" logs
local cmd=(
swift sft
--model "${model_path}"
--dataset "${TRAIN_JSONL}"
--val_dataset "${VAL_JSONL}"
--train_type "${train_type}"
--torch_dtype bfloat16
--num_train_epochs "${NUM_EPOCHS}"
--per_device_train_batch_size "${PER_DEVICE_BATCH_SIZE}"
--per_device_eval_batch_size 1
--gradient_accumulation_steps "${GRAD_ACCUM_STEPS}"
--learning_rate "${LEARNING_RATE}"
--warmup_ratio "${WARMUP_RATIO}"
--max_length "${MAX_LENGTH}"
--save_steps "${SAVE_STEPS}"
--eval_steps "${EVAL_STEPS}"
--logging_steps "${LOGGING_STEPS}"
--report_to tensorboard
--logging_dir "${tb_dir}"
--output_dir "${output_dir}"
--save_total_limit "${SAVE_TOTAL_LIMIT:-3}"
--dataloader_num_workers "${DATALOADER_NUM_WORKERS:-4}"
)
if [[ "${train_type}" == "lora" ]]; then
cmd+=(--lora_rank "${LORA_RANK}")
fi
printf '%q ' "${cmd[@]}" | tee "${log_file}.cmd"
echo
if [[ "${DRY_RUN:-0}" == "1" ]]; then
return 0
fi
"${cmd[@]}" 2>&1 | tee "${log_file}"
}

6
scripts/tensorboard.sh Executable file
View File

@@ -0,0 +1,6 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT_DIR}"
source .venv/bin/activate
tensorboard --logdir runs --host 0.0.0.0 --port "${TENSORBOARD_PORT:-6006}"

View File

@@ -0,0 +1,4 @@
#!/usr/bin/env bash
set -euo pipefail
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/swift_train_common.sh"
run_swift_train "${QWEN35_9B_MODEL_PATH:-models/qwen3.5-9b}" full qwen35_9b_full_bf16

View File

@@ -0,0 +1,4 @@
#!/usr/bin/env bash
set -euo pipefail
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/swift_train_common.sh"
run_swift_train "${QWEN35_9B_MODEL_PATH:-models/qwen3.5-9b}" lora qwen35_9b_lora_r32

View File

@@ -0,0 +1,4 @@
#!/usr/bin/env bash
set -euo pipefail
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/swift_train_common.sh"
run_swift_train "${QWEN36_27B_MODEL_PATH:-models/qwen3.6-27b}" full qwen36_27b_full_bf16

View File

@@ -0,0 +1,4 @@
#!/usr/bin/env bash
set -euo pipefail
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/swift_train_common.sh"
run_swift_train "${QWEN36_27B_MODEL_PATH:-models/qwen3.6-27b}" lora qwen36_27b_lora_r32

18
scripts/validate_setup.sh Executable file
View File

@@ -0,0 +1,18 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT_DIR}"
test -d third_party/modelscope-swift
python3 -m py_compile scripts/download_dataset.py
bash -n scripts/*.sh
DRY_RUN=1 TRAIN_JSONL=/tmp/train.jsonl VAL_JSONL=/tmp/validation.jsonl bash -c '
echo "{}" > /tmp/train.jsonl
echo "{}" > /tmp/validation.jsonl
./scripts/train_qwen35_9b_lora.sh
./scripts/train_qwen35_9b_full.sh
./scripts/train_qwen36_27b_lora.sh
./scripts/train_qwen36_27b_full.sh
'

View File

@@ -0,0 +1,4 @@
"""Utilities for the TI coding-agent SFT probe."""
__all__ = ["__version__"]
__version__ = "0.1.0"

View File

@@ -0,0 +1,16 @@
from __future__ import annotations
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[2]
DATA_DIR = REPO_ROOT / "data"
MODEL_DIR = REPO_ROOT / "models"
OUTPUT_DIR = REPO_ROOT / "outputs"
RUNS_DIR = REPO_ROOT / "runs"
LOG_DIR = REPO_ROOT / "logs"
def ensure_project_dirs() -> None:
for path in (DATA_DIR, MODEL_DIR, OUTPUT_DIR, RUNS_DIR, LOG_DIR):
path.mkdir(parents=True, exist_ok=True)

1
third_party/modelscope-swift vendored Submodule