#
Modular transformers

`transformers` is an opinionated framework; our philosophy is defined in the following conceptual guide.

The core of that philosophy is exemplified by the single model, single file aspect of the library. This component’s downside is that it limits the inheritance and importability of components from files to others in the toolkit.

As a result, model components tend to be repeated across many files. There are as many attention layers defined in `transformers` as there are models, and a significant number of those are identical to each other. The unfortunate consequence is that independent implementations tend to diverge as fixes and changes get applied to specific parts of the code.

In order to balance this issue, we introduced the concept of “copies” across the library. By adding a comment indicating that code is a copy of another, we can enforce through CI and local commands that copies do not diverge. However, while the complexity is low, this is often quite tedious to do.

And, finally, this contributes to adding a significant overhead to contributing models which we would like to remove. This approach often requires model contributions to add modeling code (~1k lines), processor (~500 lines), tests, docs, etc. Model contribution PRs rarely add less than 3-5k lines of code, with much of this code being boilerplate.

This raises the bar for contributions, and with Modular Transformers, we’re aiming to lower the bar to a much more acceptable point.

##
What is it?

Modular Transformers introduces the concept of a “modular” file to a model folder. This modular file accepts code that isn’t typically accepted in modeling/processing files, as it allows importing from neighbouring models as well as inheritance from classes to others.

This modular file defines models, processors, and the configuration class that would otherwise be defined in their respective modules.

Finally, this feature introduces a new `linter` which will “unravel” the modular file into the “single model, single file” directory structure. These files will get auto-generated every time the script is run; reducing the required contributions to the modular file, and therefore only to the changes between the contributed model and others.

Model users will end up importing and using the single-file interface, so no change is expected here. Doing this, we hope to combine the best of both worlds: enabling simple contributions while sticking to our philosophy.

This is therefore a replacement for the `# Copied from` markers, and previously contributed models can be expected to be moved to the new Modular Transformers format in the coming months.

###
Details

The “linter”, which unravels the inheritance and creates all single-files from the modular file, will flatten the inheritance while trying to be invisible to Python users. At this time, the linter flattens a single level of inheritance.

For example:

- If a configuration class inherits from another and adds/deletes an argument, the generated file will either directly reference it (in case of addition) or completely remove it (in case of deletion).
- If a class inherits from another, for example: class GemmaModel(LlamaModel):, dependencies are automatically inferred. All submodules will be automatically inferred from the superclass.
- If you define new functions in the `modular` and use them inside classes, the linter will automatically infer the

You should be able to write everything (the tokenizer, the image processor, the model, the config) in this `modular` file, and the corresponding files will be created for you.

###
Enforcement

[TODO] We are introducing a new test, that makes sure the generated content matches what is present in the `modular_xxxx.py`

###
Examples

Here is a quick example with BERT and RoBERTa. The two models are intimately related: their modeling implementation differs solely by a change in the embedding layer.

Instead of redefining the model entirely, here is what the `modular_roberta.py` file looks like for the modeling & configuration classes (for the sake of the example, the tokenizer is ignored at this time as very different).

    from torch import nn
    from ..bert.configuration_bert import BertConfig
    from ..bert.modeling_bert import (
        BertModel,
        BertEmbeddings,
        BertForMaskedLM
    )

    # The RoBERTa config is identical to BERT's config
    class RobertaConfig(BertConfig):
      model_type = 'roberta'

    # We redefine the embeddings here to highlight the padding ID difference, and we redefine the position embeddings
    class RobertaEmbeddings(BertEmbeddings):
        def __init__(self, config):
            super().__init__(config())

            self.padding_idx = config.pad_token_id
            self.position_embeddings = nn.Embedding(
                config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
            )

    # The RoBERTa model is identical to the BERT model, except for the embedding layer.
    # We redefine the embeddings above, so here there is no need to do additional work
    class RobertaModel(BertModel):
      def __init__(self, config):
        super().__init__(config)
        self.embeddings = RobertaEmbeddings(config)

    # The heads now only need to redefine the model inside to the correct `RobertaModel`
    class RobertaForMaskedLM(BertForMaskedLM):
      def __init__(self, config):
        super().__init__(config)
        self.model = RobertaModel(config)

Note that if you do not use the dependency that you defined, you will have the following error:

    ValueError: You defined `RobertaEmbeddings` in the modular_roberta.py, it should be used
                                        when you define `BertModel`, as it is one of it's direct dependencies. Make sure
                                        you use it in the `__init__` function.

Additionally, you may find a list of examples here:

##
What it is not

It is not a replacement for the modeling code (yet?), and if your model is not based on anything else that ever existed, then you can add a `modeling` file as usual.

##
Advanced usage

###
Removing attributes and functions

To remove attributes that are not used in your modular model, and that you don’t want to see in the unravelled modeling:

    class GemmaModel(LlamaModel):                 |           class GemmaModel(PreTrainedModel):
        def __init__(self, config):               |              def __init__(self, config):
            super().__init__(self, eos_token)     |                 super().__init__(config)
            del self.embed_tokens                 |                 self.padding_idx = config.pad_token_id
                                                  |                 self.vocab_size = config.vocab_size
                                                  |
                                                  |                 self.layers = nn.ModuleList(
                                                  |                     [LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
                                                  |                 )
                                                  |                 self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
                                                  |                 self.rotary_emb = LlamaRotaryEmbedding(config=config)
                                                  |                 self.gradient_checkpointing = False
                                                  |
                                                  |                 # Initialize weights and apply final processing
                                                  |                 self.post_init()

If you check the original `LlamaModel`, it has a `embed_tokens` which was removed here (as you would expect!)

Removing a function is pretty similar, you just need to write it with a `raise ValueError("")` to mimick the behaviour you actually want when you remove a parent function in python.

    class GemmaTokenizer(LlamaTokenizer):
        ...

        def get_spm_processor(self):
            raise AttributeError("Not needed for Gemma")

        def unk_token_length(self):
            raise AttributeError("Not needed for Gemma")

###
Define new functions

If you define a new function in the `modular` file to be used inside a class, say

    def my_new_function(*args, **kwargs):
      # Do something here
      pass

    class GemmaModel(LlamaModel):
        def forward(*args, **kwargs):
          # Call the function
          example = my_new_function(*args, **kwargs)
          # continue here

the `my_new_function` function (and, recursively, any other new functions called in its body) will be automatically copy-pasted in the file where it is used.

###
Calling super()

We recently shipped a few features that allow you to go from:

    class GemmaTokenizer(LlamaTokenizer, PretrainedTokenizerFast):         |           class GemmaModel(nn.Module):
        def __init__(self, eos_token="</s>"):                              |             def __init__(self):
            eos_token = AddedToken(eos_token)                              |                eos_token = AddedToken(eos_token)
            PretrainedTokenizerFast.__init__(self, eos_token)              |                super().__init__(eos_token)

This is useful want you don’t want to unravel the call to `super()`, and you want to differentiate which super init call you are doing!

###
Special naming

We now also support special cases like

    class GemmaVisionModel(CLIPModel):
        pass

where the name of your class `GemmaVision` is not the same as the modular `Gemma`. This is super useful for composite models.

< > Update on GitHub