重要提示

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

混合精度训练#

混合精度训练通过以低精度格式执行操作，同时选择性地在单精度下维护最少量数据，以在网络的关键区域中保留关键信息，从而显着提高计算效率。NeMo 框架现在通过 Transformer Engine (TE) 在大多数模型中支持 FP16、BF16 和 FP8。

半精度训练#

NeMo 框架通过 Megatron Core 和分布式优化器支持半精度 FP16 和 BF16 计算训练。此训练方案在所有层计算中使用半精度，同时将模型状态（优化器状态和主参数）保持为单精度。为了避免在每个层计算中重复数据类型转换，Megatron Core 保留了半精度参数的单独副本，该副本在每个优化器步骤后更新。

当将训练器的 plugins 设置为 fp16-mixed 或 bf16-mixed 之一时，启用半精度训练。参数梯度以相同的半精度计算，并且跨数据并行 GPU 的梯度 reduce-scatter 的精度根据训练器的精度自动设置。

实现半精度训练#

import nemo_run as run

from nemo import lightning as nl
from nemo.collections.llm.recipes.precision.mixed_precision import bf16_mixed, fp16_mixed

trainer_args = {TRAINER_ARGS}

# Set up trainer with bf16 precision
trainer_bf16 = run.Config(
  nl.Trainer,
  plugins=bf16_mixed(),
  **trainer_args,
)

# Set up trainer with fp16 precision
trainer_fp16 = run.Config(
  nl.Trainer,
  plugins=fp16_mixed(),
  **trainer_args,
)

也可以为特定方案更改精度

from functools import partial

from nemo.collections import llm
from nemo.collections.llm.recipes.precision.mixed_precision import bf16_mixed, fp16_mixed

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

# Change precision
recipe.trainer.plugins = fp16_mixed()

FP8 训练#

NVIDIA H100 GPU 引入了对新数据类型 FP8（8 位浮点）的支持，从而提高了矩阵乘法和卷积的吞吐量。NeMo 框架使用 NVIDIA TransformerEngine (TE) 来利用 FP8 的加速。下表总结了可以在 NeMo 中配置的与 FP8 相关的参数（示例配置设置）。有关更详细的概述，请参阅 TE 文档，特别是 FP8 格式 和 配方。

FP8 参数#

参数

描述

fp8

FP8 的训练方案格式可以设置为 ‘hybrid’ 或 ‘e4m3’，默认为 ‘hybrid’。在 ‘hybrid’ 格式中，激活和权重张量使用 E4M3 格式，而梯度使用 E5M2 格式，以满足后向张量的额外动态范围要求。

fp8_margin

FP8 张量的缩放因子可以使用此参数移动 $2 ^ {margin}$ 倍。

fp8_amax_history_len

amax 历史记录的窗口大小。窗口大小决定了每个张量存储的最近绝对最大值 (amaxes) 的实例数量。

fp8_amax_compute_algo

“max” 和 “most_recent” 之间的选择指定了如何从给定的历史记录中选择 amax 值。

fp8_params

指示是否在 FP8 中存储模块级参数。启用此选项可以减少内存消耗，因为它消除了在更高精度下存储权重副本的需要，这种情况发生在这些权重在外部维护时，例如优化器中的主参数。有关更多信息，请参阅 TE 中的 fp8_model_init API。

实现 FP8 训练#

import nemo_run as run

from nemo import lightning as nl
from nemo.collections.llm.recipes.precision.mixed_precision import bf16_with_fp8_mixed, fp16_with_fp8_mixed

trainer_args = {TRAINER_ARGS}
fp8_args = {FP8_ARGS}

# Set up trainer with bf16 & fp8 precision
trainer_bf16_fp8 = run.Config(
  nl.Trainer,
  plugins=bf16_with_fp8_mixed(),
  **trainer_args,
  **fp8_args,
)

# Set up trainer with fp16 & fp8 precision
trainer_fp16_fp8 = run.Config(
  nl.Trainer,
  plugins=fp16_with_fp8_mixed(),
  **trainer_args,
  **fp8_args,
)

资源#

本页内容