2 सेकंदांपेक्षा कमी वेळेत RL पोस्ट-ट्रेनिंगसाठी वेट ट्रान्सफर

अति-जलद क्रॉस-GPU मॉडेल सिंक

लेखकPerplexity Research

आम्ही नुकतीच Kimi-K2 (1T पॅरामीटर्स) साठी 1.3 सेकंदांची क्रॉस-मशीन पॅरामीटर अपडेट्स साध्य केली आहेत, ज्यामध्ये 256 ट्रेनिंग GPU (BF16) वरून 128 इन्फरन्स GPU (FP8) वर वेट्स ट्रान्सफर केले आहेत.

असिंक्रोनस रीइंफोर्समेंट लर्निंग फाईन-ट्यूनिंगमध्ये, ट्रेनिंग आणि इन्फरन्स स्वतंत्र GPU वर चालतात. प्रत्येक ट्रेनिंग स्टेप नंतर, नवीन वेट्स इन्फरन्स नोड्सवर पुश केले जाणे आवश्यक आहे. अनेक विद्यमान फ्रेमवर्क ट्रिलियन-पॅरामीटर मॉडेल्ससाठी अनेक सेकंद—किंवा मिनिटे सुद्धा—घेतात.

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

डेस्टिनेशन साईडला ट्रान्सफर्ससाठी कोणतीही सूचना मिळत नाही. यामुळे इन्फरन्स नोड्सवर कोणत्याही कंट्रोल लॉजिकशिवाय ट्रेनिंग नोड्सद्वारे चालवले जाणारे कमी-सुप्तता (low-latency), उच्च-थ्रूपुट, झिरो-कॉपी ट्रान्सफर्स मिळतात.

हाय-लेव्हल वर्कफ्लो

मेटाडाटा संकलन – कंट्रोलर सर्व ट्रेनिंग आणि इन्फरन्स GPU मधून पॅरामीटर मेटाडाटा गोळा करतो.

शिड्यूल गणना – कंट्रोलर एक स्टॅटिक वेट ट्रान्सफर शिड्यूल मोजतो, कोणता ट्रेनिंग GPU कोणत्या इन्फरन्स GPU ला कोणते पॅरामीटर पाठवतो आणि कोणत्या क्रमाने पाठवतो हे मॅप करतो.

शिड्यूल वितरण – कंट्रोलर सर्व ट्रेनिंग GPU ला शिड्यूल पाठवतो.

एक्झिक्युशन – प्रत्येक ट्रेनिंग स्टेप नंतर, कंट्रोलर ट्रेनिंग GPU ला ट्रान्सफर सुरू करण्याचे सिग्नल देतो.

वेट ट्रान्सफर एक्झिक्युशन

हाय-लेव्हल वर्कफ्लो परिभाषित केल्यानंतर, ट्रिलियन-पॅरामीटर स्केलवर वेट ट्रान्सफर्स कार्यक्षमतेने कसे एक्झिक्युट करायचे हे मुख्य आव्हान आहे. येथे आम्ही एक्झिक्युशन पाथचे तपशील वर्णन करतो.

डिव्हाइसमेश आणि मेश ग्रुप्स

ट्रेनिंगमधील पॅरामीटर्स FSDP प्लेसमेंटनुसार वितरित केले जातात. full_tensor() वापरून, DeviceMesh मधील सर्व GPU संपूर्ण पॅरामीटर पुन्हा तयार करू शकतात, म्हणूनच सर्वजण वेट ट्रान्सफरसाठी सोर्स म्हणून काम करू शकतात.

अनेक विलग (disjoint) DeviceMeshes मिळून एक मेश ग्रुप तयार होतो. एकाच ग्रुपमधील DeviceMeshes विलग असल्यामुळे, त्यांचे ट्रान्सफर्स एकमेकांमध्ये अडथळा आणत नाहीत आणि पूर्णपणे समांतरपणे चालतात. मेश ग्रुप्सच्या दरम्यान, आम्ही ऑर्डरिंग लागू करण्यासाठी एक ग्लोबल बॅरियर इन्सर्ट करतो.

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

वजन-हस्तांतरण पाइपलाइन MeshGroups वर होस्ट-टू-डिव्हाइस कॉपी, टेन्सर असेंब्ली, फ्यूजन आणि क्वांटायझेशन, तसेच RDMA हस्तांतरण यांना ओव्हरलॅप करते.

आम्ही प्रत्येक पॅरामीटर टेन्सरचे ट्रान्सफर एक टास्क मानतो. वेट ट्रान्सफर प्रक्रियेमध्ये विविध प्रकारचे हार्डवेअर सोर्स वापरले जातात, म्हणून आम्ही वेट ट्रान्सफर टास्कचे वेगवेगळ्या पाईपलाईन टप्प्यांमध्ये विभाजन करतो जे वेळेत ओव्हरलॅप होतात:

होस्ट ते डिव्हाइस memcpy — FSDP ने CPU वर वेट ऑफलोड केल्यास

पॅरामीटर तयारी — full_tensor() सह संपूर्ण वेट रिकन्स्ट्रक्ट करा, प्रोजेक्शन फ्यूजन लागू करा, आवश्यक असल्यास क्वांटाइज करा.

RDMA ट्रान्सफर — रिमोट इन्फरन्स GPU मेमरीवर झिरो-कॉपी राईट

ग्लोबल बॅरियर — सर्व full_tensor() कॉल्स पूर्ण झाल्यावर, इथरनेटद्वारे GLOO वापरून मेश ग्रुप्समध्ये सिंक्रोनाइझ करा.

अंमलबजावणीमध्ये, आम्ही प्रत्येक पाईपलाईन टप्प्यासाठी टास्कची FIFO रांगे राखतो. जेव्हा रांगेचा हेड टास्क टप्पा पूर्ण करतो, तेव्हा तो पुढील टप्प्याच्या रांगेच्या शेवटी हलवला जातो.

GPU मेमरी वापर नियंत्रण

full_tensor() आणि इतर GPU ऑपरेशन्स अतिरिक्त GPU मेमरी वापरतात. आऊट ऑफ मेमरी एरर टाळण्यासाठी, आम्ही वर्तमान ऑन-द-फ्लाय टास्क कॉन्फिगर करण्यायोग्य वॉटरमार्कपेक्षा कमी तात्पुरती GPU मेमरी वापरत असल्यास केवळ टास्कचे एक्झिक्युशन सुरू करतो.

हे जलद आणि सोपे का आहे

अनेक डिझाइन निवडींमुळे आमची सिस्टीम सामान्य ओपन-सोर्स सोल्यूशन्सपेक्षा चालवण्यासाठी लक्षणीयरीत्या जलद आणि राखण्यासाठी सोपी बनते.

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

Rank0-आधारित वजन हस्तांतरण, ते inference node ला वितरित करण्यापूर्वी वजन गोळा करते.
Point-to-point संप्रेषण दोन्ही training nodes मधून weights inference node च्या network interfaces कडे हस्तांतरित करते.

एक सामान्य पॅटर्न सर्व पॅरामीटर्सना रँक-0 GPU द्वारे चॅनल करणे आहे: ट्रेनिंग रँक-0 वर गोळा करणे, इन्फरन्स रँक-0 ला पाठवणे, नंतर पुन्हा स्कॅटर करणे. हे लवकरच एक चोक पॉइंट बनते, जे एकाच GPU च्या PCIe बँडविड्थ आणि NIC द्वारे मर्यादित असते (उदा., 400 Gbps ≈ 50 GB/s).

याच्या विपरीत, आमचे पॉइंट-टू-पॉइंट सेटअप प्रत्येक ट्रेनिंग GPU ला प्रत्येक इन्फरन्स GPU वर थेट पाठवण्याची परवानगी देते, एकाच लिंकऐवजी संपूर्ण नेटवर्क फॅब्रिक संतृप्त करते.

वन-साईड डेटा ट्रान्सफर

RPC-आधारित हस्तांतरण inference-side buffer आणि worker calls वापरते, तर one-sided transfer inference-worker सहभागाशिवाय weights थेट लिहिते.

काही सिस्टीम्स प्रत्येक टेन्सरसाठी इन्फरन्स इंजिनच्या update_weight() पद्धतीमध्ये कॉल करण्यावर अवलंबून असतात. याचा अर्थ इन्फरन्स कोडमध्ये हस्तक्षेप करणारे बदल, तसेच RPCs, सिरीयलायझेशन आणि कंट्रोल-प्लेन कोर्डीनेशनचा ओव्हरहेड असा होतो.

RDMA WRITE प्रिमिटिव्हसह, आम्ही अतिरिक्त प्रतींशिवाय इन्फरन्स GPU मेमरीवर वेट्स मूकपणे अपडेट करतो. कोणतेही कंट्रोल प्लेन संदेश आणि कोणतेही CPU कंट्रोल लॉजिक गुंतलेले नाही. इन्फरन्स इंजिनमध्ये कोणत्याही संशोधनाची आवश्यकता नाही.

पाईपलाइनिंग

वेट ट्रान्सफर प्रक्रियेमध्ये चार प्रकारचे हार्डवेअर रिसोर्सेस वापरले जाऊ शकतात: (1) होस्ट-डिव्हाइस डेटा मूव्हमेंट (2) प्रोजेक्शन फ्यूजन आणि क्वांटायझेशनसाठी GPU कॉम्प्युटेशन (3) डेटा प्लेनसाठी RDMA नेटवर्क (4) कंट्रोल प्लेनसाठी इथरनेट.

आमच्या डिझाइनने वेट ट्रान्सफर टास्कचे पाईपलाईन टप्प्यांमध्ये विभाजन केले आहे, ज्यामुळे वेगवेगळ्या हार्डवेअर रिसोर्सेसवर सहज ओव्हरलॅप करता येते.

स्टॅटिक शिड्यूल

काही अंमलबजावणी प्रत्येक ट्रेनिंग स्टेपवर ट्रान्सफर शिड्यूलची पुनर्गणना करतात, वारंवार मेटाडाटा गोळा करतात आणि सूचना वितरित करतात. यामुळे अनावश्यक कंट्रोल-प्लेन सुप्तता जोडली जाते.

आमचे शिड्यूल एकदाच इनिशिएलायझेशनवर मोजले जाते. प्रत्येक ट्रेनिंग इटेशन फक्त प्लॅन रिप्ले करते: कंट्रोलर “go” सिग्नल जारी करतो, आणि GPU त्यांचे पूर्व-नियुक्त मार्ग पाळतात. एक्झिक्युशन अंदाज लावणारे आणि हलके असते.

स्वच्छ विभाजन

संपूर्ण वेट अपडेट प्रक्रियेला एकाच मोनोलिथिक फंक्शनमध्ये गुंतवणे प्रलोभनदायक ठरू शकते: मेटाडाटा गोळा करणे, नाव मॅचिंग, इंट्रा-नोड गोळा करणे, प्रोजेक्शन फ्यूजन, क्वांटायझेशन, सबस्लायसिंग कम्युनिकेशन वर्ल्ड, इंटर-नोड नेटवर्क ट्रान्सफर. हे योग्यरित्या प्रोग्राम करणे कठीण आहे, आणि ऑप्टिमाइझ करणे आणखी कठीण आहे.

आमच्या इंजिनिअरिंगमध्ये, आम्ही या पायऱ्या स्वतंत्र घटक म्हणून वेगळ्या करतो. प्रत्येक घटक युनिट टेस्ट केला जाऊ शकतो, विचार केला जाऊ शकतो आणि स्वतंत्रपणे ऑप्टिमाइझ केला जाऊ शकतो.

निष्कर्ष

जलद, विश्वसनीय वेट ट्रान्सफर हा मोठ्या प्रमाणावरील 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}
}