[原始碼解析] PyTorch 分散式(12) ----- DistributedDataParallel 之 前向傳播
[原始碼解析] PyTorch 分散式(12) ----- DistributedDataParallel 之 前向傳播
目錄0x00 摘要
前文已經對Reducer如何構建和幾個重要場景做了介紹,本文就來分析 Reducer 如何實現前向傳播。
本系列其他文章如下:
[原始碼解析]深度學習利器之自動微分(3) --- 示例解讀
[原始碼解析]PyTorch如何實現前向傳播(1) --- 基礎類(上)
[原始碼解析]PyTorch如何實現前向傳播(2) --- 基礎類(下)
[原始碼解析] PyTorch如何實現前向傳播(3) --- 具體實現
[原始碼解析] Pytorch 如何實現後向傳播 (1)---- 呼叫引擎
[原始碼解析] Pytorch 如何實現後向傳播 (2)---- 引擎靜態結構
[原始碼解析] Pytorch 如何實現後向傳播 (3)---- 引擎動態邏輯
[原始碼解析] PyTorch 如何實現後向傳播 (4)---- 具體演算法
[原始碼解析] PyTorch 分散式(1)------歷史和概述
[原始碼解析] PyTorch 分散式(2) ----- DataParallel(上)
[原始碼解析] PyTorch 分散式(3) ----- DataParallel(下)
[原始碼解析] PyTorch 分散式(4)------分散式應用基礎概念
[原始碼解析] PyTorch分散式(5) ------ DistributedDataParallel 總述&如何使用
[原始碼解析] PyTorch分散式(6) ---DistributedDataParallel -- 初始化&store
[原始碼解析] PyTorch 分散式(7) ----- DistributedDataParallel 之程序組
[原始碼解析] PyTorch 分散式(8) -------- DistributedDataParallel之論文篇
[原始碼解析] PyTorch 分散式(9) ----- DistributedDataParallel 之初始化
[原始碼解析] PyTorch 分散式(10)------DistributedDataParallel 之 Reducer靜態架構
[原始碼解析] PyTorch 分散式(11) ----- DistributedDataParallel 之 構建Reducer和Join操作
0x01 總體邏輯
我們還是需要祭出法寶,看看論文中的DDP總體邏輯:
然後給出一個前向傳播的總體策略如下:
Forward Pass:
- 每個程序讀去自己的訓練資料,DistributedSampler確保每個程序讀到的資料不同。
- DDP 獲取輸入並將其傳遞給本地模型。
- 模型進行前向計算,結果設定為 out。現在計算都是在每個程序(CUDA裝置)上完成。
- 如果
find_unused_parameters
設定為True
,DDP 會分析本地模型的輸出,從 out 開始遍歷計算圖,把未使用引數標示為 ready,因為每次計算圖都會改變,所以每次都要遍歷。- 此模式(Mode)允許在模型的子圖上向後執行,並且 DDP 通過從模型輸出out遍歷 autograd 圖,將所有未使用的引數標記為就緒,以減少反向傳遞中涉及的引數。
- 在後向傳播期間,Reducer會規約所有桶,在此過程中,
Reducer
會等待未準備好的引數。將引數梯度標記為就緒並不能幫助 DDP 跳過桶,但它會阻止 DDP 在向後傳遞期間永遠等待不存在的梯度。 - 請注意,遍歷 autograd 圖會引入額外的開銷,因此應用程式僅在必要時才設定
find_unused_parameters
為True
。
- 返回out即可。這點與 DP不同,DDP的模型網路輸出不需要被gather到 rank 0程序。
0x02 Python 世界
我們還是從 Python 程式碼入手開始分析,程式碼位於:torch/nn/parallel/distributed.py。
我們這裡省略 join 相關,只關注主體部分,forward 方法邏輯如下:
- 儲存執行緒本地狀態。
- 如果做配置,則呼叫 reducer.prepare_for_forward 為forward做準備。
- 如果配置ddp_join_enabled,做相應處理。
- 在前向傳播之前使用 _rebuild_buckets 來重置桶。
- 在 _rebuild_buckets 函式之中,也許會在釋放舊bucket之前分配新bucket。
- 如果要節省峰值記憶體使用量,請在正向計算期間峰值記憶體使用量增加之前呼叫
_rebuild_bucket
。
- 如果需要同步,則呼叫_sync_params對前向傳播引數進行前向傳播引數。
- 進行前向傳播。
- 如果需要同步後向傳播梯度,則呼叫prepare_for_backward。
- 當DDP引數 find_unused_parameter 為 true 時,其會在 forward 結束時,啟動一個回溯,標記出所有沒被用到的 parameter,提前把這些設定為 ready,這樣 backward 就可以在一個 subgraph 之上進行,但這樣會犧牲一部分時間。
具體程式碼如下:
def forward(self, *inputs, **kwargs):
with torch.autograd.profiler.record_function("DistributedDataParallel.forward"):
# 儲存執行緒本地狀態
self.reducer.save_thread_local_state()
# 如果做配置,則呼叫 reducer 為forward做準備
if torch.is_grad_enabled() and self.require_backward_grad_sync:
self.logger.set_runtime_stats_and_log()
self.num_iterations += 1
self.reducer.prepare_for_forward()
# 如果配置ddp_join_enabled,做相應處理
if self.ddp_uneven_inputs_config.ddp_join_enabled:
ones = torch.ones(1, device=self.device)
work = dist.all_reduce(ones, group=self.process_group, async_op=True)
if self.ddp_uneven_inputs_config.ddp_join_throw_on_early_termination:
# Active ranks schedule an allreduce with zeros, inactive
# ranks schedule them with 1. If the result != 0 it
# indicates at least one rank has terminated and we should
# throw.
zeros = torch.zeros(1, device=self.device)
dist.all_reduce(zeros, group=self.process_group)
should_throw_stop_iteration = zeros.item()
if should_throw_stop_iteration:
raise RuntimeError(
"Detected at least one rank that exhausted inputs. Throwing across all ranks."
)
else:
self.reducer._set_forward_pass_work_handle( # 是join這裡用到
work,
self.ddp_uneven_inputs_config.ddp_join_divide_by_initial_world_size,
)
# Calling _rebuild_buckets before forward compuation,
# It may allocate new buckets before deallocating old buckets
# inside _rebuild_buckets. To save peak memory usage,
# call _rebuild_buckets before the peak memory usage increases
# during forward computation.
# This should be called only once during whole training period.
# 在前向傳播之前使用 _rebuild_buckets 來重置桶
# 在此函式內,也許在釋放舊bucket之前分配新bucket。
# 如果要節省峰值記憶體使用量,請在正向計算期間峰值記憶體使用量增加之前呼叫_rebuild_bucket。
# 在整個訓練期間,這隻能呼叫一次。
if torch.is_grad_enabled() and self.reducer._rebuild_buckets():
logging.info("Reducer buckets have been rebuilt in this iteration.")
# 如果需要同步前向傳播引數,則進行同步
if self.require_forward_param_sync:
self._sync_params()
if self.ddp_uneven_inputs_config.ddp_join_enabled:
# Notify joined ranks whether they should sync in backwards pass or not.
self._check_global_requires_backward_grad_sync(is_joined_rank=False)
# 進行前向傳播
if self.device_ids:
# 多卡情況
inputs, kwargs = self.to_kwargs(inputs, kwargs, self.device_ids[0])
output = self.module(*inputs[0], **kwargs[0])
else:
output = self.module(*inputs, **kwargs)
# 如果需要同步後向傳播梯度,則呼叫prepare_for_backward
if torch.is_grad_enabled() and self.require_backward_grad_sync:
# 當DDP引數 find_unused_parameter 為 true 時,其會在 forward 結束時,啟動一個回溯,標記出所有沒被用到的 parameter,提前把這些設定為 ready,這樣 backward 就可以在一個 subgraph 進行,但這樣會犧牲一部分時間。
self.require_forward_param_sync = True
# We'll return the output object verbatim since it is a freeform
# object. We need to find any tensors in this object, though,
# because we need to figure out which parameters were used during
# this forward pass, to ensure we short circuit reduction for any
# unused parameters. Only if `find_unused_parameters` is set.
if self.find_unused_parameters and not self.static_graph:
# Do not need to populate this for static graph.
self.reducer.prepare_for_backward(list(_find_tensors(output)))
else:
self.reducer.prepare_for_backward([])
else:
self.require_forward_param_sync = False
# TODO. Right now we add this sink for static_graph training only. once
# this feature is stable, we will add this sink for all cases. E.g.
# This sink can help capture more accuracte backward start time as well.
if self.static_graph and self.num_iterations == 1:
# Need to grab list of tensors from user output in order to pass
# to custom autograd function.
output_tensor_list, treespec = tree_flatten(output)
passthrough_tensor_list = _DDPSink.apply(
self.reducer,
*output_tensor_list
)
# Reconstruct output data structure.
output = tree_unflatten(passthrough_tensor_list, treespec)
return output
其中,使用 _sync_params 來同步模型引數,具體是使用 _distributed_broadcast_coalesced 進行完成。
def _sync_params(self):
with torch.no_grad():
# module buffer sync
if self.will_sync_module_buffers():
# Synchronize buffers across processes.
# If we are running DDP with the join manager, we have to agree
# upon a rank to sync module buffers from, since rank 0 may
# already have been joined and have stale module buffers.
if self.ddp_uneven_inputs_config.ddp_join_enabled:
authoritative_rank = self._find_common_rank(
self._distributed_rank, True
)
else:
# The process with rank 0 is considered the authoritative copy.
authoritative_rank = 0
self._distributed_broadcast_coalesced(
self.modules_buffers[0],
self.broadcast_bucket_size,
authoritative_rank,
)
0x03 C++世界
我們接下來進入到 C++ 世界,看看這裡如何支援前向傳播。具體分為:準備前向傳播,重建桶,準備後向傳播這幾部分。
3.1 準備前向傳播
這裡把 num_iterations_ 增加,並且記錄時間。
void Reducer::prepare_for_forward() {
std::lock_guard<std::mutex> lock(mutex_);
num_iterations_++; // 這裡會遞增
if (should_collect_runtime_stats()) {
record_forward_compute_start_time();
}
}
3.2 重建桶
接下來進行重建桶,具體分為:
- 配置各種尺寸限制。
- 計算桶的尺寸。
- 同步桶indices。
- 初始化桶。
bool Reducer::rebuild_buckets() {
// Ensure reduction for previous backwards pass is finished. If user's model
// has unused parameters for example, this will raise an error recommending to
// run with find_unused_parameters=True, instead of the size mismatch
// exception below.
std::lock_guard<std::mutex> lock(mutex_);
ensure_prior_reduction_finished();
if (!should_rebuild_buckets() || rebuilt_params_.empty()) {
return false;
}
std::vector<std::vector<size_t>> rebuilt_bucket_indices;
// 配置各種尺寸限制
std::vector<size_t> bucket_size_limits;
bucket_size_limits.push_back(kDefaultFirstBucketBytes);
bucket_size_limits.push_back(bucket_bytes_cap_);
// 計算桶的尺寸
rebuilt_bucket_indices = compute_bucket_assignment_by_size(
rebuilt_params_,
bucket_size_limits,
expect_sparse_gradients_[0],
rebuilt_param_indices_);
// For rebuilt bucket indices, it needs to be synced across all ranks.
// Broadcast the newly rebuilt bucket indices from rank 0 in default.
// After syncing up rebuilt bucket indices, initialize buckets for reducer.
// 同步桶indices
sync_bucket_indices(rebuilt_bucket_indices);
has_rebuilt_bucket_ = true;
rebuilt_params_.clear();
rebuilt_param_indices_.clear();
// 初始化桶
initialize_buckets(std::move(rebuilt_bucket_indices));
return true;
}
我們接下來具體看看如何重建。
3.2.1 計算桶尺寸
我們首先要看看compute_bucket_assignment_by_size 之中關鍵結構如下,BucketAccumulator 可以認為是實際的桶。
struct BucketAccumulator {
std::vector<size_t> indices; // 桶內容,是張量列表
size_t size = 0; // 桶大小,比如若干mb
}; // 桶的邏輯內容
// Keep vector of indices and size accumulator by tensor type and device.
std::unordered_map<BucketKey, BucketAccumulator, c10::hash<BucketKey>>
buckets; // 所有桶的列表,每一個實際桶可以認為是 BucketAccumulator
其次,我們來看看 compute_bucket_assignment_by_size的具體邏輯:
-
生成一個計算結果 result,並且使用引數tensors的大小來為result預留出空間。
-
生成一個buckets,這是所有桶的列表,每一個實際桶可以認為是 BucketAccumulator
-
遍歷傳入的所有張量,對於每一個張量:
- 如果有index,就拿到張量的index。
- 如果配置了期待sparse gradient,則把這個張量自己放入一個桶,因為沒法和其他張量放在一起。
- 使用張量資訊構建桶的key。
- 使用 key 找到對應的桶, 拿到BucketAccumulator。
- 向該桶的張量列表 indices 裡面插入新張量的index,indices 是 tensor index list。
- 增加對應桶大小。
- 如果需要,就設定成大小限制的初始值。
- 如果桶的尺寸大於最小值限制,就是說目前桶的尺寸已經達到了桶的最大限制,按說需要轉移到新桶了(實際上確實轉移到了邏輯的新桶,但是實際還是在現有桶內執行,因為 type, device 還是同樣的,還是應該在原有桶內繼續累積,不過原有桶的indice已經轉移到了result之中,就相當於清空了)。
- 把桶內容插入到返回result,就是說,當桶尺寸過大的時候,就先插入到result之中。
- 利用 BucketAccumulator() 重新生成桶,bucket是個引用,所以直接賦值,就相當於清空原有的桶,就是原來桶繼續用,但是桶內原有的indices已經轉移到了result之中。
-
把剩餘的桶內indices插入到返回值result。之前已經有些直接插入到了result之中。
-
對 result 進行排序:
- 如果 tensor_indices 非空,說明張量的順序已經是梯度準備好的順序,不需要再排序了。
- 如果 tensor_indices 是空的,依據最小張量index來排序,這裡假定張量的順序是他們使用的順序(或者說是他們梯度產生次序的反序)。這種排序可保證桶是按照連續不斷的順序準備好。
- 注意,這裡就是正序排列,等到建立Reducer的時候,才反序傳入:list(reversed(bucket_indices))
另外需要注意的是:因為 tensors就是 Python 程式碼中的引數 parameters[0],而 parameters[0] 是按照 parametes() 的返回結果來的,所以DDP最終是按model.parameters()的相反順序啟動AllReduce。
std::vector<std::vector<size_t>> compute_bucket_assignment_by_size(
const std::vector<at::Tensor>& tensors,
const std::vector<size_t>& bucket_size_limits, // 桶大小限制
const std::vector<bool>& expect_sparse_gradient,
const std::vector<int64_t>& tensor_indices) { //實際上,初始化時候沒有傳入 tensor_indices
// Either expect_sparse_gradient is not specified or it has as many elements
// as the vector with tensors.
TORCH_INTERNAL_ASSERT(
expect_sparse_gradient.empty() ||
(tensors.size() == expect_sparse_gradient.size()));
TORCH_INTERNAL_ASSERT(tensors.size() > 0);
std::vector<std::vector<size_t>> result;
result.reserve(tensors.size()); // 預留大小
// Keep iterator into the size_limit vector by tensor type and device.
// This is done so that we can use the consecutive bucket limits per type.
std::unordered_map<
BucketKey,
std::vector<size_t>::const_iterator,
c10::hash<BucketKey>>
bucket_size_limit_iterators;
// Local accumulator type for a single bucket.
struct BucketAccumulator {
std::vector<size_t> indices; // 桶內容,是張量列表
size_t size = 0; // 桶大小,比如若干mb
}; // 桶的邏輯內容
// Keep vector of indices and size accumulator by tensor type and device.
std::unordered_map<BucketKey, BucketAccumulator, c10::hash<BucketKey>>
buckets; // 所有桶的列表,每一個實際桶可以認為是 BucketAccumulator
for (size_t i = 0; i < tensors.size(); i++) { // 遍歷傳入的所有張量
const auto& tensor = tensors[i]; //拿到張量
TORCH_CHECK(!tensor.is_sparse(), "No support for sparse tensors.");
// when tensor_indices is empty, the index of tensors[i] assigned to
// bucket is i, otherwise the tensor index is tensor_indices[i].
auto tensor_index = i; // 就是給所有的tensor一個index,從0開始遞增,一直到 tensors.size()
if (!tensor_indices.empty()) {
tensor_index = tensor_indices[i]; // 如果有index,就拿到張量的index
}
// If we expect a sparse gradient to be produced for this tensor, it cannot
// be grouped together with other gradients and gets its own bucket.
// 如果配置了期待sparse gradient,則把這個張量自己放入一個桶,因為沒法和其他張量放在一起
if (!expect_sparse_gradient.empty() &&
expect_sparse_gradient[tensor_index]) {
result.push_back({tensor_index});
continue;
}
auto key = BucketKey(tensor.scalar_type(), tensor.device()); //使用張量資訊構建桶的key
auto& bucket = buckets[key]; // 找到對應的桶, 拿到BucketAccumulator
bucket.indices.push_back(tensor_index); // 該桶的張量列表裡面插入新張量的index,indices 是 tensor index list
bucket.size += tensor.numel() * tensor.element_size();// 增加對應桶大小
// Initialize bucket size limit iterator if necessary.
// 如果需要,就設定成大小限制的初始值
if (bucket_size_limit_iterators.count(key) == 0) {
bucket_size_limit_iterators[key] = bucket_size_limits.begin();
}
// bucket_size_limit_iterator 就是桶大小的範圍, 即 [_DEFAULT_FIRST_BUCKET_BYTES, int(bucket_cap_mb * 1024 * 1024)]
auto& bucket_size_limit_iterator = bucket_size_limit_iterators[key];
const auto bucket_size_limit = *bucket_size_limit_iterator; // 當前最小值限制
if (bucket.size >= bucket_size_limit) {
// 如果桶的尺寸大於最小值限制,就是說目前桶的尺寸已經達到了桶的最大限制,按說需要轉移到新桶了(實際上確實轉移到了邏輯的新桶,但是實際還是在現有桶內執行,因為 type, device 還是同樣的,還是應該在原有桶內繼續累積,不過原有桶的indice已經轉移到了result之中,就相當於清空了)
result.emplace_back(std::move(bucket.indices)); // 把桶內容插入到返回result,就是說,當桶尺寸過大的時候,就先插入到result之中。
bucket = BucketAccumulator(); // 重新生成桶,bucket是個引用,所以直接賦值,就相當於清空原有的桶,就是原來桶繼續用,但是桶內原有的indices已經轉移到了result之中。
// Advance to the next bucket size limit for this type/device.
// 前進到下一個尺寸限制
auto next = bucket_size_limit_iterator + 1;
if (next != bucket_size_limits.end()) {
bucket_size_limit_iterator = next;
}
}
}
// Add remaining buckets. 把剩餘的桶內indices插入到返回值,因為之前已經有些直接插入到了result之中
for (auto& it : buckets) {
auto& bucket = it.second;
if (!bucket.indices.empty()) {
result.emplace_back(std::move(bucket.indices));
}
}
// If tensor_indices is not empty, the order of the tensors is in the gradient
// ready order, so no need to sort.
// If tensor_indices is empty, sort resulting buckets by the minimum tensor
// index they include. We assume that the order of the tensors is the order in
// which they are used (or the reverse order in which their gradients are
// produced). This sorting step ensures that the buckets are ready in
// consecutive order.
// 如果 tensor_indices 非空,說明張量的順序已經是梯度準備好的順序,不需要再排序了
// 如果 tensor_indices 是空的,依據最小張量index來排序,這裡假定張量的順序是他們使用的順序(或者說是他們梯度產生次序的反序)。這種排序可保證桶是按照連續不斷的順序準備好。
// 注意,這裡就是正序排列,等到建立Reducer的時候,才反序傳入:list(reversed(bucket_indices))
if (tensor_indices.empty()) {
std::sort(
result.begin(),
result.end(),
[](const std::vector<size_t>& a, const std::vector<size_t>& b) {
// 對於任意兩個vector,排序的依據是:用這兩個vector之中最小index來排序
const auto amin = std::min_element(a.begin(), a.end()); // a中的最小index
const auto bmin = std::min_element(b.begin(), b.end()); // b中的最小index
return *amin < *bmin;
});
}
return result;
}
result 最終如下,裡面每個vector 都對應了一個bucket,裡面是都是 tensor 的 index,這裡都是從小到大順序排序。模型引數以(大致)Model.parameters()
與給定模型相反的順序分配到桶中 。使用相反順序的原因是因為 DDP 期望梯度在反向傳遞期間以大約該順序準備就緒。
+-----------------------------------------------------------------------+
| |
| <tensor index 1, tensor index 2, tensor index 3, tensor index 4> |
| |
| |
| <tensor index 5, tensor index 6, tensor 7> |
| |
| |
| ...... |
| |
| |
| <tensor index 8, tensor index 9, tensor index 10, tensor index 11> |
| |
+-----------------------------------------------------------------------+
3.2.2 同步桶indices
產生尺寸之後,就使用 sync_bucket_indices 同步桶的indices,其邏輯如下:
- 遍歷桶,把桶的大小都記錄到bucket_sizes。
- 配置TensorOptions。
- 把桶對應的indices和桶數目放入indices_tensor,這裡是通過 PyTorch accessor來對張量進行讀寫,accessor就像是一個張量,但它將張量的維度和 dtype 硬編碼為了模板引數,可以高效的訪問元素。
- 因為 NCCL這樣的 ProcessGroup 只支援device之間的操作,所以把indices_tensor拷貝到indices_tensor_device。
- 對 indices_tensor_device 進行廣播。
- 類似,對桶尺寸進行廣播。
- 廣播結束之後,遍歷桶,使用從rank 0收到的num_buckets, bucket_sizes_tensor 和 indices_tensor 更新傳進來的引數bucket_indices。
void Reducer::sync_bucket_indices(
std::vector<std::vector<size_t>>& bucket_indices) {
auto num_buckets = bucket_indices.size();
std::vector<size_t> bucket_sizes;
bucket_sizes.reserve(num_buckets);
int64_t total_size = 0;
//遍歷桶,把桶的大小都記錄到bucket_sizes
for (size_t i = 0; i < num_buckets; i++) {
auto bucket_size = bucket_indices.at(i).size();
bucket_sizes.push_back(bucket_size);
total_size += bucket_size;
}
// 配置TensorOptions
at::TensorOptions options;
options = options.dtype(at::kInt);
options = options.device(replicas_[0][0].device());
// Group indices and num_bucket together into indices_tensor
// Broadcast this tensor first, as its size is equal among all processes
// 把桶對應的indices和桶數目放入indices_tensor,這裡是通過 PyTorch accessor來對張量進行讀寫,accessor就像是一個張量,但它將張量的維度和 dtype 硬編碼為了模板引數,可以高效的訪問元素
auto indices_tensor = at::empty({total_size + 1}, at::kInt);
auto indices_accessor = indices_tensor.accessor<int, 1>();
auto indices_accessor_Index = 0;
for (size_t i = 0; i < num_buckets; i++) {
const auto& bucket_size = bucket_indices.at(i).size();
for (size_t j = 0; j < bucket_size; j++) {
indices_accessor[indices_accessor_Index++] = bucket_indices[i][j];
}
}
indices_accessor[indices_accessor_Index] = num_buckets;
// Copy CPU tensor to device tensor, as the process_group_ could be NCCL and
// it can only broadcast device tensors.
auto indices_tensor_device = at::empty({total_size + 1}, options);
// 因為 NCCL這樣的 ProcessGroup 只支援device之間的操作,所以把indices_tensor拷貝到indices_tensor_device
indices_tensor_device.copy_(indices_tensor, /*non_blocking=*/true);
std::vector<at::Tensor> indices_tensor_list = {indices_tensor_device};
// 對 indices_tensor_device 進行廣播
process_group_->broadcast(indices_tensor_list)->wait();
indices_tensor.copy_(indices_tensor_list.front(), /*non_blocking=*/false);
// Update num_buckets after receiving it from rank 0
num_buckets = indices_accessor[indices_accessor_Index];
// Broadcast bucket_sizes
// 類似,對桶尺寸進行廣播
auto bucket_sizes_tensor = at::empty({(int64_t)num_buckets}, at::kInt);
auto bucket_sizes_accessor = bucket_sizes_tensor.accessor<int, 1>();
for (size_t i = 0; i < num_buckets; i++) {
// For rank != 0, it is possible that local num buckets bucket_sizes.size()
// is smaller than broadcasted num_buckets
bucket_sizes_accessor[i] =
bucket_sizes.at(std::min(i, (bucket_sizes.size() - 1)));
}
auto bucket_sizes_tensor_device = at::empty({(int64_t)num_buckets}, options);
bucket_sizes_tensor_device.copy_(bucket_sizes_tensor, /*non_blocking=*/true);
std::vector<at::Tensor> bucket_sizes_tensor_list = {
bucket_sizes_tensor_device};
process_group_->broadcast(bucket_sizes_tensor_list)->wait();
bucket_sizes_tensor.copy_(
bucket_sizes_tensor_list.front(), /*non_blocking=*/false);
// Clear bucket_indices first, and then update bucket_indices using received
// num_buckets, bucket_sizes_tensor and indices_tensor from rank 0
bucket_indices.clear();
bucket_indices.reserve(num_buckets);
indices_accessor_Index = 0;
// 遍歷桶,使用從rank 0收到的num_buckets, bucket_sizes_tensor 和 indices_tensor 更新傳進來的引數bucket_indices
for (size_t i = 0; i < num_buckets; i++) {
const auto& bucket_size = bucket_sizes_accessor[i];
std::vector<size_t> bucket;
bucket.reserve(bucket_size);
for (size_t j = 0; j < bucket_size; j++) {
bucket.push_back(indices_accessor[indices_accessor_Index++]);
}
bucket_indices.emplace_back(std::move(bucket));
}
}
3.2.3 初始化桶
同步之後就是初始化桶,本部分程式碼在前文已經分析過,故此省略。
3.3 準備後向傳播
前向傳播完成之後,呼叫 prepare_for_backward 完成了後向傳播的準備。
具體大致分為兩步:重置,查詢未使用的引數。
void Reducer::prepare_for_backward(
const std::vector<torch::autograd::Variable>& outputs) {
std::lock_guard<std::mutex> lock(mutex_);
// 記錄開始時間
cpu_timer_.backward_compute_start_time = current_time_in_nanos();
if (should_collect_runtime_stats()) {
record_backward_compute_start_time();
}
// Reset accounting.
expect_autograd_hooks_ = true;
reset_bucket_counting();
// Reset unused parameter accounting.
has_marked_unused_parameters_ = false;
// Reset per iteration marked ready parameters.
perIterationReadyParams_.clear(); // 重置每次迭代的marked ready parameters
// If static graph is not set, search graph to detect unused parameters.
// When static graph is set, unused_parameters_ will be detected and will
// not change after 1st iteration.
// If static_graph_ = false and find_unused_parameters_ is false,
// we assume that autograd hooks for ALL variables will be called,
// and we don't have to search the autograd graph for presence of these hooks.
if (dynamic_graph_find_unused()) {
unused_parameters_.clear();
search_unused_parameters(outputs); // 查詢沒有使用的引數
}
}
3.3.1 重置
這裡會遍歷桶,對於每個桶,重置其副本的pending狀態,某一個模型副本pending狀態是由這個模型副本中對應桶的變數數目決定。
如果是靜態圖,則重置numGradHooksTriggeredMapPerIteration_。
void Reducer::reset_bucket_counting() {
next_bucket_ = 0;
// Reset num_buckets_ready_ at the beginning of backward computation
// in each iteration.
num_buckets_ready_ = 0;
for (auto& bucket : buckets_) { // 遍歷桶
for (auto& replica : bucket.replicas) {
replica.pending = replica.variables.size(); //對於每個桶,重置其副本的pending狀態,某一個模型副本pending,是由這個模型副本中,本桶的變數數目決定
}
bucket.pending = bucket.replicas.size(); // 重置桶的pending狀態,桶pending是由多少個模型副本決定
}
if (static_graph_) {
// 重置numGradHooksTriggeredMapPerIteration_
numGradHooksTriggeredMapPerIteration_ = numGradHooksTriggeredMap_;
}
}
3.3.2 查詢未使用的引數
search_unused_parameters 完成了 "查詢未使用的引數" 功能。
我們首先要看看 Reducer 的 find_unused_parameters_ 成員變數。如果 find_unused_parameters_ 被設定為 true,則 DDP 會在前向傳播結束時候,從指定的輸出進行回溯,遍歷autograd計算圖來找到所有沒有使用過的引數,並且一一標記為就緒 ready。
對於所有引數,DDP 都有一個指向它們的梯度累積函式的指標,但對於那些autograd圖中不存在的引數,它們將在第一次呼叫autograd鉤子時就被標記為準備就緒。
因為模型輸出可能會被忽略,所以這個操作不是立即完成的,我們只是像在torch.autograd.backward()
這裡開始執行規約操作。
大家可以發現,這麼做開銷會很大,為什麼要這麼做?這是因為計算動態圖會改變。
- 訓練時候,某次迭代可能只用到模型的一個子圖,而且因為PyTorch 是動態計算,所以子圖會在迭代期間改變,就是說,某些引數可能在下一次迭代訓練時候被跳過。
- 同時,因為所有引數在一開始就已經被分好桶,而 hook 又規定了只有整個桶 ready (即,pending == 0)之後才會進行通訊,所以如果我們不將未使用引數標記為 ready,整個通訊過程就會沒法進行。
// Traverse the autograd graph starting at the specified output.
// All parameters for which we have a pointer to their gradient accumulation
// functions, but don't show up in the autograd graph will be marked ready for
// for reduction as soon as the first autograd hook is called. This is not
// done immediately because the model output may be ignored, and we only
// want to start performing reductions on `torch.autograd.backward()`.
void Reducer::search_unused_parameters(
const std::vector<torch::autograd::Variable>& outputs) {
std::unordered_set<torch::autograd::Node*> seen;
std::vector<torch::autograd::Node*> queue;
RECORD_FUNCTION(
"torch.distributed.ddp.reducer::search_unused_parameters",
std::vector<c10::IValue>());
// Seed queue with the grad functions of all outputs.
for (const auto& output : outputs) {
const auto& grad_fn = output.grad_fn();
if (grad_fn) {
queue.push_back(grad_fn.get()); // 把所有輸出節點的梯度函式插入到queue
}
}
// Traverse the autograd graph starting at the specified output.
// 遍歷這個queue中的元素,對於每一個函式,找到其後向圖之中的後續邊,然後把後續邊指向的節點再插入queue,然後繼續迴圈,最終 seen 裡面是所有從output出發,所有節點的梯度函式
while (!queue.empty()) {
auto fn = queue.back();
queue.pop_back();
for (const auto& edge : fn->next_edges()) {
if (auto next_ptr = edge.function.get()) {
const bool was_inserted = seen.insert(next_ptr).second;
if (was_inserted) {
queue.push_back(next_ptr);
}
}
}
}
// Find accumulator functions that don't show up in this graph.
// gradAccToVariableMap_ 裡面是所有需要被規約的variable
// 遍歷gradAccToVariableMap_,如果 seen 之中沒有,就說明這個引數沒有被使用,插入到unused_parameters_
for (const auto& it : gradAccToVariableMap_) {
// If the accumulator function is present in the graph, we know
// a gradient will be computed for the corresponding parameter.
if (seen.count(it.first) == 0) {
unused_parameters_.push_back(it.second);
}
}
// Warn user about unnecessary perf hit if all parameters were used in
// forward.
if (unused_parameters_.empty()) {
TORCH_WARN_ONCE(
"find_unused_parameters=True was specified in DDP constructor, "
"but did not find any unused parameters in the forward pass. This flag "
"results in an extra traversal of the autograd graph every iteration, "
" which can adversely affect performance. If your model indeed never "
"has any unused parameters in the forward pass, consider turning this "
"flag off. Note that this warning may be a false positive if your model "
"has flow control causing later iterations to have unused parameters.");
}
}
至此,前向傳播已經結束,我們得到了如下:
- 需要計算梯度的引數已經分桶。
- 桶已經重建完畢。
- 前向傳播已經完成。
- 從指定的輸出進行回溯,遍歷autograd計算圖來找到所有沒有使用過的引數,並且一一標記為就緒 ready。
我們在下一篇就分析後向傳播。
0xFF 參考
pytorch分散式系列3——分散式訓練時,torch.utils.data.distributed.DistributedSampler做了什麼?
pytorch分散式系列1——搞清torch.distributed.launch相關的環境變數
pytorch分散式系列2——DistributedDataParallel是如何做同步的?
pytorch(分散式)資料並行個人實踐總結——DataParallel/DistributedDataParallel
https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/20
https://pytorch.org/docs/stable/distributed.html
pytorch分散式訓練(二init_process_group)
https://pytorch.org/tutorials/intermediate/ddp_tutorial.html
https://pytorch.org/docs/master/notes/ddp.html
https://pytorch.org/tutorials/intermediate/dist_tuto.html