2 സെക്കൻഡിൽ താഴെ സമയത്തിനുള്ളിൽ RL പോസ്റ്റ്-ട്രെയിനിങ്ങിനായുള്ള വെയിറ്റ് കൈമാറ്റം

ഞങ്ങൾ അടുത്തിടെ Kimi-K2 (1T പാരാമീറ്ററുകൾക്കായി) 1.3 സെക്കൻഡ് ക്രോസ്-മെഷീൻ പാരാമീറ്റർ അപ്‌ഡേറ്റുകൾ കൈവരിച്ചു, 256 ട്രെയിനിങ് GPU-കളിൽ (BF16) നിന്ന് 128 ഇൻഫറൻസ് GPU-കളിലേക്ക് (FP8) വെയിറ്റുകൾ കൈമാറി.

അസിൻക്രണസ് റീഇൻഫോഴ്‌സ്‌മെന്റ് ലേണിങ് ഫൈൻ-ട്യൂണിങ്ങിൽ, ട്രെയിനിങ്ങും ഇൻഫറൻസും വെവ്വേറെ GPU-കളിലാണ് പ്രവർത്തിക്കുന്നത്. ഓരോ ട്രെയിനിങ് ഘട്ടത്തിന് ശേഷവും പുതിയ വെയിറ്റുകൾ ഇൻഫറൻസ് നോഡുകളിലേക്ക് പുഷ് ചെയ്യേണ്ടതുണ്ട്. ട്രില്യൺ പാരാമീറ്റർ മോഡലുകൾക്കായി നിലവിലുള്ള പല ഫ്രെയിംവർക്കുകൾക്കും നിരവധി സെക്കൻഡുകളോ മിനിറ്റുകളോ വരെ എടുക്കുന്നു.

RDMA പോയിന്റ്-ടூ-പോയിന്റ് കമ്മ്യൂണിക്കേഷൻ പ്രയോജനപ്പെടുത്തുന്നതിലൂടെ, ഇൻഫറൻസ് എൻജിൻ മാറ്റാതെ തന്നെ വെയിറ്റ് കൈമാറ്റം അതിവേഗത്തിലാക്കാനും കോഡ് എഴുതാനും പരിപാലിക്കാനും എളുപ്പമാക്കാനും ഞങ്ങൾക്ക് കഴിയുന്നു.

ഞങ്ങളുടെ പരിഹാരം നിർമ്മിച്ചിരിക്കുന്നത് RDMA WRITE-നെ അടിസ്ഥാനമാക്കിയാണ്. ഉറവിടം നേരിട്ട് ലക്ഷ്യസ്ഥാനത്തിന്റെ GPU മെമ്മറിയിലേക്ക് എഴുതുന്ന ഒറ്റത്തവണയുള്ള (one-sided) ഒരു പ്രിമിറ്റീവ് ആണിത്.

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

കൈമാറ്റത്തിനായി ലക്ഷ്യസ്ഥാനത്ത് അറിയിപ്പ് പോലും ലഭിക്കില്ല. ഇത് ഇൻഫറൻസ് നോഡുകളിൽ കൺട്രോൾ ലോജിക് ഒന്നുമില്ലാതെ, ട്രെയിനിങ് നോഡുകൾ വഴി കുറഞ്ഞ ലേറ്റൻസി, ഉയർന്ന ത്രൂപുട്ട്, സീറോ-കോപ്പി കൈമാറ്റങ്ങൾ എന്നിവ നൽകുന്നു.

മെറ്റാഡാറ്റ ശേഖരണം - എല്ലാ ട്രെയിനിങ്, ഇൻഫറൻസ് GPU-കളിൽ നിന്നും കൺട്രോളർ പാരാമീറ്റർ മെറ്റാഡാറ്റ ശേഖരിക്കുന്നു.

ഷെഡ്യൂൾ കണക്കുകൂട്ടൽ - ഏത് ട്രെയിനിങ് GPU ഏത് പാരാമീറ്റർ ഏത് ഇൻഫറൻസ് GPU-ലേക്ക് അയക്കണം, ഏത് ക്രമത്തിൽ അയക്കണം എന്ന് മാപ്പ് ചെയ്ത് കൺട്രോളർ ഒരു സ്റ്റാറ്റിക് വെയിറ്റ് ട്രാൻസ്ഫർ ഷെഡ്യൂൾ കണക്കാക്കുന്നു.

ഷെഡ്യൂൾ വിതരണം - കൺട്രോളർ എല്ലാ ട്രെയിനിങ് GPU-കളിലേക്കും ഷെഡ്യൂൾ അയക്കുന്നു.

എക്സിക്യൂഷൻ - ഓരോ ട്രെയിനിങ് ഘട്ടത്തിന് ശേഷവും, കൈമാറ്റം ആരംഭിക്കാൻ കൺട്രോളർ ട്രെയിനിങ് GPU-കൾക്ക് സിഗ്നൽ നൽകുന്നു.

ഉയർന്ന തലത്തിലുള്ള വർക്ക്ഫ്ലോ നിർവചിച്ചുകഴിഞ്ഞാൽ, ട്രില്യൺ പാരാമീറ്റർ സ്കെയിലിൽ വെയിറ്റ് കൈമാറ്റങ്ങൾ കാര്യക്ഷമമായി എങ്ങനെ നടപ്പിലാക്കാം എന്നതാണ് പ്രധാന വെല്ലുവിളി. എക്സിക്യൂഷൻ പാതയുടെ വിശദാംശങ്ങൾ ഞങ്ങൾ ഇവിടെ വിവരിക്കുന്നു.

FSDP പ്ലേസ്‌മെന്റുകൾ അനുസരിച്ചാണ് ട്രെയിനിങ്ങിലെ പാരാമീറ്ററുകൾ വിതരണം ചെയ്യുന്നത്. full_tensor() ഉപയോഗിച്ച്, ഒരു DeviceMesh-ലെ എല്ലാ GPU-കൾക്കും മുഴുവൻ പാരാമീറ്ററും പുനർനിർമ്മിക്കാൻ കഴിയും, അതിനാൽ എല്ലാറ്റിനും വെയിറ്റ് കൈമാറ്റത്തിനുള്ള ഉറവിടമായി പ്രവർത്തിക്കാൻ കഴിയും.

പരസ്പര ബന്ധമില്ലാത്ത ഒന്നിലധികം DeviceMesh-കൾ ചേർന്ന് ഒരു മെഷ് ഗ്രൂപ്പ് രൂപീകരിക്കുന്നു. ഒരേ ഗ്രൂപ്പിലെ DeviceMesh-കൾ പരസ്പര ബന്ധമില്ലാത്തതായതിനാൽ, അവയുടെ കൈമാറ്റങ്ങൾ പരസ്പരം തടസ്സപ്പെടുത്തുന്നില്ല കൂടാതെ പൂർണ്ണമായും സമാന്തരമായി പ്രവർത്തിക്കാനും കഴിയും. ക്രമം നടപ്പിലാക്കുന്നതിനായി മെഷ് ഗ്രൂപ്പുകൾക്കിടയിൽ ഞങ്ങൾ ഒരു ഗ്ലോബൽ ബാരിയർ ഉൾപ്പെടുത്തുന്നു.

വെയിറ്റ്-ട്രാൻസ്ഫർ പൈപ്പ്‌ലൈൻ ഹോസ്റ്റ്-ടു-ഡിവൈസ് പകർപ്പുകൾ, ടെൻസർ അസംബ്ലി, ഫ്യൂഷൻ, ക്വാണ്ടൈസേഷൻ, കൂടാതെ MeshGroups-കുറുകെയുള്ള RDMA ട്രാൻസ്ഫറുകൾ എന്നിവ ഓവർലാപ്പ് ചെയ്യുന്നു.

ഓരോ പാരാമീറ്റർ ടെൻസറിന്റെയും കൈമാറ്റത്തെ ഞങ്ങൾ ഒരു ടാസ്കായി പരിഗണിക്കുന്നു. വെയിറ്റ് ട്രാൻസ്ഫർ പ്രക്രിയ വിവിധ തരം ഹാർഡ്‌വെയർ ഉറവിടങ്ങൾ ഉപയോഗിക്കുന്നു, അതിനാൽ ഞങ്ങൾ ഒരു വെയിറ്റ് ട്രാൻസ്ഫർ ടാസ്കിനെ സമയത്തിൽ ഓവർലാപ്പ് ചെയ്യുന്ന വ്യത്യസ്ത പൈപ്പ്‌ലൈൻ ഘട്ടങ്ങളായി തിരിക്കുന്നു:

ഹോസ്റ്റ് മുതൽ ഡിവൈസ് വരെയുള്ള memcpy — FSDP CPU-വിലേക്ക് വെയിറ്റ് ഓഫ്‌ലോഡ് ചെയ്യുകയാണെങ്കിൽ

പാരാമീറ്റർ തയ്യാറെടുപ്പ് — full_tensor() ഉപയോഗിച്ച് മുഴുവൻ വെയിറ്റും പുനർനിർമ്മിക്കുക, പ്രൊജക്ഷൻ ഫ്യൂഷൻ പ്രയോഗിക്കുക, ആവശ്യമെങ്കിൽ ക്വാണ്ടൈസ് ചെയ്യുക.

RDMA കൈമാറ്റം — റിമോട്ട് ഇൻഫറൻസ് GPU മെമ്മറിയിലേക്ക് സീറോ-കോപ്പി റൈറ്റ് ചെയ്യുക

ഗ്ലോബൽ ബാരിയർ — എല്ലാ full_tensor() കോളുകളും പൂർത്തിയായ ശേഷം, ഇഥർനെറ്റ് വഴി GLOO ഉപയോഗിച്ച് മെഷ് ഗ്രൂപ്പുകളിലുടനീളം സിൻക്രണൈസ് ചെയ്യുക.

നടപ്പിലാക്കുമ്പോൾ, ഓരോ പൈപ്പ്‌ലൈൻ ഘട്ടത്തിനും ഞങ്ങൾ ടാസ്കുകളുടെ ഒരു FIFO ക്യൂ സൂക്ഷിക്കുന്നു. ക്യൂ ടാസ്കിന്റെ ഹെഡ് ആ ഘട്ടം പൂർത്തിയാക്കുമ്പോഴെല്ലാം, അത് അടുത്ത ഘട്ട ക്യൂവിന്റെ ടെയിലിലേക്ക് നീക്കപ്പെടുന്നു.

full_tensor()-ഉം മറ്റ് GPU പ്രവർത്തനങ്ങളും അധിക GPU മെമ്മറി ഉപയോഗം കൊണ്ടുവരുന്നു. മെമ്മറി തീർന്നുപോകുന്ന പിഴവ് (out of memory error) ഒഴിവാക്കാൻ, നിലവിലുള്ള on-the-fly ടാസ്കുകൾ ക്രമീകരിക്കാവുന്ന വാട്ടർമാർക്കിനേക്കാൾ കുറഞ്ഞ താൽക്കാലിക GPU മെമ്മറി ഉപയോഗിക്കുന്നുണ്ടെങ്കിൽ മാത്രമേ ഞങ്ങൾ ഒരു ടാസ്കിന്റെ എക്സിക്യൂഷൻ ആരംഭിക്കുകയുള്ളൂ.

സാധാരണ ഓപ്പൺ സോഴ്സ് പരിഹാരങ്ങളെ അപേക്ഷിച്ച് ഞങ്ങളുടെ സിസ്റ്റം പ്രവർത്തിപ്പിക്കാൻ വളരെ വേഗതയുള്ളതും പരിപാലിക്കാൻ എളുപ്പവുമാക്കുന്ന നിരവധി ഡിസൈൻ ചോയ്സുകൾ ഉണ്ട്.

Rank0 അടിസ്ഥാനമാക്കിയ വെയിറ്റ് ട്രാൻസ്ഫർ, അവ ഇൻഫറൻസ് നോഡിലേക്ക് വിതരണം ചെയ്യുന്നതിന് മുമ്പ് വെയിറ്റുകൾ ശേഖരിക്കുന്നു.

പോയിന്റ്-ടു-പോയിന്റ് കമ്മ്യൂണിക്കേഷൻ ഇരുവശത്തുമുള്ള പരിശീലന നോഡുകളിൽ നിന്നുള്ള വെയിറ്റുകൾ ഇൻഫറൻസ് നോഡിന്റെ നെറ്റ്‌വർക്ക് ഇന്റർഫേസുകളിലേക്ക് കൈമാറുന്നു.

എല്ലാ പാരാമീറ്ററുകളും റാങ്ക്-0 GPU-കൾ വഴി ഫണൽ ചെയ്യുക എന്നത് ഒരു സാധാരണ രീതിയാണ്: ട്രെയിനിങ് റാങ്ക്-0-ൽ ശേഖരിക്കുക, ഇൻഫറൻസ് റാങ്ക്-0-ലേക്ക് അയക്കുക, തുടർന്ന് വീണ്ടും സ്കാറ്റർ ചെയ്യുക. ഇത് പെട്ടെന്ന് തന്നെ ഒരു തടസ്സമായി മാറുന്നു, ഇത് ഒരొറ്റ GPU-ന്റെ PCIe ബാൻഡ്‌വിഡ്റ്റും NIC-ഉം (ഉദാഹരണത്തിന്, 400 Gbps ≈ 50 GB/s) വഴി പരിമിതപ്പെടുത്തപ്പെടുന്നു.

ഇതിനു വിപരീതമായി, ഞങ്ങളുടെ പോയിന്റ്-ടു-പോയിന്റ് സജ്ജീകരണം ഓരോ പരിശീലന GPU-ക്കും ഓരോ ഇൻഫറൻസ് GPU-ിലേക്കും നേരിട്ട് അയയ്ക്കാൻ അനുവദിക്കുന്നു, ഒറ്റ ലിങ്കിന് പകരം മുഴുവൻ നെറ്റ്‌വർക്ക് ഫാബ്രിക്കിനെയും പരമാവധി ഉപയോഗപ്പെടുത്തിക്കൊണ്ട്.

RPC അടിസ്ഥാനമാക്കിയ ട്രാൻസ്ഫർ ഇൻഫറൻസ്-വശത്തെ ബഫറും വർക്കർ കോളുകളും ഉപയോഗിക്കുന്നു, അതേസമയം one-sided ട്രാൻസ്ഫർ ഇൻഫറൻസ്-വർക്കറുടെ പങ്കാളിത്തമില്ലാതെ വെയിറ്റുകൾ നേരിട്ട് എഴുതുന്നു.

ചില സിസ്റ്റങ്ങൾ ഓരോ ടെൻസറിനുമായി ഇൻഫറൻസ് എൻജിന്റെ update_weight() രീതിയിലേക്ക് കോൾ ചെയ്യുന്നതിനെ ആശ്രയിക്കുന്നു. അതിനർത്ഥം ഇൻഫറൻസ് കോഡിൽ വലിയ മാറ്റങ്ങൾ വരുത്തേണ്ടിവരും എന്നാണ്, കൂടാതെ RPC-കൾ, സീരിയലൈസേഷൻ, കൺട്രോൾ-പ്ലെയിൻ കോർഡിനേഷൻ എന്നിവയിൽ നിന്നുള്ള ഓവർഹെഡും ഉണ്ടാകും.

RDMA WRITE പ്രിമിറ്റീവ് ഉപയോഗിച്ച്, അധിക പകർപ്പുകൾ ഇല്ലാതെ, ഇൻഫറൻസ് GPU മെമ്മറിയിൽ ഞങ്ങൾ നിശബ്ദമായി വെയിറ്റുകൾ അപ്‌ഡേറ്റ് ചെയ്യുന്നു. കൺട്രോൾ പ്ലെയിൻ സന്ദേശമോ CPU കൺട്രോൾ ലോഗിക്കോ ഉൾപ്പെട്ടിട്ടില്ല. ഇൻഫറൻസ് എൻജിനിൽ മാറ്റങ്ങളൊന്നും വരുത്തേണ്ടതില്ല.

വെയിറ്റ് ട്രാൻസ്ഫർ പ്രക്രിയയ്ക്ക് നാല് തരം ഹാർഡ്‌വെയർ ഉറവിടങ്ങൾ പ്രയോജനപ്പെടുത്താം: (1) ഹോസ്റ്റ്-ഡിവൈസ് ഡാറ്റ മൂവ്‌മെന്റ് (2) പ്രൊജക്ഷൻ ഫ്യൂഷനും ക്വാണ്ടൈസേഷനുമുള്ള GPU കമ്പ്യൂട്ടേഷൻ (3) ഡാറ്റ പ്ലെയിനുവേണ്ടി RDMA നെറ്റ്‌വർക്ക് (4) കൺട്രോൾ പ്ലെയിനുവേണ്ടി ഇഥർനെറ്റ്.

ഞങ്ങളുടെ ഡിസൈൻ വെയിറ്റ് ട്രാൻസ്ഫർ ടാസ്കുകളെ പൈപ്പ്‌ലൈൻ ഘട്ടങ്ങളായി തിരിക്കുന്നു, ഇത് വ്യത്യസ്ത ഹാർഡ്‌വെയർ ഉറവിടങ്ങളിലുടനീളം എളുപ്പത്തിൽ ഓവർലാപ്പ് ചെയ്യാൻ അനുവദിക്കുന്നു.

ചില നടപ്പിലാക്കലുകൾ ഓരോ ട്രെയിനിങ് ഘട്ടത്തിലും ഒരു ട്രാൻസ്ഫർ ഷെഡ്യൂൾ പുനർനിർമ്മിക്കുന്നു, മെറ്റാഡാറ്റ ആവർത്തിച്ച് ശേഖരിക്കുകയും നിർദ്ദേശങ്ങൾ വിതരണം ചെയ്യുകയും ചെയ്യുന്നു. ഇത് ആവശ്യമില്ലാത്ത കൺട്രോൾ-പ്ലെയിൻ ലേറ്റൻസി വർദ്ധിപ്പിക്കുന്നു.

ഞങ്ങളുടെ ഷെഡ്യൂൾ സമാരംഭിക്കുമ്പോൾ (at initialization) ഒരിക്കൽ കണക്കാക്കപ്പെടുന്നു. ഓരോ ട്രെയിനിങ് ആവർത്തനവും പ്ലാൻ ലളിതമായി വീണ്ടും പ്രവർത്തിപ്പിക്കുന്നു: കൺട്രോളർ ഒരു "go" സിഗ്നൽ നൽകുന്നു, കൂടാതെ GPU-കൾ അവ മുൻകൂട്ടി നിശ്ചയിച്ച റൂട്ടുകൾ പിന്തുടരുന്നു. എക്സിക്യൂഷൻ പ്രവചിക്കാവുന്നതും ഭാരം കുറഞ്ഞതുമാണ്.

മെറ്റാഡാറ്റ ശേഖരിക്കൽ, പേരുകളുടെ പൊരുത്തപ്പെടുത്തൽ, ഇൻട്രാ-നോഡ് സമാഹരണം, പ്രൊജക്ഷൻ ഫ്യൂഷൻ, ക്വാണ്ടൈസേഷൻ, സബ്‌സ്ലൈസിംഗ് കമ്മ്യൂണിക്കേഷൻ വേൾഡ്, ഇന്റർ-നോഡ് നെറ്റ്‌വർക്ക് ട്രാൻസ്ഫർ എന്നിവയെല്ലാം ഒന്നായ മോണൊലിത്തിക് ഫംഗ്ഷനിൽ കുടുക്കുന്നത് ആകർഷകമായി തോന്നാം. അത് ശരിയായി പ്രോഗ്രാം ചെയ്യുന്നത് ബുദ്ധിമുട്ടാണ്, ഒപ്റ്റിമൈസ് ചെയ്യുന്നത് അതിലും ബുദ്ധിമുട്ടാണ്.

ഞങ്ങളുടെ എൻജിനീയറിങ്ങിൽ, ഞങ്ങൾ ഈ ഘട്ടങ്ങളെ വ്യക്തിഗത ഘടകങ്ങളായി വേർതിരിക്കുന്നു. ഓരോ ഘടകവും യൂണിറ്റ് ടെസ്റ്റ് ചെയ്യാനും ന്യായീകരിക്കാനും ഒറ്റപ്പെട്ട് ഒപ്റ്റിമൈസ് ചെയ്യാനും കഴിയും.

വിപുലമായ സ്കെയിലിലുള്ള RL ഫൈൻ-ട്യൂണിംഗിന് വേഗതയുള്ള, വിശ്വസനീയമായ വെയ്റ്റ് ട്രാൻസ്ഫർ ഒരു നിർണായക നിർമാണഘടകമാണ്. RDMA WRITE primitive, ഒരു static transfer schedule, കൂടാതെ pipelined execution എന്നിവ സംയോജിപ്പിച്ചതിലൂടെ, Kimi-K2-ൽ ഞങ്ങൾ ട്രില്യൺ-പരാമീറ്റർ അപ്‌ഡേറ്റുകൾ വെറും 1.3 seconds ആയി കുറച്ചു. ഈ സമീപനം തർക്കപരമായി മനസ്സിലാക്കാൻ എളുപ്പമാണ്, പരിപാലിക്കാനും എളുപ്പമാണ്, കൂടാതെ പരമ്പരാഗത രൂപകല്പനകളിലെ തടസ്സങ്ങൾ ഒഴിവാക്കുന്നു.

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: ഒറ്റത്തവണയുള്ള കൈമാറ്റങ്ങൾ (one-sided transfers)

ഉയർന്ന തലത്തിലുള്ള വർക്ക്ഫ്ലോ

വെയിറ്റ് ട്രാൻസ്ഫർ എക്സിക്യൂഷൻ

DeviceMesh-ഉം Mesh ഗ്രൂപ്പുകളും

ടാസ്ക് പൈപ്പ്‌ലൈൻ

GPU മെമ്മറി ഉപയോഗ നിയന്ത്രണം

ഇത് വേഗതയേറിയതും ലളിതവുമാകുന്നത് কেন

പോയിന്റ്-ടൂ-പോയിന്റ് കമ്മ്യൂണിക്കേഷൻ

ഒറ്റത്തവണയുള്ള ഡാറ്റ കൈമാറ്റം (One-sided data transfer)

പൈപ്പ്‌ലൈനിങ്

സ്റ്റാറ്റിക് ഷെഡ്യൂൾ

വ്യക്തമായ വേർതിരിവ്

ഉപസംഹാരം

ഉദ്ധരണി