生成式 AI 具有创建全新内容的能力，这是传统机器学习（ML）方法难以实现的。在自然语言处理（NLP）领域， 大型语言模型（LLMs） 的出现特别催生了许多创新和创造性的 AI 应用案例，包括客户支持聊天机器人、语音助手、文本摘要和翻译等——这些任务以前由人类处理。
LLMs 通过各种方法不断发展，包括增加参数数量和采用 Mixture of Experts (MoE) 等新算法。预计许多行业 (包括零售、制造和金融) 都会应用和调整 LLMs。
然而，许多目前在 LLM 排行榜上名列前茅的模型在非英语语言（包括日语）方面表现出的理解和性能不足。其中一个原因是训练语料库包含大量英语数据。例如， GPT-3 语料库中只有 0.11%是日语数据 。创建在日语（日语的训练数据比英语少）中表现良好的 LLM 模型极具挑战性。
本文介绍了在 生成式 AI 加速器挑战赛（GENIAC） 项目中训练具有 172 亿个参数的 AI 模型时获得的见解，该项目使用 NVIDIA Megatron-LM 帮助解决用于日语理解的高性能模型短缺问题。
GENIAC 上的 LLM-JP 计划
经济产业省（METI） 为了提高日本的平台模型开发能力水平，并鼓励企业等发挥创造力，启动了 GENIAC。GENIAC 提供了计算资源，支持了与企业和数据持有者的匹配，促进了与全球技术公司的合作，举办了社区活动，并对已开发的平台模型的性能进行了评估。
LLM-jp 项目旨在开发一个 完全开放的模型，拥有 172 亿个参数 (可在 Hugging Face 上获取)，并具备强大的日语能力。 LLM-jp 172B 是当时日本最大的模型开发项目 (2024 年 2 月至 8 月)，广泛分享其开发知识具有重要意义。
LLM-jp 是自然语言处理和计算机系统领域的研究人员 (主要是 NII) 发起的一项计划，旨在通过持续开发完全开放的商用模型，积累有关训练原理的数学阐明的专业知识，例如大型模型如何获得泛化性能和学习效率。目标是积累有关训练效率的专业知识。
使用 NVIDIA Megatron-LM 训练模型
Megatron-LM 是一种轻量级的研究型框架，利用 Megatron-Core 以前所未有的速度训练大型语言模型（LLMs）。Megatron-Core 是一个开源库，包含了 GPU 优化技术和大规模训练所必需的尖端系统级优化。
Megatron-Core 支持各种高级模型并行技术，包括张量、序列、流水线、上下文和 MoE 专家并行。此库提供 可定制的构建块 、训练弹性功能（例如 快速分布式检查点 ）以及许多其他创新，如 基于 Mamba 的混合模型训练 。它与所有 NVIDIA Tensor Core GPU 兼容，并支持 Transformer Engine（TE）， NVIDIA Hopper 架构 引入了 FP8 精度。
模型架构和训练设置
表 1 概述了该项目的模型架构，该架构遵循 Llama 2 架构 。
| 参数 | 价值 |
| 隐藏尺寸 | 12288 |
| FFN 中间尺寸 | 38464 |
| 层数量 | 96 |
| 注意力头数量 | 96 |
| 查询组数量 | 16 |
| 激活函数 | SwiGLU |
| 位置嵌入 | RoPE |
| 归一化 | RMS 规范 |
LLM-jp 172B 模型正在使用为该项目开发的多语种语料库中的 2.1 万亿个令牌从头开始训练，该语料库主要由日语和英语组成。训练使用 NVIDIA H100 Tensor Core GPU 在 Google Cloud A3 实例上执行，并使用 Transformer Engine 进行 FP8 混合训练。实验中使用了 Megatron-Core v0.6 和 Transformer Engine v1.4。
表 2 显示了用于训练的超参数设置。
| 参数 | 值 |
| LR | 1E-4 |
| 最小 LR | 1E-5 |
| LR 热身迭代器 | 2000 |
| 权重衰减 | 0.1 |
| 研究生剪辑 | 1.0 |
| 全局批量大小 | 1728 |
| 上下文长度 | 4096 |
此外，还整合了 PaLM 中使用的 z-loss 和 batch-skipping 技术，以稳定训练过程，并使用 flash attention 进一步加速训练过程。
要查看其他训练配置，请参阅 llm-jp/Megatron-LM 。
训练吞吐量和结果
最新的 LLM-jp 172B 模型的预训练目前正在进行中，每隔数千次迭代就会进行定期评估，以监控训练进度，并确保日语和英语下游任务的准确性成功（图 1）。到目前为止，在目标 2.1 万亿个令牌中，已经完成了 80%以上。
值得注意的是，在大约 7,000 次迭代后，TFLOP/s 急剧增加，这与从 BF16 到 FP8-hybrid 精度的过渡相对应。在本实验中，在 7,000 次迭代之前，使用 BF16 和 TE 进行训练，在 7,000 次迭代之后，使用 FP8-hybrid 和 TE。在 Megatron-LM 中，可以使用简单的选项 --fp8-format‘hybrid‘启用混合 FP8 训练。请注意，此功能是实验性功能，不久后将会进行进一步优化。
我们开始使用 BF16 和 TE 训练，然后切换到 FP8 混合，不仅是为了查看 BF16 和 FP8 之间的 tokens/sec 性能差异，也是为了使初始训练更加稳定。在训练的早期阶段，由于热身，学习率（LR）增加，导致训练不稳定。
我们选择使用 BF16 执行初始训练，在确认训练损失、优化器状态、梯度范数等值没有问题后，我们切换到 FP8 以加速训练过程。FP8 混合提高了训练速度。我们观察到使用 Megatron-LM 的训练速度为 545-553 TFLOP/s。
结束语
如上所述，LLM-jp 172B 的训练仍在使用 Megatron-LM 进行中。根据使用当前检查点数据对下游任务的评估结果，我们认为该模型已经具备出色的日语能力，但完整的模型预计将于明年初准备就绪。在需要大量数据集的 LLM 预训练中，训练时间通常是一项重大挑战。因此，像 Megatron-LM 这样的高效训练框架对于加速生成式 AI 研究和开发至关重要。对于使用 Megatron-LM 训练的 172B 模型，我们探讨了 FP8-hybrid 训练作为提高训练速度的潜在方法，训练速度从 400 TFLOP/s 提高到 550 TFLOP/s，提高了 1.4 倍。我们观察到性能加速从 400 TFLOP/s 到 550 TFLOP/s，这表明 FP8-hybrid 可能是提高大规模模型预训练效率的一种有价值的方法。
