TorchTitan：PyTorch原生的大规模LLM训练库

TorchTitan: 原生PyTorch大规模LLM训练库

TorchTitan是一个基于PyTorch的开源项目，专为大规模语言模型(LLM)训练而设计。它以简洁高效的代码实现，展示了PyTorch最新的分布式训练特性，为研究人员和开发者提供了一个易于理解和扩展的大模型训练框架。

TorchTitan的核心理念

TorchTitan的开发遵循以下几个核心原则：

设计简洁易懂，便于使用和扩展到不同的训练目的。
在应用1D、2D或3D并行时，对模型代码的修改最小化。
采用模块化组件而非单一庞大的代码库。
可以在短时间内快速上手使用。

值得注意的是，TorchTitan并非旨在替代现有的大规模LLM训练代码库（如Megatron、Megablocks、LLM Foundry、Deepspeed等），而是作为一个补充。它的目标是展示PyTorch最新的分布式训练特性，并希望这些特性能够被其他代码库迅速采纳。

TorchTitan Logo

TorchTitan的主要特性

目前，TorchTitan支持从头开始预训练各种规模的Llama 3和Llama 2模型。它与PyTorch nightly版本（torch-2.4.0.dev20240412）兼容，并提供了以下关键功能：

FSDP2（每参数分片）
张量并行（包括异步TP）
选择性层和操作符激活检查点
分布式检查点（包括异步检查点）
可检查点的数据加载，预配置了C4数据集（144M条目）
通过TensorBoard显示和记录损失、GPU内存、每秒令牌数和MFU
学习率调度器、元初始化、可选的融合RMSNorm等特性
Float8支持
torch.compile支持
所有选项都可以通过toml文件轻松配置
可互操作的检查点，可直接加载

安装和使用

要开始使用TorchTitan，首先需要克隆项目仓库并安装依赖：

git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121 # or cu118

接下来，您需要下载一个tokenizer模型。TorchTitan目前支持开箱即用地训练Llama 3（8B、70B）和Llama 2（7B、13B、70B）。请按照meta-llama仓库的官方说明确保您有权访问Llama模型权重。

确认访问权限后，您可以运行以下命令将Llama 3 / Llama 2 tokenizer下载到本地计算机：

# 获取HF token: https://huggingface.co/settings/tokens

# 下载llama3或3.1 tokenizer.model
python torchtitan/datasets/download_tokenizer.py --repo_id meta-llama/Meta-Llama-3-8B --tokenizer_path "original" --hf_token=...

# 下载llama2 tokenizer.model
python torchtitan/datasets/download_tokenizer.py --repo_id meta-llama/Llama-2-13b-hf --hf_token=...

开始训练

要在本地8个GPU上开始训练Llama 3 8B模型，可以运行以下命令：

CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh

使用TensorBoard可视化训练过程

TorchTitan支持使用TensorBoard来可视化训练指标。要在远程服务器上训练模型并在本地Web浏览器中查看TensorBoard指标，请按照以下步骤操作：

确保在模型训练中将metrics.enable_tensorboard选项设置为true（可以在.toml文件或CLI中设置）。
通过在本地CLI中运行以下命令设置SSH隧道：

ssh -L 6006:127.0.0.1:6006 [username]@[hostname]

在登录到远程服务器的SSH隧道内，进入torchtitan仓库，并启动TensorBoard后端：

tensorboard --logdir=./outputs/tb

在本地Web浏览器中，访问提供的URL或 http://localhost:6006/。

多节点训练

对于ParallelCluster/Slurm类型的配置，您可以使用multinode_trainer.slurm文件提交sbatch作业。首先，调整节点数和GPU数：

#SBATCH --ntasks=2
#SBATCH --nodes=2

然后启动运行，其中nnodes是您的总节点数，与上面的sbatch节点数匹配。如果每个节点的GPU数不是8，请在torchrun命令中调整--nproc_per_node，并在SBATCH命令部分调整#SBATCH --gpus-per-task。

TorchTitan的未来发展

TorchTitan团队计划在未来添加以下功能：

上下文并行
流水线并行（及3D并行）
HSDP（高速数据并行）

这些功能的加入将进一步提升TorchTitan在大规模LLM训练方面的能力和效率。

结语

TorchTitan作为一个展示PyTorch最新分布式训练特性的项目，为研究人员和开发者提供了一个简洁、高效的大规模LLM训练框架。它的模块化设计和易于使用的特性使得用户可以快速上手，并根据自己的需求进行扩展。虽然TorchTitan可能不会发展成为一个大型社区项目，但它的创新和实践无疑将为整个深度学习领域带来积极影响。

随着深度学习和大规模语言模型的不断发展，像TorchTitan这样的项目将继续推动技术边界，为研究人员和开发者提供更多可能性。我们期待看到TorchTitan在未来的发展，以及它如何影响和改变大规模LLM训练的格局。

🔗 相关链接：

通过使用TorchTitan，研究人员和开发者可以更轻松地进行大规模LLM的训练和实验。无论您是想复现最新的研究成果，还是开发自己的创新模型，TorchTitan都为您提供了一个强大而灵活的工具。让我们一起期待TorchTitan和整个深度学习领域的未来发展！