import os
import torch
import transformers
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TrainingArguments,
DataCollatorForLanguageModeling,
BitsAndBytesConfig,
Trainer
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_dataset
import logging
import psutil
import gc
from datetime import datetime
# === 配置区域 ===
MODEL_NAME = "/home/vipuser/ai_writer_project_final_with_fixed_output_ui/models/Yi-6B"
DATASET_PATH = "./data/train_lora_formatted.jsonl"
OUTPUT_DIR = "./yi6b-lora-optimized"
DEVICE_MAP = "auto" # 使用自动设备映射
# 确保输出目录存在
os.makedirs(OUTPUT_DIR, exist_ok=True)
# === 内存优化配置 ===
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # 减少内存碎片
torch.backends.cuda.cufft_plan_cache.clear() # 清理CUDA缓存
# === 增强的日志系统 ===
def setup_logging(output_dir):
"""配置日志系统，支持文件和TensorBoard"""
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# 文件日志处理器
file_handler = logging.FileHandler(os.path.join(output_dir, "training.log"))
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger.addHandler(file_handler)
# 控制台日志处理器
console_handler = logging.StreamHandler()
console_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
logger.addHandler(console_handler)
# TensorBoard日志目录
tensorboard_log_dir = os.path.join(output_dir, "logs", datetime.now().strftime("%Y%m%d-%H%M%S"))
os.makedirs(tensorboard_log_dir, exist_ok=True)
# 安装TensorBoard回调
tb_writer = None
try:
from torch.utils.tensorboard import SummaryWriter
tb_writer = SummaryWriter(log_dir=tensorboard_log_dir)
logger.info(f"TensorBoard日志目录: {tensorboard_log_dir}")
except ImportError:
logger.warning("TensorBoard未安装，可视化功能不可用")
return logger, tb_writer
logger, tb_writer = setup_logging(OUTPUT_DIR)
# === 量化配置 - 使用更高效的配置 ===
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
# === 加载模型 ===
logger.info("加载预训练模型...")
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
device_map=DEVICE_MAP,
quantization_config=quant_config,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
attn_implementation="flash_attention_2" # 使用FlashAttention优化内存
)
# === 分词器处理 ===
logger.info("加载分词器...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
tokenizer.padding_side = "right"
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
# === 准备模型训练 ===
model = prepare_model_for_kbit_training(
model,
use_gradient_checkpointing=True # 启用梯度检查点以节省内存
)
# === LoRA 配置 - 优化内存使用 ===
logger.info("配置LoRA...")
lora_config = LoraConfig(
r=64, # 降低rank以减少内存使用
lora_alpha=32, # 降低alpha值
target_modules=["q_proj", "v_proj"], # 减少目标模块
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)
# 记录可训练参数
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
logger.info(f"可训练参数: {trainable_params:,} / 总参数: {total_params:,} ({trainable_params/total_params:.2%})")
# === 加载并预处理数据集 ===
logger.info("加载和预处理数据集...")
dataset = load_dataset("json", data_files=DATASET_PATH, split="train")
# 文本过滤函数
def is_valid_text(example):
text = example.get("text", "")
return text is not None and len(text.strip()) > 200 # 增加最小长度要求
dataset = dataset.filter(is_valid_text)
logger.info(f"过滤后数据集大小: {len(dataset)} 条")
# 动态填充的分词函数 - 节省内存
def tokenize_function(examples):
tokenized = tokenizer(
examples["text"],
padding=True, # 使用动态填充
truncation=True,
max_length=1024, # 降低上下文长度以减少内存使用
)
# 创建 labels - 因果语言建模需要 labels = input_ids
tokenized["labels"] = tokenized["input_ids"].copy()
return tokenized
tokenized_dataset = dataset.map(
tokenize_function,
batched=True,
remove_columns=["text"],
batch_size=64, # 降低批处理大小以减少内存峰值
num_proc=4, # 减少进程数以降低内存开销
)
# === 数据整理器 ===
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer,
mlm=False # 因果语言建模
)
# === 训练参数 - 优化内存使用 ===
report_to_list = ["tensorboard"] if tb_writer else []
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
per_device_train_batch_size=4, # 大幅降低批次大小
gradient_accumulation_steps=4, # 增加梯度累积步数以保持有效批次大小
learning_rate=2e-5,
num_train_epochs=3,
logging_steps=50,
save_strategy="steps",
save_steps=500,
bf16=True,
optim="paged_adamw_32bit",
report_to=report_to_list,
warmup_ratio=0.05,
gradient_checkpointing=True, # 启用梯度检查点
fp16=False,
max_grad_norm=0.3, # 降低梯度裁剪阈值
remove_unused_columns=True, # 移除未使用的列以节省内存
dataloader_num_workers=4, # 减少数据加载工作线程
evaluation_strategy="steps",
eval_steps=500,
save_total_limit=2, # 减少保存的检查点数量
logging_dir=os.path.join(OUTPUT_DIR, "logs"),
load_best_model_at_end=True,
ddp_find_unused_parameters=False,
logging_first_step=True,
group_by_length=True,
lr_scheduler_type="cosine",
weight_decay=0.01,
)
# === GPU监控工具 ===
def monitor_gpu():
"""监控GPU使用情况"""
if torch.cuda.is_available():
device = torch.device("cuda")
mem_alloc = torch.cuda.memory_allocated(device) / 1024**3
mem_reserved = torch.cuda.memory_reserved(device) / 1024**3
mem_total = torch.cuda.get_device_properties(device).total_memory / 1024**3
return {
"allocated": f"{mem_alloc:.2f} GB",
"reserved": f"{mem_reserved:.2f} GB",
"total": f"{mem_total:.2f} GB",
"utilization": f"{mem_alloc/mem_total*100:.1f}%"
}
return {}
# === 创建训练器 ===
eval_dataset = None
if len(tokenized_dataset) > 100:
eval_dataset = tokenized_dataset.select(range(100))
trainer = Trainer(
model=model,
tokenizer=tokenizer,
args=training_args,
train_dataset=tokenized_dataset,
eval_dataset=eval_dataset,
data_collator=data_collator,
)
# === 训练前验证 ===
def validate_data_and_model():
"""验证数据和模型是否准备好训练"""
logger.info("\n=== 训练前验证 ===")
# 检查样本格式
sample = tokenized_dataset[0]
logger.info(f"样本键: {list(sample.keys())}")
logger.info(f"input_ids 长度: {len(sample['input_ids'])}")
# 创建单个样本测试批次
test_batch = data_collator([sample])
# 移动数据到设备
test_batch = {k: v.to(model.device) for k, v in test_batch.items()}
# 前向传播测试
model.train()
outputs = model(**test_batch)
loss_value = outputs.loss.item()
logger.info(f"测试批次损失: {loss_value:.4f}")
# 记录到TensorBoard
if tb_writer:
tb_writer.add_scalar("debug/test_loss", loss_value, 0)
# 反向传播测试
outputs.loss.backward()
logger.info("反向传播成功!")
# 重置梯度
model.zero_grad()
logger.info("验证完成，准备开始训练\n")
# 记录初始GPU使用情况
gpu_status = monitor_gpu()
logger.info(f"初始GPU状态: {gpu_status}")
# 记录到TensorBoard
if tb_writer:
tb_writer.add_text("system/initial_gpu", str(gpu_status), 0)
validate_data_and_model()
# === 自定义回调 - 监控资源使用 ===
class ResourceMonitorCallback(transformers.TrainerCallback):
def __init__(self, tb_writer=None):
self.tb_writer = tb_writer
self.start_time = datetime.now()
self.last_log_time = datetime.now()
def on_step_end(self, args, state, control, **kwargs):
current_time = datetime.now()
time_diff = (current_time - self.last_log_time).total_seconds()
# 每分钟记录一次资源使用情况
if time_diff > 60:
self.last_log_time = current_time
# GPU监控
gpu_status = monitor_gpu()
logger.info(f"Step {state.global_step} - GPU状态: {gpu_status}")
# CPU和内存监控
cpu_percent = psutil.cpu_percent()
mem = psutil.virtual_memory()
logger.info(f"CPU使用率: {cpu_percent}%, 内存使用: {mem.used/1024**3:.2f}GB/{mem.total/1024**3:.2f}GB")
# 记录到TensorBoard
if self.tb_writer:
# GPU显存使用
if torch.cuda.is_available():
device = torch.device("cuda")
mem_alloc = torch.cuda.memory_allocated(device) / 1024**3
self.tb_writer.add_scalar("system/gpu_mem", mem_alloc, state.global_step)
# CPU使用率
self.tb_writer.add_scalar("system/cpu_usage", cpu_percent, state.global_step)
# 系统内存使用
self.tb_writer.add_scalar("system/ram_usage", mem.used/1024**3, state.global_step)
def on_log(self, args, state, control, logs=None, **kwargs):
"""记录训练指标到TensorBoard"""
if self.tb_writer and logs is not None:
for metric_name, metric_value in logs.items():
if "loss" in metric_name or "lr" in metric_name or "grad_norm" in metric_name:
self.tb_writer.add_scalar(f"train/{metric_name}", metric_value, state.global_step)
def on_train_end(self, args, state, control, **kwargs):
"""训练结束时记录总时间"""
training_time = datetime.now() - self.start_time
logger.info(f"训练总时间: {training_time}")
if self.tb_writer:
self.tb_writer.add_text("system/total_time", str(training_time))
# 添加回调
trainer.add_callback(ResourceMonitorCallback(tb_writer=tb_writer))
# === 内存清理函数 ===
def clear_memory():
"""清理内存和GPU缓存"""
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
logger.info("内存清理完成")
# === 启动训练 ===
try:
logger.info("开始训练...")
# 分阶段训练以减少内存峰值
num_samples = len(tokenized_dataset)
chunk_size = 1000 # 每次处理1000个样本
for i in range(0, num_samples, chunk_size):
end_idx = min(i + chunk_size, num_samples)
logger.info(f"训练样本 {i} 到 {end_idx-1} / {num_samples}")
# 创建子数据集
chunk_dataset = tokenized_dataset.select(range(i, end_idx))
# 更新训练器
trainer.train_dataset = chunk_dataset
# 训练当前块
trainer.train()
# 清理内存
clear_memory()
# 保存训练指标
metrics = trainer.evaluate()
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
# 保存最佳模型
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
logger.info(f"训练完成! 模型保存在: {OUTPUT_DIR}")
# 记录最终指标到TensorBoard
if tb_writer:
for metric_name, metric_value in metrics.items():
tb_writer.add_scalar(f"final/{metric_name}", metric_value)
tb_writer.close()
except Exception as e:
logger.error(f"训练出错: {e}")
import traceback
logger.error(traceback.format_exc())
# 尝试更小批量训练
logger.info("\n尝试更小批量训练...")
small_dataset = tokenized_dataset.select(range(50))
trainer.train_dataset = small_dataset
trainer.train()
# 保存模型
trainer.save_model(f"{OUTPUT_DIR}_small")
tokenizer.save_pretrained(f"{OUTPUT_DIR}_small")
logger.info(f"小批量训练完成! 模型保存在: {OUTPUT_DIR}_small")
# 记录错误到TensorBoard
if tb_writer:
tb_writer.add_text("error/exception", traceback.format_exc())
# 清理内存
clear_memory()
# === 训练后验证 ===
def validate_final_model():
"""验证训练后的模型"""
logger.info("\n=== 训练后验证 ===")
# 加载保存的模型
from peft import PeftModel
# 仅加载基础模型配置
base_model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
device_map=DEVICE_MAP,
quantization_config=quant_config,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
load_in_4bit=True
)
# 加载LoRA适配器
peft_model = PeftModel.from_pretrained(base_model, OUTPUT_DIR)
# 合并LoRA权重
merged_model = peft_model.merge_and_unload()
# 测试生成
prompt = "中国的首都是"
inputs = tokenizer(prompt, return_tensors="pt").to(merged_model.device)
outputs = merged_model.generate(
**inputs,
max_new_tokens=50, # 减少生成长度
temperature=0.7,
top_p=0.9,
repetition_penalty=1.2,
do_sample=True
)
generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
logger.info(f"提示: {prompt}")
logger.info(f"生成结果: {generated}")
# 记录到TensorBoard
if tb_writer:
tb_writer.add_text("validation/sample", f"提示: {prompt}\n生成: {generated}")
# 更全面的测试
test_prompts = [
"人工智能的未来发展趋势是",
"如何学习深度学习？",
"写一个关于太空探索的短故事："
]
for i, test_prompt in enumerate(test_prompts):
inputs = tokenizer(test_prompt, return_tensors="pt").to(merged_model.device)
outputs = merged_model.generate(
**inputs,
max_new_tokens=100, # 减少生成长度
temperature=0.7,
top_p=0.9,
repetition_penalty=1.2,
do_sample=True
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
logger.info(f"\n提示: {test_prompt}\n生成: {generated_text}\n{'='*50}")
# 记录到TensorBoard
if tb_writer:
tb_writer.add_text(f"validation/test_{i}", f"提示: {test_prompt}\n生成: {generated_text}")
logger.info("验证完成")
# 执行验证
validate_final_model()
# 关闭TensorBoard写入器
if tb_writer:
tb_writer.close()
logger.info("TensorBoard日志已关闭")
(.venv) (base) vipuser@ubuntu22:~/ai_writer_project_final_with_fixed_output_ui$ python train_lora.py
2025-07-13 22:10:19,098 - INFO - TensorBoard日志目录: ./yi6b-lora-optimized/logs/20250713-221019
2025-07-13 22:10:19,099 - INFO - 加载预训练模型...
Traceback (most recent call last):
File "/home/vipuser/ai_writer_project_final_with_fixed_output_ui/train_lora.py", line 77, in <module>
model = AutoModelForCausalLM.from_pretrained(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vipuser/ai_writer_project_final_with_fixed_output_ui/.venv/lib/python3.11/site-packages/transformers/models/auto/auto_factory.py", line 566, in from_pretrained
return model_class.from_pretrained(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vipuser/ai_writer_project_final_with_fixed_output_ui/.venv/lib/python3.11/site-packages/transformers/modeling_utils.py", line 3590, in from_pretrained
config = cls._autoset_attn_implementation(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vipuser/ai_writer_project_final_with_fixed_output_ui/.venv/lib/python3.11/site-packages/transformers/modeling_utils.py", line 1389, in _autoset_attn_implementation
cls._check_and_enable_flash_attn_2(
File "/home/vipuser/ai_writer_project_final_with_fixed_output_ui/.venv/lib/python3.11/site-packages/transformers/modeling_utils.py", line 1480, in _check_and_enable_flash_attn_2
raise ImportError(f"{preface} the package flash_attn seems to be not installed. {install_message}")
ImportError: FlashAttention2 has been toggled on, but it cannot be used due to the following error: the package flash_attn seems to be not installed. Please refer to the documentation of https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2 to install Flash Attention 2.
本文介绍了一个与AndroidAnnotations框架相关的错误，该错误导致工程出现红叉并无法通过clear命令解决。文章详细描述了解决过程，包括定位问题源文件及最终通过删除特定自动生成文件夹来修复问题。
3835
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
扫一扫
