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# 包裝模型以包含損失計算等。 self.model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation='flash_attention_2')defforward(self,samples):outputs = self.model.model(input_ids=samples)： =假，
)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},#. .. )#...# 使用 dataloader/model/optimizer 等配置設定 trainer trainer = Trainer(train_args=TrainerArgs(#...model=model_config,optimizer=optimizer_config,dataloader=dataloader_config,#...))trainer 。跑步（）然後我們就可以啟動範例了，所有的平行化任務都會由nnScaler自動完成。
。 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 --run__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] 與我們聯絡。
