[源碼解析] PyTorch 分布式(3) ----- DataParallel(下)
目錄
0x00 摘要
本文是 PyTorch 分布式的第三篇，繼續上文，介紹 DataPrallel 的並行操作和反向傳播。
本系列其他文章如下：
[源碼解析]PyTorch如何實現前向傳播(1) --- 基礎類(上)
[源碼解析]PyTorch如何實現前向傳播(2) --- 基礎類(下)
[源碼解析] PyTorch如何實現前向傳播(3) --- 具體實現
[源碼解析] Pytorch 如何實現后向傳播 (1)---- 調用引擎
[源碼解析] Pytorch 如何實現后向傳播 (2)---- 引擎靜態結構
[源碼解析] Pytorch 如何實現后向傳播 (3)---- 引擎動態邏輯
[源碼解析] PyTorch 如何實現后向傳播 (4)---- 具體算法
[源碼解析] PyTorch 分布式(1)------歷史和概述
[源碼解析] PyTorch 分布式(2) ----- DataParallel(上)
0x01 前向操作
我們先回憶一下目前的前向圖，replicate 調用了Broadcast.forward，同時往其context 存儲了input_device和num_inputs。
+----------------------------------------------------------------------------------------+
| DataParallel.forward |
| |
| |
| replicate +---------------> parallel_apply gather |
| |
+----------------------------------------------------------------------------------------+
+---------------------------+
| Broadcast |
| |
| |
| |
| forward() +----------->
| |
| |
| +---------------------+ |
| | ctx | |
| | input_device | |
| | | |
| | num_inputs | |
| | | |
| +---------------------+ |
| |
| |
| |
| |
| |
| |
+---------------------------+
1.1 並行
目前，我們已經使用 Scatter 函數將數據從 device[0] 分配並復制到不同的卡，用 Replicate 函數將模型從 device[0] 復制到不同的卡，這樣各個卡都有了同樣的模型和不同的數據，現在就要分別調用 forward 計算損失和梯度。也就是 parallel_apply 部分。
# 分發數據
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
# 分發模型
replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
# 並行訓練
outputs = self.parallel_apply(replicas, inputs, kwargs)
對應我們傳播圖是：
parallel_apply 是基於threading 實現，用前面准備好的 replica 和輸入數據，然后for 循環啟動多線程進行前向傳播，最后輸出傳播結果。
def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
# 確保模型和輸入大小一致
assert len(modules) == len(inputs)
# 確保每個 GPU 都有相應的元數據，如沒有就空白補全
if kwargs_tup is not None:
# 在前面已經補全
assert len(modules) == len(kwargs_tup)
else:
kwargs_tup = ({},) * len(modules)
# 確保模型數目和CPU數目一致
if devices is not None:
assert len(modules) == len(devices)
else:
devices = [None] * len(modules)
devices = [_get_device_index(x, True) for x in devices]
# 基於threading多線程實現
lock = threading.Lock()
results = {}
grad_enabled, autocast_enabled = torch.is_grad_enabled(), torch.is_autocast_enabled()
# 定義 worker
def _worker(i, module, input, kwargs, device=None):
torch.set_grad_enabled(grad_enabled)
if device is None:
device = get_a_var(input).get_device()
try:
# 設置當前的設備
with torch.cuda.device(device), autocast(enabled=autocast_enabled):
# this also avoids accidental slicing of `input` if it is a Tensor
if not isinstance(input, (list, tuple)):
input = (input,)
output = module(*input, **kwargs) # 前向操作
with lock:
# 並行計算得到輸出
results[i] = output
except Exception:
with lock:
results[i] = ExceptionWrapper(
where="in replica {} on device {}".format(i, device))
if len(modules) > 1:
# 如有一個進程控制多個 GPU ，起多個線程
# 注意，這里就是每個 worker 調用了 modules 數組中的一個模型copy
threads = [threading.Thread(target=_worker,
args=(i, module, input, kwargs, device))
for i, (module, input, kwargs, device) in
enumerate(zip(modules, inputs, kwargs_tup, devices))]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
else:
# 一個GPU對應一個進程
_worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])
outputs = []
for i in range(len(inputs)):
output = results[i]
# error handle
if isinstance(output, ExceptionWrapper):
output.reraise()
outputs.append(output)
# 輸出 n 個計算結果
return outputs
此時前向傳播具體對應如下圖，現在並行操作調用了 module 的forward方法。
+----------------------------------------------------------------------------------------+
| DataParallel.forward |
| |
| 1 2 3 |
| replicate +---------------> parallel_apply gather |
| |
+----------------------------------------------------------------------------------------+
+---------------------------+ +-------------------+
| Broadcast | | module |
| | | |
| | | |
| 1 | | 2 |
| forward() +-----------> | forward() +--------->
| | | |
| | | |
| +---------------------+ | | |
| | ctx | | | |
| | input_device | | | |
| | | | | |
| | num_inputs | | | |
| | | | | |
| +---------------------+ | | |
| | | |
| | | |
| | | |
| | | |
| | | |
| | | |
+---------------------------+ +-------------------+
1.2 Gather
目前，我們已經使用 Scatter 函數將數據從 device[0] 分配並復制到不同的卡，用 Replicate 函數將模型從 device[0] 復制到不同的卡，這樣各個卡都有了同樣的模型和不同的數據，然后分別調用 forward 計算損失和梯度。也就是 parallel_apply 部分。
現在要做的就是把分布式計算的梯度合並到 device[0]，就是 self.output_device。
# 分發數據
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
# 分發模型
replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
# 並行訓練
outputs = self.parallel_apply(replicas, inputs, kwargs)
# 收集到 devices[0]
return self.gather(outputs, self.output_device)
對應我們傳播圖是：
我們看看如何把結果收集到 device[0]，以及device[0]如何作為參數服務器。
1.2.1 Python世界
gather 主要是調用 Gather.apply(target_device, dim, *outputs) 完成收集工作。
def gather(outputs, target_device, dim=0): # target_device 就是 device[0]
r"""
Gathers tensors from different GPUs on a specified device
(-1 means the CPU).
"""
def gather_map(outputs):
out = outputs[0]
if isinstance(out, torch.Tensor):
return Gather.apply(target_device, dim, *outputs) # 調用下面的 Gather
if out is None:
return None
if isinstance(out, dict):
return type(out)(((k, gather_map([d[k] for d in outputs]))
for k in out))
return type(out)(map(gather_map, zip(*outputs)))
# Recursive function calls like this create reference cycles.
# Setting the function to None clears the refcycle.
try:
res = gather_map(outputs)
finally:
gather_map = None
return res
Gather 則調用了 comm.gather 完成工作，而 comm.gather 則會帶領我們進入到 C++世界。
我們省略一些校驗代碼。
# Gather 源碼
class Gather(Function):
@staticmethod
def forward(ctx, target_device, dim, *inputs): # target_device 就是 device[0]
# 下面會往 context 內部存放幾個變量，后續會用到
target_device = _get_device_index(target_device, True)
ctx.target_device = target_device
ctx.dim = dim
ctx.input_gpus = tuple(i.get_device() for i in inputs)
if all(t.dim() == 0 for t in inputs) and dim == 0:
inputs = tuple(t.view(1) for t in inputs)
ctx.unsqueezed_scalar = True
else:
ctx.unsqueezed_scalar = False
ctx.input_sizes = tuple(i.size(ctx.dim) for i in inputs)
return comm.gather(inputs, ctx.dim, ctx.target_device) # 這里會進入C++世界
@staticmethod
def backward(ctx, grad_output): # 注意，這里后續會用到
scattered_grads = Scatter.apply(ctx.input_gpus, ctx.input_sizes, ctx.dim, grad_output)
if ctx.unsqueezed_scalar:
scattered_grads = tuple(g[0] for g in scattered_grads)
return (None, None) + scattered_grads
現在前向計算如圖：
gather 調用到了Gather的forward 函數，forward 方法在 ctx 存儲了 input_gpus, input_sizes, dim 這三個變量，這些變量后續會用到。
+-----------------------------------------------------------------------------------------+
| DataParallel.forward |
| |
| 1 2 3 |
| replicate +---------------> parallel_apply +--------------> gather |
| |
+-----------------------------------------------------------------------------------------+
+---------------------------+ +-------------------+ +--------------------+
| Broadcast | | module | |Gather |
| | | | | |
| | | | | |
| 1 | | 2 | | 3 |
| forward() +-----------> | forward() +--------> | forward() |
| | | | | |
| | | | | |
| +---------------------+ | | | | +----------------+ |
| | ctx | | | | | |ctx | |
| | input_device | | | | | | input_gpus | |
| | | | | | | | | |
| | num_inputs | | | | | | input_sizes| |
| | | | | | | | | |
| +---------------------+ | | | | | dim | |
| | | | | +----------------+ |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
+---------------------------+ +-------------------+ +--------------------+
1.2.2 C++世界
gather 函數調用了 _gather_out_impl 來完成拷貝操作。
at::Tensor gather(
at::TensorList tensors,
int64_t dim,
c10::optional<int32_t> destination_index) { // destination_index 就是 device[0] 的index
int64_t total_size = 0;
auto& first = tensors.front();
const auto first_size = first.sizes();
dim = at::maybe_wrap_dim(dim, first);
std::vector<int64_t> expected_size(first_size.begin(), first_size.end());
auto memory_format = first.suggest_memory_format();
for (size_t i = 0; i < tensors.size(); i++) {
const auto& tensor = tensors[i];
expected_size[dim] = tensor.size(dim);
total_size += tensor.size(dim);
if (memory_format != MemoryFormat::Contiguous &&
tensor.suggest_memory_format() != memory_format) {
memory_format = MemoryFormat::Contiguous;
}
}
expected_size[dim] = total_size;
at::Device device(DeviceType::CPU);
// 根據 index 得到輸出的目標設備
if (!destination_index || *destination_index != -1) {
// device 就是 GPU 0 這個設備
device = at::Device(
DeviceType::CUDA, destination_index ? *destination_index : -1);
}
//首先，構建一個空的目標tensor建立在目標設備之上，命名為result
at::Tensor result =
at::empty(expected_size, first.options().device(device), memory_format);
return _gather_out_impl(tensors, result, dim); // 然后對result進行gather
}
_gather_out_impl 執行了具體的gather 操作，就是把 輸入的tensors 拷貝到 目標 tensor 之上，即拷貝到 GPU0 之上。
// ***************** Gather *******************
//
// Gather a list of CUDA tensors on one or more devices to a target tensor or
// device, either CPU or CUDA.
// no checks
static inline at::Tensor& _gather_out_impl(
at::TensorList tensors,
at::Tensor& out_tensor,
int64_t dim) {
std::vector<int64_t> chunk_sizes;
chunk_sizes.reserve(tensors.size());
for (auto& tensor : tensors) {
chunk_sizes.push_back(tensor.size(dim));
}
auto chunks =
out_tensor.split_with_sizes(/*split_sizes=*/chunk_sizes, /*dim=*/dim);
for (size_t i = 0; i < tensors.size(); i++) { // 拷貝到GPU 0 之上
chunks[i].copy_(tensors[i], /*non_blocking=*/out_tensor.is_cuda());
}
return out_tensor;
}
0x02 計算損失
現在，我們已經把梯度收集到 device[0] 之上，現在我們開始進行反向傳播，其整體邏輯如上圖所示。首先是在 device[0] 計算損失。其實這步計算損失算是前向計算和后向傳播的中間環節，這里把它算成是反向傳播的開端，如下圖。
我們找出來示例代碼看看，里面關鍵的幾點：
- 數據已經放到了默認GPU，即GPU 0上。
- prediction 是gather到 GPU 0 的前向計算輸出。
- 使用
loss = criterion(prediction,target_var)在默認GPU之上計算loss。 - 使用 loss.backward() 開始反向傳播。
for batch_idx, (data, label) in pbar:
if args.cuda:
data,label= data.cuda(),label.cuda(); # 1. 數據已經放到了默認GPU上
data_v = Variable(data)
target_var = Variable(label)
prediction= model(data_v,target_var,args) # 2. prediction 是gather到 GPU 0 的前向計算輸出
# 到目前為止，我們完成了DataParallel.forward()
#這里的prediction 預測結果是由兩個gpu合並過的，並行計算只存在於前向傳播里
#前向傳播每個gpu計算量為 batch_size/len(device_ids),等前向傳播完了將結果聚合到主gpu里
criterion = nn.CrossEntropyLoss()
loss = criterion(prediction,target_var) # 3. 在默認GPU之上計算loss
optimizer.zero_grad()
loss.backward() # 4. 開始反向傳播
optimizer.step()
0x03 后向傳播
我們前面運行的是上面的 Forward 部分，計算損失，接下來就運行上面代碼中 loss.backward() 部分。
3.1 分發梯度
我們首先來到分發梯度部分，這部分作用是：把損失在 GPUs 之間 scatter，這樣后續才可以在每個GPU之上獨立進行后向傳播。對應下圖：
3.1.1 Gather.backward
前面有提到，prediction 是gather到 GPU 0 的前向計算輸出。而 loss 又是根據 prediction 計算出來，所以從 loss.backward() 開始反向傳播，從后向前的第一個步驟就來到了 gather 的傳播操作，對應的就是 Gather 的 backward 函數，其中的核心代碼是 Scatter.apply。
class Gather(Function):
# 這里前向傳播用到了，為了對照，我們依然貼出來
@staticmethod
def forward(ctx, target_device, dim, *inputs): # target_device 就是 device[0]
# 下面會往 context 內部存放幾個變量，后續會用到
target_device = _get_device_index(target_device, True)
ctx.target_device = target_device
ctx.dim = dim
ctx.input_gpus = tuple(i.get_device() for i in inputs)
if all(t.dim() == 0 for t in inputs) and dim == 0:
inputs = tuple(t.view(1) for t in inputs)
ctx.unsqueezed_scalar = True
else:
ctx.unsqueezed_scalar = False
ctx.input_sizes = tuple(i.size(ctx.dim) for i in inputs)
# 這里會進入C++世界，把輸出聚集到 GPU 0。
return comm.gather(inputs, ctx.dim, ctx.target_device)
@staticmethod
def backward(ctx, grad_output): # 這里現在后向傳播用到了！
# 把前向傳播在 context 之中存放的變量取出，作為 Scatter 的輸入
scattered_grads = Scatter.apply(ctx.input_gpus, ctx.input_sizes, ctx.dim, grad_output)
if ctx.unsqueezed_scalar:
scattered_grads = tuple(g[0] for g in scattered_grads)
return (None, None) + scattered_grads
具體如下，可以看到，backward 使用了之前前向傳播時候存儲的 ctx.input_gpus, ctx.input_sizes, ctx.dim, grad_output，以此調用 Scatter.apply。
圖中，最上面是前向傳播過程，最下面是反向傳播過程，中間是某些在前后傳播中都用到的代碼模塊。
+--------------------------------------------------------------------------------------+
| DataParallel.forward |
| |
| 1 2 3 |
| replicate +---------------> parallel_apply +--------------> gather |
| |
+--------------------------------------------------------------------------------------+
+---------------------------+ +-------------------+ +--------------------+
| Broadcast | | module | |Gather |
| | | | | |
| | | | | |
| 1 | | 2 | | 3 |
| forward() +-----------> | forward() +--------> | forward() |
| | | | | |
| | | | | |
| +---------------------+ | | | | +----------------+ |
| | ctx | | | | | |ctx | |
| | input_device | | | | | | input_gpus | |
| | | | | | | | | |
| | num_inputs | | | | | | input_sizes| |
| | | | | | | | | |
| +---------------------+ | | | | | dim | |
| | | | | +----------------+ |
| | | | | |
| | | | | |
| | | | <---------+ backward() |
| | | | | 3 |
| | | | | |
+---------------------------+ +-------------------+ +--------------------+
+--------------------------------------------------------------------------------------+
| loss.backward() |
| 3 |
| <--------------------+ |
| |
| |
+--------------------------------------------------------------------------------------+
3.1.2 Scatter
Scatter.apply 實際上調用到了其 forward 方法。
- 首先從上下文之中提取之前存儲的變量，這里主要是輸入設備 input_device（源設備）和 target_gpus（目標設備）。
- 獲取到目標設備的流。
- 調用 comm.scatter 把梯度分發到目標設備。
class Scatter(Function):
@staticmethod
def forward(ctx, target_gpus, chunk_sizes, dim, input):
target_gpus = [_get_device_index(x, True) for x in target_gpus]
ctx.dim = dim
ctx.input_device = input.get_device() if input.device.type != "cpu" else -1
streams = None
if torch.cuda.is_available() and ctx.input_device == -1:
# Perform CPU to GPU copies in a background stream
streams = [_get_stream(device) for device in target_gpus]
# 分發到其他GPU
outputs = comm.scatter(input, target_gpus, chunk_sizes, ctx.dim, streams)
# Synchronize with the copy stream
if streams is not None:
for i, output in enumerate(outputs):
with torch.cuda.device(target_gpus[i]):
main_stream = torch.cuda.current_stream()
main_stream.wait_stream(streams[i])
output.record_stream(main_stream)
return outputs
@staticmethod
def backward(ctx, *grad_output):
return None, None, None, Gather.apply(ctx.input_device, ctx.dim, *grad_output)
3.1.3 C++
上面python代碼 outputs = comm.scatter(input, target_gpus, chunk_sizes, ctx.dim, streams) 會直接進入到C++世界。具體代碼位於 torch/csrc/cuda/comm.cpp。
scatter 的作用就是把tensor進行split，然后分發給各個設備的流。
std::vector<at::Tensor> scatter(
const at::Tensor& tensor,
at::IntArrayRef devices,
const c10::optional<std::vector<int64_t>>& chunk_sizes,
int64_t dim,
const c10::optional<std::vector<c10::optional<at::cuda::CUDAStream>>>&
streams) {
dim = at::maybe_wrap_dim(dim, tensor);
// 把tensor進行split
std::vector<at::Tensor> chunks = chunk_sizes
? tensor.split_with_sizes(/*split_sizes=*/*chunk_sizes, /*dim=*/dim)
: tensor.chunk(/*chunks=*/devices.size(), /*dim=*/dim);
at::cuda::OptionalCUDAStreamGuard cuda_guard;
for (size_t i = 0; i < chunks.size(); ++i) {
const auto device_index = static_cast<int16_t>(devices[i]);
if (device_index != tensor.get_device()) {
if (i < (streams ? streams->size() : 0U) && (*streams)[i]) {
cuda_guard.reset_stream(*(*streams)[i]);
}
// 發送給各個設備的流
chunks[i] = chunks[i].to(
{DeviceType::CUDA, device_index},
/*non_blocking=*/true,
/*copy=*/false,
/*memory_format=*/at::MemoryFormat::Preserve);
}
}
return chunks;
}
3.2 並行后向傳播
現在梯度已經分發到各個 GPU，接下來正式進入並行后向傳播，這部分作用是：在各個GPU之上並行運行后向傳播，計算參數梯度。對應下圖：
這部分調用到了原始模型的 backward，具體如下圖中的數值 4：
+--------------------------------------------------------------------------------------+
| DataParallel.forward |
| |
| 1 2 3 |
| replicate +---------------> parallel_apply +--------------> gather |
| |
+--------------------------------------------------------------------------------------+
+---------------------------+ +-------------------+ +--------------------+
| Broadcast | | module | |Gather |
| | | | | |
| | | | | |
| 1 | | 2 | | 3 |
| forward() +-----------> | forward() +--------> | forward() |
| | | | | |
| | | | | |
| +---------------------+ | | | | +----------------+ |
| | ctx | | | | | |ctx | |
| | input_device | | | | | | input_gpus | |
| | | | | | | | | |
| | num_inputs | | | | | | input_sizes| |
| | | | | | | | | |
| +---------------------+ | | | | | dim | |
| | | | | +----------------+ |
| | | | | |
| | | | | |
| | <---------+ backward() | <---------+ backward() |
| | | 4 | | 3 |
| | | | | |
+---------------------------+ +-------------------+ +--------------------+
+--------------------------------------------------------------------------------------+
| loss.backward() |
| 4 3 |
| <------------------+ <--------------------+ |
| |
| |
+--------------------------------------------------------------------------------------+
3.3 歸並梯度
這部分作用是 在 GPU 0 之上歸並梯度，總體流程拓展對應下圖：
3.3.1 Broadcast.backward
這部分對應了 Broadcast 的 反向傳播。
class Broadcast(Function):
@staticmethod
def forward(ctx, target_gpus, *inputs):
target_gpus = [_get_device_index(x, True) for x in target_gpus]
# 前向傳播時候，向上下文存入了一些變量
ctx.target_gpus = target_gpus
if len(inputs) == 0:
return tuple()
ctx.num_inputs = len(inputs)
# input 放在 device[0]，所以 input_device 就是 GPU 0
ctx.input_device = inputs[0].get_device()
# 和 detach 的情形一樣
outputs = comm.broadcast_coalesced(inputs, ctx.target_gpus)
non_differentiables = []
# 在上下文中設置哪些不需要梯度
for idx, input_requires_grad in enumerate(ctx.needs_input_grad[1:]):
if not input_requires_grad:
for output in outputs:
non_differentiables.append(output[idx])
ctx.mark_non_differentiable(*non_differentiables)
return tuple([t for tensors in outputs for t in tensors])
@staticmethod
def backward(ctx, *grad_outputs):
# 反向傳播來到這里，取出之前在上下文存放的變量作為輸入。ctx.input_device 就是之前存儲的 GPU 0。
return (None,) + ReduceAddCoalesced.apply(ctx.input_device, ctx.num_inputs, *grad_outputs)
因此，我們可以拓展流程圖：
+--------------------------------------------------------------------------------------+
| DataParallel.forward |
| |
| 1 2 3 |
| replicate +---------------> parallel_apply +--------------> gather |
| |
+--------------------------------------------------------------------------------------+
+---------------------------+ +-------------------+ +--------------------+
| Broadcast | | module | |Gather |
| | | | | |
| | | | | |
| 1 | | 2 | | 3 |
| forward() +-----------> | forward() +--------> | forward() |
| | | | | |
| | | | | |
| +---------------------+ | | | | +----------------+ |
| | ctx | | | | | |ctx | |
| | input_device | | | | | | input_gpus | |
| | | | | | | | | |
| | num_inputs | | | | | | input_sizes| |
| | | | | | | | | |
| +---------------------+ | | | | | dim | |
| | | | | +----------------+ |
| | | | | |
| | | | | |
| backward() | <---------+ backward() | <---------+ backward() |
| 5 | | 4 | | 3 |
| | | | | |
+---------------------------+ +-------------------+ +--------------------+
+--------------------------------------------------------------------------------------+
| loss.backward() |
| 5 4 3 |
| <------------------------+ <------------------+ <--------------------+ |
| |
| |
+--------------------------------------------------------------------------------------+
3.3.2 ReduceAddCoalesced
Broadcast.backward 調用了 ReduceAddCoalesced.apply，其對應了 ReduceAddCoalesced 的 forward 方法，目的是把梯度歸並到目標設備 destination，就是GPU 0。
class ReduceAddCoalesced(Function):
@staticmethod
# 會調用到這里，destination 是GPU 0
def forward(ctx, destination, num_inputs, *grads):
# 從梯度之中提取所在的設備
ctx.target_gpus = [grads[i].get_device() for i in range(0, len(grads), num_inputs)]
grads_ = [grads[i:i + num_inputs]
for i in range(0, len(grads), num_inputs)]
# 把梯度歸並到目標設備 destination，就是GPU 0
return comm.reduce_add_coalesced(grads_, destination)
@staticmethod
def backward(ctx, *grad_outputs):
return (None, None,) + Broadcast.apply(ctx.target_gpus, *grad_outputs)
3.3.3 c++
看注釋就是：從多個 GPU 來相加梯度，代碼之中就是歸並相加。
def reduce_add_coalesced(inputs, destination=None, buffer_size=10485760):
"""Sums tensors from multiple GPUs.
Small tensors are first coalesced into a buffer to reduce the number
of synchronizations.
Args:
inputs (Iterable[Iterable[Tensor]]): iterable of iterables that
contain tensors from a single device.
destination (int, optional): a device on which the output will be
placed (default: current device).
buffer_size (int): maximum size of the buffer used for coalescing
Returns:
A tuple of tensors containing an elementwise sum of each group of
inputs, placed on the ``destination`` device.
"""
dense_tensors: List[List] = [[] for _ in inputs] # shape (num_gpus, num_tensors)
output = []
ref_order = []
# process sparse ones first since they may have different sizes on different gpus
for tensor_at_gpus in zip(*inputs):
if all(t.is_sparse for t in tensor_at_gpus):
# 進行歸並
result = reduce_add(tensor_at_gpus, destination) # this will be sparse too
output.append(result)
ref_order.append(tensor_at_gpus[0])
else:
for coll, t in zip(dense_tensors, tensor_at_gpus):
coll.append(t.to_dense() if t.is_sparse else t)
ref_order.append(dense_tensors[0][-1])
itrs = [_take_tensors(tensors, buffer_size) for tensors in dense_tensors]
# now the dense ones, which have consistent sizes
for chunks in zip(*itrs):
flat_tensors = [_flatten_dense_tensors(chunk) for chunk in chunks] # (num_gpus,)
# 進行歸並
flat_result = reduce_add(flat_tensors, destination)
for t in _unflatten_dense_tensors(flat_result, chunks[0]):
# The unflattened tensors do not share storage, and we don't expose
# base flat tensor anyways, so give them different version counters.
# See NOTE [ Version Counter in comm.*_coalesced ]
output.append(t.data)
return tuple(_reorder_tensors_as(output, ref_order))
3.4 更新模型參數
這部分功能是：更新梯度參數。進行梯度下降，並更新主GPU上的模型參數。
另外，由於模型參數僅在主GPU上更新，而其他從屬GPU此時並不是同步更新的，所以需要將更新后的模型參數復制到剩余的從屬 GPU 中，以此來實現並行。這就是在下一次for循環之中進行，以此循環反復。
對應示例代碼是：
for batch_idx, (data, label) in pbar: # 6. 下一次迭代會繼續從分發開始
if args.cuda:
data,label= data.cuda(),label.cuda(); # 1. 數據已經放到了默認GPU上
data_v = Variable(data)
target_var = Variable(label)
prediction= model(data_v,target_var,args) # 2. prediction 是gather到 GPU 0 的前向計算輸出
# 到目前為止，我們完成了DataParallel.forward()
#這里的prediction 預測結果是由兩個gpu合並過的，並行計算只存在在前向傳播里
#前向傳播每個gpu計算量為 batch_size/len(device_ids),等前向傳播完了將結果和到主gpu里
criterion = nn.CrossEntropyLoss()
loss = criterion(prediction,target_var) # 3. 在默認GPU之上計算loss
optimizer.zero_grad()
loss.backward() # 4. 開始反向傳播
optimizer.step() # 5. 更新模型
0x04 總結
我們總結一下流程，起初數據和模型被放入到默認GPU，就是 GPU 0，然后迭代如下：
- scatter 會把數據分發到其他 GPU。
- replicate 會把模型分發到其他 GPU。
- parallel_apply 會啟動多個線程進行前向計算。
- gather 會把計算輸出收集到 GPU 0。
- GPU 0 會計算損失。
- 把梯度 scatter 到其他 GPU。
- 模型調用 backward 計算。
- 把梯度歸並到 GPU 0。
- optimizer.step 更新模型。
具體對應下圖之中的數字。
+-----+ +-------+
|GPU1 | | GPU1 |
main thread +-----+ +-------+
+-----> Forward----> scatter +--------------> replicate-------> parallel_apply +--------> gather +---------+
+ + + |
1 | 2 | 3 | |
| | | |
| +---------+----------+---+ | |
| | | | | |
+---------+----------+ | +--------------------+ |
| | | | | | | | | |
| | 2 | | 2 | | 2 thread|1 thread 2 thread 3 |
1 | | 1 | | 1 | | | | | |
| v | v | v | | | |
v v v v v v |
+--+---+ +--+---+ +--+---+ +--+---+ +--+---+ +--+---+ +-------+ |
| GPU1 | | GPU2 | | GPU3 | | GPU1 | | GPU2 | | GPU3 | | GPU1 | |
+------+ +------+ +------+ +--+---+ +-+----+ +---+--+ +-+-+--++ |
| | | ^ ^ ^ |
| | | 4 | | | |
| | ----------^ | | |
| | 4 | | |
| +------------------------+ | |
| | |
+------------------------------------+ |
+------------------------------------------------------------------------------------------------------+
| +------+
| | GPU1 |
| +------+ main thread
+-> loss = criterion(...)+-----> scatter +--------------> model.backward() +----------> reduce gradient +-------> optimizer.step
+ + + +------+ 9
| 5 | 6 | 7 | GPU1 |
| | | +--+---+
| v---------------v +--------------------+ ^
| | | | | | | | 8
| | | | thread 1 thread 2 thread 3 |
| | | | + | | +-------------+
| | | | | | | | | |
v v v v v v v | | |
+--+---+ +---+-+ +--+--+ +-+---+ +--+--+ +---+--+ +--+--+ +--+--+ +-+--+ +-+---+
| GPU1 | | GPU1| | GPU2| |GPU3 | | GPU1| | GPU2 | |GPU3 | | GPU1| |GPU2| | GPU3|
+------+ +-----+ +-----+ +-----+ +-----+ +------+ +-----+ +-----+ +----+ +-----+
手機如下：
至此，DP 分析完畢，我們下一篇要介紹 DDP 的一些相關知識。
0xFF 參考
PyTorch 源碼解讀之 torch.optim：優化算法接口詳解
pytorch(分布式)數據並行個人實踐總結——DataParallel/DistributedDataParallel
https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/20
[原創][深度][PyTorch] DDP系列第二篇：實現原理與源代碼解析
Pytorch踩坑記：賦值、淺拷貝、深拷貝三者的區別以及model.state_dict()和model.load_state_dict()的坑點