SwitchTransformers

## [](#overview)Overview

The SwitchTransformers model was proposed in [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by William Fedus, Barret Zoph, Noam Shazeer.

The Switch Transformer model uses a sparse T5 encoder-decoder architecture, where the MLP are replaced by a Mixture of Experts (MoE). A routing mechanism (top 1 in this case) associates each token to one of the expert, where each expert is a dense MLP. While switch transformers have a lot more weights than their equivalent dense models, the sparsity allows better scaling and better finetuning performance at scale. During a forward pass, only a fraction of the weights are used. The routing mechanism allows the model to select relevant weights on the fly which increases the model capacity without increasing the number of operations.

The abstract from the paper is the following:

_In deep learning, models typically reuse the same parameters for all inputs. Mixture of Experts (MoE) defies this and instead selects different parameters for each incoming example. The result is a sparsely-activated model — with outrageous numbers of parameters — but a constant computational cost. However, despite several notable successes of MoE, widespread adoption has been hindered by complexity, communication costs and training instability — we address these with the Switch Transformer. We simplify the MoE routing algorithm and design intuitive improved models with reduced communication and computational costs. Our proposed training techniques help wrangle the instabilities and we show large sparse models may be trained, for the first time, with lower precision (bfloat16) formats. We design models based off T5-Base and T5-Large to obtain up to 7x increases in pre-training speed with the same computational resources. These improvements extend into multilingual settings where we measure gains over the mT5-Base version across all 101 languages. Finally, we advance the current scale of language models by pre-training up to trillion parameter models on the “Colossal Clean Crawled Corpus” and achieve a 4x speedup over the T5-XXL model._

Tips:

-   SwitchTransformers uses the [T5Tokenizer](/docs/transformers/v4.30.0/en/model_doc/t5#transformers.T5Tokenizer), which can be loaded directly from each model’s repository.
-   The released weights are pretrained on English [Masked Language Modeling](https://moon-ci-docs.huggingface.co/docs/transformers/pr_19323/en/glossary#general-terms) task, and should be finetuned.

This model was contributed by [Younes Belkada](https://huggingface.co/ybelkada) and [Arthur Zucker](https://huggingface.co/ArtZucker) . The original code can be found [here](https://github.com/google/flaxformer/tree/main/flaxformer/architectures/moe).

## [](#resources)Resources

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

## [](#transformers.SwitchTransformersConfig)SwitchTransformersConfig

### class transformers.SwitchTransformersConfig

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

( vocab\_size = 32128d\_model = 768d\_kv = 64d\_ff = 2048expert\_capacity = 64num\_layers = 12num\_sparse\_encoder\_layers = 3num\_decoder\_layers = 12num\_sparse\_decoder\_layers = 3num\_heads = 12num\_experts = 8router\_bias = Falserouter\_jitter\_noise = 0.01router\_dtype = 'float32'router\_ignore\_padding\_tokens = Falserelative\_attention\_num\_buckets = 32relative\_attention\_max\_distance = 128dropout\_rate = 0.1layer\_norm\_epsilon = 1e-06router\_z\_loss\_coef = 0.001router\_aux\_loss\_coef = 0.001initializer\_factor = 1.0feed\_forward\_proj = 'relu'is\_encoder\_decoder = Trueadd\_router\_probs = Falseuse\_cache = Truepad\_token\_id = 0eos\_token\_id = 1\*\*kwargs )

This is the configuration class to store the configuration of a [SwitchTransformersModel](/docs/transformers/v4.30.0/en/model_doc/switch_transformers#transformers.SwitchTransformersModel). It is used to instantiate a SwitchTransformers 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 SwitchTransformers [google/switch-base-8](https://huggingface.co/google/switch-base-8) 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.

## [](#transformers.SwitchTransformersTop1Router)SwitchTransformersTop1Router

### class transformers.SwitchTransformersTop1Router

[](#transformers.SwitchTransformersTop1Router)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L131)

( config: SwitchTransformersConfig )

Router using tokens choose top-1 experts assignment.

This router uses the same mechanism as in Switch Transformer ([https://arxiv.org/abs/2101.03961](https://arxiv.org/abs/2101.03961)) and V-MoE ([https://arxiv.org/abs/2106.05974](https://arxiv.org/abs/2106.05974)): tokens choose their top experts. 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.

#### \_compute\_router\_probabilities

[](#transformers.SwitchTransformersTop1Router._compute_router_probabilities)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L151)

( hidden\_states: Tensor ) → router\_probabilities (`torch.Tensor`)

Parameters

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

Returns

router\_probabilities (`torch.Tensor`)

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`): Logits tensor of shape (batch\_size, sequence\_length, num\_experts) corresponding to raw router logits. This is used later for computing router z-loss.

Computes router probabilities from input hidden states.

#### forward

[](#transformers.SwitchTransformersTop1Router.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L201)

( hidden\_states: Tensor )

Parameters

-   [](#transformers.SwitchTransformersTop1Router.forward.hidden_states)**hidden\_states** (`torch.Tensor`) — \[num\_groups, tokens\_per\_group, hidden\_dim\] inputs to send to experts.

Generic forward function for every Router class. Each Router expects to have the same input hidden states (`hidden_states`) corresponding to the hidden states for each token, the `expert_capacity` corresponding to the number of tokens the Router will send to each expert, some Routers can send up to few tokens to each expert.

Each Router works as the following: it expects the hidden states for each token, gets the `router_probs` and `router_logits` from the `router_weights`. This will assign for each token, the raw probability to be assigned to an expert. Then each Router class will have to define its own `_compute_routing_instructions`.

## [](#transformers.SwitchTransformersSparseMLP)SwitchTransformersSparseMLP

( config: SwitchTransformersConfigexpert\_class: Module = <class 'transformers.models.switch\_transformers.modeling\_switch\_transformers.SwitchTransformersDenseActDense'> )

Implementation of the Switch Transformers Sparse MLP module.

Hold on, this will be slightly tricky to understand In the correct order, a MoE layer does the following:

1- Gets the `router_mask` from the router. The shape of the mask is `(batch_size, sequence_length, num_expert)` and corresponds to the argmax of the `router_probs`. The probabilities are needed in the computation of the hidden states : they are broadcasted to the hidden states values (can be interpreted as a scaling factor).

2- Dispatch the tokens to its associated experts. We do a classic for loop over the experts and assign for each expert the corresponding hidden states.

## [](#transformers.SwitchTransformersModel)SwitchTransformersModel

### class transformers.SwitchTransformersModel

[](#transformers.SwitchTransformersModel)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L1339)

( config: SwitchTransformersConfig )

Parameters

-   [](#transformers.SwitchTransformersModel.config)**config** ([SwitchTransformersConfig](/docs/transformers/v4.30.0/en/model_doc/switch_transformers#transformers.SwitchTransformersConfig)) — 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 SWITCH\_TRANSFORMERS Model transformer outputting raw hidden-states without any specific head on top.

The SWITCH\_TRANSFORMERS model was proposed in [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by [William Fedus](https://arxiv.org/search/cs?searchtype=author&query=Fedus%2C+W), [Barret Zoph](https://arxiv.org/search/cs?searchtype=author&query=Zoph%2C+B), and [Noam Shazeer](https://arxiv.org/search/cs?searchtype=author&query=Shazeer%2C+N). It’s an encoder-decoder T5-like model with sparse Feed Forward that stands for Mixture of Experts (MoE) architecture.

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.SwitchTransformersModel.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L1385)

( input\_ids: typing.Optional\[torch.LongTensor\] = Noneattention\_mask: typing.Optional\[torch.FloatTensor\] = Nonedecoder\_input\_ids: typing.Optional\[torch.LongTensor\] = Nonedecoder\_attention\_mask: typing.Optional\[torch.BoolTensor\] = Nonehead\_mask: typing.Optional\[torch.FloatTensor\] = Nonedecoder\_head\_mask: typing.Optional\[torch.FloatTensor\] = 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.Tensor\] = Nonedecoder\_inputs\_embeds: typing.Optional\[torch.Tensor\] = 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 [SwitchTransformersModel](/docs/transformers/v4.30.0/en/model_doc/switch_transformers#transformers.SwitchTransformersModel) 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, SwitchTransformersModel

>>> tokenizer = AutoTokenizer.from_pretrained("google/switch-base-8")
>>> model = SwitchTransformersModel.from_pretrained("google/switch-base-8")

>>> 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.SwitchTransformersForConditionalGeneration)SwitchTransformersForConditionalGeneration

### class transformers.SwitchTransformersForConditionalGeneration

[](#transformers.SwitchTransformersForConditionalGeneration)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L1507)

( config: SwitchTransformersConfig )

Parameters

-   [](#transformers.SwitchTransformersForConditionalGeneration.config)**config** ([SwitchTransformersConfig](/docs/transformers/v4.30.0/en/model_doc/switch_transformers#transformers.SwitchTransformersConfig)) — 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.

SWITCH\_TRANSFORMERS Model with a `language modeling` head on top.

The SWITCH\_TRANSFORMERS model was proposed in [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by [William Fedus](https://arxiv.org/search/cs?searchtype=author&query=Fedus%2C+W), [Barret Zoph](https://arxiv.org/search/cs?searchtype=author&query=Zoph%2C+B), and [Noam Shazeer](https://arxiv.org/search/cs?searchtype=author&query=Shazeer%2C+N). It’s an encoder-decoder T5-like model with sparse Feed Forward that stands for Mixture of Experts (MoE) architecture.

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.SwitchTransformersForConditionalGeneration.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L1563)

( input\_ids: typing.Optional\[torch.LongTensor\] = Noneattention\_mask: typing.Optional\[torch.FloatTensor\] = Nonedecoder\_input\_ids: typing.Optional\[torch.LongTensor\] = Nonedecoder\_attention\_mask: typing.Optional\[torch.BoolTensor\] = Nonehead\_mask: typing.Optional\[torch.FloatTensor\] = Nonedecoder\_head\_mask: typing.Optional\[torch.FloatTensor\] = Nonecross\_attn\_head\_mask: typing.Optional\[torch.Tensor\] = Noneencoder\_outputs: typing.Optional\[typing.Tuple\[typing.Tuple\[torch.Tensor\]\]\] = Nonepast\_key\_values: typing.Optional\[typing.Tuple\[typing.Tuple\[torch.Tensor\]\]\] = 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\] = Truereturn\_dict: typing.Optional\[bool\] = None ) → `transformers.modeling_outputs.Seq2SeqMoEOutput` or `tuple(torch.FloatTensor)`

The [SwitchTransformersForConditionalGeneration](/docs/transformers/v4.30.0/en/model_doc/switch_transformers#transformers.SwitchTransformersForConditionalGeneration) 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.

Examples:

```
>>> from transformers import AutoTokenizer, SwitchTransformersForConditionalGeneration

>>> tokenizer = AutoTokenizer.from_pretrained("google/switch-base-8")
>>> model = SwitchTransformersForConditionalGeneration.from_pretrained("google/switch-base-8")

>>> 
>>> input_ids = tokenizer("The <extra_id_0> walks in <extra_id_1> park", return_tensors="pt").input_ids
>>> labels = tokenizer("<extra_id_0> cute dog <extra_id_1> the <extra_id_2>", return_tensors="pt").input_ids
>>> outputs = model(input_ids=input_ids, labels=labels)
>>> loss = outputs.loss
>>> logits = outputs.logits

>>> 
>>> input_ids = tokenizer(
...     "summarize: studies have shown that owning a dog is good for you", return_tensors="pt"
... ).input_ids  
>>> outputs = model.generate(input_ids)
>>> 
>>> ```

## [](#transformers.SwitchTransformersEncoderModel)SwitchTransformersEncoderModel

### class transformers.SwitchTransformersEncoderModel

[](#transformers.SwitchTransformersEncoderModel)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L1824)

( config: SwitchTransformersConfig )

Parameters

-   [](#transformers.SwitchTransformersEncoderModel.config)**config** ([SwitchTransformersConfig](/docs/transformers/v4.30.0/en/model_doc/switch_transformers#transformers.SwitchTransformersConfig)) — 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 SWITCH\_TRANSFORMERS Model transformer outputting encoder’s raw hidden-states without any specific head on top.

The SWITCH\_TRANSFORMERS model was proposed in [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://arxiv.org/abs/2101.03961) by [William Fedus](https://arxiv.org/search/cs?searchtype=author&query=Fedus%2C+W), [Barret Zoph](https://arxiv.org/search/cs?searchtype=author&query=Zoph%2C+B), and [Noam Shazeer](https://arxiv.org/search/cs?searchtype=author&query=Shazeer%2C+N). It’s an encoder-decoder T5-like model with sparse Feed Forward that stands for Mixture of Experts (MoE) architecture.

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.SwitchTransformersEncoderModel.forward)[< source \>](https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L1860)

( input\_ids: typing.Optional\[torch.LongTensor\] = Noneattention\_mask: typing.Optional\[torch.FloatTensor\] = Nonehead\_mask: typing.Optional\[torch.FloatTensor\] = Noneinputs\_embeds: typing.Optional\[torch.FloatTensor\] = Noneoutput\_attentions: typing.Optional\[bool\] = Noneoutput\_hidden\_states: typing.Optional\[bool\] = Noneoutput\_router\_logits: typing.Optional\[bool\] = Truereturn\_dict: typing.Optional\[bool\] = None ) → `transformers.modeling_outputs.MoEModelOutput` or `tuple(torch.FloatTensor)`

The [SwitchTransformersEncoderModel](/docs/transformers/v4.30.0/en/model_doc/switch_transformers#transformers.SwitchTransformersEncoderModel) 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, SwitchTransformersEncoderModel

>>> tokenizer = AutoTokenizer.from_pretrained("google/switch-base-8")
>>> model = SwitchTransformersEncoderModel.from_pretrained("google/switch-base-8")
>>> input_ids = tokenizer(
...     "Studies have been shown that owning a dog is good for you", return_tensors="pt"
... ).input_ids  
>>> outputs = model(input_ids=input_ids)
>>> last_hidden_states = outputs.last_hidden_state```