MetaRoCE: Giao thức truyền tải RDMA mới tối ưu cho hạ tầng AI quy mô lớn

Tại Meta, chúng tôi là động lực mạnh mẽ đằng sau sự đồng thuận ngày càng tăng của ngành rằng Ethernet nên là hạ tầng mạng (fabric) được lựa chọn cho cơ sở hạ tầng AI. Chúng tôi đã chứng minh rằng RoCE có thể hỗ trợ huấn luyện AI phân tán ở quy mô lớn. Hiện tại, chúng tôi đang phát triển dựa trên nền tảng đó với MetaRoCE, một giao thức được thiết kế từ đầu cho Ethernet ở quy mô hàng triệu GPU.

Chúng tôi đã mở rộng các cụm hàng trăm nghìn GPU, trải rộng trên nhiều trung tâm dữ liệu và khu vực. Cho dù các cụm này đang huấn luyện mô hình tiên phong tiếp theo hay phục vụ suy luận ở quy mô toàn cầu, mạng lưới luôn nằm trong lộ trình quan trọng.

Các hoạt động tập thể (collective operations) như all-reduce và all-to-all đồng bộ hóa hàng nghìn bộ tăng tốc trong quá trình huấn luyện, và việc truyền dữ liệu chậm nhất sẽ quyết định tốc độ cho toàn bộ công việc. Trong suy luận, giao tiếp độ trễ thấp giữa các phân đoạn mô hình phân tán ảnh hưởng trực tiếp đến thời gian phản hồi cho hàng trăm triệu người dùng. Ngay cả một lượng nhỏ ma sát mạng cũng trực tiếp làm lãng phí đáng kể năng lực tính toán.

RoCE tiêu chuẩn yêu cầu mạng phải phân phối mọi khung dữ liệu theo đúng thứ tự, tận dụng PFC và hạn chế việc rải gói tin (packet spraying) vốn mang lại hiệu suất trong các mạng đa mặt phẳng (multiplane) và quy mô lớn. MetaRoCE được xây dựng để cung cấp băng thông cao, độ trễ đuôi thấp và sự đơn giản trong vận hành khi mạng lưới tăng lên về số lượng bộ tăng tốc và khoảng cách giữa chúng.

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

Thông tin cốt lõi của MetaRoCE rất đơn giản: Fabric nhìn thấy các gói tin, nhưng NIC nhìn thấy mục đích. Các kiến trúc truyền thống tập trung trí thông minh vào fabric, dựa vào các switch để thực thi tính không mất dữ liệu và duy trì thứ tự.

Bằng cách chuyển trí thông minh sang điểm cuối (endpoint), MetaRoCE phân tách mạng thành nhiều đường dẫn logic chi tiết, mỗi đường dẫn có phép đo từ xa (telemetry) thời gian thực riêng – RTT theo từng đường dẫn, trạng thái ECN và mức độ sử dụng. Khả năng hiển thị này mở ra các tính năng khó đạt được với RDMA truyền thống.

Phân phối không theo thứ tự (Out-of-Order Delivery) tự nhiên

MetaRoCE rải các gói tin qua nhiều đường dẫn, vì vậy chúng đến nơi không theo thứ tự ngay từ thiết kế. Lớp truyền tải coi việc đến không theo thứ tự là trường hợp bình thường. Mỗi gói tin mang theo đích đến của riêng nó, vì vậy dữ liệu được ghi thẳng vào vị trí bộ nhớ cuối cùng khi nó đến nơi, không cần bộ đệm sắp xếp lại (reorder buffer) và không bị chặn đầu dòng (head-of-line blocking).

Các lệnh ghi mang theo đích đến trong mỗi gói tin. Các lệnh gửi (Send) mang theo sự khớp với bộ đệm nhận đã đăng ký, vì vậy một lệnh Send sẽ đến đúng vị trí ngay cả khi các thông điệp phía trước nó chưa đến, và không cần một vòng lặp (round trip) để tìm hiểu dữ liệu đi đâu. Các thư viện tập thể có thể sử dụng nhắn tin hai chiều (two-sided messaging) ở nơi phù hợp thay vì giảm mọi thứ thành lệnh Write.

Đa đường dẫn (Multipathing) tự nhiên

MetaRoCE cung cấp cho mỗi kết nối các đường dẫn hạng nhất và rải gói tin trên đó theo từng gói. Mỗi đường dẫn mang một cổng nguồn UDP riêng biệt làm entropy ECMP, mà NIC có thể thay đổi bất cứ lúc nào để chuyển lưu lượng ra khỏi tuyến đường xấu. Trên các fabric đa mặt phẳng, việc chọn mặt phẳng hoàn toàn thuộc về NIC, và fabric chỉ được sử dụng tốt như cách NIC rải gói tin. Vì mỗi đường dẫn giữ cửa sổ và ước tính vòng lặp riêng, lớp truyền tải có thể phân biệt tắc nghẽn với lỗi và tái cân bằng một cách rõ ràng, do đó một liên kết nóng hoặc bị hỏng chỉ làm chậm một đường dẫn thay vì làm đình trệ toàn bộ kết nối.

Khả năng chịu lỗi theo thiết kế

MetaRoCE coi fabric Ethernet là loại có thể mất dữ liệu và không yêu cầu nó phải khác đi – không PFC, không khung tạm dừng (pause frames). Vì mỗi đường dẫn mang theo chuỗi thứ tự riêng, một khoảng trống trong bitvector xác nhận chọn lọc (selective acknowledgment) 256-bit của nó là bằng chứng của việc mất dữ liệu thay vì sắp xếp lại. Trong các giao thức khác, SACK chủ yếu tránh gửi lại dữ liệu đã đến; ở đây, nó kích hoạt truyền lại chính xác gói tin bị thiếu, trên đường dẫn đã làm mất nó, ngay khi khoảng trống xuất hiện.

Kiểm soát tắc nghẽn từ cả hai phía

MetaRoCE kết hợp kiểm soát tắc nghẽn AIMD dựa trên ECN thông thường, do người gửi điều khiển với các gợi ý tốc độ chia sẻ công bằng do người nhận điều khiển. Các cửa sổ được giữ theo từng đường dẫn cũng như theo từng kết nối, vì vậy một dấu hiệu tắc nghẽn sẽ cắt giảm đường dẫn đã thấy nó và hướng các gói tin tiếp theo về phía các đường dẫn thông thoáng. Trong mỗi xác nhận, người nhận trả về phần băng thông đầu vào mà nó đã phân bổ cho người gửi đó, vì vậy người gửi tiếp cận tốc độ phù hợp trực tiếp thay vì phải tìm kiếm nó. Incast được giải quyết trong một hoặc hai vòng lặp, với sự công bằng tốt hơn và độ trễ đuôi thấp hơn.

Độc lập với cấu trúc liên kết (Topology)

MetaRoCE yêu cầu fabric hai thứ mà mọi switch đều đã có: đánh dấu ECN và ECMP. Nó không yêu cầu cắt gói tin, đo từ xa trong mạng, kiểm soát luồng dựa trên tín dụng hoặc rải gói tin phía switch, và nó không bị lỗi khi fabric cung cấp các tính năng đó. Cùng một lớp truyền tải chạy trên các fabric fat-tree, đa mặt phẳng, bộ đệm sâu và bộ đệm nông, và trên các đám mây của nhà cung cấp mà bạn không kiểm soát cấu hình. Không có gì độc quyền liên quan, vì vậy fabric vẫn tự do tối ưu hóa chi phí và hệ thống cáp.

Kết nối hợp nhất ở quy mô lớn

Một cặp hàng đợi (QP) mang cả luồng thông điệp có thứ tự và băng thông. RDMA truyền thống có được nhiều hơn một trong hai bằng cách mở thêm nhiều QP (hàng chục trên mỗi cặp nút), mỗi QP có một cửa sổ tắc nghẽn không biết gì về các QP còn lại và trạng thái riêng trên NIC.

MetaRoCE tách biệt hai yếu tố này. Một kết nối duy nhất mang nhiều luồng có thứ tự độc lập ở phía trên, mỗi luồng cho một communicator hoặc collective, và nhiều đường dẫn ở phía dưới, dưới một bộ điều khiển tắc nghẽn. Trạng thái kết nối không còn tăng theo tính song song của khối lượng công việc.

Lớp ứng dụng hầu như không bị ảnh hưởng – các API RDMA Verbs và ngăn xếp phần mềm hiện có hoạt động mà không cần sửa đổi. Các tính năng nâng cao như hỗ trợ đa mặt phẳng được hỗ trợ thông qua các API mở rộng.

MetaRoCE trong thực tế

Để tăng tốc xác thực phần cứng, chúng tôi đã làm việc với AMD để triển khai MetaRoCE trên các NIC lập trình được Pensando của họ.

Trên một cụm 64 GPU AMD chạy các tập thể RCCL, chúng tôi đã so sánh trực tiếp MetaRoCE với RoCEv2 trên các hoạt động all-reduce và all-to-all. Kết quả nhất quán với các mục tiêu thiết kế:

MetaRoCE liên tục mang lại băng thông cao hơn và thời gian hoàn thành luồng thấp hơn so với RoCEv2.

Trong các điều kiện mất gói tin vốn sẽ làm giảm hiệu suất của RoCEv2, MetaRoCE duy trì ~86% băng thông ở mức mất gói tin 1% và tiếp tục cung cấp băng thông hữu ích ngay cả ở tỷ lệ mất gói tin cực đoan 10% – hội tụ một cách ổn định thay vì sụp đổ.

Xác thực đa mặt phẳng trên các cấu trúc liên kết 4 mặt phẳng và 8 mặt phẳng với tối đa 4.000 kết nối đồng thời xác nhận rằng băng thông tăng tuyến tính theo số lượng mặt phẳng.

Trong các mô phỏng lỗi mặt phẳng, giao thức thể hiện khả năng phục hồi tự động ổn định – lưu lượng tự phân phối lại mà không cần sự tham gia của ứng dụng hoặc sự can thiệp của người vận hành.

Những kết quả này củng cố lựa chọn thiết kế cốt lõi của MetaRoCE. Bằng cách thiết kế để chịu được mất mát ngay từ ngày đầu và đẩy trí thông minh ra biên, bạn có được một lớp truyền tải hoạt động tốt hơn trong điều kiện lý tưởng và suy giảm một cách ổn định khi có sự cố xảy ra.

Mở theo thiết kế

Cơ sở hạ tầng AI được hưởng lợi từ các tiêu chuẩn chung giúp tăng tốc đổi mới trên toàn hệ sinh thái. MetaRoCE mở rộng cùng một triết lý mở, đa nhà cung cấp mà sáng kiến Ethernet Scalable Unified Network (ESUN) của Open Compute Project (OCP) đã thiết lập cho fabric vào lớp truyền tải.

Đó là lý do tại sao chúng tôi mở MetaRoCE:

Đặc tả mở thông qua OCP: Đặc tả giao thức đầy đủ đang được đóng góp cho OCP, sẵn sàng cho bất kỳ nhà cung cấp nào triển khai và xây dựng phần cứng có khả năng tương tác.

Nhiều triển khai NIC: MetaRoCE được thiết kế để chạy trên các kiến trúc NIC đa dạng – cả lập trình được và chức năng cố định. Chúng tôi đã chứng minh điều đó trên phần cứng AMD Pensando, với các triển khai bổ sung đang được thực hiện từ các nhà cung cấp khác.