Добавление модели с модульными трансформерами

Modular transformers сокращает объём кода, необходимого для добавления модели, за счёт использования импортов и наследования, в отличие от политики одна модель — один файл. Вместо повторения компонентов модели в разных файлах добавьте файл modular в папку вашей модели и наследуйтесь от существующих классов.

Конвертер генерирует автономные файлы из модульного файла. Пользователи получают тот же самый интерфейс с одним файлом, который они уже знают.

Модульные трансформеры не предназначены для замены устаревшего кода моделирования. Если ваша модель не основана на существующей, добавьте файл modeling.py вручную. То же самое касается файлов конфигурации, токенизации или обработки, которые не могут чисто наследоваться от похожего файла.

Единого правильного порядка здесь тоже нет. Некоторые участники сначала пишут модульный файл и генерируют из него остальное. Другие начинают с вручную написанного modeling.py и позже рефакторят его в модульный файл. Оба подхода работают.

Реализация модульного файла

Начните с поиска модели в Transformers, похожей на вашу. Хорошей отправной точкой могут служить Mistral, Qwen2, Cohere и Cohere2, а также Llama. В таблице ниже приведено соответствие между общими компонентами и моделями, от которых можно наследовать.

Используйте инструмент modular-detector-v2, чтобы найти существующие реализации, от которых можно наследовать. Вставьте фрагмент кода, и он вернёт наиболее похожие методы, уже присутствующие в Transformers, что позволит определить оптимальный родительский класс до начала написания кода.

Не изменяйте существующую модель только ради того, чтобы наследование работало для вашей новой. Если переименование или создание подкласса родительского класса слишком неудобно, просто скопируйте нужный код напрямую.

Создайте src/transformers/models/<name>/modular_<name>.py, где <name> совпадает с именем каталога модели в формате snake_case. В этом разделе показано, как реализовать Olmo2 на основе Olmo с использованием модульного подхода (см. исходный файл modular_olmo2.py).

Конфигурация

Есть два отличия Olmo2Config от OlmoConfig.

1. Добавлен новый аргумент rms_norm_eps.
2. Аргумент clip_qkv больше не используется.

Объявляйте новые аргументы как аннотации типов на уровне класса со значением по умолчанию. Для удалённых аргументов присвойте AttributeError(), чтобы подавить унаследованный атрибут в сгенерированном файле (см. Удаление атрибутов).

- @auto_docstring(checkpoint="allenai/OLMo-7B-hf")
+ @auto_docstring(checkpoint="allenai/Olmo2-7B-1124-hf")
+ @strict
- class OlmoConfig(PreTrainedConfig):
+ class Olmo2Config(OlmoConfig):
      ...
-     model_type = "olmo"
+     model_type = "olmo2"
      ...
+     rms_norm_eps: float = 1e-5
-     clip_qkv: float | None = None
+     clip_qkv = AttributeError()

@auto_docstring автоматически генерирует стандартные описания аргументов (см. руководство по @auto_docstring). @strict отклоняет неизвестные kwargs при инициализации, позволяя на раннем этапе выявлять опечатки и устаревшие аргументы. Добавляйте оба декоратора в каждый класс конфигурации, так как они не наследуются от родительского класса. Объявляйте их явно, даже если родительский класс конфигурации уже их содержит.

Чтобы задать производный атрибут или обработать логику обратной совместимости, используйте __post_init__ вместо __init__. Например, Cohere2 вычисляет head_dim и выводит layer_types при инициализации.

def __post_init__(self, **kwargs):
    if self.num_key_value_heads is None:
        self.num_key_value_heads = self.num_attention_heads
    self.head_dim = self.hidden_size // self.num_attention_heads
    super().__post_init__(**kwargs)

Для моделей с поддержкой тензорного или конвейерного параллелизма определите base_model_tp_plan и base_model_pp_plan как словари на уровне класса в конфигурации. Оба словаря определяют, как разбить модель на части для распределения по устройствам. Примеры можно найти в существующих конфигурациях, например Olmo2 или Cohere2.

class MyNewModelConfig(PreTrainedConfig):
    model_type = "my_new_model"

    # Tensor parallelism: maps layer name patterns to sharding strategies.
    # Use "colwise" / "rowwise" for standard sharding, or the "gather_output" /
    # "split_input" variants when an extra op (e.g. a QK norm) prevents fusing.
    base_model_tp_plan = {
        "layers.*.self_attn.q_proj": "colwise",
        "layers.*.self_attn.k_proj": "colwise",
        "layers.*.self_attn.v_proj": "colwise",
        "layers.*.self_attn.o_proj": "rowwise",
        "layers.*.mlp.gate_proj": "colwise",
        "layers.*.mlp.up_proj": "colwise",
        "layers.*.mlp.down_proj": "rowwise",
    }

    # Pipeline parallelism: maps submodule names to their (input, output) tensor names.
    base_model_pp_plan = {
        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
        "norm": (["hidden_states"], ["hidden_states"]),
    }

Норма

Чтобы скопировать родительский класс без изменений, унаследуйте его с помощью pass. Линтер копирует содержимое родительского класса и переименовывает все ссылки, чтобы они соответствовали новой модели.

from ..olmo.modeling_olmo import OlmoRotaryEmbedding

class Olmo2RotaryEmbedding(OlmoRotaryEmbedding):
    pass

Чтобы изменить конкретное поведение, наследуйте и переопределяйте только то, что отличается. Olmo2RMSNorm отличается от LlamaRMSNorm одной строкой. Умножение выполняется до приведения обратно к типу данных входных данных, а не после.

  from ..llama.modeling_llama import LlamaRMSNorm

  class Olmo2RMSNorm(LlamaRMSNorm):
      def forward(self, hidden_states):
          input_dtype = hidden_states.dtype
          hidden_states = hidden_states.to(torch.float32)
          variance = hidden_states.pow(2).mean(-1, keepdim=True)
          hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-         return self.weight * hidden_states.to(input_dtype)
+         return (self.weight * hidden_states).to(input_dtype)

Внимание

Внимание в Olmo2 идентично вниманию в Olmo, за исключением того, что к запросам и ключам применяется RMSNorm, а обрезание qkv удалено. Вызов super().__init__(...) копирует тело родительского класса и добавляет две новые строки нормализации. Метод forward полностью переопределён, поскольку теперь запросы и ключи проходят через нормализацию перед проекцией. Линтер также переносит в сгенерированный файл все импортированные функции, включая apply_rotary_pos_emb, eager_attention_forward и их зависимости.

  class Olmo2Attention(OlmoAttention):
      def __init__(self, config: Olmo2Config, layer_idx: int | None = None):
          super().__init__(config, layer_idx=layer_idx)
+         self.q_norm = Olmo2RMSNorm(config.num_attention_heads * self.head_dim, config.rms_norm_eps)
+         self.k_norm = Olmo2RMSNorm(config.num_key_value_heads * self.head_dim, config.rms_norm_eps)

      def forward(self, ...):
          ...
-         query_states = self.q_proj(hidden_states)
-         key_states = self.k_proj(hidden_states)
+         query_states = self.q_norm(self.q_proj(hidden_states))
+         key_states = self.k_norm(self.k_proj(hidden_states))
          value_states = self.v_proj(hidden_states)

-         if self.config.clip_qkv is not None:
-             query_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
-             key_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
-             value_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
-
          ...

DecoderLayer

После super().__init__(...) перезапишите атрибуты нормализации экземплярами Olmo2RMSNorm и переназначьте self.self_attn на новый класс Olmo2Attention. Оператор del self.input_layernorm удаляет присваивание input_layernorm из родительского класса, поскольку Olmo2 применяет нормализацию после внимания, а не до него. Подробности о том, что удаляет, а что оставляет оператор del, см. в разделе Удаление атрибутов.

Метод forward переписан с учётом нового расположения нормализации после внимания. Переписывание forward требуется только при переименовании атрибута, а не при изменении только его типа.

  class Olmo2DecoderLayer(OlmoDecoderLayer):
      def __init__(self, config: Olmo2Config, layer_idx: int):
          super().__init__(config, layer_idx=layer_idx)
-         self.self_attn = OlmoAttention(config=config, layer_idx=layer_idx)
-         self.input_layernorm = OlmoLayerNorm(config.hidden_size)
-         self.post_attention_layernorm = OlmoLayerNorm(config.hidden_size)
+         self.self_attn = Olmo2Attention(config=config, layer_idx=layer_idx)
+         del self.input_layernorm
+         self.post_attention_layernorm = Olmo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+         self.post_feedforward_layernorm = Olmo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

      def forward(self, ...):
          residual = hidden_states
-         hidden_states = self.input_layernorm(hidden_states)
          # Self Attention
          hidden_states, _ = self.self_attn(...)
-         hidden_states = residual + hidden_states
+         hidden_states = self.post_attention_layernorm(hidden_states)
+         hidden_states = residual + hidden_states

          # Fully Connected
          residual = hidden_states
-         hidden_states = self.post_attention_layernorm(hidden_states)
          hidden_states = self.mlp(hidden_states)
-         hidden_states = residual + hidden_states
+         hidden_states = self.post_feedforward_layernorm(hidden_states)
+         hidden_states = residual + hidden_states
          return hidden_states

Модель

Здесь меняется только тип self.norm. Метод forward идентичен методу родительского класса, поэтому линтер переносит его автоматически.

  class Olmo2Model(OlmoModel):
      def __init__(self, config: Olmo2Config):
          super().__init__(config)
-         self.layers = nn.ModuleList(
-             [OlmoDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-         )
-         self.norm = OlmoLayerNorm(config.hidden_size)
+         self.norm = Olmo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+         self.layers = nn.ModuleList(
+             [Olmo2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+         )

Голова модели

Логика идентична логике OlmoForCausalLM, поэтому изменения не требуются.

from ..olmo.modeling_olmo import OlmoForCausalLM

class Olmo2ForCausalLM(OlmoForCausalLM):
    pass

Другие классы

Файл modeling_olmo2.py, сгенерированный линтером, также содержит классы (Olmo2MLP, Olmo2RotaryEmbedding, Olmo2PreTrainedModel), которые не были явно определены в modular_olmo2.py.

Линтер подтягивает любой класс, от которого зависит наследуемый класс, если вы не переопределяете его явно. Импортированные функции, такие как apply_rotary_pos_emb, подчиняются тому же правилу.

Например, OlmoDecoderLayer содержит self.mlp = OlmoMLP(config). Olmo2MLP не был определён в модульном файле, поэтому линтер создаёт его автоматически, что эквивалентно использованию pass.

from ..olmo.modeling_olmo import OlmoMLP

class Olmo2MLP(OlmoMLP):
    pass

Если вы хотите, чтобы Olmo2MLP наследовался от другой модели, укажите это явно.

# switch to Mistral definition
from ..mistral.modeling_mistral import MistralMLP

class Olmo2MLP(MistralMLP):
    pass

Завершение файла

Каждый модульный файл должен объявлять logger и список __all__ на уровне модуля.

logger = logging.get_logger(__name__)

__all__ = [
    "Olmo2Config",
    "Olmo2ForCausalLM",
    "Olmo2Model",
    "Olmo2PreTrainedModel",
]

__all__ должен перечислять все публичные классы в файле. Конвертер и последующие импорты зависят от него. Класс, отсутствующий в __all__, не будет экспортирован корректно.

Сгенерировать файлы моделирования

Скрипт modular_model_converter.py генерирует автономные файлы modeling.py, configuration.py и другие на основе вашего модульного файла. Для каждого наследуемого класса он копирует тело родительского класса в дочерний, переименовывает все ссылки в соответствии с новой моделью и подтягивает вспомогательные функции или классы, от которых зависят эти родительские классы.

В выходных файлах отсутствуют импорты между моделями и наследование из других каталогов моделей. Linter сворачивает наследование до одного уровня. Если Olmo2Attention наследуется от OlmoAttention, то сгенерированный Olmo2Attention является полностью самодостаточным. Однако, если OlmoAttention сам наследовался от чего-то другого, linter не инлайнит этого «прадеда».

Выполните команду ниже, чтобы сгенерировать файлы из модульного файла.

python utils/modular_model_converter.py your_model

Никогда не редактируйте сгенерированные файлы напрямую, так как любые изменения будут перезаписаны при следующем запуске.

Шаблоны для модульных файлов

Разделы ниже описывают типичные сценарии использования, такие как удаление атрибутов или переопределение декорированных методов, при работе с модульным файлом.

Удаление атрибутов

Способ удаления унаследованного атрибута зависит от того, работаете ли вы с классом конфигурации или с подклассом nn.Module.

Для класса конфигурации присвойте AttributeError() атрибуту на уровне класса.

class MyNewConfig(ParentConfig):
    removed_attr = AttributeError()

Инструмент проверки кода полностью удаляет объявление атрибута из сгенерированного файла конфигурации. Классы конфигурации используют структуру в стиле dataclass без __init__, поэтому присваивание AttributeError() на уровне класса является правильным подходом.

Для подкласса nn.Module используйте del self.attribute после super().__init__(...).

class MyNewModel(ParentModel):
    def __init__(self, config: MyNewConfig):
        super().__init__(config)
        del self.attribute

del self.attribute удаляет только строку присваивания self.attribute = ... из скопированного тела родительского класса. Она не удаляет другие строки, в которых используется self.attribute. Если в методе forward родителя или в других методах также есть ссылки на этот атрибут, переопределите и эти методы.

class DummyModel(nn.Module):
    def __init__(self, config: DummyConfig):
        super().__init__()
        self.attribute = config.attribute
        if self.attribute:
            # do more stuff with `self.attribute` here
            ...

class MyNewDummyModel(DummyModel):
    def __init__(self, config: MyNewDummyConfig):
        super().__init__(config)
        del self.attribute
        # 'self.attribute = config.attribute' is removed, but the 'if self.attribute:' block remains.
        # Override forward() or any other method that references self.attribute.

Работа с super()

super().__init__(config) указывает конвертеру скопировать тело родительского класса в дочерний. Два паттерна позволяют переопределить это поведение.

• Вызывайте конкретный родительский класс напрямую, если сгенерированный вывод должен вызывать класс-прадедушку (nn.Module.__init__), а не модульный родительский класс.
• Используйте **super_kwargs, чтобы унаследовать полную сигнатуру родительского метода, добавив при этом собственный docstring или заменив декоратор.

Вызов класса-предка напрямую

Будьте явны в том, какой класс вы вызываете, когда вам нужно, чтобы super() обращался к родительскому классу сгенерированного класса, а не к родительскому модульному классу. В примере ниже напрямую вызывается nn.Module.__init__(self). DummyModule сам по себе является nn.Module, поэтому конвертер записывает его как super().__init__() в сгенерированном MyNewDummyModule.

class MyNewDummyModule(DummyModule):                   |     class MyNewDummyModule(nn.Module):
                                                       |
  def __init__(self):                                  |       def __init__(self):
    nn.Module.__init__(self)                           |         super().__init__()
    self.foo = config.foo                              |         self.foo = config.foo
    ...                                                |         ...

super_kwargs

Используйте **super_kwargs, чтобы унаследовать полную сигнатуру родительского метода, добавив при этом собственный docstring или заменив декоратор. В переопределённой сигнатуре этот параметр указывает линтеру развернуть все аргументы родительского метода в сгенерированном результате.

Наиболее частое применение — добавление docstring, специфичного для модели, например, документирование аргумента labels, без переписывания полной сигнатуры.

# modular_gemma.py
class GemmaForCausalLM(LlamaForCausalLM):
    def forward(**super_kwargs):
        r"""
        Example:

        ```python
        >>> from transformers import AutoTokenizer, GemmaForCausalLM
        >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b")
        ...
        ```"""
        return super().forward(**super_kwargs)

Сгенерированный GemmaForCausalLM.forward имеет полную сигнатуру LlamaForCausalLM, поэтому ручное копирование не требуется.

**super_kwargs — это сокращение для узкоспециализированных случаев. Если вы изменяете поведение, лучше указать полную сигнатуру.

Удаление неиспользуемых методов

Удалите родительский метод, переопределив его с помощью оператора raise AttributeError(""). Линтер удалит метод из сгенерированного файла.

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")

Переопределение декорированных методов

При переопределении метода родительского класса, имеющего декоратор, декоратор родителя автоматически переносится на переопределённый метод. Если вы добавите собственный декоратор, он заменит декоратор родителя.

В библиотеке используются два декоратора: один для захвата промежуточных выходов модели, а другой — для автоматической генерации docstrings.

В примере ниже подкласс переопределяет декорированный метод родительского класса, не добавляя собственный декоратор. Декоратор родительского метода переносится на переопределённый метод.

class NewModel(DummyModel):       |   class NewModel(nn.Module):
  ...                             |     ...
                                  |
  def forward(...):               |     @decorator(...)
    ...                           |     def forward(...):
                                  |       ...

Если вы добавите новый декоратор, он заменит декоратор родительского класса.

class NewModel(DummyModel):       |   class NewModel(nn.Module):
  ...                             |     ...
                                  |
  @my_new_decorator(...)          |     @my_new_decorator(...)
  def forward(...):               |     def forward(...):
    ...                           |       ...

Особое именование

Инструмент проверки кода автоматически переименовывает всё при наследовании от класса. Используйте одинаковый префикс имени класса для всех классов в одном файле.

Избегайте смешивания префиксов, как в примере ниже. MyModelIncredibleMLP нарушает соглашения об именовании, и линтер не сможет определить, использовать ли MyModelIncredible или MyModel при переименовании зависимостей более высокого уровня.

class MyModelIncredibleMLP(LlamaMLP):
    ...

class MyModelDecoderLayer(LlamaDecoderLayer):
    ...

Благодаря отсутствию неявных зависимостей можно переименовать один класс локально. Явно переопределите каждое другое упоминание этого класса с новым именем. В противном случае линтер добавит нежелательный класс MyModelMLP рядом с MyModelIncredibleMLP.

Линтер выдает предупреждение, если обнаруживает неоднозначный префикс.

We detected multiple prefix names when inheriting from transformers.models.llama.modeling_llama: ('Emu3Text', 'Emu3'). We will only use the most used 'Emu3' prefix when grabbing args and dependencies. Make sure to subclass the intermediate classes with the prefix you want (if different from 'Emu3') or use a single prefix in all the modular (best).

Двусмысленные префиксы чаще всего встречаются в мультимодальных моделях, где имена классов содержат модальный квалификатор, например Text. Чтобы задать конкретный префикс для зависимости, явно переименуйте её с помощью pass.

class Emu3TextMLP(LlamaMLP):
    pass

Строки документации конфигурации

Линтер пока не поддерживает частичное наследование docstring. При добавлении или удалении атрибутов конфигурации добавляйте полный docstring непосредственно в модульном файле под определением класса.

Конвертация контрольных точек

После того как вы сгенерировали файлы моделирования, убедитесь, что реальные веса загружаются корректно. Напишите скрипт конвертации, который преобразует формат чекпоинта из исходного репозитория в формат, совместимый с Transformers, а затем сохраните его на Hub.

Напишите скрипт конвертации

Добавьте файл convert_<model>_to_hf.py в src/transformers/models/<model>/. Скрипт загружает веса из исходного репозитория, переименовывает и изменяет размерность ключей, чтобы они соответствовали именам параметров вашего модуля, и сохраняет результат с помощью save_pretrained().

Ищите существующий скрипт, который можно скопировать и адаптировать. В моделях, расположенных в src/transformers/models/, есть convert_*_to_hf.py, который можно использовать в качестве отправной точки.

После запуска скрипта загрузите сохранённую контрольную точку с помощью from_pretrained() и убедитесь, что все ожидаемые веса загружены корректно. Непроизведённые ключи контрольной точки указывают на несоответствие имён, поэтому выводите их, чтобы выявлять проблемы на раннем этапе.

model = YourModelForTask.from_pretrained("path/to/output/")

Проверяйте совпадение формы и имени при итерации по ключам. Несоответствие форм обычно означает, что параметр в вашем конфиге указан неверно, архитектура отличается от исходной или вес необходимо транспонировать.

for key, tensor in original_state_dict.items():
    hf_tensor = hf_model.state_dict().get(mapped_key)
    assert hf_tensor.shape == tensor.shape, (
        f"Shape mismatch for {key}: expected {tensor.shape}, got {hf_tensor.shape}"
    )

Устраняйте проблемы, последовательно переключаясь между модульным файлом, сгенерированным файлом модели и скриптом конвертации, пока все веса не будут загружены без ошибок.

После успешной загрузки чекпоинта отправьте его на Hub с помощью push_to_hub(). Подробнее см. в руководстве по обмену моделями.

model.push_to_hub("username/your-model-name")

Карта преобразований во время выполнения

Добавьте маппинг времени выполнения в src/transformers/conversion_mapping.py, если опубликованные веса не соответствуют структуре параметров вашего модуля. Типичные случаи включают слиянные веса, хранящиеся отдельно, и тензоры экспертов MoE, требующие стекирования. Маппинг позволяет from_pretrained() загрузить чекпоинт с Hub без отдельного шага экспорта.

См. руководство по динамической загрузке весов, чтобы узнать, как написать правила WeightRenaming и WeightConverter и зарегистрировать их для вашего model_type.

Дальнейшие шаги

• Правила структуры модели — это статические правила, применяемые ко всем файлам modeling_*.py, modular_*.py и configuration_*.py. Запустите make typing, чтобы проверить их перед созданием PR.
• Добавление компонентов обработки изображений описывает, как добавить image processor, video processor и processor для мультимодальной модели.
• Автоматическая генерация docstrings показывает, как использовать @auto_docstring, чтобы не писать вручную документацию аргументов для общих API моделей.
• Создание тестов для моделей описывает, как написать интеграционные тесты для новой модели и запустить её локально.
• Проверки pull request описывает проверки CI, которые должен пройти ваш PR перед слиянием, и объясняет, как воспроизвести и исправить их локально.