概述

MaxText是一个高性能、高度可扩展、开源的LLM，完全用Python/Jax编写，目标是Google Cloud的TPU和GPU进行训练和推理。借助Jax和XLA编译器的强大功能，MaxText实现了高MFU，并且可以从单主机扩展到非常大的集群，同时保持简单且“无优化”。

MaxText旨在成为研究和生产中雄心勃勃的LLM项目的起点。我们鼓励用户首先直接使用MaxText进行实验，然后分叉并修改MaxText以满足他们的需求。

我们使用MaxText 演示了int8的高性能、良好收敛训练以及将训练规模扩展到约51K芯片。

支持的主要功能：

TPU和GPU（预览中）
训练和推理（预览中）
模型：Llama2、Mistral和Gemma

快速开始

首次运行MaxText时，我们提供了详细的说明。

MaxText支持各种开源模型的训练和推理。请参阅快速入门文件夹中的用户指南以了解更多信息。

一些额外的有用指南：

Gemma：一个由Google DeepMind基于Gemini研究和技术开发的开源权重大型语言模型（LLM）家族。您可以使用这些说明进行解码和微调。
Llama2：一个由Meta开发的开源权重大型语言模型（LLM）家族。您可以使用这些说明进行解码和微调。
Mixtral：一个由Mistral AI开发的开源权重稀疏专家混合（MoE）模型家族。您可以使用这些说明进行解码和微调。

除了快速入门指南外，还有许多不断添加的MaxText功能！完整的端到端测试套件在end_to_end中。我们以每夜的频率运行它们。它们可以成为理解MaxText的良好来源。或者，您可以查看几乎持续运行的单元测试。

运行时性能结果

有关复现这些结果的更多细节，请参阅MaxText/configs/README.md。

TPU v5p

参数数量	加速器类型	TFLOP/芯片/秒	模型运算利用率（MFU）
32B	v5p-128	3.28e+02	71.47%
64B	v5p-128	3.23e+02	70.31%
128B	v5p-256	3.15e+02	68.68%
128B	v5p-512	3.15e+02	68.53%
256B	v5p-1024	3.16e+02	68.82%
512B	v5p-1024	2.94e+02	63.99%
1024B	v5p-2048	2.49e+02	64.05%
1024B	v5p-4096	2.97e+02	64.80%
1160B	v5p-7680	2.95e+02	64.27%
1160B	v5p-12288	3.04e+02	66.23%

TPU v5e

适用于16B、32B、64B和128B模型。完整的运行配置见MaxText/configs/v5e/中的16b.sh、32b.sh、64b.sh、128b.sh。

硬件	16B TFLOP/秒/芯片	16B MFU	32B TFLOP/秒/芯片	32B MFU	64B TFLOP/秒/芯片	64B MFU	128B TFLOP/秒/芯片	128B MFU
1x v5e-256	120	61.10%	132	66.86%	118	59.90%	110	56.06%
2x v5e-256	117	59.37%	128	64.81%	112	56.66%	110	55.82%
4x v5e-256	117	59.14%	126	64.10%	110	55.85%	108	54.93%
8x v5e-256	115	58.27%	125	63.67%	108	54.96%	104	52.93%
16x v5e-256	111	56.56%	123	62.26%	105	53.29%	100	50.86%
32x v5e-256	108	54.65%	119	60.40%	99	50.18%	91	46.25%

与其他方案的比较

MaxText受到了MinGPT/NanoGPT的启发，这是一种优雅的独立GPT实现，使用PyTorch编写，目标是Nvidia GPU。MaxText更复杂，支持更多行业标准模型并扩展到数万个芯片。最终，MaxText的MFU比最近报告的代码库中的17%高出三倍多，具有大规模扩展性，并实现了用于高效自回归解码的键值缓存。

MaxText更类似于Nvidia/Megatron-LM，后者是一个非常优化的LLM实现，目标是Nvidia GPU。这两个实现实现了可比的MFU。代码库的差异突显了不同的编程策略。MaxText是纯Python，主要依赖于XLA编译器来实现高性能。相比之下，Megatron-LM是Python和CUDA的混合体，依赖于优化良好的CUDA内核来实现高性能。

MaxText还可以与Pax进行比较。与Pax一样，MaxText在Jax中提供了高性能和可扩以下配置将有助于调试故障或在程序卡住或挂起时通过收集堆栈跟踪进行调试。在 MaxText/configs/base.yml 中相应更改参数值：

设置 collect_stack_trace: True 以在故障或程序挂起时启用堆栈跟踪收集。此设置将定期转储程序的跟踪信息以帮助调试。要禁用此功能，设置 collect_stack_trace: False。
设置 stack_trace_to_cloud: False 以在控制台上显示堆栈跟踪。stack_trace_to_cloud: True 会在 TPUs 的 /tmp/debugging 中创建一个临时文件来存储堆栈跟踪。有一个代理在 TPU 虚拟机上运行，会定期将临时目录中的跟踪信息上传到 gcp 项目的云日志。您可以使用以下查询在云日志的日志资源管理器中查看跟踪信息：
```
logName="projects/<project_name>/logs/tpu.googleapis.com%2Fruntime_monitor"
jsonPayload.verb="stacktraceanalyzer"
```
stack_trace_interval_seconds 表示每次堆栈跟踪收集事件之间的秒数。在 stack_trace_interval_seconds: 600 中设置将会每 600 秒（10 分钟）收集一次堆栈跟踪。

这是相关的 PyPI 包：https://pypi.org/project/cloud-tpu-diagnostics。

提前编译（AOT）

为了提前编译您的训练任务，我们提供了一个工具 train_compile.py。此工具允许您为目标硬件提前编译 train.py 中的主要 train_step（例如大量 v5e 设备）而无需使用整个集群。

TPU 支持

您可以仅使用 CPU 或从不同家族的单个 VM 来为 TPU 集群进行预编译。此编译有助于实现以下两个主要目标：

它将标记任何内存不足（OOM）的信息，例如当 per_device_batch_size 设置过高时，与在目标硬件上编译时相同的 OOM 堆栈跟踪。
提前编译可以保存并加载，以实现目标硬件上的快速启动和重新启动时间。

工具 train_compile.py 与 train.py 紧密相关，并使用相同的配置文件 configs/base.yml。虽然您不需要在 TPU 上运行，但需要安装 jax[tpu] 以及其他依赖项，因此我们建议运行 setup.sh 来安装这些依赖项（如果尚未安装）。

示例 AOT 1：提前编译基础

在安装了上述依赖项后，您已经准备好提前编译：

# 在单台机器上运行，例如 CPU
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=v5e-256 compile_topology_num_slices=2 \
global_parameter_scale=16 per_device_batch_size=4

这将在两个 v5e 节点上编译一个 16B 参数的 MaxText 模型。

示例 AOT 2：保存已编译的函数，然后加载并运行它

以下示例保存并加载已编译的 train_step，从保存开始：

步骤 1：运行 AOT 并保存已编译的函数

# 在单台机器上运行，例如 CPU
export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true"
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=v5e-256 \
compile_topology_num_slices=2 \
compiled_trainstep_file=my_compiled_train.pickle global_parameter_scale=16 \
per_device_batch_size=4 steps=10000 learning_rate=1e-3

步骤 2：运行 train.py 并加载已编译的函数

要加载已编译的 train_step，只需将 compiled_trainstep_file=my_compiled_train.pickle 传入 train.py：

# 在目标硬件的每个主机上运行，例如在两个 v5e-256 薄片的每个主机上
export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true"
python3 MaxText/train.py MaxText/configs/base.yml run_name=example_load_compile \
compiled_trainstep_file=my_compiled_train.pickle \
global_parameter_scale=16  per_device_batch_size=4 steps=10000 learning_rate=1e-3 \
base_output_directory=gs://my-output-bucket dataset_path=gs://my-dataset-bucket

在上述示例 2 的保存步骤中，我们包括了导出编译器标志 LIBTPU_INIT_ARGS 和 learning_rate，因为这些会影响到已编译的对象 my_compiled_train.pickle。当您通过 compile_train.py 初次编译时，模型的大小（例如 global_parameter_scale、max_sequence_length 和 per_device_batch）是固定的，如果您尝试使用与编译时不同的大小运行保存的已编译对象，您会看到一个大小错误。然而，一个细微的备注是 学习率计划 也是在运行 compile_train 时固定的 - 这由 steps 和 learning_rate 共同决定。优化器参数如 adam_b1 仅作为形状对象传递给编译器 - 因此它们的实际值是在运行 train.py 时确定的，而不是在编译期间。如果您传递了不同形状（例如 per_device_batch），您会得到一个明确的错误消息，报告已编译签名与输入的形状期望不同。如果您尝试在与通过 compile_topology 请求的编译目标不同的硬件上运行，您会收到一个错误，提示无法将设备从已编译映射到您的实际设备。使用与编译时不同的 XLA 标志或 LIBTPU 可能会在编译环境中无错误地运行。然而，在这种情况下，没有保证的行为；您应该在与编译环境相同的环境中运行。

GPU 支持

GPU 也支持提前编译，但与 TPU 有一些不同之处：

GPU 不支持跨硬件编译：仍需要 GPU 主机运行提前编译，但单个 GPU 主机可以为更大的相同硬件集群编译程序。
对于 A3 云 GPU，最大“分片”大小是一台主机，而 compile_topology_num_slices 参数代表预编译的 A3 机器数量。

示例

此示例说明了用于多主机 GPU 编译的标志，目标是 4 个 A3 主机的集群：

步骤 1：运行 AOT 并保存已编译的函数

# 在单个 A3 机器上运行
export XLA_FLAGS="--xla_gpu_enable_async_collectives=true"
python3 MaxText/train_compile.py MaxText/configs/base.yml compile_topology=a3 \
compile_topology_num_slices=4 \
compiled_trainstep_file=my_compiled_train.pickle global_parameter_scale=16 \
attention=dot_product per_device_batch_size=4 steps=10000 learning_rate=1e-3

步骤 2：运行 train.py 并加载已编译的函数

要加载已编译的 train_step，只需将 compiled_trainstep_file=my_compiled_train.pickle 传入 train.py：

# 在四个目标 A3 主机上的每个主机上运行
export XLA_FLAGS="--xla_gpu_enable_async_collectives=true"
python3 MaxText/train.py MaxText/configs/base.yml run_name=example_load_compile \
compiled_trainstep_file=my_compiled_train.pickle \
attention=dot_product global_parameter_scale=16  per_device_batch_size=4 steps=10000 learning_rate=1e-3 \
base_output_directory=gs://my-output-bucket dataset_path=gs://my-dataset-bucket

如同 TPU 的情况一样，注意编译环境必须与执行环境匹配，在这种情况下，通过设置相同的 XLA_FLAGS。