nnScaler 是一个并行化引擎，它将专为单 GPU 执行而设计的深度神经网络 (DNN) 模型编译为能够跨多个 GPU 并行运行的程序。
易于使用：只需更改几行代码即可实现自动并行化。
Pythonic：并行化输出为PyTorch代码，方便用户理解，方便进一步开发或定制。
可扩展性：nnScaler 公开 API 来支持新兴模型的新运算符。
可靠性：通过各种端到端培训课程验证，nnScaler 是一个可靠的系统。
性能：通过探索大的并行化空间，nnScaler 可以显着增强并行训练性能。
对于DNN 科学家来说，他们可以在单个 GPU 上使用 PyTorch 专注于模型设计，同时将并行化复杂性留给 nnScaler。它引入了创新的并行技术，在性能上超越了现有方法。此外，nnScaler 支持使用新结构或执行模式扩展 DNN 模块，使用户能够并行化其自定义 DNN 模型。
对于DNN 系统专家来说，他们可以利用 nnScaler 探索新兴模型的新 DNN 并行化机制和策略。通过为 nnScaler 无法识别的新运算符提供用户定义的函数，它确保了新型 DNN 模型的无缝并行化。例如，促进法学硕士中的长序列支持。
在安装nnScaler之前安装以下软件包：
Python >= 3.8, < 3.11 (3.10 is recommanded) PyTorch >= 2.0, < 2.4 (2.2.0 is recommanded)
在nnScaler目录中执行以下命令：
pip install -r requirements.txt pip install -e .
此外，为了避免cppimport错误，还需要在环境变量PYTHONPATH中包含 nnScaler 目录：
export NNSCALER_HOME=$(pwd)
export PYTHONPATH=${NNSCALER_HOME}:$PYTHONPATH安装运行 Llama-3 所需的软件包。此外，安装flash-attn时还需要特定版本的CUDA库。例如，如果使用 PyTorch 2.20，则需要 CUDA V11.8。
python -m pip install transformers==4.40.0 flash-attn==2.5.5 tensorboard
从 HuggingFace 获取 Llama-3 模型的访问权限，您将收到一个访问令牌，该令牌应设置为环境变量：
export HF_TOKEN=<HUGGINGFACE_ACCESS_TOKEN>
您可以在examples/llama3_8B_128K找到所有示例代码。如下图所示，用户需要：
包装模型：包括损失计算和其他必要的组件。
配置组件：设置模型、优化器和数据加载器。
初始化并开始：在主函数中，使用上述配置创建一个nnScaler训练器并开始训练过程。
# 导入 nnScaler 内置并行化训练器 from nnscaler.cli.trainer import Trainer# 包装模型以包含损失计算等。class WrapperModel(torch.nn.Module):def __init__(self, model_id):super(). __init__()self.model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation='flash_attention_2')defforward(self,samples):outputs = self.model.model(input_ids=samples['net_input']['src_tokens'],use_cache=False,return_dict=False,
)loss = torch.sum(chunk_linear_cross_entropy(outputs[0], self.model.lm_head.weight,样本['target'],...))返回损失,样本['ntokens'],样本['nsentences'] def main(args):# 数据配置dataloader_config = ... # 模型配置model_config = ModelConfig(type=WrapperModel,args={'model_id': args.model_id,
},
)# 优化器超参数 optimizationr_config = OptimizerConfig(type=MixedPrecisionAdamW,args={'lr': 2e-5, 'betas': (0.9, 0.95), 'weight_decay': 0.0, 'fused': True},#... )#...# 使用数据加载器/模型/优化器等配置设置训练器。 trainer =训练器(train_args=TrainerArgs(#...model=model_config,optimizer=optimizer_config,dataloader=dataloader_config,#...))trainer.run()然后我们就可以启动示例了，所有的并行化任务都会由nnScaler自动完成。
cd Examples/llama3_8B_128K# 准备训练数据:python bookcorpus.py --data_path_or_name bookcorpus/bookcorpus --tokenizer_path_or_name meta-llama/Meta-Llama-3-8B-Instruct --save_path ./bookcorpus_llama3_4K --sequence_length 4096# 构建迷你模型python create_mini_model.py --model_id meta-llama/Meta-Llama-3-8B-Instruct --output_id ./llama3_mini#使用数据并行编译并运行 + Zero1torchrun --nproc_per_node=2 train.py --plan_ngpus 1 --runtime_ngpus 2 --name llama3_debug -- model_id ./llama3_mini --dataset_path ./bookcorpus_llama3_4K
我们还提供了一个示例来演示如何通过 nnScaler 中的 PyTorch Lightning 兼容接口并行化模型。
在 nnScaler 存储库中找到 nanoGPT 示例：
CD 示例/nanogpt
安装nanoGPT的依赖项：
pip install -r 要求.txt
准备数据集：
python nanoGPT/data/shakespeare_char/prepare.py
使用单 GPU 进行测试
现在您可以使用torchrun <https://pytorch.org/docs/stable/elastic/run.html>运行train_nnscaler.py ：
torchrun --nproc_per_node=1 train_nnscaler.py nanoGPT/config/train_shakespeare_char.py
这将在单个 GPU 上训练一个婴儿 GPT 模型。这将需要几分钟的时间，最佳验证损失将在 1.47 左右。
使用多 GPU 进行测试
默认情况下，nnScaler 通过 GPU 与数据并行性并行化模型。如果一个节点上有 4 个 GPU：
torchrun --nproc_per_node=4 train_nnscaler.py nanoGPT/config/train_shakespeare_char.py
或者，如果您有多个节点，例如 2 个节点，每个节点有 4 个 GPU：
# on each node torchrun --nnodes=2 --nproc_per_node=4 --rdzv-id=NNSCALER_NANOGPT --rdzv-backend=c10d --rdzv-endpoint=<IP> train_nnscaler.py nanoGPT/config/train_shakespeare_char.py
注意：本地批量大小默认是固定的，因此使用更多的工作线程将导致更大的全局批量大小。
？对于高级用法，请继续关注我们未来的版本。
nnScaler 已被多个项目采用，包括产品和研究探索：
(YOCO)你只缓存一次：语言模型的解码器-解码器架构
LongRoPE：将 LLM 上下文窗口扩展至超过 200 万个令牌
Phi-3 系列长上下文版本的后期培训
您可以在此处找到 OSDI'24 的工件评估和指导。如果 nnScaler 对您的研究有帮助，请在您的出版物中引用它：
@inproceedings{lin2024nnscaler,
title = {nnScaler: Constraint-Guided Parallelization Plan Generation for Deep Learning Training},
author={Lin, Zhiqi and Miao, Youshan and Zhang, Quanlu and Yang, Fan and Zhu, Yi and Li, Cheng and Maleki, Saeed and Cao, Xu and Shang, Ning and Yang, Yilei and Xu, Weijiang and Yang, Mao and Zhang, Lintao and Zhou, Lidong},
booktitle={18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24)},
pages={347--363},
year={2024}
}该项目欢迎贡献和建议。大多数贡献都要求您同意贡献者许可协议 (CLA)，声明您有权并且实际上授予我们使用您的贡献的权利。有关详细信息，请访问 https://cla.opensource.microsoft.com。
当您提交拉取请求时，CLA 机器人将自动确定您是否需要提供 CLA 并适当地修饰 PR（例如，状态检查、评论）。只需按照机器人提供的说明进行操作即可。您只需使用我们的 CLA 在所有存储库中执行一次此操作。
该项目采用了微软开源行为准则。有关详细信息，请参阅行为准则常见问题解答或联系 [email protected] 提出任何其他问题或意见。
该项目可能包含项目、产品或服务的商标或徽标。 Microsoft 商标或徽标的授权使用须遵守且必须遵循 Microsoft 的商标和品牌指南。在此项目的修改版本中使用 Microsoft 商标或徽标不得引起混淆或暗示 Microsoft 赞助。对第三方商标或徽标的任何使用均须遵守这些第三方的政策。
您可以从 https://github.com/microsoft/nnscaler 或 Microsoft 内部存储库 https://aka.ms/ms-nnscaler 找到我们的公共存储库。如有任何疑问或疑问，请通过 [email protected] 与我们联系。
