Truyền trọng số cho hậu huấn luyện RL dưới 2 giây

Chúng tôi gần đây đã đạt được cập nhật tham số xuyên máy 1,3 giây cho Kimi-K2 (1T tham số), truyền trọng số từ 256 GPU huấn luyện (BF16) sang 128 GPU suy luận (FP8).

Trong quá trình tinh chỉnh học tăng cường bất đồng bộ, việc huấn luyện và suy luận chạy trên các GPU riêng biệt. Sau mỗi bước huấn luyện, các trọng số mới phải được đẩy đến các nút suy luận. Nhiều framework hiện tại mất vài giây—hoặc thậm chí vài phút—đối với các mô hình hàng nghìn tỷ tham số.

Bằng cách tận dụng giao tiếp điểm-điểm RDMA, chúng tôi có thể giúp việc truyền trọng số diễn ra cực kỳ nhanh chóng mà không cần thay đổi công cụ suy luận, đồng thời giúp mã nguồn dễ viết và dễ duy trì hơn.

Giải pháp của chúng tôi được xây dựng trên RDMA WRITE, một nguyên hàm một chiều nơi nguồn trực tiếp ghi vào bộ nhớ GPU của đích.

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

Phía đích thậm chí sẽ không nhận được thông báo về việc truyền dữ liệu. Điều này mang lại cho chúng tôi việc truyền dữ liệu có độ trễ thấp, thông lượng cao, không cần sao chép (zero-copy) được điều khiển bởi các nút huấn luyện mà không cần bất kỳ logic điều khiển nào trên các nút suy luận.

Thu thập siêu dữ liệu – Bộ điều khiển thu thập siêu dữ liệu tham số từ tất cả các GPU huấn luyện và suy luận.

Tính toán lịch trình – Bộ điều khiển tính toán lịch trình truyền trọng số tĩnh, ánh xạ GPU huấn luyện nào gửi tham số nào đến GPU suy luận nào, và theo thứ tự nào.

Phân phối lịch trình – Bộ điều khiển gửi lịch trình đến tất cả các GPU huấn luyện.

Thực thi – Sau mỗi bước huấn luyện, bộ điều khiển báo hiệu cho các GPU huấn luyện bắt đầu truyền.

Với quy trình làm việc cấp cao được xác định, thách thức chính là cách thực thi truyền trọng số một cách hiệu quả ở quy mô nghìn tỷ tham số. Tại đây, chúng tôi mô tả chi tiết về đường dẫn thực thi.

Các tham số trong quá trình huấn luyện được phân phối theo các vị trí FSDP. Bằng cách sử dụng full_tensor(), tất cả các GPU trong một DeviceMesh có thể tái tạo lại toàn bộ tham số, do đó tất cả đều có thể đóng vai trò là nguồn truyền trọng số.

Nhiều DeviceMesh rời rạc tạo thành một mesh group. Do các DeviceMesh trong cùng một nhóm rời rạc với nhau, việc truyền dữ liệu của chúng không bị can thiệp và có thể chạy hoàn toàn song song. Giữa các mesh group, chúng tôi chèn một rào cản toàn cục (global barrier) để thực thi thứ tự.

Pipeline chuyển giao trọng số chồng lấp các thao tác sao chép từ host sang device, lắp ghép tensor, hợp nhất và lượng tử hóa, cùng các lần truyền RDMA qua MeshGroups.

Chúng tôi coi việc truyền từng tensor tham số là một tác vụ. Quá trình truyền trọng số sử dụng nhiều loại nguồn phần cứng, do đó chúng tôi chia một tác vụ truyền trọng số thành các giai đoạn đường ống khác nhau chồng chéo nhau về thời gian:

Host to device memcpy — Nếu FSDP giảm tải trọng số sang CPU

Chuẩn bị tham số — Tái tạo toàn bộ trọng số bằng full_tensor(), áp dụng gộp phép chiếu (projection fusion), lượng tử hóa nếu cần.

Truyền RDMA — Ghi không cần sao chép vào bộ nhớ GPU suy luận từ xa

Global barrier — Sau khi tất cả các lệnh gọi full_tensor() hoàn tất, đồng bộ hóa giữa các mesh group bằng GLOO qua Ethernet.

Trong quá trình thực thi, chúng tôi duy trì một hàng đợi FIFO các tác vụ cho từng giai đoạn đường ống. Bất cứ khi nào tác vụ đứng đầu hàng đợi hoàn thành giai đoạn, nó sẽ được chuyển đến cuối hàng đợi giai đoạn tiếp theo.

full_tensor() và các hoạt động GPU khác giới thiệu mức sử dụng bộ nhớ GPU bổ sung. Để tránh lỗi hết bộ nhớ (out of memory), chúng tôi chỉ bắt đầu thực thi một tác vụ nếu các tác vụ đang chạy hiện tại chiếm ít bộ nhớ GPU tạm thời hơn một ngưỡng có thể cấu hình.

Một số lựa chọn thiết kế giúp hệ thống của chúng tôi chạy nhanh hơn đáng kể và dễ bảo trì hơn so com các giải pháp mã nguồn mở thông thường.

Chuyển giao trọng số dựa trên Rank0 thu thập các trọng số trước khi phân phối chúng tới nút suy luận.

Giao tiếp điểm-điểm truyền các trọng số từ cả hai nút huấn luyện tới các giao diện mạng của nút suy luận.

Một mô hình phổ biến là dồn tất cả các tham số qua các GPU rank-0: tập hợp ở rank-0 huấn luyện, gửi đến rank-0 suy luận, sau đó phân tán lại. Điều này nhanh chóng trở thành một điểm nghẽn, bị giới hạn bởi băng thông PCIe và NIC của một GPU duy nhất (ví dụ: 400 Gbps ≈ 50 GB/s).

Ngược lại, thiết lập điểm-điểm của chúng tôi cho phép mọi GPU huấn luyện gửi trực tiếp đến mọi GPU suy luận, làm bão hòa toàn bộ cấu trúc mạng thay vì chỉ một liên kết duy nhất.

Chuyển giao dựa trên RPC sử dụng bộ đệm phía suy luận và các lời gọi của worker, trong khi chuyển giao một phía ghi trực tiếp các trọng số mà không có sự tham gia của worker suy luận.

Một số hệ thống dựa vào việc gọi phương thức update_weight() của công cụ suy luận cho mỗi tensor. Điều đó có nghĩa là các thay đổi xâm lấn đối với mã suy luận, cộng với chi phí chung từ RPC, tuần tự hóa và phối hợp mặt phẳng điều khiển.

Với nguyên hàm RDMA WRITE, chúng tôi cập nhật trọng số một cách âm thầm trên bộ nhớ GPU suy luận mà không cần thêm các bản sao. Không có thông điệp mặt phẳng điều khiển và không có logic điều khiển CPU nào liên quan. Không cần sửa đổi công cụ suy luận.

Quá trình truyền trọng số có thể tận dụng bốn loại tài nguyên phần cứng: (1) Di chuyển dữ liệu host-device (2) Tính toán GPU cho gộp phép chiếu và lượng tử hóa (3) Mạng RDMA cho mặt phẳng dữ liệu (4) Ethernet cho mặt phẳng điều khiển.

Thiết kế của chúng tôi chia các tác vụ truyền trọng số thành các giai đoạn đường ống, cho phép chồng chéo dễ dàng trên các tài nguyên phần cứng khác nhau.

Một số triển khai tính toán lại lịch trình truyền ở mỗi bước huấn luyện, liên tục thu thập siêu dữ liệu và phân phối hướng dẫn. Điều này làm tăng độ trễ mặt phẳng điều khiển không cần thiết.

Lịch trình của chúng tôi được tính toán một lần khi khởi tạo. Mỗi vòng lặp huấn luyện chỉ đơn giản là phát lại kế hoạch: bộ điều khiển phát tín hiệu "tiếp tục", và các GPU tuân theo các tuyến đường đã được gán trước của chúng. Quá trình thực thi có thể dự đoán được và nhẹ nhàng.

Thật dễ bị cám dỗ khi làm rối tung toàn bộ quá trình cập nhật trọng số trong một hàm nguyên khối: thu thập siêu dữ liệu, khớp tên, tập hợp nội bộ nút, gộp phép chiếu, lượng tử hóa, chia nhỏ thế giới giao tiếp, truyền mạng liên nút. Việc lập trình đúng là rất khó, và việc tối ưu hóa lại càng khó hơn.

Trong kỹ thuật của chúng tôi, chúng tôi tách các bước này thành các thành phần riêng lẻ. Mỗi thành phần có thể được kiểm thử đơn vị, suy luận và tối ưu hóa một cách cô lập.

Truyền trọng số nhanh chóng, đáng tin cậy là một khối xây dựng quan trọng cho việc tinh chỉnh RL quy mô lớn. Bằng cách kết hợp nguyên hàm RDMA WRITE, lịch trình truyền tĩnh và thực thi theo đường ống, chúng tôi đã giảm thời gian cập nhật hàng nghìn tỷ tham số xuống chỉ còn 1,3 giây trên Kimi-K2. Cách tiếp cận này dễ suy luận, dễ bảo trì và tránh được các điểm nghẽn của các thiết kế truyền thống.

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: truyền dữ liệu một chiều

Quy trình làm việc cấp cao

Thực thi truyền trọng số

DeviceMesh và Mesh Groups

Đường ống tác vụ

Kiểm soát sử dụng bộ nhớ GPU

Tại sao nó nhanh và đơn giản

Giao tiếp điểm-điểm

Truyền dữ liệu một chiều

Xử lý đường ống (Pipelining)

Lịch trình tĩnh

Phân tách rõ ràng

Kết luận

Trích dẫn