hupenglong1
309 lines
11 KiB
Python
309 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
从 ai-planning 组装 SFT 训练/验证集，并启动 zk_trainer SFT 训练。
|
||
|
||
用法:
|
||
# 仅组装数据
|
||
python prepare_and_train_sft.py prepare --output_dir /mnt/wangsenhao/autoresearch-zk/sft_data
|
||
|
||
# 启动训练（需先 prepare）
|
||
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__)
|
||
|
||
AI_PLANNING = Path(__file__).resolve().parent / "ai-planning"
|
||
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(Path(__file__).resolve().parent / "zk_trainer")
|
||
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()
|