AWS Trainium

AWS Trainium 인스턴스는 torch-neuronx 패키지를 통해 PyTorch를 네이티브로 실행하며, 이 패키지는 neuron 장치를 등록합니다. Neuron은 확인하는 각 입력 형식에 대해 모델을 하나 이상의 그래프로 컴파일하므로, 입력 형식을 정적으로 유지하면 재컴파일을 방지할 수 있습니다.

설치

AWS Neuron 설정 가이드를 따라 torch-neuronx를 설치합니다. 그런 다음 PyTorch가 Neuron 장치를 인식하는지 확인합니다.

import torch

print(torch.neuron.is_available(), torch.neuron.device_count())

True와 Neuron 장치의 수를 출력합니다. 예를 들어 True 4입니다.

추론 실행

device_map="auto"는 Neuron을 감지하지 않습니다. 자동 배치에는 사용 가능한 장치를 나열하는 Accelerate에 의존하지만, Accelerate는 Neuron 장치를 나열하지 않으므로 모델이 CPU에 조용히 로드됩니다. from_pretrained()에 device_map="neuron"을 명시적으로 전달하거나, Pipeline에 device="neuron"을 전달하세요.

.to(model.device)를 호출하는 대신 생성 입력을 CPU로 전달합니다. generate()는 forward가 소비하는 텐서만 Neuron 장치로 이동합니다. 각 단계에서 커지는 텐서(출력 시퀀스, attention_mask)는 CPU에 유지되므로 컴파일된 그래프가 안정적으로 유지됩니다. generate는 입력과 동일한 장치에서 출력을 반환합니다. 자세한 내용은 생성 텐서를 CPU에 유지하기를 참조하세요.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="neuron"
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=False, max_new_tokens=50, cache_implementation="static"
)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The French Bread Law states that the French government must provide a certain amount of bread for the citizens of the country. The law is passed by the French parliament. The French parliament is a body of the French government. The French government is a body of the French state. The

정적 캐시를 사용하는 경우, generate()는 Neuron에서 torch.compile을 사용하여 forward pass를 자동으로 컴파일합니다. generate()에 disable_compile=True를 전달하여 이 기능을 비활성화할 수 있습니다.

Trainer로 학습

TrainingArguments는 Neuron을 감지하여 장치를 자동으로 neuron으로 설정합니다. 모든 Trainium 세대가 bf16을 지원합니다.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
)

모델 컴파일

torch_compile=True를 설정하면 모델이 torch.compile로 래핑됩니다. 이렇게 하면 모델의 forward pass가 컴파일되고, optimizer step은 즉시 실행됩니다. TrainingArguments는 Neuron을 감지하면 torch_compile_backend를 자동으로 "neuron"으로 설정합니다.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
    torch_compile=True,
)

각 새로운 입력 형상은 새로운 컴파일을 트리거합니다. 입력을 고정된 길이로 패딩(padding="max_length"을 토크나이저에서 설정)하고, 마지막 배치의 크기가 달라지지 않도록 dataloader_drop_last=True를 설정하세요.

분산 학습

Neuron은 Transformers의 분산 API를 지원합니다. 원하는 병렬화 크기를 지정하여 DistributedConfig를 생성하고 from_pretrained()에 전달합니다. Transformers는 로드 시점에 모델을 샤딩하며, Trainer는 이를 그대로 사용합니다.

import torch
from transformers import AutoModelForCausalLM, DistributedConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-1.7B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=2, fsdp_size=2),
)

torchrun로 스크립트를 실행합니다.

torchrun --nproc_per_node=4 train.py

사용 가능한 옵션은 Tensor parallelism, Expert parallelism, FSDP를 참고하세요.

한계

device_map="auto"는 모델을 Neuron 대신 CPU에 로드합니다. device_map="neuron"을 명시적으로 전달하세요.
동적 shape는 재컴파일을 유발합니다. 입력 shape가 단계별로 변경되면 생성과 학습 모두 속도가 느려집니다. 생성에는 정적 캐시를, 학습에는 고정 길이 패딩을 사용하세요.

Update on GitHub