Optimum Neuron

HuggingFace Optimum for training and inference on AWS Trainium and Inferentia

View on GitHub

Overview

Optimum Neuron is HuggingFace’s integration library for AWS Trainium and Inferentia chips. It enables training and inference of HuggingFace models on AWS custom silicon with minimal code changes.

Key Features

  • Drop-in replacement β€” Use familiar HuggingFace Trainer API on Trainium
  • Pre-compiled models β€” Optimized graph compilation for common architectures
  • Distributed training β€” Multi-node training on Trainium clusters
  • Inference optimization β€” Model compilation for Inferentia2
  • PEFT support β€” LoRA and other parameter-efficient methods on Neuron