143 lines
7.6 KiB
Python
143 lines
7.6 KiB
Python
#!/usr/bin/env python3
|
||
"""铸渊全自动蒸馏流水线
|
||
代码模型完成 → 全自动蒸馏两个1.5B → 深度微调 → 上传COS
|
||
|
||
使用方法：
|
||
export ZY_OSS_KEY=... ZY_OSS_SECRET=...
|
||
nohup python3 -u scripts/auto_distill_pipeline.py > distill_pipeline.log 2>&1 &
|
||
"""
|
||
import os, json, time, sys, subprocess, glob
|
||
sys.stdout.reconfigure(line_buffering=True)
|
||
|
||
COS_BUCKET = "sy-finetune-corpus-1317346199"
|
||
COS_REGION = "ap-guangzhou"
|
||
COS_ID = os.environ.get("ZY_OSS_KEY")
|
||
COS_KEY = os.environ.get("ZY_OSS_SECRET")
|
||
if not COS_ID or not COS_KEY:
|
||
print("❌ 需要环境变量：export ZY_OSS_KEY=... ZY_OSS_SECRET=...")
|
||
sys.exit(1)
|
||
|
||
WORK = "/root/autodl-tmp"
|
||
TEACHER_MOTHER = f"{WORK}/output/qwen25-7b-sft/final"
|
||
TEACHER_CODER = f"{WORK}/output/qwen25-coder-7b-sft/final"
|
||
STUDENT = f"{WORK}/cache/Qwen/Qwen2___5-1___5B"
|
||
STUDENT_CODER = f"{WORK}/cache/Qwen/Qwen2___5-Coder-1___5B"
|
||
|
||
def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}"); sys.stdout.flush()
|
||
def run(cmd): r = subprocess.run(cmd, shell=True, capture_output=True, text=True); return r.returncode == 0
|
||
|
||
def cos_check(prefix):
|
||
s = f'from qcloud_cos import CosConfig,CosS3Client; c=CosS3Client(CosConfig(Region="{COS_REGION}",SecretId="{COS_ID}",SecretKey="{COS_KEY}")); r=c.list_objects(Bucket="{COS_BUCKET}",Prefix="{prefix}",MaxKeys=1); print("OK" if "Contents" in r else "NO")'
|
||
return subprocess.run(f'python3 -c "{s}"', shell=True,capture_output=True,text=True).stdout.strip()=="OK"
|
||
|
||
def cos_dl(key, local):
|
||
s = f'from qcloud_cos import CosConfig,CosS3Client;c=CosS3Client(CosConfig(Region="{COS_REGION}",SecretId="{COS_ID}",SecretKey="{COS_KEY}"));r=c.get_object(Bucket="{COS_BUCKET}",Key="{key}");open("{local}","wb").write(r["Body"].get_raw_stream().read())'
|
||
return run(f'python3 -c "{s}"')
|
||
|
||
def cos_upload(local, prefix):
|
||
if not os.path.isdir(local): return False
|
||
s = f'import os,glob;from qcloud_cos import CosConfig,CosS3Client;c=CosS3Client(CosConfig(Region="{COS_REGION}",SecretId="{COS_ID}",SecretKey="{COS_KEY}"));count=0\nfor f in glob.glob("{local}/**/*",recursive=True):\n if os.path.isfile(f):\n rel=os.path.relpath(f,"{local}")\n c.upload_file(Bucket="{COS_BUCKET}",Key="{prefix}/"+rel,LocalFilePath=f)\n count+=1\nprint(f"OK:{count}")'
|
||
return run(f'python3 -c "{s}"')
|
||
|
||
def check_coder():
|
||
log("检查代码模型...")
|
||
if os.path.isdir(TEACHER_CODER): log(" ✅ 本地已存在"); return True
|
||
if cos_check("models/qwen25-coder-7b-sft/final/"): log(" ✅ COS上存在，需手动下载到GPU"); return False
|
||
log(" ⏳ 代码模型尚未完成"); return False
|
||
|
||
def download_students():
|
||
for name, path in [("Qwen2.5-1.5B",STUDENT),("Qwen2.5-Coder-1.5B",STUDENT_CODER)]:
|
||
if not os.path.isdir(path):
|
||
log(f"下载{name}...")
|
||
run(f"pip3 install modelscope -q && python3 -c \"from modelscope import snapshot_download; snapshot_download('{name.replace('-1.5B','-1___5B').replace('.','/')}', cache_dir='{WORK}/cache')\"")
|
||
|
||
def run_distill(script, logfile, desc):
|
||
if os.path.isdir(f"{WORK}/output/{desc}/final"): log(f" ✅ {desc}已存在，跳过"); return True
|
||
log(f"启动{desc}蒸馏...")
|
||
p = f"{WORK}/_run_{desc}.py"
|
||
with open(p,'w') as f: f.write(open(script).read())
|
||
return run(f"cd {WORK} && python3 -u {p} > {logfile} 2>&1")
|
||
|
||
def run_deepsft(student_path, data, output, logfile, desc):
|
||
if os.path.isdir(f"{output}/final"): log(f" ✅ {desc}已存在，跳过"); return True
|
||
log(f"启动{desc}深度微调...")
|
||
s = f'''#!/usr/bin/env python3
|
||
import os,json,torch,sys;os.environ["CUDA_VISIBLE_DEVICES"]="0"
|
||
from transformers import AutoModelForCausalLM,AutoTokenizer,TrainingArguments,Trainer
|
||
from datasets import Dataset;from tqdm import tqdm
|
||
M="{student_path}";D="{data}";O="{output}"
|
||
E,B,G,L,ML=3,4,8,5e-6,2048
|
||
os.makedirs(O,exist_ok=True)
|
||
raw=[json.loads(l) for l in open(D)]
|
||
tok=AutoTokenizer.from_pretrained(M,trust_remote_code=True);tok.pad_token=tok.eos_token
|
||
model=AutoModelForCausalLM.from_pretrained(M,trust_remote_code=True,torch_dtype=torch.bfloat16,attn_implementation="sdpa").cuda()
|
||
model.config.use_cache=False;model.gradient_checkpointing_enable()
|
||
p=[]
|
||
for d in tqdm(raw):
|
||
i,l=[],[]
|
||
for m in d["messages"]:
|
||
c=m["content"]
|
||
if not c.strip(): continue
|
||
t=f"<|im_start|>{m['role']}\\n{c}<|im_end|>\\n"
|
||
tk=tok.encode(t,add_special_tokens=False)
|
||
i.extend(tk);l.extend(tk if m["role"]=="assistant" else [-100]*len(tk))
|
||
if len(i)>ML:i,l=i[:ML],l[:ML]
|
||
p.append({{"input_ids":i,"labels":l,"attention_mask":[1]*len(i)}})
|
||
def collate(f):
|
||
ml=max(len(x["input_ids"]) for x in f);b={{}}
|
||
for k in["input_ids","labels","attention_mask"]:
|
||
pad=tok.pad_token_id if k!="labels" else -100
|
||
b[k]=torch.tensor([x[k]+[pad]*(ml-len(x[k])) for x in f])
|
||
return b
|
||
args=TrainingArguments(output_dir=O,num_train_epochs=E,per_device_train_batch_size=B,gradient_accumulation_steps=G,learning_rate=L,warmup_ratio=0.05,lr_scheduler_type="cosine",bf16=True,tf32=True,logging_steps=10,save_strategy="epoch",save_total_limit=3,remove_unused_columns=False,gradient_checkpointing=True,optim="adamw_torch",report_to="none")
|
||
Trainer(model=model,args=args,train_dataset=Dataset.from_list(p),data_collator=collate).train()
|
||
f=os.path.join(O,"final");model.save_pretrained(f);tok.save_pretrained(f)
|
||
print("DONE!")
|
||
'''
|
||
with open(f"{WORK}/_run_{desc}_sft.py",'w') as f: f.write(s)
|
||
return run(f"cd {WORK} && python3 -u _run_{desc}_sft.py > {logfile} 2>&1")
|
||
|
||
def main():
|
||
log("="*60); log("铸渊全自动蒸馏流水线 v1.0"); log("="*60)
|
||
|
||
if not check_coder(): log("代码模型未就绪，终止"); return
|
||
download_students()
|
||
|
||
# Phase 1: 霜砚蒸馏
|
||
s = os.path.join(os.path.dirname(os.path.abspath(__file__)),"distill_mother.py")
|
||
if os.path.exists(s): run_distill(s,"distill_mother.log","shuangyan-distill")
|
||
else: log("⚠️ distill_mother.py 不在本地，请从仓库获取")
|
||
|
||
# Phase 2: 铸渊蒸馏 (下载语料+执行)
|
||
log("下载铸渊语料...")
|
||
os.makedirs(f"{WORK}/corpus",exist_ok=True)
|
||
cos_dl("corpus/zhuyuan_full_corpus.jsonl",f"{WORK}/corpus/zhuyuan_full_corpus.jsonl")
|
||
|
||
s2 = os.path.join(os.path.dirname(os.path.abspath(__file__)),"distill_coder.py")
|
||
if os.path.exists(s2):
|
||
c = open(s2).read().replace('DATA = "/root/autodl-tmp/corpus/zhuyuan_deep_finetune.jsonl"','DATA = "/root/autodl-tmp/corpus/zhuyuan_full_corpus.jsonl"')
|
||
with open(f"{WORK}/_run_zhuyuan-distill.py",'w') as f: f.write(c)
|
||
run_distill(f"{WORK}/_run_zhuyuan-distill.py","distill_coder.log","zhuyuan-distill")
|
||
else: log("⚠️ distill_coder.py 不在本地")
|
||
|
||
# Phase 3: 铸渊深度SFT
|
||
cos_dl("corpus/zhuyuan_deep_finetune.jsonl",f"{WORK}/corpus/zhuyuan_deep_finetune.jsonl")
|
||
# 合并语料
|
||
with open(f"{WORK}/corpus/zhuyuan_combined.jsonl",'w') as out:
|
||
for src in ["zhuyuan_full_corpus.jsonl","zhuyuan_deep_finetune.jsonl"]:
|
||
p=f"{WORK}/corpus/{src}"
|
||
if os.path.exists(p):
|
||
for l in open(p): out.write(l)
|
||
run_deepsft(f"{WORK}/output/zhuyuan-15b-distill/final",f"{WORK}/corpus/zhuyuan_combined.jsonl",f"{WORK}/output/zhuyuan-15b-deep-sft","zhuyuan_deep_sft.log","zhuyuan")
|
||
|
||
# Phase 4: 上传
|
||
for local,cos in [("output/shuangyan-15b-distill/final","models/shuangyan-15b-distill/final"),("output/zhuyuan-15b-distill/final","models/zhuyuan-15b-distill/final"),("output/zhuyuan-15b-deep-sft/final","models/zhuyuan-15b-deep-sft/final")]:
|
||
p = f"{WORK}/{local}"
|
||
if os.path.isdir(p):
|
||
log(f"上传 {local}...")
|
||
cos_upload(p,cos)
|
||
|
||
log("="*60); log("🎉 全自动蒸馏流水线完成！"); log("="*60)
|
||
|
||
if __name__ == "__main__": main()
|