使用 Transformer Engine 加速 Hugging Face Llama 2 和 Llama 3 模型

目标

本教程展示了如何通过使用来自 Transformer Engine 库 的 TransformerLayer 在 BF16 和 FP8 精度下加速微调 Hugging Face 的完整 Llama 2 或 Llama 3 模型。

本教程的依赖项

有效运行本教程需要以下文件和媒体

1. te_llama.py

   • 此文件包含将 Hugging Face Llama 2 或 Llama 3 检查点加载到 Transformer Engine 的 TransformerLayer 而不是 Hugging Face 的 LlamaDecoderLayer 的代码。这用于本教程的以下两个部分 - “改进 1”和“改进 2”。

2. utils.py

   • 此文件包含与数据加载、超参数、设置模型/优化器/加速器、模型训练以及其他杂项任务（如从单元格内重启 jupyter notebook）相关的代码。

3. media/

   • 此目录包含以下教程中使用的图像。

运行本教程需要以下软件包：pytorch、transformer_engine、accelerate、transformers、peft、datasets。

关于使用 Llama 3 权重运行教程的说明

本教程显示了使用 Llama 2 7B 权重运行时单元格的输出。它可以通过简单地提供包含这些权重（Hugging Face 格式）的目录而不是 Llama 2 7B 权重来使用 Llama 3 8B 权重运行。这两个模型几乎相同，最大的区别在于模型维度（最小的 Llama 3 模型有 8B 参数，而最小的 Llama 2 模型有 7B），这使得本教程可以用于这两个模型。

目录

1. 从“Transformer”到“Llama”

2. Hugging Face 的 LlamaModel

   • Hugging Face 的 LlamaDecoderLayer

3. [基线] 运行 HF LlamaModel（精度：BF16）

4. [改进 1] 将 HF 的 LlamaDecoderLayer 替换为 TE 的 TransformerLayer（精度：BF16）

   • Transformer Engine 的 TransformerLayer

   • TransformerLayer 选项说明

   • 将权重从 HF 的 LlamaDecoderLayer 映射到 TE 的 TransformerLayer

5. [改进 2] 将 HF 的 LlamaDecoderLayer 替换为 TE 的 TransformerLayer（精度：FP8）

6. 结论

从“Transformer”到“Llama”

图 1：Llama 可视化为 Transformer。（使用 Nvidia 的 AI 基础模型 生成）

回顾

• 2017 年：“Attention Is All You Need” 论文介绍了开创性的“Transformer”架构，并永远改变了 NLP 领域。

• 2018-2020 年：GPT 模型系列的出现，表明因果解码器架构非常适合预训练、少样本和零样本学习。

• 快进到 2023-2024 年：继 GPT-3/GPT-4 成功案例之后，研究人员和公司竞相生产下一个最佳预训练模型，该模型可以进一步微调以用于特定应用场景。

• 2023 年 2 月：Meta 发布 Llama 2 模型（大型语言模型 Meta AI）。

  • 这些模型的参数范围从 7B 到 70B。

  • LLaMA 2 在 2 万亿个 tokens 上进行了预训练。

• 2024 年 4 月：Meta 发布 Llama 3 模型。

  • 这些模型的参数范围从 8B 到 70B。

  • LLaMA 3 在 15 万亿个 tokens 上进行了预训练。

有关 Llama 2 的更多信息，请考虑阅读 Huggingface 教程。 简而言之，以下是传统 Transformer 解码器架构与 Llama 2 架构之间的一些重要区别

1. 仅解码器模型（因果语言建模和下一个词预测）

2. 使用 RMSNorm 代替 LayerNorm

3. SwiGLU 激活函数

4. RoPE 作为位置嵌入

5. 70B 模型的 Grouped Query Attention

6. 在 4K 上下文长度上训练

Hugging Face 还发布了关于 Llama 3 的教程。 关键点是

1. 使用更大的 tokenizer - 128256 vs 32K。

2. 较小的 8B 模型也使用 Grouped Query Attention。

3. 所有模型的上下文长度都增加到 8K。

4. Llama 3 的训练数据量是 Llama 2 的 8 倍。

图 2：GPT 和 Llama 架构的比较。

Hugging Face 的 LlamaModel

Hugging Face 在 modeling_llama.py 中提供了 Llama 模型的开源实现。

这是一个框图，显示了 Llama 模型如何在 Hugging Face repo 中实现。 请注意模块化的封装形式以及模型实现核心的 LlamaDecoderLayer。

图 3：因果 Llama 模型框图。

上面的图表转换为 PyTorch 中模型的以下文本输出。 请注意，模型的核心有 32 个 LlamaDecoderLayer。

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096, padding_idx=0)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaFlashAttention2(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)

Hugging Face 的 LlamaDecoderLayer

让我们仔细看看 LlamaDecoderLayer。 它由 input_layernorm、self_attn、post_attention_layernorm 和 mlp 模块组成。 每个模块都有相关的权重，如图所示。

图 4：因果 Llama 模型框图（简化说明了 LlamaDecoderLayer）。

Self_Attn 层

为了简化“self_attn”框的框图说明，我们省略了“Grouped Query Attention”操作，仅展示了具有相关权重的模块。

MLP 层

SwiGLU 是在 Hugging Face github repo 的 modeling_llama.py 文件中定义的激活函数，如下所示

"""
1. `self.up_proj`, `self.gate_proj` and `self.down_proj` are "Linear" layers
2. `self.act_fn` is a "Swish" function

"""
down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

与传统“MLP”层（例如传统 Transformer 或 GPT 架构）中的 2 个权重相比，它需要一组 3 个权重。 这也在下图中说明

图 5：带有 swiglu 激活函数的前馈层内部结构。

[基线] 运行 HF LlamaModel（精度：BF16）

Llama 2 权重被加载到 Hugging Face 原生实现 LlamaForCausalLM 中（请参阅 modeling_llama.py）。

对于本次和其他后续运行，batch_size 为 8。 基线中的 LlamaDecoderLayer 保持不变，如下所示

图 6：重新审视“LlamaDecoderLayer”。

注意

基线实现将在 BF16 精度下运行。

注意

本教程加载并训练 Llama 3 8B 或 Llama 2 7B 模型，这会占用大部分 GPU 内存，因此，我们需要在运行以下部分之前每次都重启 jupyter notebook。 随附的 utils.py 文件中定义了一个小型实用程序方法 restart_jupyter_notebook。 此函数重启 jupyter notebook，以便在再次从检查点加载模型之前刷新 GPU 内存，以避免遇到 OOM（内存不足）错误。

如果该实用程序不起作用，请注释掉以下单元格中的此行 restart_jupyter_notebook()，并在运行单元格之前手动重启 jupyter notebook。 对于本教程中的其他部分，请重复相同的操作。

# Restart the notebook (to flush the GPU memory)
from utils import restart_jupyter_notebook
restart_jupyter_notebook()


# Import necessary packages, methods and variables
from utils import *


# Provide Huggingface Access Token
hyperparams.hf_access_token = ""
assert hyperparams.hf_access_token, "Provide a HF API Access Token!"

# Provide a directory to cache weights in to avoid downloading them every time.
# (By default, weights are cached in `~/.cache/huggingface/hub/models`)
hyperparams.weights_cache_dir = ""

# For Llama 2, uncomment this line (also set by default)
hyperparams.model_name = "meta-llama/Llama-2-7b-hf"

# For Llama 3, uncomment this line
# hyperparams.model_name = "meta-llama/Meta-Llama-3-8B"

hyperparams.mixed_precision = "bf16"


# Init the model and accelerator wrapper
model = init_baseline_model(hyperparams)
accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator(model, hyperparams)


# Finetune the model
finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)

10 finetuning steps complete!
Average time taken per step: 248 milliseconds

让我们在表格中添加此信息，并将其与未来部分中可能的一些改进进行比较

模型	精度	步长时间（或每批次毫秒数）	加速比（相对于基线）
HF（基线）	BF16	248	1

[改进 1] 将 HF 的 LlamaDecoderLayer 替换为 TE 的 TransformerLayer（精度：BF16）

除了 Linear 和 LayerNorm 等基本层之外，Transformer Engine 还提供更大的模块，如 MultiheadAttention（结合了“LayerNorm”和“Self Attention”）和 LayerNormMLP（结合了“LayerNorm”和“MLP”），它们可以替换 LlamaDecoderLayer 中的对应模块，并可能提供加速。 Transformer Engine 还提供完整的 TransformerLayer（进一步结合了 MultiheadAttention 和 LayerNormMLP 层），它可以替换 LlamaDecoderLayer 并提供加速（由于权重名称对于这两个层不同，因此需要仔细映射权重）。 让我们仔细看看 Transformer Engine 的 TransformerLayer。

Transformer Engine 的 TransformerLayer

在更高的层面上，TE 的 TransformerLayer 可以被视为 LlamaDecoderLayer 的合适替代品。 但 TransformerLayer 的内部结构组织方式略有不同。

图 7：Transformer Engine 的 TransformerLayer

就像 Hugging Face 的 LlamaDecoderLayer 一样，Transformer Engine 的 TransformerLayer 封装了 self_attention（作为 MultiheadAttention）和 mlp（作为 LayerNormMLP）。 主要区别在于，两个 Norm 都包含在 MultiheadAttention 和 LayerNormMLP 层中，如下面的输出提示所示

TransformerLayer(
    (self_attention): MultiheadAttention(
      (layernorm_qkv): LayerNormLinear()
      (core_attention): DotProductAttention()
      (proj): Linear()
    )
    (layernorm_mlp): LayerNormMLP()
)

另一个区别是，Transformer Engine 实现了带有 SwiGLU 的前馈层的有效版本，其中 up_proj 和 gate_proj 模块的权重合并在一起，并使用自定义融合内核应用 SwiGLU。 这样做是为了仅向 GPU 发出一个大型高效的矩阵乘法运算，而不是两个较小的运算。

图 8：Transformer Engine 中 SwiGLU 实现的抽象图示。

TransformerLayer 选项说明

注意

在这里，我们将介绍本教程中需要的 TransformerLayer 中的一些选项。 有关选项的完整列表，请参阅 TransformerLayer API 文档。

在随附的 te_llama.py 文件中，TELlamaDecoderLayer 被定义为 TE 的 TransformerLayer 的包装器，其中包含一些必要的选项，使 TransformerLayer 成为 HF 的 LlamaDecoderLayer 的插件替代品。

class TELlamaDecoderLayer(te.pytorch.TransformerLayer):
    def __init__(self, config):
        super().__init__(
            config.hidden_size,
            config.intermediate_size,
            config.num_attention_heads,
            bias=False,
            layernorm_epsilon=config.rms_norm_eps,
            hidden_dropout=0,
            attention_dropout=0,
            fuse_qkv_params=False,
            normalization="RMSNorm",
            activation="swiglu",
            attn_input_format="bshd",
            num_gqa_groups=config.num_key_value_heads,
        )
        te_rope = RotaryPositionEmbedding(config.hidden_size//config.num_attention_heads)
        self.te_rope_emb = te_rope(max_seq_len=config.max_position_embeddings).cuda()

以下列表简要总结了每个选项

1. hidden_size：每个输入样本的大小。

2. ffn_hidden_size：样本投影到的中间大小。

3. num_attention_heads：Transformer 层中注意力头的数量。

4. bias：切换以向子模块层添加加性偏差。

5. layernorm_epsilon：添加到层归一化分母的值，以提高数值稳定性。 默认为 1e-5。

6. hidden_dropout：FC2 层（全连接层 2 号）之后 dropout 操作的 dropout 概率。 默认为 0.1。

7. attention_dropout：多头注意力期间 dropout 操作的 dropout 概率。 默认为 0.1。

8. fuse_qkv_params：如果设置为 True，则 TransformerLayer 模块为 query-key-value 公开单个融合参数。 这可以实现优化，例如无需连接/拆分的 QKV 融合，并且还启用了参数 fuse_wgrad_accumulation。

9. normalization：应用的归一化类型。 默认为 LayerNorm。

10. activation：MLP 块中使用的激活类型。 默认为 gelu。

11. attn_input_format：控制中间隐藏状态的维度是“批优先”（“bshd”）还是“序列优先”（“sbhd”）。 s 代表序列长度，b 批大小，h 头数，d 头大小。 请注意，这些格式与 MultiHeadAttention 和 DotProductAttention 模块中的 qkv_format 非常相关。

12. num_gqa_groups：Transformer 层中 GQA 组的数量。 本文 中描述了 Grouped Query Attention。 这仅影响键和值，而不影响查询。 GQA-1 等效于 Multi-Query Attention (MQA)，而 GQA-H 等效于 MultiHead Attention，即 num_gqa_groups = num_attention_heads。

此外，请注意，RotaryPositionEmbedding 被定义为 TELlamaDecoderLayer（TE 的 TransformerLayer 的包装器）本身的一部分，因为它期望如果模型中使用 RoPE，则使用此 rope 缓存。

让我们回顾一下 LlamaDecoderLayer 如何在 HF 的 llama 实现中形成解码器层堆栈的核心

ModuleList(
  (0-31): 32 x LlamaDecoderLayer(
    (self_attn): LlamaAttention(
      (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
      (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
      (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
      (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
      (rotary_emb): LlamaRotaryEmbedding()
    )
    (mlp): LlamaMLP(
      (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
      (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
      (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
      (act_fn): SiLU()
    )
    (input_layernorm): LlamaRMSNorm()
    (post_attention_layernorm): LlamaRMSNorm()
  )
)

Hugging Face 模型实现的大部分（32 个 LlamaDecoderLayer 层）可能会被 Transformer Engine 的 TransformerLayer 层替换。 让我们看看如何实现这一点。

将权重从 HF 的 LlamaDecoderLayer 映射到 TE 的 TransformerLayer

请参阅随附文件 te_llama.py，该文件提供了在使用 TE 的 TransformerLayer 替换 HF 的 LlamaDecoderLayer 后创建 Llama 2 模型的参考。

简而言之，以下代码片段被放在一起

1. TELlamaDecoderLayer 被添加为 TransformerLayer 的包装器。

class TELlamaDecoderLayer(te.pytorch.TransformerLayer):
    """
    Wrapper class over TE's `TransformerLayer`. This makes the wrapper very
    similar to HF's `LlamaDecoderLayer` and easier to replace it in the code.

    Args:
        config: LlamaConfig
        args: positional args (for compatibility with `LlamaDecoderLayer`)
        kwargs: keyword args (for compatibility with `LlamaDecoderLayer`)
    """
    def __init__(self, config, *args, **kwargs):
        super().__init__(
            hidden_size=config.hidden_size,
            ffn_hidden_size=config.intermediate_size,
            num_attention_heads=config.num_attention_heads,
            bias=False,
            layernorm_epsilon=config.rms_norm_eps,
            hidden_dropout=0,
            attention_dropout=0,
            fuse_qkv_params=False,
            normalization="RMSNorm",
            activation="swiglu",
            attn_input_format="bshd",
        )
        te_rope = RotaryPositionEmbedding(config.hidden_size//config.num_attention_heads)
        self.te_rope_emb = te_rope(max_seq_len=config.max_position_embeddings).cuda()

    def forward(self,
                hidden_states,
                *args,
                attention_mask,
                **kwargs):
        """
        Custom forward to make sure we only pass relevant arguments to the
        forward pass of the `TransformerLayer`. Also, make sure the output
        format matches the output of the HF's `LlamaDecoderLayer`.
        """
        return (super().forward(hidden_states, attention_mask=attention_mask, rotary_pos_emb=self.te_rope_emb),)

2. 在创建 LlamaForCausalLM 之前，使用 replace_decoder 上下文管理器将 LlamaDecoderLayer monkey-patch 为 TELlamaDecoderLayer。

@contextmanager
def replace_decoder(te_decoder_cls):
    """
    Replace `LlamaDecoderLayer` with custom `TELlamaDecoderLayer`.
    """
    original_llama_decoder_cls = transformers.models.llama.modeling_llama.LlamaDecoderLayer
    transformers.models.llama.modeling_llama.LlamaDecoderLayer = te_decoder_cls
    try:
        yield
    finally:
        transformers.models.llama.modeling_llama.LlamaDecoderLayer = original_llama_decoder_cls
.
.
.
class TELlamaForCausalLM:
    """
    Causal LM created with `LlamaModel`. The underlying `LlamaDecoderLayer`
    class is monkey-patched with `TELlamaDecoderLayer` class before
    initializing the causal LM with `LlamaForCausalLM`.

    Args:
        config: LlamaConfig
    """

    def __new__(cls, config: LlamaConfig):
        with replace_decoder(te_decoder_cls=TELlamaDecoderLayer):
            llama_for_causal_lm = LlamaForCausalLM(config)
        return llama_for_causal_lm
.
.
.

3. 添加了一个自定义的 pretrained_from_local 方法，通过仔细地将权重从 LlamaDecoderLayer (HF) 映射到 TransformerLayer (TE)，将权重从检查点（用于 HF Llama 实现）复制到修改后的 TELlamaForCausalLM。 方法 replace_params 将合适的权重从 LlamaDecoderLayer 映射并复制到 TransformerLayer。 有关更多详细信息，请参阅下图。

def replace_params(hf_state_dict, te_state_dict):
    # collect all layer prefixes to update
    all_layer_prefixes = set()
    for param_key in hf_state_dict.keys():
        layer_prefix_pat = 'model.layers.\d+.'
        m = re.match(layer_prefix_pat, param_key)
        if m is not None:
            all_layer_prefixes.add(m.group())

    for layer_prefix in all_layer_prefixes:
        # When loading weights into models with less number of layers, skip the
        # copy if the corresponding layer doesn't exist in TE model
        if layer_prefix + 'self_attention.layernorm_qkv.layer_norm_weight' in te_state_dict:
            te_state_dict[layer_prefix + 'self_attention.layernorm_qkv.layer_norm_weight'].data[:] = hf_state_dict[layer_prefix + 'input_layernorm.weight'].data[:]

        if layer_prefix + 'self_attention.layernorm_qkv.query_weight' in te_state_dict:
            te_state_dict[layer_prefix + 'self_attention.layernorm_qkv.query_weight'].data[:] = hf_state_dict[layer_prefix + 'self_attn.q_proj.weight'].data[:]

        if layer_prefix + 'self_attention.layernorm_qkv.key_weight' in te_state_dict:
            te_state_dict[layer_prefix + 'self_attention.layernorm_qkv.key_weight'].data[:] = hf_state_dict[layer_prefix + 'self_attn.k_proj.weight'].data[:]
    .
    .
    .

    return all_layer_prefixes

下图显示了权重如何从 HF 的 LlamaDecoderLayer 映射到 TE 的 TransformerLayer。

图 9：将 LlamaDecoderLayer 替换为 TransformerLayer。

以这种方式初始化修改后的 Llama 模型后，核心解码器层将更改为 TELlamaDecoderLayer（TransformerLayer 的包装器），如下面的输出所示

ModuleList(
  (0-31): 32 x TELlamaDecoderLayer(
    (self_attention): MultiheadAttention(
      (layernorm_qkv): LayerNormLinear()
      (core_attention): DotProductAttention(
        (flash_attention): FlashAttention()
        (fused_attention): FusedAttention()
        (unfused_attention): UnfusedDotProductAttention(
          (scale_mask_softmax): FusedScaleMaskSoftmax()
          (attention_dropout): Dropout(p=0, inplace=False)
        )
      )
      (proj): Linear()
    )
    (layernorm_mlp): LayerNormMLP()
  )
)

总而言之，模型按如下方式更改，其中大部分实现（核心解码器层）来自 Transformer Engine。

图 10：将 HF 的 LlamaDecoderLayer 替换为 TE 的 TransformerLayer 后的语言模型。

注意

让我们首先在 BF16 精度下运行这个 “TELlama” 实现。

# Restart the notebook (to flush the GPU memory)
from utils import restart_jupyter_notebook
restart_jupyter_notebook()


# Import necessary packages, methods and variables
from utils import *


# Provide Huggingface Access Token
hyperparams.hf_access_token = ""
assert hyperparams.hf_access_token, "Provide a HF API Access Token!"

# Provide a directory to cache weights in to avoid downloading them every time.
# (By default, weights are cached in `~/.cache/huggingface/hub/models`)
hyperparams.weights_cache_dir = ""

# For Llama 2, uncomment this line (also set by default)
hyperparams.model_name = "meta-llama/Llama-2-7b-hf"

# For Llama 3, uncomment this line
# hyperparams.model_name = "meta-llama/Meta-Llama-3-8B"

hyperparams.mixed_precision = "bf16"


# Init the model and accelerator wrapper
model = init_te_llama_model(hyperparams)
accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator(model, hyperparams)


# Finetune the model
finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)

10 finetuning steps complete!
Average time taken per step: 185 milliseconds

与“基线”实现相比，我们看到即使仅使用 BF16 精度，使用 Transformer Engine 的 TransformerLayer 代替 Huggging Face 的 LlamaDecoderLayer 也能获得 34% 的加速！

模型	精度	步长时间（或每批次毫秒数）	加速比（相对于基线）
HF（基线）	BF16	248	1
TE（将 LlamaDecoderLayer 替换为 TE.TransformerLayer）	BF16	185	1.34

[改进 2] 将 HF 的 LlamaDecoderLayer 替换为 TE 的 TransformerLayer (精度: FP8)

现在，大部分 HF Llama 模型实现 (LlamaDecoderLayers) 已被 Transformer Engine 实现 (TELlamaDecoderLayer 或 TransformerLayer) 替换，让我们看看在 FP8 精度下进行微调如何帮助提高性能。

如何在 FP8 精度下运行模型

替换后，可以通过对之前的 BF16 运行进行以下更改，在 FP8 精度下运行模型。（有关更多信息，请参阅随附 utils.py 文件中的相应 wrap_with_accelerator 函数）。

# Specify the `FP8RecipeKwargs` (additional argument required to run in `fp8` precision)
fp8_kwarg_handler = [FP8RecipeKwargs(backend="te")]

# Pass the `FP8RecipeKwargs` to the `Accelerator` init call
accelerator = Accelerator(
    ...
    kwargs_handlers=fp8_kwarg_handler
)

# Restart the notebook (to flush the GPU memory)
from utils import restart_jupyter_notebook
restart_jupyter_notebook()


# Import necessary packages, methods and variables
from utils import *


# Provide Huggingface Access Token
hyperparams.hf_access_token = ""
assert hyperparams.hf_access_token, "Provide a HF API Access Token!"

# Provide a directory to cache weights in to avoid downloading them every time.
# (By default, weights are cached in `~/.cache/huggingface/hub/models`)
hyperparams.weights_cache_dir = ""

# For Llama 2, uncomment this line (also set by default)
hyperparams.model_name = "meta-llama/Llama-2-7b-hf"

# For Llama 3, uncomment this line
# hyperparams.model_name = "meta-llama/Meta-Llama-3-8B"

hyperparams.mixed_precision = "fp8"


# Init the model and accelerator wrapper
model = init_te_llama_model(hyperparams)
accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator(model, hyperparams)


# Finetune the model
finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)

10 finetuning steps complete!
Average time taken per step: 160 milliseconds

模型	精度	步长时间（或每批次毫秒数）	加速比（相对于基线）
HF（基线）	BF16	248	1
TE（将 LlamaDecoderLayer 替换为 TE.TransformerLayer）	BF16	185	1.34
TE（将 LlamaDecoderLayer 替换为 TE.TransformerLayer）	FP8	160	1.55

在开启 FP8 精度后，我们获得了更高的加速，达到 55% (使用 Llama 2 7B)！

Llama 3 性能结果

使用 Llama 3 8B 运行相同的教程会产生以下性能数据

模型	精度	步长时间（或每批次毫秒数）	加速比（相对于基线）
HF（基线）	BF16	270	1
TE（将 LlamaDecoderLayer 替换为 TE.TransformerLayer）	BF16	217	1.24
TE（将 LlamaDecoderLayer 替换为 TE.TransformerLayer）	FP8	185	1.46

对于 Llama 3 8B，我们在 FP8 精度下获得了最高的 46% 加速！

结论

使用 Transformer Engine 的 TransformerLayer 模块作为 Hugging Face 的 LlamaDecoderLayer 的替代品，可以提供比 Hugging Face 原生 Llama 2 和 Llama 3 实现更高的速度。这需要仔细初始化模型，以便将模型权重（用于 LlamaDecoderLayer）正确映射到 TE 的 TransformerLayer 中的对应部分。即使使用 BF16 精度，TransformerLayer 也比基线实现提供加速。使用 FP8 精度，加速效果更加明显！