TorchNeuron：在AWS Trainium上运行HF模型比以往更轻松

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

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)

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

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