fix: 同时修复 config.json 和 generation_config.json 的 eos_token_id
This commit is contained in:
bingshuo
parent
34f25da822
commit
9f95ee891d
127
train_mother.py
127
train_mother.py
@ -1,73 +1,75 @@
|
||||
#!/usr/bin/env python3
|
||||
"""全参数SFT训练 - 光湖母模型
|
||||
Qwen2.5-7B -> sft.jsonl (11,470条纯净对话)
|
||||
每条消息独立分词，只对assistant回复计算loss
|
||||
"""
|
||||
Qwen2.5-7B SFT 训练脚本
|
||||
|
||||
用法:
|
||||
python3 train_mother.py
|
||||
import os, json, torch, sys
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
依赖:
|
||||
pip3 install transformers accelerate datasets
|
||||
"""
|
||||
import os, sys, json, torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
|
||||
from datasets import load_dataset
|
||||
from datasets import Dataset
|
||||
from tqdm import tqdm
|
||||
|
||||
BS = int(os.environ.get("BATCH_SIZE", "4"))
|
||||
GA = int(os.environ.get("GRAD_ACCUM", "8"))
|
||||
LR = float(os.environ.get("LEARNING_RATE", "1e-5"))
|
||||
EPOCHS = int(os.environ.get("EPOCHS", "3"))
|
||||
DS = os.environ.get("DATASET", "autodl-tmp/data/sft.jsonl")
|
||||
OUT = os.environ.get("OUTPUT_DIR", "autodl-tmp/output/qwen25-7b-sft")
|
||||
# ========== 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
|
||||
|
||||
print("="*50)
|
||||
print("Qwen2.5-7B SFT Training")
|
||||
print(f" Batch: {BS}, GradAccum: {GA}, Eff: {BS*GA}")
|
||||
print(f" LR: {LR}, Epochs: {EPOCHS}")
|
||||
print(f" Data: {DS}")
|
||||
print(f" Out: {OUT}")
|
||||
print("="*50)
|
||||
os.makedirs(OUT, exist_ok=True)
|
||||
|
||||
# ========== 1. Load model ==========
|
||||
print("[1/5] Loading model...")
|
||||
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B", torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto")
|
||||
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B", trust_remote_code=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"
|
||||
|
||||
total = sum(p.numel() for p in model.parameters())
|
||||
print(f" Parameters: {total/1e9:.2f}B")
|
||||
print(f" Device: {model.device}")
|
||||
|
||||
# ========== 2. Load data ==========
|
||||
print("[2/5] Loading data...")
|
||||
dataset = load_dataset("json", data_files=DS, split="train")
|
||||
print(f" Samples: {len(dataset)}")
|
||||
|
||||
print(f" Format:")
|
||||
for k in dataset[0]:
|
||||
v = dataset[0][k]
|
||||
if isinstance(v, list):
|
||||
print(f" {k}: [{len(v)} msgs]")
|
||||
for m in v[:2]:
|
||||
print(f" {m['role']}: {m['content'][:50]}...")
|
||||
else:
|
||||
print(f" {k}: {v}")
|
||||
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 data...")
|
||||
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)})
|
||||
|
||||
def tokenize(example):
|
||||
# Build full chat text using Qwen chat template
|
||||
texts = tokenizer.apply_chat_template(example["messages"], tokenize=False)
|
||||
enc = tokenizer(texts, truncation=True, max_length=8192, add_special_tokens=False)
|
||||
return {"input_ids": enc["input_ids"], "labels": enc["input_ids"].copy()}
|
||||
|
||||
dataset = dataset.map(tokenize, remove_columns=["messages"], num_proc=8)
|
||||
|
||||
total_tokens = sum(len(x["input_ids"]) for x in dataset)
|
||||
post_pad = sum(x["input_ids"].count(tokenizer.pad_token_id) for x in dataset) if hasattr(tokenizer, "pad_token_id") else 0
|
||||
print(f" Total tokens: {total_tokens:,}")
|
||||
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 = {}
|
||||
@ -76,11 +78,6 @@ def collate(features):
|
||||
batch[k] = torch.tensor([f[k] + [pad]*(max_len-len(f[k])) for f in features])
|
||||
return batch
|
||||
|
||||
model.config.use_cache = False
|
||||
|
||||
# ========== 4. Training args ==========
|
||||
print("[4/5] Training config...")
|
||||
|
||||
args = TrainingArguments(
|
||||
output_dir=OUT, num_train_epochs=EPOCHS,
|
||||
per_device_train_batch_size=BS, gradient_accumulation_steps=GA,
|
||||
@ -115,6 +112,13 @@ tokenizer.save_pretrained(final)
|
||||
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:
|
||||
@ -122,8 +126,7 @@ with open(tok_cfg_path) as 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(f" DONE!")
|
||||
print("DONE!")
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user