2 सेकंड से कम समय में RL पोस्ट-ट्रेनिंग के लिए वेट ट्रांसफ़र

अल्ट्रा-फ़ास्ट क्रॉस-GPU मॉडल सिंक

लेखकPerplexity Research

हमने हाल ही में Kimi-K2 (1T पैरामीटर्स) के लिए 1.3-सेकंड का क्रॉस-मशीन पैरामीटर अपडेट हासिल किया है, जिसमें 256 ट्रेनिंग GPUs (BF16) से 128 इंफ़रेंस GPUs (FP8) में वेट्स (weights) ट्रांसफ़र किए गए।

असिंक्रोनस रीइन्फोर्समेंट लर्निंग फ़ाइन-ट्यूनिंग में, ट्रेनिंग और इंफ़रेंस अलग-अलग GPUs पर चलते हैं। प्रत्येक ट्रेनिंग चरण के बाद, नए वेट्स को इंफ़रेंस नोड्स पर पुश करना आवश्यक होता है। कई मौजूदा फ़्रेमवर्क में ट्रिलियन-पैरामीटर मॉडल के लिए इसमें कई सेकंड—या मिनटों तक का समय लगता है।

RDMA पॉइंट-टू-पॉइंट कम्युनिकेशन का लाभ उठाकर, हम इंफ़रेंस इंजन को बदले बिना वेट ट्रांसफ़र को अत्यधिक तेज़ बनाने में सक्षम हैं, और कोड को लिखना और बनाए रखना आसान बना सकते हैं।

RDMA WRITE: वन-साइडेड ट्रांसफ़र

हमारा समाधान RDMA WRITE पर आधारित है, जो एक वन-साइडेड प्रिमिटिव है जहाँ सोर्स सीधे डेस्टिनेशन की GPU मेमोरी में लिखता है।

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. ...

ट्रांसफ़र के लिए डेस्टिनेशन साइड को सूचित भी नहीं किया जाता है। यह हमें ट्रेनिंग नोड्स द्वारा संचालित लो-लेटेंसी, हाई-थ्रूपुट, ज़ीरो-कॉपी ट्रांसफ़र प्रदान करता है, जिसमें इंफ़रेंस नोड्स पर किसी कंट्रोल लॉजिक की आवश्यकता नहीं होती।

उच्च-स्तरीय वर्कफ़्लो

मेटाडेटा कलेक्शन – कंट्रोलर सभी ट्रेनिंग और इंफ़रेंस GPUs से पैरामीटर मेटाडेटा इकट्ठा करता है।

शेड्यूल कंप्यूटेशन – कंट्रोलर एक स्टैटिक वेट ट्रांसफ़र शेड्यूल की गणना करता है, जो यह मैप करता है कि कौन सा ट्रेनिंग GPU किस पैरामीटर को किस इंफ़रेंस GPU को भेजेगा, और किस क्रम में।

शेड्यूल डिस्ट्रीब्यूशन – कंट्रोलर सभी ट्रेनिंग GPUs को शेड्यूल भेजता है।

एक्जीक्यूशन – प्रत्येक ट्रेनिंग चरण के बाद, कंट्रोलर ट्रेनिंग GPUs को ट्रांसफ़र शुरू करने का सिग्नल देता है।

वेट ट्रांसफ़र एक्जीक्यूशन

उच्च-स्तरीय वर्कफ़्लो निर्धारित होने के साथ, मुख्य चुनौती यह है कि ट्रिलियन-पैरामीटर स्केल पर वेट ट्रांसफ़र को कुशलतापूर्वक कैसे निष्पादित किया जाए। यहाँ हम एक्जीक्यूशन पाथ का विवरण दे रहे हैं।

DeviceMesh और मेश ग्रुप्स

ट्रेनिंग में पैरामीटर्स FSDP प्लेसमेंट के अनुसार वितरित किए जाते हैं। full_tensor() का उपयोग करके, एक DeviceMesh के सभी GPUs पूरे पैरामीटर का पुनर्निर्माण कर सकते हैं, इसलिए सभी वेट ट्रांसफ़र के लिए सोर्स के रूप में कार्य कर सकते हैं।

कई अलग-थलग DeviceMeshes एक मेश ग्रुप बनाते हैं। चूंकि एक ही ग्रुप में DeviceMeshes अलग-थलग होते हैं, इसलिए उनके ट्रांसफ़र में कोई बाधा नहीं आती और वे पूरी तरह से समानांतर चल सकते हैं। मेश ग्रुप्स के बीच, हम ऑर्डर लागू करने के लिए एक ग्लोबल बैरियर डालते हैं।

टास्क पाइपलाइन

वेट-ट्रांसफर पाइपलाइन होस्ट-टू-डिवाइस कॉपी, टेंसर असेम्बली, फ्यूज़न और क्वांटाइज़ेशन, तथा MeshGroups के बीच RDMA ट्रांसफर को ओवरलैप करती है.

हम प्रत्येक पैरामीटर टेंसर के ट्रांसफ़र को एक टास्क मानते हैं। वेट ट्रांसफ़र प्रक्रिया कई प्रकार के हार्डवेयर संसाधनों का उपयोग करती है, इसलिए हम एक वेट ट्रांसफ़र टास्क को अलग-अलग पाइपलाइन चरणों में विभाजित करते हैं जो समय में ओवरलैप होते हैं:

होस्ट टू डिवाइस मेमकॉपी (memcpy) — यदि FSDP वेट को CPU पर ऑफ़लोड करता है

पैरामीटर तैयारी — full_tensor() के साथ पूर्ण वेट का पुनर्निर्माण करें, प्रोजेक्शन फ़्यूज़न लागू करें, यदि आवश्यक हो तो क्वांटाइज़ करें।

RDMA ट्रांसफ़र — रिमोट इंफ़रेंस GPU मेमोरी में ज़ीरो-कॉपी राइट

ग्लोबल बैरियर — सभी full_tensor() कॉल पूरे होने के बाद, ईथरनेट के माध्यम से GLOO का उपयोग करके मेश ग्रुप्स के बीच सिंक्रोनाइज़ करें।

कार्यान्वयन में, हम प्रत्येक पाइपलाइन चरण के लिए टास्क की एक FIFO कतार बनाए रखते हैं। जब भी कतार का हेड टास्क चरण पूरा करता है, तो उसे अगली चरण की कतार के अंत में ले जाया जाता है।

GPU मेमोरी उपयोग नियंत्रण

full_tensor() और अन्य GPU ऑपरेशन अतिरिक्त GPU मेमोरी उपयोग उत्पन्न करते हैं। आउट-ऑफ़-मेमोरी त्रुटि से बचने के लिए, हम किसी टास्क का निष्पादन केवल तभी शुरू करते हैं जब वर्तमान ऑन-द-फ़्लाई टास्क एक विन्यास योग्य वॉटरमार्क से कम अस्थायी GPU मेमोरी का उपयोग करते हैं।

यह तेज़ और सरल क्यों है

कई डिज़ाइन विकल्प हमारे सिस्टम को सामान्य ओपन-सोर्स समाधानों की तुलना में चलाने में काफी तेज़ और बनाए रखने में आसान बनाते हैं।

पॉइंट-टू-पॉइंट कम्युनिकेशन

Rank0-आधारित वेट ट्रांसफर, उन्हें inference node पर वितरित करने से पहले वेट्स को इकट्ठा करता है.
पॉइंट-टू-पॉइंट संचार, दोनों training nodes से वेट्स को inference node के नेटवर्क interfaces तक स्थानांतरित करता है.

एक सामान्य पैटर्न सभी पैरामीटर्स को रैंक-0 GPUs के माध्यम से फ़नल करना है: ट्रेनिंग रैंक-0 पर इकट्ठा करना, इंफ़रेंस रैंक-0 को भेजना, फिर वापस स्कैटर करना। यह जल्दी ही एक बाधा बन जाता है, जो सिंगल GPU की PCIe बैंडविड्थ और NIC (उदा., 400 Gbps ≈ 50 GB/s) द्वारा सीमित होता है।

इसके विपरीत, हमारा पॉइंट-टू-पॉइंट सेटअप हर ट्रेनिंग GPU को सीधे हर इंफ़रेंस GPU पर भेजने की अनुमति देता है, जो एक सिंगल लिंक के बजाय पूरे नेटवर्क फैब्रिक को संतृप्त करता है।

वन-साइडेड डेटा ट्रांसफ़र

RPC-आधारित ट्रांसफर एक inference-side buffer और worker calls का उपयोग करता है, जबकि one-sided transfer inference-worker की भागीदारी के बिना सीधे वेट्स लिखता है.

कुछ सिस्टम प्रत्येक टेंसर के लिए इंफ़रेंस इंजन के update_weight() मेथड को कॉल करने पर निर्भर करते हैं। इसका मतलब है इंफ़रेंस कोड में दखल देने वाले बदलाव, साथ ही RPCs, सीरियलाइज़ेशन और कंट्रोल-प्लेन समन्वय से ओवरहेड।

RDMA WRITE प्रिमिटिव के साथ, हम अतिरिक्त प्रतियों के बिना इंफ़रेंस GPU मेमोरी पर चुपचाप वेट्स अपडेट करते हैं। इसमें कोई कंट्रोल प्लेन संदेश और कोई CPU कंट्रोल लॉजिक शामिल नहीं है। इंफ़रेंस इंजन में किसी संशोधन की आवश्यकता नहीं है।

पाइपलाइनिंग

वेट ट्रांसफ़र प्रक्रिया चार प्रकार के हार्डवेयर संसाधनों का लाभ उठा सकती है: (1) होस्ट-डिवाइस डेटा मूवमेंट (2) प्रोजेक्शन फ़्यूज़न और क्वांटाइज़ेशन के लिए GPU कंप्यूटेशन (3) डेटा प्लेन के लिए RDMA नेटवर्क (4) कंट्रोल प्लेन के लिए ईथरनेट।

हमारा डिज़ाइन वेट ट्रांसफ़र टास्क को पाइपलाइन चरणों में विभाजित करता है, जिससे विभिन्न हार्डवेयर संसाधनों में आसानी से ओवरलैप हो सकता है।

स्टैटिक शेड्यूल

कुछ कार्यान्वयन हर ट्रेनिंग चरण में ट्रांसफ़र शेड्यूल की पुनर्गणना करते हैं, बार-बार मेटाडेटा इकट्ठा करते हैं और निर्देश वितरित करते हैं। यह अनावश्यक कंट्रोल-प्लेन लेटेंसी को जोड़ता है।

हमारा शेड्यूल इनिशियलाइज़ेशन पर एक बार गणना किया जाता है। प्रत्येक ट्रेनिंग इटरेशन केवल योजना को दोहराता है: कंट्रोलर एक “गो” सिग्नल जारी करता है, और GPUs अपने पूर्व-निर्धारित रास्तों का पालन करते हैं। एक्जीक्यूशन अनुमानित और हल्का है।

साफ़ अलगाव

पूरी वेट अपडेट प्रक्रिया को एक ही मोनोलिथिक फ़ंक्शन में उलझाना आकर्षक लगता है: मेटाडेटा इकट्ठा करना, नाम मिलान, इंट्रा-नोड गैदरिंग, प्रोजेक्शन फ़्यूज़न, क्वांटाइज़ेशन, सब्सलाइसिंग कम्युनिकेशन वर्ल्ड, इंटर-नोड नेटवर्क ट्रांसफ़र। इसे सही ढंग से प्रोग्राम करना कठिन है, और ऑप्टिमाइज़ करना और भी कठिन है।

हमारी इंजीनियरिंग में, हम इन चरणों को अलग-अलग घटकों के रूप में अलग करते हैं। प्रत्येक घटक को अलग से यूनिट टेस्ट, तर्क और ऑप्टिमाइज़ किया जा सकता है।

निष्कर्ष

तेज़, विश्वसनीय वेट ट्रांसफ़र बड़े पैमाने पर RL फ़ाइन-ट्यूनिंग के लिए एक महत्वपूर्ण बिल्डिंग ब्लॉक है। RDMA WRITE प्रिमिटिव, एक स्टैटिक ट्रांसफ़र शेड्यूल, और पाइपलाइंड एक्जीक्यूशन को मिलाकर, हमने Kimi-K2 पर ट्रिलियन-पैरामीटर अपडेट को केवल 1.3 सेकंड तक कम कर दिया है। यह दृष्टिकोण तर्क करने में सरल, बनाए रखने में आसान है, और पारंपरिक डिज़ाइनों की बाधाओं से बचता है।

उद्धरण

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}
}