使用Qlora训练open_llama_7b_v2：大批次尺寸致训练时长激增咨询

问题描述 #

使用QLoRA方法训练openlm-research/open_llama_7b_v2模型时，发现增大per_device_train_batch_size会大幅延长训练时长：

• 当per_device_train_batch_size=1时，单GPU显存占用3GB，训练耗时4小时
• 当per_device_train_batch_size=128时，单GPU显存占用8GB（无显存溢出），训练耗时增至40小时

配置信息 #

Model: openlm-research/open_llama_7b_v2
Method: Qlora
Batch size: 1 (3GB GPU, 4hrs) vs. 128 (8GB GPU, 40hrs)
gradient_accumulation_steps: 未设置
GPU: 4X A100 40GB
Transformers version: 4.30.0.dev0

代码实现 #

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(
    config['model_id'],
    quantization_config=bnb_config,
    device_map=config.get('device_map', 'auto')
)

# Initialize tokenizer
tokenizer = LlamaTokenizer.from_pretrained(config['model_id'])
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
model.resize_token_embeddings(len(tokenizer))

# Define LoRA configuration
qlora_config = LoraConfig(
    r=config['lora_r'],#16,
    lora_alpha=config['lora_alpha'],#32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# Initialize SFTTrainer
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=config['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
)

callbacks = [EarlyStoppingCallback(early_stopping_patience=2)]

supervised_finetuning_trainer = SFTTrainer(
    model,
    train_dataset=train_dataset,
    args=training_args,
    tokenizer=tokenizer,
    peft_config=qlora_config,
    dataset_text_field="text",
    max_seq_length=config['max_seq_length'],
    # data_collator=DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True),
    callbacks=callbacks
)

# Train the model
supervised_finetuning_trainer.train()

原因分析 #

1. 总训练样本量大幅增加：若max_steps保持固定，增大per_device_train_batch_size会直接导致总处理的样本量变为原来的128倍（max_steps * per_device_train_batch_size * GPU数量），训练计算量剧增，时长自然延长。
2. 梯度累积未匹配调整：配置中未设置gradient_accumulation_steps，无法通过梯度累积抵消批次增大带来的梯度更新频率变化，导致每步梯度更新的样本量远大于原设置，总迭代次数不变的情况下，总计算量翻倍。
3. 数据加载瓶颈凸显：大批次尺寸下，数据预处理、加载到GPU的开销上升，若未启用多进程数据加载，会导致GPU频繁等待数据，利用率下降，拉长训练时间。
4. 8bit优化器额外开销：使用paged_adamw_8bit优化器时，大批次下的内存管理和量化计算会产生额外开销，相比小批次，每步的优化器计算耗时增加。
5. 多GPU通信占比上升：4卡A100模型并行时，大批次会增加GPU间的数据通信量，通信耗时占比上升，整体训练效率下降。

解决建议 #

1. 保持等效梯度更新批次：调整gradient_accumulation_steps，确保per_device_train_batch_size * gradient_accumulation_steps的乘积与原设置（1*原梯度步数）一致，比如原梯度步数为128，批次调至128后将梯度步数设为1，这样每步梯度更新的等效样本量不变，总计算量与原训练相当。
2. 调整max_steps匹配总样本量：若目标是训练相同的总样本数，将max_steps设为原数值的1/128，避免因批次增大导致总训练数据量翻倍。
3. 优化数据加载：在TrainingArguments中设置dataloader_num_workers为CPU核心数（如8或16），启用多进程数据加载，减少GPU等待时间。
4. 更换优化器测试：尝试换成adamw_bnb_8bit或常规adamw（显存允许的话），对比大批次下的训练速度，排查是否为8bit优化器的开销问题。
5. 定位性能瓶颈：用nvidia-smi或PyTorch Profiler监控训练时的GPU利用率、数据加载耗时、通信耗时，明确瓶颈所在后针对性优化。

内容的提问来源于stack exchange，提问作者JulesR