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.
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
translationtask accepts sequences up to 8192 tokens by default (--max-source-positions/--max-target-positions) --wandb-run-namesets the Weights & Biases run name- architectures
transformer_IT2,transformer_IT2_dist,roformer_IT2androformer_IT2_dist(IndicTrans2)
Checkpoints trained with earlier versions of this fork load and produce the same outputs with the current code.
- 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=hostor--shm-sizeas command line options tonvidia-docker run.
fairseq(-py) is MIT-licensed.
The license applies to the pre-trained models as well.
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.
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.
