Amazon đã công bố TorchNeuron, một chương trình phụ trợ PyTorch nguồn mở dành cho dòng chip AI Trainium của mình, trong một động thái nhằm loại bỏ một trong những rào cản thực tế lớn nhất trong việc áp dụng silicon tùy chỉnh cho khối lượng công việc học máy. Theo một bài đăng chính thức trên blog PyTorch được xuất bản trước Hội nghị PyTorch Bắc Mỹ vào ngày 20 tháng 10, các nhà phát triển hiện có thể chạy mã mô hình PyTorch hiện có trên Trainium chỉ với một thay đổi một dòng, thay vì viết lại quy trình công việc của họ xung quanh công cụ độc quyền.

Thông báo này đã thu hút sự chú ý ngay lập tức vì nó nhắm đến điều mà các nhà phân tích từ lâu đã mô tả là lợi thế lâu bền nhất của Nvidia: không phải hiệu suất chip thô mà là hệ sinh thái phần mềm CUDA gắn kết các nhà phát triển với GPU của họ. Reuters đưa tin về việc phát hành khi Amazon cho biết các chip AI của họ sẽ hoạt động với khung phát triển phổ biến nhất trong AI, trong khi TechCrunch gọi TorchNeuron là cú hích lớn nhất của công ty nhằm vào Nvidia. Để có cái nhìn rộng hơn về cách thị trường máy gia tốc đang thay đổi, tin tức về phần cứng AI đang diễn ra của chúng tôi theo dõi những diễn biến này khi chúng diễn ra.

TorchNeuron thực sự thay đổi điều gì

Lời hứa cốt lõi của TorchNeuron là sự đơn giản. Trong ví dụ được đăng trên blog PyTorch, một vòng đào tạo hiện có chạy trên Trainium không gì khác hơn là chuyển mô hình sang thiết bị mới bằng `.to('neuron')`. Theo bài đăng, mã được viết dựa trên `torch.accelerator`, API không xác định thiết bị được đề xuất của PyTorch, chạy chưa được sửa đổi. Đối với khối lượng công việc được biên dịch, nhà phát triển đặt phần phụ trợ mặc định một lần bằng `torch.compiler.set_default_backend("neuron")` rồi gọi `model.compile()` mà không thay đổi gì thêm.

Dưới vỏ bọc, TorchNeuron đăng ký Trainium thông qua PrivateUse1, điểm mở rộng chính thức của PyTorch cho các chương trình phụ trợ phần cứng tùy chỉnh, có nghĩa là chip hoạt động như một thiết bị PyTorch tiêu chuẩn chứ không phải là một phần phụ được gắn vào khung. AWS mô tả việc tích hợp bao gồm nhiều chế độ thực thi:

  • Háo hức thực thi để nghiên cứu và gỡ lỗi, với "chế độ háo hức thích ứng" thử nghiệm giúp giảm chi phí khởi chạy
  • torch.compile để cải thiện hiệu suất sản xuất, thu thập các mô hình thành biểu đồ mà trình biên dịch Neuron tối ưu hóa, với sự hỗ trợ riêng cho FlexAttention
  • Giao diện hạt nhân Neuron (NKI) dành cho các nhà nghiên cứu cần hạt nhân cấp thấp, tùy chỉnh

Theo thông báo, phần phụ trợ này hoạt động trên cả chip Trainium2 và chip Trainium3 mới hơn.

Một đòn tấn công trực tiếp vào hào phần mềm của Nvidia

Sự nắm bắt của Nvidia đối với điện toán AI chưa bao giờ chỉ dựa vào phần cứng nhanh hơn. CUDA, nền tảng phần mềm có tuổi đời hai thập kỷ của công ty, là nơi hầu hết các mã nghiên cứu AI được viết, điều chỉnh và triển khai, và mọi con chip thách thức — từ dòng Instinct của AMD đến TPU của Google và Trainium của Amazon — đều phải đối mặt với chi phí yêu cầu các nhà phát triển rời bỏ nó. Tin tức về buổi ra mắt TorchNeuron trên GeekWire, Benzinga và GuruFocus đều đóng khung thông báo như một cuộc tấn công vào con hào đó.

Amazon không cố gắng kéo các nhà phát triển ra khỏi PyTorch để thu phục họ. Thay vào đó, nó đưa Trainium vào PyTorch, để cơ sở người dùng khổng lồ của khung có thể sử dụng chip mà không cần thay đổi thói quen. Bài đăng trên blog PyTorch bao gồm lời xác nhận từ Meta, nơi duy trì khuôn khổ: "Hỗ trợ gốc cho AWS Trainium mở rộng lựa chọn đó, cho phép các nhà nghiên cứu và kỹ sư thử nghiệm và lặp lại bằng cách sử dụng các API phân tán và thực thi háo hức mà họ đã biết", Jana van Greunen, Giám đốc Kỹ thuật PyTorch tại Meta, cho biết trong bài đăng.

Khung đó rất quan trọng đối với hoạt động kinh doanh trên nền tảng đám mây của AWS. Phiên bản Trainium đã được tiếp thị như một giải pháp thay thế chi phí thấp hơn cho GPU Nvidia để đào tạo và suy luận, đồng thời các con chip này đã thu hút được nhiều khách hàng AI. Làm cho các con chip có cảm giác tự nhiên đối với PyTorch sẽ trực tiếp giải quyết vấn đề phản đối đã làm chậm quá trình áp dụng rộng rãi hơn: nỗ lực chuyển đổi.

Phạm vi kỹ thuật, từ sổ ghi chép đến đào tạo quy mô cụm

Các khả năng được liệt kê trên blog PyTorch cho thấy Amazon đang hướng tới mục tiêu vượt xa các thử nghiệm mang tính sở thích. TorchNeuron bao gồm các bộ toán tử được thấy trên các mô hình văn bản, đa phương thức và ngôn ngữ tầm nhìn cũng như tạo hình ảnh và video, mô hình không gian trạng thái và mô hình thế giới, với độ chính xác BF16, FP32 và FP16; Hỗ trợ TorchAO và FP8 đang ở giai đoạn thử nghiệm riêng tư trước khi có sẵn rộng rãi trong phiên bản tiếp theo.

Đào tạo phân tán, nguyên tắc biến mã GPU đơn thành các hoạt động quy mô cụm, được hỗ trợ thông qua API PyTorch tiêu chuẩn — DDP, FSDP2, SimpleFSDP, tensor và pipe song song cũng như phân mảnh kiểu ZeRO, cùng với DTensor và DeviceMesh cho các nhóm quy trình đa chiều. Bài đăng cho biết các luồng và sự kiện ánh xạ ngữ nghĩa của PyTorch vào hàng đợi phần cứng của Trainium.

Để suy luận, các ngăn xếp Neuron của Amazon được xây dựng trên cùng một nền tảng: vLLM Neuron phục vụ các mô hình ngôn ngữ lớn, trong khi vLLM Omni Neuron phục vụ các mô hình khuếch tán để tạo hình ảnh và video. Giao diện hạt nhân Neuron cho phép các nhà nghiên cứu viết các toán tử tùy chỉnh cho cả quá trình thực thi mong muốn và được biên dịch, đồng thời AWS là nguồn mở NKI-lib, một tập hợp các hạt nhân được tối ưu hóa Trainium bao gồm các tính toán theo khối FlashAttention và MoE cho các kịch bản LLM chính thống và đa phương thức.

Công cụ hiệu suất cũng được bao gồm. Neuron Explorer mở rộng `torch.profiler` tiêu chuẩn để theo dõi quá trình thực thi từ khung đến phần cứng, liên kết mọi sự kiện trở lại dòng nguồn của nó và có thể xuất các số liệu ở định dạng Parquet mà các tác nhân có thể sử dụng ghi chú bài đăng để phân tích tắc nghẽn tự động.

Tính khả dụng và điều gì sẽ xảy ra tiếp theo

Theo bài đăng trên blog, TorchNeuron được thiết lập để ra mắt nguồn mở tại PyTorch Conference Bắc Mỹ vào ngày 20 tháng 10, theo bài đăng trên blog, với khả năng tương thích mở rộng cho Hugging Face Transformers v5. AWS mô tả Trainium là một danh mục chip được thiết kế từ đầu để đào tạo và suy luận AI trên quy mô đám mây.

Bối cảnh chiến lược là một thị trường trong đó mọi nhà cung cấp đám mây lớn đều đồng thời là người mua lớn nhất các bộ tăng tốc Nvidia và là đối thủ cạnh tranh tiềm năng. Microsoft, Google và Meta đều phát triển silicon tùy chỉnh và các công ty Trung Quốc bao gồm cả Huawei đang xây dựng kho phần mềm của riêng họ để thoát khỏi lực hấp dẫn của CUDA. Bằng cách gặp gỡ các nhà phát triển bên trong PyTorch thay vì yêu cầu họ di chuyển, Amazon đã chọn con đường ít gây gián đoạn nhất cho cuộc thi đó — và con đường có nhiều khả năng nhất, nếu kinh nghiệm của nhà phát triển tiếp tục, sẽ chuyển khối lượng công việc đào tạo thực sự ra khỏi phần cứng Nvidia.

Liệu TorchNeuron có thực sự làm xói mòn lợi thế của CUDA hay không sẽ phụ thuộc vào hiệu suất trên quy mô lớn và sự tồn tại lâu dài của các mô hình trong thế giới thực, những câu hỏi mà vật liệu phóng không trả lời được. Nhưng rào cản khiến hầu hết các nhà phát triển PyTorch không thể thử Trainium - một bản viết lại - giờ đây đã được hạ xuống chỉ còn một dòng mã.

---

Đi trước AI

Nhận tin tức, phân tích và đột phá mới nhất về AI — tất cả ở cùng một nơi.

Đọc thêm tin tức về AI →