Python PyTorch DistributedDataParallel用法及代碼示例

本文簡要介紹python語言中 torch.nn.parallel.DistributedDataParallel 的用法。

用法:

class torch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=None, dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25, find_unused_parameters=False, check_reduction=False, gradient_as_bucket_view=False)

參數：

• module(torch.nn.Module) -要並行化的模塊

• device_ids(Python列表：int或者torch.device) -

  CUDA 設備。 1)對於單設備模塊，device_ids可以隻包含一個設備id，代表該進程對應的輸入模塊所在的唯一CUDA設備。或者，device_ids 也可以是 None 。 2) 對於多設備模塊和 CPU 模塊，device_ids 必須為 None 。

  對於這兩種情況，當 device_ids 為 None 時，前向傳遞的輸入數據和實際模塊都必須放置在正確的設備上。 (默認：None)

• output_device(int或者torch.device) -單設備 CUDA 模塊的輸出的設備位置。對於多設備模塊和CPU模塊，它必須是None，並且模塊本身決定輸出位置。 (對於單設備模塊，默認值：device_ids[0])

• broadcast_buffers(bool) -在 forward 函數開頭啟用模塊同步(廣播)緩衝區的標誌。 (默認值：True)

• process_group-用於分布式數據的進程組all-reduction。如果是 None ，則將使用由 torch.distributed.init_process_group() 創建的默認進程組。 (默認：None)

• bucket_cap_mb-DistributedDataParallel 將參數分桶到多個桶中，以便每個桶的梯度減少可能與後向計算重疊。 bucket_cap_mb 控製存儲桶大小，單位為MegaBytes (MB)。 (默認值：25)

• find_unused_parameters(bool) -從包裝模塊的 forward 函數的返回值中包含的所有張量遍曆 autograd 圖。不接收梯度作為該圖的一部分的參數被預先標記為準備減少。此外，可能已在包裝模塊的 forward 函數中使用但不屬於損失計算的一部分，因此也不會接收梯度的參數被搶先標記為準備減少。 (默認：False)

• check_reduction-此論點已棄用。

• gradient_as_bucket_view(bool) -當設置為 True 時，梯度將是指向 allreduce 通信桶的不同偏移量的視圖。這可以減少峰值內存使用量，其中節省的內存大小將等於總梯度大小。此外，它避免了在梯度和allreduce 通信桶之間複製的開銷。當漸變是視圖時，不能在漸變上調用detach_()。如果遇到此類錯誤，請參考torch/optim/optimizer.py 中的 zero_grad() 函數作為解決方案來修複它。

變量：

~DistributedDataParallel.module(torch.nn.Module) -要並行化的模塊。

在模塊級別實現基於torch.distributed 包的分布式數據並行性。

此容器通過在批處理維度中分塊將輸入拆分到指定的設備，從而並行化給定模塊的應用程序。該模塊在每台機器和每台設備上複製，每個這樣的副本處理輸入的一部分。在反向傳播期間，來自每個節點的梯度被平均。

批量大小應大於本地使用的 GPU 數量。

另請參閱：基礎知識和使用 nn.parallel.DistributedDataParallel 而不是多處理或 nn.DataParallel。對輸入的限製與 torch.nn.DataParallel 中的相同。

創建此類需要通過調用 torch.distributed.init_process_group() 來初始化 torch.distributed。

對於 single-node multi-GPU 數據並行訓練，DistributedDataParallel 被證明比 torch.nn.DataParallel 快得多。

要在具有 N 個 GPU 的主機上使用 DistributedDataParallel，您應該生成 N 進程，確保每個進程僅在從 0 到 N-1 的單個 GPU 上工作。這可以通過為每個進程設置 CUDA_VISIBLE_DEVICES 或調用：

>>> torch.cuda.set_device(i)

其中 i 是從 0 到 N-1。在每個過程中，您應該參考以下內容來構建此模塊：

>>> torch.distributed.init_process_group(
>>>     backend='nccl', world_size=N, init_method='...'
>>> )
>>> model = DistributedDataParallel(model, device_ids=[i], output_device=i)

為了在每個節點生成多個進程，您可以使用 torch.distributed.launch 或 torch.multiprocessing.spawn 。

注意

分布式訓練相關的所有特性請參考PyTorch Distributed Overview簡單介紹。

注意

DistributedDataParallel 可以與 torch.distributed.optim.ZeroRedundancyOptimizer 結合使用，以減少 per-rank 優化器狀態內存占用。更多詳情請參閱ZeroRedundancyOptimizer recipe。

注意

nccl 後端是目前使用 GPU 時速度最快且強烈推薦的後端。這適用於 single-node 和 multi-node 分布式訓練。

注意

該模塊還支持mixed-precision分布式訓練。這意味著您的模型可以具有不同類型的參數，例如 fp16 和 fp32 的混合類型，這些混合類型的參數的梯度減少將正常工作。

注意

如果您在一個進程上使用torch.save 來檢查模塊，並在其他一些進程上使用torch.load 來恢複它，請確保為每個進程正確配置map_location。如果沒有 map_location ， torch.load 會將模塊恢複到保存模塊的設備。

注意

當使用 batch=N 在 M 節點上訓練模型時，如果對損失求和(不像平常那樣平均)，則與使用 batch=M*N 在單個節點上訓練的相同模型相比，梯度將小 M 倍批次中的跨實例(因為不同節點之間的梯度是平均的)。當您想要獲得與本地訓練對應的數學上等效的訓練過程時，您應該考慮到這一點。但在大多數情況下，您可以將 DistributedDataParallel 包裝模型、DataParallel 包裝模型和單個 GPU 上的普通模型視為相同(例如，對等效批量大小使用相同的學習率)。

注意

參數永遠不會在進程之間廣播。該模塊對梯度執行 all-reduce 步驟，並假設優化器將在所有進程中以相同的方式修改它們。在每次迭代中，緩衝區(例如 BatchNorm stats)從處於 0 級進程的模塊廣播到係統中的所有其他副本。

注意

如果您將 DistributedDataParallel 與分布式 RPC 框架結合使用，則應始終使用 torch.distributed.autograd.backward() 來計算梯度，並使用 torch.distributed.optim.DistributedOptimizer 來優化參數。

例子：

>>> import torch.distributed.autograd as dist_autograd
>>> from torch.nn.parallel import DistributedDataParallel as DDP
>>> from torch import optim
>>> from torch.distributed.optim import DistributedOptimizer
>>> from torch.distributed.rpc import RRef
>>>
>>> t1 = torch.rand((3, 3), requires_grad=True)
>>> t2 = torch.rand((3, 3), requires_grad=True)
>>> rref = rpc.remote("worker1", torch.add, args=(t1, t2))
>>> ddp_model = DDP(my_model)
>>>
>>> # Setup optimizer
>>> optimizer_params = [rref]
>>> for param in ddp_model.parameters():
>>>     optimizer_params.append(RRef(param))
>>>
>>> dist_optim = DistributedOptimizer(
>>>     optim.SGD,
>>>     optimizer_params,
>>>     lr=0.05,
>>> )
>>>
>>> with dist_autograd.context() as context_id:
>>>     pred = ddp_model(rref.to_here())
>>>     loss = loss_func(pred, loss)
>>>     dist_autograd.backward(context_id, loss)
>>>     dist_optim.step()

注意

要讓非 DDP 模型從 DDP 模型加載狀態字典，需要應用consume_prefix_in_state_dict_if_present() 以在加載前去除 DDP 狀態字典中的前綴 “module.”。

警告

構造函數、前向方法和輸出的微分(或本模塊輸出的函數)是分布式同步點。考慮到這一點，以防不同的進程可能正在執行不同的代碼。

警告

該模塊假定所有參數在模型創建時都已注冊在模型中。以後不應添加或刪除任何參數。同樣適用於緩衝區。

警告

該模塊假設所有參數都注冊在模型中，每個分布式進程的順序相同。模塊本身將按照模型注冊參數的相反順序進行梯度allreduce。換句話說，用戶有責任確保每個分布式進程具有完全相同的模型，從而具有完全相同的參數注冊順序。

警告

此模塊允許具有非行主要連續步幅的參數。例如，您的模型可能包含一些 torch.memory_format 為 torch.contiguous_format 的參數和其他格式為 torch.channels_last 的參數。但是，不同進程中對應的參數必須具有相同的步幅。

警告

此模塊不適用於 torch.autograd.grad() (即，它僅在漸變將在參數的.grad 屬性中累積時才有效)。

警告

如果您計劃將此模塊與 nccl 後端或 gloo 後端(使用 Infiniband)以及使用多個工作線程的 DataLoader 一起使用，請將多處理啟動方法更改為 forkserver(僅限 Python 3) )或spawn。不幸的是，Gloo(使用 Infiniband)和 NCCL2 不是分叉安全的，如果不更改此設置，您可能會遇到死鎖。

警告

module 及其子模塊上定義的前向和後向掛鉤將不再被調用，除非掛鉤在 forward() 方法中被初始化。

警告

在使用 DistributedDataParallel 包裝模型後，您永遠不應該嘗試更改模型的參數。因為，當用 DistributedDataParallel 包你的模型時，DistributedDataParallel 的構造函數將在構建時在模型本身的所有參數上注冊額外的梯度減少函數。如果之後更改模型的參數，梯度縮減函數將不再匹配正確的參數集。

警告

將DistributedDataParallel 與分布式 RPC 框架結合使用是實驗性的，可能會發生變化。

例子：

>>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...')
>>> net = torch.nn.parallel.DistributedDataParallel(model, pg)