Привет, Хабр! Меня зовут Владимир, и это пятая часть цикла статей по написанию и обучению небольшой decoder-only LLM с нуля. В прошлых частях мы собрали и обучили модель LinguaLaboratoriumMechanicus — миниатюрную GPT-like LLM во вселенной Warhammer 40K. В этой части упакуем её в формат Hugging Face: сделаем конфиг и обёртку под transformers, сохраним и зальём модель на Hub.

Содержание цикла

1. Подготовка и токенизация данных

2. Трансформер

3. Сборка и обучение LLM

4. SFT этап - дообучение на вопрос-ответ

5. Интеграция с Hugging Face (вы находитесь здесь)

Содержание может меняться и дополняться ссылками по мере написания

Зачем это надо

Модель у нас уже есть. Веса сохранены, токенизатор лежит рядом, функцию generate мы написали ещё в третьей части. Вроде бы всё готово - открывай чекпоинт и общайся. На практике без общей “упаковки” это быстро превращается в ручную сборку LEGO по памяти: надо вспомнить архитектуру, аккуратно создать модель с теми же гиперпараметрами, подгрузить state_dict и свой generate. Примерно так:

import torch
from transformers import AutoTokenizer
from llm import LinguaLaboratoriumMechanicus
from llm.generation import generate

ckpt = torch.load('checkpoints/checkpoint_epoch10.pt', map_location='cpu')
model = LinguaLaboratoriumMechanicus(
    vocab_size=50257, emb_dim=768, n_layers=12, n_heads=12, max_context_length=1024,
).to('cuda')
model.load_state_dict(ckpt['model_state_dict'])
tokenizer = AutoTokenizer.from_pretrained('tokenizer/tokenizer_config')

print(generate(model, tokenizer, 'Гибельный шторм - это', device='cuda'))

Код рабочий, но он привязан к конкретному репозиторию, собственным классам и локальным путям. Если захочется дать кому-то поиграть в твою модельку или просто через месяц снова запустить - придётся тащить за собой кусок проекта и помнить, что в чекпоинте лежит словарь с ключом model_state_dict, а не “голые” веса. Универсальности ноль.

Hugging Face Hub как раз про слой удобства. Один раз упаковал модель в привычный формат transformers - и дальше везде одно и то же:

from transformers import AutoTokenizer, AutoModelForCausalLM

repo_id = 'GoldenGekko/LinguaLaboratoriumMechanicus'
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True).to('cuda')

Плюс уже готовый метод генерации токенов с различными методами внесения разнообразия в генерацию (внимательный читатель помнит, что мне было лениво метод top_p реализовывать, а тут уже есть). Плюс к этому на странице загрузки модели появляется блок Use this model: Hugging Face сам подставляет готовые сниппеты, варианты запуска и интеграции. Не нужно объяснять “скачай вот этот .pt и скопируй мой llm/” - модель становится обычным артефактом экосистемы, а не локальной поделкой на диске.

И когда все плюсы перевесили сотни полторы лишних строк кода интеграции, перейдём к реализации.

Интеграция с HF

Не знаю, кто такой John6666 и как ему в голову пришло завести блог в карточке датасета на Hugging Face, но факт есть факт - в этом необычном месте я нашел достаточно внятную инструкцию как залить кастомную архитектуру в Hub. Понадобятся всего две вещи - конфиг модели, унаследованный от transformers.PretrainedConfig и обёртка модели, унаследованная от transformers.PreTrainedModel.

LinguaLaboratoriumMechanicusConfig

Если почитать документацию, то transformers.PretrainedConfig - это базовый класс конфигурации, который управляет несколькими параметрами, общими для конфигураций всех моделей, а также методами загрузки/выгрузки/сохранения конфигураций.

На практике это контейнер: сюда складывают архитектуру модели (vocab_size, число слоёв, размер эмбеддингов и т.д.), а библиотека умеет сохранить всё в config.json и потом восстановить. Веса модели при этом не передаются - конфиг отвечает только за “как собрать модель”, а не за “какие у неё веса”.

Чтобы AutoConfig / AutoModel поняли, что это именно наша LLM, а не GPT или Qwen, у конфигурации должен быть свой уникальный model_type - строковый идентификатор архитектуры. В init необходимо записать параметры своей архитектуры атрибутами класса.

Ещё один момент - алиасы. Hugging Face GenerationMixin (от которого в дальнейшем будем наследовать класс LLM для получения model.generate(...)) ожидает привычные имена - vocab_size, hidden_size, num_attention_heads и num_hidden_layers. Несоответствия в конфиге я устранил простым дублированием полей - так и генерация работает, и в коде модели по-прежнему “свои” поля:

from transformers import PretrainedConfig

class LinguaLaboratoriumMechanicusConfig(PretrainedConfig):
    model_type = "lingua_laboratorium_mechanicus"

    def __init__(
        self, vocab_size: int = 50257, emb_dim: int = 768, n_layers: int = 12, n_heads: int = 12,
        max_context_length: int = 1024, dropout: float = 0.1, qkv_bias: bool = False, **kwargs,
    ):
        super().__init__(**kwargs)

        self.vocab_size = vocab_size
        self.emb_dim = emb_dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.max_context_length = max_context_length
        self.dropout = dropout
        self.qkv_bias = qkv_bias

        self.num_hidden_layers = n_layers
        self.hidden_size = emb_dim
        self.num_attention_heads = n_heads

Конфиг есть. Следующий шаг - обернуть нашу LinguaLaboratoriumMechanicus в класс, который библиотека transformers умеет сохранять, грузить и запускать generate.

LLMForCausalLM

Конфиг сам по себе - это паспорт архитектуры: он описывает, какой модель должна быть, но не содержит ни слоёв, ни весов, ни forward. В третьей части у нас уже была нормальная nn.Module - LinguaLaboratoriumMechanicus. Для Hub этого мало: AutoModelForCausalLM.from_pretrained не умеет разговаривать с произвольным torch.nn.Module, ему нужен класс в терминах transformers.

Поэтому поверх нашей архитектуры делается обёртка. По сути это тот же Transformer из части 3, только завёрнутый в интерфейс, который библиотека уже знает: умеет сохраняться, грузиться с диска/Hub и отдавать готовый generate. В transformers для этого есть два базовых класса, и наша обёртка наследуется сразу от обоих:

1. PreTrainedModel - надстройка вокруг nn.Module, от которой наследуются save_pretrained / from_pretrained, работа с config, инициализация весов, перенос на device и прочая обязательная обвязка.

2. GenerationMixin - предоставляет метод model.generate(...) со всякими температурами, top_k, top_p и остальными радостями, которые мы уже частично реализовали руками.

Связка конфиг - модель объявляется явно: атрибут класса config_class = LinguaLaboratoriumMechanicusConfig. По нему transformers понимает, какой JSON читать и в какой класс собирать модель:

from transformers import GenerationMixin, PreTrainedModel
from .configuration_llm import LinguaLaboratoriumMechanicusConfig

class LLMForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = LinguaLaboratoriumMechanicusConfig
    _no_split_modules = ['TransformerBlock']

nosplit_modules = ['TransformerBlock'] - подсказка для шардинга “этот модуль не режь пополам”, для модели на 163M не принципиально, но как бы это хорошая практика.

В конструкторе мы повторяем архитектуру нашего LinguaLaboratoriumMechanicus. Единственное отличие - в первоначальной инициализации весов. Для этого у Hugging Face есть свой стандартный метод инициализации - post_init:

class LLMForCausalLM(PreTrainedModel, GenerationMixin):
    # Предыдущий код
    def __init__(self, config: LinguaLaboratoriumMechanicusConfig):
        super().__init__(config)

        self.vocab_size = config.vocab_size
        self.emb_dim = config.emb_dim
        self.max_context_length = config.max_context_length

        self.token_emb = nn.Embedding(config.vocab_size, config.emb_dim)
        self.pos_emb = nn.Embedding(config.max_context_length, config.emb_dim)
        self.drop_emb = nn.Dropout(config.dropout)

        self.blocks = nn.Sequential(*[
            TransformerBlock(
                emb_dim=config.emb_dim,
                n_heads=config.n_heads,
                context_length=config.max_context_length,
                dropout=config.dropout,
                qkv_bias=config.qkv_bias,
            ) for _ in range(config.n_layers)
        ])

        self.final_norm = nn.LayerNorm(config.emb_dim)
        self.out_head = nn.Linear(config.emb_dim, config.vocab_size, bias=False)

        self.post_init()   

Осталось определить два метода - forward и prepare_inputs_for_generation. Начнём с первого.

Метод forward определяет порядок прямого прохода данных по слоям модели. Фактически, он повторяет метод forward базового класса за одним исключением - метод должен вернуть не голый тензор логитов, а тензор логитов, упакованный в transformers.modeling_outputs.CausalLMOutputWithPast. Это необходимо для корректной работы метода generate и прочих утилит библиотеки transformers:

class LLMForCausalLM(PreTrainedModel, GenerationMixin):
    # Предыдущий код
    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        if input_ids is None:
            raise ValueError('input_ids обязателен')

        _, n_tokens = input_ids.size()
        if n_tokens > self.max_context_length:
            raise ValueError(
                f'Длина входной последовательности ({n_tokens}) превышает максимальную заданную ({self.max_context_length}).')

        x = self.drop_emb(
            self.token_emb(input_ids) +
            self.pos_emb(
                torch.arange(n_tokens, device=input_ids.device).unsqueeze(0)
            )
        )

        x = self.blocks(x)
        logits = self.out_head(self.final_norm(x))

        return CausalLMOutputWithPast(logits=logits, past_key_values=None)

Из входных параметров мы используем только input_ids, остальные нужны для совместимости - во взрослых моделях они используются для передачи маски PAD токенов (attention_mask) и KV-кеша (past_key_values и use_cache).

Метод prepare_inputs_for_generation, как видно из названия, используется для подготовки данных перед генерацией. Обычно, в нём идёт работа с KV-кешем, но так как мы ничего не реализовывали, то будет:

class LLMForCausalLM(PreTrainedModel, GenerationMixin):
    # Предыдущий код
    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {'input_ids': input_ids}

Готово. Сохраняем код конфигуратора и обёртки в файлы configuration_*.py и modeling_*.py соответственно, и переходим к заключительному этапу - преобразованию сохранённых torch-тензоров в Hugging Face структуры.

Подготовка Hugging Face модели

Алгоритм преобразования torch модели в Hugging Face достаточно простой - загружаем модель из чекпоинта, создаём объект нашей обёртки, “наполняем” её весами из обученной torch модели и сохраняем во временное хранилище для дальнейшей загрузки через консоль:

def main(checkpoint: Path, tokenizer_path: Path, out_dir: Path) -> None:
    ckpt = torch.load(checkpoint, map_location='cpu', weights_only=True)
    state = ckpt['model_state_dict']
    vocab_size = ckpt['vocab_size']

    model = LinguaLaboratoriumMechanicus(
        vocab_size=vocab_size, emb_dim=EMB_DIM, n_layers=N_LAYERS, n_heads=N_HEADS,
        max_context_length=MAX_CONTEXT, dropout=DROPOUT, qkv_bias=QKV_BIAS)
    model.load_state_dict(state, strict=True)

    hf_config = LinguaLaboratoriumMechanicusConfig(
        vocab_size=vocab_size, emb_dim=EMB_DIM, n_layers=N_LAYERS, n_heads=N_HEADS,
        max_context_length=MAX_CONTEXT, dropout=DROPOUT, qkv_bias=QKV_BIAS)

    hf_model = LLMForCausalLM(hf_config)
    hf_model.load_state_dict(model.state_dict(), strict=True)

    LinguaLaboratoriumMechanicusConfig.register_for_auto_class()
    hf_model.register_for_auto_class('AutoModelForCausalLM')
    hf_config.save_pretrained(out_dir)
    hf_model.save_pretrained(out_dir)

    shutil.copytree('llm', out_dir / 'llm', ignore=shutil.ignore_patterns('__pycache__', '*.pyc'))

    for name in ('tokenizer_config.json', 'tokenizer.json', 'special_tokens_map.json'):
        src = tokenizer_path / name
        if not src.exists():
            raise FileNotFoundError(f'Нет файла токенизатора: {src}')
        shutil.copy2(src, out_dir / name)

Метод .register_for_auto_class, условно, прописывает в config.json модели поле auto_map вида:

  "auto_map": {
    "AutoConfig": "configuration_llm.LinguaLaboratoriumMechanicusConfig",
    "AutoModelForCausalLM": "modeling_llm.LLMForCausalLM"
  },

Данное поле как раз и позволит в дальнейшем использовать AutoConfig и AutoModelForCausalLM.

shutil.copytree(...) копирует код нашей кастомной LLM для дальнейшей загрузки на Hub. Без этого упадёт импорт llm.transformer. Ну или можно было скопипастить весь код ещё раз в классе LLMForCausalLM.

Собственно, это все подготовительные операции.

Загрузка

Для загрузки модели надо, внезапно, быть зарегистрированным на Hugging Face и иметь токен с правами на запись в Хаб. Далее необходимо залогиниться:

uvx hf auth login

и следовать указаниям на экране. Если не работает, то можно использовать uvx --from huggingface_hub hf ...

Далее можно проверить, что процедура прошла успешно:

uvx hf auth whoami

Ну и посмотреть зарегистрированные токены

uvx hf auth list

В результате получится что-то типа

Ну и сама команда загрузки модели:

uvx hf upload <ВАШ НИКНЕЙМ>/<ВАШ РЕПОЗИТОРИЙ> <ПАПКА С МОДЕЛЬЮ>

После завершения загрузки необходимо протестировать, что модель успешно загрузилась в Хаб (на примере Base модели):

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = 'GoldenGekko/LinguaLaboratoriumMechanicus'
device = 'cuda' if torch.cuda.is_available() else 'cpu'

tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True).to(device)

prompt = 'В 31-м тысячелетии Империум'
inputs = tokenizer(prompt, return_tensors='pt').to(device)
output_ids = model.generate(inputs['input_ids'], max_new_tokens=100)
print(tokenizer.decode(output_ids[0], skip_special_tokens=True))

trust_remote_code=True надо, так как для Hugging Face наш самописный код - чужой, а этим флагом мы покажем, что доверяем “левым” скриптам.

Итоги

Ну что, за эти пять статей мы прошли большой путь - насобирали данных, собрали их в датасет, обучили на них токенизатор, разработали (ну почти) и обучили модель, запушили её на Hugging Face Hub.

Вот тут надо бы анимацию с надписью “Конец” и проявляющимся многозначительным знаком вопроса в конце

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

Конечно, это будет не то, что модель, которая не видела нашего мира, но зато она уже будет знать основы грамматики и сносно уметь говорить по-русски. Останется только немного “сдвинуть” её мысли в сторону вселенной Боевого молота сорокового тысячелетия…

Код обучения тут

Базовая модель тут

Instruct версия тут