Fix bugs in pipeline tutorial with respect to batch size. (pytorch#1459)

* Fix bugs in pipeline tutorial with respect to batch size.

Summary: As described in
pytorch/pytorch#55036, certain modules were
not handling batch size correctly.

* Rebase

Co-authored-by: pritam <[email protected]>
Co-authored-by: Brian Johnson <[email protected]>

Author: 3 people

intermediate_source/pipeline_tutorial.py

@@ -56,7 +56,6 @@
 class Encoder(nn.Module):
     def __init__(self, ntoken, ninp, dropout=0.5):
         super(Encoder, self).__init__()
-        self.src_mask = None
         self.pos_encoder = PositionalEncoding(ninp, dropout)
         self.encoder = nn.Embedding(ntoken, ninp)
         self.ninp = ninp
@@ -66,17 +65,9 @@ def init_weights(self):
         initrange = 0.1
         self.encoder.weight.data.uniform_(-initrange, initrange)
 
-    def _generate_square_subsequent_mask(self, sz):
-        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
-        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
-        return mask
-
     def forward(self, src):
-        if self.src_mask is None or self.src_mask.size(0) != src.size(0):
-            device = src.device
-            mask = self._generate_square_subsequent_mask(src.size(0)).to(device)
-            self.src_mask = mask
-
+        # Need (S, N) format for encoder.
+        src = src.t()
         src = self.encoder(src) * math.sqrt(self.ninp)
         return self.pos_encoder(src)
 
@@ -92,7 +83,8 @@ def init_weights(self):
         self.decoder.weight.data.uniform_(-initrange, initrange)
 
     def forward(self, inp):
-        return self.decoder(inp)
+        # Need batch dimension first for output of pipeline.
+        return self.decoder(inp).permute(1, 0, 2)
 
 
 ######################################################################
@@ -221,7 +213,8 @@ def get_batch(source, i):
     seq_len = min(bptt, len(source) - 1 - i)
     data = source[i:i+seq_len]
     target = source[i+1:i+1+seq_len].view(-1)
-    return data, target
+    # Need batch dimension first for pipeline parallelism.
+    return data.t(), target
 
 ######################################################################
 # Model scale and Pipe initialization
@@ -297,7 +290,8 @@ def get_batch(source, i):
 from torch.distributed.pipeline.sync import Pipe
 
 # Build the pipeline.
-model = Pipe(torch.nn.Sequential(*module_list), chunks = 8)
+chunks = 8
+model = Pipe(torch.nn.Sequential(*module_list), chunks = chunks)
 
 
 def get_total_params(module: torch.nn.Module):