Install

pip install --upgrade-strategy eager optimum-neuron[neuronx]

pip install --upgrade-strategy eager optimum-neuron[neuron]

pip install git+https://github.com/huggingface/optimum-neuron.git

Quick Start

Training

importtorchimporttorch_xla.runtimeasxrfromdatasetsimportload_datasetfromtransformersimportAutoTokenizer# Optimum Neuron's drop-in replacements for standard training componentsfromoptimum.neuronimportNeuronSFTConfig,NeuronSFTTrainer,NeuronTrainingArgumentsfromoptimum.neuron.models.trainingimportNeuronModelForCausalLMdefformat_dolly_dataset(example):"""Format Dolly dataset into instruction-following format."""instruction=f"### Instruction\n{example['instruction']}"context=f"### Context\n{example['context']}"ifexample["context"]elseNoneresponse=f"### Answer\n{example['response']}"# Combine all parts with double newlinesparts= [instruction,context,response]return"\n\n".join(partforpartinpartsifpart)defmain():# 📊 Load instruction-following datasetdataset=load_dataset("databricks/databricks-dolly-15k",split="train")# 🔧 Model configurationmodel_id="Qwen/Qwen3-1.7B"output_dir="qwen3-1.7b-finetuned"# 🔤 Setup tokenizertokenizer=AutoTokenizer.from_pretrained(model_id)tokenizer.pad_token=tokenizer.eos_token# ⚙️ Configure training for Trainiumtraining_args=NeuronTrainingArguments(learning_rate=1e-4,tensor_parallel_size=8,# Split model across 8 acceleratorsper_device_train_batch_size=1,# Batch size per devicegradient_accumulation_steps=8,logging_steps=1,output_dir=output_dir,    )# 🧠 Load model optimized for Trainiummodel=NeuronModelForCausalLM.from_pretrained(model_id,training_args.trn_config,torch_dtype=torch.bfloat16,use_flash_attention_2=True,# Enable fast attention    )# 📝 Setup supervised fine-tuningsft_config=NeuronSFTConfig(max_seq_length=2048,packing=True,# Pack multiple samples for efficiency**training_args.to_dict(),    )# 🚀 Initialize trainer and start trainingtrainer=NeuronSFTTrainer(model=model,args=sft_config,tokenizer=tokenizer,train_dataset=dataset,formatting_func=format_dolly_dataset,    )trainer.train()# 🤗 Share your model with the communitytrainer.push_to_hub(commit_message="Fine-tuned on Databricks Dolly dataset",blocking=True,model_name=output_dir,    )ifxr.local_ordinal()==0:print(f"✅ Training complete! Model saved to{output_dir}")if__name__=="__main__":main()

NEURON_CC_FLAGS="--model-type transformer" neuron_parallel_compile torchrun --nproc_per_node 32 sft_finetune_qwen3.py

NEURON_CC_FLAGS="--model-type transformer" torchrun --nproc_per_node 32 sft_finetune_qwen3.py

Inference

optimum-cliexport neuron \  --model distilbert-base-uncased-finetuned-sst-2-english \  --batch_size 1 \  --sequence_length 32 \  --auto_cast matmul \  --auto_cast_type bf16 \  distilbert_base_uncased_finetuned_sst2_english_neuron/

from transformers import AutoTokenizer-from transformers import AutoModelForSequenceClassification+from optimum.neuron import NeuronModelForSequenceClassification# PyTorch checkpoint-model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")+model = NeuronModelForSequenceClassification.from_pretrained("distilbert_base_uncased_finetuned_sst2_english_neuron")tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")inputs = tokenizer("Hamilton is considered to be the best musical of past years.", return_tensors="pt")logits = model(**inputs).logitsprint(model.config.id2label[logits.argmax().item()])# 'POSITIVE'

Documentation