Update Fairseq LSTM to jitable version (#2016)

Summary:
Pull Request resolved: #2016

It is to update Fairseq LSTM to jitable version

Reviewed By: cndn

Differential Revision: D20937370

fbshipit-source-id: 26f677fcb58bbeaa507d303e9a81060ff78f0502

Author: Xianfeng Rui

fairseq/models/fairseq_decoder.py

@@ -59,7 +59,7 @@ def get_normalized_probs(
         self,
         net_output: Tuple[Tensor, Optional[Dict[str, List[Optional[Tensor]]]]],
         log_probs: bool,
-        sample: Optional[Dict[str, Tensor]],
+        sample: Optional[Dict[str, Tensor]] = None,
     ):
         """Get normalized probabilities (or log probs) from a net's output."""
 

fairseq/models/lstm.py

@@ -16,10 +16,14 @@
     register_model_architecture,
 )
 from fairseq.modules import AdaptiveSoftmax
+from torch import Tensor
+from typing import Dict, List, Optional, Tuple
+
 
 DEFAULT_MAX_SOURCE_POSITIONS = 1e5
 DEFAULT_MAX_TARGET_POSITIONS = 1e5
 
+
 @register_model('lstm')
 class LSTMModel(FairseqEncoderDecoderModel):
     def __init__(self, encoder, decoder):
@@ -176,6 +180,19 @@ def load_pretrained_embedding_from_file(embed_path, dictionary, embed_dim):
         )
         return cls(encoder, decoder)
 
+    def forward(
+        self,
+        src_tokens,
+        src_lengths,
+        prev_output_tokens,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+    ):
+        encoder_out = self.encoder(src_tokens, src_lengths=src_lengths)
+        decoder_out = self.decoder(
+            prev_output_tokens, encoder_out=encoder_out, incremental_state=incremental_state
+        )
+        return decoder_out
+
 
 class LSTMEncoder(FairseqEncoder):
     """LSTM encoder."""
@@ -213,13 +230,17 @@ def __init__(
         if bidirectional:
             self.output_units *= 2
 
-    def forward(self, src_tokens, src_lengths):
+    def combine_bidir(self, outs, bsz: int):
+        out = outs.view(self.num_layers, 2, bsz, -1).transpose(1, 2).contiguous()
+        return out.view(self.num_layers, bsz, -1)
+
+    def forward(self, src_tokens, src_lengths: Tensor):
         if self.left_pad:
             # nn.utils.rnn.pack_padded_sequence requires right-padding;
             # convert left-padding to right-padding
             src_tokens = utils.convert_padding_direction(
                 src_tokens,
-                self.padding_idx,
+                torch.zeros_like(src_tokens).fill_(self.padding_idx),
                 left_to_right=True,
             )
 
@@ -233,7 +254,7 @@ def forward(self, src_tokens, src_lengths):
         x = x.transpose(0, 1)
 
         # pack embedded source tokens into a PackedSequence
-        packed_x = nn.utils.rnn.pack_padded_sequence(x, src_lengths.data.tolist())
+        packed_x = nn.utils.rnn.pack_padded_sequence(x, src_lengths.data)
 
         # apply LSTM
         if self.bidirectional:
@@ -245,34 +266,32 @@ def forward(self, src_tokens, src_lengths):
         packed_outs, (final_hiddens, final_cells) = self.lstm(packed_x, (h0, c0))
 
         # unpack outputs and apply dropout
-        x, _ = nn.utils.rnn.pad_packed_sequence(packed_outs, padding_value=self.padding_idx)
+        x, _ = nn.utils.rnn.pad_packed_sequence(packed_outs, padding_value=self.padding_idx*1.0)
         x = F.dropout(x, p=self.dropout_out, training=self.training)
         assert list(x.size()) == [seqlen, bsz, self.output_units]
 
         if self.bidirectional:
 
-            def combine_bidir(outs):
-                out = outs.view(self.num_layers, 2, bsz, -1).transpose(1, 2).contiguous()
-                return out.view(self.num_layers, bsz, -1)
-
-            final_hiddens = combine_bidir(final_hiddens)
-            final_cells = combine_bidir(final_cells)
+            final_hiddens = self.combine_bidir(final_hiddens, bsz)
+            final_cells = self.combine_bidir(final_cells, bsz)
 
         encoder_padding_mask = src_tokens.eq(self.padding_idx).t()
 
         return {
             'encoder_out': (x, final_hiddens, final_cells),
-            'encoder_padding_mask': encoder_padding_mask if encoder_padding_mask.any() else None
+            'encoder_padding_mask': (encoder_padding_mask, torch.empty(0), torch.empty(0)),
         }
 
-    def reorder_encoder_out(self, encoder_out, new_order):
-        encoder_out['encoder_out'] = tuple(
-            eo.index_select(1, new_order)
-            for eo in encoder_out['encoder_out']
-        )
-        if encoder_out['encoder_padding_mask'] is not None:
-            encoder_out['encoder_padding_mask'] = \
-                encoder_out['encoder_padding_mask'].index_select(1, new_order)
+    def reorder_encoder_out(self, encoder_out: Dict[str, Tuple[Tensor, Tensor, Tensor]], new_order):
+        encoder_out['encoder_out'] = (
+                encoder_out['encoder_out'][0].index_select(1, new_order),
+                encoder_out['encoder_out'][1].index_select(1, new_order),
+                encoder_out['encoder_out'][2].index_select(1, new_order))
+        if encoder_out['encoder_padding_mask'][0] is not None:
+            encoder_out['encoder_padding_mask'] = (
+                encoder_out['encoder_padding_mask'][0].index_select(1, new_order),
+                encoder_out['encoder_padding_mask'][1],
+                encoder_out['encoder_padding_mask'][2])
         return encoder_out
 
     def max_positions(self):
@@ -329,6 +348,7 @@ def __init__(
         self.share_input_output_embed = share_input_output_embed
         self.need_attn = True
         self.max_target_positions = max_target_positions
+        self.num_layers = num_layers
 
         self.adaptive_softmax = None
         num_embeddings = len(dictionary)
@@ -369,36 +389,51 @@ def __init__(
         elif not self.share_input_output_embed:
             self.fc_out = Linear(out_embed_dim, num_embeddings, dropout=dropout_out)
 
-    def forward(self, prev_output_tokens, encoder_out=None, incremental_state=None, **kwargs):
+    def get_cached_state(self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]]):
+
+        cached_state = self.get_incremental_state(incremental_state, 'cached_state')
+        assert cached_state is not None
+        prev_hiddens_ = cached_state["prev_hiddens"]
+        assert prev_hiddens_ is not None
+        prev_cells_ = cached_state["prev_cells"]
+        assert prev_cells_ is not None
+        prev_hiddens = [prev_hiddens_[i] for i in range(self.num_layers)]
+        prev_cells = [prev_cells_[j] for j in range(self.num_layers)]
+        input_feed = cached_state["input_feed"]
+        assert input_feed is not None
+        return prev_hiddens, prev_cells, input_feed
+
+    def forward(self, prev_output_tokens, encoder_out: Dict[str, Tuple[Tensor, Tensor, Tensor]],
+                incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]]):
         x, attn_scores = self.extract_features(
             prev_output_tokens, encoder_out, incremental_state
         )
         return self.output_layer(x), attn_scores
 
     def extract_features(
-        self, prev_output_tokens, encoder_out, incremental_state=None
+            self, prev_output_tokens, encoder_out: Dict[str, Tuple[Tensor, Tensor, Tensor]],
+            incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None
     ):
         """
         Similar to *forward* but only return features.
         """
         if encoder_out is not None:
-            encoder_padding_mask = encoder_out['encoder_padding_mask']
+            encoder_padding_mask, _, _ = encoder_out['encoder_padding_mask']
             encoder_out = encoder_out['encoder_out']
         else:
             encoder_padding_mask = None
             encoder_out = None
 
-        if incremental_state is not None:
+        if incremental_state is not None and len(incremental_state) > 0:
             prev_output_tokens = prev_output_tokens[:, -1:]
+
         bsz, seqlen = prev_output_tokens.size()
 
         # get outputs from encoder
+        srclen = 0
         if encoder_out is not None:
             encoder_outs, encoder_hiddens, encoder_cells = encoder_out[:3]
             srclen = encoder_outs.size(0)
-        else:
-            srclen = None
-
         # embed tokens
         x = self.embed_tokens(prev_output_tokens)
         x = F.dropout(x, p=self.dropout_in, training=self.training)
@@ -407,27 +442,25 @@ def extract_features(
         x = x.transpose(0, 1)
 
         # initialize previous states (or get from cache during incremental generation)
-        cached_state = utils.get_incremental_state(self, incremental_state, 'cached_state')
-        if cached_state is not None:
-            prev_hiddens, prev_cells, input_feed = cached_state
+
+        if incremental_state is not None and len(incremental_state) > 0:
+            prev_hiddens, prev_cells, input_feed = self.get_cached_state(incremental_state)
         elif encoder_out is not None:
             # setup recurrent cells
-            num_layers = len(self.layers)
-            prev_hiddens = [encoder_hiddens[i] for i in range(num_layers)]
-            prev_cells = [encoder_cells[i] for i in range(num_layers)]
+            prev_hiddens = [encoder_hiddens[i] for i in range(self.num_layers)]
+            prev_cells = [encoder_cells[i] for i in range(self.num_layers)]
             if self.encoder_hidden_proj is not None:
-                prev_hiddens = [self.encoder_hidden_proj(x) for x in prev_hiddens]
-                prev_cells = [self.encoder_cell_proj(x) for x in prev_cells]
+                prev_hiddens = [self.encoder_hidden_proj(y) for y in prev_hiddens]
+                prev_cells = [self.encoder_cell_proj(y) for y in prev_cells]
             input_feed = x.new_zeros(bsz, self.hidden_size)
         else:
             # setup zero cells, since there is no encoder
-            num_layers = len(self.layers)
             zero_state = x.new_zeros(bsz, self.hidden_size)
-            prev_hiddens = [zero_state for i in range(num_layers)]
-            prev_cells = [zero_state for i in range(num_layers)]
+            prev_hiddens = [zero_state for i in range(self.num_layers)]
+            prev_cells = [zero_state for i in range(self.num_layers)]
             input_feed = None
 
-        assert srclen is not None or self.attention is None, \
+        assert srclen > 0 or self.attention is None, \
             "attention is not supported if there are no encoder outputs"
         attn_scores = x.new_zeros(srclen, seqlen, bsz) if self.attention is not None else None
         outs = []
@@ -451,6 +484,7 @@ def extract_features(
 
             # apply attention using the last layer's hidden state
             if self.attention is not None:
+                assert attn_scores is not None
                 out, attn_scores[:, j, :] = self.attention(hidden, encoder_outs, encoder_padding_mask)
             else:
                 out = hidden
@@ -463,11 +497,14 @@ def extract_features(
             # save final output
             outs.append(out)
 
-        # cache previous states (no-op except during incremental generation)
-        utils.set_incremental_state(
-            self, incremental_state, 'cached_state',
-            (prev_hiddens, prev_cells, input_feed),
-        )
+        # Stack all the necessary tensors together and store
+        prev_hiddens_tensor = torch.stack(prev_hiddens)
+        prev_cells_tensor = torch.stack(prev_cells)
+        cache_state = torch.jit.annotate(
+            Dict[str, Optional[Tensor]],
+            {"prev_hiddens": prev_hiddens_tensor, "prev_cells": prev_cells_tensor, "input_feed": input_feed})
+        self.set_incremental_state(
+            incremental_state, 'cached_state', cache_state)
 
         # collect outputs across time steps
         x = torch.cat(outs, dim=0).view(seqlen, bsz, self.hidden_size)
@@ -478,9 +515,9 @@ def extract_features(
         if hasattr(self, 'additional_fc') and self.adaptive_softmax is None:
             x = self.additional_fc(x)
             x = F.dropout(x, p=self.dropout_out, training=self.training)
-
         # srclen x tgtlen x bsz -> bsz x tgtlen x srclen
         if not self.training and self.need_attn and self.attention is not None:
+            assert attn_scores is not None
             attn_scores = attn_scores.transpose(0, 2)
         else:
             attn_scores = None
@@ -495,22 +532,22 @@ def output_layer(self, x):
                 x = self.fc_out(x)
         return x
 
-    def reorder_incremental_state(self, incremental_state, new_order):
-        super().reorder_incremental_state(incremental_state, new_order)
-        cached_state = utils.get_incremental_state(self, incremental_state, 'cached_state')
-        if cached_state is None:
-            return
-
-        def reorder_state(state):
-            if isinstance(state, list):
-                return [reorder_state(state_i) for state_i in state]
-            elif state is not None:
-                return state.index_select(0, new_order)
-            else:
-                return None
+    def reorder_state(self, state: List[Tensor], new_order):
+        return [state_i.index_select(0, new_order) for state_i in state]
 
-        new_state = tuple(map(reorder_state, cached_state))
-        utils.set_incremental_state(self, incremental_state, 'cached_state', new_state)
+    def reorder_incremental_state(self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]], new_order):
+        if incremental_state is None or len(incremental_state) == 0:
+            return
+        prev_hiddens, prev_cells, input_feed = self.get_cached_state(incremental_state)
+        cached_state = (prev_hiddens, prev_cells, [input_feed])
+        new_state = [self.reorder_state(state, new_order) for state in cached_state]
+        prev_hiddens_tensor = torch.stack(new_state[0])
+        prev_cells_tensor = torch.stack(new_state[1])
+        cached_state_new = torch.jit.annotate(
+            Dict[str, Optional[Tensor]],
+            {"prev_hiddens": prev_hiddens_tensor, "prev_cells": prev_cells_tensor, "input_feed": new_state[2][0]})
+        self.set_incremental_state(incremental_state, 'cached_state', cached_state_new),
+        return
 
     def max_positions(self):
         """Maximum output length supported by the decoder."""

fairseq/utils.py

@@ -212,7 +212,7 @@ def buffered_arange(max):
 
 
 def convert_padding_direction(
-    src_tokens, padding_idx, right_to_left=False, left_to_right=False
+    src_tokens, padding_idx, right_to_left: bool = False, left_to_right: bool = False
 ):
     assert right_to_left ^ left_to_right
     pad_mask = src_tokens.eq(padding_idx)
@@ -226,7 +226,10 @@ def convert_padding_direction(
         # already left padded
         return src_tokens
     max_len = src_tokens.size(1)
-    range = buffered_arange(max_len).type_as(src_tokens).expand_as(src_tokens)
+    buffered = torch.empty(0).long()
+    if max_len > 0:
+        torch.arange(max_len, out=buffered)
+    range = buffered.type_as(src_tokens).expand_as(src_tokens)
     num_pads = pad_mask.long().sum(dim=1, keepdim=True)
     if right_to_left:
         index = torch.remainder(range - num_pads, max_len)