本文參考 maskrcnn-benchmark 和 MMAction2 (MMCV) 中的實現方式，大概是一年前寫的。在最近整理時還發現一篇 Distributed data parallel training in Pytorch，參考了一下，就如這篇裡面所述，即便現在 PyTorch 官網都沒有一個清晰的多卡訓練教程，只能自己看 code。

至於混合精度，之前 NVIDIA 出了 apex，在目前 PyTorch 1.6 中已經把 AMP (Automatic mixed precision) 整合進去了，可以很方便調用，本文只是簡述一下原理。

【2021 更新】最近看到 open-mmlab 寫了一系列文，關於 PyTorch 內部實現的，會更關注 PyTorch 的代碼實現：

多 GPU 和 Process 支持

目前譬如 Object Detection，在 inference 單卡一般可行，train 的話數據一多就會很慢。

原理

nn.DataParallel

早期 PyTorch 中多 GPU 訓練的方式一般為使用 torch.nn.DataParallel()（或 torch.multiprocessing），只需 model = nn.DataParallel(model).cuda()。Model 首先被加載到主 GPU 上，然後複製到其它 GPU 中（DataParallel，多線程）。輸入數據按 batch 維度進行劃分，每個 GPU 分配到的 batch 數量等於輸入總的 batch 除以 GPU 個數（不一定需要整除）。每個 GPU 對各自的輸入數據獨立進行 forward 計算，return 之前又被 DataParallel concat 在一起，因此後續計算 loss 等無需額外 code。loss/gradient 計算後，更新主 GPU 上的模型參數，DataParallel 再將更新後的模型參數複製到其餘 GPU 中，這樣就完成了一次迭代計算。

nn.DistributedDataParallel

Mask R-CNN 的訓練使用 PyTorch 1.0 新加入的多進程方式，參考 GitHub/torch.distributed.launch 和 PyTorch/Docs。根據官方文檔中的這部分，這種方式有以下好處：

forward 計算之後不合併 batch，每個 GPU 上對自己的 batch 單獨計算 loss/gradient，但隨後梯度會在各個進程之間合併/計算總和（synchronize），而每個進程都有自己的 optimizer，獨立地進行參數更新，沒有了 parameters/batch tensor 在各 GPU 之間傳輸的時間。
每個 GPU 都有獨立的 Python interpreter，避免了 Python 在多線程時的問題，對於 Python 比 C++ 更影響運行時間的模型幫助很大（比如有很多小模塊組成的複雜模型）
可以用於多機 (node) 多 GPU 訓練

總結，這種方法會給每個 GPU 啟動一個獨立的進程，這個進程獨立地運行著所有的代碼，它們之間的通信可以只有 Gradient 的傳輸（當然比如為了 logging 的需要，有時候會寫一點傳輸 loss 值的代碼）。為了達到上面的目標：

首先，需要讓每個進程知道自己的 rank 和 world_size，也就是第幾個進程和總共有幾個進程。對應 torch.distributed.launch
還需要一個 Data Sampler，它拿到 rank 之後，給每個進程中的 model，在一個 unique (non-overlap) 的數據集的部分中 sample，使得每個進程處理自己的一部分數據。nn.utils.data.DistributedSampler 就是幹這事的
還需要 Gradients 之間的平均，處理通信和所有 param 的平均值計算之類的，使得每次 backward 之後所有 GPU 上的梯度都是一個相同的平均值（即 all-reduce）。nn.DistributedDataParallel 就是一個 model wrapper 負責這件事情
最後，需要設定隨機數種子，不然 model 初始化的時候參數都不一樣，梯度平均沒意義

如果對 Data Sampler 不了解，應該看前一篇，鏈接在開頭。如果對 all-reduce 想了解更多，見下文。

使用

對於單機多 GPU，啟動方式為使用 torch.distributed.launch，只需要 --nproc_per_node=$NGPUS，後方跟訓練腳本位置和它的參數：

1 2	export NGPUS=4 python -m torch.distributed.launch --nproc_per_node=$NGPUS train_net.py --config-file "config.yaml"

torch.distributed.launch 會在 os.environ 中添加 ["WORLD_SIZE"] = $NGPUS 並使用 subprocess.Popen() 啟動 NGPUS 次訓練腳本（並在參數中傳入不同的 Rank）並捕獲 return code。

在訓練腳本內，需要在開頭加上下面的代碼獲得 torch.distributed.launch 傳入的 local_rank。

# parser.add_argument ...
# 接收分配給這個進程的 process rank
parser.add_argument("--local_rank", type=int, default=0)
args = parser.parse_args()

num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
args.distributed = num_gpus > 1

if args.distributed:
    # 設定執行當前进程的 GPU 并初始化 group，nccl 作為通信後端，需要先裝好
    # 註：單 GPU 訓練中指定使用的 GPU，
    # 官方更推薦 `os.environ["CUDA_VISIBLE_DEVICES"] = "0"`
    rank = args.local_rank  # or: rank = int(os.environ['RANK'])
    torch.cuda.set_device(rank)  # rank % torch.cuda.device_count()
    torch.distributed.init_process_group(
        backend="nccl", init_method="env://"
    )
    # synchronize()

init_process_group 註冊了一系列信息，包括如何找到 Process 0，其中同步的 synchronize() 主要是 dist.barrier() 使執行到這裡的進程等待，當所有進程都執行到這裡再繼續往下。

import torch.distributed as dist

def synchronize():  # get_dist_info()
    if not dist.is_available():
        return
    if not dist.is_initialized():
        return
    else:
        rank = dist.get_rank()
        world_size = dist.get_world_size()
    if world_size == 1:
        return
    dist.barrier()
    return rank, world_size

init_process_group (dist.is_initialized()) 之後可以用 dist.get_world_size() dist.get_rank() 等方法獲得需要信息，不需要依靠 args 了。

在實例化完 model 後使用

# model = model.to(f'cuda:{local_rank}')
model = nn.parallel.DistributedDataParallel(
    model, device_ids=[local_rank], output_device=local_rank
)

對所有參數註冊一個 gradient reduction functions 用於求和梯度。這樣操作時候，每個 Process 獨立進行 forward，在 backward 的時候，gradient 會自動在所有 GPU 之間 all-reduce，且由於 backward 是單向依賴和進行的，所以這個通信和 backward 是同時進行的，進一步加快速度和減少帶寬壓力。

然後配置 DataLoader 的 Sampler，使不同進程的 model 獲得不同的 input batch。

train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset,
    num_replicas=args.world_size,
    rank=rank,
    shuffle=True,
)

train_loader = torch.utils.data.DataLoader(
    dataset=train_dataset,
    batch_size=batch_size_per_gpu,
    shuffle=False,  # shuffle should set in sampler if dist
    sampler=train_sampler
)

這一部分只適用與單機多卡 (single node multi gpu)，多幾多卡似乎還要額外對 backward 和 optimizer 進行配置，然後加入 os.environ['MASTER_ADDR'] 和 os.environ['MASTER_PORT'] 設置 node 0 的 ip 和 port，slurm 啟動等，我沒用過，不知道具體。

注意事項

learning rate

每個 GPU 的 memory 是一定的，所以 batch size per GPU 的上限是一定的，如果在 NGPUS = 2 batch size = 2 訓練了模型得到一個精度結果，那麼想在 NGPUS = 1 的情況下使用相同 config 得到基本相同結果，需要進行改動：batch size = 2，但是只有一個 GPU，那麼等同於總 batch size 除以了 2，所以總的 iteration 也需要 ×2，lr 也要減小 N 倍，因為 grad 小 N 倍。

reduce/gather

雖然 grad 被自動在所有 GPU 求和了，但是為了 log 的需要，一般還要將每個 process 中得到的 loss 傳輸到 process 0，求得平均後輸出。WRITING DISTRIBUTED APPLICATIONS WITH PYTORCH 中介紹了各種 Collective Communication。

https://zhuanlan.zhihu.com/p/100012827

logger/save checkpoint

參考 mmcv/logging，對於 dist.get_rank() ！= 0 的進程不添加 FileHandler，並設置 log_level=logging.ERROR。from mmcv.utils import get_logger -> get_logger('taskname', log_file) 。

混合精度訓練

According to Mixed Precision Training, the steps of fp16 optimizer is as follows.

Prepare model, .half() all modules except batchnorm (groupnorm)
Scale the loss value by a scale factor and convert from fp32 to fp16.
BP in the fp16 model.
Copy gradients from fp16 model to fp32 weights.
Update fp32 weights.
Copy updated parameters from fp32 weights to fp16 model.

https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/hooks/optimizer.py

雖然 params 有兩份儲存了，但是最消耗 memory 的是對圖像求導的部分，因此仍然可以降低

PyTorch

參考 PyTorch 源码解读之 torch.cuda.amp: 自动混合精度详解。

Apex (Deprecated)

Apex 是 Nvidia 的 PyTorch 混合精度訓練工具，對於 FP16 安全的操作在訓練中會 Cast 到 FP16，反之則使用 FP32。

FP16 的優缺點可參考 Quora/What is the difference between FP16 and FP32 when doing deep learning?。

1 2	from apex import amp from apex.parallel import DistributedDataParallel

在構建完 model 和 optimizer 之後，DistributedDataParallel(model) 之前，使用 amp.initialize 初始化

1	model, optimizer = amp.initialize(model, optimizer, opt_level=amp_opt_level)

opt_level 可參考官方文檔，大致 00 就是不啟用混合精度，工作在 FP32，01 則是啟用。

在 loss bp 前

1
2
3

with amp.scale_loss(losses, optimizer) as scaled_losses:
    scaled_losses.backward()
optimizer.step()

其它無需改變。

PyTorch 中多卡及混合精度使用方法

多 GPU 和 Process 支持

原理

使用

注意事項

混合精度訓練

PyTorch

Apex (Deprecated)

Tang Huan