Weight Transfer for RL Post-Training in under 2 seconds (research.perplexity.ai)

A recent breakthrough in reinforcement learning (RL) has achieved weight transfer in just 1.3 seconds for the Kimi-K2 model, utilizing 256 training GPUs and 128 inference GPUs. This significant advancement leverages RDMA WRITE, a one-sided data transfer method that allows direct writing into GPU memory without requiring notification or control logic on the inference side. By employing a static transfer schedule and separating various stages of the weight transfer process, the system efficiently manages trillion-parameter model updates, a task previously hampered by lengthy transfer times. The implications for the AI/ML community are profound. Fast and reliable weight transfer is vital for scaling RL fine-tuning, allowing models to adapt quickly to new data. The approach minimizes latency and maximizes throughput by bypassing traditional communication bottlenecks, making it simpler to maintain and optimize. The design promotes clean separations of components for better testing and understanding, facilitating advancements in large-scale AI systems. This innovative solution sets a new standard for efficient parameter updates in complex models, potentially accelerating the development cycle of AI applications across various domains.