nanotron

<h1 align="center">⚡️ Nanotron</h1> <p align="center"> <a href="https://github.com/huggingface/nanotron/releases"> <img alt="GitHub release" src="https://yellow-cdn.veclightyear.com/835a84d5/d92ac2f8-d4c0-47b8-80b9-420cd2de2d67.svg"> </a> <a href="https://github.com/huggingface/nanotron/blob/master/LICENSE"> <img alt="License" src="https://yellow-cdn.veclightyear.com/835a84d5/e3087f4b-c1ff-4d03-8379-b161a0ffcd82.svg?color=green"> </a> </p> <h4 align="center"> <p> <a href="#installation">安装</a> • <a href="#quick-start">快速开始</a> • <a href="#features">特性</a> • <a href="CONTRIBUTING.md">贡献</a> <p> </h4> <h3 align="center"> <a href="https://huggingface.co/nanotron"><img style="float: middle; padding: 10px 10px 10px 10px;" width="60" height="55" src="https://yellow-cdn.veclightyear.com/835a84d5/be59f932-7db1-4e54-a799-7745078fe6bb.png" /></a> </h3> <h3 align="center"> <p>简化模型预训练 </h3>

Nanotron是一个用于预训练transformer模型的库。它提供了一个简单灵活的API，可以在自定义数据集上预训练模型。Nanotron旨在易于使用、快速且可扩展。它的设计基于以下原则：

• 简单性：Nanotron设计为易于使用。它提供了一个简单灵活的API，可以在自定义数据集上预训练模型。
• 性能：为速度和可扩展性而优化，Nanotron使用最新技术更快、更高效地训练模型。

安装

# 要求：Python>=3.10
git clone https://github.com/huggingface/nanotron
cd nanotron
pip install --upgrade pip
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
pip install -e .

# 如果你想使用示例脚本，请安装依赖
pip install datasets transformers
pip install triton "flash-attn>=2.5.0" --no-build-isolation

[!注意] 如果你遇到undefined symbol: ncclCommRegister错误，你应该安装torch 2.1.2：pip install torch==2.1.2 --index-url https://download.pytorch.org/whl/cu121

[!提示] 如果安装了wandb，我们会自动记录日志。为此，你可以使用pip install wandb。如果你不想使用wandb，可以运行wandb disabled。

快速开始

训练一个小型Llama模型

以下命令将在单节点8个GPU上训练一个小型Llama模型。模型将保存在配置文件中指定的checkpoints目录中。

CUDA_DEVICE_MAX_CONNECTIONS=1 torchrun --nproc_per_node=8 run_train.py --config-file examples/config_tiny_llama.yaml

从你的检查点运行生成

torchrun --nproc_per_node=1 run_generate.py --ckpt-path checkpoints/10/ --tp 1 --pp 1
# 我们可以设置更大的TP以加快生成速度，对于非常大的模型，可以设置更大的PP。

自定义示例

你可以在/examples目录中找到更多示例：

<!-- 制作我们支持的示例表格 -->

示例	描述
custom-dataloader	将自定义数据加载器插入nanotron
datatrove	使用datatrove库加载数据
doremi	使用DoReMi加速训练
mamba	训练一个Mamba模型示例
moe	训练一个混合专家（MoE）模型示例
mup	使用谱µTransfer扩展你的模型

我们正在努力添加更多示例！欢迎提交PR添加你自己的示例。🚀

特性

我们目前支持以下特性：

• 3D并行（DP+TP+PP）
• MoE的专家并行
• PP的AFAB和1F1B调度
• TP和PP的显式API，便于调试
• ZeRO-1优化器
• FP32梯度累积
• 参数绑定/分片
• 大模型的自定义模块检查点
• 用于扩展神经网络的谱µTransfer参数化
• Mamba示例

我们的路线图包括：

• FP8训练
• ZeRO-3优化器（即FSDP）
• torch.compile支持
• 环形注意力
• 交错1f1b调度

致谢

我们要感谢所有从事LLM工作的人，特别是那些公开分享他们的工作并给予我们极大启发的人：Nvidia的Megatron-LM/apex，Microsoft的DeepSpeed，HazyResearch的flash-attn等。