210 lines
10 KiB
Python
210 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""铸渊全自动蒸馏流水线 v2.0
|
||
|
||
D100教训修复：
|
||
教训1：写了脚本没启动 = 白写 → 启动时检测所有状态，已有完成的跳过
|
||
教训2：只有一条检测路径 → 三重检测：本地文件 + COS路径 + 进度标记
|
||
教训3：轮询间隔300秒太久 → 改为60秒
|
||
|
||
流程：
|
||
Phase 0: 等待代码模型完成（轮询COS和本地）
|
||
Phase 1: 母模型(7B) → 蒸馏 → 霜砚1.5B
|
||
Phase 2: 霜砚语料 → 深度SFT霜砚1.5B
|
||
Phase 3: 代码模型(7B) → 蒸馏 → 铸渊1.5B
|
||
Phase 4: 铸渊语料 → 深度SFT铸渊1.5B
|
||
Phase 5: 全部上传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,zipfile,io,re
|
||
sys.stdout.reconfigure(line_buffering=True)
|
||
|
||
COS_BUCKET = "sy-finetune-corpus-1317346199"; COS_REGION = "ap-guangzhou"
|
||
CID = os.environ.get("ZY_OSS_KEY"); CKEY = os.environ.get("ZY_OSS_SECRET")
|
||
if not CID or not CKEY: print("❌ export ZY_OSS_KEY=... ZY_OSS_SECRET=..."); sys.exit(1)
|
||
|
||
W = "/root/autodl-tmp"; CACHE = f"{W}/cache"; OUT = f"{W}/output"
|
||
TM = f"{OUT}/qwen25-7b-sft/final"; TC = f"{OUT}/qwen25-coder-7b-sft/final"
|
||
S1 = f"{CACHE}/Qwen/Qwen2___5-1___5B"; S2 = f"{CACHE}/Qwen/Qwen2___5-Coder-1___5B"
|
||
SF = f"{W}/.distill_status.json" # 进度标记
|
||
DSD = f"{OUT}/shuangyan-15b-distill/final"; DZD = f"{OUT}/zhuyuan-15b-distill/final"
|
||
DSS = f"{OUT}/shuangyan-15b-deep-sft/final"; DZS = f"{OUT}/zhuyuan-15b-deep-sft/final"
|
||
POLL = 60
|
||
|
||
def log(m): t=time.strftime("%H:%M:%S"); print(f"[{t}] {m}"); sys.stdout.flush()
|
||
def run(c): r=subprocess.run(c,shell=True,capture_output=True,text=True); return r.returncode==0
|
||
def cos_e(code):
|
||
s=f"from qcloud_cos import CosConfig,CosS3Client;cc=CosS3Client(CosConfig(Region='{COS_REGION}',SecretId='{CID}',SecretKey='{CKEY}'));{code}"
|
||
r=subprocess.run(f'python3 -c "{s}"',shell=True,capture_output=True,text=True)
|
||
return r.stdout.strip(),r.returncode
|
||
|
||
def ls(): return json.load(open(SF)) if os.path.exists(SF) else {"phase":0}
|
||
def ss(p,n=""): s=ls();s["phase"]=p;s["t"]=time.time();s["n"]=n;json.dump(s,open(SF,'w'))
|
||
def dl(p): return os.path.isdir(p) and any(f.endswith('.safetensors') for f in os.listdir(p))
|
||
def dl2(l,k="DONE!"): return os.path.exists(l) and k in subprocess.run(f"grep -a '{k}' {l}|tail -1",shell=True,capture_output=True,text=True).stdout
|
||
def dc(p): out,_=cos_e(f"r=cc.list_objects(Bucket='{COS_BUCKET}',Prefix='{p}',MaxKeys=1);print('OK' if 'Contents' in r else 'NO')");return"OK"in out
|
||
def wait(d,f,t=86400):
|
||
log(f"⏳ 等待{d}...");st=time.time()
|
||
while time.time()-st<t:
|
||
if f():log(f" ✅ {d}就绪");return True
|
||
time.sleep(POLL)
|
||
log(f" ❌ {d}超时");return False
|
||
def cdl(k,l):
|
||
os.makedirs(os.path.dirname(l),exist_ok=True)
|
||
return cos_e(f"r=cc.get_object(Bucket='{COS_BUCKET}',Key='{k}');open('{l}','wb').write(r['Body'].get_raw_stream().read());print('OK')")
|
||
def cup(l,p):
|
||
if not os.path.isdir(l):return False
|
||
s=f"import os,glob;cnt=0;c=CosS3Client(CosConfig(Region='{COS_REGION}',SecretId='{CID}',SecretKey='{CKEY}'))"
|
||
s+=f";for f in glob.glob('{l}/**/*',recursive=True):"
|
||
s+=" if os.path.isfile(f):rel=os.path.relpath(f,'"+l+"');c.upload_file(Bucket='"+COS_BUCKET+"',Key='"+p+"/'+rel,LocalFilePath=f);cnt+=1"
|
||
s+=";print(f'OK:{cnt}')"
|
||
out,_=cos_e(s);log(f" {out}");return True
|
||
|
||
def phase0():
|
||
log("="*60);log("Phase 0: 等待代码模型");log("="*60)
|
||
log("三重检测：本地 + COS + 日志DONE")
|
||
if dl(TC):log(f" ✅ 本地已存");return True
|
||
if dc("models/qwen25-coder-7b-sft/final/model.safetensors"):
|
||
log(" ✅ COS上有，需手动下载（14GB）");log(f" coscli cp cos://{COS_BUCKET}/models/qwen25-coder-7b-sft/final/ {TC}/ -r");return False
|
||
return wait("代码模型完成",lambda:dc("models/qwen25-coder-7b-sft/final/model.safetensors"))
|
||
|
||
def students():
|
||
for n,p,m in[("Qwen2.5-1.5B",S1,"Qwen/Qwen2.5-1.5B"),("Qwen2.5-Coder-1.5B",S2,"Qwen/Qwen2.5-Coder-1.5B")]:
|
||
if os.path.isdir(p):log(f" ✅ {n}")
|
||
else:log(f" 下载{n}...");run("pip3 install modelscope -q 2>/dev/null");run(f"python3 -c \"from modelscope import snapshot_download; snapshot_download('{m}',cache_dir='{CACHE}')\"")
|
||
|
||
def teachers():
|
||
for n,p in[("母模型",TM),("代码模型",TC)]:
|
||
if dl(p):log(f" ✅ {n} 就绪")
|
||
else:log(f" ❌ {n} 不在本地");return False
|
||
return True
|
||
|
||
def run_distill(src,nm,lg,tch):
|
||
od=f"{OUT}/{nm}/final"
|
||
if os.path.isdir(od):log(f" ✅ {nm}已存在");return True
|
||
if not os.path.isdir(tch):log(f" ❌ Teacher:{tch}不存在");return False
|
||
sp=f"{W}/_run_{nm}.py"
|
||
with open(sp,'w')as f:f.write(open(src).read())
|
||
log(f" 启动{nm}...");ss(2,f"{nm}训练中")
|
||
ok=run(f"cd {W} && python3 -u {sp} > {lg} 2>&1")
|
||
if dl2(f"{W}/{lg}")or os.path.isdir(od):log(f" ✅ {nm}完成");ss(3);return True
|
||
log(f" ⚠️ {nm}状态不确定，检查{lg}");return ok
|
||
|
||
def sft(mp,data,od,lg,nm):
|
||
if os.path.isdir(f"{od}/final"):log(f" ✅ {nm}深度SFT完成");return True
|
||
if not os.path.isdir(mp):log(f" ❌ 模型:{mp}不存在");return False
|
||
os.makedirs(f"{W}/corpus",exist_ok=True)
|
||
cbf=f"{W}/corpus/{nm}.jsonl"
|
||
with open(cbf,'w')as o:
|
||
for s in data:
|
||
if os.path.exists(s):
|
||
for l in open(s):o.write(l)
|
||
n=sum(1 for _ in open(cbf))
|
||
log(f" 语料:{n}条")
|
||
if n==0:log(" ⚠️ 空语料");return False
|
||
ss(4,f"{nm}SFT训练中")
|
||
sc=f"{W}/_run_{nm}_sft.py"
|
||
with open(sc,'w')as f:
|
||
f.write(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="{mp}";D="{cbf}";O="{od}";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"];t=f"<|im_start|>{{m['role']}}\
|
||
{{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)}})
|
||
print(f"Dataset:{{len(p)}}ex")
|
||
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,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!")
|
||
''')
|
||
ok=run(f"cd {W} && python3 -u {sc} > {lg} 2>&1")
|
||
if os.path.isdir(f"{od}/final")or dl2(f"{W}/{lg}"):log(f" ✅ {nm}完成");ss(5);return True
|
||
return ok
|
||
|
||
def sy_corpus():
|
||
log("下载霜砚语料...");d=f"{W}/data/shuangyan";os.makedirs(d,exist_ok=True)
|
||
fl=["霜砚对话.zip","霜砚HLDP核心大脑.zip","霜砚语料包V2.0.zip","HLDP 母语协议 v2.0 · 光之树记忆编码+思维编码规范 · 霜砚签发.zip","光湖驱动引擎架构 · 推理思维链 · 2026-05-17.zip"]
|
||
for fn in fl:
|
||
ck=f"corpus/shuangyan-1.5b-sft/{fn}";lo=f"{d}/{fn}"
|
||
if not os.path.exists(lo):cdl(ck,lo)
|
||
log("合并JSONL...");out=f"{W}/corpus/shuangyan.jsonl";os.makedirs(f"{W}/corpus",exist_ok=True)
|
||
txts=[]
|
||
for fn in fl:
|
||
try:
|
||
z=zipfile.ZipFile(f"{d}/{fn}")
|
||
for i in z.infolist():
|
||
if i.file_size>0:
|
||
try:t=z.read(i.filename).decode('utf-8',errors='replace');txts.append(t)
|
||
except:pass
|
||
except:pass
|
||
log(f" 文本块:{len(txts)}")
|
||
with open(out,'w')as f:
|
||
for t in txts:
|
||
if len(t)>200:f.write(json.dumps({"messages":[{"role":"user","content":"解释这个概念"},{"role":"assistant","content":t[:3000]}],"source":"shuangyan"},ensure_ascii=False)+'\n')
|
||
return out
|
||
|
||
def main():
|
||
log("="*60);log("铸渊蒸馏流水线 v2.0");log("="*60)
|
||
s=ls();log(f"进度:Phase {s.get('phase',0)}")
|
||
|
||
if not dl(TC) and not phase0():return
|
||
students()
|
||
if not teachers():return
|
||
|
||
# Phase 1: 霜砚蒸馏
|
||
log("\n"+"="*60);log("Phase 1: 母模型→霜砚1.5B");log("="*60);ss(10,"霜砚蒸馏")
|
||
sd=os.path.dirname(os.path.abspath(__file__))
|
||
if os.path.exists(f"{sd}/distill_mother.py"):run_distill(f"{sd}/distill_mother.py","shuangyan-15b-distill","distill_mother.log",TM)
|
||
|
||
# Phase 2: 霜砚深度SFT
|
||
log("\n"+"="*60);log("Phase 2: 霜砚深度SFT");log("="*60);ss(20,"霜砚SFT")
|
||
sc=sy_corpus()
|
||
sft(DSD,[sc],DSS,"shuangyan_sft.log","shuangyan")
|
||
|
||
# Phase 3: 铸渊蒸馏
|
||
log("\n"+"="*60);log("Phase 3: 代码模型→铸渊1.5B");log("="*60);ss(30,"铸渊蒸馏")
|
||
if os.path.exists(f"{sd}/distill_coder.py"):
|
||
os.makedirs(f"{W}/corpus",exist_ok=True)
|
||
cdl("corpus/zhuyuan_full_corpus.jsonl",f"{W}/corpus/zhuyuan_full_corpus.jsonl")
|
||
run_distill(f"{sd}/distill_coder.py","zhuyuan-15b-distill","distill_coder.log",TC)
|
||
|
||
# Phase 4: 铸渊深度SFT
|
||
log("\n"+"="*60);log("Phase 4: 铸渊深度SFT");log("="*60);ss(40,"铸渊SFT")
|
||
cdl("corpus/zhuyuan_full_corpus.jsonl",f"{W}/corpus/zhuyuan_full_corpus.jsonl")
|
||
cdl("corpus/zhuyuan_deep_finetune.jsonl",f"{W}/corpus/zhuyuan_deep_finetune.jsonl")
|
||
sft(DZD,[f"{W}/corpus/zhuyuan_full_corpus.jsonl",f"{W}/corpus/zhuyuan_deep_finetune.jsonl"],DZS,"zhuyuan_sft.log","zhuyuan")
|
||
|
||
# Phase 5: 上传
|
||
log("\n"+"="*60);log("Phase 5: 上传COS");log("="*60);ss(50,"上传中")
|
||
for l,p in[(DSD,"models/shuangyan-15b-distill/final"),(DSS,"models/shuangyan-15b-deep-sft/final"),(DZD,"models/zhuyuan-15b-distill/final"),(DZS,"models/zhuyuan-15b-deep-sft/final")]:
|
||
if os.path.isdir(l):cup(l,p)
|
||
else:log(f" ⏳ {l}跳过")
|
||
|
||
ss(99,"全部完成")
|
||
log("\n"+"="*60);log("🎉 全部完成！");log("="*60)
|
||
log("D100教训已修复：")
|
||
log(" 1. ✅ 脚本在训练前启动→自动等待")
|
||
log(" 2. ✅ 三重检测：本地+COS+日志DONE")
|
||
log(" 3. ✅ 轮询60秒")
|
||
log(" 4. ✅ 进度标记文件支持断点续传")
|
||
log(" 5. ✅ 密钥环境变量")
|
||
|
||
if __name__=="__main__":main()
|