Fully quantize Fairseq transformer

[cndn]

Summary:
F.linear -> nn.Linear so FBGEMM backend could quantize the linear projection. We observed 3x+ speedup.

Backward compatibility code is added to upgrade_state_dict_named. Locally it worked.

Testing loading OSS checkpoints.

Differential Revision: D20967830

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D20967830

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D20967830

[erip]

Nicely done! This should resolve #1943, too!

[erip]

Looks like there might be some backwards compatibility issues. I tried loading the wmt14.en-fr.joined-dict.transformer pretrained model:

My test script:

$ cat test.py
#!/usr/bin/env python

import torch

from fairseq.sequence_generator import SequenceGenerator
from fairseq.models.transformer import TransformerModel
from fairseq.data import Dictionary

if __name__ == "__main__":
    tgt_dict = Dictionary.load(open('dict.fr.txt'))
    model = TransformerModel.from_pretrained('.', 'model.pt', '.', bpe='subword_nmt', bpe_vocab='bpecodes')

    generator = SequenceGenerator(model.models, tgt_dict)

    scripted_gen = torch.jit.script(generator)
    scripted_gen.save('generator.pt')

And the error:

$ ./test.py
Traceback (most recent call last):
  File "./test.py", line 11, in <module>
    model = TransformerModel.from_pretrained('.', 'model.pt', '.', bpe='subword_nmt', bpe_vocab='bpecodes')
  File "/Users/erippeth/miniconda3/envs/fairseq-dev/lib/python3.6/site-packages/fairseq/models/fairseq_model.py", line 218, in from_pretrained
    **kwargs,
  File "/Users/erippeth/miniconda3/envs/fairseq-dev/lib/python3.6/site-packages/fairseq/hub_utils.py", line 73, in from_pretrained
    arg_overrides=kwargs,
  File "/Users/erippeth/miniconda3/envs/fairseq-dev/lib/python3.6/site-packages/fairseq/checkpoint_utils.py", line 210, in load_model_ensemble_and_task
    model.load_state_dict(state["model"], strict=strict, args=args)
  File "/Users/erippeth/miniconda3/envs/fairseq-dev/lib/python3.6/site-packages/fairseq/models/fairseq_model.py", line 93, in load_state_dict
    return super().load_state_dict(new_state_dict, strict)
  File "/Users/erippeth/miniconda3/envs/fairseq-dev/lib/python3.6/site-packages/torch/nn/modules/module.py", line 855, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for TransformerModel:
	Missing key(s) in state_dict: "decoder.output_projection.weight", "decoder.output_projection.bias".

[erip]

It looks like there's at least an issue here -- specifically I think the property is ...embed_tokens.weight, not ...embed_tokens.weights. Once I correct that, the stack trace is just the missing bias term:

  File "./test.py", line 11, in <module>
    model = TransformerModel.from_pretrained('.', 'model.pt', '.', bpe='subword_nmt', bpe_vocab='bpecodes')
  File "/Users/erippeth/miniconda3/envs/fairseq-dev/lib/python3.6/site-packages/fairseq/models/fairseq_model.py", line 218, in from_pretrained
    **kwargs,
  File "/Users/erippeth/miniconda3/envs/fairseq-dev/lib/python3.6/site-packages/fairseq/hub_utils.py", line 73, in from_pretrained
    arg_overrides=kwargs,
  File "/Users/erippeth/miniconda3/envs/fairseq-dev/lib/python3.6/site-packages/fairseq/checkpoint_utils.py", line 210, in load_model_ensemble_and_task
    model.load_state_dict(state["model"], strict=strict, args=args)
  File "/Users/erippeth/miniconda3/envs/fairseq-dev/lib/python3.6/site-packages/fairseq/models/fairseq_model.py", line 93, in load_state_dict
    return super().load_state_dict(new_state_dict, strict)
  File "/Users/erippeth/miniconda3/envs/fairseq-dev/lib/python3.6/site-packages/torch/nn/modules/module.py", line 855, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for TransformerModel:
	Missing key(s) in state_dict: "decoder.output_projection.bias".

---

Edit:

the original F.linear projection adds no bias by way of default arguments, so it's safe to add bias=False to the newly added linear projections, too. Since I can't push commits to your repo, here's the diff:

diff --git a/fairseq/models/transformer.py b/fairseq/models/transformer.py
index 98c4ab5..6b74129 100644
--- a/fairseq/models/transformer.py
+++ b/fairseq/models/transformer.py
@@ -683,11 +683,11 @@ class TransformerDecoder(FairseqIncrementalDecoder):
 
         if self.share_input_output_embed:
             self.output_projection = nn.Linear(
-                self.embed_tokens.weight.shape[1], self.embed_tokens.weight.shape[0]
+                self.embed_tokens.weight.shape[1], self.embed_tokens.weight.shape[0], bias=False
             )
         else:
             self.output_projection = nn.Linear(
-                self.output_embed_dim, len(dictionary)
+                self.output_embed_dim, len(dictionary), bias=False
             )
 
     def build_decoder_layer(self, args, no_encoder_attn=False):
@@ -891,7 +891,7 @@ class TransformerDecoder(FairseqIncrementalDecoder):
                 "{}.embed_positions._float_tensor".format(name)
             ] = torch.FloatTensor(1)
 
-        embed_tokens_weights_key = f"{name}.embed_tokens.weights"
+        embed_tokens_weights_key = f"{name}.embed_tokens.weight"
         embed_out_key = f"{name}.embed_out"
         if embed_tokens_weights_key in state_dict:
             state_dict[f"{name}.output_projection.weight"] = state_dict[

[cndn]

Thanks @erip ! I will address it very soon.

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D20967830

[facebook-github-bot]

This pull request has been merged in 6379573.