Fully quantize Fairseq transformer
Summary:
Pull Request resolved:
F.linear -> nn.Linear so FBGEMM backend could quantize the linear projection. We observed 3x+ speedup.
Add backward compatibility code.
@@ -669,11 +669,6 @@ class TransformerDecoder(FairseqIncrementalDecoder):
factor=args.adaptive_softmax_factor,
tie_proj=args.tie_adaptive_proj,
)
- elif not self.share_input_output_embed:
- self.embed_out = nn.Parameter(
- torch.Tensor(len(dictionary), self.output_embed_dim)
- )
- nn.init.normal_(self.embed_out, mean=0, std=self.output_embed_dim ** -0.5)
if args.decoder_normalize_before and not getattr(
args, "no_decoder_final_norm", False
@@ -686,6 +681,16 @@ class TransformerDecoder(FairseqIncrementalDecoder):
else:
self.layernorm_embedding = None
+ if self.share_input_output_embed:
+ self.output_projection = nn.Linear(
+ 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), bias=False
+ )
+ nn.init.normal_(self.output_projection.weight, mean=0, std=self.output_embed_dim ** -0.5)
+
def build_decoder_layer(self, args, no_encoder_attn=False):
return TransformerDecoderLayer(args, no_encoder_attn)
@@ -852,10 +857,7 @@ class TransformerDecoder(FairseqIncrementalDecoder):
"""Project features to the vocabulary size."""
if self.adaptive_softmax is None:
# project back to size of vocabulary
- if self.share_input_output_embed:
- return F.linear(features, self.embed_tokens.weight)
- else:
- return F.linear(features, self.embed_out)
+ return self.output_projection(features)
else:
return features
@@ -890,6 +892,18 @@ class TransformerDecoder(FairseqIncrementalDecoder):
"{}.embed_positions._float_tensor".format(name)
] = torch.FloatTensor(1)
+ embed_tokens_weights_key = f"{name}.embed_tokens.weights"
+ embed_out_key = f"{name}.embed_out"
+ if embed_tokens_weights_key in state_dict:
+ state_dict[f"{name}.output_projection.weight"] = state_dict[
+ embed_tokens_weights_key
+ ]
+ if embed_out_key in state_dict:
+ state_dict[f"{name}.output_projection.weight"] = state_dict[
+ embed_out_key
+ ]
+ del state_dict[embed_out_key]
+
for i in range(self.num_layers):
# update layer norms
layer_norm_map = {
