AliceAI-T5-35B-A0.6B-Base

Описание

AliceAI-T5 — базовая языковая модель с архитектурой encoder-decoder и разреженными MoE-слоями: 34,35 млрд уникальных параметров, 512 экспертов в каждом MoE-слое, из которых для каждого токена выбираются 8. Подробно про эту модель можно прочитать в статье на Хабре.

Бенчмарки

Бенчмарки pretrain моделей, посчитанные во внутренней инфраструктуре.

Как использовать

Установите зависимости из директории модели:

pip install -r requirements.txt

Загрузка пользовательской архитектуры требует trust_remote_code=True.

Пример генерации через Hugging Face Transformers:

import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model_path = "yandex/AliceAI-T5-35B-A0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForSeq2SeqLM.from_pretrained(
    model_path,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    attn_implementation="eager", # flash_attention_2
    device_map={"": "cuda:0"},
).eval()

prompt = "Question: What is the capital of France?\nAnswer:"
mode_id = tokenizer.convert_tokens_to_ids("[_S_]")
span_id = tokenizer.convert_tokens_to_ids("<SPAN#0>")
content_ids = tokenizer(prompt, add_special_tokens=False).input_ids
input_ids = torch.tensor([[mode_id, *content_ids, span_id]], device=model.device)
decoder_input_ids = torch.tensor([[model.config.decoder.bos_token_id, span_id]], device=model.device)

with torch.inference_mode():
    output_ids = model.generate(
        input_ids=input_ids,
        attention_mask=torch.ones_like(input_ids),
        decoder_input_ids=decoder_input_ids,
        do_sample=False,
        max_new_tokens=64,
    )

completion_ids = output_ids[0, decoder_input_ids.shape[1] :]
text = tokenizer.decode(completion_ids, skip_special_tokens=True)
print(text.split("<SPAN#", 1)[0])

Отдельный энкодер

AliceAIT5MoEEncoderModel загружает только энкодер из полного чекпойнта.

import torch
from transformers import AutoConfig, AutoTokenizer
from transformers.dynamic_module_utils import get_class_from_dynamic_module

model_path = "yandex/AliceAI-T5-35B-A0.6B"
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
config.is_encoder_decoder = False
EncoderClass = get_class_from_dynamic_module(
    "modeling_aliceai_t5_moe.AliceAIT5MoEEncoderModel", model_path
)
encoder = EncoderClass.from_pretrained(
    model_path,
    config=config,
    dtype=torch.bfloat16,
    attn_implementation="eager",
    device_map={"": "cuda:0"},
).eval()
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
mode_id = tokenizer.convert_tokens_to_ids("[_S_]")
span_id = tokenizer.convert_tokens_to_ids("<SPAN#0>")
content_ids = tokenizer("Текст для энкодера", add_special_tokens=False).input_ids
input_ids = torch.tensor([[mode_id, *content_ids, span_id]], device=encoder.device)

with torch.inference_mode():
    states = encoder(
        input_ids=input_ids,
        attention_mask=torch.ones_like(input_ids),
    ).last_hidden_state
print(states.shape)  # [batch, sequence_length, 1536]

Как дообучить под свои задачи

Минимальный пример Transformers + PEFT LoRA — finetune_example.py.

Данные в примере берутся из tatsu-lab/alpaca.

Запустите из директории модели в новом процессе:

pip install -r requirements-training.txt
CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=2 MKL_NUM_THREADS=2 TOKENIZERS_PARALLELISM=false python finetune_example.py

import torch
from peft import PeftModel
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

adapter_path = "/path/to/output-adapter"
tokenizer = AutoTokenizer.from_pretrained(adapter_path)
base = AutoModelForSeq2SeqLM.from_pretrained(
    "yandex/AliceAI-T5-35B-A0.6B",
    trust_remote_code=True,
    dtype=torch.bfloat16,
    attn_implementation="eager",
    device_map={"": "cuda:0"},
)
model = PeftModel.from_pretrained(base, adapter_path).eval()
prompt = "Question: What is the capital of France?\nAnswer:"
mode_id, span_id = tokenizer.convert_tokens_to_ids(["[_S_]", "<SPAN#0>"])
input_ids = torch.tensor(
    [[mode_id, *tokenizer(prompt, add_special_tokens=False).input_ids, span_id]],
    device="cuda:0",
)
decoder_input_ids = torch.tensor([[model.config.decoder.bos_token_id, span_id]], device="cuda:0")
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
    output_ids = model.generate(
        input_ids=input_ids,
        attention_mask=torch.ones_like(input_ids),
        decoder_input_ids=decoder_input_ids,
        use_cache=True,
        do_sample=False,
        max_new_tokens=64,
    )
print(tokenizer.decode(output_ids[0, 2:], skip_special_tokens=True).split("<SPAN#", 1)[0])