Rebasing tpu branch on a more recent fairseq upstream commit

fairseq/checkpoint_utils.py

@@ -17,6 +17,8 @@
 
 from fairseq.models import FairseqEncoder, FairseqDecoder
 
+import torch_xla.core.xla_model as xm
+
 
 def save_checkpoint(args, trainer, epoch_itr, val_loss):
     from fairseq import distributed_utils, meters
@@ -62,15 +64,17 @@ def is_better(a, b):
         extra_state.update({'best': save_checkpoint.best})
 
     checkpoints = [os.path.join(args.save_dir, fn) for fn, cond in checkpoint_conds.items() if cond]
+
     if len(checkpoints) > 0:
         trainer.save_checkpoint(checkpoints[0], extra_state)
         for cp in checkpoints[1:]:
             try:
                 from fairseq.fb_pathmgr import fb_pathmgr
-                fb_pathmgr.copy(checkpoints[0], cp, True)
+                if getattr(args, 'use_gpu', True) or xm.is_master_ordinal():
+                    fb_pathmgr.copy(checkpoints[0], cp, True)
             except (ModuleNotFoundError, ImportError):
-                shutil.copyfile(checkpoints[0], cp)
-
+                if getattr(args, 'use_gpu', True) or xm.is_master_ordinal():
+                    shutil.copyfile(checkpoints[0], cp)
         write_timer.stop()
         print('| saved checkpoint {} (epoch {} @ {} updates) (writing took {} seconds)'.format(
             checkpoints[0], epoch, updates, write_timer.sum))
@@ -97,7 +101,7 @@ def is_better(a, b):
 def load_checkpoint(args, trainer, data_selector=None):
     """Load a checkpoint and restore the training iterator."""
     # only one worker should attempt to create the required dir
-    if args.distributed_rank == 0:
+    if args.distributed_rank == 0 or xm.is_master_ordinal():
         os.makedirs(args.save_dir, exist_ok=True)
 
     if args.restore_file == 'checkpoint_last.pt':
@@ -210,7 +214,8 @@ def checkpoint_paths(path, pattern=r'checkpoint(\d+)\.pt'):
 def torch_persistent_save(*args, **kwargs):
     for i in range(3):
         try:
-            return torch.save(*args, **kwargs)
+            save_func = xm.save if kwargs.pop('xla', False) else torch.save
+            return save_func(*args, **kwargs)
         except Exception:
             if i == 2:
                 logging.error(traceback.format_exc())
@@ -256,14 +261,17 @@ def save_state(
         state_dict['criterion'] = criterion.state_dict()
     if not args.no_save_optimizer_state:
         state_dict['last_optimizer_state'] = convert_state_dict_type(optimizer.state_dict())
-
     try:
         from fairseq.fb_pathmgr import fb_pathmgr
         with fb_pathmgr.open(filename, "wb") as f:
-            torch_persistent_save(state_dict, f)
+            torch_persistent_save(
+                state_dict, f, xla=not getattr(args, 'use_gpu', True)
+            )
     except (ModuleNotFoundError, ImportError):
         # if path manager not found, continue with local file.
-        torch_persistent_save(state_dict, filename)
+        torch_persistent_save(
+            state_dict, filename, xla=not getattr(args, 'use_gpu', True)
+        )
 
 
 def _upgrade_state_dict(state):

fairseq/criterions/label_smoothed_cross_entropy.py

@@ -17,8 +17,8 @@ def label_smoothed_nll_loss(lprobs, target, epsilon, ignore_index=None, reduce=T
     smooth_loss = -lprobs.sum(dim=-1, keepdim=True)
     if ignore_index is not None:
         non_pad_mask = target.ne(ignore_index)
-        nll_loss = nll_loss[non_pad_mask]
-        smooth_loss = smooth_loss[non_pad_mask]
+        nll_loss.masked_fill_(~non_pad_mask, 0.0)
+        smooth_loss.masked_fill_(~non_pad_mask, 0.0)
     else:
         nll_loss = nll_loss.squeeze(-1)
         smooth_loss = smooth_loss.squeeze(-1)
@@ -57,8 +57,8 @@ def forward(self, model, sample, reduce=True):
         loss, nll_loss = self.compute_loss(model, net_output, sample, reduce=reduce)
         sample_size = sample['target'].size(0) if self.args.sentence_avg else sample['ntokens']
         logging_output = {
-            'loss': utils.item(loss.data) if reduce else loss.data,
-            'nll_loss': utils.item(nll_loss.data) if reduce else nll_loss.data,
+            'loss': loss.data,
+            'nll_loss': nll_loss.data,
             'ntokens': sample['ntokens'],
             'nsentences': sample['target'].size(0),
             'sample_size': sample_size,

fairseq/data/data_utils.py

@@ -26,9 +26,25 @@ def infer_language_pair(path):
     return src, dst
 
 
-def collate_tokens(values, pad_idx, eos_idx=None, left_pad=False, move_eos_to_beginning=False):
+def get_pad_size(values, input_shapes):
+    if input_shapes is None:
+        return max(v.size(0) for v in values)
+    for batch_size, padlen in input_shapes:
+        if len(values) == batch_size:
+            return padlen
+    else:
+        raise IndexError(
+            'Encountered values with invalid length {}, input shapes were {}'
+            .format(len(values), input_shapes)
+        )
+
+
+def collate_tokens(
+    values, pad_idx, eos_idx=None, left_pad=False,
+    move_eos_to_beginning=False, input_shapes=None,
+):
     """Convert a list of 1d tensors into a padded 2d tensor."""
-    size = max(v.size(0) for v in values)
+    size = get_pad_size(values, input_shapes)
     res = values[0].new(len(values), size).fill_(pad_idx)
 
     def copy_tensor(src, dst):
@@ -227,10 +243,25 @@ def batch_by_size(
 
     if isinstance(indices, types.GeneratorType):
         indices = np.fromiter(indices, dtype=np.int64, count=-1)
-
     return batch_by_size_fast(indices, num_tokens_fn, max_tokens, max_sentences, bsz_mult)
 
 
+def batch_by_size_tpu(
+    indices, num_tokens_fn, input_shapes
+):
+    batches = [[] for _ in input_shapes]
+    for idx in indices:
+        sample_len = num_tokens_fn(idx)
+        for j, (batch_size, padlen) in enumerate(input_shapes):
+            if padlen < sample_len:
+                continue
+            batches[j].append(idx)
+            if len(batches[j]) == batch_size:
+                yield batches[j]
+                batches[j] = []
+            break
+
+
 def process_bpe_symbol(sentence: str, bpe_symbol: str):
     if bpe_symbol == 'sentencepiece':
         sentence = sentence.replace(' ', '').replace('\u2581', ' ').strip()

fairseq/data/language_pair_dataset.py

@@ -11,15 +11,15 @@
 
 def collate(
     samples, pad_idx, eos_idx, left_pad_source=True, left_pad_target=False,
-    input_feeding=True,
+    input_feeding=True, input_shapes=None,
 ):
     if len(samples) == 0:
         return {}
 
     def merge(key, left_pad, move_eos_to_beginning=False):
         return data_utils.collate_tokens(
-            [s[key] for s in samples],
-            pad_idx, eos_idx, left_pad, move_eos_to_beginning,
+            [s[key] for s in samples], pad_idx,
+            eos_idx,left_pad, move_eos_to_beginning, input_shapes,
         )
 
     def check_alignment(alignment, src_len, tgt_len):
@@ -154,7 +154,8 @@ def __init__(
         shuffle=True, input_feeding=True,
         remove_eos_from_source=False, append_eos_to_target=False,
         align_dataset=None,
-        append_bos=False
+        append_bos=False,
+        input_shapes=None,
     ):
         if tgt_dict is not None:
             assert src_dict.pad() == tgt_dict.pad()
@@ -178,6 +179,7 @@ def __init__(
         if self.align_dataset is not None:
             assert self.tgt_sizes is not None, "Both source and target needed when alignments are provided"
         self.append_bos = append_bos
+        self.input_shapes = input_shapes
 
     def __getitem__(self, index):
         tgt_item = self.tgt[index] if self.tgt is not None else None
@@ -249,7 +251,7 @@ def collater(self, samples):
         return collate(
             samples, pad_idx=self.src_dict.pad(), eos_idx=self.src_dict.eos(),
             left_pad_source=self.left_pad_source, left_pad_target=self.left_pad_target,
-            input_feeding=self.input_feeding,
+            input_feeding=self.input_feeding, input_shapes=self.input_shapes,
         )
 
     def num_tokens(self, index):

fairseq/models/transformer.py

@@ -355,8 +355,8 @@ def forward(self, src_tokens, src_lengths, cls_input=None, return_all_hiddens=Fa
 
         # compute padding mask
         encoder_padding_mask = src_tokens.eq(self.padding_idx)
-        if not encoder_padding_mask.any():
-            encoder_padding_mask = None
+        #if not encoder_padding_mask.any():
+        #    encoder_padding_mask = None
 
         encoder_states = [] if return_all_hiddens else None
 
@@ -596,8 +596,8 @@ def extract_features(
         x = x.transpose(0, 1)
 
         self_attn_padding_mask = prev_output_tokens.eq(self.padding_idx)
-        if not self_attn_padding_mask.any() and not self.cross_self_attention:
-            self_attn_padding_mask = None
+        # if not self_attn_padding_mask.any() and not self.cross_self_attention:
+        #     self_attn_padding_mask = None
 
         # decoder layers
         attn = None

fairseq/modules/multihead_attention.py

@@ -39,9 +39,17 @@ def __init__(self, embed_dim, num_heads, kdim=None, vdim=None, dropout=0., bias=
         assert not self.self_attention or self.qkv_same_dim, 'Self-attention requires query, key and ' \
                                                              'value to be of the same size'
 
-        self.k_proj = nn.Linear(self.kdim, embed_dim, bias=bias)
-        self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias)
-        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        if self.qkv_same_dim:
+            self.in_proj_weight = Parameter(torch.Tensor(3 * embed_dim, embed_dim))
+        else:
+            self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))
+            self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))
+            self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
+
+        if bias:
+            self.in_proj_bias = Parameter(torch.Tensor(3 * embed_dim))
+        else:
+            self.register_parameter('in_proj_bias', None)
 
         self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
 
@@ -57,11 +65,12 @@ def __init__(self, embed_dim, num_heads, kdim=None, vdim=None, dropout=0., bias=
 
         self.onnx_trace = False
 
+        # XXX: (taylanbil) try F.multi...
         self.enable_torch_version = False
-        if hasattr(F, "multi_head_attention_forward"):
-            self.enable_torch_version = True
-        else:
-            self.enable_torch_version = False
+        # if hasattr(F, "multi_head_attention_forward"):
+        #     self.enable_torch_version = True
+        # else:
+        #     self.enable_torch_version = False
 
     def prepare_for_onnx_export_(self):
         self.onnx_trace = True
@@ -70,15 +79,15 @@ def reset_parameters(self):
         if self.qkv_same_dim:
             # Empirically observed the convergence to be much better with
             # the scaled initialization
-            nn.init.xavier_uniform_(self.k_proj.weight, gain=1/math.sqrt(2))
-            nn.init.xavier_uniform_(self.v_proj.weight, gain=1/math.sqrt(2))
-            nn.init.xavier_uniform_(self.q_proj.weight, gain=1/math.sqrt(2))
+            nn.init.xavier_uniform_(self.in_proj_weight, gain=1/math.sqrt(2))
         else:
-            nn.init.xavier_uniform_(self.k_proj.weight)
-            nn.init.xavier_uniform_(self.v_proj.weight)
-            nn.init.xavier_uniform_(self.q_proj.weight)
+            nn.init.xavier_uniform_(self.k_proj_weight)
+            nn.init.xavier_uniform_(self.v_proj_weight)
+            nn.init.xavier_uniform_(self.q_proj_weight)
 
         nn.init.xavier_uniform_(self.out_proj.weight)
+        if self.in_proj_bias is not None:
+            nn.init.constant_(self.in_proj_bias, 0.)
         nn.init.constant_(self.out_proj.bias, 0.)
         if self.bias_k is not None:
             nn.init.xavier_normal_(self.bias_k)
@@ -146,23 +155,19 @@ def forward(
             saved_state = None
 
         if self.self_attention:
-            q = self.q_proj(query)
-            k = self.k_proj(query)
-            v = self.v_proj(query)
+            q, k, v = self.in_proj_qkv(query)
         elif self.encoder_decoder_attention:
             # encoder-decoder attention
-            q = self.q_proj(query)
+            q = self.in_proj_q(query)
             if key is None:
                 assert value is None
                 k = v = None
             else:
-                k = self.k_proj(key)
-                v = self.v_proj(key)
+                k = self.in_proj_k(key)
+                v = self.in_proj_v(key)
 
         else:
-            q = self.q_proj(query)
-            k = self.k_proj(key)
-            v = self.v_proj(value)
+            raise
         q *= self.scaling
 
         if self.bias_k is not None:
@@ -242,10 +247,9 @@ def forward(
         if key_padding_mask is not None:
             # don't attend to padding symbols
             attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
-            attn_weights = attn_weights.masked_fill(
-                key_padding_mask.unsqueeze(1).unsqueeze(2),
-                float('-inf'),
-            )
+            attn_weights = attn_weights.transpose(0, 2)
+            attn_weights.masked_fill_(key_padding_mask, float('-inf'))
+            attn_weights = attn_weights.transpose(0, 2)
             attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
 
         if before_softmax:
@@ -330,3 +334,43 @@ def upgrade_state_dict_named(self, state_dict, name):
 
         for key, value in items_to_add.items():
             state_dict[key] = value
+
+    def in_proj_qkv(self, query):
+        return self._in_proj(query).chunk(3, dim=-1)
+
+    def in_proj_q(self, query):
+        if self.qkv_same_dim:
+            return self._in_proj(query, end=self.embed_dim)
+        else:
+            bias = self.in_proj_bias
+            if bias is not None:
+                bias = bias[:self.embed_dim]
+            return F.linear(query, self.q_proj_weight, bias)
+
+    def in_proj_k(self, key):
+        if self.qkv_same_dim:
+            return self._in_proj(key, start=self.embed_dim, end=2 * self.embed_dim)
+        else:
+            weight = self.k_proj_weight
+            bias = self.in_proj_bias
+            if bias is not None:
+                bias = bias[self.embed_dim:2 * self.embed_dim]
+            return F.linear(key, weight, bias)
+
+    def in_proj_v(self, value):
+        if self.qkv_same_dim:
+            return self._in_proj(value, start=2 * self.embed_dim)
+        else:
+            weight = self.v_proj_weight
+            bias = self.in_proj_bias
+            if bias is not None:
+                bias = bias[2 * self.embed_dim:]
+            return F.linear(value, weight, bias)
+
+    def _in_proj(self, input, start=0, end=None):
+        weight = self.in_proj_weight
+        bias = self.in_proj_bias
+        weight = weight[start:end, :]
+        if bias is not None:
+            bias = bias[start:end]
+        return F.linear(input, weight, bias)

fairseq/tasks/fairseq_task.py

@@ -146,10 +146,17 @@ def get_batch_iterator(
             )
 
         # create mini-batches with given size constraints
-        batch_sampler = data_utils.batch_by_size(
-            indices, dataset.num_tokens, max_tokens=max_tokens, max_sentences=max_sentences,
-            required_batch_size_multiple=required_batch_size_multiple,
-        )
+        if getattr(self.args, 'use_gpu', True):
+            batch_sampler = data_utils.batch_by_size(
+                indices, dataset.num_tokens, max_tokens=max_tokens,
+                max_sentences=max_sentences,
+                required_batch_size_multiple=required_batch_size_multiple,
+            )
+        else:
+            batch_sampler = data_utils.batch_by_size_tpu(
+                indices, dataset.num_tokens,
+                getattr(self.args, 'input_shapes', None)
+            )
 
         # return a reusable, sharded iterator
         epoch_iter = iterators.EpochBatchIterator(