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

Описание

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

Параметр	Значение
Размер словаря / скрытого состояния	135 040 / 1 536
Энкодер	16 слоёв, 12 голов внимания
Декодер	12 слоёв, 12 query-голов, 4 KV-головы
Размерность головы	128
Эксперты в MoE-слое	512, маршрутизация top-8
Промежуточная размерность эксперта	512
Активация / нормализация	SiLU / RMSNorm
Позиционные представления	RoPE с YaRN, контекст 128к токенов
Эмбеддинги	Общие для энкодера и декодера, связаны с LM head

Бенчмарки

Benchmark	T5 Gemma 2 4B-4B Base	Gemma 4 E4B Base	Qwen 3.5 2B Base	Qwen 3.5 4B Base	Qwen 3.5 35B-A3B Base	AliceAI-T5-35B-A0.6B Base
Factuality
(ya) CultCat (4-shot)	23,2	44,0	31,0	39,7	59,2	68,0
(ya) WikiWebFacts (5 shot)	33,6	47,2	23,6	42,1	62,4	81,3
TriviaQA (5-shot)	53,3	65,0	32,4	50,4	71,4	60,5
General Tasks
MMLU (5-shot)	54,5	71,7	65,4	77,0	84.4	78,5
MMLU Pro (CoT, 5-shot)	32,9	37,4	36,5	51,1	63,2	56,3
GPQA (5-shot)	30,0	34,0	31,6	38,8	47,1	41,3
Math and Code
GSM8K (CoT, 8-shot)	51,3	58,4	69,5	84,5	90,4	84,7
MATH 500 (CoT, 4-shot)	17,0	23,5	43.4	69.5	81.9	62,9
HumanEval (5-shot)	31,9	42.2	48,7	74.6	88,3	69,3
MBPP (3-shot)	52,4	54,4	42,6	61.1	75,4	71,9
Extract and Long Context
(ya) YExtract (4-shot)	18,1	19,8	12,7	28,4	42,4	40,4
Ruler 32K	81,9	89,6	83,6	90,2	93,0	94,7
Ruler 128K	57,5	81,3	74,6	84,4	90,1	81,4

Бенчмарки 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])