seqax = sequence modeling + JAX

seqax是一个用于小型到中型LLM预训练研究的代码库。整个训练程序 - 包括模型实现、优化器、多主机FSDP和张量并行分区 - 只有500行代码,可以很好地扩展到~100个GPU或TPU¹,通常可以达到30-50%的良好MFU[性能]。

seqax采用了一种使重要信息可见的编程风格,而不是被抽象和间接隐藏,或者自动和不可预测地推断。这体现在:

数学。seqax实现了训练步骤的所有数学公式,而不是调用外部库。如果你想理解或改变数学公式,它就在那里!
内存。所有进入磁盘上模型检查点的张量都是明确的。所有占用大量内存的张量,包括为反向传播保存的激活函数,都是明确的。你可以直接从源代码中读取内存占用情况。
分区和通信。所有张量和操作的分区布局都是明确的。所有芯片间通信都是明确的。

入门

安装

从你的系统包管理器安装 graphviz: 例如 brew install graphviz 或 apt install graphviz。
在虚拟环境中安装Python依赖项: python -m pip install -r requirements-cpu.txt。

注意: requirements-cpu.txt是为CPU安装配置的。对于GPU或TPU安装,你可能需要安装不同版本的JAX和jaxlib。请参考JAX安装文档。如果你的GPU环境已经安装了Torch-GPU,你可能需要切换到Torch-CPU以避免与JAX-GPU发生冲突。

在CPU上运行本地开发

对于开发和测试,你可以在CPU上运行。通常你会使用我们的合成数据集(已检入此存储库中的/synthetic_dataset.zarr)或Huggingface数据加载器,并设置XLA标志来模拟多个设备,以便测试并行性是否按预期工作:

XLA_FLAGS=--xla_force_host_platform_device_count=8 python -m train --config-name=local_test_synthetic +paths.model_name=synthetic_000

paths.model_name标志指定在磁盘上(/tmp)写入模型检查点的子目录。你通常会在开始新的模型运行时更改这个标志。

在GPU上运行

我们已经配置了一系列模型大小,用于在C4数据集上训练,使用Llama分词器。浏览configs/目录以选择您喜欢的配置文件。每个配置文件顶部都列出了如何运行它。

你通常会希望将paths.model_name设置为每次训练运行的唯一名称。这个路径指定要在磁盘上写入模型检查点的子目录。

性能

最近在A100集群上的基准测试结果:

单主机A100x8

模型大小	MFU
84m	14
270m	24
540m	35
1b	41.6
2b	50.66

在4台A100x8主机上通过Infiniband连接

模型大小	MFU
1b	32.4
2b	39.0

数据加载器

seqax可以直接从Huggingface流式传输训练数据(参见示例配置文件),或者可以先将训练数据转换为磁盘上的预分词格式,我们称之为flat-tokens(参见示例配置文件)。从Huggingface流式传输数据允许你快速尝试不同的数据集,但它没有为在作业中断后从检查点恢复训练提供有效的方式,而且在批处理边界会浪费一些数据集中的标记。flat-tokens格式支持从检查点高效地恢复训练,100%利用训练标记,并且在训练期间也消耗更少的CPU时间。

要预分词训练数据,你可以运行huggingface_to_flat_tokens.py。你需要先安装/tools/requirements.txt中的依赖项,然后你可以调用/tools/configs/c4_en.yaml顶部列出的命令。在现代CPU上,这个脚本每分钟可以处理约100M个标记。你可以使用配置标志限制它处理的输出标记数量。

使用`shardlib`表示分区和通信

seqax附带了一个新的库shardlib,用于使用JAX表示分区和通信,建立在jaxtyping、einops、equinox和shard_map的思想和风格之上。下面我们演示了它的核心思想,实现了一个简单的全连接神经网络的完全分布式数据并行(FSDP)。

# XLA_FLAGS=--xla_force_host_platform_device_count=8 python -m shardlib_example
from shardlib.shardtypes import bf16, bool_, f32, pytree_dataclass, typed_shard_map, u32, make_shardings
from shardlib import shardtypes
shardtypes.register_with_typeguard()
import shardlib.shardops as shardops
from jax.sharding import Mesh
from jax.experimental import mesh_utils
import jax
import jax.numpy as jnp

# 我们设置一个设备网格,其中"d"表示"数据并行"轴。
MESH = Mesh(mesh_utils.create_device_mesh([8], jax.devices()), ('d'))

# 在静止状态下,权重完全分散在数据并行轴上。
#
# `hidden1/d`语法表示第二个轴的大小为`hidden1`,被分散在设备轴`d`上。
# 等价地,你可以将其视为每个设备的形状为`(in, hidden1/d)`,其中`/`表示除法。
@pytree_dataclass
class Weights:
  w1: f32['in hidden1/d']
  w2: f32['hidden1 hidden2/d']
  w3: f32['hidden2/d']

with MESH:
  # 创建dummy权重。
  w = Weights(
    w1=jnp.zeros((8, 8), dtype=jnp.float32),
    w2=jnp.zeros((8, 8), dtype=jnp.float32),
    w3=jnp.zeros((8,), dtype=jnp.float32),
  )
  # 对权重应用分片。分片规范是从Weights类的类型注释中推断出来的。
  w = jax.tree.map(jax.device_put, w, make_shardings(Weights))

我们使用 typed_shard_map 来允许我们编写具有显式通信的每设备代码。

与无类型的 jax.shard_map 相比，in_specs 和 out_specs 不需要指定:它们是从函数签名上的分片中推断出来的。

@typed_shard_map def forward_pass(x: f32[b'batch/d in'], w: Weights) -> f32[b'batch/d']: # Weights 在使用之前都会被全部收集。（这是完全分片数据并行的核心思想。） # in hidden1/d -> in hidden1 语法表达了这个全部收集操作应该做什么:它删除了 hidden1 轴上的 d 分片,resulting in a fully replicated output. w1 = shardops.all_gather('in hidden1/d -> in hidden1', w.w1) # einsum_unreduced 操作是一个 chip-local einsum。与 jnp.einsum 不同,它支持分片语法,并且它使用当前的类型环境进行形状检查,因此如果你在一个函数中以不同的方式使用 batch,它会抛出一个错误。

# 我们称这个 einsum 为"unreduced",因为它不执行任何跨 chip 的折减,即使它们是必需的。
# 例如,在一个 `a b/d, b/d c -> a c` einsum 中,需要对 `d` 分片轴进行跨芯片折减,这是调用者的责任。
y = jax.nn.relu(shardops.einsum_unreduced('batch/d in, in hidden1 -> batch/d hidden1', x, w1))
w2 = shardops.all_gather('hidden1 hidden2/d -> hidden1 hidden2', w.w2)
z = jax.nn.relu(shardops.einsum_unreduced('batch/d hidden1, hidden1 hidden2 -> batch/d hidden2', y, w2))
w3 = shardops.all_gather('hidden2/d -> hidden2', w.w3)
return shardops.einsum_unreduced('batch/d hidden2, hidden2 -> batch/d', z, w3)

x = forward_pass(jnp.zeros((32, 8), dtype=jnp.float32), w) assert(x.shape == (32,))