TorchNeuron: AWS Trainium платформасында Hugging Face модельдерін іске қосудың жаңа мүмкіндіктері

Жасанды интеллект пен машиналық оқыту саласындағы технологиялар қарқынды дамуда. Мәлімет өңдеу мен модельдерді жұмысқа енгізу үдерісін оңтайландыру үшін арнайы аппараттық шешімдер пайда болуда. AWS Trainium — мұндай аппараттық құралдардың бірі, ол тиімді әрі арзан есептеу қуатын ұсынады. TorchNeuron — AWS Trainium үшін арнайы әзірленген, жоғары өнімділік пен қуатты интеграция құралы, ол Hugging Face кітапханаларымен тікелей жұмыс істеуге мүмкіндік береді. Бұл мақалада TorchNeuron-ның маңызды тұстары мен оның жасанды интеллект модельдерін іске қосудағы рөлі қарастырылады.

Кіріспе

Қазіргі кезде ірі модельдер мен үлкен көлемді деректерді өңдеу үшін жылдам әрі үнемді есептеу ресурстары қажет. Бұл талаптарға сай AWS Trainium арнайы жасалған, ол үлкен нейрондық желілерді оқыту мен іске қосуға бағытталған. TorchNeuron — осы платформаның тиімділігін арттыру үшін әзірленген құрал, ол PyTorch экожүйесіне толықтай интеграцияланады және Hugging Face кітапханаларының мүмкіндіктерін кеңейтеді.

Негізгі түсіндірме

TorchNeuron — PyTorch үшін әзірленген плагин, ол Trainium-да модельдерді іске қосуды оңайлатады. Бұрынғы қолданыстағы Optimum Neuron кітапханасы мұқият қарастырғанымен, оның жұмысы компиляцияға тәуелді болды, яғни жай белсенді орындалу орнына «жалқау» тензорлар қолданылды және статикалық пішімдермен шектелді. Бұның салдарынан модельді қайта құрастыру мен күрделі теңшеулер қажет болды.

TorchNeuron осы шектеулерді жояды. Ол PyTorch-қа тікелей қосылады, Trainium-ды жеке құрылғы ретінде қарастырады және белсенді орындалуды толық қолдайды. Сонымен қатар PyTorch-тың таратылған есептеу функциялары (мысалы, DTensor, FSDP2) және torch.compile сияқты жаңа құралдарды ендіреді. Hugging Face кітапханалары (Transformers, Diffusers, Accelerate, PEFT, TRL) Trainium-де стандартты түрде жұмыс істейді, арнайы қосымша кодсыз және ерекше теңшелімсіз.

Контекст және мысалдар

Transformers кітапханасында TorchNeuron арқасында мәтін генерациясы функциялары 100-дің ең жиі қолданылатын моделінің 89%-ын қолдайды, бұрын Optimum Neuron версиясымен бұл көрсеткіш 46% еді. Мысалы, TorchNeuron көмегімен Qwen сияқты үлкен тіл модельдерін параллель есептеу арқылы оңай іске қосуға болады. Пайдаланушы тек құрылғы картасын (device_map=neuron) көрсету керек, қалған істер автоматты түрде орындалады.

Төменде Hugging Face Transformers кітапханасындағы GPT-стильдегі тіл моделін tensor parallelism қолдана отырып іске қосу мысалы берілген. Бұл тәсіл модельдің барлық Trainium ядроларынан толық пайдаланылуын қамтамасыз етеді.

import torch
import torch.distributed as dist
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
from torch.distributed import DeviceMesh

dist.init_process_group(backend="neuron")
world_size = dist.get_world_size()
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3.8-27B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=world_size),
    device_mesh=DeviceMesh("neuron", list(range(world_size))),
).eval()

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3.8-27B")
inputs = tokenizer.apply_chat_template(
    [{"role": "user", "content": "The future of AI is"}],
    add_generation_prompt=True,
    padding=True,
    return_tensors="pt",
)

generation_config = GenerationConfig(
    max_new_tokens=128,
    use_cache=True,
    cache_implementation="static",
    compile_config=CompileConfig(backend="neuron")
)

with torch.no_grad():
    output_ids = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        generation_config=generation_config
    )
print(tokenizer.decode(output_ids[0], skip_special_tokens=True))

Оқу процесі де Trainium-да тікелей Hugging Face құралдарымен жүзеге асады, бұл бұрынғыдай арнайы тренер немесе модификацияланған модель кодтарын қажет етпейді. Есептеу үдерісі модельдің масштабталуын, параметрлер мен градиенттерді автоматты түрде бөлу мүмкіндігін қамтиды.

Практикалық маңызы

TorchNeuron-ның артықшылығы — AWS Trainium-дағы есептеулер мен оқу үдерісін оңай және тиімді ұйымдастыру. Құрал PyTorch-тің байтында құралдарға ұқсас жұмыс істейді, бұл модельдерді бірлік кодпен әртүрлі аппараттық платформаға оңай ауыстыруға мүмкіндік береді. Сонымен қатар TorchNeuron арқасында әзірлеушілер күрделі компиляциялық әдістер мен кешенді теңшелімдерге қажетсіз төтеп бере алады.

Модельді оқу мен іске қосу ортасы біркелкі болғандықтан, зерттеушілер мен мамандар жаңа технологияларды игеруге көп уақыт жұмсамайды. Бұл әсіресе ірі тіл модельдері мен диффузиялық модельдерді жетілдіруде маңызды. TorchNeuron-ның ашық бета нұсқасының пайда болуы бұл процесті барынша жеңілдетеді және қолжетімді етеді.

Қорытынды

TorchNeuron — PyTorch-қа арналған жаңа backend, ол AWS Trainium аппараттық платформасын оңтайлы пайдаланып, Hugging Face модельдерін іске қосу мен оқуды едәуір жеңілдетеді. Бұл құрал тек есептеу тиімділігін арттырып қоймай, оқыту мен генерация үдерістерінің қарапайымдылығын қамтамасыз етеді. TorchNeuron көмегімен Trainium мойындалған және қол жетімді платформаға айналады, бұл өңдеу қуатын қажет ететін жобалар мен зерттеулер үшін тың мүмкіндіктер ашады.

TAGS: PyTorch, AWS Trainium, TorchNeuron, Hugging Face, нейрондық желілер, жасанды интеллект, машиналық оқыту

Дереккөз: Hugging Face Blog