分散式機器學習:非同步SGD和Hogwild!演算法(Pytorch)

orion發表於2023-02-13

1 導引

我們在部落格《分散式機器學習:同步並行SGD演算法的實現與複雜度分析(PySpark)》和部落格《分散式機器學習:模型平均MA與彈性平均EASGD(PySpark) 》中介紹的都是同步演算法。同步演算法的共性是所有的節點會以一定的頻率進行全域性同步。然而,當工作節點的計算效能存在差異,或者某些工作節點無法正常工作(比如當機)的時候,分散式系統的整體執行效率不好,甚至無法完成訓練任務。為了解決此問題,人們提出了非同步的並行演算法。在非同步的通訊模式下,各個工作節點不需要互相等待,而是以一個或多個全域性伺服器做為中介,實現對全域性模型的更新和讀取。這樣可以顯著減少通訊時間,從而獲得更好的多機擴充套件性。

2 非同步SGD

2.1 演算法描述與實現

非同步SGD[9]是最基礎的非同步演算法,其流暢如下圖所示。粗略地講,ASGD的引數更新發生在工作節點,而模型的更新發生在伺服器端。當引數伺服器接收到來自某個工作節點的引數梯度時,就直接將其加到全域性模型上,而無需等待其它工作節點的梯度資訊。

分散式機器學習:非同步SGD和Hogwild!演算法(Pytorch)

下面我們用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}\),而已經是被更新過的版本。

分散式機器學習:非同步SGD和Hogwild!演算法(Pytorch)

我們將上面這種現象稱為梯度和模型的失配,也即我們用一個比較舊的引數計算了梯度,而將這個“延遲”的梯度更新到了模型引數上。這種延遲使得ASGD和SGD之間在引數更新規則上存在偏差,可能導致模型在某些特定的更新點上出現嚴重抖動,設定最佳化過程出錯,無法收斂。後面我們會介紹克服延遲問題的手段。

3 Hogwild!演算法

3.1 演算法描述與實現

非同步並行演算法既可以在多機叢集上開展,也可以在多核系統下透過多執行緒開展。當我們把ASGD演算法應用在多執行緒環境中時,因為不再有引數伺服器這一角色,演算法的細節會發生一些變化。特別地,因為全域性模型儲存在共享記憶體中,所以當非同步的模型更新發生時,我們需要討論是否將記憶體加鎖,以保證模型寫入的一致性。

Hogwild!演算法[2]為了提高訓練過程中的資料吞吐量,選擇了無鎖的全域性模型訪問,其工作邏輯如下所示:

分散式機器學習:非同步SGD和Hogwild!演算法(Pytorch)

這裡使用我們在《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\)是由一系列稀疏子函式組合而來的:

\[l(w) = \sum_{e\in E}f_e(w_e) \]

也就是說,實際的學習過程中,每個訓練樣本涉及的引數組合\(e\)只是全體引數集合中的一個很小的子集。我們可以用一個超圖\(G=(V, E)\)來表述這個學習過程中引數和引數之間的關係,其中節點\(v\)表示引數,而超邊\(e\)表示訓練樣本涉及的引數組合。那麼,稀疏性可以用下面幾個統計量加以表示:

\[\Omega:=\max_{e\in E}|e|\\ \Delta:=\frac{\underset{1\leqslant v \leqslant n}{\max}|\{e\in E: v\in e\}|}{|E|}\\ \rho:=\frac{\underset{e\in E}{\max}|\{\hat{e}\in E: \hat{e}\cap e \neq \emptyset \}|}{|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.

相關文章