分佈式機器學習 — 系統、工程與實戰

本書主要講解分佈式機器學習算法和開源框架，讀者既可以從宏觀的設計上瞭解分佈式機器學習的概念和理論，也可以深入核心技術的細節設計中，對分佈式機器學習形成深刻而直觀的認識，做到學以致用。 本書共分為5篇，第1篇是分佈式基礎，首先介紹了分佈式機器學習的概念、基礎設施，以及機器學習並行化技術、框架和軟件系統，然後對集合通信和參數服務器PS-Lite進行了介紹。第2篇是數據並行，以PyTorch和Horovod為主對數據並行進行分析，讀者可以瞭解在具體工程領域內實現數據並行有哪些挑戰和解決方案。第3篇是流水線並行，講解了除模型劃分之外，還通過引入額外的流水線來提高效率，以GPipe / PyTorch / PipeDream為例進行分析。第4篇是模型並行，首先對NVIDIA Megatron進行分析，講解如何進行層內分割模型並行，然後講解PyTorch 如何支持模型並行，最後介紹分佈式優化器。第5篇是TensorFlow分佈式，前面幾篇以PyTorch為綱，結合其他框架/庫來穿插完成，本篇帶領大家進入TensorFlow分佈式領域。

目 錄

第1篇 分佈式基礎
第1章 分佈式機器學習 2
1.1 機器學習概念 2
1.2 機器學習的特點 3
1.3 分佈式訓練的必要性 3
1.4 分佈式機器學習研究領域 6
1.4.1 分佈式機器學習的目標 6
1.4.2 分佈式機器學習的分類 6
1.5 從模型角度看如何並行 8
1.5.1 並行方式 8
1.5.2 數據並行 9
1.5.3 模型並行 10
1.5.4 流水線並行 11
1.5.5 比對 12
1.6 從訓練並發角度看如何並行 12
1.6.1 參數分佈和通信拓撲 13
1.6.2 模型一致性和通信模式 14
1.6.3 訓練分佈 19
1.7 分佈式機器學習編程接口 19
1.7.1 手動同步更新 20
1.7.2 指定任務和位置 20
1.7.3 猴子補丁優化器 21
1.7.4 Python作用域 21
1.8 PyTorch分佈式 22
1.8.1 歷史脈絡 22
1.8.2 基本概念 23
1.9 總結 24
第2章 集合通信 26
2.1 通信模式 26
2.2 點對點通信 26
2.3 集合通信 28
2.3.1 Broadcast 29
2.3.2 Scatter 29
2.3.3 Gather 30
2.3.4 All-Gather 30
2.3.5 All-to-All 30
2.3.6 Reduce 31
2.3.7 All-Reduce 31
2.3.8 Reduce-Scatter 32
2.4 MPI_AllReduce 32
2.5 Ring All-Reduce 33
2.5.1 特點 34
2.5.2 策略 34
2.5.3 結構 35
2.5.4 Reduce-Scatter 35
2.5.5 All-Gather 38
2.5.6 通信性能 40
2.5.7 區別 40
第3章 參數服務器之PS-Lite 41
3.1 參數服務器 41
3.1.1 概念 41
3.1.2 歷史淵源 42
3.1.3 問題 43
3.2 基礎模塊Postoffice 44
3.2.1 基本邏輯 44
3.2.2 系統啟動 45
3.2.3 功能實現 47
3.3 通信模塊Van 51
3.3.1 功能概述 51
3.3.2 定義 51
3.3.3 初始化 52
3.3.4 接收消息 53
3.4 代理人Customer 59
3.4.1 基本思路 59
3.4.2 基礎類 61
3.4.3 Customer 62
3.4.4 功能函數 66
3.5 應用節點實現 67
3.5.1 SimpleApp 67
3.5.2 KVServer 68
3.5.3 KVWorker 68
3.5.4 總結 70
第2篇 數據並行 73
第4章 PyTorch DataParallel 74
4.1 綜述 74
4.2 示例 76
4.3 定義 77
4.4 前向傳播 78
4.5 計算損失 87
4.6 反向傳播 88
4.7 總結 91
第5章 PyTorch DDP的基礎架構 93
5.1 DDP總述 93
5.1.1 DDP的運行邏輯 93
5.1.2 DDP的使用 94
5.2 設計理念 97
5.2.1 系統設計 97
5.2.2 梯度歸約 98
5.2.3 實施 99
5.3 基礎概念 101
5.3.1 初始化方法 101
5.3.2 Store類 102
5.3.3 TCPStore類 104
5.3.4 進程組概念 107
5.3.5 構建進程組 109
5.4 架構和初始化 111
5.4.1 架構與迭代流程 111
5.4.2 初始化DDP 114
第6章 PyTorch DDP的動態邏輯 122
6.1 Reducer類 122
6.1.1 調用Reducer類 122
6.1.2 定義Reducer類 122
6.1.3 Bucket類 124
6.1.4 BucketReplica類 126
6.1.5 查詢數據結構 128
6.1.6 梯度累積相關成員變量 131
6.1.7 初始化 135
6.1.8 靜態圖 141
6.1.9 Join操作 142
6.2 前向/反向傳播 143
6.2.1 前向傳播 143
6.2.2 反向傳播 149
第7章 Horovod 161
7.1 從使用者角度切入 161
7.1.1 機制概述 161
7.1.2 示例代碼 162
7.1.3 運行邏輯 163
7.2 horovodrun 167
7.2.1 入口點 167
7.2.2 運行訓練Job 168
7.2.3 Gloo實現 169
7.2.4 MPI實現 174
7.2.5 總結 174
7.3 網絡基礎和Driver 174
7.3.1 總體架構 175
7.3.2 基礎網絡服務 176
7.3.3 Driver服務 177
7.3.4 Task服務 178
7.3.5 總結 180
7.4 DistributedOptimizer 181
7.4.1 問題點 181
7.4.2 解決思路 182
7.4.3 TensorFlow 1.x 183
7.5 融合框架 191
7.5.1 總體架構 191
7.5.2 算子類體系 192
7.5.3 後台線程 194
7.5.4 執行線程 195
7.5.5 總結 196
7.6 後台線程架構 198
7.6.1 設計要點 198
7.6.2 總體代碼 201
7.6.3 業務邏輯 202
第3篇 流水線並行
第8章 GPipe 210
8.1 流水線基本實現 210
8.1.1 流水線並行 210
8.1.2 GPipe概述 211
8.1.3 計算內存 213
8.1.4 計算算力 213
8.1.5 自動並行 214
8.2 梯度累積 218
8.2.1 基本概念 218
8.2.2 PyTorch實現 219
8.2.3 GPipe實現 223
8.3 Checkpointing 225
8.3.1 問題 225
8.3.2 解決方案 225
8.3.3 OpenAI 226
8.3.4 PyTorch實現 228
8.3.5 GPipe實現 240
第9章 PyTorch流水線並行 243
9.1 如何劃分模型 243
9.1.1 使用方法 244
9.1.2 自動平衡 245
9.1.3 模型劃分 247
9.2 切分數據和Runtime系統 249
9.2.1 分發小批量 249
9.2.2 Runtime 250
9.3 前向計算 255
9.3.1 設計 255
9.3.2 執行順序 260
9.4 計算依賴 265
9.4.1 反向傳播依賴 266
9.4.2 前向傳播依賴 270
9.5 並行計算 274
9.5.1 總體架構 274
9.5.2 並行復制和計算 276
9.5.3 重計算 278
第10章 PipeDream之基礎架構 280
10.1 總體思路 280
10.1.1 目前問題 280
10.1.2 1F1B策略概述 282
10.1.3 流水線方案 283
10.2 profile階段 285
10.3 計算分區階段 288
10.3.1 構建圖 288
10.3.2 構建反鏈 289
10.3.3 計算分區 295
10.3.4 分析分區 302
10.3.5 輸出 305
10.4 轉換模型階段 305
10.4.1 分離子圖 306
10.4.2 轉換模型 307
10.4.3 融合模型 308
第11章 PipeDream之動態邏輯 312
11.1 Runtime引擎 312
11.1.1 功能 312
11.1.2 總體邏輯 313
11.1.3 加載模型 314
11.1.4 實現 314
11.2 通信模塊 323
11.2.1 類定義 324
11.2.2 構建 325
11.2.3 發送和接收 331
11.3 1F1B策略 333
11.3.1 設計思路 333
11.3.2 權重問題 335
11.3.3 實現 340
第4篇 模型並行
第12章 Megatron 346
12.1 設計思路 346
12.1.1 背景 346
12.1.2 張量模型並行 348
12.1.3 並行配置 354
12.1.4 結論 354
12.2 模型並行實現 354
12.2.1 並行MLP 355
12.2.2 ColumnParallelLinear 358
12.2.3 RowParallelLinear 363
12.2.4 總結 367
12.3 如何設置各種並行 367
12.3.1 初始化 368
12.3.2 起始狀態 371
12.3.3 設置張量模型並行 373
12.3.4 設置流水線並行 375
12.3.5 設置數據並行 378
12.3.6 模型並行組 380
12.3.7 如何把模型分塊到GPU上 381
12.4 Pipedream的流水線刷新 383
第13章 PyTorch如何實現模型並行 387
13.1 PyTorch模型並行 387
13.1.1 PyTorch特點 387
13.1.2 示例 387
13.2 分佈式自動求導之設計 389
13.2.1 分佈式RPC框架 389
13.2.2 自動求導記錄 390
13.2.3 分佈式自動求導上下文 391
13.2.4 分佈式反向傳播算法 392
13.2.5 分佈式優化器 396
13.3 RPC基礎 396
13.3.1 RPC代理 396
13.3.2 發送邏輯 396
13.3.3 接收邏輯 398
13.4 上下文相關 399
13.4.1 設計脈絡 400
13.4.2 AutogradMetadata 401
13.4.3 DistAutogradContainer 403
13.4.4 DistAutogradContext 403
13.4.5 前向傳播交互過程 408
13.5 如何切入引擎 411
13.5.1 反向傳播 411
13.5.2 SendRpcBackward 415
13.5.3 總結 417
13.6 自動求導引擎 417
13.6.1 原生引擎 417
13.6.2 分佈式引擎 419
13.6.3 總體執行 421
13.6.4 驗證節點和邊 421
13.6.5 計算依賴 422
13.6.6 執行GraphTask 429
13.6.7 RPC調用閉環 433
13.6.8 DistAccumulateGradCapture-
Hook 436
13.6.9 等待完成 442
第14章 分佈式優化器 443
14.1 原生優化器 443
14.2 DP的優化器 445
14.3 DDP的優化器 446
14.3.1 流程 446
14.3.2 優化器狀態 446
14.4 Horovod的優化器 447
14.4.1 利用鉤子同步梯度 448
14.4.2 利用step()函數同步梯度 449
14.5 模型並行的分佈式問題 450
14.6 PyTorch分佈式優化器 451
14.6.1 初始化 452
14.6.2 更新參數 453
14.7 PipeDream分佈式優化器 455
14.7.1 如何確定優化參數 456
14.7.2 優化 458
第5篇 TensorFlow分佈式
第15章 分佈式運行環境之靜態架構 462
15.1 總體架構 462
15.1.1 集群角度 462
15.1.2 分佈式角度 463
15.1.3 系統角度 465
15.1.4 圖操作角度 467
15.1.5 通信角度 468
15.2 Server 469
15.2.1 邏輯概念 469
15.2.2 GrpcServer 471
15.3 Master的靜態邏輯 474
15.3.1 總述 474
15.3.2 接口 474
15.3.3 LocalMaster 476
15.3.4 GrpcRemoteMaster 478
15.3.5 GrpcMasterService 478
15.3.6 業務實現Master類 480
15.4 Worker的靜態邏輯 481
15.4.1 邏輯關系 481
15.4.2 GrpcRemoteWorker 483
15.4.3 GrpcWorkerService 483
15.4.4 Worker 487
15.4.5 GrpcWorker 488
第16章 分佈式運行環境之動態邏輯 489
16.1 Session機制 489
16.1.1 概述 489
16.1.2 GrpcSession 491
16.1.3 MasterSession 492
16.1.4 WorkerSession 494
16.2 Master動態邏輯 495
16.2.1 Client如何調用 495
16.2.2 Master業務邏輯 495
16.3 Worker動態邏輯 501
16.3.1 概述 501
16.3.2 註冊子圖 501
16.3.3 運行子圖 502
16.3.4 分佈式計算流程總結 504
16.4 通信機制 505
16.4.1 協調機制 505
16.4.2 發送流程 508
16.4.3 接收流程 508
16.4.4 總結 509
第17章 分佈式策略基礎 511
17.1 使用TensorFlow進行分佈式訓練 511
17.1.1 概述 511
17.1.2 策略類型 511
17.2 DistributedStrategy基礎 515
17.2.1. StrategyBase 515
17.2.2 讀取數據 518
17.3 分佈式變量 523
17.3.1 MirroredVariable 523
17.3.2 ShardedVariable 530
第18章 MirroredStrategy 535
18.1 MirroredStrategy集合通信 535
18.1.1 設計思路 535
18.1.2 實現 536
18.1.3 更新分佈式變量 538
18.2 MirroredStrategy分發計算 540
18.2.1 運行 540
18.2.2 mirrored_run 541
18.2.3 Context 544
18.2.4 通信協議 546
18.2.5 EagerService 547
18.2.6 在遠端運行訓練代碼 551
18.2.7 總結 552
第19章 ParameterServerStrategy 554
19.1 ParameterServerStrategyV1 554
19.1.1 思路 554
19.1.2 數據 556
19.1.3 作用域和變量 557
19.1.4 運行 559
19.2 ParameterServerStrategyV2 560
19.2.1 如何使用 560
19.2.2 運行 561
19.3 ClusterCoordinator 561
19.3.1 使用 561
19.3.2 定義 563
19.3.3 數據 565
19.3.4 Cluster 566
19.3.5 Closure 568
19.3.6 隊列 570
19.3.7 Worker類 570
19.3.8 Failover 573
19.3.9 總結 574