Merge pull request #1133 from pytorch/xlmr_mlm_two_shards

pbelevich · web-flow · commit 213f0efefe77 · 2021-02-08T16:48:08.000-05:00
DistCrossLingualMLMTask with two shards
diff --git a/examples/BERT/README.md b/examples/BERT/README.md
@@ -151,6 +151,10 @@ To run the workflow with 3000 lines from each of the 100 languages (CC-100 datas
 
     python cross_lingual_mlm_task.py --num_lines 3000
 
+To run the distributed training use '--dist' flag, to specify world size use '--world_size=N', the default world size is 3 for one master and 2 worker nodes.
+
+    python cross_lingual_mlm_task.py --num_lines 3000 --dist
+
 To Run the reference XLM-R model from fairseq, download and unzip the pretrained model from [link](https://dl.fbaipublicfiles.com/fairseq/models/xlmr.large.tar.gz).
 
     python cross_lingual_mlm_task.py --eval_ref ./xlmr.large 
diff --git a/examples/BERT/cross_lingual_mlm_task.py b/examples/BERT/cross_lingual_mlm_task.py
@@ -1,15 +1,24 @@
 import argparse
-import time
 import math
+import os
+import time
+from typing import List
+
 import torch
+import torch.distributed.autograd as dist_autograd
+import torch.distributed.rpc as rpc
+import torch.multiprocessing as mp
 import torch.nn as nn
+import torch.optim as optim
+from torch.distributed.optim import DistributedOptimizer
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader
+
 from data import CC100
+from dist_model import DistCrossLingualMLMTask
 from model import CrossLingualMLMTask
-from torch.utils.data import DataLoader
 from torchtext.experimental.transforms import sentencepiece_tokenizer
 from transforms import PretrainedSPVocab
-from torch.nn.utils.rnn import pad_sequence
-from typing import List
 
 
 def collate_batch(batch_data, args, mask_id, pad_id, text_transform):
@@ -47,8 +56,28 @@ def evaluate(data_source, model, mask_id, pad_id, ntokens, criterion, args, devi
     return total_loss / (len(data_source) - 1)  # Set batch # to 1 for inference
 
 
+def step(model, data, targets, criterion, optimizer, ntokens):
+    optimizer.zero_grad()
+    output = model(data)
+    loss = criterion(output.view(-1, ntokens), targets.view(-1))
+    loss.backward()
+    torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
+    optimizer.step()
+    return loss
+
+
+def dist_step(model, data, targets, criterion, optimizer, ntokens):
+    with dist_autograd.context() as context_id:
+        output = model(data)
+        loss = criterion(output.view(-1, ntokens), targets.view(-1))
+        dist_autograd.backward(context_id, [loss])
+        # torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
+        optimizer.step(context_id)
+        return loss
+
+
 def train(model, mask_id, pad_id, train_loss_log, train_data, text_transform,
-          optimizer, criterion, ntokens, epoch, scheduler, args, device, rank=None):
+          optimizer, criterion, ntokens, epoch, last_lr, args, device, step_impl):
     model.train()
     total_loss = 0.
     start_time = time.time()
@@ -57,33 +86,25 @@ def train(model, mask_id, pad_id, train_loss_log, train_data, text_transform,
                             shuffle=False, collate_fn=lambda b: collate_batch(b, args, mask_id, pad_id, text_transform))
 
     for batch, (data, targets) in enumerate(dataloader):
-        optimizer.zero_grad()
-        data = data.to(device)
-        targets = targets.to(device)
-        output = model(data)
-        loss = criterion(output.view(-1, ntokens), targets.view(-1))
-        loss.backward()
-        torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
-        optimizer.step()
+        loss = step_impl(model, data.to(device), targets.to(device), criterion, optimizer, ntokens)
+
         total_loss += loss.item()
         if batch % args.log_interval == 0 and batch > 0:
             cur_loss = total_loss / args.log_interval
             elapsed = time.time() - start_time
-            if (rank is None) or rank == 0:
-                train_loss_log[-1] = cur_loss
-                print('| epoch {:3d} | {:5d}/{:5d} batches | lr {:05.5f} | ms/batch {:5.2f} | '
-                      'loss {:5.2f} | ppl {:8.2f}'.format(epoch, batch,
-                                                          len(train_data) // args.batch_size,
-                                                          scheduler.get_last_lr()[0],
-                                                          elapsed * 1000 / args.log_interval,
-                                                          cur_loss, math.exp(cur_loss)))
+            train_loss_log[-1] = cur_loss
+            print('| epoch {:3d} | {:5d}/{:5d} batches | lr {:05.5f} | ms/batch {:5.2f} | '
+                    'loss {:5.2f} | ppl {:8.2f}'.format(epoch, batch,
+                                                        len(train_data) // args.batch_size,
+                                                        last_lr,
+                                                        elapsed * 1000 / args.log_interval,
+                                                        cur_loss, math.exp(cur_loss)))
             total_loss = 0
             start_time = time.time()
 
 
-def run_main(args, rank=None):
+def run_main(args):
     torch.manual_seed(args.seed)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     # Set up tokenizer and vocab
     tokenizer = sentencepiece_tokenizer(args.spm_path)
@@ -95,11 +116,23 @@ def text_transform(x: str) -> List:
     pad_id = vocab(['pad'])[0]
     ntokens = len(vocab)
 
-    model = CrossLingualMLMTask(ntokens, args.emsize, args.nhead, args.nhid, args.nlayers, args.dropout)
-    model = model.to(device)
+    if not args.dist:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = CrossLingualMLMTask(ntokens, args.emsize, args.nhead, args.nhid, args.nlayers, args.dropout)
+        model = model.to(device)
+        optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.75)
+    else:
+        device = "cpu"
+        model = DistCrossLingualMLMTask(args.split_size, ["worker1", "worker2"], ntokens, args.emsize, args.nhead, args.nhid, args.nlayers, args.dropout)
+        optimizer = DistributedOptimizer(
+            optim.Adam,
+            model.parameter_rrefs(),
+            lr=args.lr,
+        )
+        scheduler = None
+
     criterion = nn.CrossEntropyLoss(ignore_index=pad_id)
-    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
-    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.75)
     best_val_loss = None
     train_loss_log, val_loss_log = [], []
 
@@ -110,8 +143,9 @@ def text_transform(x: str) -> List:
         val_data = [(17, item) for item in val_data if item != ' \n']  # english language type is 17 in CC100 dataset
 
         epoch_start_time = time.time()
+        last_lr = scheduler.get_last_lr()[0] if scheduler is not None else args.lr
         train(model, mask_id, pad_id, train_loss_log, train_data, text_transform,
-              optimizer, criterion, ntokens, epoch, scheduler, args, device, rank)
+              optimizer, criterion, ntokens, epoch, last_lr, args, device, step if not args.dist else dist_step)
 
         # Turn on evaluation mode which disables dropout.
         model.eval()
@@ -122,12 +156,13 @@ def text_transform(x: str) -> List:
               'valid ppl {:8.2f}'.format(epoch, (time.time() - epoch_start_time),
                                          val_loss, math.exp(val_loss)))
         print('-' * 89)
-        if not best_val_loss or val_loss < best_val_loss:
+        if not args.dist and not best_val_loss or val_loss < best_val_loss:
             with open(args.save, 'wb') as f:
                 torch.save(model, f)
             best_val_loss = val_loss
         else:
-            scheduler.step()
+            if scheduler is not None:
+                scheduler.step()
 
     # Run reference XLM-R model from fairseq
     if args.eval_ref != 'None':
@@ -159,6 +194,32 @@ def _forward(x):
         print('-' * 89)
 
 
+def run_worker(rank, args):
+    os.environ['MASTER_ADDR'] = 'localhost'
+    os.environ['MASTER_PORT'] = '29500'
+    options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256)
+
+    if rank == 0:
+        rpc.init_rpc(
+            "master",
+            rank=rank,
+            world_size=args.world_size,
+            rpc_backend_options=options
+        )
+        run_main(args)
+    else:
+        rpc.init_rpc(
+            f"worker{rank}",
+            rank=rank,
+            world_size=args.world_size,
+            rpc_backend_options=options
+        )
+        pass
+
+    # block until all rpcs finish
+    rpc.shutdown()
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='PyTorch Cross-lingual XLM MLM')
     parser.add_argument('--emsize', type=int, default=768,
@@ -197,6 +258,15 @@ def _forward(x):
                         help='path to load the reference model for evaluation')
     parser.add_argument('--mask_frac', type=float, default=0.15,
                         help='the fraction of masked tokens')
+    parser.add_argument('--dist', action='store_true',
+                        help='run distributed version')
+    parser.add_argument('--world_size', type=int, default=3,
+                        help='world_size')
+    parser.add_argument('--split_size', type=int, default=8,
+                        help='split the input batch into micro-batches')
     args = parser.parse_args()
 
-    run_main(args)
+    if args.dist:
+        mp.spawn(run_worker, args=(args,), nprocs=args.world_size, join=True)
+    else:
+        run_main(args)
diff --git a/examples/BERT/dist_model.py b/examples/BERT/dist_model.py
@@ -0,0 +1,114 @@
+import threading
+
+import torch
+import torch.distributed.rpc as rpc
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributed.rpc import RRef
+from torch.nn import Linear, LayerNorm
+
+from model import XLMREmbedding, TransformerEncoderLayer, TransformerEncoder
+
+
+def get_cuda_if_available(i):
+    assert i >= 0
+    if torch.cuda.is_available():
+        return f"cuda:{min(i, torch.cuda.device_count() - 1)}"
+    else:
+        return "cpu"
+
+
+class CrossLingualMLMTaskBase(nn.Module):
+    def __init__(self, device):
+        super(CrossLingualMLMTaskBase, self).__init__()
+        self.device = device
+        self._lock = threading.Lock()
+
+    def forward(self, x_rref):
+        x = x_rref.to_here().to(self.device)
+        with self._lock:
+            out = self._forward(x)
+        return out.cpu()
+
+    def parameter_rrefs(self):
+        r"""
+        Create one RRef for each parameter in the given local module, and return a
+        list of RRefs.
+        """
+        return [RRef(p) for p in self.parameters()]
+
+
+class CrossLingualMLMTaskShard1(CrossLingualMLMTaskBase):
+    def __init__(self, device, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
+        super(CrossLingualMLMTaskShard1, self).__init__(device)
+        self.xlmr_embed = XLMREmbedding(ntoken, ninp, dropout).to(device)
+        encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)
+        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers // 2).to(device)
+
+    def _forward(self, src):
+        output = self.xlmr_embed(src)
+        output = self.transformer_encoder(output)
+        return output
+
+
+class CrossLingualMLMTaskShard2(CrossLingualMLMTaskBase):
+    def __init__(self, device, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
+        super(CrossLingualMLMTaskShard2, self).__init__(device)
+        encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)
+        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers // 2).to(device)
+        self.mlm_span = Linear(ninp, ninp).to(device)
+        self.activation = F.gelu
+        self.norm_layer = LayerNorm(ninp, eps=1e-12).to(device)
+        self.mlm_head = Linear(ninp, ntoken).to(device)
+
+    def _forward(self, src):
+        output = self.transformer_encoder(src)
+        output = self.mlm_span(output)
+        output = self.activation(output)
+        output = self.norm_layer(output)
+        output = self.mlm_head(output)
+        return output
+
+
+class DistCrossLingualMLMTask(nn.Module):
+    """Two shards CrossLingualMLMTask"""
+
+    def __init__(self, split_size, workers, *args, **kwargs):
+        super(DistCrossLingualMLMTask, self).__init__()
+
+        self.split_size = split_size
+
+        # Put the first part of the ResNet50 on workers[0]
+        self.p1_rref = rpc.remote(
+            workers[0],
+            CrossLingualMLMTaskShard1,
+            args=(get_cuda_if_available(0),) + args,
+            kwargs=kwargs
+        )
+
+        # Put the second part of the ResNet50 on workers[1]
+        self.p2_rref = rpc.remote(
+            workers[1],
+            CrossLingualMLMTaskShard2,
+            args=(get_cuda_if_available(1),) + args,
+            kwargs=kwargs
+        )
+
+    def forward(self, xs):
+        # Split the input batch xs into micro-batches, and collect async RPC
+        # futures into a list
+        out_futures = []
+        for x in iter(xs.split(self.split_size, dim=0)):
+            x_rref = RRef(x)
+            y_rref = self.p1_rref.remote().forward(x_rref)
+            z_fut = self.p2_rref.rpc_async().forward(y_rref)
+            out_futures.append(z_fut)
+
+        # collect and cat all output tensors into one tensor.
+        return torch.cat(torch.futures.wait_all(out_futures))
+
+    def parameter_rrefs(self):
+        remote_params = []
+        remote_params.extend(self.p1_rref.remote().parameter_rrefs().to_here())
+        remote_params.extend(self.p2_rref.remote().parameter_rrefs().to_here())
+        return remote_params
