Giao tiếp Điểm-Điểm RDMA cho Hệ thống LLM

Các hệ thống LLM từ trước đến nay vẫn dựa vào giao tiếp tập thể cho các mô hình song song có cấu trúc trong quá trình huấn luyện và suy luận. Tuy nhiên, các khối lượng công việc mới nổi, chẳng hạn như suy luận phân tách, định tuyến Chuyên gia hỗn hợp (MoE), và tinh chỉnh học tăng cường bất đồng bộ, đòi hỏi sự linh hoạt của giao tiếp điểm-điểm.

Bài báo nghiên cứu mới nhất của chúng tôi giới thiệu TransferEngine, một thư viện giao tiếp dựa trên RDMA có khả năng di động, hoạt động trên các Bộ điều khiển Giao diện Mạng (NIC) khác nhau, bao gồm NVIDIA ConnectX-7 và AWS Elastic Fabric Adapter (EFA). Chúng tôi chứng minh hiệu quả của thư viện này trong ba hệ thống sản xuất tại Perplexity AI:

Tiền nạp và Giải mã phân tách: Truyền KvCache theo từng lớp bằng cách sử dụng các lệnh ghi được phân trang và trình theo dõi UVM, cho phép các phiên bản tiền nạp và giải mã mở rộng quy mô độc lập.

Truyền trọng số cho quá trình hậu huấn luyện RL trong vòng chưa đầy 2 giây: Đồng bộ hóa trọng số cực nhanh giữa các nút huấn luyện và suy luận.

Kích hoạt các mô hình nghìn tỷ tham số trên AWS EFA: Các hoạt động phân phối và kết hợp MoE độ trễ thấp sử dụng các nguyên hàm phân tán và rào cản, đạt được độ trễ tiên tiến trên ConnectX-7 và là triển khai khả thi đầu tiên tương thích với AWS EFA.

Bài viết này tập trung vào bản thân TransferEngine, thư viện giao tiếp hỗ trợ các hệ thống này. Hãy đọc toàn bộ bài báo nghiên cứu trên arXiv và khám phá mã nguồn trên GitHub.

Giao tiếp tập thể xuất sắc trong các mô hình song song có cấu trúc như Song song Tensor và Song song Dữ liệu, nhưng chúng áp đặt các ràng buộc làm hạn chế khả năng ứng dụng đối với các khối lượng công việc mới nổi:

Thành viên cố định: Giao tiếp tập thể yêu cầu thiết lập một thế giới giao tiếp toàn cục tại thời điểm khởi tạo, với quá trình thiết lập được đồng bộ hóa trên tất cả những người tham gia. Đối với suy luận phân tách, điều này có vấn đề vì số lượng phiên bản tiền nạp và giải mã phải mở rộng quy mô động để đáp ứng các mẫu lưu lượng truy cập.

Hình dạng tensor đồng nhất: Giao tiếp tập thể giả định tất cả những người tham gia trao đổi tensor có hình dạng và kiểu dữ liệu giống hệt nhau. Trong định tuyến MoE, số lượng token được gửi đến mỗi chuyên gia phụ thuộc vào dữ liệu và được phát hiện tại thời điểm chạy, buộc các giao tiếp all-to-all dày đặc phải đệm đến kích thước bộ đệm trường hợp xấu nhất.

Thứ tự hoạt động: Các API giao tiếp tập thể thực thi ngữ nghĩa đồng bộ hóa, ngay cả khi các ứng dụng chỉ cần biết rằng một tập hợp các hoạt động đã hoàn tất. Ví dụ, trong việc truyền KvCache, tất cả các trang từ phiên bản tiền nạp đến phiên bản giải mã có thể đến theo bất kỳ thứ tự nào. Tương tự, trong các bản cập nhật trọng số RL, tất cả các tensor trọng số có thể được truyền độc lập mà không có ràng buộc về thứ tự.

Giao tiếp điểm-điểm dựa trên RDMA giải quyết tự nhiên các giới hạn này, nhưng ba rào cản đã ngăn cản việc áp dụng nó:

Giới hạn thư viện: Mặc dù torch.distributed và NCCL cung cấp các API điểm-điểm như Send/Recv, các hoạt động này vẫn kế thừa ngữ nghĩa và giới hạn giao tiếp tập thể.

Sự đa dạng về phần cứng: Các NIC NVIDIA ConnectX với truyền tải Kết nối Đáng tin cậy (RC) là thiết lập RDMA phổ biến nhất, đảm bảo phân phối theo đúng thứ tự. Các nhà cung cấp đám mây như AWS sử dụng phần cứng khác, Elastic Fabric Adapter (EFA) với truyền tải Datagram Đáng tin cậy Có thể mở rộng (SRD) phân phối các thông điệp không theo thứ tự. EFA cũng yêu cầu kết hợp bốn NIC 100 Gbps (hoặc hai NIC 200 Gbps) để đạt băng thông 400 Gbps.

Tính di động của phần mềm: Các thư viện RDMA hiện có chịu sự ràng buộc của nhà vendor. DeepEP yêu cầu RDMA do GPU khởi tạo (IBGDA) chỉ có sẵn trên ConnectX. NVSHMEM cho thấy sự suy giảm hiệu suất nghiêm trọng trên EFA. Các thư viện như Mooncake và NIXL thiếu hỗ trợ EFA hoặc vẫn đang trong giai đoạn phát triển đầu. Kết quả là, không có giải pháp liên nhà cung cấp khả thi nào tồn tại cho việc suy luận LLM.

TransferEngine giải quyết các rào cản này thông qua bốn năng lực cốt lõi:

Các giả định truyền tải chung: Thay vì nhắm mục tiêu vào một NIC cụ thể, TransferEngine chỉ giả định rằng quá trình truyền tải bên dưới là đáng tin cậy nhưng không theo thứ tự, tức là phần giao của các đảm bảo được cung cấp bởi cả hai bộ truyền tải RC và SRD.

Đồng bộ hóa ImmCounter: Một nguyên hàm mới cho thông báo hoàn thành hoạt động mà không phụ thuộc vào thứ tự thông điệp, cho phép cùng một đoạn mã chạy trên cả bộ truyền tải có thứ tự (RC) và không theo thứ tự (SRD).

Hỗ trợ đa NIC: Quản lý minh bạch nhiều NIC cho mỗi GPU, tự động phân mảnh các quá trình truyền để tối đa hóa băng thông.

Bề mặt API tối thiểu: Một tập hợp nhỏ các hoạt động đủ để xây dựng các mô hình giao tiếp phức tạp.

Thiết kế hệ thống TransferEngine, cho thấy các đường đi của yêu cầu và callback, các worker CPU, bộ đệm UVM, GPU và NIC trên hai miền GPU/NUMA.

TransferEngine tạo ra một chuỗi tác nhân cho mỗi GPU. Mỗi tác nhân quản lý một DomainGroup điều phối từ 1 đến 4 RDMA NIC, tùy thuộc vào nền tảng phần cứng (một ConnectX-7 duy nhất cung cấp 400 Gbps, trong khi EFA yêu cầu bốn bộ chuyển đổi 100 Gbps hoặc hai bộ chuyển đổi 200 Gbps). Trong một DomainGroup, mỗi Domain là cụ thể theo phần cứng và quản lý một NIC duy nhất, xử lý quản lý cặp hàng đợi, gửi công việc và thăm dò hoàn thành.

Mỗi thể hiện TransferEngine hiển thị một địa chỉ chính duy nhất để nhận dạng và phát hiện thiết bị ngang hàng. Các thiết bị ngang hàng từ xa trao đổi cấu trúc NetAddr mã hóa các địa chỉ mạng và có thể được tuần tự hóa để giao tiếp ngoài băng tần. Tất cả các thiết bị ngang hàng trong một hệ thống phải sử dụng cùng một số lượng NIC cho mỗi GPU. Cấu hình đồng nhất này cho phép TransferEngine phân mảnh các quá trình truyền một cách thông minh: khi bắt đầu truyền, cả nguồn và đích đều biết chính xác những NIC nào khả dụng, cho phép cân bằng tải tự động trên tất cả các bộ chuyển đổi.

TransferEngine hiển thị một API tối thiểu giúp trừu tượng hóa độ phức tạp của các giao diện RDMA bên dưới.

rust

#[serde] struct NetAddr(Bytes);
#[serde] struct MrDesc{ ptr: u64, rkeys: Vec<(NetAddr, u64)> }
struct MrHandle(NonNull<c_void>); type Offset = u64;
struct Pages{ indices: Vec<u32>, stride: u64, offset: Offset }
struct PeerGroupHandle(u64);
struct ScatterDst{ len: u64, src: Offset, dst: (MrDesc,Offset)}
enum OnDone { Callback(fn () -> ()), Flag(Atomic<bool>) }

trait TransferEngine {
  fn main_address() -> NetAddr;
  // Memory Region Management
  fn reg_mr(ptr, len, device) -> (MrHandle, MrDesc);
  // Two-sided Send/Recv
  fn submit_send(addr: NetAddr, msg: &[u8], cb: fn () -> ());
  fn submit_recvs(len: u64, cnt: u64, cb: fn (&[u8]) -> ());
  // One-sided Write
  fn expect_imm_count(imm: u32, count: u32, cb: fn () -> ());
  fn submit_single_write(len: u64, imm: Option<u32>
      src: (MrHandle, Offset), dst: (MrDesc, Offset), OnDone);
  fn submit_paged_writes(page_len: u64, imm: Option<u32>,
      src: (MrHandle, Pages), dst: (MrDesc, Pages), OnDone);
  // One-sided Write to a group of peers
  fn add_peer_group(addrs: Vec<NetAddr>) -> PeerGroupHandle;
  fn submit_scatter(h: Option<PeerGroupHandle>, OnDone,
      imm: Option<u32>, src: MrHandle, dst: Vec<ScatterDst>);
  fn submit_barrier(h: Option<PeerGroupHandle>, OnDone,
      imm: u32, dst: Vec<MrDesc>);
  // Watcher for CPU-GPU synchronization
  fn alloc_uvm_watcher(cb: fn(u64,u64) -> ()) -> NonNull<u64>;
}

API bao gồm một số thành phần chính:

Các hoạt động RDMA yêu cầu các vùng bộ nhớ phải được đăng ký với NIC trước khi truyền dữ liệu. Hàm reg_mr nhận một con trỏ, chiều dài và thiết bị, trả về hai tay cầm:

MrHandle: Một tay cầm cục bộ được sử dụng làm nguồn cho các quá trình truyền

MrDesc: Một mô tả có thể tuần tự hóa có thể được trao đổi với các thiết bị ngang hàng từ xa để cho phép chúng ghi vào vùng bộ nhớ đã đăng ký

Cấu trúc MrDesc chứa một danh sách địa chỉ mạng từ xa và khóa từ xa (mỗi NIC một khóa) cấp quyền truy cập từ xa vào vùng bộ nhớ. Thiết kế này xử lý tự nhiên các thiết lập đa NIC.

Đối với các thông điệp nhỏ và thông điệp mặt phẳng điều khiển theo kiểu RPC, TransferEngine cung cấp submit_send và submit_recvs:

submit_send: Gửi thông điệp đến một thiết bị ngang hàng từ xa được nhận dạng bởi NetAddr. Thông điệp được sao chép nội bộ, vì vậy người gọi có thể ngay lập tức tái sử dụng hoặc giải phóng bộ đệm.

submit_recvs: Đăng một nhóm bộ đệm nhận. Khi một thông điệp đến, một hàm gọi lại được gọi với dữ liệu đã nhận. Bộ đệm được tự động trả về nhóm sau khi hàm gọi lại hoàn tất.

Đối với việc truyền dữ liệu lớn, TransferEngine cung cấp các hoạt động ghi một chiều không sao chép:

submit_single_write: Truyền một vùng liên tục từ nguồn đến đích

submit_paged_writes: Truyền các trang không liên tục được chỉ định bởi các chỉ mục, bước nhảy và độ lệch

Cả hai hoạt động đều chấp nhận một giá trị tức thời 32-bit tùy chọn được phân phối đến bên nhận một cách nguyên tử khi hoàn thành.

Công cụ này tự động phân mảnh các quá trình truyền lớn trên nhiều NIC để tối đa hóa băng thông.

Đối với các hoạt động nhắm mục tiêu vào nhiều thiết bị ngang hàng (phổ biến trong định tuyến MoE), TransferEngine cung cấp các nguyên hàm được tối ưu hóa:

add_peer_group: Đăng ký trước một tập hợp các địa chỉ ngang hàng để tra cứu nhanh

submit_scatter: Gửi các lát dữ liệu khác nhau đến từng thiết bị ngang hàng trong nhóm

submit_barrier: Gửi một thông báo chỉ có giá trị tức thời đến tất cả các thiết bị ngang hàng (truyền không byte)

Các hoạt động này được tối ưu hóa cho độ trễ thấp bằng cách lưu bộ nhớ đệm các cấu trúc dữ liệu nội bộ và tạo mẫu cho các yêu cầu công việc RDMA, tránh chi phí thiết lập lặp đi lặp lại.

Một thách thức chính trong các hệ thống LLM là điều phối các quá trình truyền RDMA với tiến trình của GPU. Lý tưởng nhất, việc truyền dữ liệu nên bắt đầu ngay khi GPU tạo ra đầu ra, nhưng việc thăm dò GPU từ CPU là không hiệu quả.

TransferEngine giải quyết vấn đề này bằng Trình theo dõi UVM: các vị trí bộ nhớ ảo thống nhất có thể được cập nhật bởi các kernel GPU (ngay cả trong đồ thị CUDA) và được thăm dò hiệu quả bởi CPU sử dụng GDRCopy. Khi một trình theo dõi UVM được cấp phát, một hàm gọi lại được đăng ký sẽ được kích hoạt bất cứ khi nào GPU cập nhật giá trị.

Điều này cho phép các mô hình như truyền KvCache theo từng lớp: sau khi mỗi lớp transformer hoàn tất, GPU tăng giá trị của trình theo dõi, kích hoạt quá trình truyền RDMA của KvCache lớp đó trong khi các lớp tiếp theo tiếp tục tính toán.

TransferEngine cung cấp các cơ chế thông báo hoàn thành cho cả phía khởi tạo (người gọi submit_write) và phía nhận (mục tiêu của các hoạt động Ghi).

Thông báo ở phía khởi tạo: Khi một hoạt động ghi được gửi đi, người gọi có thể chọn nhận thông báo khi hoàn thành thông qua một hàm gọi lại hoặc một cờ nguyên tử. Thông báo này xác nhận rằng dữ liệu đã được truyền thành công từ bộ đệm nguồn.

Thông báo ở phía nhận: Đổi mới cốt lõi là ImmCounter, một thành phần theo dõi các bộ đếm được lập chỉ mục bởi các số tức thời 32-bit. Khi một hoạt động ghi mang theo một giá trị tức thời, ImmCounter của bên nhận được tăng lên một cách nguyên tử khi giao dữ liệu. Hàm expect_imm_count cho phép bên nhận đăng ký một hàm gọi lại được kích hoạt khi một giá trị tức thời cụ thể đạt đến số lượng mục tiêu. Các bộ đếm có thể được đồng bộ hóa một cách minh bạch với GPU thông qua GDRCopy, được thăm dò trực tiếp hoặc được giám sát thông qua các hàm gọi lại.

Cơ chế phía nhận này cung cấp sự đồng bộ hóa đáng tin cậy mà không yêu cầu thứ tự thông điệp. Nó hoạt động tốt như nhau trên các NIC ConnectX (đảm bảo phân phối theo đúng thứ tự qua truyền tải RC) và các NIC EFA (phân phối thông điệp không theo thứ tự qua truyền tải SRD).

TransferEngine được triển khai bằng ngôn ngữ Rust, tối ưu hóa cẩn thận việc cấp phát, tạo chuỗi và đồng bộ hóa để giảm thiểu độ trễ và đạt được thông lượng cao. Côngع này tạo ra một chuỗi tác nhân cho mỗi DomainGroup, mỗi chuỗi được ghim vào một lõi CPU trên nút NUMA mà các thiết bị được gắn vào, giúp giảm thiểu độ trễ lập lịch và truy cập bộ nhớ. Các cấu trúc dữ liệu dành riêng cho một miền được cấp phát sau khi ghim để đảm bảo bộ nhớ được dành riêng trong nút NUMA chính xác. Một chuỗi tác nhân xử lý tối đa 4 Domain (mỗi miền quản lý một NIC duy nhất), trong khi một chuỗi chuyên i khác thăm dò GPU để cập nhật các trình theo dõi UVM. Giao tiếp giữa các chuỗi được thực hiện thông qua các hàng đợi không khóa.

API chuyển tiếp các yêu cầu đến DomainGroup phục vụ thiết bị tương thích, với thiết bị đầu tiên cũng phục vụ máy chủ. Trong một vòng lặp chặt chẽ, tác nhân miền thăm dò công việc mới, ưu tiên gửi các yêu cầu mới. Tùy thuộc vào phần cứng và cấu hình, các yêu cầu được phân mảnh và cân bằng tải trên các Domain khả dụng. Lệnh Ghi đầu tiên của một yêu cầu hỗn hợp được gửi ngay lập tức đến hàng đợi gửi của NIC. Khi các yêu cầu mới đã cạn kiệt, tác nhân tiến hành xử lý các hoạt động đang chờ xử lý, gửi các lệnh ghi để làm đầy đường ống phần cứng. Sau đó, các hàng đợi hoàn thành được thăm dò để truy vấn các quá trình truyền đã hoàn tất và việc tăng bộ đếm tức thời. Các sự kiện được tổng hợp để cung cấp các thông báo cho mỗi lần truyền, bàn giao quá trình truyền cho một chuỗi hàm gọi lại chuyên dụng được chia sẻ bởi tất cả các nhóm.

Việc phân mảnh bên trong một DomainGroup rất linh hoạt. Các quá trình truyền có thể nhắm mục tiêu vào các NIC cụ thể theo chỉ mục. Một lệnh Ghi duy nhất có thể được chia nhỏ. Các quá trình truyền được phân trang, phân tán và các hoạt động rào cản, tất cả đều dịch thành nhiều lệnh Ghi, có thể phân mảnh trên tất cả các NIC.

Các Domain trong TransferEngine được chuyên óa và tối ưu hóa cho phần cứng dưới quyền kiểm soát của chúng:

AWS EFA: Việc triển khai EFA sử dụng libfabric, quản lý một miền fabric cho mỗi NIC trong DomainGroup. Các quyết định triển khai và tối ưu hóa quan trọng bao gồm:

Ghi kích thước bằng không: EFA khác với thông số kỹ thuật RDMA bằng cách yêu cầu các mô tả mục tiêu hợp lệ ngay cả đối với các lần ghi có kích thước bằng không. Do đó, TransferEngine thực thi các mô tả hợp lệ cho tất cả các quá trình truyền.

Tạo mẫu yêu cầu công việc: Điền trước và giữ lại các trường chung của cấu trúc libfabric trước khi gửi, làm giảm chi phí chung cho mỗi hoạt động đối với các quá trình truyền lớn và các hoạt động nhóm ngang hàng

NVIDIA ConnectX-7: Việc triển khai ConnectX-7 sử dụng libibverbs với truyền tải RC. Các quyết định triển khai và tối ưu hóa quan trọng bao gồm:

Cặp hàng đợi UD cho bắt tay: Mỗi Domain sử dụng một cặp hàng đợi Datagram Không đáng tin cậy (UD) để trao đổi quá trình bắt tay kết nối RC.

Cặp hàng đợi RC riêng biệt: Tạo 2 cặp hàng đợi RC cho mỗi thiết bị ngang hàng, một cho các hoạt động Gửi/Nhận hai chiều và một cho các hoạt động Ghi và WriteImm một chiều. Sự tách biệt này là cần thiết vì cả hoàn thành Recv và WriteImm đều tiêu thụ các yêu cầu công việc theo thứ tự đăng, cho phép ngữ nghĩa Recv cấp cao trong khi hỗ trợ WriteImm mà không bị nhiễu.

Tạo mẫu yêu cầu công việc: Điền trước các trường chung của cấu trúc yêu cầu công việc cho các quá trình truyền lớn và các nhóm ngang hàng.

Chuỗi yêu cầu công việc: Liên kết tối đa 4 yêu cầu công việc thông qua con trỏ next của ibv_send_wr, làm giảm số lần chuông cửa (doorbell ring) đến NIC.

Thứ tự linh hoạt: Cho phép IBV_ACCESS_RELAXED_ORDERING cho phép các giao dịch PCIe không theo thứ tự giữa NIC và bộ nhớ GPU.

Bằng cách nhắm mục tiêu vào phần giao của các đảm bảo trên các bộ truyền tải RDMA khác nhau và giới thiệu tính đồng bộ hóa độc lập với thứ tự thông qua ImmCounter, TransferEngine cung cấp một nền tảng có thể di động cho giao tiếp điểm-điểm trong các hệ thống LLM. Thư viện là mã nguồn mở và có sẵn trên GitHub, và chúng tôi khuyến khích cộng đồng khám phá và xây dựng dựa trên nó cho các mô hình giao tiếp mới nổi trong các hệ thống học máy quy mô lớn.

BibTeX

@online{perplexity2025RdmaPointToPointCommunicationForLlmSystems,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{RDMA point-to-point communication for LLM systems}},
  date = {2025-11-05},
  url = {https://www.perplexity.ai/hub/blog/rdma-point-to-point-communication-for-llm-systems}
}

Động lực

TransferEngine

Thiết kế API

Đăng ký Bộ nhớ

Gửi/Nhận Hai chiều

Ghi Một chiều

Nhóm Ngang hàng và Hoạt động Lớn

Trình theo dõi UVM: Báo hiệu GPU-đến-CPU

Thông báo Hoàn thành

Triển khai

Tối ưu hóa Riêng cho Phần cứng

Kết luận

Trích dẫn