QLoRA训练优化：解决大批次尺寸耗时问题

在使用 QLoRA (Quantization-aware Low-Rank Adaptation) 技术微调大型语言模型时，可能会遇到一些意想不到的问题。其中一个常见问题是，当增加 per_device_train_batch_size 时，训练时间会不成比例地增加，即使 GPU 内存可以容纳更大的批量尺寸。 本文将深入探讨这个问题，并提供可能的解决方案。

如摘要所述，问题通常在于训练步数 (max_steps) 和 epoch 之间的关系。 让我们更详细地了解这一点。

理解 max_steps 和 Epoch

在训练机器学习模型时，max_steps 和 epoch 是两个重要的参数，它们决定了训练过程的持续时间。

• Epoch: 一个 epoch 表示模型训练数据集的完整一次迭代。 例如，如果你的训练数据集包含 1000 个样本，并且你设置 num_epochs=2，那么模型将遍历整个数据集两次。

• max_steps: max_steps 定义了训练过程中的最大更新步数。 这是一种更精细的控制训练过程的方式，尤其是在你希望限制训练时间或在特定步数后停止训练的情况下。

在 transformers 库中，如果你同时指定了 num_epochs 和 max_steps，那么 max_steps 将覆盖 num_epochs。 如果仅指定 num_epochs，则训练将持续到所有 epoch 完成。

问题分析

当增加 per_device_train_batch_size 时，每个 epoch 的迭代次数会减少。 这是因为模型在每个步骤中处理更多的数据。 如果 max_steps 的值保持不变，那么实际上训练的 epoch 数会减少，导致模型训练不足。

例如，假设你的训练数据集包含 10000 个样本，并且你设置了 max_steps=1000。

• 如果 per_device_train_batch_size=1，那么每个 epoch 将包含 10000 步，因此训练将持续 0.1 个 epoch (1000 / 10000)。
• 如果 per_device_train_batch_size=100，那么每个 epoch 将包含 100 步，因此训练将持续 10 个 epoch (1000 / 100)。

在这种情况下，当 per_device_train_batch_size 从 1 增加到 100 时，训练的 epoch 数从 0.1 增加到 10。 这意味着模型实际上训练了更多次，从而导致训练时间显着增加。

解决方案

要解决这个问题，你需要确保 max_steps 的值与预期的训练 epoch 数相匹配。

1. 确定目标 Epoch 数： 首先，确定你希望模型训练多少个 epoch。 这取决于你的数据集大小、模型复杂性和训练目标。

2. 计算所需的 max_steps： 使用以下公式计算所需的 max_steps 值：

   max_steps = (num_samples / per_device_train_batch_size) * num_epochs

   其中：

   • num_samples 是训练数据集中的样本数量。
   • per_device_train_batch_size 是每个设备的训练批量大小。
   • num_epochs 是你希望模型训练的 epoch 数。

3. 更新 TrainingArguments： 在你的 TrainingArguments 中，将 max_steps 设置为计算出的值。

   training_args = TrainingArguments(
       output_dir=config['output_dir'],
       per_device_train_batch_size=config['per_device_train_batch_size'],
       gradient_accumulation_steps=config['gradient_accumulation_steps'],
       learning_rate=float(config['learning_rate']),
       max_steps=calculated_max_steps, # 使用计算出的 max_steps
       optim="paged_adamw_8bit",
       fp16=True,
       load_best_model_at_end = True,
       save_strategy="epoch",  # Save at the end of each epoch
       evaluation_strategy="epoch",
       save_total_limit=1  # Keep only the last 2 checkpoints
   )

示例代码

假设你的训练数据集包含 10000 个样本，你希望模型训练 3 个 epoch，并且你使用 per_device_train_batch_size=128。 那么，你需要将 max_steps 设置为：

num_samples = 10000
per_device_train_batch_size = 128
num_epochs = 3

calculated_max_steps = (num_samples / per_device_train_batch_size) * num_epochs
print(f"Calculated max_steps: {calculated_max_steps}") # 输出: Calculated max_steps: 234.375

# 由于 max_steps 必须是整数，通常向上取整
calculated_max_steps = int(calculated_max_steps + 0.5) # 四舍五入
print(f"Rounded max_steps: {calculated_max_steps}") # 输出: Rounded max_steps: 234

training_args = TrainingArguments(
    output_dir=config['output_dir'],
    per_device_train_batch_size=config['per_device_train_batch_size'],
    gradient_accumulation_steps=config['gradient_accumulation_steps'],
    learning_rate=float(config['learning_rate']),
    max_steps=calculated_max_steps,
    optim="paged_adamw_8bit",
    fp16=True,
    load_best_model_at_end = True,
    save_strategy="epoch",  # Save at the end of each epoch
    evaluation_strategy="epoch",
    save_total_limit=1  # Keep only the last 2 checkpoints
)

注意事项

• 确保 max_steps 是一个整数。 如果计算出的 max_steps 不是整数，请将其四舍五入到最接近的整数。
• 监控训练过程，并根据需要调整 max_steps 的值。
• 考虑使用验证集来评估模型的性能，并防止过度拟合。

总结

当使用 QLoRA 对大型语言模型进行微调时，max_steps 的设置至关重要。 通过确保 max_steps 的值与预期的训练 epoch 数相匹配，你可以避免训练时间过长的问题，并确保模型得到充分的训练。 通过本文提供的分析和解决方案，你可以更好地理解和解决在使用 QLoRA 时遇到的训练时间问题。 记住，细致地调整训练参数是获得最佳模型性能的关键。