Huấn luyện Qwen3 trên AWS Trainium với torchtitan và các kernel tùy chỉnh

AWS đã phát hành một hướng dẫn kỹ thuật chi tiết minh họa cách huấn luyện mô hình ngôn ngữ lớn Qwen3 trên bộ tăng tốc Trainium2 của họ bằng khung làm việc torchtitan. Được công bố vào tháng 10 năm 2026, hướng dẫn này phác thảo ba cấu hình khác nhau dành cho các kỹ sư muốn tận dụng các giao diện PyTorch gốc trên phần cứng chuyên dụng. Tài liệu đóng vai trò là tham chiếu thực tế cho các nhóm học máy nhằm tối ưu hóa quy trình tiền huấn luyện (pre-training) và tinh chỉnh (fine-tuning) trên cơ sở hạ tầng Amazon Web Services.

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

Bài viết cung cấp hướng dẫn triển khai từng bước để chạy các mô hình Qwen3 trên một instance trn2.48xlarge, chứa 16 bộ tăng tốc Trainium2 và 128 NeuronCore vật lý. Các tác giả chứng minh rằng nhà phát triển có thể sử dụng torchtitan trên Trainium tương tự như cách họ dùng trên các cụm GPU tiêu chuẩn, nhờ TorchNeuron đăng ký Trainium như một thiết bị PyTorch gốc. Điều này cho phép khung làm việc phát hiện phần cứng khi khởi động và biên dịch các thao tác theo chế độ eager hoặc thông qua torch.compile.

Hướng dẫn bao gồm ba cấu hình huấn luyện cụ thể. Phương pháp đầu tiên sử dụng Hugging Face Transformers v5 làm backend mô hình, cho phép người dùng tải trực tiếp các mô hình từ Hub bằng AutoModelForCausalLM. Cách tiếp cận thứ hai chuyển sang triển khai Qwen3 gốc của torchtitan, mang lại khả năng kiểm soát tốt hơn đối với kiến trúc mô hình và chiến lược phân mảnh (sharding). Phương pháp thứ ba giới thiệu tính tùy biến nâng cao bằng cách tích hợp các kernel mixture-of-experts của Neuron Kernel Interface (NKI), cho phép nhà phát triển gắn các toán tử tùy chỉnh cho những lớp quan trọng về hiệu suất.

Mỗi cấu hình đều sử dụng cùng một mẫu khởi chạy torchrun, đảm bảo tính nhất quán giữa các thiết lập thử nghiệm khác nhau. Hướng dẫn nhấn mạnh vào khả năng tái tạo bằng cách chỉ định chính xác các phiên bản phần mềm, bao gồm PyTorch 2.14.0, Transformers 5.16.0 và bản phát hành torchtitan-neuron-2.33. Nó cũng chi tiết hóa cách quản lý các checkpoint, xác minh rằng các trọng số được lưu ở định dạng safetensors của Hugging Face có thể được nạp lại mà không thiếu khóa nào, qua đó khẳng định tính toàn vẹn của quy trình huấn luyện.

Cơ chế hoạt động

Cơ chế cốt lõi dựa vào sự tích hợp giữa torchtitan và stack AWS Neuron. Khi torch_neuronx được cài đặt, nó sẽ hiển thị Trainium như một thiết bị neuron. Torchtitan sau đó sử dụng các nguyên thủy huấn luyện phân tán chuẩn của PyTorch, chẳng hạn như Fully Sharded Data Parallel (FSDP) và song song tensor, để phân bổ khối lượng công việc trên 64 NeuronCore logic có sẵn trong instance. Khung làm việc quản lý việc phân mảnh tham số mô hình, gradient và trạng thái optimizer, trong khi song song chuỗi (sequence parallelism) xử lý việc phân mảnh kích hoạt dọc theo chiều chuỗi.

Đối với việc tích hợp kernel tùy chỉnh, hướng dẫn minh họa cách sử dụng registry ghi đè (override registry) của torchtitan. Nhà phát triển có thể viết các kernel bằng NKI, ngôn ngữ kernel của Neuron, và bọc chúng thành các toán tử PyTorch bằng cách sử dụng wrap_nki và nki_op. Các toán tử này sau đó được đăng ký để thay thế các module cụ thể trong đồ thị mô hình. Ví dụ, một kernel rotary embedding tùy chỉnh có thể thay thế triển khai PyTorch mặc định. Hệ thống xác thực các ghi đè này khi khởi động, đảm bảo rằng mỗi lần thay thế nhắm đúng vào tên đầy đủ (fully qualified names) trong cấu trúc mô hình.

Vòng lặp huấn luyện báo cáo các chỉ số chính như loss, norm gradient, mức sử dụng bộ nhớ băng thông cao và tỷ lệ sử dụng FLOPs của mô hình. Khả năng hiển thị này cho phép các kỹ sư giám sát mức tiêu thụ tài nguyên và điều chỉnh các thiết lập song song một cách linh hoạt. Quá trình biên dịch diễn ra trong bước huấn luyện đầu tiên, mất nhiều thời gian hơn so với các bước tiếp theo, nhưng giúp tối ưu hóa việc thực thi cho phần còn lại của quá trình chạy.

Chi tiết quan trọng

• Thiết lập yêu cầu một instance trn2.48xlarge với 64 NeuronCore logic được cấu hình ở LNC=2.
• Các phụ thuộc phần mềm bao gồm PyTorch 2.14.0, Transformers 5.16.0 và torchtitan-neuron 0.2.2+neuron.1.
• Việc nạp checkpoint Hugging Face xác minh tính chính xác bằng cách kiểm tra xem loss ở bước 1 có gần giá trị 6.54 hay không đối với mô hình Qwen3-4B-Instruct-2507.
• Các kernel NKI tùy chỉnh phải tuân thủ các ràng buộc cụ thể về độ dài chuỗi, chẳng hạn như phải là bội số của 128 * LNC.
• Song song tensor phân mảnh trọng số của các lớp giữa các rank, trong khi song song dữ liệu phân mảnh tham số và trạng thái optimizer.
• Registry ghi đè ngăn ngừa xung đột bằng cách dừng lại nếu hai ghi đè cùng tuyên bố một node và cảnh báo nếu không có node nào khớp.

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

Đối với các kỹ sư phần mềm và chuyên gia ML, hướng dẫn này hạ thấp rào cản gia nhập cho việc huấn luyện các mô hình lớn trên phần cứng không phải NVIDIA. Bằng cách duy trì khả năng tương thích với các giao diện PyTorch tiêu chuẩn và hệ sinh thái Hugging Face, nó giảm bớt nhu cầu refactor mã nguồn mở rộng. Các đội ngũ có thể tận dụng các kho mô hình và định dạng checkpoint hiện có, giảm thiểu chi phí vận hành liên quan đến việc áp dụng các kiến trúc bộ tăng tốc mới.

Khả năng chèn các kernel NKI tùy chỉnh đặc biệt quan trọng đối với việc tối ưu hóa hiệu suất. Khi các mô hình ngày càng phức tạp, các triển khai chung có thể không tận dụng hết khả năng của phần cứng. Cho phép nhà phát triển thay thế các thao tác cụ thể, chẳng hạn như cơ chế attention hoặc rotary embeddings, bằng các kernel được tinh chỉnh thủ công giúp kiểm soát hiệu quả tính toán tốt hơn. Sự linh hoạt này rất cần thiết để đạt được thông lượng cạnh tranh và hiệu quả chi phí trong các công việc huấn luyện quy mô lớn.

Hơn nữa, các bước xác minh chi tiết cung cấp sự tự tin vào độ tin cậy của quy trình huấn luyện. Việc xác nhận rằng các checkpoint xuất khẩu được nạp đúng cách đảm bảo rằng các mô hình được huấn luyện trên Trainium có thể được triển khai liền mạch trong các môi trường sản xuất dựa trên các stack suy luận tiêu chuẩn. Việc xác thực end-to-end này hỗ trợ các thực hành MLOps vững chắc và thúc đẩy sự cộng tác giữa các đội ngũ sử dụng các backend phần cứng khác nhau.

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

• Thiết lập một instance trn2.48xlarge với AWS Deep Learning AMI dành cho Neuron hoặc container Neuron PyTorch.
• Cài đặt các phiên bản được chỉ định của torchtitan-neuron và các phụ thuộc bằng cách sử dụng tệp requirements.txt được cung cấp.
• Tải xuống các artifact mô hình Qwen3 từ Hugging Face Hub về một thư mục cục bộ để phục vụ việc huấn luyện.
• Cấu hình script huấn luyện để sử dụng phân mảnh song song dữ liệu trên tất cả 64 rank bằng cách đặt data_parallel_shard_degree = -1.
• Xác minh đầu ra của kernel tùy chỉnh so với các triển khai baseline bằng cách sử dụng đầu vào float32 để đảm bảo độ chính xác số học.
• Giám sát mức sử dụng bộ nhớ băng thông cao theo từng rank để điều chỉnh kích thước batch và độ dài chuỗi phù hợp.