Skip to content

CMU 11-868 HW5：在兩張 GPU 上自己寫資料平行與管線平行

2026年9月30日1 分鐘
首頁›ai
TL;DRCMU 11-868 第五份作業改用 PyTorch 與 Hugging Face 的 GPT-2，要你只用 torch.distributed 和 torch.multiprocessing 寫出資料平行（切資料、建 process group、平均梯度，50 分），再寫出 GPipe 式的管線平行（切模型、產生時脈排程、用 worker thread 跑 micro-batch，50 分）。兩部分都要在至少兩張 GPU 上做 benchmark 並交圖：資料平行要在 2 張卡上達到至少 1.5 倍加速，管線平行要比單純模型平行快。2026 春季版 3/25 截止。
#cmu-11868#cmu#ai-course#homework#distributed-training#pytorch#gpu
系列: CMU 11-868 LLM Systems 導讀 (14 / 23)
CMU 11-868 L18：ZeRO 怎麼把資料平行的記憶體切掉——optimizer state、梯度、參數各切一刀CMU 11-868 L19–L20 模型量化：GPTQ 省下的是記憶體，速度是順便換來的

🌏 English version

本文依據 CMU 11-868 LLM Systems 的 2026 春季版。 這是 CMU 11-868 LLM Systems 導讀系列的第 14 篇。作業頁與 repo 依 2026-09-30 所見；作業站是跨學期共用的，之後可能被 Fall 2026 修改。

這份作業把第 11 到 13 篇講的東西拿來動手：資料平行（L14–L15）和管線平行（L16 上半）。上一篇的 ZeRO 不在這份作業裡，要到 HW6 才用 DeepSpeed 實際跑。

基本資料

項目內容
作業頁Assignment 5: Distributed Training and Parallelism
起始碼頁面連到 llmsys_f25_hw5，會 301 轉到 llmsystem/llmsys_hw5
截止2026 春季版 3/25（Syllabus 第 11 週）；發放日 Syllabus 沒寫
配分Problem 1 Data Parallel 50、Problem 2 Pipeline Parallel 50
硬體至少 2 張 GPU；作業頁強烈建議用 PSC
環境conda 建 Python 3.9 環境；requirements.txt 固定 torch 2.2.0、transformers 4.37.2、datasets 3.6.0
對應講次L14–L15（資料平行）、L16（pipeline）；3/20 的 Recitation 6 主題是 Distributed Training

和前四份作業最大的不同：這份作業不再用 MiniTorch。 起始碼的 project/run_data_parallel.py 直接載入 Hugging Face 的預訓練 GPT2LMHeadModel，資料是 bbaaaa/iwslt14-de-en-preprocess，訓練集只取前 5,000 筆，預設跑 10 個 epoch。重點從「自己寫框架」換成「自己寫分散式的那一層」。

規則：只能用兩個套件做通訊

Problem 1 限定只能用 torch.distributed 和 torch.multiprocessing.Process 做 GPU 通訊，不准新增 import。Problem 2 只能用起始碼裡已經 import 的套件。所有要實作的地方都用 BEGIN_HW5_* 和 END_HW5_* 註解標起來，修改必須留在標記區塊內。作業頁說會人工檢查程式碼是否遵守套件限制。

Problem 1：資料平行（50 分）

1.1 切資料。 在 data_parallel/dataset.py 實作三樣東西：

  • Partition：依一組索引回傳資料的 dataset 類別
  • DataPartitioner：把索引打亂後依比例切成幾份（例如四張卡時是 [0.25, 0.25, 0.25, 0.25]），use(rank) 取出某一份
  • partition_dataset：算出每張卡的 batch size（總 batch 128、四張卡時每張 32），取出本 rank 的那份資料，包成 DataLoader

測試 a5_1_1 只檢查各份資料有沒有重疊。

1.2 建 process group、平均梯度。 在 project/run_data_parallel.py 實作 setup：把 MASTER_ADDR 設為 localhost、MASTER_PORT 設為 11868，呼叫 init_process_group；主程式要依 world_size 建立、啟動並正確收掉多個 process。再實作 average_gradients，走過模型所有參數，用 torch.distributed 的函式把梯度彙總，並在 project/utils.py 的 train 裡 backward 之後呼叫它。

測試方式是先用 world_size 2 跑一個 batch，把每個 rank 的梯度存成 model{rank}_gradients.pth，再用 a5_1_2 比較兩張卡的梯度是否一致。

1.3 benchmark。 比較單卡（batch 64）和兩張卡（總 batch 128）的訓練時間與 tokens per second。作業頁規定的算法：訓練時間取多張卡各自的平均；吞吐量把各卡的 tokens per second 相加。建議丟掉第一個 epoch 或至少先暖機一次。兩個指標各畫一張圖存進 submit_figures。在 2 張 GPU 上達到至少 1.5 倍加速（訓練時間與吞吐量）即得滿分。

Problem 2：管線平行（50 分）

這部分實作的是 L16 講的 GPipe 式排程：把模型按層切到不同 GPU，把輸入切成 micro-batch，讓各 stage 同時處理不同份。

2.1 切模型與排程。 在 pipeline/partition.py 實作 _split_module，把一個 nn.Sequential 按層切成幾段、每段放一張 GPU。在 pipeline/pipe.py 實作 _clock_cycles(num_batches, num_partitions)，產生每個時間步要由哪個 stage 處理哪個 micro-batch 的排程。測試 a5_2_1。

2.2 Pipe 模組。 先讀懂 worker.py，再實作 Pipe.forward 和 Pipe.compute。Pipe 是一個通用包裝，能把任何 nn.Sequential 變成管線化模組。起始碼在 create_workers 裡替每張 GPU 開一個 worker thread，各自從 in_queue 取工作、把結果放進 out_queue；你要把計算包成 Task 放進對應裝置的佇列，再取回結果。作業頁特別提醒：forward 的結果要放在最後一張卡上，放在輸入 x 所在的卡會讓訓練失敗。測試 a5_2_2。

2.3 接上 GPT-2。 在 pipeline/model_parallel.py 實作 _prepare_pipeline_parallel：GPT-2 的各個 block 已經由 GPT2ModelCustom.parallelize 分到不同 GPU，你要把 self.h 裡的 transformer block 取出來包成 nn.Sequential 交給 Pipe。作業頁提醒 GPT2Block 回傳的是 tuple 而不是張量，只需要 hidden states；若自己加了沒有參數的輔助模組，要用 WithDevice 包起來，否則它會被判定在 CPU 上。

最後分別用 --model_parallel_mode='model_parallel' 和 'pipeline_parallel' 在兩張卡上訓練，畫圖比較。管線平行的訓練時間更短、吞吐量更高即得滿分。

繳交與評分

作業頁提供 scripts/create_submission_zip.sh 打包，以及 scripts/run_benchmarks.py 產生效能紀錄、submit_figures/performance_summary.json 與圖。助教會編譯執行程式碼、檢查圖，並人工檢查套件限制。

作業頁與起始碼的不一致

2026-09-30 對照作業頁與 llmsys_hw5 main 分支時，看到三處對不上：

  1. partition_dataset 的簽名：作業頁寫 partition_dataset(dataset, batch_size=128, collate_fn=None)，起始碼是 partition_dataset(rank, world_size, dataset, batch_size=128, collate_fn=None)。以起始碼為準。
  2. 繳交檔名：Submission 一節寫「submit the whole llmsys_s25_hw5 as a zip on canvas」，還是 2025 春季的 repo 名。
  3. 加速門檻：Problem 1.3 寫訓練時間與吞吐量都要 1.5 倍，但 grader 模式的範例指令是 --dp-time-threshold 1.2 --dp-throughput-threshold 1.5（管線平行兩項都是 1.0）。實際評分用哪個值，官方頁面沒有說明。

另外，main 分支最後一筆 commit 是 2026-04-30，仍是春季狀態；repo 另有一個 merged-hw5-hw6 分支。Fall 2026 若修改作業，可能出現在那裡或 main 上。要對齊春季版，checkout 2026-03-25 之前的 commit 最保險。

校外讀者會卡在哪

  • 兩張 GPU：這是整個系列第一份硬性要求多卡的作業。PSC 是 Logistics 寫明提供給修課學生的運算資源，校外讀者要自己租雲端的雙卡機器；Logistics 也提醒 PSC 採排程制，不保證工作何時開始跑
  • 1.5 倍加速不是自動的：兩張卡的梯度同步有成本（L14–L15 講的就是這個），單卡 batch 64 對兩卡總 batch 128 的設定下，加速多少取決於你的互連與實作
  • 沒有 Canvas 與人工評分：測試只驗證正確性（資料不重疊、梯度一致、排程正確），效能與套件限制要自己檢查

怎麼做：沒修課也能練的三件事

  1. 先在 CPU 上跑通資料平行：init_process_group 用 gloo backend、兩個 process，就能在筆電上驗證 DataPartitioner 和 average_gradients 的正確性，再上 GPU 量加速。
  2. 手寫一次時脈排程：在紙上畫 3 個 stage、4 個 micro-batch 的 GPipe forward 排程，每個時間步寫出 (micro-batch, stage) 配對，再和 _clock_cycles 的測試對照。
  3. 量一次 bubble：管線平行跑通後，用 run_pipeline.py 的 --n_chunk（micro-batch 數，預設 4）改變切法，記錄吞吐量怎麼變，對照 L16 的 bubble 公式 O((K−1)/(M+K−1))。

延伸閱讀

參考資料

相關文章

CMU 11-868 L14–L15：分散式訓練與資料平行，梯度同步的代價從哪裡來1 minCMU 11-868 作業一：用 CUDA 寫 MiniTorch 的 map、zip、reduce 與 matmul1 minCMU 11-868 L18：ZeRO 怎麼把資料平行的記憶體切掉——optimizer state、梯度、參數各切一刀1 min