【大模型训练】forward_backward_func返回多个micro batch 损失

您询问的是 self.forward_backward_func 是否在一次执行中处理多个 microbatch。

答案

是的,self.forward_backward_func 在一次执行中处理多个 microbatch 。

详细说明

在 MegatronTrainStrategy.train_step 中,forward_backward_func 被调用一次,但内部会循环处理 num_microbatches 个 microbatch:[link to Repo alibaba/ROLL: roll/distributed/strategy/megatron_strategy.py:601-609]

关键参数:

num_microbatches : 指定要处理的 microbatch 数量,等于 gradient_accumulation_steps[link to Repo alibaba/ROLL: roll/distributed/strategy/megatron_strategy.py:593-596]
data_iterator : 提供所有 microbatch 的数据

内部执行流程

forward_backward_func 是 Megatron-Core 提供的函数,它的内部逻辑是:

循环 num_microbatches 次
每次迭代 :
- 从 data_iterator 获取一个 microbatch
- 调用 forward_step_func (即 inner_forward_step) 执行前向传播
- 执行反向传播,梯度累积到参数的 .grad 中
- 收集该 microbatch 的 metrics
返回所有 microbatch 的 metrics 列表

在 SFT Trainer 中的类似实现

在 McaTrainer.training_step 中也是相同的模式:[link to Repo alibaba/ROLL: mcore_adapter/src/mcore_adapter/trainer/trainer.py:414-422]

同样是调用一次 forward_backward_func,内部处理 gradient_accumulation_steps 个 microbatch。

执行流程图

Yes No train_step 调用一次 forward_backward_func 执行一次内部循环: i=0 to num_microbatches-1 获取 microbatch i forward_step_func backward pass
梯度累积收集 metrics_i i < num_microbatches-1? 返回 metrics_tensors 列表 optimizer.step 一次性更新

Notes

forward_backward_func 只被调用一次 ,但内部会处理多个 microbatch
返回的 metrics_tensors 是一个列表,长度等于 num_microbatches,每个元素对应一个 microbatch 的 metrics[link to Repo alibaba/ROLL: roll/distributed/strategy/megatron_strategy.py:601-609]
所有 microbatch 的梯度累积完成后,才调用一次 optimizer.step() 更新参数[link to Repo alibaba/ROLL: roll/distributed/strategy/megatron_strategy.py:613-619]

Wiki pages you might want to explore:

Strategy Pattern and Backend Abstraction (alibaba/ROLL)

Wiki pages you might want to explore:

Strategy Pattern and Backend Abstraction (alibaba/ROLL)您询问的是 train_step 是否指用一个 batch 更新一次参数。

答案

是的,train_step 指的是用一个 batch 更新一次参数 。

详细说明

在 MegatronTrainStrategy.train_step 中,一个完整的训练步骤包括:[link to Repo alibaba/ROLL: roll/distributed/strategy/megatron_strategy.py:582-643]

处理一个 batch 的所有 microbatch :通过 forward_backward_func 处理 num_microbatches 个 microbatch,梯度在这些 microbatch 间累积[link to Repo alibaba/ROLL: roll/distributed/strategy/megatron_strategy.py:601-609]
一次性更新参数 :所有 microbatch 处理完后,调用一次 optimizer.step() 更新参数[link to Repo alibaba/ROLL: roll/distributed/strategy/megatron_strategy.py:613-619]
更新学习率 :调用 scheduler.step() 更新学习率[link to Repo alibaba/ROLL: roll/distributed/strategy/megatron_strategy.py:618-618]
清空梯度 :为下一个训练步骤准备[link to Repo alibaba/ROLL: roll/distributed/strategy/megatron_strategy.py:622-624]

Batch 的组成

这里的"一个 batch"实际上是:

Global batch: 所有数据并行 rank 上的数据总和
Local batch : 当前 rank 上的数据,会被分成 gradient_accumulation_steps 个 microbatch

全局 batch size 的计算公式:[link to Repo alibaba/ROLL: docs_roll/docs/QuickStart/qa_issues.md:69-73]

复制代码

global_batch_size = per_device_train_batch_size * gradient_accumulation_steps * world_size / 
                    tensor_model_parallel_size / pipeline_model_parallel_size / context_parallel_size

Notes

一个 train_step = 一次参数更新 = 处理一个 global batch
虽然内部处理了多个 microbatch,但这些 microbatch 的梯度会累积,最后只更新一次参数
这种设计允许使用更大的有效 batch size,同时保持 GPU 内存在可控范围内

Wiki pages you might want to explore:

Strategy Pattern and Backend Abstraction (alibaba/ROLL)

Wiki pages you might want to explore:

Strategy Pattern and Backend Abstraction (alibaba/ROLL)