跳到正文
Hugging Face Blog·· 22 小時前精選AI 評分68

TorchNeuron 為 AWS Trainium 提供 PyTorch 原生後端

TorchNeuron: running HF models on AWS Trainium is easier than ever

AI 導讀

TorchNeuron 為 AWS Trainium 提供 PyTorch 原生後端，令 Transformers、Diffusers 等 Hugging Face 函式庫可直接使用 Trainium。

推薦理由

文章比較 Optimum Neuron 與 TorchNeuron 的執行方式，並提供推理、微調及圖像生成範例，呈現 Trainium 接入 Hugging Face 工作流程的具體變化。

正文 · 繁體中文

TorchNeuron 為 AWS Trainium 上的 PyTorch 提供全新的原生後端。Transformers、Diffusers 和 Kernels 原生支援這個後端，讓用戶比以往更輕易在 Trainium 上執行模型，並在 HF hub 上分享模型。

TorchNeuron

TorchNeuron 是 PyTorch 外掛，為 Trainium 帶來全新後端，支援 eager 模式、torch.compile，以及 FSDP 和 TP 等原生分散式 API。 閲讀 Pytorch 的網誌

從 2026 年 10 月 20 日起，TorchNeuron 將以公開測試版形式提供。研究人員和開發人員可透過 pip install torch-neuronx 安裝。

新後端，HF 程式庫中的原生介面

迄今為止，Hugging Face 用戶在 AWS Trainium 和 Inferentia 上的使用體驗，一直透過 optimum-neuron 提供。這兩款 AI 晶片分別為具成本效益的大規模訓練和推理而設計。這個程式庫依賴 torch-neuronx 套件，而該套件以 PyTorch XLA 為基礎。當時這是自然的選擇，但也意味著要以編譯器的方式思考：使用惰性張量而非即時執行、只支援靜態形狀，而且每次圖中斷都要重新編譯。Optimum Neuron 的存在就是為了隱藏這種複雜性，並透過維護自訂程式碼（建模類別、訓練器等）來達成這一點。

TorchNeuron 改變了這一點。它直接整合至 PyTorch，並將 neuron 公開為原生裝置，支援即時執行、PyTorch 自身的分散式基本元件（DTensor、FSDP2），以及 torch.compile。

對 Hugging Face 程式庫而言，Trainium 就如同其他裝置一樣：Transformers、Diffusers、Accelerate、PEFT 和 TRL 都可直接在其上執行，使用的類別和指令碼與在其他加速器上相同。從現在起，我們建議 Neuron 用戶直接使用這些程式庫，而非透過 Optimum Neuron。

Transformers

推理

Transformers 現已原生支援多種模型的文字生成。使用 TorchNeuron 時，最常用的 100 種文字生成模型架構中，有 89% 開箱即用便獲支援；Optimum Neuron v0.4.5 的支援比例則為 46%。只要在程式碼中傳入 device_map=neuron 參數，模型就會移至 Neuron 裝置上。

以下是使用張量平行處理進行文字生成的範例，讓你能充分運用 Neuron 執行個體。請注意，範例使用靜態快取進行生成，因為我們強烈建議在 Neuron 上使用靜態快取，以避免頻繁重新編譯。

import os

import torch
import torch.distributed as dist
import torch_neuronx  # noqa: F401
from torch.distributed import DeviceMesh
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
from transformers.distributed import DistributedConfig
from transformers.generation.configuration_utils import CompileConfig

MODEL_ID = "Qwen/Qwen3.8-27B"
MAX_SEQ_LEN = 1024
PREFILL_CHUNK_SIZE = 512

dist.init_process_group(backend="neuron")
world_size = dist.get_world_size()

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=world_size),
    device_mesh=DeviceMesh("neuron", list(range(world_size))),
).eval()

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

inputs = tokenizer.apply_chat_template(
    [{"role": "user", "content": "The future of AI is"}],
    add_generation_prompt=True,
    padding=True,
    tokenizer_kwargs={"pad_to_multiple_of": PREFILL_CHUNK_SIZE},
    return_tensors="pt",
)

generation_config = GenerationConfig(
    max_new_tokens=128,
    do_sample=False,
    use_cache=True,
    pad_token_id=tokenizer.pad_token_id,
    eos_token_id=tokenizer.eos_token_id,
    cache_implementation="static",
    max_cache_len=MAX_SEQ_LEN,
    prefill_chunk_size=PREFILL_CHUNK_SIZE,
    compile_config=CompileConfig(backend="neuron", fullgraph=False, dynamic=False),
)

with torch.no_grad():
    output_ids = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        generation_config=generation_config,
    )

if dist.get_rank() == 0:
    prompt_len = inputs["input_ids"].shape[1]
    print(tokenizer.decode(output_ids[0, prompt_len:], skip_special_tokens=True))

訓練

Neuron 現在是 PyTorch 原生裝置，因此在 Trainium 上訓練時，可使用與其他加速器相同的 Hugging Face 程式庫：Transformers、Accelerate、PEFT 和 TRL。你不再需要專用訓練器，也不需要 Neuron 專用的模型程式碼分支。標準的 TRL 微調指令碼只需作少量 Neuron 專用設定，便可在 Trainium 執行個體上執行。

以下是以 TRL 的 SFTTrainer 為基礎編寫的監督式微調指令碼。如果你曾使用 TRL 在 GPU 上微調模型，這段指令碼應該會很熟悉。

import os

import torch
from datasets import load_dataset
from transformers.distributed import DistributedConfig

from trl import ModelConfig, ScriptArguments, SFTConfig, SFTTrainer, TrlParser


def main(script_args, training_args, model_args):
    # The Neuron backend can't compile Dynamo's automatic dynamic shapes yet (https://github.com/huggingface/trl/issues/7660)
    torch._dynamo.config.automatic_dynamic_shapes = False

    training_args.model_init_kwargs = {
        "revision": model_args.model_revision,
        "trust_remote_code": training_args.trust_remote_code,
        "attn_implementation": model_args.attn_implementation,
        "dtype": model_args.dtype,
        "device_map": None,
    }
    if int(os.environ.get("WORLD_SIZE", "1")) > 1:
        # The model is sharded on the TP x FSDP grid as its weights are loaded
        training_args.model_init_kwargs["distributed_config"] = DistributedConfig(
            tp_size=int(os.environ.get("TP_SIZE", "1")),
            fsdp_size=int(os.environ.get("FSDP_SIZE", "1")),
        )

    # Prepare dataset
    train_dataset = load_dataset(
        script_args.dataset_name,
        name=script_args.dataset_config,
        split=script_args.dataset_train_split,
    )
    # UltraChat also has a string `prompt` column, which makes TRL treat it as a prompt-completion dataset.
    train_dataset = train_dataset.select_columns(["messages"])

    # Training
    trainer = SFTTrainer(
        model=model_args.model_name_or_path,
        args=training_args,
        train_dataset=train_dataset,
    )

    trainer.train()
    trainer.save_model(training_args.output_dir)
    trainer.end()


if __name__ == "__main__":
    parser = TrlParser((ScriptArguments, SFTConfig, ModelConfig))
    script_args, training_args, model_args = parser.parse_args_and_config()
    main(script_args, training_args, model_args)

平行處理由 Transformers 原生處理。你只需在 model_init_kwargs 中加入包含張量平行處理（TP）和 FSDP 尺寸的 DistributedConfig，再由 SFTTrainer 將其傳給 from_pretrained，模型便會在載入權重時分片。Trainer 會識別模型已經分散式處理，因此不會再次包裝模型。你不需要 --fsdp 旗標或 Accelerate FSDP 設定。

指令碼中沒有任何部分專為 Trainium 而設：相同程式碼可在任何硬件上執行。在 Trainium 上執行時，有幾點值得留意：

  • 固定形狀。 指令碼會在 torch.compile 中停用自動動態形狀，並使用環繞式打包產生最多包含 2048 個 token 的序列。使用 TRL ≥ 1.16 時，Neuron 上的批次會填充至 128 個 token 的倍數，從而限制輸入形狀的數量。後續步驟如使用相同形狀，便會重用已編譯的計算圖。
  • torch.compile 已啟用。 只需傳入 --torch_compile true 即可：在 Trainium 執行個體上，torch.compile 會偵測裝置並自動選用 Neuron 後端，無需進行其他設定。
  • 已關閉非同步權重載入。 Transformers 內置的非同步權重載入不支援 Neuron，因此目前我們會使用 HF_DEACTIVATE_ASYNC_LOAD=1 將其停用。

以下指令會使用 torchrun 啟動訓練，每個 NeuronCore 會啟動一個程序。TP_SIZE 和 FSDP_SIZE 用來設定平行處理配置，兩者的乘積就是程序數量：此處為 4 x 8 = 32，與 --nproc_per_node 相符。這 32 個程序組成 4 x 8 的網格。張量平行處理會將每一層拆分至 TP 羣組中的 4 個核心，而 FSDP 會將參數、梯度及最佳化器狀態分片至 8 個 TP 羣組。

TP_SIZE=4 FSDP_SIZE=8 HF_DEACTIVATE_ASYNC_LOAD=1 uv run torchrun --nproc_per_node=32 sft_neuron_example.py \
    --model_name_or_path Qwen/Qwen3-8B \
    --dataset_name HuggingFaceH4/ultrachat_200k \
    --dataset_train_split train_sft \
    --output_dir Qwen3-8B-sft-tp4-fsdp8 \
    --dtype bfloat16 \
    --max_length 2048 \
    --packing true \
    --packing_strategy wrapped \
    --learning_rate 2e-5 \
    --warmup_steps 0.03 \
    --num_train_epochs 1 \
    --logging_steps 2 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --save_only_model true \
    --torch_compile true

我們在 8 個 Trainium2 晶片（32 個 NeuronCore）上執行，這是 trn2.48xlarge 執行個體的一半。

訓練會執行至完成，而 trainer.save_model 會寫入標準的 Transformers 檢查點。檢查點內容不受 Neuron 限制：你可以原樣將其推送至 Hugging Face Hub，並在任何其他硬件上載入。

Diffusers

推理

Diffusers 程式庫中的許多管線都支援在 TorchNeuron 的急切模式和編譯模式下運行。FLUX.1-dev、FLUX.2、Qwen-Image、Qwen-Image-2.1 和 H3 亦支援張量平行處理，讓你充分運用 Neuron 裝置。以下範例展示如何使用 TP：

import torch, torch.distributed as dist, torch_neuronx # TP
from torch.distributed.device_mesh import DeviceMesh # TP
from diffusers import QwenImage21Pipeline, QwenImage21Transformer2DModel, TensorParallelConfig

# Distributed devices initialization
dist.init_process_group(backend="neuron") # TP
dev = torch.neuron.current_device()
mesh = DeviceMesh("neuron", list(range(8))) # TP=8

# Pipeline initialization
model = "Qwen/Qwen-Image-2.1"
transformer = QwenImage21Transformer2DModel.from_pretrained(
    model,
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
    parallel_config=TensorParallelConfig(mesh=mesh), # TP
)
pipe = QwenImage21Pipeline.from_pretrained(
    model,
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)
pipe.vae.to(dev)
pipe.transformer.compile_repeated_blocks(backend="neuron", fullgraph=True, dynamic=False)

# Image generation
emb, mask, _ = pipe.encode_prompt(prompt="A capybara wearing a wizard hat, oil painting", device="cpu")  # text encoder stays on CPU
image = pipe(
    prompt_embeds=emb.to(dev),
    prompt_embeds_mask=mask.to(dev) if mask is not None else None,
    height=1024,
    width=1024,
    num_inference_steps=40,
    generator=torch.Generator("cpu").manual_seed(0), # same noise on every rank
).images[0]
if dist.get_rank() == 0:
    image.save("out.png")

你可以使用以下指令執行這段程式碼：NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py。

將會有愈來愈多管線支援 TP。如果你喜歡的管線尚未支援，也可以按照本指南為其新增支援。

核心程式

記得查看 Neuron 團隊提供的核心程式，愈來愈多 Neuron 核心程式將會在 Hub 上推出，以加快模型運行。 Neuron 核心程式可使用 Neuron Kernel Interface（NKI）編寫，該介面提供更底層的控制，以最佳化運算。用戶可執行自行建立的自訂運算，並在 Hub 上分享。

例如，你可以使用這個核心程式加快視覺 Transformer 的運行，在 Trainium 上最高可達 x2.4 的加速效果！

致謝

特別感謝 AWS 團隊在 TorchNeuron 上的工作，以及他們迅速回應我們的意見，盡力完善 HF 整合。

來源：Hugging Face Blog · huggingface.co

#產品更新#部署/工程#Hugging Face