AWS Trainium

AWS Trainium instances chạy PyTorch một cách bản địa thông qua gói torch-neuronx, trong đó đăng ký thiết bị neuron. Neuron biên dịch mô hình của bạn thành một hoặc nhiều đồ thị cho mỗi hình dạng đầu vào mà nó nhận được, do đó việc giữ hình dạng đầu vào tĩnh sẽ tránh được việc biên dịch lại.

Cài đặt

Cài đặt torch-neuronx bằng cách làm theo hướng dẫn thiết lập AWS Neuron. Sau đó, kiểm tra xem PyTorch có nhìn thấy các thiết bị Neuron hay không.

import torch

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

Nó in ra True và số lượng thiết bị Neuron, ví dụ True 4.

Chạy suy luận

device_map="auto" không phát hiện được Neuron. Việc đặt thiết bị tự động dựa vào Accelerate để liệt kê các thiết bị khả dụng, và Accelerate không liệt kê các thiết bị Neuron, nên mô hình sẽ được nạp lên CPU một cách im lặng. Hãy truyền device_map="neuron" cho from_pretrained() một cách tường minh, hoặc device="neuron" cho Pipeline.

Truyền các đầu vào sinh trên CPU thay vì gọi .to(model.device). generate() chỉ di chuyển các tensor mà forward sử dụng sang thiết bị Neuron. Các tensor tăng kích thước ở mỗi bước (chuỗi đầu ra, attention_mask) vẫn ở trên CPU, nhờ đó đồ thị đã biên dịch được giữ ổn định. generate trả về đầu ra trên cùng thiết bị với đầu vào. Xem Giữ các tensor sinh trên CPU để biết thêm chi tiết.

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

Với bộ nhớ đệm tĩnh, generate() tự động biên dịch bước truyền thẳng bằng torch.compile trên Neuron. Truyền disable_compile=True vào generate() để từ chối tính năng này.

Huấn luyện với Trainer

TrainingArguments phát hiện Neuron và tự động đặt thiết bị thành neuron. Tất cả các thế hệ Trainium đều hỗ trợ bf16.

from transformers import TrainingArguments

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

Biên dịch mô hình

Đặt torch_compile=True để bọc model bằng torch.compile. Điều này biên dịch forward pass của model, trong khi bước optimizer chạy theo chế độ eager. TrainingArguments tự động đặt torch_compile_backend thành "neuron" khi phát hiện Neuron.

from transformers import TrainingArguments

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

Mỗi hình dạng đầu vào mới sẽ kích hoạt một lần biên dịch mới. Hãy đệm các đầu vào về một độ dài cố định (padding="max_length" trong tokenizer) và đặt dataloader_drop_last=True để lô cuối cùng không có kích thước khác biệt.

Huấn luyện phân tán

Neuron hỗ trợ API phân tán của Transformers. Tạo một DistributedConfig với các kích thước song song mong muốn và truyền nó vào from_pretrained(). Transformers sẽ chia nhỏ mô hình khi tải, và Trainer sẽ sử dụng nó như vậy.

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),
)

Chạy script bằng torchrun.

torchrun --nproc_per_node=4 train.py

Xem Tensor parallelism, Expert parallelism và FSDP để biết các tùy chọn khả dụng.

Hạn chế

device_map="auto" nạp mô hình lên CPU thay vì Neuron. Hãy truyền device_map="neuron" một cách tường minh.
Hình dạng động kích hoạt việc biên dịch lại. Cả quá trình sinh và huấn luyện đều bị chậm lại khi hình dạng đầu vào thay đổi từ bước này sang bước khác. Hãy dùng bộ nhớ đệm tĩnh cho quá trình sinh và đệm có độ dài cố định cho quá trình huấn luyện.

Update on GitHub