重要

您正在查看 NeMo 2.0 文档。此版本引入了 API 的重大更改和一个新库 NeMo Run。我们目前正在将所有功能从 NeMo 1.0 移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档,请参阅 NeMo 24.07 文档

批处理#

批大小是您应该调整的首批参数之一。为了效率和收敛,我们建议首先最大化每个 GPU 的批大小,以充分利用您的 GPU RAM。

NeMo 框架使用以下参数

参数

描述

微批大小

每个数据并行 rank 的示例数量。

全局批大小

全局批大小计算公式为:global batch size = micro_batch_size * data_parallel_size * gradient_accumulation_steps。有关 data_parallel_size 的详细信息,请参阅此页面

梯度累积

此参数支持使用大批量大小进行训练,同时保持固定的内存占用,但这需要额外的计算。accumulate_grad_batches 由 PyTorch Lightning 自动管理。

设置批处理参数#

以下示例展示了如何为 LLaMA-3 8B 模型设置预训练配方和批处理参数

from nemo.collections import llm
from functools import partial

# Load train recipe
recipe = partial(llm.llama3_8b.pretrain_recipe)()

# Set micro and global batch size
recipe.data.micro_batch_size = 4
recipe.data.global_batch_size = 16

# Set accumulate_grad_batches
recipe.trainer.accumulate_grad_batches = 1

直接从 CLI 设置批处理参数

nemo llm pretrain --factory llama3_8b data.micro_batch_size=4 data.global_batch_size=16 trainer.accumulate_grad_batches=1