commit 73e9752331889991bcb166e08df274a21584f66b Author: Codex Date: Wed Jun 24 21:55:23 2026 +0800 Add SWIFT coding agent probe experiment scripts diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..cf99021 --- /dev/null +++ b/.gitignore @@ -0,0 +1,10 @@ +.venv/ +__pycache__/ +*.py[cod] +.pytest_cache/ +data/ +models/ +outputs/ +runs/ +logs/ +*.log diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..1f35bf0 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "third_party/modelscope-swift"] + path = third_party/modelscope-swift + url = https://github.com/modelscope/ms-swift.git diff --git a/README.md b/README.md new file mode 100644 index 0000000..a716518 --- /dev/null +++ b/README.md @@ -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=/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= +export QWEN36_27B_MODEL_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//`。SWIFT/Transformers 的 TensorBoard 标量通常包括 loss、learning rate、eval loss、runtime、samples/sec、steps/sec 等 throughput 指标；同时 stdout 会保存在 `logs/.log`。 diff --git a/SKILL.md b/SKILL.md new file mode 100644 index 0000000..0e462d3 --- /dev/null +++ b/SKILL.md @@ -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=/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= +export QWEN36_27B_MODEL_ID= +export QWEN35_9B_MODEL_PATH= +export QWEN36_27B_MODEL_PATH= +``` diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..8319c0d --- /dev/null +++ b/pyproject.toml @@ -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"] diff --git a/scripts/download_dataset.py b/scripts/download_dataset.py new file mode 100755 index 0000000..ee0863a --- /dev/null +++ b/scripts/download_dataset.py @@ -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()) diff --git a/scripts/download_models.sh b/scripts/download_models.sh new file mode 100755 index 0000000..8c78213 --- /dev/null +++ b/scripts/download_models.sh @@ -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 diff --git a/scripts/run_all_experiments.sh b/scripts/run_all_experiments.sh new file mode 100755 index 0000000..a8121d9 --- /dev/null +++ b/scripts/run_all_experiments.sh @@ -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" diff --git a/scripts/setup_env.sh b/scripts/setup_env.sh new file mode 100755 index 0000000..f8aed56 --- /dev/null +++ b/scripts/setup_env.sh @@ -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 diff --git a/scripts/swift_train_common.sh b/scripts/swift_train_common.sh new file mode 100755 index 0000000..0fc1f57 --- /dev/null +++ b/scripts/swift_train_common.sh @@ -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}" +} diff --git a/scripts/tensorboard.sh b/scripts/tensorboard.sh new file mode 100755 index 0000000..d06f826 --- /dev/null +++ b/scripts/tensorboard.sh @@ -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}" diff --git a/scripts/train_qwen35_9b_full.sh b/scripts/train_qwen35_9b_full.sh new file mode 100755 index 0000000..e849ead --- /dev/null +++ b/scripts/train_qwen35_9b_full.sh @@ -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 diff --git a/scripts/train_qwen35_9b_lora.sh b/scripts/train_qwen35_9b_lora.sh new file mode 100755 index 0000000..6af9186 --- /dev/null +++ b/scripts/train_qwen35_9b_lora.sh @@ -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 diff --git a/scripts/train_qwen36_27b_full.sh b/scripts/train_qwen36_27b_full.sh new file mode 100755 index 0000000..042ee2c --- /dev/null +++ b/scripts/train_qwen36_27b_full.sh @@ -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 diff --git a/scripts/train_qwen36_27b_lora.sh b/scripts/train_qwen36_27b_lora.sh new file mode 100755 index 0000000..d7b6a02 --- /dev/null +++ b/scripts/train_qwen36_27b_lora.sh @@ -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 diff --git a/scripts/validate_setup.sh b/scripts/validate_setup.sh new file mode 100755 index 0000000..89480e6 --- /dev/null +++ b/scripts/validate_setup.sh @@ -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 +' diff --git a/src/ti_coding_agent_probe/__init__.py b/src/ti_coding_agent_probe/__init__.py new file mode 100644 index 0000000..d8d0b79 --- /dev/null +++ b/src/ti_coding_agent_probe/__init__.py @@ -0,0 +1,4 @@ +"""Utilities for the TI coding-agent SFT probe.""" + +__all__ = ["__version__"] +__version__ = "0.1.0" diff --git a/src/ti_coding_agent_probe/paths.py b/src/ti_coding_agent_probe/paths.py new file mode 100644 index 0000000..46cd56c --- /dev/null +++ b/src/ti_coding_agent_probe/paths.py @@ -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) diff --git a/third_party/modelscope-swift b/third_party/modelscope-swift new file mode 160000 index 0000000..4b77fcb --- /dev/null +++ b/third_party/modelscope-swift @@ -0,0 +1 @@ +Subproject commit 4b77fcbdbc711abd0fe9eabd1a34ab430b764a1e