From 71354d00fec211d7dbe4f994732c539ffcd9c7a6 Mon Sep 17 00:00:00 2001 From: bingshuo <565183519@qq.com> Date: Mon, 18 May 2026 14:03:38 +0800 Subject: [PATCH] =?UTF-8?q?D101=20=E6=96=B0=E5=A2=9E=EF=BC=9A=E5=85=A8?= =?UTF-8?q?=E8=87=AA=E5=8A=A8=E8=92=B8=E9=A6=8F=E6=B5=81=E6=B0=B4=E7=BA=BF?= =?UTF-8?q?=E8=84=9A=E6=9C=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scripts/auto_distill_pipeline.py | 142 +++++++++++++++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 scripts/auto_distill_pipeline.py diff --git a/scripts/auto_distill_pipeline.py b/scripts/auto_distill_pipeline.py new file mode 100644 index 0000000..54f1e55 --- /dev/null +++ b/scripts/auto_distill_pipeline.py @@ -0,0 +1,142 @@ +#!/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()