1 導引
我們在部落格《分散式機器學習：同步並行SGD演算法的實現與複雜度分析（PySpark）》和部落格《分散式機器學習：模型平均MA與彈性平均EASGD（PySpark） 》中介紹的都是同步演算法。同步演算法的共性是所有的節點會以一定的頻率進行全域性同步。然而，當工作節點的計算效能存在差異，或者某些工作節點無法正常工作（比如當機）的時候，分散式系統的整體執行效率不好，甚至無法完成訓練任務。為了解決此問題，人們提出了非同步的並行演算法。在非同步的通訊模式下，各個工作節點不需要互相等待，而是以一個或多個全域性伺服器做為中介，實現對全域性模型的更新和讀取。這樣可以顯著減少通訊時間，從而獲得更好的多機擴充套件性。
2 非同步SGD
2.1 演算法描述與實現
非同步SGD[9]是最基礎的非同步演算法，其流暢如下圖所示。粗略地講，ASGD的引數更新發生在工作節點，而模型的更新發生在伺服器端。當引數伺服器接收到來自某個工作節點的引數梯度時，就直接將其加到全域性模型上，而無需等待其它工作節點的梯度資訊。
下面我們用Pytorch實現的訓練程式碼（採用RPC進行程式間通訊）。首先，我們設定初始化多個程式，其中0號程式做為引數伺服器，其餘程式做為worker來對模型進行訓練，則總的通訊域（world_size）大小為workers的數量+1。這裡我們設定引數伺服器IP地址為localhost，埠號29500。
def run(rank, world_size):
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'
options=rpc.TensorPipeRpcBackendOptions(
num_worker_threads=16,
rpc_timeout=0 # infinite timeout
)
if rank == 0:
rpc.init_rpc(
"ps",
rank=rank,
world_size=world_size,
rpc_backend_options=options
)
run_ps([f"trainer{r}" for r in range(1, world_size)])
else:
rpc.init_rpc(
f"trainer{rank}",
rank=rank,
world_size=world_size,
rpc_backend_options=options
)
# trainer passively waiting for ps to kick off training iterations
# block until all rpcs finish
rpc.shutdown()
if __name__=="__main__":
world_size = n_workers + 1
mp.spawn(run, args=(world_size, ), nprocs=world_size, join=True)
下面我們定義引數伺服器的所要完成工作流程，包括將訓練資料劃分到各個worker，非同步呼叫所有worker的訓練流程，最後訓練完畢後在引數伺服器完成對模型的評估。
def run_trainer(ps_rref, train_dataset):
trainer = Trainer(ps_rref)
trainer.train(train_dataset)
def run_ps(trainers):
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST('./data', train=True, download=True,
transform=transform)
local_train_datasets = dataset_split(train_dataset, n_workers)
print(f"{datetime.now().strftime('%H:%M:%S')} Start training")
ps = ParameterServer()
ps_rref = rpc.RRef(ps)
futs = []
for idx, trainer in enumerate(trainers):
futs.append(
rpc.rpc_async(trainer, run_trainer, args=(ps_rref, local_train_datasets[idx]))
)
torch.futures.wait_all(futs)
print(f"{datetime.now().strftime('%H:%M:%S')} Finish training")
ps.evaluation()
這裡資料集的劃分程式碼採用我們在《Pytorch：單卡多程式並行訓練》中所述的資料劃分方式：
class CustomSubset(Subset):
'''A custom subset class with customizable data transformation'''
def __init__(self, dataset, indices, subset_transform=None):
super().__init__(dataset, indices)
self.subset_transform = subset_transform
def __getitem__(self, idx):
x, y = self.dataset[self.indices[idx]]
if self.subset_transform:
x = self.subset_transform(x)
return x, y
def __len__(self):
return len(self.indices)
def dataset_split(dataset, n_workers):
n_samples = len(dataset)
n_sample_per_workers = n_samples // n_workers
local_datasets = []
for w_id in range(n_workers):
if w_id < n_workers - 1:
local_datasets.append(CustomSubset(dataset, range(w_id * n_sample_per_workers, (w_id + 1) * n_sample_per_workers)))
else:
local_datasets.append(CustomSubset(dataset, range(w_id * n_sample_per_workers, n_samples)))
return local_datasets
以下是引數伺服器類ParameterServer的定義：
class ParameterServer(object):
def __init__(self, n_workers=n_workers):
self.model = Net().to(device)
self.lock = threading.Lock()
self.future_model = torch.futures.Future()
self.n_workers = n_workers
self.curr_update_size = 0
self.optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9)
for p in self.model.parameters():
p.grad = torch.zeros_like(p)
self.test_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=32, shuffle=True)
def get_model(self):
# TensorPipe RPC backend only supports CPU tensors,
# so we move your tensors to CPU before sending them over RPC
return self.model.to("cpu")
@staticmethod
@rpc.functions.async_execution
def update_and_fetch_model(ps_rref, grads):
self = ps_rref.local_value()
for p, g in zip(self.model.parameters(), grads):
p.grad += g
with self.lock:
self.curr_update_size += 1
fut = self.future_model
if self.curr_update_size >= self.n_workers:
for p in self.model.parameters():
p.grad /= self.n_workers
self.curr_update_size = 0
self.optimizer.step()
self.optimizer.zero_grad()
fut.set_result(self.model)
self.future_model = torch.futures.Future()
return fut
def evaluation(self):
self.model.eval()
self.model = self.model.to(device)
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in self.test_loader:
output = self.model(data.to(device))
test_loss += F.nll_loss(output, target.to(device), reduction='sum').item() # sum up batch loss
pred = output.max(1)[1] # get the index of the max log-probability
correct += pred.eq(target.to(device)).sum().item()
test_loss /= len(self.test_loader.dataset)
print('\nTest result - Accuracy: {}/{} ({:.0f}%)\n'.format(
correct, len(self.test_loader.dataset), 100. * correct / len(self.test_loader.dataset)))
以下是Trainer類的定義：
class Trainer(object):
def __init__(self, ps_rref):
self.ps_rref = ps_rref
self.model = Net().to(device)
def train(self, train_dataset):
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
model = self.ps_rref.rpc_sync().get_model().cuda()
pid = os.getpid()
for epoch in range(epochs):
for batch_idx, (data, target) in enumerate(train_loader):
output = model(data.to(device))
loss = F.nll_loss(output, target.to(device))
loss.backward()
model = rpc.rpc_sync(
self.ps_rref.owner(),
ParameterServer.update_and_fetch_model,
args=(self.ps_rref, [p.grad for p in model.cpu().parameters()]),
).cuda()
if batch_idx % log_interval == 0:
print('{}\tTrain Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
pid, epoch + 1, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
完整程式碼我已經上傳到了GitHub倉庫 [Distributed-Algorithm-PySpark]，感興趣的童鞋可以前往檢視。
執行該程式碼得到的評估結果為：
Test result - Accuracy: 9696/10000 (97%)
可見該訓練演算法是收斂的，但在10個epoch下在測試集上只能達到97%的精度，不如我們下面提到的在10個epoch就能在測試集上達到99%精度的Hogwild!演算法。注意，ASGD和Hogwild都是非同步演算法，但ASGD是分散式演算法（當然我們這裡是單機多程式模擬），程式間採用RPC通訊，不會出現同步錯誤的問題，根本不需要考慮加不加鎖。而Hogwild!演算法是單機演算法，程式/執行緒間採用共享記憶體通訊，需要考慮加不加鎖的問題，不過Hogwild！演算法為了提高訓練過程中的資料吞吐量，直接採用了無鎖的全域性記憶體訪問。
2.2 收斂性分析
ASGD避開了同步開銷，但會給模型更新增加一些延遲。我們下面將ASGD的工作流程用下圖加以剖析來解釋這一點。用\(\text{worker}(k)\)來代表第\(k\)個工作節點，用\(w^t\)來代表第\(t\)輪迭代時服務端的全域性模型。按照時間順序，首先\(\text{worker}(k)\)先從引數伺服器獲取全域性模型\(w^t\)，再根據本地資料計算模型梯度\(g(w_t)\)並將其發往引數伺服器。一段時間後，\(\text{worker}(k')\)也從引數伺服器取回當時的全域性模型\(w^{t+1}\)，並同樣依據它的本地資料計算模型的梯度\(f(w^{t+1})\)。注意，在\(\text{worker}(k')\)取回引數並進行計算的過程中，其它工作節點（比如\(\text{worker}(k)\)）可能已經將它的梯度提交給伺服器並進行更新了。所以當\(\text{worker}(k')\)將其梯度\(g(w^{t+1})\)發給伺服器時，全域性模型已經不再是\(w^{t+1}\)，而已經是被更新過的版本。
我們將上面這種現象稱為梯度和模型的失配，也即我們用一個比較舊的引數計算了梯度，而將這個“延遲”的梯度更新到了模型引數上。這種延遲使得ASGD和SGD之間在引數更新規則上存在偏差，可能導致模型在某些特定的更新點上出現嚴重抖動，設定最佳化過程出錯，無法收斂。後面我們會介紹克服延遲問題的手段。
3 Hogwild!演算法
3.1 演算法描述與實現
非同步並行演算法既可以在多機叢集上開展，也可以在多核系統下透過多執行緒開展。當我們把ASGD演算法應用在多執行緒環境中時，因為不再有引數伺服器這一角色，演算法的細節會發生一些變化。特別地，因為全域性模型儲存在共享記憶體中，所以當非同步的模型更新發生時，我們需要討論是否將記憶體加鎖，以保證模型寫入的一致性。
Hogwild！演算法[2]為了提高訓練過程中的資料吞吐量，選擇了無鎖的全域性模型訪問，其工作邏輯如下所示：
這裡使用我們在《Pytorch：單卡多程式並行訓練》所提到的torch.multiprocessing來進行多程式並行訓練。多程式原本記憶體空間是獨立的，這裡我們顯式呼叫model.share_memory()來講模型設定在共享記憶體中以進行程式間通訊。不過注意，如果我們採用GPU訓練，則GPU直接就做為了多程式的共享記憶體，此時model.share_memory()實際上為空操作（no-op）。
我們用Pytorch實現的訓練程式碼如下：
from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.multiprocessing as mp
from torchvision import datasets, transforms
import os
import torch
import torch.optim as optim
import torch.nn.functional as F
batch_size = 64 # input batch size for training
test_batch_size = 1000 # input batch size for testing
epochs = 10 # number of global epochs to train
lr = 0.01 # learning rate
momentum = 0.5 # SGD momentum
seed = 1 # random seed
log_interval = 10 # how many batches to wait before logging training status
n_workers = 4 # how many training processes to use
cuda = True # enables CUDA training
mps = False # enables macOS GPU training
dry_run = False # quickly check a single pass
def train(rank, model, device, dataset, dataloader_kwargs):
torch.manual_seed(seed + rank)
train_loader = torch.utils.data.DataLoader(dataset, **dataloader_kwargs)
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
for epoch in range(1, epochs + 1):
model.train()
pid = os.getpid()
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = model(data.to(device))
loss = F.nll_loss(output, target.to(device))
loss.backward()
optimizer.step()
if batch_idx % log_interval == 0:
print('{}\tTrain Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
pid, epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
if dry_run:
break
def test(model, device, dataset, dataloader_kwargs):
torch.manual_seed(seed)
test_loader = torch.utils.data.DataLoader(dataset, **dataloader_kwargs)
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
output = model(data.to(device))
test_loss += F.nll_loss(output, target.to(device), reduction='sum').item() # sum up batch loss
pred = output.max(1)[1] # get the index of the max log-probability
correct += pred.eq(target.to(device)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Global loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
if __name__ == '__main__':
use_cuda = cuda and torch.cuda.is_available()
use_mps = mps and torch.backends.mps.is_available()
if use_cuda:
device = torch.device("cuda")
elif use_mps:
device = torch.device("mps")
else:
device = torch.device("cpu")
print(device)
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST('../data', train=True, download=True,
transform=transform)
test_dataset = datasets.MNIST('../data', train=False,
transform=transform)
kwargs = {'batch_size': batch_size,
'shuffle': True}
if use_cuda:
kwargs.update({'num_workers': 1,
'pin_memory': True,
})
torch.manual_seed(seed)
mp.set_start_method('spawn', force=True)
model = Net().to(device)
model.share_memory() # gradients are allocated lazily, so they are not shared here
processes = []
for rank in range(n_workers):
p = mp.Process(target=train, args=(rank, model, device,
train_dataset, kwargs))
# We first train the model across `n_workers` processes
p.start()
processes.append(p)
for p in processes:
p.join()
# Once training is complete, we can test the model
test(model, device, test_dataset, kwargs)
執行得到的評估結果為：
Test set: Global loss: 0.0325, Accuracy: 9898/10000 (99%)
可見該訓練演算法是收斂的。
3.2 收斂性分析
當採用不帶鎖的多執行緒的寫入（即在更新\(w_j\)的時候不用先獲取對\(w_j\)的訪問許可權），而這可能會導致導致同步錯誤[10]的問題。比如線上程\(1\)載入全域性引數\(w^t_j\)之後，執行緒\(2\)還沒等執行緒\(1\)儲存全域性引數更新後的值，就也對全域性引數\(w^t_j\)進行載入，這樣導致每個執行緒都會儲存值為\(w^t_j - \eta^t g(w^t_j)\)的更新後的全域性引數值，這樣就導致其中一個執行緒的更新實際上在做“無用功”。直觀的感覺是這應該會對學習的過程產生負面影響。不過，當我們對模型訪問的稀疏性（sparity）做一定的限定後，這種訪問衝突實際上是非常有限的。這正是Hogwild！演算法收斂性存在的理論依據。
假設我們要最小化的損失函式為\(l: \mathcal{W}\rightarrow \mathbb{R}\)，對於特定的訓練樣本集合，損失函式\(l\)是由一系列稀疏子函式組合而來的：
也就是說，實際的學習過程中，每個訓練樣本涉及的引數組合\(e\)只是全體引數集合中的一個很小的子集。我們可以用一個超圖\(G=(V, E)\)來表述這個學習過程中引數和引數之間的關係，其中節點\(v\)表示引數，而超邊\(e\)表示訓練樣本涉及的引數組合。那麼，稀疏性可以用下面幾個統計量加以表示：
其中，\(\Omega\)表達了最大超邊的大小，也就是單個樣本最多涉及的引數個數；\(\Delta\)反映的是一個引數最多可以涉及多少個不同的超邊；而\(\rho\)則反映了給定任意一個超邊，與其共享引數的超邊個數。這三個值的取值越小，則最佳化問題越稀疏。在\(\Omega\)、\(\Delta\)、\(\rho\)都比較小的條件下，Hogwild！演算法的收斂性保證還需要假設損失函式是凸函式，並且是Lipschitz連續的，詳細的理論證明和定量關係請參考文獻[2]。
參考
-
[1] Agarwal A, Duchi J C. Distributed delayed stochastic optimization[J]. Advances in neural information processing systems, 2011, 24.
-
[2] Recht B, Re C, Wright S, et al. Hogwild!: A lock-free approach to parallelizing stochastic gradient descent[J]. Advances in neural information processing systems, 2011, 24.
-
[3] 劉浩洋，戶將等. 最最佳化：建模、演算法與理論[M]. 高教出版社, 2020.
-
[4] 劉鐵巖，陳薇等. 分散式機器學習：演算法、理論與實踐[M]. 機械工業出版社, 2018.
-
[5] Stanford CME 323: Distributed Algorithms and Optimization (Lecture 7)
-
[6] Bryant R E等.《深入理解計算機系統》[M]. 機械工業出版社, 2015.