Merge pull request huggingface#10 from nottombrown/multi-gpu

[WIP] Add multi-gpu support

Author: thomwolf

Diff for: README.md

@@ -61,6 +61,3 @@ The single run test accuracy of this PyTorch version is 85.84%, while the author
 The authors implementations uses 8 GPU and can thus accomodate a batch of 64 samples while the present implementation is single GPU and is in consequence limited to 20 instances on a K80 for memory reasons. In our test, increasing the batch size from 8 to 20 samples increased the test accuracy by 2.5 points. A better accuracy may be obtained by using a multi-GPU setting (on the TO-DO list).
 
 The previous SOTA on the ROCStories dataset is 77.6% ("Hidden Coherence Model" of Chaturvedi et al. published in "Story Comprehension for Predicting What Happens Next" EMNLP 2017, which is a very nice paper too!)
-
-### TO-DO list
-- [ ] Add Multi-GPU training logic

Diff for: model_pytorch.py

@@ -11,10 +11,12 @@
 
 
 def gelu(x):
-    return 0.5*x*(1+torch.tanh(math.sqrt(2/math.pi)*(x+0.044715*torch.pow(x, 3))))
+    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+
 
 def swish(x):
-    return x*torch.sigmoid(x)
+    return x * torch.sigmoid(x)
+
 
 ACT_FNS = {
     'relu': nn.ReLU,
@@ -25,6 +27,7 @@ def swish(x):
 
 class LayerNorm(nn.Module):
     "Construct a layernorm module in the OpenAI style (epsilon inside the square root)."
+
     def __init__(self, n_state, e=1e-5):
         super(LayerNorm, self).__init__()
         self.g = nn.Parameter(torch.ones(n_state))
@@ -43,12 +46,12 @@ def __init__(self, nf, rf, nx):
         super(Conv1D, self).__init__()
         self.rf = rf
         self.nf = nf
-        if rf == 1: #faster 1x1 conv
+        if rf == 1:  # faster 1x1 conv
             w = torch.empty(nx, nf)
             nn.init.normal_(w, std=0.02)
             self.w = Parameter(w)
             self.b = Parameter(torch.zeros(nf))
-        else: #was used to train LM
+        else:  # was used to train LM
             raise NotImplementedError
 
     def forward(self, x):
@@ -64,9 +67,9 @@ def forward(self, x):
 class Attention(nn.Module):
     def __init__(self, nx, n_ctx, cfg, scale=False):
         super(Attention, self).__init__()
-        n_state = nx # in Attention: n_state=768 (nx=n_embd)
-        #[switch nx => n_state from Block to Attention to keep identical to TF implem]
-        assert n_state % cfg.n_head==0
+        n_state = nx  # in Attention: n_state=768 (nx=n_embd)
+        # [switch nx => n_state from Block to Attention to keep identical to TF implem]
+        assert n_state % cfg.n_head == 0
         self.register_buffer('b', torch.tril(torch.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))
         self.n_head = cfg.n_head
         self.split_size = n_state
@@ -80,19 +83,19 @@ def _attn(self, q, k, v):
         w = torch.matmul(q, k)
         if self.scale:
             w = w / math.sqrt(v.size(-1))
-        w = w * self.b + -1e9*(1-self.b) # TF implem method: mask_attn_weights
+        w = w * self.b + -1e9 * (1 - self.b)  # TF implem method: mask_attn_weights
         w = nn.Softmax(dim=-1)(w)
         w = self.attn_dropout(w)
         return torch.matmul(w, v)
 
     def merge_heads(self, x):
         x = x.permute(0, 2, 1, 3).contiguous()
         new_x_shape = x.size()[:-2] + (x.size(-2) * x.size(-1),)
-        return x.view(*new_x_shape) # in Tensorflow implem: fct merge_states
+        return x.view(*new_x_shape)  # in Tensorflow implem: fct merge_states
 
     def split_heads(self, x, k=False):
-        new_x_shape = x.size()[:-1] + (self.n_head, x.size(-1)//self.n_head)
-        x = x.view(*new_x_shape) # in Tensorflow implem: fct split_states
+        new_x_shape = x.size()[:-1] + (self.n_head, x.size(-1) // self.n_head)
+        x = x.view(*new_x_shape)  # in Tensorflow implem: fct split_states
         if k:
             return x.permute(0, 2, 3, 1)
         else:
@@ -112,7 +115,7 @@ def forward(self, x):
 
 
 class MLP(nn.Module):
-    def __init__(self, n_state, cfg): # in MLP: n_state=3072 (4 * n_embd)
+    def __init__(self, n_state, cfg):  # in MLP: n_state=3072 (4 * n_embd)
         super(MLP, self).__init__()
         nx = cfg.n_embd
         self.c_fc = Conv1D(n_state, 1, nx)
@@ -132,19 +135,20 @@ def __init__(self, n_ctx, cfg, scale=False):
         nx = cfg.n_embd
         self.attn = Attention(nx, n_ctx, cfg, scale)
         self.ln_1 = LayerNorm(nx)
-        self.mlp = MLP(4*nx, cfg)
+        self.mlp = MLP(4 * nx, cfg)
         self.ln_2 = LayerNorm(nx)
 
     def forward(self, x):
         a = self.attn(x)
-        n = self.ln_1(x+a)
+        n = self.ln_1(x + a)
         m = self.mlp(n)
-        h = self.ln_2(n+m)
+        h = self.ln_2(n + m)
         return h
 
 
 class Model(nn.Module):
     """ Transformer model """
+
     def __init__(self, cfg, vocab=40990, n_ctx=512):
         super(Model, self).__init__()
         self.vocab = vocab
@@ -153,8 +157,8 @@ def __init__(self, cfg, vocab=40990, n_ctx=512):
         block = Block(n_ctx, cfg, scale=True)
         self.h = nn.ModuleList([copy.deepcopy(block) for _ in range(cfg.n_layer)])
         self.decoder = nn.Linear(cfg.n_embd, vocab, bias=False)
-        self.decoder.weight = self.embed.weight # Tied weights
-        self.clf_dropout = nn.Dropout2d(cfg.clf_pdrop) # To reproduce the noise_shape parameter of TF implementation
+        self.decoder.weight = self.embed.weight  # Tied weights
+        self.clf_dropout = nn.Dropout2d(cfg.clf_pdrop)  # To reproduce the noise_shape parameter of TF implementation
 
         nn.init.normal_(self.embed.weight, std=0.02)
 
@@ -169,25 +173,27 @@ def forward(self, x):
 
 class LMHead(nn.Module):
     """ Language Model Head for the transformer """
+
     def __init__(self, model, cfg):
         super(LMHead, self).__init__()
         self.n_embd = cfg.n_embd
-        self.decoder = lambda x: F.linear(x, model.embed.weight) # Tied weights
+        self.decoder = lambda x: F.linear(x, model.embed.weight)  # Tied weights
 
     def forward(self, h):
         # Truncated Language modeling logits (we remove the last token)
-        h_trunc = h[:, :-1].contiguous().view(-1, self.n_embd) # Shape: 252, 768
+        h_trunc = h[:, :-1].contiguous().view(-1, self.n_embd)  # Shape: 252, 768
         lm_logits = self.decoder(h_trunc)
         return lm_logits
 
 
 class ClfHead(nn.Module):
     """ Classifier Head for the transformer """
+
     def __init__(self, clf_token, cfg):
         super(ClfHead, self).__init__()
         self.n_embd = cfg.n_embd
         self.clf_token = clf_token
-        self.dropout = nn.Dropout2d(cfg.clf_pdrop) # To reproduce the noise_shape parameter of TF implementation
+        self.dropout = nn.Dropout2d(cfg.clf_pdrop)  # To reproduce the noise_shape parameter of TF implementation
         self.linear = nn.Linear(cfg.n_embd, 1)
         nn.init.normal_(self.linear.weight, std=0.02)
         nn.init.normal_(self.linear.bias, 0)
@@ -196,17 +202,30 @@ def forward(self, h, x):
         # Classification logits
         clf_h = h.view(-1, self.n_embd)
         flat = x[:, :, :, 0].contiguous().view(-1)
-        #pool_idx = torch.eq(x[:, :, 0].contiguous().view(-1), self.clf_token)
-        clf_h = clf_h[flat == self.clf_token, :] #.index_select(0, pool_idx)
+        # pool_idx = torch.eq(x[:, :, 0].contiguous().view(-1), self.clf_token)
+        clf_h = clf_h[flat == self.clf_token, :]  # .index_select(0, pool_idx)
         clf_h = clf_h.view(-1, 2, self.n_embd, 1)
         clf_h = self.dropout(clf_h)
         clf_h = clf_h.view(-1, self.n_embd)
         clf_logits = self.linear(clf_h)
         return clf_logits.view(-1, 2)
 
 
-def load_openai_pretrained_model(model, n_ctx=-1, n_special=-1, n_transfer=12, n_embd=768, path='./model/', path_names='./'):
+class DataParallelWithEmbed(torch.nn.DataParallel):
+    """DataParallel that proxies the embed property to the wrapped module"""
+
+    def __init__(self, model):
+        super(DataParallelWithEmbed, self).__init__(model)
+
+    @property
+    def embed(self):
+        return self.module.embed
+
+
+def load_openai_pretrained_model(model, n_ctx=-1, n_special=-1, n_transfer=12, n_embd=768, path='./model/',
+                                 path_names='./'):
     # Load weights from TF model
+    print("Loading weights...")
     names = json.load(open(path_names + 'parameters_names.json'))
     shapes = json.load(open(path + 'params_shapes.json'))
     offsets = np.cumsum([np.prod(shape) for shape in shapes])
@@ -216,32 +235,40 @@ def load_openai_pretrained_model(model, n_ctx=-1, n_special=-1, n_transfer=12, n
     if n_ctx > 0:
         init_params[0] = init_params[0][:n_ctx]
     if n_special > 0:
-        init_params[0] = np.concatenate([init_params[1],
-                                         (np.random.randn(n_special, n_embd)*0.02).astype(np.float32),
-                                         init_params[0]
-                                        ], 0)
+        init_params[0] = np.concatenate(
+            [init_params[1],
+             (np.random.randn(n_special, n_embd) * 0.02).astype(np.float32),
+             init_params[0]
+             ], 0)
     else:
-        init_params[0] = np.concatenate([init_params[1],
-                                         init_params[0]
-                                        ], 0)
+        init_params[0] = np.concatenate(
+            [init_params[1],
+             init_params[0]
+             ], 0)
     del init_params[1]
     if n_transfer == -1:
         n_transfer = 0
     else:
-        n_transfer = 1+n_transfer*12
+        n_transfer = 1 + n_transfer * 12
     init_params = [arr.squeeze() for arr in init_params]
+
     try:
         assert model.embed.weight.shape == init_params[0].shape
     except AssertionError as e:
         e.args += (model.embed.weight.shape, init_params[0].shape)
         raise
+
     model.embed.weight.data = torch.from_numpy(init_params[0])
+
+    # Load the weights into our torch module
+    module = model.module
+
     for name, ip in zip(names[1:n_transfer], init_params[1:n_transfer]):
-        name = name[6:] # skip "model/"
+        name = name[6:]  # skip "model/"
         assert name[-2:] == ":0"
         name = name[:-2]
         name = name.split('/')
-        pointer = model
+        pointer = module
         for m_name in name:
             if re.fullmatch(r'[A-Za-z]+\d+', m_name):
                 l = re.split(r'(\d+)', m_name)
@@ -258,12 +285,14 @@ def load_openai_pretrained_model(model, n_ctx=-1, n_special=-1, n_transfer=12, n
             raise
         pointer.data = torch.from_numpy(ip)
 
+
 class dotdict(dict):
     """dot.notation access to dictionary attributes"""
     __getattr__ = dict.get
     __setattr__ = dict.__setitem__
     __delattr__ = dict.__delitem__
 
+
 DEFAULT_CONFIG = dotdict({
     'n_embd': 768,
     'n_head': 12,

Diff for: train.py

@@ -10,7 +10,7 @@
 
 from analysis import rocstories as rocstories_analysis
 from datasets import rocstories
-from model_pytorch import Model, LMHead, ClfHead, load_openai_pretrained_model
+from model_pytorch import Model, LMHead, ClfHead, load_openai_pretrained_model, DataParallelWithEmbed
 from opt import OpenAIAdam
 from text_utils import TextEncoder
 from utils import (encode_dataset, iter_data,
@@ -237,6 +237,7 @@ def run_epoch():
     encoder = text_encoder.encoder
     n_vocab = len(text_encoder.encoder)
 
+    print("Encoding dataset...")
     (trX1, trX2, trX3, trY), (vaX1, vaX2, vaX3, vaY), (teX1, teX2, teX3) = encode_dataset(
         rocstories(data_dir, n_valid=args.n_valid), encoder=text_encoder)
     n_y = 2
@@ -266,6 +267,8 @@ def run_epoch():
     n_updates_total = (n_train // n_batch_train) * args.n_iter
 
     model = Model(args, vocab, n_ctx)
+    model = DataParallelWithEmbed(model).cuda()
+
     lm_head = LMHead(model, args)
     clf_head = ClfHead(clf_token, args)