<think>我们遇到了一个CUDA内存不足的错误。错误信息显示我们试图分配688 MiB，但GPU上只有652.12 MiB空闲。同时，已经分配了76.57 GiB，还有1.42 GiB是预留给PyTorch但未分配的。
考虑到我们使用的是4位量化的模型（Yi-6B），理论上模型本身占用的显存应该大大减少（大约在6B*0.5=3GB左右，加上一些开销）。但是，训练过程中需要存储梯度、优化器状态和激活值，这些都会增加显存使用。
我们当前的训练参数：
per_device_train_batch_size=16
max_length=2048
每个样本的显存占用大约是：2048 tokens * 每个token的隐藏层维度（例如4096）* 每个参数4位（0.5字节）* 3（用于梯度、激活等）？但实际计算更复杂。
为了解决这个问题，我们可以尝试以下策略：
1. 减少批处理大小（per_device_train_batch_size）。这是最直接的方法。
2. 使用梯度累积（gradient_accumulation_steps）。这样可以在不增加显存的情况下模拟更大的批次。
3. 启用梯度检查点（gradient_checkpointing）。这会用计算时间换取显存。
4. 减少最大序列长度（max_length）。如果我们的数据不需要2048，可以适当减少。
5. 使用更小的LoRA rank（r）值。
由于我们使用的是4位量化和LoRA，已经非常节省显存了，所以现在的问题可能是批次太大。
我们当前的批次大小是16，我们可以尝试降低到8或者4，同时增加梯度累积步数以保持有效批次大小不变。
另外，我们还可以尝试启用梯度检查点。
让我们调整参数：
原有效批次大小 = per_device_train_batch_size * gradient_accumulation_steps * 设备数
我们目前gradient_accumulation_steps=1，设备数=1（假设单卡）。
如果我们希望保持有效批次大小不变（比如16），我们可以：
将per_device_train_batch_size设为8，gradient_accumulation_steps设为2。
但是，由于我们现在已经OOM，我们可能需要先降低有效批次大小。我们可以尝试：
per_device_train_batch_size=8
gradient_accumulation_steps=2 # 这样有效批次大小=8*2=16，和原来一样。
或者，如果我们想先确保能运行，可以设置更小：
per_device_train_bump_size=4
gradient_accumulation_steps=4 # 有效批次大小=16
另外，我们还可以启用梯度检查点：
gradient_checkpointing=True
注意：梯度检查点会降低训练速度（约20%），但可以节省大量显存。
另外，我们还可以尝试减少序列长度，比如从2048降到1024，但这可能会影响模型性能。
考虑到我们的数据是长文本，可能确实需要2048，所以先不调整长度。
我们尝试以下调整：
1. 设置gradient_accumulation_steps=2，同时将per_device_train_batch_size减半到8。
2. 启用梯度检查点：gradient_checkpointing=True
修改TrainingArguments中的相应参数。
另外，在准备模型的时候，我们之前关闭了梯度检查点，现在要开启，所以需要修改：
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True) # 开启梯度检查点
同时，在TrainingArguments中设置gradient_checkpointing=True。
注意：梯度检查点需要同时在模型准备和训练参数中启用。
另外，我们还可以尝试使用更小的LoRA rank（比如64），但作为第二选择。
我们首先调整批次和梯度检查点。
修改代码中的以下部分：
1. 在准备模型时启用梯度检查点：
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
2. 在TrainingArguments中：
gradient_checkpointing=True,
per_device_train_batch_size=8, # 原来是16
gradient_accumulation_steps=2, # 原来是1
同时，为了确保安全，我们可以先尝试更小的设置，比如batch_size=4，accumulation_steps=4，然后再逐步增加。
但是，考虑到我们之前已经OOM，而且显存只剩下600MB左右，我们需要更激进的措施。
另外，我们还可以考虑优化数据加载，比如设置pin_memory=True，但这不是主要问题。
我们按照以下步骤修改代码：
步骤1：在准备模型时启用梯度检查点。
步骤2：调整训练参数，降低批次大小，增加梯度累积步数。
步骤3：如果仍然OOM，则考虑降低LoRA的rank。
由于错误发生在训练开始之前（0%），可能是前向传播的初始批次就OOM，所以我们先降低批次大小。
我们尝试将批次大小设为4，梯度累积步数为4，这样有效批次大小是16（与原来相同），同时启用梯度检查点。
修改后的TrainingArguments部分：
training_args = TrainingArguments(
...,
per_device_train_batch_size=4, # 从16降到4
gradient_accumulation_steps=4, # 从1增加到4
gradient_checkpointing=True, # 原来是False，现在改为True
...
)
同时，在准备模型时启用梯度检查点：
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
另外，我们还可以尝试使用更小的序列长度（比如1024）进行测试，但为了不改变模型能力，我们先不改变。
我们进行这些修改，然后重新运行。
如果仍然出现OOM，我们可以进一步降低批次大小（比如2）或者降低序列长度。
另外，我们注意到在训练前验证中，我们使用了两个样本进行测试，这个测试批次也可能很大。我们可以将测试批次改为一个样本，但这不是主要问题。
我们进行修改，然后重新运行。
注意：梯度检查点开启后，训练速度会变慢，但可以节省显存。
此外，我们还可以考虑使用`batch_size=1`和更大的梯度累积步数，但这样会导致优化不稳定，所以尽量保持有效批次大小不变。
修改后的代码部分如下：</think>您的代码出现了内存不足错误（`torch.OutOfMemoryError`），这在使用大型语言模型训练时很常见。以下是优化后的解决方案，专注于内存管理和训练效率的提升：
```python
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日志已关闭")
```
### 主要内存优化策略：
1. **内存配置优化**：
```python
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # 减少内存碎片
torch.backends.cuda.cufft_plan_cache.clear() # 清理CUDA缓存
```
2. **模型加载优化**：
```python
model = AutoModelForCausalLM.from_pretrained(
...,
attn_implementation="flash_attention_2" # 使用FlashAttention优化内存
)
```
3. **LoRA配置优化**：
```python
lora_config = LoraConfig(
r=64, # 降低rank以减少内存使用
lora_alpha=32, # 降低alpha值
target_modules=["q_proj", "v_proj"], # 减少目标模块
)
```
4. **梯度检查点启用**：
```python
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
training_args = TrainingArguments(..., gradient_checkpointing=True)
```
5. **批次大小优化**：
```python
training_args = TrainingArguments(
per_device_train_batch_size=4, # 大幅降低批次大小
gradient_accumulation_steps=4, # 增加梯度累积步数
)
```
6. **动态填充和序列长度优化**：
```python
def tokenize_function(examples):
tokenized = tokenizer(
examples["text"],
padding=True, # 使用动态填充
truncation=True,
max_length=1024, # 降低上下文长度
)
```
7. **分块训练策略**：
```python
num_samples = len(tokenized_dataset)
chunk_size = 1000
for i in range(0, num_samples, chunk_size):
chunk_dataset = tokenized_dataset.select(range(i, end_idx))
trainer.train_dataset = chunk_dataset
trainer.train()
clear_memory() # 清理内存
```
8. **内存清理函数**：
```python
def clear_memory():
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
```
### 内存优化原理：
1. **梯度检查点**：通过牺牲计算时间换取内存，只存储部分激活值而非全部
2. **FlashAttention**：使用优化的注意力机制实现，减少内存占用
3. **动态填充**：避免对所有样本填充到最大长度，减少不必要的内存分配
4. **分块训练**：将大数据集分成小块处理，避免一次性加载全部数据
5. **LoRA优化**：减少可训练参数数量和目标模块数量
6. **梯度累积**：使用小批次但累积梯度，保持有效批次大小
7. **内存碎片管理**：通过环境变量配置减少内存碎片
这些优化措施综合使用，可以显著降低训练过程中的内存峰值，解决CUDA内存不足的问题。
本文介绍了一种在Android应用中处理Out of Memory (OOM) 错误的方法，特别是在加载大型位图图像时出现的问题。文章详细展示了如何通过调整图像大小来避免内存溢出，并提供了具体的代码实现。
扫一扫
944
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
