Commit 48d9afbe authored by Myle Ott's avatar Myle Ott Committed by Facebook Github Bot
Browse files

Speed improvements (#531) 

Summary:
* Add FusedLayerNorm and FusedAdam
* Softmax and zero grad optimizations
Pull Request resolved: https://github.com/pytorch/fairseq/pull/531

Differential Revision: D14218457

Pulled By: myleott

fbshipit-source-id: 5656b2d0152cd85f77dc21ec0e1439ec04b9fa89
parent a24880bd

README.md

+10 −4
Original line number Diff line number Diff line
@@ -36,12 +36,12 @@ translation and language modeling datasets. 
![Model](fairseq.gif)



# Requirements and Installation

* A [PyTorch installation](http://pytorch.org/)



* [PyTorch](http://pytorch.org/) version >= 1.0.0

* Python version >= 3.6

* For training new models, you'll also need an NVIDIA GPU and [NCCL](https://github.com/NVIDIA/nccl)

* Python version 3.6



Currently fairseq requires PyTorch version >= 1.0.0.

Please follow the instructions here: https://github.com/pytorch/pytorch#installation.

Please follow the instructions here to install PyTorch: https://github.com/pytorch/pytorch#installation.



If you use Docker make sure to increase the shared memory size either with

`--ipc=host` or `--shm-size` as command line options to `nvidia-docker run`.
@@ -60,6 +60,12 @@ cd fairseq 
pip install --editable .

```



**Improved training speed**



Training speed can be further improved by installing NVIDIA's

[apex](https://github.com/NVIDIA/apex) library with the `--cuda_ext` option.

fairseq will automatically switch to the faster modules provided by apex.



# Getting Started



The [full documentation](https://fairseq.readthedocs.io/) contains instructions

fairseq/distributed_utils.py

+9 −7
Original line number Diff line number Diff line
@@ -122,8 +122,10 @@ def all_gather_list(data, group=None, max_size=16384): 
    if not hasattr(all_gather_list, '_buffer') or \

            all_gather_list._buffer.numel() < buffer_size:

        all_gather_list._buffer = torch.cuda.ByteTensor(buffer_size)

        all_gather_list._cpu_buffer = torch.ByteTensor(max_size).pin_memory()

    buffer = all_gather_list._buffer

    buffer.zero_()

    cpu_buffer = all_gather_list._cpu_buffer



    enc = pickle.dumps(data)

    enc_size = len(enc)
@@ -131,10 +133,12 @@ def all_gather_list(data, group=None, max_size=16384): 
        raise ValueError('encoded data exceeds max_size: {}'.format(enc_size + 2))

    assert max_size < 255*256



    buffer_rank = buffer[rank * max_size : (rank + 1) * max_size]

    buffer_rank[0] = enc_size // 255  # this encoding works for max_size < 65k

    buffer_rank[1] = enc_size % 255

    buffer_rank[2:enc_size+2] = torch.ByteTensor(list(enc))

    cpu_buffer[0] = enc_size // 255  # this encoding works for max_size < 65k

    cpu_buffer[1] = enc_size % 255

    cpu_buffer[2 : enc_size + 2] = torch.ByteTensor(list(enc))

    start = rank * max_size

    size = enc_size + 2

    buffer[start : start + size].copy_(cpu_buffer[:size])



    all_reduce(buffer, group=group)
@@ -144,9 +148,7 @@ def all_gather_list(data, group=None, max_size=16384): 
            out_buffer = buffer[i * max_size : (i + 1) * max_size]

            size = (255 * utils.item(out_buffer[0])) + utils.item(out_buffer[1])

            if size > 0:

                result.append(

                    pickle.loads(bytes(out_buffer[2:size+2].tolist()))

                )

                result.append(pickle.loads(bytes(out_buffer[2 : size + 2].tolist())))

        return result

    except pickle.UnpicklingError:

        raise Exception(

fairseq/models/fairseq_decoder.py

+9 −4
Original line number Diff line number Diff line
@@ -6,7 +6,8 @@ 
# can be found in the PATENTS file in the same directory.



import torch.nn as nn

import torch.nn.functional as F



from fairseq import utils





class FairseqDecoder(nn.Module):
@@ -15,6 +16,7 @@ class FairseqDecoder(nn.Module): 
    def __init__(self, dictionary):

        super().__init__()

        self.dictionary = dictionary

        self.onnx_trace = False



    def forward(self, prev_output_tokens, encoder_out):

        """
@@ -33,6 +35,9 @@ class FairseqDecoder(nn.Module): 
        """

        raise NotImplementedError



    def prepare_for_onnx_export_(self):

        self.onnx_trace = True



    def get_normalized_probs(self, net_output, log_probs, sample):

        """Get normalized probabilities (or log probs) from a net's output."""
@@ -45,11 +50,11 @@ class FairseqDecoder(nn.Module): 
            out = self.adaptive_softmax.get_log_prob(net_output[0], target=target)

            return out.exp_() if not log_probs else out



        logits = net_output[0].float()

        logits = net_output[0]

        if log_probs:

            return F.log_softmax(logits, dim=-1)

            return utils.log_softmax(logits, dim=-1, onnx_trace=self.onnx_trace)

        else:

            return F.softmax(logits, dim=-1)

            return utils.softmax(logits, dim=-1, onnx_trace=self.onnx_trace)



    def max_positions(self):

        """Maximum input length supported by the decoder."""

fairseq/models/fconv_self_att.py

+6 −6
Original line number Diff line number Diff line
@@ -13,8 +13,8 @@ import torch.nn as nn 
import torch.nn.functional as F



from fairseq.modules import (

    DownsampledMultiHeadAttention, GradMultiply, LearnedPositionalEmbedding,

    LinearizedConvolution,

    DownsampledMultiHeadAttention, GradMultiply, LayerNorm,

    LearnedPositionalEmbedding, LinearizedConvolution,

)

from fairseq import utils
@@ -351,13 +351,13 @@ class FConvDecoder(FairseqDecoder): 
            # pretrained and trained models are joined

            self.joining = nn.Sequential(

                Linear(out_embed_dim*2, out_embed_dim*2),

                nn.LayerNorm(out_embed_dim*2),

                LayerNorm(out_embed_dim*2),

                nn.GLU(),

                Linear(out_embed_dim, out_embed_dim*2),

                nn.LayerNorm(out_embed_dim*2),

                LayerNorm(out_embed_dim*2),

                nn.GLU(),

                Linear(out_embed_dim, out_embed_dim),

                nn.LayerNorm(out_embed_dim)

                LayerNorm(out_embed_dim)

            )

            # pretrained model contains an output layer that is nhid -> vocab size

            # but the models are combined in their hidden state
@@ -470,7 +470,7 @@ class SelfAttention(nn.Module): 
        self.in_proj_q = Linear(out_channels, embed_dim)

        self.in_proj_k = Linear(out_channels, embed_dim)

        self.in_proj_v = Linear(out_channels, embed_dim)

        self.ln = nn.LayerNorm(out_channels)

        self.ln = LayerNorm(out_channels)



    def forward(self, x):

        residual = x

fairseq/models/lightconv.py

+6 −12
Original line number Diff line number Diff line
@@ -11,17 +11,16 @@ import torch 
import torch.nn as nn

import torch.nn.functional as F



from fairseq import options

from fairseq import utils



from fairseq import options, utils

from fairseq.modules import (

    AdaptiveInput, AdaptiveSoftmax, CharacterTokenEmbedder, LearnedPositionalEmbedding, MultiheadAttention,

    SinusoidalPositionalEmbedding, DynamicConv1dTBC, LightweightConv1dTBC

    AdaptiveInput, AdaptiveSoftmax, CharacterTokenEmbedder, LayerNorm,

    LearnedPositionalEmbedding, MultiheadAttention, SinusoidalPositionalEmbedding,

    DynamicConv1dTBC, LightweightConv1dTBC,

)



from . import (

    FairseqIncrementalDecoder, FairseqEncoder, FairseqLanguageModel, FairseqModel, register_model,

    register_model_architecture,

    FairseqIncrementalDecoder, FairseqEncoder, FairseqLanguageModel,

    FairseqModel, register_model, register_model_architecture,

)
@@ -771,11 +770,6 @@ def Embedding(num_embeddings, embedding_dim, padding_idx): 
    return m





def LayerNorm(embedding_dim):

    m = nn.LayerNorm(embedding_dim)

    return m





def Linear(in_features, out_features, bias=True):

    m = nn.Linear(in_features, out_features, bias)

    nn.init.xavier_uniform_(m.weight)