Меня зовут Владимир, и это пятая часть цикла статей по написанию и обучению небольшой decoder‑only LLM с нуля. В прошлых частях мы собрали и обучили модель LinguaLaboratoriumMechanicus — миниатюрную GPT‑like LLM во вселенной Warhammer 40K. В этой части упакуем её в формат Hugging Face: сделаем конфиг и обёртку под transformers, сохраним и зальём модель на Hub.
Содержание цикла
-
Интеграция с 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 для этого есть два базовых класса, и наша обёртка наследуется сразу от обоих:
-
PreTrainedModel— надстройка вокругnn.Module, от которой наследуютсяsave_pretrained/from_pretrained, работа сconfig, инициализация весов, перенос на device и прочая обязательная обвязка. -
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"
},Немножко душноты
На самом деле поле auto_map прописывает save_pretrained, а .register_for_auto_class просто создаёт намерение, но для упрощения пойдёт
Данное поле как раз и позволит в дальнейшем использовать 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 версия тут
Автор: GoldenGekko
