Huấn luyện mô hình attention tuyến tính trên AWS Trainium2 với PyTorch gốc

Các nhà nghiên cứu tại i14 và Trung tâm Đổi mới Sáng tạo AI Tạo sinh của AWS (AWS Generative AI Innovation Center) đã huấn luyện và triển khai thành công mô hình Laplace 3B sử dụng bộ tăng tốc AWS Trainium2. Được công bố vào tháng 10 năm 2026, nghiên cứu này minh chứng cho thấy kiến trúc attention tuyến tính có thể được tối ưu hóa cho phần cứng chuyên dụng mà không cần dựa vào các kernel GPU tùy chỉnh.

Sự hợp tác này đánh dấu bước chuyển dịch khỏi cơ chế softmax attention truyền thống trong transformer, vốn gặp vấn đề về khả năng mở rộng theo cấp số nhân khi huấn luyện và chi phí bộ nhớ ngày càng tăng khi suy luận. Bằng cách tận dụng hỗ trợ PyTorch gốc thông qua TorchNeuron Native, nhóm đã đạt được các chỉ số hiệu suất đáng kể trong khi vẫn duy trì mức sử dụng bộ nhớ cố định cho mỗi chuỗi.

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

Các mô hình ngôn ngữ lớn thường dựa trên kiến trúc transformer sử dụng softmax attention. Trong quá trình suy luận, các mô hình này duy trì một bộ nhớ đệm key-value (KV cache) phát triển tuyến tính theo độ dài của cuộc hội thoại hoặc tài liệu. Sự gia tăng này hạn chế khả năng xử lý đồng thời, làm tăng chi phí phục vụ và thường buộc ứng dụng phải loại bỏ ngữ cảnh lịch sử. Để giải quyết những hạn chế này, i14 đã phát triển Laplace 3B, một mô hình dựa trên nguyên lý attention tuyến tính.

Khác với các transformer tiêu chuẩn, Laplace 3B sử dụng một trạng thái hồi quy (recurrent state) luôn giữ ở mức cố định 29,5 MB cho mỗi chuỗi, bất kể độ dài ngữ cảnh. Nhóm đã phối hợp với AWS để điều chỉnh kiến trúc này cho phần cứng Trainium2. Họ sử dụng TorchNeuron Native, một tính năng beta cho phép chạy mã PyTorch gốc trực tiếp trên các bộ tăng tốc AWS Trainium. Cách tiếp cận này bảo toàn thiết kế Gated Linear Attention (GLA) của mô hình cũng như việc sử dụng bộ tối ưu hóa Muon.

Kết quả cho thấy những cải thiện hiệu suất đáng kể. Trên một instance trn2.48xlarge đơn lẻ, mô hình được huấn luyện với tốc độ khoảng 120.000 token/giây. Trong quá trình suy luận, nó phục vụ 11.544 token đầu ra/giây trên 1.280 chuỗi đồng thời. Những con số này chứng minh rằng các mô hình attention tuyến tính có thể đạt được thông lượng cao trên phần cứng chuyên dụng khi stack phần mềm được căn chỉnh đúng với kiến trúc nền tảng.

Cơ chế hoạt động

Các mô hình attention tuyến tính tổ chức lại quá trình tính toán xoay quanh một trạng thái hồi quy thay vì so sánh từng cặp token. Trong các transformer truyền thống, điểm attention tăng theo cấp số nhân với độ dài chuỗi trong quá trình huấn luyện. Laplace 3B thay thế cơ chế này bằng một phương pháp cập nhật trạng thái ma trận có kích thước cố định cho mỗi token. Điều này làm giảm độ phức tạp huấn luyện từ O(T^2) xuống O(T) và giữ cho bộ nhớ suy luận không đổi. Mô hình nén lịch sử vào trạng thái này, đòi hỏi phải đánh giá cẩn thận chất lượng truy xuất cho các khối lượng công việc cụ thể.

Việc điều chỉnh mô hình này cho Trainium2 yêu cầu ba thay đổi kỹ thuật cụ thể. Thứ nhất, nhóm đã thay thế các kernel Triton được tối ưu hóa cho GPU bằng các toán tử PyTorch tiêu chuẩn. Họ cấu trúc kernel GLA để sử dụng 16 chunk gồm 128 token, phù hợp với 128 phân vùng của bộ đệm on-chip NeuronCore. Thứ hai, họ thiết kế lại giao tiếp cho bộ tối ưu hóa Muon. Bằng cách sử dụng hai tập hợp all-to-all thay vì các phương pháp gather-and-broadcast, họ đã giảm lưu lượng truyền tải trên mỗi rank đi 24 lần và tăng thông lượng lên 1,63 lần. Thứ ba, họ thực hiện unroll tĩnh cho quá trình backward pass tuần tự để phù hợp với giới hạn kích thước đồ thị và bộ nhớ của trình biên dịch.

Kiến trúc Trainium2 đóng vai trò then chốt trong các tối ưu hóa này. Mỗi NeuronCore có một State Buffer (SBUF) do phần mềm quản lý với dung lượng 28 MiB và 128 phân vùng. Khác với cache, trình biên dịch lập lịch mọi chuyển dữ liệu trước, đảm bảo không xảy ra tình trạng eviction hay miss. Bằng cách ánh xạ mỗi chunk 128 token vào một phân vùng, nhóm đã đạt được mức tăng 14% trong thông lượng đo được. Sự căn chỉnh giữa kích thước chunk của mô hình và các phân vùng vật lý của phần cứng là chìa khóa để giải phóng hiệu suất.

Chi tiết chính

• Laplace 3B duy trì trạng thái hồi quy cố định 29,5 MB cho mỗi chuỗi, độc lập với độ dài ngữ cảnh.
• Huấn luyện trên một instance trn2.48xlarge đạt khoảng 120.000 token/giây.
• Suy luận phục vụ 11.544 token đầu ra/giây trên 1.280 chuỗi đồng thời.
• Triển khai sử dụng TorchNeuron Native để chạy các toán tử PyTorch tiêu chuẩn trên Trainium2.
• Lưu lượng giao tiếp cho bộ tối ưu hóa Muon được giảm 24 lần nhờ sử dụng các tập hợp all-to-all.
• Mô hình sử dụng 16 chunk gồm 128 token để căn chỉnh với 128 phân vùng SBUF trên các NeuronCore.

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

Đối với các kỹ sư xây dựng sản phẩm AI, chi phí ngày càng tăng của KV cache là một nút thắt cổ chai lớn. Khi ngữ cảnh trở nên dài hơn, mức sử dụng bộ nhớ tăng vọt, hạn chế số lượng người dùng đồng thời mà máy chủ có thể xử lý. Các mô hình attention tuyến tính như Laplace 3B cung cấp một lối đi hướng tới chi phí phục vụ dự đoán được. Vì kích thước trạng thái không tăng theo ngữ cảnh, các ứng dụng có thể lưu giữ lịch sử dài mà không cần tăng tỷ lệ thuận việc cấp phát bộ nhớ. Điều này đặc biệt có giá trị đối với các quy trình làm việc agentic và phân tích tài liệu lớn, nơi bộ nhớ bền vững là yếu tố then chốt.

Hơn nữa, nghiên cứu này chứng minh rằng phần cứng chuyên dụng như Trainium2 có thể hỗ trợ các kiến trúc mới mà không cần viết lại kernel ở tầng thấp. Bằng cách sử dụng PyTorch gốc thông qua TorchNeuron Native, các nhà phát triển có thể tận dụng các tối ưu hóa cụ thể của phần cứng trong khi vẫn ở trong một framework quen thuộc. Thành công của chiến lược giao tiếp all-to-all cho bộ tối ưu hóa Muon cũng cung cấp một khuôn mẫu cho việc huấn luyện phân tán trên các hệ thống sharded. Nó cho thấy việc giảm bớt chi phí giao tiếp có thể mang lại lợi ích thông lượng đáng kể, ngay cả khi số phép tính FLOPs không thay đổi.

Bạn có thể làm gì

• Đánh giá các mô hình attention tuyến tính cho các trường hợp sử dụng yêu cầu cửa sổ ngữ cảnh dài và mức sử dụng bộ nhớ dự đoán được.
• Thử nghiệm TorchNeuron Native nếu bạn đang triển khai các kiến trúc tùy chỉnh trên các instance AWS Trainium.
• Căn chỉnh kích thước chunk dữ liệu với các phân vùng bộ đệm phần cứng để tối đa hóa thông lượng trên các bộ tăng tốc chuyên dụng.
• Cân nhắc các mẫu giao tiếp all-to-all cho các bộ tối ưu hóa phân tán để giảm bớt nút thắt cổ chai lưu lượng mạng.
• Giám sát chất lượng truy xuất khi chuyển từ softmax attention sang các cơ chế nén trạng thái.
• Profiling engine suy luận của bạn để xác định xem sự tăng trưởng của KV cache có đang hạn chế mức độ đồng thời của bạn hay không.