Truyền bộ nhớ GPU hiệu năng cao trên AWS Sagemaker Hyperpod

Cơ sở hạ tầng học sâu hiện đại thường yêu cầu truyền lượng lớn dữ liệu giữa các GPU qua các máy chủ. Tại Perplexity, chúng tôi gặp phải một thách thức kỹ thuật độc đáo: truyền hiệu quả các vùng bộ nhớ GPU không liền kề giữa các máy chủ với tốc độ tối đa có thể. Nền tảng mục tiêu của chúng tôi, các phiên bản AWS p5, cung cấp băng thông mạng ấn tượng 3200 Gbps thông qua 32 card mạng. Bài viết này chia sẻ hành trình xây dựng giải pháp mạng hiệu năng cao tùy chỉnh đạt 97.1% băng thông lý thuyết này.

Thách thức kỹ thuật

Trường hợp sử dụng của chúng tôi đặt ra một số yêu cầu kỹ thuật quan trọng:

• Truyền băng thông cao giữa các GPU từ xa của các khối bộ nhớ không liền kề
• Khả năng thêm hoặc xóa động các nút khỏi các triển khai Kubernetes mà không làm gián đoạn các hoạt động đang diễn ra
• Hỗ trợ các mô hình giao tiếp ngang hàng (peer-to-peer)

Mặc dù thư viện NCCL của NVIDIA là tiêu chuẩn thực tế cho học sâu phân tán, nhưng nó không lý tưởng cho trường hợp sử dụng của chúng tôi:

• NCCL xuất sắc trong giao tiếp tập thể nhưng yêu cầu thiết lập một "thế giới" tĩnh, đòi hỏi phải khởi động lại toàn bộ cụm khi điều chỉnh các nút tham gia.
• Mô hình giao tiếp đồng bộ của NCCL làm tăng độ phức tạp cho khối lượng công việc bất đồng bộ của chúng tôi
• Chúng tôi muốn kiểm soát trực tiếp các mẫu truyền bộ nhớ của mình để tối ưu hóa
• Xây dựng giải pháp của riêng mình mang lại cơ hội học hỏi quý giá

Mạng hiệu năng cao hiện đại

Để hiểu giải pháp của chúng tôi, trước tiên hãy khám phá cách các mạng hiệu năng cao hiện đại khác với mạng truyền thống.

Hầu hết các mạng chúng ta sử dụng hàng ngày đều dựa trên các giao thức TCP/IP, nơi các ứng dụng giao tiếp với card mạng thông qua nhân hệ điều hành sử dụng socket. Tuy nhiên, các mạng hiệu năng cao sử dụng RDMA (Truy cập bộ nhớ trực tiếp từ xa) - một ngăn xếp phần cứng và phần mềm hoàn toàn khác cho phép truy cập bộ nhớ trực tiếp giữa các máy mà không cần sự can thiệp của CPU.

AWS cung cấp Elastic Fabric Adapter (EFA), một giao diện mạng tùy chỉnh thực hiện giao thức tùy chỉnh của Amazon gọi là Scalable Reliable Datagram (SRD). Khác với mạng TCP/IP truyền thống nơi dữ liệu phải được sao chép nhiều lần giữa không gian người dùng, không gian nhân và bộ đệm mạng, EFA với RDMA cho phép truyền dữ liệu trực tiếp giữa bộ nhớ GPU và card mạng, bỏ qua hoàn toàn CPU.

Triết lý thiết kế mạng hiệu năng cao

Xây dựng hệ thống mạng hiệu năng cao đòi hỏi phải suy nghĩ lại một số giả định cơ bản:

1. Quyền sở hữu bộ đệm: Khác với các socket truyền thống nơi nhân quản lý bộ đệm mạng và yêu cầu sao chép giữa không gian người dùng và không gian nhân, RDMA yêu cầu các ứng dụng tự quản lý bộ đệm của chúng. Khi một ứng dụng khởi tạo một thao tác mạng, nó chuyển quyền sở hữu bộ đệm cho card mạng cho đến khi thao tác hoàn tất, loại bỏ nhu cầu sao chép dữ liệu.
2. Đăng ký bộ nhớ: Các ứng dụng phải đăng ký các vùng bộ nhớ với nhân hệ điều hành. Nhân thiết lập các ánh xạ địa chỉ ảo cho phép CPU, GPU và card mạng đều hiểu cùng một địa chỉ ảo. Việc đăng ký này là một thao tác một lần cho phép các lần truyền dữ liệu không sao chép tiếp theo.
3. Mặt phẳng điều khiển so với Mặt phẳng dữ liệu: Mạng hiệu năng cao tách các thao tác thành hai danh mục:
   1. Các thao tác mặt phẳng điều khiển (như thiết lập kết nối và đăng ký bộ nhớ) đi qua nhân để đảm bảo bảo mật
   2. Các thao tác mặt phẳng dữ liệu (truyền dữ liệu thực tế) bỏ qua nhân để đạt hiệu năng tối đa
4. Nhận trước khi truyền: Không có bộ nhớ đệm do nhân quản lý, các ứng dụng phải đăng ký trước các thao tác nhận, chỉ định nơi dữ liệu đến nên được đặt. Đây là một sự thay đổi cơ bản so với mô hình socket nơi các ứng dụng có thể nhận dữ liệu bất cứ lúc nào.
5. Hoàn thành dựa trên thăm dò: Thay vì chờ thông báo từ nhân thông qua các cơ chế như epoll, các ứng dụng trực tiếp thăm dò các hàng đợi hoàn thành của phần cứng. Điều này loại bỏ chi phí chung của lệnh gọi hệ thống và cho phép phản ứng tức thì với các thao tác đã hoàn thành.
6. Nhận thức cấu trúc phần cứng: Hiểu và tối ưu hóa cấu trúc phần cứng là rất quan trọng để đạt được hiệu năng tối đa.

Hiểu về cấu trúc phần cứng

Các phiên bản AWS p5 có kiến trúc nội bộ tinh vi. Như được hiển thị bên dưới, mỗi phiên bản chứa hai socket CPU tạo thành hai nút NUMA, với mỗi nút NUMA kết nối với bốn switch PCIe:

Phóng to

Dưới mỗi switch PCIe, chúng ta tìm thấy:

• Bốn card mạng EFA 100 Gbps
• Một GPU NVIDIA H100
• Một ổ SSD NVMe

Các đường dữ liệu cho các chuyển đổi TCP/IP và RDMA chứng minh sự khác biệt cơ bản trong cách tiếp cận của chúng:

Phóng to

Với TCP/IP (phía bên trái), dữ liệu phải được sao chép nhiều lần:

1. Từ bộ nhớ GPU đến bộ đệm ứng dụng trong không gian người dùng
2. Từ bộ đệm không gian người dùng đến bộ đệm nhân
3. Từ bộ đệm nhân đến card mạng
4. Quá trình ngược lại diễn ra ở phía nhận

Mỗi thao tác sao chép tiêu thụ các chu kỳ CPU và băng thông bộ nhớ. Ứng dụng cũng phải chuyển đổi ngữ cảnh giữa không gian người dùng và không gian nhân cho mỗi thao tác mạng.

Ngược lại, RDMA (phía bên phải) cung cấp truyền dữ liệu không sao chép thực sự:

1. Card mạng đọc trực tiếp từ bộ nhớ GPU
2. Dữ liệu truyền trực tiếp đến card mạng từ xa
3. Card mạng từ xa ghi trực tiếp vào bộ nhớ GPU đích
4. Ứng dụng chỉ cần kiểm tra hàng đợi hoàn thành (CQ) trong không gian người dùng để biết khi nào quá trình truyền hoàn tất

Với RDMA và ghép nối phần cứng phù hợp, việc truyền dữ liệu giữa hai GPU chỉ yêu cầu đi qua switch PCIe cục bộ và mạng:

Phóng to

Ngược lại, các tệp truyền TCP/IP phải sao chép dữ liệu nhiều lần qua bộ nhớ chính, gây ra tắc nghẽn bus PCIe đáng kể:

Phóng to

Xây dựng với libfabric

Chúng tôi đã sử dụng libfabric - một framework cung cấp giao diện chung cho các dịch vụ fabric. Việc triển khai của chúng tôi sử dụng hai loại thao tác RDMA:

• RDMA hai chiều (SEND/RECV) cho các thông điệp điều khiển mang siêu dữ liệu về các vùng bộ nhớ (ví dụ: độ lệch và kích thước).
• RDMA WRITE một chiều cho việc truyền dữ liệu thực tế, trong đó mỗi thao tác WRITE xử lý một khối bộ nhớ liên tục.

Hành trình đến hiệu năng đỉnh cao

Quá trình phát triển của chúng tôi trải qua một số giai đoạn:

1. Đã triển khai truyền thông điệp một chiều cơ bản sử dụng SEND/RECV
2. Đã mở rộng sang giao tiếp hai chiều
3. Đã thêm GPUDirect RDMA WRITE để truyền trực tiếp GPU-GPU
4. Đã mở rộng để xử lý nhiều luồng truyền đồng thời
5. Đã giới thiệu xếp hàng thao tác để tăng độ vững chắc
6. Đã đạt được mức sử dụng băng thông 97.4% trên một card mạng duy nhất

Khi mở rộng quy mô lên 32 card mạng, chúng tôi đã thực hiện một số tối ưu hóa quan trọng:

1. Xếp hàng thao tác: Thay vì gửi trực tiếp các thao tác đến card mạng, chúng tôi duy trì một hàng đợi cấp ứng dụng. Điều này mang lại sự vững chắc chống lại tắc nghẽn mạng và đơn giản hóa mô hình lập trình.
2. Làm ấm mạng: Thiết lập trước các kết nối giúp cải thiện hiệu năng khởi động.
3. Đa luồng: Các chuỗi chuyên dụng cho các thao tác mạng của từng GPU.
4. Gắn kết lõi CPU: Ràng buộc các chuỗi với các lõi CPU cụ thể để tránh các hiệu ứng NUMA và lỗi bộ nhớ đệm.
5. Phân mảnh trạng thái: Giảm sự tranh chấp giữa các chuỗi.
6. Gộp nhóm thao tác: Gửi nhiều thao tác cùng nhau.
7. Đăng thao tác trễ: Các thao tác được xếp hàng trong ứng dụng trước. Sau khi thăm dò hàng đợi hoàn thành, chúng tôi cố gắng gửi các thao tác đang chờ xử lý đến card mạng, đảm bảo sử dụng hiệu quả tài nguyên mạng.
8. Phân bổ tài nguyên nhận thức NUMA: Phân bổ tài nguyên libfabric như hàng đợi hoàn thành trên nút NUMA chính xác để giảm thiểu độ trễ truy cập bộ nhớ.

Thông qua các tối ưu hóa này, chúng tôi đã đạt được hiệu năng cuối cùng là 3.108 Gbps - 97.1% băng thông tối đa lý thuyết.

Video bên dưới cho thấy chương trình dòng lệnh của chúng tôi hoạt động thực tế. Nó truyền các khối bộ nhớ GPU không liền kề đến một nút từ xa, đạt tốc độ truyền 3108.283 Gbps - chứng minh mức sử dụng băng thông gần như lý thuyết của cơ sở hạ tầng mạng:

Kết luận

Xây dựng một hệ thống mạng hiệu năng cao đòi hỏi phải hiểu cả kiến trúc phần cứng và nguyên tắc thiết kế hệ thống. Mặc dù các thư viện như NCCL cung cấp các giải pháp tuyệt vời cho các mẫu phổ biến, nhưng đôi khi các giải pháp tùy chỉnh là cần thiết cho các yêu cầu cụ thể.

Hành trình của chúng tôi chứng minh rằng việc đạt được hiệu năng mạng gần như lý thuyết là có thể với sự chú ý cẩn thận đến kiến trúc hệ thống, cấu trúc phần cứng và các kỹ thuật tối ưu hóa khác nhau. Chìa khóa không chỉ là hiểu các thành phần riêng lẻ, mà là cách chúng tương tác để tạo thành một hệ thống hoàn chỉnh.

Phân tích kỹ thuật chi tiết đầy đủ của hành trình này, bao gồm chi tiết triển khai và các ví dụ mã, có sẵn trong kho lưu trữ mã nguồn mở của chúng tôi.

Trích dẫn

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}