Mạng RoCE cho huấn luyện AI phân tán quy mô lớn tại Meta

Giới thiệu

Trong những năm qua, Meta đã xây dựng một trong những mạng lưới AI lớn nhất thế giới để hỗ trợ khối lượng công việc huấn luyện AI phân tán quy mô lớn. Bài báo này sẽ thảo luận chi tiết về cách chúng tôi thiết kế, triển khai và vận hành mạng lưới này, tập trung vào việc sử dụng RDMA over Converged Ethernet version 2 (RoCEv2) làm phương thức truyền thông liên nút cho phần lớn năng lực AI của chúng tôi.

Sự phổ biến ngày càng tăng của AI đã tạo ra một kỷ nguyên mới về nhu cầu truyền thông. Huấn luyện phân tán, đặc biệt, tạo ra sức ép đáng kể nhất lên cơ sở hạ tầng mạng của trung tâm dữ liệu. Ví dụ, một tác vụ AI tạo sinh (GenAI) điển hình có thể cần phối hợp chặt chẽ hàng chục nghìn GPU trong vài tuần. Việc xây dựng một cơ sở hạ tầng mạng đáng tin cậy, hiệu suất cao, có khả năng đáp ứng nhu cầu ngày càng tăng này, đòi hỏi phải đánh giá lại thiết kế mạng trung tâm dữ liệu.

Khi Meta giới thiệu phương pháp huấn luyện dựa trên GPU phân tán, chúng tôi đã quyết định xây dựng các mạng trung tâm dữ liệu chuyên biệt được thiết kế riêng cho các cụm GPU này. Chúng tôi đã chọn RoCEv2 làm phương thức truyền thông liên nút cho phần lớn năng lực AI của mình.

Chúng tôi đã mở rộng thành công mạng RoCE của mình, phát triển từ nguyên mẫu sang triển khai nhiều cụm, mỗi cụm chứa hàng nghìn GPU. Các cụm RoCE này hỗ trợ một loạt các công việc huấn luyện GPU phân tán trong sản xuất, bao gồm xếp hạng, đề xuất nội dung, hiểu nội dung, xử lý ngôn ngữ tự nhiên và huấn luyện mô hình GenAI, cùng với các khối lượng công việc khác.

Kiến trúc mạng

Sự phân tách

Cụm huấn luyện dựa trên hai mạng độc lập: mạng frontend (FE) cho các tác vụ như nhập dữ liệu, lập điểm kiểm tra và ghi nhật ký và mạng backend (BE) để huấn luyện, như được mô tả bên dưới.

Một rack huấn luyện được kết nối với cả FE và BE của mạng trung tâm dữ liệu. FE có hệ thống phân cấp các tầng mạng – bộ chuyển mạch rack (RSW), bộ chuyển mạch fabric (FSW), v.v. – chứa kho lưu trữ, cung cấp cho GPU dữ liệu đầu vào cần thiết cho khối lượng công việc huấn luyện. Chúng tôi đảm bảo rằng có đủ băng thông ingress trên bộ chuyển mạch rack để không cản trở khối lượng công việc huấn luyện.

BE là một kết cấu chuyên biệt kết nối tất cả các NIC RDMA trong một kiến trúc không chặn, cung cấp băng thông cao, độ trễ thấp và truyền tải không mất dữ liệu giữa hai GPU bất kỳ trong cụm, bất kể vị trí thực của chúng. Kết cấu phụ trợ này sử dụng giao thức RoCEv2, đóng gói dịch vụ RDMA trong các gói UDP để truyền tải qua mạng.

Khu vực AI

Các mạng BE của chúng tôi đã trải qua một số thay đổi. Ban đầu, các cụm GPU của chúng tôi sử dụng cấu trúc liên kết hình sao đơn giản với một vài rack AI được kết nối với bộ chuyển mạch Ethernet trung tâm chạy giao thức RoCEv1 không định tuyến được. Thiết lập này có những hạn chế rõ ràng về quy mô GPU và khả năng dự phòng chuyển đổi. Do đó, chúng tôi đã nhanh chóng chuyển sang kiến trúc dựa trên kết cấu để mở rộng khả năng mở rộng và tính khả dụng cao hơn.

Chúng tôi đã thiết kế cấu trúc liên kết Clos hai tầng cho các rack AI, được gọi là Khu vực AI. Bộ chuyển mạch huấn luyện rack (RTSW), đóng vai trò là bộ chuyển mạch lá, cung cấp khả năng kết nối mở rộng cho GPU trong rack bằng cáp DAC dựa trên đồng. Tầng spine, bao gồm các bộ chuyển mạch huấn luyện cụm mô-đun (CTSW), cung cấp khả năng kết nối mở rộng giữa tất cả các rack trong cụm. CTSW có bộ đệm sâu được chia tĩnh trên các cổng trong khung. Các RTSW kết nối với CTSW thông qua cáp quang đơn mode và bộ thu phát có thể cắm 400G.

Các Khu vực AI được thiết kế để hỗ trợ một lượng lớn GPU được kết nối với nhau một cách không chặn. Tuy nhiên, những tiến bộ mới nổi của AI, chẳng hạn như LLM như Llama, yêu cầu quy mô GPU lớn hơn mức mà một Khu vực AI cung cấp. Để thích ứng với điều này, chúng tôi đã thiết kế một tầng bộ chuyển mạch huấn luyện tổng hợp (ATSW) kết nối các CTSW trong tòa nhà trung tâm dữ liệu, mở rộng miền RoCE ra ngoài một Khu vực AI duy nhất.

Lưu ý, kết nối chéo Khu vực AI được đăng ký quá mức theo thiết kế, với lưu lượng truy cập mạng được cân bằng bằng ECMP. Để giảm thiểu nút thắt cổ chai hiệu suất đối với lưu lượng truy cập chéo Khu vực AI, chúng tôi đã cải thiện trình lên lịch tác vụ huấn luyện để tìm lát cắt tối thiểu khi chia các nút huấn luyện thành các Khu vực AI khác nhau, giảm lưu lượng truy cập chéo Khu vực AI và do đó rút ngắn thời gian hoàn thành chung. Trình lên lịch thực hiện việc này bằng cách tìm hiểu vị trí của máy chủ GPU trong cấu trúc liên kết logic để đề xuất phân bổ hạng.

Định tuyến

Việc mở rộng quy mô sức mạnh tính toán và cấu trúc liên kết mạng đã thảo luận ở trên dẫn đến câu hỏi làm thế nào để cân bằng và định tuyến hiệu quả lưu lượng truy cập huấn luyện khổng lồ. Cụ thể, khối lượng công việc huấn luyện AI có một số đặc điểm đầy thách thức:

Độ hỗn loạn thấp (Low entropy): So với khối lượng công việc của trung tâm dữ liệu truyền thống, số lượng và sự đa dạng của luồng cho khối lượng công việc AI nhỏ hơn nhiều và các mẫu luồng thường lặp đi lặp lại và có thể dự đoán được.
Tính bùng nổ (Burstiness): Trên khía cạnh thời gian, các luồng thường thể hiện bản chất bật và tắt trong độ chi tiết thời gian tính bằng mili giây.
Luồng voi (Elephant flows): Đối với mỗi lần bùng nổ, cường độ của mỗi luồng có thể đạt tới tốc độ đường truyền của NIC.

ECMP và ghim đường dẫn

Ban đầu, chúng tôi đã xem xét ECMP được áp dụng rộng rãi, đặt các luồng một cách ngẫu nhiên dựa trên hàm băm trên năm thành phần: IP nguồn và đích, cổng UDP nguồn và đích và giao thức. Tuy nhiên, và như mong đợi, ECMP đã mang lại hiệu suất kém cho khối lượng công việc huấn luyện do độ hỗn loạn luồng thấp.

Ngoài ra, chúng tôi đã thiết kế và triển khai lược đồ ghim đường dẫn trong những năm đầu triển khai. Lược đồ này định tuyến các gói tin đến các đường dẫn cụ thể dựa trên lát đích (chỉ số của đường xuống RTSW). Điều này hoạt động tốt nếu mỗi rack được chỉ định đầy đủ cho cùng một công việc và không có lỗi nào trong mạng. Tuy nhiên, điều này hiếm khi đúng. Chúng tôi thấy rằng rack có thể được phân bổ một phần cho một công việc, chỉ có một trong hai máy chủ trong rack sử dụng băng thông đường lên. Việc vị trí công việc bị phân mảnh này gây ra phân phối lưu lượng không đồng đều và tắc nghẽn trên đường lên của RTSW cụ thể và làm giảm hiệu suất đào tạo lên tới hơn 30%. Hơn nữa, lỗi mạng trên đường lên hoặc CTSW khiến các luồng bị ảnh hưởng được chỉ định lại không đồng đều cho các CTSW khác bởi ECMP. Những luồng được chỉ định lại đó đã va chạm với các luồng hiện có khác và làm chậm toàn bộ công việc huấn luyện.

Chúng tôi đã giảm thiểu tác động tức thời của các va chạm luồng này bằng cách nâng cấp băng thông của băng thông đường lên RTSW lên gấp 2 lần. Do đó, chúng tôi đã cho phép dung lượng đường lên RTSW được đăng ký chưa đủ 1:2 so với dung lượng đường xuống RTSW. Mặc dù điều này đã giảm thiểu tác động đến hiệu suất tức thời, nhưng đây là một giải pháp tốn kém vì nó yêu cầu dung lượng mạng gấp 2 lần. Do đó, chúng tôi đã nhận ra đây là một biện pháp giảm thiểu ngắn hạn và tiến hành các giai đoạn tiếp theo của quá trình phát triển định tuyến.

Cân bằng hàng đợi (Queue pair scaling)

Tiếp theo, chúng tôi đã xem xét lại ECMP với mục đích tăng số lượng luồng cho các tập thể phân cấp thông qua tính năng phần mềm mở rộng quy mô cặp hàng đợi (QP) trong thư viện tập thể.

Để giải thích cho điều này, chúng tôi đã định cấu hình các bộ chuyển mạch để thực hiện ECMP nâng cao (E-ECMP) để băm thêm vào trường QP đích của gói RoCE bằng cách sử dụng khả năng UDF của ASIC chuyển đổi. Điều này đã làm tăng entropy và, so với ECMP cơ sở mà không cần mở rộng quy mô QP, chúng tôi quan sát thấy rằng E-ECMP cùng với việc mở rộng quy mô QP cho thấy hiệu suất được cải thiện lên đến 40% cho tập thể AllReduce.

Chúng tôi đã đánh giá hai chiến lược mở rộng quy mô QP. Cách đầu tiên liên quan đến việc chia mỗi thông báo có nghĩa là được đăng trên một QP duy nhất, thay vào đó là nhiều QP dẫn đến nhiều luồng. Nhưng nó cũng tạo ra kích thước thông báo nhỏ hơn trên kết cấu cũng như nhiều ACK. Cách tiếp cận thứ hai liên quan đến việc đăng từng thông báo vào một hàng đợi khác nhau, theo kiểu Round-robin. Đối với kích thước thông báo NIC được thể hiện trong quá trình sản xuất của chúng tôi với NCCL, chúng tôi quan sát thấy cách thứ hai hoạt động tốt. Tính năng này rất quan trọng đối với khả năng mở rộng ECMP bằng cách tăng lưu lượng mạng cho các tập thể phân cấp như AllReduce.

Mặc dù chúng tôi đã cải thiện hiệu suất ECMP với việc mở rộng quy mô QP, nhưng bản chất xác suất cơ bản của hàm băm là nhược điểm dai dẳng của lược đồ định tuyến này. Ngoài ra, nhu cầu tùy chỉnh hệ số масштабирования QP và phương pháp dựa trên loại khối lượng công việc, trong khi có thể hoạt động trong ngắn hạn, đã tạo ra sự phức tạp về hoạt động lâu dài.

Kiểm soát tắc nghẽn

DCQCN

Khi chúng tôi chuyển sang triển khai 400G, chúng tôi đã cố gắng điều chỉnh DCQCN để điều chỉnh tốc độ mạng và cấu trúc liên kết mới. Tuy nhiên, với cài đặt DCQCN mặc định và ngưỡng ECN tăng gấp đôi so với mạng 200G, hiệu suất đã bị giảm sút. Điều tra thêm cho thấy việc triển khai DCQCN trong phần sụn đã thay đổi, tạo ra lỗi và giảm khả năng hiển thị với các sự cố liên quan đến việc đếm CNP chính xác.

Chúng tôi đã tiến hành mà không có DCQCN cho các triển khai 400G của chúng tôi. Tại thời điểm này, chúng tôi đã có hơn một năm kinh nghiệm chỉ với PFC để kiểm soát luồng, không có bất kỳ kiểm soát tắc nghẽn cấp độ truyền tải nào khác. Chúng tôi đã quan sát thấy hiệu suất ổn định và không có tắc nghẽn liên tục đối với các tập thể huấn luyện.

Nhập lưu lượng hướng người nhận

Để giảm thiểu tắc nghẽn cho 400G trở lên, chúng tôi đã đồng thiết kế thư viện tập thể và truyền tải RoCE để thực thi quyền truy cập lưu lượng do người nhận điều khiển để có hiệu suất tốt hơn. Biểu đồ bên dưới cho thấy kiến trúc giao tiếp GPU-to-GPU trong các cụm huấn luyện sản xuất của chúng tôi chủ yếu sử dụng sao chép hai giai đoạn và giao tiếp do người nhận khởi tạo thông qua thư viện tập thể NCCL. Bộ nhớ băng thông cao (HBM) của mỗi GPU duy trì nhiều kênh để truyền tải song song các thông báo tập thể được chia thành khối. Các luồng GPU của người gửi trước tiên sao chép dữ liệu từ bộ đệm tính toán sang bộ đệm kênh khả dụng. Luồng proxy CPU của người gửi chỉ có thể đăng yêu cầu ghi RDMA sau khi nhận được gói clear-to-send (CTS) từ người nhận, bao gồm thông tin về kích thước và bộ nhớ. Sau đó, các luồng GPU của người nhận sao chép nội dung bộ đệm kênh vào bộ đệm tính toán đích. Cuối cùng, các luồng proxy CPU ở cả hai bên tái chế bộ đệm kênh và proxy CPU của người nhận gửi một gói CTS khác sau khi bộ đệm kênh sẵn sàng.

Chúng tôi tận dụng hiệu quả cơ chế này như một phương thức nhập lưu lượng do người nhận điều khiển để hạn chế lượng lưu lượng đang bay trên mạng, đặc biệt là khi bắt đầu xuất hiện tắc nghẽn. Tuy nhiên, việc định cấu hình cài đặt phù hợp có thể là một thách thức vì:

Số lượng kênh bị hạn chế do tranh chấp tài nguyên trên các luồng GPU với các hoạt động tính toán đồng thời;
Việc đặt kích thước bộ đệm kênh yêu cầu sự cân bằng cẩn thận hơn giữa việc lan truyền tắc nghẽn và sử dụng băng thông kém hơn so với Infiniband do kiểm soát luồng chi tiết hơn của RoCE và khả năng chậm của máy chủ cuối.

Do đó, chúng tôi đã thực hiện hai bước để cải thiện hiệu suất. Đầu tiên, chúng tôi đã xác định bằng thực nghiệm các cài đặt tham số phù hợp cho số lượng kênh và kích thước bộ đệm kênh trên các loại tập thể và kích thước công việc đào tạo khác nhau. Thứ hai, chúng tôi đã triển khai xếp hàng ưu tiên cao tại các bộ chuyển mạch cho các gói CTS để xúc tiến các thông báo và giảm thiểu khả năng thiếu hụt băng thông.

Kiểm soát tắc nghẽn là trọng tâm nghiên cứu trong mạng RDMA. DCQCN là tiêu chuẩn vàng cho các mạng tập trung vào lưu trữ. Tuy nhiên, kinh nghiệm của chúng tôi với khối lượng công việc huấn luyện AI phân tán cung cấp một quan điểm khác về việc điều chỉnh các thuật toán kiểm soát tắc nghẽn. Mặc dù đã tắt DCQCN và nhiều trường hợp RTSW gửi PFC đến CTSW bộ đệm sâu, nhưng chúng tôi vẫn chưa gặp phải trường hợp nào trong bốn năm qua, khi lưu lượng truy cập huấn luyện AI sản xuất khiến CTSW gửi PFC đến RTSW một cách liên tục.

Giải pháp hiện tại của chúng tôi phụ thuộc vào sự phối hợp cẩn thận giữa thư viện giao tiếp tập thể và mạng. Nó có thể phụ thuộc vào thông lượng tương đối giữa GPU và mạng, điều này có thể không áp dụng cho mọi trường hợp. Chúng tôi khuyến khích cộng đồng nghiên cứu tập trung hơn vào chủ đề này.

Hướng phát triển

Việc thiết kế và vận hành mạng RoCE quy mô lớn cho khối lượng công việc huấn luyện AI phân tán đã phát triển để đáp ứng nhu cầu ngày càng tăng về mật độ và quy mô tính toán. Bằng cách tách biệt các mạng FE và BE, sử dụng các lược đồ định tuyến khác nhau và tối ưu hóa các mẫu lưu lượng tập thể, chúng tôi đã có thể xây dựng một cơ sở hạ tầng mạng hiệu suất cao và đáng tin cậy. Những thiết kế và hiểu biết sâu sắc này nhấn mạnh tầm quan trọng của việc hiểu sâu sắc khối lượng công việc đào tạo và chuyển những hàm ý này thành thiết kế thành phần mạng, cuối cùng góp phần vào sự tiến bộ của cơ sở hạ tầng đào tạo AI phân tán.

Với xu hướng phát triển nhanh chóng của khối lượng công việc GenAI, cơ sở hạ tầng mạng của chúng tôi sẽ phát triển nhanh chóng.

Kết luận

Bài báo này đã giới thiệu về mạng RoCE được sử dụng để huấn luyện AI phân tán quy mô lớn tại Meta. Chúng tôi đã thảo luận về những thách thức trong việc xây dựng một mạng như vậy và các giải pháp được triển khai để giải quyết các vấn đề về cấu trúc liên kết, định tuyến và kiểm soát tắc nghẽn.

Với sự phát triển không ngừng của AI, chúng tôi tin rằng việc chia sẻ kinh nghiệm và giải pháp của mình sẽ đóng góp tích cực vào cộng đồng nghiên cứu và thúc đẩy sự đổi mới trong lĩnh vực mạng lưới AI.

Mạng RoCE cho huấn luyện AI phân tán quy mô lớn tại Meta