Hugging Face bổ sung hỗ trợ PyTorch gốc cho chip AWS Trainium

Hugging Face đã ra mắt TorchNeuron, một backend gốc mới cho phép các mô hình PyTorch chạy trực tiếp trên phần cứng AWS Trainium. Được phát hành dưới dạng beta công khai vào ngày 20 tháng 10 năm 2026, bản cập nhật này loại bỏ nhu cầu về các quy trình biên dịch chuyên biệt, giúp lập trình viên sử dụng các API PyTorch tiêu chuẩn để huấn luyện và suy luận trên các bộ tăng tốc AI của Amazon.

Điều gì đã xảy ra

Trong nhiều năm, việc chạy các mô hình Hugging Face trên chip AWS Trainium và Inferentia đòi hỏi phải sử dụng Optimum Neuron. Thư viện này đóng vai trò là cầu nối giữa PyTorch và phần cứng, dựa trên gói torch-neuronx được xây dựng trên nền tảng PyTorch XLA. Mặc dù hoạt động hiệu quả, cách tiếp cận này buộc lập trình viên phải thích nghi với tư duy tập trung vào trình biên dịch. Họ phải xử lý các tensor lười (lazy tensors) thay vì thực thi ngay lập tức, giới hạn ở các hình dạng tĩnh (static shapes) và quản lý việc biên dịch lại thường xuyên mỗi khi đồ thị tính toán thay đổi. Optimum Neuron duy trì mã tùy chỉnh, bao gồm các lớp mô hình hóa và trainer cụ thể, để che giấu sự phức tạp này khỏi người dùng.

TorchNeuron thay đổi kiến trúc này bằng cách tích hợp trực tiếp vào chính PyTorch. Nó cung cấp neuron như một thiết bị gốc, tương tự như cách lập trình viên tương tác với GPU CUDA hoặc thiết bị CPU. Sự chuyển đổi này cho phép thực thi eager, nghĩa là các thao tác được thực hiện ngay lập tức khi được gọi, thay vì xếp hàng chờ để biên dịch sau đó. Nó cũng hỗ trợ các API phân tán gốc như Fully Sharded Data Parallel (FSDP) và Tensor Parallelism (TP), cùng với các nguyên thủy phân tán riêng của PyTorch như DTensor và FSDP2. Ngoài ra, nó hoạt động liền mạch với torch.compile, cho phép tối ưu hóa hiệu suất mà không cần rời khỏi quy trình làm việc PyTorch tiêu chuẩn.

Kết quả là, các thư viện lớn của Hugging Face bao gồm Transformers, Diffusers, Accelerate, PEFT và TRL giờ đây có thể chạy trực tiếp trên Trainium. Lập trình viên có thể sử dụng cùng các lớp và script mà họ sẽ dùng trên bất kỳ bộ tăng tốc nào khác. Hugging Face hiện khuyến nghị người dùng di chuyển khỏi Optimum Neuron và sử dụng các thư viện này trực tiếp với backend mới. Bản cập nhật nhằm mục đích đơn giản hóa trải nghiệm lập trình viên và giúp việc chia sẻ mô hình trên Hugging Face Hub trở nên dễ dàng hơn đối với những người sử dụng cơ sở hạ tầng AWS.

Cách thức hoạt động

Cơ chế cốt lõi đằng sau TorchNeuron là việc tích hợp nó như một backend PyTorch gốc. Thay vì dịch các thao tác PyTorch thông qua một lớp trình biên dịch bên ngoài như XLA, TorchNeuron cho phép PyTorch nhận diện neuron như một thiết bị hạng nhất (first-class device). Điều này kích hoạt chế độ thực thi eager, nơi mỗi thao tác được thực hiện ngay lập tức trên phần cứng. Điều này trái ngược với mô hình đánh giá lười trước đây, vốn yêu cầu xây dựng một đồ thị tính toán hoàn chỉnh trước khi quá trình thực thi có thể bắt đầu.

Bằng cách hỗ trợ các nguyên thủy phân tán gốc, TorchNeuron tận dụng các công cụ tích hợp sẵn của PyTorch để mở rộng quy mô mô hình trên nhiều chip. Các tính năng như FSDP2 cho phép sử dụng bộ nhớ hiệu quả trong quá trình huấn luyện bằng cách phân đoạn các tham số mô hình trên nhiều thiết bị. Đối với suy luận, backend hỗ trợ song song hóa tensor (tensor parallelism), chia nhỏ các lớp mô hình trên nhiều bộ tăng tốc để đẩy nhanh tốc độ xử lý. Hệ thống cũng khuyến khích sử dụng cache tĩnh cho các tác vụ tạo sinh để giảm thiểu chi phí biên dịch lại, đảm bảo hiệu suất vẫn cao ngay cả với tính linh hoạt của thực thi eager.

Chi tiết quan trọng

• TorchNeuron khả dụng dưới dạng beta công khai bắt đầu từ ngày 20 tháng 10 năm 2026.
• Lập trình viên có thể cài đặt gói bằng lệnh pip install torch-neuronx.
• Backend mới hỗ trợ thực thi eager, torch.compile và các API phân tán gốc như FSDP và TP.
• Trong Transformers, 89% trong số 100 kiến trúc mô hình tạo sinh văn bản được sử dụng nhiều nhất được hỗ trợ sẵn, tăng so với mức 46% của Optimum Neuron v0.4.5.
• Người dùng có thể chuyển mô hình sang thiết bị Neuron bằng cách truyền đối số device_map=neuron trong mã của họ.
• Hugging Face khuyến nghị sử dụng trực tiếp Transformers, Diffusers và các thư viện khác thay vì thông qua Optimum Neuron.

Tại sao điều này quan trọng

Bản cập nhật này đáng kể hạ thấp rào cản gia nhập cho các kỹ sư muốn sử dụng AWS Trainium. Trước đây, việc cần phải suy nghĩ như một trình biên dịch và quản lý các hình dạng tĩnh đã tạo ra đường cong học tập dốc đứng. Bằng cách đồng bộ trải nghiệm Trainium với các quy trình làm việc PyTorch tiêu chuẩn, Hugging Face cho phép các nhóm tái sử dụng mã và chuyên môn hiện có. Điều này giảm bớt nỗ lực kỹ thuật cần thiết để chuyển đổi mô hình sang phần cứng AWS và giảm thiểu rủi ro lỗi liên quan đến các bước biên dịch tùy chỉnh.

Sự gia tăng số lượng kiến trúc mô hình được hỗ trợ cũng mở rộng tính hữu ích của Trainium cho các khối lượng công việc sản xuất. Với 89% các mô hình tạo sinh văn bản hàng đầu được hỗ trợ gốc, các tổ chức có thể triển khai phạm vi rộng hơn các mô hình tiên tiến nhất mà không cần chờ đợi các tích hợp tùy chỉnh. Sự tương đương với các bộ tăng tốc khác có nghĩa là lựa chọn phần cứng có thể được thúc đẩy nhiều hơn bởi chi phí và tính sẵn có thay vì các ràng buộc về tương thích phần mềm.

Những gì bạn có thể làm

• Cài đặt backend mới bằng cách chạy pip install torch-neuronx trong môi trường của bạn.
• Cập nhật các script suy luận của bạn để sử dụng device_map=neuron khi tải mô hình với Transformers.
• Thay thế các phụ thuộc Optimum Neuron bằng các lời gọi trực tiếp đến các thư viện Hugging Face như Transformers và Diffusers.
• Triển khai cache tĩnh cho các tác vụ tạo sinh văn bản để ngăn chặn việc biên dịch lại thường xuyên và cải thiện hiệu suất.
• Kiểm tra các thiết lập huấn luyện phân tán hiện có của bạn bằng cách sử dụng các nguyên thủy PyTorch gốc như FSDP2 và DTensor trên các phiên bản Trainium.
• Xem blog của PyTorch để biết thêm chi tiết kỹ thuật về cách TorchNeuron tích hợp với hệ sinh thái PyTorch rộng lớn hơn.