Skip to content

About

Facebook AI Research Sequence-to-Sequence Toolkit written in Python.

Topics

Resources

Code of conduct

Contributing

Stars

26 stars

Watchers

1 watching

Forks

 
 

Repository files navigation



MIT License


Fairseq(-py) is a sequence modeling toolkit that allows researchers and developers to train custom models for translation, summarization, language modeling and other text generation tasks.

This fork tracks the upstream fairseq and adds the features below to the transformer model and the translation task. Everything else behaves as in upstream.

Features

Add these flags to fairseq-train / fairseq-generate / fairseq-interactive:

Feature Description Flags Source
Knowledge Distillation (Hinton et al., Kim & Rush, Wang et al., Gumma et al.) Word-level KD from a pretrained teacher (forward KL). lambda weighs the KD loss against the NLL loss: loss = lambda * KD + (1 - lambda) * NLL, with lambda in (0, 1]. See the example. --task seq2seq_lm_distillation --criterion lm_distillation_loss --teacher-checkpoint-path $path --kd-args '{"lambda": 1.0}' Selective Distillation
Recurrent Stacking (Dabre & Fujita) Extreme parameter sharing: one encoder/decoder layer applied N times --encoder-recurrent-stacking 6 --decoder-recurrent-stacking 6 -
Low-Rank Adaptation (LoRA) (Hu et al.) Trains low-rank adapters on the listed modules and freezes the rest. Checkpoints store the merged weights, so they load into a plain model for inference (with --load-checkpoint-liberally) and also resume LoRA training. --lora-args '{"r": 8, "alpha": 16, "dropout": 0.05, "bias": "none", "target_modules": "k_proj,v_proj", "rank_scaled": false}' --load-checkpoint-liberally LoRA
Rotary Positional Embedding (RoPE) (Su et al.) Rotates half of each attention head by its absolute position, which makes attention depend on relative positions --rope-args '{"theta": 10000, "scaling_factor": 1.0}' --attn-implementation fast --no-token-positional-embeddings --load-checkpoint-liberally rotary-embedding-torch
Attention with Linear Biases (ALiBi) (Press et al.) Biases attention scores with a penalty proportional to the query-key distance. type is symmetrical or asymmetrical (encoder). --alibi-args '{"type": "symmetrical"}' --no-token-positional-embeddings --load-checkpoint-liberally ALiBi for fairseq
Gated Linear Unit (GLU) (Shazeer) GLU feed-forward block (e.g. SwiGLU with --activation-fn silu); --*-ffn-embed-dim is the inner dimension --encoder-use-glu --decoder-use-glu GLU
RMSNorm (Zhang & Sennrich) RMSNorm instead of LayerNorm --encoder-use-rmsnorm --decoder-use-rmsnorm RMSNorm
Factorized Embeddings (Lan et al.) Embeddings through a low-dimensional bottleneck. Not compatible with tied embeddings (--share-all-embeddings, --share-decoder-input-output-embed). --encoder-factorized-embed-dim 128 --decoder-factorized-embed-dim 128 -
Fast attention Attention on top of scaled_dot_product_attention. Same parameters as the default, so checkpoints are interchangeable. Attention weights are not returned (use the default implementation for --print-alignment). --attn-implementation fast (default: fairseq) -
Grouped Query Attention (GQA) (Ainslie et al.) Groups of query heads share key/value heads; the decoding cache only holds the key/value heads --attn-implementation fast_gqa --encoder-kv-attention-heads 2 --decoder-kv-attention-heads 2 GQA
Fused projections Single q/k/v projection matrix for fast attention or GQA. Parameter names differ, so fused and unfused checkpoints are not interchangeable. --attn-implementation fast_fused or --attn-implementation fast_gqa_fused -
StableAdamW (Wortsman et al.) Adam with per-tensor update clipping, so gradient clipping is not needed --optimizer adam --adam-stable --clip-norm 0 -
Cut Cross-Entropy (Wijmans et al.) (Label-smoothed) cross entropy computed from the decoder features and the output projection, without materializing the (tokens x vocab) logits: same loss and gradients as label_smoothed_cross_entropy, with much less memory for large vocabularies. Needs pip install cut-cross-entropy; the fast kernel needs an Ampere or newer GPU (elsewhere a torch.compile version is used). Transformer models only; not for KD; no --report-accuracy. --criterion cut_cross_entropy --label-smoothing 0.1 ml-cross-entropy
Length-normalized loss Normalizes the loss of each sentence by its length (e.g. for mixing sentence- and document-level data). Works with cross_entropy, label_smoothed_cross_entropy and lm_distillation_loss. --length-normalize-loss -
Sanity validation Runs a full validation pass before training starts --run-sanity-validation-steps -
Torch compile for inference Compiles the encoder and decoder for fairseq-generate / fairseq-interactive; $mode is a torch.compile mode --torch-compile $mode -
BF16 --bf16 no longer implies --tpu, so it works on GPUs. Models pretrained with fp16 can behave very differently in bf16. --bf16 -

Other differences from upstream:

  • the translation task accepts sequences up to 8192 tokens by default (--max-source-positions / --max-target-positions)
  • --wandb-run-name sets the Weights & Biases run name
  • architectures transformer_IT2, transformer_IT2_dist, roformer_IT2 and roformer_IT2_dist (IndicTrans2)

Checkpoints trained with earlier versions of this fork load and produce the same outputs with the current code.

Requirements and Installation

  • PyTorch >= 2.1.0
  • Python >= 3.8, <= 3.12
  • For training new models, you'll also need an NVIDIA GPU and NCCL

To install and develop locally:

git clone https://github.com/VarunGumma/fairseq
cd fairseq
pip install -e ./

or to install directly:

pip install git+https://github.com/VarunGumma/fairseq.git
  • For faster training install NVIDIA's apex library (fused Adam, LayerNorm and RMSNorm are then used automatically):
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" \
  --global-option="--deprecated_fused_adam" --global-option="--xentropy" \
  --global-option="--fast_multihead_attn" ./
  • For large datasets install PyArrow: pip install pyarrow
  • 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.

License

fairseq(-py) is MIT-licensed. The license applies to the pre-trained models as well.

Citation

Please cite as:

@inproceedings{gumma-etal-2023-empirical,
    title = "An Empirical Study of Leveraging Knowledge Distillation for Compressing Multilingual Neural Machine Translation Models",
    author = "Gumma, Varun  and
      Dabre, Raj  and
      Kumar, Pratyush",
    editor = "Nurminen, Mary  and
      Brenner, Judith  and
      Koponen, Maarit  and
      Latomaa, Sirkku  and
      Mikhailov, Mikhail  and
      Schierl, Frederike  and
      Ranasinghe, Tharindu  and
      Vanmassenhove, Eva  and
      Vidal, Sergi Alvarez  and
      Aranberri, Nora  and
      Nunziatini, Mara  and
      Escart{\'\i}n, Carla Parra  and
      Forcada, Mikel  and
      Popovic, Maja  and
      Scarton, Carolina  and
      Moniz, Helena",
    booktitle = "Proceedings of the 24th Annual Conference of the European Association for Machine Translation",
    month = jun,
    year = "2023",
    address = "Tampere, Finland",
    publisher = "European Association for Machine Translation",
    url = "https://aclanthology.org/2023.eamt-1.11",
    pages = "103--114",
    abstract = "Knowledge distillation (KD) is a well-known method for compressing neural models. However, works focusing on distilling knowledge from large multilingual neural machine translation (MNMT) models into smaller ones are practically nonexistent, despite the popularity and superiority of MNMT. This paper bridges this gap by presenting an empirical investigation of knowledge distillation for compressing MNMT models. We take Indic to English translation as a case study and demonstrate that commonly used language-agnostic and language-aware KD approaches yield models that are 4-5x smaller but also suffer from performance drops of up to 3.5 BLEU. To mitigate this, we then experiment with design considerations such as shallower versus deeper models, heavy parameter sharing, multistage training, and adapters. We observe that deeper compact models tend to be as good as shallower non-compact ones and that fine-tuning a distilled model on a high-quality subset slightly boosts translation quality. Overall, we conclude that compressing MNMT models via KD is challenging, indicating immense scope for further research.",
}
@inproceedings{ott2019fairseq,
  title = {fairseq: A Fast, Extensible Toolkit for Sequence Modeling},
  author = {Myle Ott and Sergey Edunov and Alexei Baevski and Angela Fan and Sam Gross and Nathan Ng and David Grangier and Michael Auli},
  booktitle = {Proceedings of NAACL-HLT 2019: Demonstrations},
  year = {2019},
}

and please add a footnote url to this repository.

Final Note

I will try my best to keep this repo synced with the upstream fairseq repository. Feel free to raise issues or pull requests to report bugs or propose new features.

About

Facebook AI Research Sequence-to-Sequence Toolkit written in Python.

Topics

Resources

Code of conduct

Contributing

Stars

26 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages