Meta ra mắt MTIA 300: Chip huấn luyện AI đầu tiên tích hợp sẵn NIC và công cụ giảm tải truyền thông

Các mô hình deep learning có thể phân phối nội dung cá nhân hóa—từ video ngắn đến bài đăng của bạn bè—đến người dùng trên các ứng dụng. Khi các mô hình này ngày càng trở nên phức tạp, tầm quan trọng của năng lực tính toán để huấn luyện chúng và mạng lưới kết nối các bộ tăng tốc (accelerator) đó cũng tăng theo.

Huấn luyện các mô hình gợi ý là một thách thức về hạ tầng độc đáo. Không giống như các mô hình ngôn ngữ lớn (LLM) vốn cần lưu lượng tính toán dấu phẩy động khổng lồ, các mô hình gợi ý bị nghẽn bởi nhu cầu giao tiếp nhanh và hiệu quả giữa các bộ tăng tốc dùng để huấn luyện chúng. Các bảng embedding của chúng có thể chứa hơn 99% tham số của mô hình, đòi hỏi kỹ thuật song song hỗn hợp (hybrid parallelism) tạo ra các tập lệnh AllReduce, AllToAll và AllGather thường xuyên trên hàng trăm bộ tăng tốc. Trên các con chip như GPU, các thao tác giao tiếp này cạnh tranh tài nguyên với các tác vụ tính toán huấn luyện, thường dẫn đến việc phần cứng đắt tiền bị khai thác không hiệu quả.

Chúng tôi đã giải quyết thách thức này bắt đầu từ Meta Training and Inference Accelerator (MTIA), dòng chip AI tự phát triển của chúng tôi, với MTIA 300 – sản phẩm đầu tiên trong dòng MTIA được tối ưu hóa để huấn luyện các mô hình gợi ý và xếp hạng. Bằng cách đồng thiết kế MTIA 300 với HCCL, một thư viện giao tiếp được thiết kế song song với phần cứng ngay từ đầu, chúng tôi đã biến giao tiếp thành một thành phần ưu tiên trong thiết kế chip, thay vì là một yếu tố phụ được xử lý bởi các nhân tính toán đa năng.

Tích hợp mạng trực tiếp trên chip

Với MTIA 300, giao diện mạng nằm ngay bên trong gói chip (xem Hình 1). Hai chiplet mạng, mỗi chiplet chứa sáu NIC RDMA 800 Gbps tùy chỉnh, cung cấp tổng băng thông I/O là 1,2 TB/s mà không cần đi qua bus PCIe. Điều này loại bỏ nút thắt cổ chai giữa host-device-NIC vốn tồn tại trong các kiến trúc GPU truyền thống, nơi CPU phải làm trung gian giữa bộ tăng tốc và mạng. (Thông tin chi tiết về thiết kế silicon có sẵn trong bài báo gần đây của chúng tôi tại hội nghị ISCA 26).

Vì chúng tôi sử dụng cùng 12 NIC dựa trên Ethernet cho cả giao tiếp scale-up (trong một rack gồm 16 node, tốc độ lên tới 1 TB/s) và giao tiếp scale-out (giữa các rack, tốc độ 200 GB/s), chúng tôi có thể phân chia các NIC một cách linh hoạt để điều chỉnh theo nhu cầu thay đổi.

Khi các yêu cầu về mô hình thay đổi, chúng tôi có thể cấu hình lại sự phân chia này bằng cách cấu hình lại mạng thay vì thay đổi phần cứng. Để giảm thiểu độ trễ trên mỗi giao dịch, chúng tôi đã giới thiệu các "express doorbell". Bản thân yêu cầu ghi công việc đóng vai trò như một doorbell, loại bỏ một lần đọc bộ nhớ bổ sung và tiết kiệm khoảng 800 ns cho mỗi thao tác.

Giảm tải giao tiếp khỏi lưới tính toán

Trên GPU, các thư viện như NCCL thực thi giao tiếp tập thể dưới dạng các GPU kernel tiêu tốn tài nguyên của các streaming multiprocessor—cùng loại phần cứng cần thiết cho tính toán huấn luyện. Khi các tập lệnh giao tiếp và kernel huấn luyện chạy đồng thời, cả hai đều bị chậm lại.

MTIA 300 áp dụng một cách tiếp cận khác. Bên cạnh lưới 12×6 gồm các phần tử xử lý (PE) dành cho tính toán, con chip này còn bao gồm 16 công cụ thông điệp (ME) chuyên dụng xử lý tất cả các giao tiếp một cách độc lập. Mỗi ME bao gồm:

Được đặt ở các cạnh chip bên cạnh HBM và bộ nhớ đệm, các NMC cùng nhau cung cấp hơn 2,8 TB băng thông giảm thiểu (reduction throughput)—gấp đôi băng thông I/O—cho phép thực thi các tập lệnh AllReduce và ReduceScatter ở tốc độ đường truyền mà không cần chạm vào lưới tính toán.

Kết quả là sự cô lập gần như hoàn hảo. Việc chạy các tác vụ GEMM lớn đồng thời với các thao tác tập thể chỉ gây ra mức suy giảm hiệu suất tính toán dưới 0,5%, trái ngược với các GPU truyền thống có thể bị suy giảm hơn 20% do giao tiếp được xử lý bởi cùng một tài nguyên GPU.

Mô hình giao tiếp được biên dịch

Thư viện giao tiếp của chúng tôi, HCCL, được đồng thiết kế với MTIA 300. Thay vì điều khiển giao tiếp từ host trong quá trình thực thi, HCCL biên dịch mỗi tập lệnh tập thể thành một tập hợp hoàn chỉnh các subgraph—các mảng mục nhập hàng đợi công việc với các phụ thuộc rõ ràng—được gửi đến các ME để thực thi hoàn toàn tự động. Khi công việc đến thiết bị, host không còn tham gia vào quá trình này nữa. Hình 2 cho thấy cách CPU không còn tham gia sau khi sao chép các chỉ dẫn vào HBM.

Mô hình biên dịch này tích hợp tự nhiên với các giao diện c10d và torchcomms của PyTorch. Các tập lệnh tập thể được truy vết thông qua torch.compile sẽ được biên dịch thành một đồ thị duy nhất cùng với các toán tử tính toán. HCCL chọn các thuật toán nhận biết cấu trúc liên kết (topology-aware) để khai thác băng thông bất đối xứng giữa scale-up và scale-out, giảm thiểu lưu lượng giữa các rack nơi băng thông bị hạn chế. Đối với các tác vụ suy luận, chúng tôi đã phát triển các đường dẫn bổ sung: giao tiếp một chiều nơi các PE gửi công việc trực tiếp thông qua express doorbell, và các tập lệnh tập thể do thiết bị kích hoạt, nơi các kernel tính toán báo hiệu cho giao tiếp đã được giảm tải phần cứng trên một luồng song song mà không làm gián đoạn việc thực thi đồ thị.

Hiệu suất trong thực tế sản xuất

HCCL đạt băng thông giao tiếp lên tới 940 GB/s trong một rack duy nhất. Trên một mô hình gợi ý sản xuất với 150 tỷ tham số chạy trên 40 bộ tăng tốc, tổng thời gian giao tiếp của MTIA 300 nhanh gấp 3,9 lần so với cụm GPU tương đương.

Thiết kế của MTIA 300 cho phép các chiến lược đồng thiết kế sâu hơn: 216 GB HBM3E cho phép kích thước batch cục bộ lớn hơn (giảm số lượng bộ huấn luyện và chi phí giao tiếp); tỷ lệ 1:1 giữa CPU và bộ tăng tốc cho phép CPU giảm tải các thao tác tối ưu hóa đòi hỏi tính toán số học chuyên sâu; và băng thông mạng cao cho phép chúng tôi sử dụng các kiểu dữ liệu có độ chính xác cao hơn để duy trì độ chính xác.

Hướng tới tương lai

Mặc dù MTIA 300 được thiết kế để huấn luyện các mô hình gợi ý, các nguyên tắc kiến trúc—mạng tích hợp, thực thi giao tiếp được giảm tải và đồng thiết kế cấp hệ thống giữa tính toán và giao tiếp—giúp nó sẵn sàng cho một tập hợp các tác vụ rộng lớn hơn. Khi suy luận AI phát triển theo hướng suy luận (reasoning), tác nhân (agentic) và các trường hợp sử dụng ngữ cảnh dài, nhu cầu giao tiếp sẽ thay đổi: Các thông điệp trở nên nhỏ hơn, thường xuyên hơn và nhạy cảm với độ trễ, với ngân sách khắt khe hơn cho mỗi tập lệnh tập thể.

Một kiến trúc coi mạng lưới là một ràng buộc hệ thống ưu tiên, tối ưu hóa không chỉ băng thông mà còn cả độ trễ và tốc độ thông điệp, rất phù hợp để đáp ứng những nhu cầu mới nổi này. Các mô hình được thiết lập trong MTIA 300 và HCCL là nền tảng cho thế hệ silicon AI tiếp theo của Meta.

Tìm hiểu thêm về MTIA 300

Để tìm hiểu thêm về thiết kế silicon của MTIA 300 và các công việc được trình bày chi tiết tại đây, hãy đọc các bài báo của chúng tôi: