Add model iteration skill
This commit is contained in:
武阳
@@ -0,0 +1,317 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
从 ai-planning 组装 SFT 训练/验证集，并启动 zk_trainer SFT 训练。
|
||||
|
||||
用法:
|
||||
# 仅组装数据
|
||||
AUTORESEARCH_ROOT=/mnt/wangsenhao/autoresearch-zk \
|
||||
python prepare_and_train_sft.py prepare --output_dir /mnt/wangsenhao/autoresearch-zk/sft_data
|
||||
|
||||
# 启动训练（需先 prepare）
|
||||
AUTORESEARCH_ROOT=/mnt/wangsenhao/autoresearch-zk \
|
||||
python prepare_and_train_sft.py train \
|
||||
--data_dir /mnt/wangsenhao/autoresearch-zk/sft_data \
|
||||
--model_path /mnt/wangsenhao/verl_zk/Qwen3-4B-Instruct-2507 \
|
||||
--model_type qwen3 \
|
||||
--train_output /mnt/wangsenhao/autoresearch-zk/sft_output
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _path_from_env(name: str, fallback: Path) -> Path:
|
||||
return Path(os.environ.get(name, str(fallback))).expanduser().resolve()
|
||||
|
||||
|
||||
AUTORESEARCH_ROOT = _path_from_env("AUTORESEARCH_ROOT", Path("/mnt/wangsenhao/autoresearch-zk"))
|
||||
AI_PLANNING = _path_from_env("AI_PLANNING_DIR", AUTORESEARCH_ROOT / "ai-planning")
|
||||
ZK_TRAINER_DIR = _path_from_env("ZK_TRAINER_DIR", AUTORESEARCH_ROOT / "zk_trainer")
|
||||
TRAIN_SET_DIR = AI_PLANNING / "data" / "train_set"
|
||||
DATA_TRAIN_DIR = AI_PLANNING / "data_train"
|
||||
GENERATE_TRAIN_SCRIPT = DATA_TRAIN_DIR / "generate_train.py"
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────
|
||||
# 1. 数据组装
|
||||
# ──────────────────────────────────────────────
|
||||
|
||||
def regen_jsonl_from_csv():
|
||||
"""调用 data_train/generate_train.py 从 CSV 重新生成 all_train.jsonl。"""
|
||||
if not GENERATE_TRAIN_SCRIPT.exists():
|
||||
log.error(f"generate_train.py 不存在: {GENERATE_TRAIN_SCRIPT}")
|
||||
return False
|
||||
log.info("从 CSV 重新生成 JSONL ...")
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(GENERATE_TRAIN_SCRIPT)],
|
||||
cwd=str(DATA_TRAIN_DIR),
|
||||
)
|
||||
if result.returncode != 0:
|
||||
log.error("generate_train.py 执行失败")
|
||||
return False
|
||||
log.info("CSV → JSONL 完成")
|
||||
return True
|
||||
|
||||
|
||||
def scan_train_set():
|
||||
"""扫描 train_set/ 下所有子目录，返回 (train_files, valid_files)。"""
|
||||
train_files, valid_files = [], []
|
||||
if not TRAIN_SET_DIR.exists():
|
||||
log.error(f"train_set 目录不存在: {TRAIN_SET_DIR}")
|
||||
return train_files, valid_files
|
||||
|
||||
for subdir in sorted(TRAIN_SET_DIR.iterdir()):
|
||||
if not subdir.is_dir() or subdir.name.startswith("."):
|
||||
continue
|
||||
for f in sorted(subdir.iterdir()):
|
||||
if not f.is_file() or f.suffix != ".jsonl":
|
||||
continue
|
||||
if "_valid" in f.name:
|
||||
valid_files.append(f)
|
||||
else:
|
||||
train_files.append(f)
|
||||
return train_files, valid_files
|
||||
|
||||
|
||||
def count_lines(path):
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
return sum(1 for _ in f)
|
||||
|
||||
|
||||
def prepare_data(output_dir: str, regen: bool = False, dataset_name: str = "zk_sft_new_structure"):
|
||||
"""
|
||||
组装训练集和验证集到 output_dir，生成 zk_trainer 所需的目录结构:
|
||||
output_dir/
|
||||
train/<dataset_name>/merged_train.jsonl
|
||||
validation/<dataset_name>/part-0.jsonl
|
||||
"""
|
||||
output = Path(output_dir).resolve()
|
||||
train_dir = output / "train" / dataset_name
|
||||
valid_dir = output / "validation" / dataset_name
|
||||
|
||||
if regen:
|
||||
regen_jsonl_from_csv()
|
||||
|
||||
train_files, valid_files = scan_train_set()
|
||||
if not train_files:
|
||||
log.error("未找到训练文件")
|
||||
return None
|
||||
|
||||
log.info(f"训练文件 ({len(train_files)}):")
|
||||
for f in train_files:
|
||||
log.info(f" {f.relative_to(AI_PLANNING)} ({count_lines(f)} 条)")
|
||||
log.info(f"验证文件 ({len(valid_files)}):")
|
||||
for f in valid_files:
|
||||
log.info(f" {f.relative_to(AI_PLANNING)} ({count_lines(f)} 条)")
|
||||
|
||||
# 清理旧数据
|
||||
for d in [train_dir, valid_dir]:
|
||||
if d.exists():
|
||||
shutil.rmtree(d)
|
||||
d.mkdir(parents=True)
|
||||
|
||||
# 合并训练文件
|
||||
merged_path = train_dir / "merged_train.jsonl"
|
||||
total_train = 0
|
||||
with open(merged_path, "w", encoding="utf-8") as out:
|
||||
for fp in train_files:
|
||||
with open(fp, "r", encoding="utf-8") as inp:
|
||||
for line in inp:
|
||||
out.write(line)
|
||||
total_train += 1
|
||||
log.info(f"合并训练集: {total_train} 条 → {merged_path}")
|
||||
|
||||
# 合并验证文件
|
||||
merged_valid_path = valid_dir / "part-0.jsonl"
|
||||
total_valid = 0
|
||||
with open(merged_valid_path, "w", encoding="utf-8") as out:
|
||||
for fp in valid_files:
|
||||
with open(fp, "r", encoding="utf-8") as inp:
|
||||
for line in inp:
|
||||
out.write(line)
|
||||
total_valid += 1
|
||||
log.info(f"验证集: {total_valid} 条 → {merged_valid_path}")
|
||||
|
||||
log.info("=" * 50)
|
||||
log.info(f"数据目录: {output}")
|
||||
log.info(f" train: {total_train} 条")
|
||||
log.info(f" validation: {total_valid} 条")
|
||||
|
||||
return str(output)
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────
|
||||
# 2. SFT 训练
|
||||
# ──────────────────────────────────────────────
|
||||
|
||||
def run_sft_training(data_dir: str, model_path: str, model_type: str, train_output: str,
|
||||
epochs: int = 3, lr: float = 1e-5, max_seq_length: int = 1024,
|
||||
per_device_batch: int = 1, grad_accum: int = 1):
|
||||
"""基于 zk_trainer 的 accelerate launch 方式启动 SFT 训练。"""
|
||||
zk_trainer_dir = str(ZK_TRAINER_DIR)
|
||||
if not os.path.isdir(zk_trainer_dir):
|
||||
log.error(f"zk_trainer 目录不存在: {zk_trainer_dir}")
|
||||
return None
|
||||
|
||||
data_dir = str(Path(data_dir).resolve())
|
||||
train_output = str(Path(train_output).resolve())
|
||||
os.makedirs(train_output, exist_ok=True)
|
||||
|
||||
run_name = f"{model_type}-zk"
|
||||
|
||||
env = os.environ.copy()
|
||||
env["PYTHONPATH"] = f"{zk_trainer_dir}:{env.get('PYTHONPATH', '')}"
|
||||
env["WANDB_DISABLED"] = "true"
|
||||
env["WANDB_MODE"] = "offline"
|
||||
env.setdefault("VOLUME_PREFIX", "/mnt/xiaoai-zk-model-train-tj5")
|
||||
|
||||
# Step 1: 用 get_accelerate_config.py 动态生成 accelerate config
|
||||
num_gpu_result = subprocess.run(
|
||||
["nvidia-smi", "--query-gpu=gpu_name", "--format=csv,noheader"],
|
||||
capture_output=True, text=True
|
||||
)
|
||||
num_processes = len(num_gpu_result.stdout.strip().split("\n")) if num_gpu_result.returncode == 0 else 1
|
||||
log.info(f"检测到 {num_processes} 张 GPU")
|
||||
|
||||
acc_config_dir = zk_trainer_dir
|
||||
gen_config_cmd = [
|
||||
sys.executable, f"{zk_trainer_dir}/llm/train/get_accelerate_config.py",
|
||||
"--dt_framework=fsdp",
|
||||
"--fsdp_sharding_strategy=1",
|
||||
"--machine_rank=0",
|
||||
"--num_machines=1",
|
||||
f"--num_processes={num_processes}",
|
||||
f"--output_dir={acc_config_dir}",
|
||||
"--main_process_ip=127.0.0.1",
|
||||
"--main_process_port=56390",
|
||||
"--mixed_precision=bf16",
|
||||
f"--model_type={model_type}",
|
||||
"--ds_stage=3",
|
||||
"--fsdp_version=2",
|
||||
"--fsdp_wrap_cls=",
|
||||
]
|
||||
log.info("生成 accelerate config ...")
|
||||
result = subprocess.run(gen_config_cmd, env=env, cwd=zk_trainer_dir)
|
||||
if result.returncode != 0:
|
||||
log.error("生成 accelerate config 失败")
|
||||
return None
|
||||
|
||||
acc_config = os.path.join(acc_config_dir, "accelerate_config_0.yaml")
|
||||
if not os.path.exists(acc_config):
|
||||
log.error(f"accelerate config 未生成: {acc_config}")
|
||||
return None
|
||||
log.info(f"accelerate config: {acc_config}")
|
||||
|
||||
# Step 2: 启动训练
|
||||
train_cmd = [
|
||||
"accelerate", "launch",
|
||||
"--main_process_port", "12345",
|
||||
"--config_file", acc_config,
|
||||
f"{zk_trainer_dir}/llm/train/trainer_general.py",
|
||||
"--model_name_or_path", model_path,
|
||||
"--model_type", model_type,
|
||||
"--use_auto_class", "true",
|
||||
"--output_dir", train_output,
|
||||
"--dataset_type", "zk_sft",
|
||||
"--train_data_dir", os.path.join(data_dir, "train"),
|
||||
"--valid_data_dir", os.path.join(data_dir, "validation"),
|
||||
"--do_eval",
|
||||
"--fp16_full_eval",
|
||||
"--per_device_eval_batch_size", "1",
|
||||
"--include_for_metrics", "inputs",
|
||||
"--batch_eval_metrics",
|
||||
"--concat_samples", "False",
|
||||
"--mix_at_eval", "False",
|
||||
"--num_train_epochs", str(epochs),
|
||||
"--learning_rate", str(lr),
|
||||
"--gradient_accumulation_steps", str(grad_accum),
|
||||
"--ignore_data_skip", "False",
|
||||
"--resume_train_if_ckpt_exists", "false",
|
||||
"--save_strategy", "no",
|
||||
"--logging_steps", "5",
|
||||
"--save_total_limit", "0",
|
||||
"--max_seq_length", str(max_seq_length),
|
||||
"--adam_beta1", "0.9",
|
||||
"--adam_beta2", "0.95",
|
||||
"--adam_epsilon", "1e-9",
|
||||
"--warmup_ratio", "0.1",
|
||||
"--lr_scheduler_type", "cosine",
|
||||
"--use_shuffle", "true",
|
||||
"--batch_eval_metrics",
|
||||
"--run_name", run_name,
|
||||
"--report_to", "none",
|
||||
"--bf16",
|
||||
]
|
||||
|
||||
log.info("=" * 50)
|
||||
log.info("启动 SFT 训练")
|
||||
log.info(f" 模型: {model_path}")
|
||||
log.info(f" 数据: {data_dir}")
|
||||
log.info(f" 输出: {train_output}")
|
||||
log.info(f" GPU: {num_processes}")
|
||||
log.info(f" epochs: {epochs}, lr: {lr}, max_seq: {max_seq_length}")
|
||||
|
||||
result = subprocess.run(train_cmd, env=env, cwd=zk_trainer_dir)
|
||||
|
||||
if result.returncode != 0:
|
||||
log.error(f"训练失败，返回码: {result.returncode}")
|
||||
return None
|
||||
|
||||
log.info(f"训练完成，输出: {train_output}")
|
||||
return train_output
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────
|
||||
# CLI
|
||||
# ──────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="ai-planning SFT 数据组装 & 训练")
|
||||
subparsers = parser.add_subparsers(dest="command")
|
||||
|
||||
# prepare
|
||||
p_prepare = subparsers.add_parser("prepare", help="仅组装数据")
|
||||
p_prepare.add_argument("--output_dir", required=True, help="数据输出目录")
|
||||
p_prepare.add_argument("--regen_from_csv", action="store_true", help="从 CSV 重新生成 JSONL")
|
||||
p_prepare.add_argument("--dataset_name", default="zk_sft_new_structure")
|
||||
|
||||
# train
|
||||
p_train = subparsers.add_parser("train", help="启动训练（需先 prepare）")
|
||||
p_train.add_argument("--data_dir", required=True, help="已组装的数据目录（prepare 的 output_dir）")
|
||||
p_train.add_argument("--model_path", required=True, help="基模路径")
|
||||
p_train.add_argument("--model_type", default="qwen3", help="模型类型")
|
||||
p_train.add_argument("--train_output", required=True, help="训练输出目录")
|
||||
p_train.add_argument("--epochs", type=int, default=3)
|
||||
p_train.add_argument("--lr", type=float, default=1e-5)
|
||||
p_train.add_argument("--max_seq_length", type=int, default=1024)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "prepare":
|
||||
prepare_data(args.output_dir, regen=args.regen_from_csv, dataset_name=args.dataset_name)
|
||||
|
||||
elif args.command == "train":
|
||||
run_sft_training(
|
||||
data_dir=args.data_dir,
|
||||
model_path=args.model_path,
|
||||
model_type=args.model_type,
|
||||
train_output=args.train_output,
|
||||
epochs=args.epochs,
|
||||
lr=args.lr,
|
||||
max_seq_length=args.max_seq_length,
|
||||
)
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user