JAX序列建模框架简化小型到中型LLM预训练研究
seqax是基于JAX的序列建模框架,面向小型到中型LLM预训练研究。它以500行简洁代码实现模型、优化器、多主机FSDP和张量并行分区,具备良好扩展性和性能。seqax注重代码可读性和透明度,使数学计算、内存使用和芯片间通信等关键细节清晰可见,便于理解和修改。
seqax是一个用于小型到中型LLM预训练研究的代码库。整个训练程序 - 包括模型实现、优化器、多主机FSDP和张量并行分区 - 只有500行代码,可以很好地扩展到~100个GPU或TPU1,通常可以达到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上运行。通常你会使用我们的合成数据集(已检入此存储库中的/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)写入模型检查点的子目录。你通常会在开始新的模型运行时更改这个标志。
我们已经配置了一系列模型大小,用于在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,))
要在更大的规模上获得良好的性能,需要使用管道并行性(我们尚未实现)。在这个规模上,你可能还需要使用定制的内核来进一步提高性能,代价是代码的简单性。 ↩
AI小说写作助手,一站式润色、改写、扩写
蛙蛙写作—国内先进的AI写作平台,涵盖小说、学术、社交媒体等多场景。提供续写、改写、润色等功能,助力创作者高效优化写作流程。界面简洁,功能全面,适合各类写作者提升内容品质和工作效率。
字节跳动发布的AI编程神器IDE
Trae是一种自适应的集成开发环境(IDE),通过自动化和多元协作改变开发流程。利用Trae,团队能够更快速、精确地编写和部署代码,从而提高编程效率和项目交付速度。Trae具备上下文感知和代码自动完成功能,是提升开发效率的理想工具。
全能AI智能助手,随时解答生活与工作的多样问题
问小白,由元石科技研发的AI智能助手,快速准确地解答各种生活和工作问题,包括但不限于搜索、规划和社交互动,帮助用户在日常生活中提高效率,轻松管理个人事务。
实时语音翻译/同声传译工具
Transly是一个多场景的AI大语言模型驱动的同声传译、专业翻译助手,它拥有超精准的音频识别翻译能力,几乎零延迟的使用体验和支持多国语言可以让你带它走遍全球,无论你是留学生、商务人士、韩剧美剧爱好者,还是出国游玩、多国会议、跨国追星等等,都可以满足你所有需要同传的场景需求,线上线下通用,扫除语言障碍,让全世界的语言交流不再有国界。
一键生成PPT和Word,让学习生活更轻松
讯飞智文是一个利用 AI 技术的项目,能够帮助用户生成 PPT 以及各类文档。无论是商业领域的市场分析报告、年度目标制定,还是学生群体的职业生涯规划、实习避坑指南,亦或是活动策划、旅游攻略等内容,它都能提供支持,帮助用户精准表达,轻松呈现各种信息。
深度推理能力全新升级,全面对标OpenAI o1
科大讯飞的星火大模型,支持语言理解、知识问答和文本创作等多功能,适用于多种文件和业务场景,提升办公和日常生活的效率。讯飞星火是一个提供丰富智能服务的平台,涵盖科技资讯、图像创作、写作辅助、编程解答、科研文献解读等功能,能为不同需求的用户提供便捷高效的帮助,助力用户轻松获取信息、解决问题,满足多样化使用场景。
一种基于大语言模型的高效单流解耦语音令牌文本到语音合成模型
Spark-TTS 是一个基于 PyTorch 的开源文本到语音合成项目,由多个知名机构联合参与。该项目提供了高效的 LLM(大语言模型)驱动的语音合成方案,支持语音克隆和语音创建功能,可通过命令行界面(CLI)和 Web UI 两种方式使用。用户可以根据需求调整语音的性别、音高、速度等参数,生成高质量的语音。该项目适用于多种场景,如有声读物制作、智能语音助手开发等。
AI助力,做PPT更简单!
咔片是一款轻量化在线演示设计工具,借助 AI 技术,实现从内容生成到智能设计的一站式 PPT 制作服务。支持多种文档格式导入生成 PPT,提供海量模板、智能美化、素材替换等功能,适用于销售、教师、学生等各类人群,能高效制作出高品质 PPT,满足不同场景演示需求。
选题、配图、成文,一站式创作,让内容运营更高效
讯飞绘文,一个AI集成平台,支持写作、选题、配图、排版和发布。高效生成适用于各类媒体的定制内容,加速品牌传播,提升内容营销效果。
专业的AI公文写作平台,公文写作神器
AI 材料星,专业的 AI 公文写作辅助平台,为体制内工作人员提供高效的公文写作解决方案。拥有海量公文文库、9 大核心 AI 功能,支持 30 + 文稿类型生成,助力快速完成领导讲话、工作总结、述职报告等材料,提升办公效率,是体制打工人的得力写作神器。
最新AI工具、AI资讯
独家AI资源、AI项目落地
微信扫一扫关注公众号