133 lines
4.8 KiB
Python
133 lines
4.8 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""全参数SFT训练 - 光湖母模型
|
|||
|
|
Qwen2.5-7B -> sft.jsonl (11,470条纯净对话)
|
|||
|
|
每条消息独立分词，只对assistant回复计算loss
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os, json, torch, sys
|
|||
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
|||
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|||
|
|
|
|||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
|
|||
|
|
from datasets import Dataset
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
|
|||
|
|
# ========== Config ==========
|
|||
|
|
MODEL_PATH = "/root/autodl-tmp/cache/Qwen/Qwen2___5-7B"
|
|||
|
|
DATA = "/root/autodl-tmp/data/sft.jsonl"
|
|||
|
|
OUT = "/root/autodl-tmp/output/qwen25-7b-sft"
|
|||
|
|
EPOCHS = 3
|
|||
|
|
BS = 1
|
|||
|
|
GA = 8
|
|||
|
|
LR = 2e-5
|
|||
|
|
MAX_LEN = 2048
|
|||
|
|
|
|||
|
|
os.makedirs(OUT, exist_ok=True)
|
|||
|
|
|
|||
|
|
# ========== 1. Load data ==========
|
|||
|
|
print("[1/5] Loading data...")
|
|||
|
|
with open(DATA) as f:
|
|||
|
|
raw = [json.loads(line) for line in f]
|
|||
|
|
raw = [{"messages": [m for m in obj["messages"] if m["role"] != "system"]} for obj in raw if any(m["role"] != "system" for m in obj["messages"])]
|
|||
|
|
print(f" {len(raw)} examples (system filtered)")
|
|||
|
|
|
|||
|
|
# ========== 2. Load model ==========
|
|||
|
|
print(f"[2/5] Loading Qwen/Qwen2.5-7B...")
|
|||
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
|
|||
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|||
|
|
tokenizer.padding_side = "right"
|
|||
|
|
|
|||
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|||
|
|
MODEL_PATH, trust_remote_code=True,
|
|||
|
|
torch_dtype=torch.bfloat16, attn_implementation="sdpa",
|
|||
|
|
).cuda()
|
|||
|
|
model.config.use_cache = False
|
|||
|
|
model.gradient_checkpointing_enable()
|
|||
|
|
print(f" Params: {sum(p.numel() for p in model.parameters())/1e9:.2f}B (full FT)")
|
|||
|
|
|
|||
|
|
# ========== 3. Tokenize ==========
|
|||
|
|
print("[3/5] Tokenizing...")
|
|||
|
|
processed = []
|
|||
|
|
for d in tqdm(raw, desc="Tokenize"):
|
|||
|
|
ids, labs = [], []
|
|||
|
|
for msg in d["messages"]:
|
|||
|
|
c = msg["content"]
|
|||
|
|
if not c.strip():
|
|||
|
|
continue
|
|||
|
|
t = f"<|im_start|>{msg['role']}\n{c}<|im_end|>\n"
|
|||
|
|
tok = tokenizer.encode(t, add_special_tokens=False)
|
|||
|
|
ids.extend(tok)
|
|||
|
|
labs.extend(tok if msg["role"] == "assistant" else [-100] * len(tok))
|
|||
|
|
if len(ids) > MAX_LEN:
|
|||
|
|
ids, labs = ids[:MAX_LEN], labs[:MAX_LEN]
|
|||
|
|
processed.append({"input_ids": ids, "labels": labs, "attention_mask": [1]*len(ids)})
|
|||
|
|
|
|||
|
|
ds = Dataset.from_list(processed)
|
|||
|
|
total_tok = sum(len(d["input_ids"]) for d in processed)
|
|||
|
|
loss_tok = sum(sum(1 for l in d["labels"] if l != -100) for d in processed)
|
|||
|
|
print(f" Dataset: {len(ds)} ex, {total_tok:,} tokens, {loss_tok:,} loss ({loss_tok/max(total_tok,1)*100:.1f}%)")
|
|||
|
|
sys.stdout.flush()
|
|||
|
|
|
|||
|
|
# ========== 4. Train ==========
|
|||
|
|
print("[4/5] Training config...")
|
|||
|
|
def collate(features):
|
|||
|
|
max_len = max(len(f["input_ids"]) for f in features)
|
|||
|
|
batch = {}
|
|||
|
|
for k in ["input_ids", "labels", "attention_mask"]:
|
|||
|
|
pad = tokenizer.pad_token_id if k != "labels" else -100
|
|||
|
|
batch[k] = torch.tensor([f[k] + [pad]*(max_len-len(f[k])) for f in features])
|
|||
|
|
return batch
|
|||
|
|
|
|||
|
|
args = TrainingArguments(
|
|||
|
|
output_dir=OUT, num_train_epochs=EPOCHS,
|
|||
|
|
per_device_train_batch_size=BS, gradient_accumulation_steps=GA,
|
|||
|
|
learning_rate=LR, 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, dataloader_num_workers=4,
|
|||
|
|
gradient_checkpointing=True, optim="adamw_torch",
|
|||
|
|
report_to="none", ddp_find_unused_parameters=False,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
trainer = Trainer(model=model, args=args, train_dataset=ds, data_collator=collate)
|
|||
|
|
|
|||
|
|
# ========== 5. Go ==========
|
|||
|
|
print("[5/5] Starting training!")
|
|||
|
|
gpu = torch.cuda.get_device_name(0)
|
|||
|
|
mem = torch.cuda.get_device_properties(0).total_memory / 1e9
|
|||
|
|
print(f" GPU: {gpu} ({mem:.1f}GB) | Epochs: {EPOCHS} | Eff batch: {BS*GA} | LR: {LR}")
|
|||
|
|
sys.stdout.flush()
|
|||
|
|
|
|||
|
|
trainer.train()
|
|||
|
|
|
|||
|
|
# ========== 6. Save ==========
|
|||
|
|
print("Saving model...")
|
|||
|
|
final = os.path.join(OUT, "final")
|
|||
|
|
trainer.save_model(final)
|
|||
|
|
tokenizer.save_pretrained(final)
|
|||
|
|
|
|||
|
|
# ⚠️ 关键修复：Qwen chat template 使用 <|im_end|> (151645) 作为对话EOS
|
|||
|
|
# 但 config.json 中默认 eos_token_id=151643 (<|endoftext|>)
|
|||
|
|
# 不修复会导致部署时模型无限生成 → 死循环乱码
|
|||
|
|
model.config.eos_token_id = 151645
|
|||
|
|
model.config.save_pretrained(final)
|
|||
|
|
|
|||
|
|
# ⚠️ generation_config.json 也必须修复！
|
|||
|
|
# HuggingFace 的 model.generate() 读取 generation_config.json，不是 config.json
|
|||
|
|
# 不修这个 → 部署后仍然无限生成 → 乱码
|
|||
|
|
model.generation_config.eos_token_id = 151645
|
|||
|
|
model.generation_config.pad_token_id = 151645
|
|||
|
|
model.generation_config.save_pretrained(final)
|
|||
|
|
|
|||
|
|
# 修复 tokenizer 默认system prompt
|
|||
|
|
tok_cfg_path = os.path.join(final, "tokenizer_config.json")
|
|||
|
|
with open(tok_cfg_path) as f:
|
|||
|
|
tok_cfg = json.load(f)
|
|||
|
|
tok_cfg["default_system"] = ""
|
|||
|
|
with open(tok_cfg_path, "w") as f:
|
|||
|
|
json.dump(tok_cfg, f, indent=2, ensure_ascii=False)
|
|||
|
|
peak = torch.cuda.max_memory_allocated() / 1e9
|
|||
|
|
print(f" Model: {final}")
|
|||
|
|
print(f" Peak VRAM: {peak:.2f}GB / {mem:.1f}GB")
|
|||
|
|
print("DONE!")
|