Adding multi-gpu backward ops using pytorch-encoding

Author: codertimo

bert_pytorch/trainer/pretrain.py

@@ -3,6 +3,8 @@
 from torch.optim import Adam
 from torch.utils.data import DataLoader
 
+from encoding.parallel import DataParallelModel, DataParallelCriterion
+
 from ..model import BERTLM, BERT
 
 import tqdm
@@ -47,7 +49,7 @@ def __init__(self, bert: BERT, vocab_size: int,
         # Distributed GPU training if CUDA can detect more than 1 GPU
         if torch.cuda.device_count() > 1:
             print("Using %d GPUS for BERT" % torch.cuda.device_count())
-            self.model = nn.DataParallel(self.model)
+            self.model = DataParallelModel(self.model)
 
         # Setting the train and test data loader
         self.train_data = train_dataloader
@@ -57,7 +59,7 @@ def __init__(self, bert: BERT, vocab_size: int,
         self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay)
 
         # Using Negative Log Likelihood Loss function for predicting the masked_token
-        self.criterion = nn.NLLLoss(ignore_index=0)
+        self.criterion = DataParallelCriterion(nn.NLLLoss(ignore_index=0))
 
         self.log_freq = log_freq
 

requirements.txt

@@ -1,3 +1,4 @@
 tqdm
 numpy
 torch>=0.4.0
+torch-encoding

setup.py

@@ -3,7 +3,7 @@
 import os
 import sys
 
-__version__ = "0.0.1a3"
+__version__ = "0.0.1a4"
 
 with open("requirements.txt") as f:
     require_packages = [line[:-1] for line in f]