Levanter: 开源高性能大规模语言模型训练框架

Levanter简介

Levanter是一个由斯坦福大学CRFM(基础模型研究中心)开发的开源大规模语言模型训练框架。它的名字源自一句名言:"你无法阻止雷暴,但你可以利用电力;你无法指引风向,但你可以调整帆的角度,无论风从哪个方向吹来,都能推动你的船只前进。"这恰如其分地描述了Levanter的设计理念 - 尽管我们无法完全掌控大规模语言模型的训练过程,但我们可以通过精心设计的框架来更好地利用和控制这一过程。

Levanter的三大核心目标是:

1. 可读性(Legible): 利用命名张量库Haliax编写易于理解和组合的深度学习代码,同时保持高性能。

2. 可扩展性(Scalable): 支持大规模模型训练,能够在各种硬件(包括GPU和TPU)上进行训练。

3. 可重现性(Reproducible): 实现逐位确定性,确保相同的配置总是产生相同的结果,即使在中断和恢复的情况下也能保持一致。

Levanter基于JAX、Equinox和Haliax构建,充分利用了这些优秀框架的特性。

Levanter的主要特性

分布式训练

Levanter支持在TPU(很快也将支持GPU)上进行分布式训练,包括全尺寸数据并行(FSDP)和张量并行。这使得Levanter能够高效地利用大规模计算资源,加速模型训练过程。

兼容性

Levanter与Hugging Face生态系统高度兼容,支持导入和导出模型、分词器和数据集。通过SafeTensors技术,Levanter能够无缝对接Hugging Face的各种资源,极大地增强了其实用性和生态系统集成能力。

性能

在性能方面,Levanter可以与商业支持的框架(如MosaicML的Composer或Google的MaxText)相媲美。这意味着研究人员和开发者可以使用Levanter获得与商业级解决方案相当的训练效率,而无需支付高昂的费用。

缓存式按需数据预处理

Levanter采用在线预处理语料库的方式,但会缓存预处理结果,使得恢复训练变得更快,后续运行也更加高效。一旦缓存的第一部分完成,Levanter就会开始训练,实现了数据处理和模型训练的高效并行。

优化器支持

Levanter支持新型Sophia优化器,相比Adam可以提高2倍的训练速度。同时,它也集成了Optax库,支持AdamW等常用优化器,为用户提供了丰富的优化策略选择。

日志记录

在日志记录方面,Levanter支持多种后端,包括WandB和TensorBoard。用户可以轻松添加新的日志后端,而且Levanter甚至支持在JAX jit编译的函数内部进行日志记录,为实验监控和调试提供了强大的工具。

可重现性

在TPU上,Levanter实现了逐位确定性,这意味着相同的配置总是会产生相同的结果,即使在中断和恢复的情况下也是如此。这对于科研实验的可重复性和结果验证至关重要。

分布式检查点

Levanter通过Google的TensorStore库支持分布式检查点。训练甚至可以在不同数量的主机上恢复,尽管目前这会破坏可重现性。这一特性极大地提高了训练的灵活性和容错能力。

Levanter的安装与使用

安装Levanter

在安装Levanter之前,需要先安装适合您平台的JAX版本。然后,可以通过以下命令安装Levanter:

pip install levanter

如果想使用最新的GitHub版本,可以执行:

git clone https://github.com/stanford-crfm/levanter.git
cd levanter
pip install -e .
wandb login  # 可选,用于日志记录

对于同时开发Haliax和Levanter的用户,可以采用以下方式:

git clone https://github.com/stanford-crfm/levanter.git
cd levanter
pip install -e .
cd ..
git clone https://github.com/stanford-crfm/haliax.git
cd haliax
pip install -e .
cd ../levanter

训练示例

训练GPT2-nano模型

作为入门示例,可以使用以下命令训练一个GPT2-nano模型:

python -m levanter.main.train_lm --config_path config/gpt2_nano.yaml

这将在WikiText-103数据集上训练一个小型GPT2模型。

使用自定义数据集训练GPT2-small模型

如果想使用自己的数据集训练GPT2-small模型,可以修改配置文件中的数据源:

python -m levanter.main.train_lm --config_path config/gpt2_small.yaml --data.id openwebtext

还可以指定自定义的分词器和缓存目录:

python -m levanter.main.train_lm --config_path config/gpt2_small.yaml --data.id openwebtext --data.tokenizer "EleutherAI/gpt-neox-20b" --data.cache_dir "gs://path/to/cache/dir"

如果数据源是URL列表,可以这样指定:

python -m levanter.main.train_lm --config_path config/gpt2_small.yaml --data.train_urls ["https://path/to/train/data_*.jsonl.gz"] --data.validation_urls ["https://path/to/val/data_*.jsonl.gz"]

自定义配置文件

Levanter使用YAML格式的配置文件来定义模型结构、训练参数等。以下是一个gpt2_small.yaml配置文件的示例:

data:
  train_urls:
      - "gs://pubmed-mosaic/openwebtext-sharded/openwebtext_train.{1..128}-of-128.jsonl.gz"
  validation_urls:
      - "gs://pubmed-mosaic/openwebtext-sharded/openwebtext_val.{1..8}-of-8.jsonl.gz"
  cache_dir: "gs://pubmed-mosaic/tokenized/openwebtext/"
model:
  gpt2:
    hidden_dim: 768
    num_heads: 12
    num_layers: 12
    seq_len: 1024
    gradient_checkpointing: true
    scale_attn_by_inverse_layer_idx: true
trainer:
  tracker:
    type: wandb
    project: "levanter"
    tags: [ "openwebtext", "gpt2"]
  mp: p=f32,c=bfloat16
  model_axis_size: 1
  per_device_parallelism: 4
  train_batch_size: 512
optimizer:
  learning_rate: 6E-4
  weight_decay: 0.1
  min_lr_ratio: 0.1

用户可以根据需要修改这些参数,以适应不同的训练需求。

Levanter支持的模型架构

目前,Levanter支持以下模型架构:

• GPT-2
• LLama 1和2
• Backpacks
• MosaicML的MPT

未来,Levanter团队计划添加更多的模型架构支持。

使用Llama 1或Llama 2继续预训练

如果想在OpenWebText数据集上继续预训练Llama 1或Llama 2模型,可以使用以下命令:

python -m levanter.main.train_lm --config_path config/llama2_7b_continued.yaml

分布式和云训练

Levanter支持在TPU Cloud VM和CUDA环境中进行分布式训练。对于TPU用户,可以参考TPU入门指南来设置环境并运行Levanter。对于CUDA用户,可以查看CUDA入门指南了解如何在GPU环境中使用Levanter。

贡献与社区

Levanter是一个开源项目,欢迎社区贡献。贡献者可以参考CONTRIBUTING.md文件了解如何参与项目开发。此外,Levanter社区还在Jax LLM Discord的#levanter频道中活跃,欢迎加入讨论。

结语

Levanter作为一个强大而灵活的大规模语言模型训练框架,为研究人员和开发者提供了一个优秀的工具。它的可读性、可扩展性和可重现性使其在竞争激烈的AI领域中脱颖而出。通过支持多种模型架构、优化技术和分布式训练,Levanter为推动大规模语言模型的研究和应用提供了坚实的基础。

随着AI技术的不断发展,像Levanter这样的开源框架将在推动技术民主化和加速创新方面发挥越来越重要的作用。无论您是研究人员、开发者还是AI爱好者,Levanter都为您提供了一个强大的平台,让您能够探索和挑战大规模语言模型的前沿。

🚀 立即开始使用Levanter,成为下一代AI技术革命的一部分吧!