237 lines
8.6 KiB
Python
237 lines
8.6 KiB
Python
#!/usr/bin/env python3
|
||
# Copyright (c) 2026 XingChina
|
||
# SPDX-License-Identifier: BSD-3-Clause
|
||
# 本代码采用 BSD 3-Clause 许可证，详见项目根目录的 LICENSE 文件。
|
||
# 为4090显卡进行优化，保证普通显卡也能训练（而不是找大显存显卡）
|
||
# 已经拿4090测试过，全程没有崩溃
|
||
|
||
import os
|
||
import json
|
||
import signal
|
||
import sys
|
||
import torch
|
||
from transformers import (
|
||
GPT2Config,
|
||
GPT2LMHeadModel,
|
||
GPT2Tokenizer,
|
||
Trainer,
|
||
TrainingArguments,
|
||
DataCollatorForLanguageModeling,
|
||
)
|
||
from datasets import Dataset
|
||
from bitsandbytes.optim import Adam8bit
|
||
import logging
|
||
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ================== 模型配置 ==================
|
||
MODEL_CONFIG = {
|
||
"n_embd": 1024,
|
||
"n_layer": 16,
|
||
"n_head": 16,
|
||
"n_positions": 4096,
|
||
"vocab_size": 50257,
|
||
"n_ctx": 4096,
|
||
"resid_pdrop": 0.3,
|
||
"embd_pdrop": 0.3,
|
||
"attn_pdrop": 0.3,
|
||
}
|
||
|
||
# ================== 训练参数（为4090优化，保证普通显卡能复现） ==================
|
||
TRAIN_ARGS = {
|
||
"output_dir": "./checkpoints_chunmengdie_gpt2",
|
||
"per_device_train_batch_size": 2,
|
||
"gradient_accumulation_steps": 8,
|
||
"num_train_epochs": 1,
|
||
"learning_rate": 3e-4,
|
||
"weight_decay": 0.01,
|
||
"warmup_steps": 500,
|
||
"logging_steps": 10,
|
||
"save_steps": 500,
|
||
"save_total_limit": 3,
|
||
"bf16": True,
|
||
"report_to": "none",
|
||
"dataloader_num_workers": 0,
|
||
"eval_strategy": "steps",
|
||
"eval_steps": 500,
|
||
"load_best_model_at_end": True,
|
||
"metric_for_best_model": "eval_loss",
|
||
"greater_is_better": False,
|
||
"gradient_checkpointing": True,
|
||
}
|
||
|
||
MAX_SEQ_LEN = 4096
|
||
CHECKPOINT_DIR = TRAIN_ARGS["output_dir"]
|
||
|
||
# ====== 数据文件列表（这里假设数据已经清理过，实际用的数据确实清理过） ======
|
||
# mengdie_train由于包含其他数据集，没有开源
|
||
DATA_FILES = [
|
||
"mengdie_train.json",
|
||
"Belle_open_source_0.5M.json",
|
||
]
|
||
|
||
# ================== 数据加载 ==================
|
||
def load_data(file_list):
|
||
texts = []
|
||
for fpath in file_list:
|
||
if not os.path.exists(fpath):
|
||
logger.warning(f"File not found: {fpath}, skipped")
|
||
continue
|
||
with open(fpath, 'r', encoding='utf-8') as f:
|
||
first_char = f.read(1)
|
||
f.seek(0)
|
||
if first_char == '[':
|
||
data = json.load(f)
|
||
logger.info(f"Loaded {len(data)} samples from {fpath} (JSON array)")
|
||
else:
|
||
data = []
|
||
for line in f:
|
||
line = line.strip()
|
||
if line:
|
||
try:
|
||
data.append(json.loads(line))
|
||
except json.JSONDecodeError:
|
||
pass
|
||
logger.info(f"Loaded {len(data)} samples from {fpath} (JSONL)")
|
||
|
||
for item in data:
|
||
try:
|
||
if "instruction" in item and "input" in item and "output" in item:
|
||
inst = item["instruction"].strip()
|
||
inp = item.get("input", "").strip()
|
||
out = item["output"].strip()
|
||
user_text = f"{inst}\n{inp}" if inp else inst
|
||
assistant_text = out
|
||
elif "input" in item and "output" in item:
|
||
user_text = item["input"].strip()
|
||
assistant_text = item["output"].strip()
|
||
else:
|
||
continue
|
||
texts.append(f"用户：{user_text}\n猫娘：{assistant_text}")
|
||
except Exception:
|
||
continue
|
||
logger.info(f"Total texts: {len(texts)}")
|
||
return texts
|
||
|
||
# ================== Tokenization==================
|
||
def tokenize_function(examples):
|
||
tokenized = tokenizer(
|
||
examples["text"],
|
||
truncation=True,
|
||
max_length=MAX_SEQ_LEN - 1,
|
||
padding=False,
|
||
return_attention_mask=False,
|
||
)
|
||
tokenized["input_ids"] = [ids + [tokenizer.eos_token_id] for ids in tokenized["input_ids"]]
|
||
return tokenized
|
||
|
||
# ================== 紧急保存 ==================
|
||
def emergency_save(sig, frame):
|
||
logger.info("\n🛑 Saving emergency checkpoint...")
|
||
try:
|
||
torch.cuda.synchronize()
|
||
model_to_save = model.cpu()
|
||
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
|
||
model_to_save.save_pretrained(os.path.join(CHECKPOINT_DIR, "emergency"))
|
||
tokenizer.save_pretrained(os.path.join(CHECKPOINT_DIR, "emergency"))
|
||
logger.info("✅ Emergency checkpoint saved.")
|
||
except Exception as e:
|
||
logger.error(f"Emergency save failed: {e}")
|
||
sys.exit(0)
|
||
|
||
# ================== 主程序 ==================
|
||
if __name__ == "__main__":
|
||
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
|
||
tokenizer.pad_token = tokenizer.eos_token
|
||
|
||
raw_texts = load_data(DATA_FILES)
|
||
dataset = Dataset.from_dict({"text": raw_texts})
|
||
|
||
tokenized_dataset = dataset.map(
|
||
tokenize_function,
|
||
batched=True,
|
||
remove_columns=["text"],
|
||
num_proc=4,
|
||
load_from_cache_file=True,
|
||
)
|
||
|
||
# 验证 EOS 是否添加成功
|
||
sample_ids = tokenized_dataset[0]["input_ids"]
|
||
logger.info(f"Sample last 5 tokens: {sample_ids[-5:]}")
|
||
logger.info(f"Last token is EOS: {sample_ids[-1] == tokenizer.eos_token_id}")
|
||
|
||
# 9:1 划分
|
||
total_len = len(tokenized_dataset)
|
||
train_size = int(0.9 * total_len)
|
||
eval_size = total_len - train_size
|
||
train_dataset, eval_dataset = torch.utils.data.random_split(
|
||
tokenized_dataset, [train_size, eval_size]
|
||
)
|
||
logger.info(f"Train: {train_size}, Eval: {eval_size}")
|
||
|
||
logger.info("Initializing GPT-2 model from scratch...")
|
||
config = GPT2Config(**MODEL_CONFIG)
|
||
model = GPT2LMHeadModel(config)
|
||
|
||
optimizer = Adam8bit(model.parameters(), lr=TRAIN_ARGS["learning_rate"])
|
||
|
||
training_args = TrainingArguments(
|
||
output_dir=CHECKPOINT_DIR,
|
||
per_device_train_batch_size=TRAIN_ARGS["per_device_train_batch_size"],
|
||
gradient_accumulation_steps=TRAIN_ARGS["gradient_accumulation_steps"],
|
||
num_train_epochs=TRAIN_ARGS["num_train_epochs"],
|
||
learning_rate=TRAIN_ARGS["learning_rate"],
|
||
weight_decay=TRAIN_ARGS["weight_decay"],
|
||
warmup_steps=TRAIN_ARGS["warmup_steps"],
|
||
logging_steps=TRAIN_ARGS["logging_steps"],
|
||
save_steps=TRAIN_ARGS["save_steps"],
|
||
save_total_limit=TRAIN_ARGS["save_total_limit"],
|
||
bf16=TRAIN_ARGS["bf16"],
|
||
report_to=TRAIN_ARGS["report_to"],
|
||
dataloader_num_workers=TRAIN_ARGS["dataloader_num_workers"],
|
||
optim="adamw_8bit",
|
||
eval_strategy=TRAIN_ARGS["eval_strategy"],
|
||
eval_steps=TRAIN_ARGS["eval_steps"],
|
||
load_best_model_at_end=TRAIN_ARGS["load_best_model_at_end"],
|
||
metric_for_best_model=TRAIN_ARGS["metric_for_best_model"],
|
||
greater_is_better=TRAIN_ARGS["greater_is_better"],
|
||
gradient_checkpointing=TRAIN_ARGS["gradient_checkpointing"],
|
||
)
|
||
|
||
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
|
||
|
||
trainer = Trainer(
|
||
model=model,
|
||
args=training_args,
|
||
train_dataset=train_dataset,
|
||
eval_dataset=eval_dataset,
|
||
data_collator=data_collator,
|
||
optimizers=(optimizer, None),
|
||
)
|
||
|
||
signal.signal(signal.SIGINT, emergency_save)
|
||
|
||
# 自动恢复最新 checkpoint（中间checkpoint保留调度器，最终权重无调度器，这是手滑删了调度器的主要原因）
|
||
latest_checkpoint = None
|
||
if os.path.exists(CHECKPOINT_DIR):
|
||
checkpoints = [d for d in os.listdir(CHECKPOINT_DIR) if d.startswith("checkpoint-")]
|
||
if checkpoints:
|
||
latest_checkpoint = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))[-1]
|
||
latest_checkpoint = os.path.join(CHECKPOINT_DIR, latest_checkpoint)
|
||
logger.info(f"✅ Found checkpoint: {latest_checkpoint}, will resume from there.")
|
||
else:
|
||
if os.path.exists(os.path.join(CHECKPOINT_DIR, "emergency")):
|
||
latest_checkpoint = os.path.join(CHECKPOINT_DIR, "emergency")
|
||
logger.info(f"✅ Found emergency checkpoint, will resume from there.")
|
||
else:
|
||
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
|
||
|
||
logger.info("🚀 Starting training...")
|
||
trainer.train(resume_from_checkpoint=latest_checkpoint)
|
||
|
||
final_path = os.path.join(CHECKPOINT_DIR, "final_model")
|
||
model.save_pretrained(final_path)
|
||
tokenizer.save_pretrained(final_path)
|
||
logger.info(f"✅ Final model saved to {final_path}")
|
||
logger.info("📌 Convert to GGUF: python /path/to/llama.cpp/convert_hf_to_gguf.py ./checkpoints_chunmengdie_gpt2/final_model --outfile model.gguf --outtype q8_0") |