NLLB-MOE

## [](#overview)Overview

The NLLB model was presented in [No Language Left Behind: Scaling Human-Centered Machine Translation](https://arxiv.org/abs/2207.04672) by Marta R. Costa-jussà, James Cross, Onur Çelebi, Maha Elbayad, Kenneth Heafield, Kevin Heffernan, Elahe Kalbassi, Janice Lam, Daniel Licht, Jean Maillard, Anna Sun, Skyler Wang, Guillaume Wenzek, Al Youngblood, Bapi Akula, Loic Barrault, Gabriel Mejia Gonzalez, Prangthip Hansanti, John Hoffman, Semarley Jarrett, Kaushik Ram Sadagopan, Dirk Rowe, Shannon Spruit, Chau Tran, Pierre Andrews, Necip Fazil Ayan, Shruti Bhosale, Sergey Edunov, Angela Fan, Cynthia Gao, Vedanuj Goswami, Francisco Guzmán, Philipp Koehn, Alexandre Mourachko, Christophe Ropers, Safiyyah Saleem, Holger Schwenk, and Jeff Wang.

The abstract of the paper is the following:

_Driven by the goal of eradicating language barriers on a global scale, machine translation has solidified itself as a key focus of artificial intelligence research today. However, such efforts have coalesced around a small subset of languages, leaving behind the vast majority of mostly low-resource languages. What does it take to break the 200 language barrier while ensuring safe, high quality results, all while keeping ethical considerations in mind? In No Language Left Behind, we took on this challenge by first contextualizing the need for low-resource language translation support through exploratory interviews with native speakers. Then, we created datasets and models aimed at narrowing the performance gap between low and high-resource languages. More specifically, we developed a conditional compute model based on Sparsely Gated Mixture of Experts that is trained on data obtained with novel and effective data mining techniques tailored for low-resource languages. We propose multiple architectural and training improvements to counteract overfitting while training on thousands of tasks. Critically, we evaluated the performance of over 40,000 different translation directions using a human-translated benchmark, Flores-200, and combined human evaluation with a novel toxicity benchmark covering all languages in Flores-200 to assess translation safety. Our model achieves an improvement of 44% BLEU relative to the previous state-of-the-art, laying important groundwork towards realizing a universal translation system._

Tips:

-   M2M100ForConditionalGeneration is the base model for both NLLB and NLLB MoE
-   The NLLB-MoE is very similar to the NLLB model, but it’s feed forward layer is based on the implementation of SwitchTransformers.
-   The tokenizer is the same as the NLLB models.

This model was contributed by [Arthur Zucker](https://huggingface.co/ArtZucker). The original code can be found [here](https://github.com/facebookresearch/fairseq).

## [](#implementation-differences-with-switchtransformers)Implementation differences with SwitchTransformers

The biggest difference is the way the tokens are routed. NLLB-MoE uses a \`top-2-gate\` which means that for each input, only the top two experts are selected based on the highest predicted probabilities from the gating network, and the remaining experts are ignored. In \`SwitchTransformers\`, only the top-1 probabilities are computed, which means that tokens have less probability of being forwarded. Moreover, if a token is not routed to any expert, \`SwitchTransformers\` still adds its unmodified hidden states (kind of like a residual connection) while they are masked in \`NLLB\`'s top-2 routing mechanism.

## [](#generating-with-nllbmoe)Generating with NLLB-MoE

The avalable checkpoints requires around 350GB of storage. Make sure to use \`accelerate\` if you do not have enough RAM on your machine.

While generating the target text set the `forced_bos_token_id` to the target language id. The following example shows how to translate English to French using the _facebook/nllb-200-distilled-600M_ model.

Note that we’re using the BCP-47 code for French `fra_Latn`. See [here](https://github.com/facebookresearch/flores/blob/main/flores200/README.md#languages-in-flores-200) for the list of all BCP-47 in the Flores 200 dataset.

```
>>> from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-moe-54b")
>>> model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-moe-54b")

>>> article = "Previously, Ring's CEO, Jamie Siminoff, remarked the company started when his doorbell wasn't audible from his shop in his garage."
>>> inputs = tokenizer(article, return_tensors="pt")

>>> translated_tokens = model.generate(
...     **inputs, forced_bos_token_id=tokenizer.lang_code_to_id["fra_Latn"], max_length=50
... )
>>> tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
"Auparavant, le PDG de Ring, Jamie Siminoff, a fait remarquer que la société avait commencé lorsque sa sonnette n'était pas audible depuis son magasin dans son garage."```

### [](#generating-from-any-other-language-than-english)Generating from any other language than English

English (`eng_Latn`) is set as the default language from which to translate. In order to specify that you’d like to translate from a different language, you should specify the BCP-47 code in the `src_lang` keyword argument of the tokenizer initialization.

See example below for a translation from romanian to german:

```
>>> from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-moe-54b", src_lang="ron_Latn")
>>> model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-moe-54b")

>>> article = "Şeful ONU spune că nu există o soluţie militară în Siria"
>>> inputs = tokenizer(article, return_tensors="pt")

>>> translated_tokens = model.generate(
...     **inputs, forced_bos_token_id=tokenizer.lang_code_to_id["deu_Latn"], max_length=30
... )
>>> tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]```

## [](#documentation-resources)Documentation resources

-   [Translation task guide](./tasks/translation)
-   [Summarization task guide](./tasks/summarization)

## [](#transformers.NllbMoeConfig)NllbMoeConfig

### class transformers.NllbMoeConfig

[](#transformers.NllbMoeConfig)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/nllb_moe/configuration_nllb_moe.py#L27)

( vocab\_size = 128112max\_position\_embeddings = 1024encoder\_layers = 12encoder\_ffn\_dim = 4096encoder\_attention\_heads = 16decoder\_layers = 12decoder\_ffn\_dim = 4096decoder\_attention\_heads = 16encoder\_layerdrop = 0.05decoder\_layerdrop = 0.05use\_cache = Trueis\_encoder\_decoder = Trueactivation\_function = 'relu'd\_model = 1024dropout = 0.1attention\_dropout = 0.1activation\_dropout = 0.0init\_std = 0.02decoder\_start\_token\_id = 2scale\_embedding = Truerouter\_bias = Falserouter\_dtype = 'float32'router\_ignore\_padding\_tokens = Falsenum\_experts = 128expert\_capacity = 64encoder\_sparse\_step = 4decoder\_sparse\_step = 4router\_z\_loss\_coef = 0.001router\_aux\_loss\_coef = 0.001second\_expert\_policy = 'all'normalize\_router\_prob\_before\_dropping = Falsebatch\_prioritized\_routing = Falsemoe\_eval\_capacity\_token\_fraction = 1.0moe\_token\_dropout = 0.2pad\_token\_id = 1bos\_token\_id = 0eos\_token\_id = 2output\_router\_logits = False\*\*kwargs )

This is the configuration class to store the configuration of a [NllbMoeModel](/docs/transformers/v4.30.0/en/model_doc/nllb-moe#transformers.NllbMoeModel). It is used to instantiate an NLLB-MoE model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of the NLLB-MoE [facebook/nllb-moe-54b](https://huggingface.co/facebook/nllb-moe-54b) architecture.

Configuration objects inherit from [PretrainedConfig](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig) and can be used to control the model outputs. Read the documentation from [PretrainedConfig](/docs/transformers/v4.30.0/en/main_classes/configuration#transformers.PretrainedConfig) for more information.

Example:

```
>>> from transformers import NllbMoeModel, NllbMoeConfig

>>> 
>>> configuration = NllbMoeConfig()

>>> 
>>> model = NllbMoeModel(configuration)

>>> 
>>> configuration = model.config```

## [](#transformers.NllbMoeTop2Router)NllbMoeTop2Router

### class transformers.NllbMoeTop2Router

[](#transformers.NllbMoeTop2Router)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/nllb_moe/modeling_nllb_moe.py#L249)

( config: NllbMoeConfig )

Router using tokens choose top-2 experts assignment.

This router uses the same mechanism as in NLLB-MoE from the fairseq repository. Items are sorted by router\_probs and then routed to their choice of expert until the expert’s expert\_capacity is reached. **There is no guarantee that each token is processed by an expert**, or that each expert receives at least one token.

The router combining weights are also returned to make sure that the states that are not updated will be masked.

#### route\_tokens

[](#transformers.NllbMoeTop2Router.route_tokens)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/nllb_moe/modeling_nllb_moe.py#L290)

( router\_logits: Tensorinput\_dtype: dtype = torch.float32padding\_mask: typing.Optional\[torch.LongTensor\] = None )

Computes the `dispatch_mask` and the `dispatch_weights` for each experts. The masks are adapted to the expert capacity.

#### forward

[](#transformers.NllbMoeTop2Router.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/nllb_moe/modeling_nllb_moe.py#L375)

( hidden\_states: Tensorpadding\_mask: typing.Optional\[torch.LongTensor\] = None ) → top\_1\_mask (`torch.Tensor` of shape (batch\_size, sequence\_length))

Parameters

-   [](#transformers.NllbMoeTop2Router.forward.hidden_states)**hidden\_states** (`torch.Tensor`) — (batch\_size, sequence\_length, hidden\_dim) from which router probabilities are computed.

Returns

top\_1\_mask (`torch.Tensor` of shape (batch\_size, sequence\_length))

Index tensor of shape \[batch\_size, sequence\_length\] corresponding to the expert selected for each token using the top1 probabilities of the router. router\_probabilities (`torch.Tensor` of shape (batch\_size, sequence\_length, nump\_experts)): Tensor of shape (batch\_size, sequence\_length, num\_experts) corresponding to the probabilities for each token and expert. Used for routing tokens to experts. router\_logits (`torch.Tensor` of shape (batch\_size, sequence\_length))): Logits tensor of shape (batch\_size, sequence\_length, num\_experts) corresponding to raw router logits. This is used later for computing router z-loss.

The hidden states are reshaped to simplify the computation of the router probabilities (combining weights for each experts.)

## [](#transformers.NllbMoeSparseMLP)NllbMoeSparseMLP

### class transformers.NllbMoeSparseMLP

[](#transformers.NllbMoeSparseMLP)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/nllb_moe/modeling_nllb_moe.py#L426)

( config: NllbMoeConfigffn\_dim: intexpert\_class: Module = <class 'transformers.models.nllb\_moe.modeling\_nllb\_moe.NllbMoeDenseActDense'> )

Implementation of the NLLB-MoE sparse MLP module.

#### forward

[](#transformers.NllbMoeSparseMLP.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/nllb_moe/modeling_nllb_moe.py#L442)

( hidden\_states: Tensorpadding\_mask: typing.Optional\[torch.Tensor\] = False ) → hidden\_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_dim)`)

Parameters

-   [](#transformers.NllbMoeSparseMLP.forward.hidden_states)**hidden\_states** (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_dim)`) — The hidden states
-   [](#transformers.NllbMoeSparseMLP.forward.padding_mask)**padding\_mask** (`torch.Tensor`, _optional_, defaults to `False`) — Attention mask. Can be in the causal form or not.

Returns

hidden\_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_dim)`)

Updated hidden states router\_logits (`torch.Tensor` of shape `(batch_size, sequence_length, num_experts)`): Needed for computing the loss

The goal of this forward pass is to have the same number of operation as the equivalent `NllbMoeDenseActDense` (mlp) layer. This means that all of the hidden states should be processed at most twice ( since we are using a top\_2 gating mecanism). This means that we keep the complexity to O(batch\_size x sequence\_length x hidden\_dim) instead of O(num\_experts x batch\_size x sequence\_length x hidden\_dim).

1- Get the `router_probs` from the `router`. The shape of the `router_mask` is `(batch_size X sequence_length, num_expert)` and corresponds to the boolean version of the `router_probs`. The inputs are masked using the `router_mask`.

2- Dispatch the hidden\_states to its associated experts. The router probabilities are used to weight the contribution of each experts when updating the masked hidden states.

## [](#transformers.NllbMoeModel)NllbMoeModel

### class transformers.NllbMoeModel

[](#transformers.NllbMoeModel)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/nllb_moe/modeling_nllb_moe.py#L1503)

( config: NllbMoeConfig )

Parameters

-   [](#transformers.NllbMoeModel.config)**config** ([NllbMoeConfig](/docs/transformers/v4.30.0/en/model_doc/nllb-moe#transformers.NllbMoeConfig)) — Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [from\_pretrained()](/docs/transformers/v4.30.0/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) method to load the model weights.

The bare NllbMoe Model outputting raw hidden-states without any specific head on top. This model inherits from [PreTrainedModel](/docs/transformers/v4.30.0/en/main_classes/model#transformers.PreTrainedModel). Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.)

This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior.

#### forward

[](#transformers.NllbMoeModel.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/nllb_moe/modeling_nllb_moe.py#L1539)

( input\_ids: typing.Optional\[torch.LongTensor\] = Noneattention\_mask: typing.Optional\[torch.Tensor\] = Nonedecoder\_input\_ids: typing.Optional\[torch.LongTensor\] = Nonedecoder\_attention\_mask: typing.Optional\[torch.LongTensor\] = Nonehead\_mask: typing.Optional\[torch.Tensor\] = Nonedecoder\_head\_mask: typing.Optional\[torch.Tensor\] = Nonecross\_attn\_head\_mask: typing.Optional\[torch.Tensor\] = Noneencoder\_outputs: typing.Optional\[typing.Tuple\[typing.Tuple\[torch.FloatTensor\]\]\] = Nonepast\_key\_values: typing.Optional\[typing.Tuple\[typing.Tuple\[torch.FloatTensor\]\]\] = Noneinputs\_embeds: typing.Optional\[torch.FloatTensor\] = Nonedecoder\_inputs\_embeds: typing.Optional\[torch.FloatTensor\] = Noneuse\_cache: typing.Optional\[bool\] = Noneoutput\_attentions: typing.Optional\[bool\] = Noneoutput\_hidden\_states: typing.Optional\[bool\] = Noneoutput\_router\_logits: typing.Optional\[bool\] = Nonereturn\_dict: typing.Optional\[bool\] = None ) → `transformers.modeling_outputs.Seq2SeqMoEModelOutput` or `tuple(torch.FloatTensor)`

The [NllbMoeModel](/docs/transformers/v4.30.0/en/model_doc/nllb-moe#transformers.NllbMoeModel) forward method, overrides the `__call__` special method.

Although the recipe for forward pass needs to be defined within this function, one should call the `Module` instance afterwards instead of this since the former takes care of running the pre and post processing steps while the latter silently ignores them.

The [NllbMoeModel](/docs/transformers/v4.30.0/en/model_doc/nllb-moe#transformers.NllbMoeModel) forward method, overrides the `__call__` special method.

Although the recipe for forward pass needs to be defined within this function, one should call the `Module` instance afterwards instead of this since the former takes care of running the pre and post processing steps while the latter silently ignores them.

Example:

```
>>> from transformers import AutoTokenizer, NllbMoeModel

>>> tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/random-nllb-moe-2-experts")
>>> model = SwitchTransformersModel.from_pretrained("hf-internal-testing/random-nllb-moe-2-experts")

>>> input_ids = tokenizer(
...     "Studies have been shown that owning a dog is good for you", return_tensors="pt"
... ).input_ids  
>>> decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids  

>>> 
>>> decoder_input_ids = model._shift_right(decoder_input_ids)

>>> 
>>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids)
>>> last_hidden_states = outputs.last_hidden_state```

## [](#transformers.NllbMoeForConditionalGeneration)NllbMoeForConditionalGeneration

### class transformers.NllbMoeForConditionalGeneration

[](#transformers.NllbMoeForConditionalGeneration)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/nllb_moe/modeling_nllb_moe.py#L1642)

( config: NllbMoeConfig )

Parameters

-   [](#transformers.NllbMoeForConditionalGeneration.config)**config** ([NllbMoeConfig](/docs/transformers/v4.30.0/en/model_doc/nllb-moe#transformers.NllbMoeConfig)) — Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [from\_pretrained()](/docs/transformers/v4.30.0/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) method to load the model weights.

The NllbMoe Model with a language modeling head. Can be used for summarization. This model inherits from [PreTrainedModel](/docs/transformers/v4.30.0/en/main_classes/model#transformers.PreTrainedModel). Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.)

This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior.

#### forward

[](#transformers.NllbMoeForConditionalGeneration.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/nllb_moe/modeling_nllb_moe.py#L1682)

( input\_ids: typing.Optional\[torch.LongTensor\] = Noneattention\_mask: typing.Optional\[torch.Tensor\] = Nonedecoder\_input\_ids: typing.Optional\[torch.LongTensor\] = Nonedecoder\_attention\_mask: typing.Optional\[torch.LongTensor\] = Nonehead\_mask: typing.Optional\[torch.Tensor\] = Nonedecoder\_head\_mask: typing.Optional\[torch.Tensor\] = Nonecross\_attn\_head\_mask: typing.Optional\[torch.Tensor\] = Noneencoder\_outputs: typing.Optional\[typing.Tuple\[typing.Tuple\[torch.FloatTensor\]\]\] = Nonepast\_key\_values: typing.Optional\[typing.Tuple\[typing.Tuple\[torch.FloatTensor\]\]\] = Noneinputs\_embeds: typing.Optional\[torch.FloatTensor\] = Nonedecoder\_inputs\_embeds: typing.Optional\[torch.FloatTensor\] = Nonelabels: typing.Optional\[torch.LongTensor\] = Noneuse\_cache: typing.Optional\[bool\] = Noneoutput\_attentions: typing.Optional\[bool\] = Noneoutput\_hidden\_states: typing.Optional\[bool\] = Noneoutput\_router\_logits: typing.Optional\[bool\] = Nonereturn\_dict: typing.Optional\[bool\] = None ) → `transformers.modeling_outputs.Seq2SeqMoEOutput` or `tuple(torch.FloatTensor)`

The [NllbMoeForConditionalGeneration](/docs/transformers/v4.30.0/en/model_doc/nllb-moe#transformers.NllbMoeForConditionalGeneration) forward method, overrides the `__call__` special method.

Although the recipe for forward pass needs to be defined within this function, one should call the `Module` instance afterwards instead of this since the former takes care of running the pre and post processing steps while the latter silently ignores them.

Translation example:

```
>>> from transformers import AutoTokenizer, NllbMoeForConditionalGeneration

>>> model = NllbMoeForConditionalGeneration.from_pretrained("facebook/nllb-moe-54b")
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-moe-54b")

>>> text_to_translate = "Life is like a box of chocolates"
>>> model_inputs = tokenizer(text_to_translate, return_tensors="pt")

>>> 
>>> gen_tokens = model.generate(**model_inputs, forced_bos_token_id=tokenizer.get_lang_id("eng_Latn"))
>>> print(tokenizer.batch_decode(gen_tokens, skip_special_tokens=True))```