കാര്യക്ഷമവും പോർട്ടബിളുമായ മിക്സ്ചർ ഓഫ് എക്സ്പെർട്ട്സ് ആശയവിനിമയം

സ്റ്റാൻഡേർഡ് ഓൾ-ടു-ഓൾ കമ്മ്യൂണിക്കേഷൻ പ്രിമിറ്റീവുകളെ അപേക്ഷിച്ച് 10 മടങ്ങ് വേഗതയേറിയ പ്രകടനം കൈവരിക്കുന്ന മിക്സ്ചർ ഓഫ് എക്സ്പെർട്ട്സ് (MoE) ആശയവിനിമയത്തിനായി ഞങ്ങൾ ഉയർന്ന പ്രകടനമുള്ള, പോർട്ടബിൾ, ഓപ്പൺ സോഴ്സ് ലൈബ്രറി അവതരിപ്പിക്കുന്നു.

മെച്ചപ്പെട്ട MoE ആശയവിനിമയ കാര്യക്ഷമത നൽകുന്ന നിരവധി പ്രധാന സാങ്കേതിക നവീകരണങ്ങൾ ഞങ്ങളുടെ നടപ്പാക്കലിൽ ഉണ്ട്:

GPU-ആരംഭിച്ച ആശയവിനിമയം (IBGDA): CPU പങ്കാളിത്തം ഒഴിവാക്കിക്കൊണ്ട് ലേറ്റൻസി ഗണ്യമായി കുറയ്ക്കുന്നു, നേരിട്ടുള്ള GPU-ടു-NIC ആശയവിനിമയത്തെ പിന്തുണയ്ക്കുന്നു

ആശയവിനിമയവും കമ്പ്യൂട്ടേഷനും ഓവർലാപ്പ് ചെയ്യുക: പ്രത്യേക സെൻഡ്, റിസീവ് ഘട്ടങ്ങളുള്ള സ്പ്ലിറ്റ് കേർണൽ ആർക്കിടെക്ചർ, നെറ്റ്‌വർക്ക് കൈമാറ്റങ്ങൾക്കൊപ്പം ഒരേസമയം കമ്പ്യൂട്ടേഷൻ മുന്നോട്ട് കൊണ്ടുപോകാൻ സഹായിക്കുന്നു

ഏറ്റവും വേഗതയേറിയ സിംഗിൾ-നോഡ് പ്രകടനം: സിംഗിൾ-നോഡ് കോൺഫിഗറേഷനുകളിൽ മുമ്പ് ഏറ്റവും വേഗതയേറിയ നടപ്പാക്കലിനേക്കാൾ 2.5 മടങ്ങ് കുറഞ്ഞ ലേറ്റൻസി

കാര്യക്ഷമവും പോർട്ടബിളുമായ മൾട്ടി-നോഡ് പ്രകടനം: ഞങ്ങളുടെ നടപ്പാക്കൽ സാധാരണ ഓൾ-ടു-ഓൾ ആശയവിനിമയത്തേക്കാൾ 10 മടങ്ങ് വേഗത കൈവരിക്കുന്നു. ഉയർന്ന പ്രത്യേക നടപ്പാക്കലുകളേക്കാൾ ഏകദേശം 2 മടങ്ങ് മന്ദഗതിയിലാണെങ്കിലും, NVSHMEM പതിപ്പുകളിലും നെറ്റ്‌വർക്ക് പരിതസ്ഥിതികളിലും (NVLink, CX-7, EFA) ഞങ്ങളുടെ സമീപനം മികച്ച പോർട്ടബിലിറ്റി വാഗ്ദാനം ചെയ്യുന്നു.

ഈ ലൈബ്രറി പൂർണ്ണമായും ഓപ്പൺ സോഴ്സ് ആണ് കൂടാതെ https://github.com/ppl-ai/pplx-kernels എന്നതിൽ ലഭ്യമാണ്.

ഈ ലേഖനത്തിൽ, വലിയ തോതിലുള്ള MoE മോഡലുകളിലെ എക്സ്പെർട്ട് പാരലലിസത്തിന്റെ വെല്ലുവിളികളെക്കുറിച്ച് ഞങ്ങൾ പഠിക്കുകയും, വിതരണം ചെയ്ത GPU പരിതസ്ഥിതികളിൽ കാര്യക്ഷമമായ ടോക്കൺ ഡിസ്പാച്ചും കോമ്പിനേഷനുമുള്ള ഞങ്ങളുടെ സമീപനം വിവരിക്കുകയും ചെയ്യുന്നു.

ആമുഖം

ഡീപ്സീക്ക് R1, മിക്സ്ട്രൽ 8x7B പോലെയുള്ള മിക്സ്ചർ ഓഫ് എക്സ്പെർട്ട്സ് (MoE) മോഡലുകൾ, ഓരോ ടോക്കണും സജീവമാക്കുന്ന ഭാരങ്ങളുടെ എണ്ണം പരിമിതപ്പെടുത്തിക്കൊണ്ട് ഡെൻസ് മോഡലുകളെ മെച്ചപ്പെടുത്തുന്നു. ലെയറുകളുടെ എണ്ണം വർദ്ധിപ്പിച്ചോ അല്ലെങ്കിൽ ഓരോ ഡീകോഡർ ലെയറിന്റെയും മൾട്ടി-ലെയർ പെർസെപ്ട്രോണിന്റെ (MLP) ലീനിയർ ലെയറുകളുടെ വലിപ്പം കൂട്ടിയോ ഉയർന്ന പാരാമീറ്റർ എണ്ണം കൈവരിക്കുന്നതിന് പകരം, MoE മോഡലുകൾ പരമ്പരാഗത MLP-ക്ക് പകരം ഒന്നിലധികം എക്സ്പെർട്ടുകളെയും ഒരു റൂട്ടറെയും ഉപയോഗിക്കുന്നു. ഉദാഹരണത്തിന്, ഡീപ്സീക്ക് R1-ന്റെ 671B പാരാമീറ്ററുകളിൽ, അനുമാന സമയത്ത് (inference) നൽകിയിരിക്കുന്ന ടോക്കണുമായി ഗുണിക്കുന്നത് 37B മാത്രം ആണ്. ഇത് ഒരു ടോക്കൺ ഡീകോഡ് ചെയ്യാൻ ആവശ്യമായ കമ്പ്യൂട്ടേഷന്റെ അളവ് ഗണ്യമായി കുറയ്ക്കുകയും ഡെൻസ് മോഡലുമായി താരതമ്യം ചെയ്യുമ്പോൾ ലേറ്റൻസി മെച്ചപ്പെടുത്തുകയും ചെയ്യുന്നു.

ഡെൻസ് മോഡലുകളുമായി താരതമ്യം ചെയ്യുമ്പോൾ MoE മോഡലുകൾ അനുമാനത്തിനായി (inference) ചില അധിക വെല്ലുവിളികൾ അവതരിപ്പിക്കുന്നു. എക്സ്പെർട്ടുകൾ സ്വയം ചെറിയ MLP ലെയറുകൾ ആണെങ്കിലും, ഓരോ ഡീകോഡർ ലെയറിലും ഒരു ടോക്കൺ ഏത് എക്സ്പെർട്ടുകളിലേക്കാണ് അയക്കേണ്ടതെന്ന് തീരുമാനിക്കുന്ന ഒരു റൂട്ടർ ഉൾപ്പെടുന്നു, ഓരോ ടോക്കണും ഒന്നിലധികം എക്സ്പെർട്ടുകളിലേക്ക് അയക്കുന്നു. റൂട്ടർ സാധാരണയായി ഒരു പ്രൊബബിലിറ്റി ഡിസ്ട്രിബ്യൂഷൻ നിർമ്മിക്കുന്ന ചെറിയൊരു ലീനിയർ ലെയറാണ്. സാധാരണയായി ടോപ്-K സ്കോറുകൾ ഉള്ള എക്സ്പെർട്ടുകൾ തിരഞ്ഞെടുക്കപ്പെടുന്നു, കൂടാതെ റൂട്ടറിന്റെ പ്രൊബബിലിറ്റി ഡിസ്ട്രിബ്യൂഷനിൽ നിന്ന് ഉരുത്തിരിഞ്ഞ ഭാരം ഉപയോഗിച്ച് എക്സ്പെർട്ട് ഔട്ട്പുട്ടുകളെ ഗുണിച്ച് കൂട്ടിയുള്ള വെയ്റ്റഡ് ആവറേജ് ആയി അന്തിമ ആക്റ്റിവേഷൻ കണക്കാക്കപ്പെടുന്നു.

വിതരണം ചെയ്ത MoE സിസ്റ്റങ്ങളിലെ ലേറ്റൻസി കുറയ്ക്കുന്നതിന്, ഒന്നിലധികം ഉപകരണങ്ങളിൽ സമാന്തരത പ്രയോജനപ്പെടുത്താം, എന്നാൽ ഇത് ആശയവിനിമയ വെല്ലുവിളികൾ അവതരിപ്പിക്കുന്നു. മിക്സ്ട്രൽ 8x7B അല്ലെങ്കിൽ Llama-70B പോലുള്ള മോഡലുകൾ ഒരൊറ്റ നോഡിലുടനീളമുള്ള 8 ഉപകരണങ്ങളിൽ ഉൾക്കൊള്ളുന്നു, ഇത് വേഗതയേറിയതും കുറഞ്ഞ ലേറ്റൻസിയുള്ളതുമായ NVLink ഇന്റർകണക്ഷനുകളിൽ (900Gbps വരെ) നിന്ന് പ്രയോജനം നേടുന്നു. എന്നിരുന്നാലും, വലിയ മോഡലുകൾക്ക് InfiniBand ഉപയോഗിച്ച് ഒന്നിലധികം നോഡുകളിൽ എക്സ്പെർട്ടുകൾ ഷാർഡ് ചെയ്യേണ്ടതുണ്ട് (പരമാവധി 400Gbps), ഇത് അധിക ലേറ്റൻസി വെല്ലുവിളികൾ അവതരിപ്പിക്കുന്നു.

വലുതാക്കുക

ഈ ലേഖനത്തിൽ, നിയുക്ത എക്സ്പെർട്ടുകളെ ধারণ ചെയ്യുന്ന ഉപകരണങ്ങളിലേക്ക് ടോക്കണുകൾ കാര്യക്ഷമമായി അയക്കുകയും എക്സ്പെർട്ട് ഔട്ട്പുട്ടുകളുടെ വെയ്റ്റഡ് സം കണക്കാക്കി അവയെ തിരികെ ശേഖരിക്കുകയും ചെയ്യുന്ന GPU കേർണലുകളുടെ ഞങ്ങളുടെ നടപ്പാക്കൽ വിവരിച്ചുകൊണ്ട്, എക്സ്പെർട്ട് പാരലലിസം എന്ന പ്രശ്നം ഞങ്ങൾ പരിശോധിക്കുന്നു. നിർദ്ദിഷ്ട നെറ്റ്‌വർക്ക് ഹാർഡ്‌വെയറിനെ ലക്ഷ്യമിട്ടുള്ള ഒപ്റ്റിമൈസ് ചെയ്ത നടപ്പാക്കലുകൾ നിലവിലുണ്ടെങ്കിലും, NVSHMEM പ്രിമിറ്റീവുകളുടെ ഒരു ചെറിയ സെറ്റിനെ ആശ്രയിക്കുന്ന എന്നാൽ സാധാരണ ഓൾ-ടു-ഓൾ ആശയവിനിമയത്തേക്കാൾ 10 മടങ്ങ് വേഗതയേറിയ പ്രകടനം നൽകുന്ന ഉയർന്ന പോർട്ടബിളായ ഒരു നടപ്പാക്കൽ ഞങ്ങൾ അവതരിപ്പിക്കുന്നു. വിവിധ നെറ്റ്‌വർക്ക് പരിതസ്ഥിതികളിൽ മികച്ച പോർട്ടബിലിറ്റി നിലനിർത്തിക്കൊണ്ട് ഞങ്ങളുടെ നടപ്പാക്കൽ സിംഗിൾ-നോഡ് കോൺഫിഗറേഷനുകളിൽ അത്യാധുനിക പ്രകടനം കൈവരിക്കുന്നു.

ഷാർഡിംഗും പാരലലിസവും

കാര്യക്ഷമമായ അനുമാനത്തിനായി (inference), ഒരു മോഡലിന്റെ ഭാരങ്ങൾ ഡിവൈസ് മെമ്മറിയിൽ സൂക്ഷിക്കണം, അതേസമയം മോഡലിലൂടെയുള്ള ഫോർവേഡ് പാസിന് ആവശ്യമായ ആക്റ്റിവേഷനുകൾ, KV കാഷുകൾ, മറ്റ് ബഫറുകൾ എന്നിവയ്ക്കുള്ള മതിയായ സ്ഥലവും ബാക്കി വയ്ക്കണം. ഏറ്റവും ശക്തമായ മോഡലുകൾ പോലും ഏറ്റവും ശക്തമായ GPU-കളുടെ ശേഷിയേക്കാൾ കവിയുന്നു, അതിനാൽ അനുമാനം ഒന്നിലധികം ഉപകരണങ്ങളിലായി വ്യാപിപ്പിക്കേണ്ടതുണ്ട്, അവയ്ക്ക് കൂട്ടായി മോഡലിന്റെ ഭാരങ്ങൾ സൂക്ഷിക്കാൻ കഴിയും. ഭാരങ്ങളുടെ ഷാർഡിംഗ് സ്കീമുകളെ അടിസ്ഥാനമാക്കി, കമ്പ്യൂട്ടേഷൻ ഓഫ്‌ലോഡ് ചെയ്യാനും ഉപകരണങ്ങൾ സിൻക്രണൈസ് ചെയ്യാനും വ്യത്യസ്ത ആശയവിനിമയ, കമ്പ്യൂട്ടേഷൻ സ്കീമുകൾ ഉപയോഗിക്കണം.

വലുതാക്കുക

ആദ്യ ചിത്രത്തിൽ ചിത്രീകരിച്ചിരിക്കുന്ന എക്സ്പെർട്ട് പാരലലിസം (EP), എക്സ്പെർട്ട് കമ്പ്യൂട്ടേഷനെ മാത്രം സമാന്തരമാക്കുന്നു. വ്യത്യസ്ത എക്സ്പെർട്ടുകൾ അവയുടെ ഭാരങ്ങൾ ধারণ ചെയ്യുന്ന വ്യത്യസ്ത ഉപകരണങ്ങൾക്ക് നൽകിയിരിക്കുന്നു. റൂട്ടിംഗിന് ശേഷം, ടോക്കണുകൾ ബന്ധപ്പെട്ട ഉപകരണത്തിലേക്ക് അയയ്ക്കപ്പെടുന്നു, ഫലങ്ങൾ ശേഖരിക്കുകയും തുടർന്ന് കുമുലേറ്റു ചെയ്യുകയും ചെയ്യുന്നു. മോഡലിന്റെ മറ്റ് ഭാഗങ്ങളിലെ സമാന്തരതയുടെ അളവിനെ ആശ്രയിച്ചിരിക്കുന്നു റൂട്ടിംഗിന്റെ സങ്കീർണ്ണത: മറ്റ് ലെയറുകൾ റെപ്ലിക്കേറ്റ് ചെയ്യുന്നത് റൂട്ടിംഗിന്റെ ആവശ്യകത പൂർണ്ണമായും ഇല്ലാതാക്കും, കാരണം ഓരോ റാങ്കിനും ലോക്കലായി റെപ്ലിക്കേറ്റ് ചെയ്ത റൂട്ടിംഗ് ടേബിളിൽ നിന്ന് ടോക്കണുകൾ തിരഞ്ഞെടുക്കാനാകും. എന്നിരുന്നാലും, ഒരു റാങ്ക് മാത്രമാണ് റൂട്ടിംഗ് നടത്തുന്നതെങ്കിൽ, ടോക്കണുകൾ, സൂചികകൾ (indices), ഭാരങ്ങൾ എന്നിവ അതത് എക്സ്പെർട്ടുകളിലേക്ക് അയക്കാൻ ഒരു ബ്രോഡ്കാസ്റ്റ് ആവശ്യമാണ്. അവസാനമായി, ഒരു ഓൾ-ഗാതർ അല്ലെങ്കിൽ ഓൾ-ടു-ഓൾ ബ്രോഡ്കാർസ്റ്റ് ഔട്ട്പുട്ട് ടോക്കണുകളെ മോഡലിന്റെ എക്സിക്യൂഷൻ തുടരുന്ന റാങ്കുമായി സിൻക്രണൈസ് ചെയ്യുന്നു. ടോർച്ച് ഇതിനകം ആവശ്യമായ പ്രിമിറ്റീവുകൾ നൽകുന്നതിനാൽ അത്തരം ഒരു നടപ്പാക്കൽ താരതമ്യേന ലളിതമാണ്, എങ്കിലും എക്സ്പെർട്ടുകളിൽ നിന്ന് ടോക്കണുകൾ ശേഖരിക്കുന്ന ഗാതർ ഓപ്പറേഷനുമായി അക്യുമുലേഷൻ ഫ്യൂസ് ചെയ്യുന്നതിലൂടെ ചില നേട്ടങ്ങൾ ലഭിച്ചേക്കാം.

വലുതാക്കുക

മോഡൽ നോൺ-എക്സ്പെർട്ട് ലെയറുകൾ, അതായത് അറ്റൻഷൻ, നോം, സാംപ്ലിംഗ് എന്നിവ പ്രവർത്തിപ്പിക്കുമ്പോൾ ക്ലസ്റ്ററിലെ നോഡുകൾ നിഷ്‌ക്രിയമായിരിക്കാൻ സാധ്യതയുള്ളതിനാൽ, എക്സ്പെർട്ട്-മാത്രം സമാന്തരത അനുയോജ്യമായ രീതിയിൽ സ്കെയിൽ ചെയ്യുന്നില്ല. എന്നിരുന്നാലും, ഈ ലെയറുകളുടെ കമ്പ്യൂട്ടേഷനും, പ്രധാനമായും അറ്റൻഷനും, ടെൻസർ പാരലലിസത്തിൽ (TP) നിന്ന് പ്രയോജനം നേടാം. മിക്ക മോഡലുകളും മൾട്ടി-ഹെഡ് അറ്റൻഷനെ ആശ്രയിക്കുന്നു, അതായത് അറ്റൻഷൻ ഹെഡുകളും അവയുടെ ബന്ധപ്പെട്ട Q, K, V പ്രൊജക്ഷനുകളും ഉപകരണങ്ങളിലുടനീളം ഷാർഡ് ചെയ്യാനും, വിവിധ ലെയറുകൾക്കിടയിലുള്ള സ്ലൈസുകൾ റെപ്ലിക്കേറ്റ് ചെയ്യാനോ ശേഖരിക്കാനോ കഴിയും. എല്ലാ ഉപകരണങ്ങളിലും അറ്റൻഷൻ വ്യാപിപ്പിച്ചിട്ടുണ്ടെങ്കിൽ, ഒരു ഓൾ-ഗാതർ ആക്റ്റിവേഷനുകളെ സിൻക്രണൈസ് ചെയ്യാൻ അനുവദിക്കുന്നു, ഇത് റൂട്ടിംഗ് റെപ്ലിക്കേറ്റ് ചെയ്യാൻ ആവശ്യപ്പെടുന്നു, കാര്യക്ഷമമായ ഒരു നടപ്പാക്കലിനായി എക്സ്പെർട്ട് പാരലലിസം കേസിന് സമാനമായ സിൻക്രണൈസേഷൻ പ്രിമിറ്റീവുകൾ ആവശ്യമാണ്. എന്നിരുന്നാലും, ഈ തലത്തിലുള്ള പാരലലിസത്തിന് പരിമിതികളുണ്ട്, കാരണം അറ്റൻഷൻ ഹെഡുകളുടെ എണ്ണം ഒരു നിശ്ചിത പരിധിക്ക് താഴെയായി കുറയ്ക്കുന്നത് കുറയുന്ന റിട്ടേണുകൾ നൽകും.

വലുതാക്കുക

എല്ലാ ഉപകരണങ്ങളും മികച്ച രീതിയിൽ ഉപയോഗിക്കുന്നതിനും 128 അല്ലെങ്കിൽ 256 GPU-കൾ വരെയുള്ള വളരെ ഉയർന്ന അളവിലുള്ള എക്സ്പെർട്ട് പാരലലിസത്തെ പിന്തുണയ്ക്കുന്നതിനും, ഡാറ്റാ പാരലലിസം (DP) ആവശ്യമാണ്. ഈ സ്കീമിന് കീഴിൽ, ഒരു മോഡൽ വിഭജിച്ചിരിക്കുന്ന ഉപകരണങ്ങൾ അഭ്യർത്ഥനകൾ ഒരേസമയം കൈകാര്യം ചെയ്യാൻ ഗ്രൂപ്പുചെയ്യപ്പെടുന്നു, അറ്റൻഷൻ കണക്കാക്കുകയും അവരുടെ ലോക്കൽ ഗ്രൂപ്പിലുടനീളം ഷാർഡ് ചെയ്ത KV കാഷുകൾ പരിപാലിക്കുകയും ചെയ്യുന്നു. ഈ പാരലൽ ഗ്രൂപ്പുകളുടെ ഒന്നിലധികം ഇൻസ്റ്റൻസുകൾ എക്സ്പെർട്ട് വിലയിരുത്തലിൽ സഹകരിക്കുന്നു, ഓരോന്നും എക്സ്പെർട്ട് ഭാരങ്ങളുടെ വ്യത്യസ്ത ഉപസെറ്റ് ഹോസ്റ്റ് ചെയ്യുന്നു. അറ്റൻഷൻ ഹെഡുകളുടെ എണ്ണത്തെ ആശ്രയിച്ച്, നോഡ്-ഇൻട്രാ ആശയവിനിമയം വേഗതയേറിയതായതിനാൽ, ഒരു ഗ്രൂപ്പിന് സാധാരണയായി ഒരു മുഴുവൻ നോഡിന്റെ വലിപ്പം വരെ വ്യാപിക്കാം. ഉദാഹരണത്തിന്, മുകളിലുള്ള ചിത്രത്തിൽ, ഒരു DP റാങ്ക് രണ്ട് അഭ്യർത്ഥനകൾ, A, B എന്നിവ സ്വതന്ത്രമായി സേവിക്കുന്നു, അറ്റൻഷൻ, നോം, കൂടാതെ അഭ്യർത്ഥനകൾക്കായുള്ള മറ്റ് ബുക്ക് കീപ്പിംഗ് എന്നിവ കൈകാര്യം ചെയ്യുന്നു. മറ്റ് DP റാങ്ക് വ്യത്യസ്തമായ ഒരു അഭ്യർത്ഥന C പ്രോസസ്സ് ചെയ്യുന്നു. എന്നിരുന്നാലും, ആദ്യത്തെ നോഡ് എക്സ്പെർട്ടുകളുടെ പകുതിയും മറ്റേ നോഡ് മറുപകുതിയും ഹോസ്റ്റ് ചെയ്യുന്നു, അതിനാൽ A, B എന്നിവയിൽ നിന്നുള്ള റൂട്ടിംഗ് ടോക്കണുകൾക്ക് ശേഷം രണ്ടാമത്തെ നോഡിലേക്കും തിരിച്ചും അയച്ചേക്കാം. ഇത് ഒരു സ്പാഴ്സ് കമ്മ്യൂണിക്കേഷൻ പ്രശ്നത്തിലേക്ക് നയിക്കുന്നു: ഓരോ ഉപകരണവും മറ്റേതെങ്കിലും ഡെസ്റ്റിനേഷൻ റാങ്കിലേക്ക് വ്യത്യസ്ത എണ്ണം ടോക്കണുകൾ അയച്ചേക്കാം. ടോർച്ചിൽ നിന്നുള്ള നിലവിലുള്ള പ്രിമിറ്റീവുകൾ, പ്രധാനമായും all_to_all, പ്രത്യേകിച്ചും നന്നായി അനുയോജ്യമല്ല, കാരണം അവയ്ക്ക് ചിലതരം പാഡിംഗ് അല്ലെങ്കിൽ GPU-ടു-CPU സിൻക്രണൈസേഷനും മെറ്റാഡാറ്റ ബ്രോഡ്കാസ്റ്റും ആവശ്യമായി വന്നേക്കാം. ആശയവിനിമയം ഫലപ്രദമായി നടപ്പിലാക്കാൻ, റൂട്ടിംഗ് വിവരങ്ങൾ പ്രോസസ്സ് ചെയ്യുന്നതിനും CPU ഇടപെടലില്ലാതെ ഉപകരണങ്ങളിൽ നിന്ന് ആശയവിനിമയം ആരംഭിക്കുന്നതിനും കസ്റ്റം കേർണലുകൾ ആവശ്യമാണ്. റൂട്ടിംഗിന് ശേഷം, ഒരു ഡിസ്പാച്ച് കേർണൽ ടോക്കണുകൾ അവ റൂട്ട് ചെയ്ത റാങ്കുകളിലേക്ക് അയക്കണം, അതേസമയം കോംബൈൻ വശത്ത് നിലവിലെ DP ഗ്രൂപ്പിലെ അഭ്യർത്ഥനകളുമായി ബന്ധപ്പെട്ട ആക്റ്റിവേഷനുകൾ ശേഖരിക്കണം. കൂടാതെ, തുടർന്നുള്ള ലെയറുകൾക്കായി എല്ലാ റാങ്കുകളിലും ടോക്കണുകൾ ശരിയായി റെപ്ലിക്കേറ്റ് ചെയ്യുന്നതിനും ബ്രോഡ്കാസ്റ്റ് കുറയ്ക്കുന്നതിനും ഒന്നിലധികം ഉപകരണങ്ങൾ ഉണ്ടായിരിക്കാവുന്ന ഒരു DP ഗ്രൂപ്പിനുള്ളിൽ ജോലി ബാലൻസ് ചെയ്യണം.

NVSHMEM

NVSHMEM എന്നത് NVIDIA-നിർദ്ദിഷ്ട OpenSHMEM നടപ്പാക്കലാണ്, ഇത് അടിസ്ഥാന ഹാർഡ്‌വെയറിന്റെ സങ്കീർണ്ണത അമൂർത്തമാക്കുന്ന പോർട്ടബിൾ ഡിവൈസ്-ഇന്റർ-ഡിവൈസ് ആശയവിനിമയ സൗകര്യങ്ങൾ നൽകുന്നു. API-ക്ക് ഡിവൈസ്-ടു-ഡിവൈസ് റീഡുകളും റൈറ്റുകളും പ്രകടിപ്പിക്കാൻ കഴിയും, അവ വ്യക്തിഗത ട്രാൻസ്പോർട്ട് ലെയറുകളുടെ പ്രിമിറ്റീവുകളിലേക്ക് മാപ്പ് ചെയ്യപ്പെടുന്നു. NVLink ഉം RDMA ഉം പിന്തുണയ്ക്കുന്നു, ഇത് പോർട്ടബിലിറ്റിയുടെ ഒരു അളവ് നൽകുകയും ഒരു നോഡിനുള്ളിലോ ഒന്നിലധികം നോഡുകളുടെ ക്ലസ്റ്ററിലോ പ്രവർത്തിക്കാൻ NVSHMEM കേർണലുകളെ അനുവദിക്കുകയും ചെയ്യുന്നു. ഞങ്ങളുടെ കേർണലുകളിൽ, ഡാറ്റ കൈമാറുന്നതിനും പങ്കെടുക്കുന്ന ഉപകരണങ്ങളെ സിൻക്രണൈസ് ചെയ്യുന്നതിനും ഞങ്ങൾ റിമോട്ട് മെമ്മറി റൈറ്റുകളും സിഗ്നലിംഗ് പ്രവർത്തനങ്ങളും ഉപയോഗിക്കുന്നു.

NVSHMEM പ്രവർത്തനങ്ങൾ സിമ്മട്രിക് മെമ്മറി എന്ന ആശയത്തിന് ചുറ്റുമാണ് നിർമ്മിച്ചിരിക്കുന്നത്: ഡിവൈസ്-ഇന്റർ ആശയവിനിമയത്തിൽ പങ്കെടുക്കുന്ന എല്ലാ ഉപകരണങ്ങളിലും അലോക്കേറ്റ് ചെയ്തിട്ടുള്ള ബഫറുകളിലാണ് അവ പ്രവർത്തിക്കുന്നത്. ഒരു പ്രവർത്തനത്തിന്റെ ലക്ഷ്യസ്ഥാനം വ്യക്തമാക്കാൻ ഒരു ലോക്കൽ ഡിവൈസിന് അതിന്റെ സ്വന്തം ലോക്കൽ ബഫറിനുള്ളിൽ നിന്ന് ഉരുത്തിരിഞ്ഞ വിലാസവും ടാർഗെറ്റ് റാങ്കിന്റെ സൂചികയും (index) ഉപയോഗിക്കാം. താഴെയുള്ള ചിത്രം ഈ ആശയം ചിത്രീകരിക്കുന്നു: രണ്ട് GPU-കളും ഒരേ വലിപ്പമുള്ള സിമ്മട്രിക് ബഫറുകൾ അലോക്കേറ്റ് ചെയ്യുന്നു, അവയിലേക്ക് src പോയിന്ററുകൾ നിലനിർത്തുന്നു. ആദ്യത്തെ റാങ്ക് രണ്ടാമത്തേതിലേക്ക് 3 പൂർണ്ണസംഖ്യകൾ (integers) അയക്കാൻ ആഗ്രഹിക്കുന്നു, അവ ബഫറിന്റെ തുടക്കത്തിൽ സ്ഥാപിക്കുന്നു: nvshmem_int_put_nbi ടാർഗെറ്റ് ഡിവൈസ് വ്യക്തമാക്കിക്കൊണ്ട് ലോക്കൽ ബഫറിൽ നിന്ന് സ്റ്റാർട്ട് അഡ്രസ് ഉരുത്തിരിഞ്ഞെടുക്കുന്നു. രണ്ടാമത്തെ റാങ്ക് അതിന്റെ സ്വന്തം ബഫറിൽ നിന്ന് ഒരു ഓഫ്സെറ്റ് ഉരുത്തിരിഞ്ഞെടുക്കുന്നു, ഒരു ഘടകം (element) ആദ്യ ഡിവൈസിലേക്ക് അയക്കുന്നു, ഒന്ന് ഓഫ്സെറ്റ് ചെയ്യുന്നു. ഡെസ്റ്റിനേഷൻ വിലാസങ്ങൾ എപ്പോഴും nvshmem_alloc ഉപയോഗിച്ച് അലോക്കേറ്റ് ചെയ്ത സിമ്മട്രിക് ബഫറുകൾ ആയിരിക്കണം, അതേസമയം സോഴ്സ് ബഫറുകൾ NVSHMEM-ൽ മുൻകൂട്ടി രജിസ്റ്റർ ചെയ്തിട്ടുണ്ടെങ്കിൽ, ഡിവൈസ് മെമ്മറിയുടെ ആർബിട്രറി മേഖലകൾ ആകാം.

NVSHMEM വൈവിധ്യമാർന്ന പ്രിമിറ്റീവുകൾ നൽകുമ്പോൾ, ഞങ്ങളുടെ കേർണലുകൾ 3 ഫങ്ഷനുകളെ മാത്രം ആശ്രയിക്കുന്നു, എല്ലാ സിൻക്രണൈസേഷനും ഫെൻസിങ്ങും അവയ്ക്ക് മുകളിൽ നിർമ്മിക്കുന്നു.

വലുതാക്കുക

nvshmemx_putmem_signal_nbi_warp: റിമോട്ട് ഡിവൈസിൽ ഒരു ഫ്ലാഗ് സജ്ജീകരിച്ചുകൊണ്ട്, ഒരു ഡിവൈസിൽ നിന്ന് മറ്റൊന്നിലേക്ക് ഡാറ്റയുടെ ഒരു ബ്ലോക്ക് കൈമാറുന്നു. പ്രവർത്തനം ഒന്നുകിൽ സജ്ജമാക്കുന്നു (NVSHMEM_SIGNAL_SET) അല്ലെങ്കിൽ 64-ബിറ്റ് ലൊക്കേഷൻ വർദ്ധിപ്പിക്കുന്നു (NVSHMEM_SIGNAL_ADD). മുഴുവൻ മെമ്മറി ബ്ലോക്കും കൈമാറിയ ശേഷമാണ് ഫ്ലാഗ് അപ്ഡേറ്റ് ചെയ്യുന്നത്. ഫ്ലാഗിലെ മാറ്റം റിമോട്ട് ഡിവൈസ് നിരീക്ഷിച്ചിട്ടുണ്ടെങ്കിൽ, അതിന് സ്വന്തം മെമ്മറിയിലെ ബഫർ സുരക്ഷിതമായി ആക്സസ് ചെയ്യാൻ കഴിയും. ഡാറ്റ കൈമാറ്റവും സിൻക്രണൈസേഷനും ഒരേസമയം നടപ്പിലാക്കാൻ ഈ ഫങ്ഷൻ ഉപയോഗപ്രദമാണ്.

nvshmemx_signal_op ഒരു ഒരൊറ്റ മെമ്മറി ലൊക്കേഷനിൽ പ്രവർത്തിക്കുന്നു, സാധാരണയായി 64-ബിറ്റ് ഫ്ലാഗ്, അത് ആറ്റോമിക്കായി സജ്ജീകരിക്കുകയോ വർദ്ധിപ്പിക്കുകയോ ചെയ്യുന്നു. മെറ്റാഡാറ്റ അയക്കുന്നതിനും ഉപകരണങ്ങൾ സിൻക്രണൈസ് ചെയ്യുന്നതിനും ഇത് ഉപയോഗപ്രദമാണ്.

nvshmem_uint64_wait_until റിമോട്ട് അപ്ഡേറ്റ് ചെയ്യുന്നതുവരെ ഒരു ഫ്ലാഗ് പോൾ ചെയ്യാൻ, സിഗ്നലിന്റെ സ്വീകരിക്കുന്ന അന്തിമഭാഗത്ത് ഉപയോഗിക്കുന്നു.

nvshmem_quiet അല്ലെങ്കിൽ nvshmem_fence വഴിയുള്ള പ്രവർത്തനങ്ങൾക്കിടയിലുള്ള ഏത് തരത്തിലുള്ള ഓർഡറിംഗും ഉറപ്പാക്കുന്നത് ചെലവേറിയതാണ്, അതിനാൽ കേർണലുകളുടെ ഞങ്ങളുടെ നടപ്പാക്കൽ മറ്റ് NVSHMEM പ്രിമിറ്റീവുകളുടെ ഉപയോഗം ഒഴിവാക്കുകയും വ്യക്തിഗത സെൻഡ്-റിസീവ് ജോഡികൾ ഒഴികെയുള്ള ഏതെങ്കിലും ഡിപൻഡൻസികൾ കുറയ്ക്കുകയും ചെയ്യുന്നു. ഡാറ്റ അയക്കുമ്പോൾ, ലോക്കൽ റാങ്കിൽ നിന്ന് ഡാറ്റ പുറത്തേക്ക് അയക്കുന്നതുവരെ പോലും കാത്തിരിക്കാതെ, ഞങ്ങൾ എപ്പോഴും ഫങ്ഷനുകളുടെ നോൺ-ബ്ലോക്കിംഗ് പതിപ്പ് ഉപയോഗിക്കുന്നു. കേർണലുകളുടെ ഗ്രൂപ്പിലുടനീളമുള്ള മറ്റ് സൂക്ഷ്മമായ സിൻക്രണൈസേഷൻ, NIC-കൾക്ക് അവയിലേക്ക് ആക്സസ് ആവശ്യമായി വരുമ്പോൾ ഈ ബഫറുകൾ നശിപ്പിക്കുന്ന രീതിയിൽ ഓവർറൈറ്റ് ചെയ്യപ്പെടുന്നില്ലെന്ന് ഉറപ്പാക്കുന്നു.

CPU-കളെ ഉൾപ്പെടുത്താതെ തന്നെ ഡാറ്റ കൈമാറ്റത്തിനായി NIC-ക്ക് GPU മെമ്മറി നേരിട്ട് റീഡ് ചെയ്യാനും റൈറ്റ് ചെയ്യാനും കഴിയുന്ന തരത്തിൽ NVSHMEM GPUDirect RDMA ഉപയോഗിക്കുന്നു. കൂടാതെ, ConnectX NIC-കളിൽ, NVSHEMEM GPU-ആരംഭിച്ച ആശയവിനിമയത്തെ പിന്തുണയ്ക്കുന്നു (ഇത് ഇൻഫിനിബാൻഡ് GPUDirect Async അല്ലെങ്കിൽ IBGDA എന്നും അറിയപ്പെടുന്നു), ഇത് GPU-ന് NIC-ലേക്ക് നേരിട്ട് RDMA പ്രവർത്തനങ്ങൾ പോസ്റ്റ് ചെയ്യാനും NIC-ന്റെ കംപ്ലീഷൻ ക്യൂ നേരിട്ട് പോൾ ചെയ്യാനും അനുവദിക്കുന്നു. GPU-ആരംഭിച്ച ആശയവിനിമയത്തെ പിന്തുണയ്ക്കാത്ത പ്ലാറ്റ്‌ഫോമുകളിൽ, ആശയവിനിമയം ആരംഭിക്കുന്നതിനും GPU-ന് വേണ്ടി കംപ്ലീഷൻ പോൾ ചെയ്യുന്നതിനും NVSHMEM ഹോസ്റ്റ് CPU-ൽ ഒരു “പ്രോക്സി” ത്രെഡ് സൃഷ്ടിക്കുന്നു. GPU-ആരംഭിച്ച ആശയവിനിമയം പിന്തുണയ്ക്കുന്നുണ്ടോ ഇല്ലയോ എന്നത് പരിഗണിക്കാതെ NVSHMEM പ്രോഗ്രാം പോർട്ടബിൾ ആണ്. എന്നിരുന്നാലും, GPU-ആരംഭിച്ച ആശയവിനിമയം ലേറ്റൻസി ഗണ്യമായി കുറയ്ക്കുന്നു കാരണം ഇത് CPU-ലേക്കുള്ള വഴിമാറലിനെ പൂർണ്ണമായും ഒഴിവാക്കുന്നു.

പോർട്ടബിൾ കേർണലുകൾ

ഡിസ്പാച്ച്, കോംബൈൻ കേർണലുകളുടെ ഒരു ജോഡിയിലൂടെ ഞങ്ങൾ MoE ആശയവിനിമയം നടപ്പിലാക്കുന്നു. ഓരോ റാങ്കിലും ടോക്കണുകളും റൂട്ടിംഗ് വിവരങ്ങളും റീഡ് ചെയ്യുന്നതിനും, അനുയോജ്യമായ എക്സ്പെർട്ടുകളിലേക്ക് അവയെ ഡിസ്പാച്ച് ചെയ്യുന്നതിനും ഡിസ്പാച്ച് കേർണലുകൾ ഉത്തരവാദികളാണ്. കോംബൈൻ കേർണലുകൾ എക്സ്പെർട്ടുകൾ നിർമ്മിക്കുന്ന ആക്റ്റിവേഷനുകൾ ശേഖരിക്കുകയും അവയെ അവയുടെ സോഴ്സ് റാങ്കുകളിലേക്ക് തിരികെ അയക്കുകയും ചെയ്യുന്നു, അതേസമയം റൂട്ടർ കണക്കാക്കിയ ഭാരങ്ങളെ അടിസ്ഥാനമാക്കി തിരഞ്ഞെടുത്ത എക്സ്പെർട്ടുകളിൽ നിന്നുള്ള വെയ്റ്റഡ് ആവറേജ് കണക്കാക്കുകയും ചെയ്യുന്നു. ഡാറ്റ കൈമാറ്റങ്ങൾ കമ്പ്യൂട്ടേഷനുമായി ഓവർലാപ്പ് ചെയ്യാൻ അനുവദിക്കുന്നതിനായി കേർണലുകളെ സെൻഡ്, റിസീവ് ഘടകങ്ങളായി കൂടുതൽ വിഭജിച്ചിരിക്കുന്നു. സെൻഡ് കേർണലുകൾ എല്ലാം നോൺ-ബ്ലോക്കിംഗും നോൺ-സിൻക്രണൈസിംഗുമാണ്: അവ റിമോട്ടുകളിലേക്ക് എല്ലാ റൈറ്റുകളും ഡിസ്പാച്ച് ചെയ്യുന്നു. മറുവശത്ത്, കോംബൈൻ കേർണലുകൾ മെമ്മറിയിൽ നിന്ന് മാത്രം റീഡ് ചെയ്യുന്നു, ആവശ്യമായ എല്ലാ ഡാറ്റയും കൈമാറുന്നത് വരെ കാത്തിരിക്കുന്നു. NIC-കളിലേക്ക് ജോലി ഡിസ്പാച്ച് ചെയ്ത ശേഷം, വയറിലൂടെ ഡാറ്റ അസിൻക്രണസ് ആയി കൈമാറുന്ന സമയത്ത്, GPU-കൾക്ക് ലോക്കലായി മറ്റ് ഉപയോഗപ്രദമായ ജോലികൾ ചെയ്യാൻ കഴിയും, അതായത് ഷെയേർഡ് എക്സ്പെർട്ടുകൾ പ്രയോഗിക്കുകയോ അറ്റൻഷൻ കണക്കാക്കുകയോ ചെയ്യുക.

ഓരോ kernel ജോഡിക്കും സ്വന്തം സമമിത മെമ്മറി ബഫറുകൾ ഉണ്ടാകും; ഓരോ rank-ും മറ്റ് എല്ലാ rank-കളിൽ നിന്നുമുള്ള ഡാറ്റ സ്വീകരിക്കാൻ സ്വകാര്യ സംഭരണം നിയോഗിക്കുന്നു. DP ഗ്രൂപ്പിൽ നിന്ന് dispatch ചെയ്യാൻ കഴിയുന്ന ടോക്കണുകളുടെ പരമാവധി എണ്ണം rank-കളിലാകെ സ്ഥിരമാണ്, ഇത് ഓരോ local expert-നുമായി ബന്ധപ്പെട്ട് ഓരോ DP rank-ൽ നിന്നും ഓരോ rank-ും സ്വീകരിക്കാവുന്ന പരമാവധി പരിധിയെയും നിർണ്ണയിക്കുന്നു. ഇതുവഴി sender rank-കൾക്ക് പരസ്പരം ഏതെങ്കിലും synchronization ആവശ്യമില്ലാതെ, ലക്ഷ്യ rank-ിൽ എഴുതാൻ ഒരു സവിശേഷ വിലാസം കണ്ടെത്താൻ കഴിയും. ഡാറ്റ ലഭിച്ചതിനുശേഷം, സ്വീകരിക്കുന്ന kernel-ുകൾ token-കളെ memory-യിൽ optimal format-ിൽ ക്രമപ്പെടുത്തുന്നതിനായി അവയെ shuffle ചെയ്യും. ബഫറുകൾക്ക് വലിയ dimension ഉണ്ടെങ്കിലും, അവ model-ന്റെ എല്ലാ sequential layer-ുകളിലും വീണ്ടും ഉപയോഗിക്കപ്പെടുന്നു.

വലുതാക്കുക

മുകളിലുള്ള ചിത്രത്തിൽ ചിത്രീകരിച്ചിരിക്കുന്നതുപോലെ, ആഗോള സിൻക്രണൈസേഷന്റെ ഒരേയൊരു രൂപം, കോംബൈൻ-റിസീവ് കേർണലിലാണ് നടപ്പിലാക്കുന്നത്. സെൻഡിൽ നിന്ന് ഡാറ്റ ഡിസ്പാച്ച് ചെയ്തുകഴിഞ്ഞാൽ, മറ്റ് റാങ്കുകളിൽ നിന്ന് ആവശ്യമായ എല്ലാ ടോക്കണുകളും ലഭിക്കുമ്പോൾ ഒരു റാങ്ക് മടങ്ങി വന്ന് എക്സിക്യൂഷൻ തുടരുന്നു, കൂടാതെ പ്രോസസ്സ് ചെയ്ത ഉടൻ തന്നെ കോംബൈൻ സെൻഡ് കേർണലുകളിൽ പ്രവേശിക്കുകയും ചെയ്യുന്നു. കോംബൈൻ-റിസീവറിലെ ബാരിയർ, മറ്റ് ഏതെങ്കിലും റാങ്ക് ഇപ്പോഴും ഡാറ്റ സ്വീകരിക്കാൻ കാത്തിരിക്കുമ്പോൾ ഒരു റാങ്കിനും മുന്നോട്ട് പോയി ഡിസ്പാച്ച്-സെൻഡ് ആരംഭിക്കാൻ കഴിയില്ലെന്ന് ഉറപ്പാക്കുന്നു, ഇത് സാധ്യതയുള്ള വിനാശകരമായ ഓവർലാപ്പിന് കാരണമാകും. ലാളിത്യത്തിനായി കോംബൈൻ കേർണലിലാണ് സിൻക്രണൈസേഷൻ ചെയ്യുന്നത്, ഒരു ഒപ്റ്റിമൈസേഷൻ എന്ന നിലയിൽ ഡിസ്പാച്ച് റിസീവിനും കോംബൈൻ സെൻഡ് കേർണലുകൾക്കുമിടയിൽ സമാനമായ ബാരിയർ നടപ്പിലാക്കാം, മറ്റ് കേർണലുകൾ പ്രവർത്തിക്കുമ്പോൾ കൈമാറ്റങ്ങൾ തീർപ്പാക്കാൻ NIC-കൾക്ക് കൂടുതൽ സമയം അനുവദിച്ചുകൊണ്ട് ലേറ്റൻസി കൂടുതൽ കുറയ്ക്കാം.

രണ്ട് കേർണലുകളും ഉപകരണങ്ങളുടെ ലഭ്യമായ എല്ലാ SM-കളിലും വിഭജിക്കപ്പെട്ടിരിക്കുന്നു: ഡിസ്പാച്ച് സെൻഡ്, കോംബൈൻ റിസീവ് കേർണലുകൾ പെർ-റാങ്ക് പരമാവധി ടോക്കൺ കൗണ്ടിലുടനീളം (max_m) സമാന്തരമാക്കണം, അതേസമയം ഡിസ്പാച്ച് റിസീവ്, കോംബൈൻ സെൻഡ് കേർണലുകൾ ഓരോ റാങ്കിൽ നിന്നും ലോക്കൽ എക്സ്പെർട്ടിൽ നിന്നും അത്രയും ടോക്കണുകൾ കൈകാര്യം ചെയ്യണം, പരമാവധി max_m * (num_experts // EP) * (EP // TP).

ഡിസ്പാച്ച്

dispatch kernel-ന്റെ ഉത്തരവാദിത്വം tokens-നെയും block scaling factors-നെയും, tokens route ചെയ്തിരിക്കുന്ന experts-നെ host ചെയ്യുന്ന peers-ലേക്ക് dispatch ചെയ്യുന്നതാണ്. sender side, രണ്ട് ജോലികൾ parallelize ചെയ്യാൻ warp specialization-നെ ആശ്രയിക്കുന്നു: ഓരോ rank-ിൽ നിന്നുമുള്ള എത്ര token-കൾ അയക്കപ്പെടുന്നു എന്ന് നിർണ്ണയിക്കാൻ routing വിവരങ്ങൾ കൂട്ടിച്ചേർക്കുക, കൂടാതെ remote-ലേക്ക് dispatch ചെയ്യപ്പെടേണ്ട സന്ദേശങ്ങളിലേക്കായി token-കൾ പാക്ക് ചെയ്യുക. receiver side ആദ്യം എല്ലാ token count-കളും ലഭിക്കുന്നതു വരെ കാത്തിരിക്കുന്നു, തുടർന്ന് shared memory-യിൽ നിന്ന് payload-കൾ വായിച്ച് unpack ചെയ്യുകയും അവ downstream Grouped GEMM kernels-ലേക്ക് കൈമാറപ്പെടുന്ന tensors-ലേക്ക് ഇടുകയും ചെയ്യുന്നു. ഓരോ token എവിടെ തിരികെ അയക്കണം എന്ന് സൂചിപ്പിക്കാൻ combine kernels-ുമായി പങ്കിട്ടിരിക്കുന്ന buffers-ലും അത് വിവരങ്ങൾ സംഭരിക്കുന്നു. ഈ വിവരങ്ങൾ ആവശ്യമാണ്, കാരണം receive kernel token-കളെ non-deterministic ക്രമത്തിൽ contiguous buffers-ൽ shuffle ചെയ്യുന്നു.

വലുതാക്കുക

സെൻഡർ ഭാഗത്ത്, ഓരോ ബ്ലോക്കിലെയും 32 ത്രെഡുകളുടെ ഒരൊറ്റ വാർപ്പ് റൂട്ടിംഗ് വിവരങ്ങൾ റീഡ് ചെയ്യുന്നതിനും ഈ റാങ്കിൽ നിന്ന് ഓരോ ഡെസ്റ്റിനേഷൻ എക്സ്പെർട്ടും സ്വീകരിക്കുമെന്ന് പ്രതീക്ഷിക്കുന്ന ടോക്കണുകളുടെ എണ്ണം കണക്കാക്കുന്നതിനും ഉത്തരവാദിയാണ്, nvshmem_signal_op ഉപയോഗിച്ച് കൗണ്ട് പ്ലസ് വൺ അയക്കുന്നു. റിമോട്ട് അറ്റത്ത് പൂജ്യത്തിൽ നിന്ന് പൂജ്യമാല്ലാത്തതിലേക്കുള്ള മാറ്റം ഒരു സോഴ്സ് റാങ്കിൽ നിന്നുള്ള കൗണ്ടുകൾ ലഭിച്ചതിനെ സൂചിപ്പിക്കുന്നതിനാൽ, കൗണ്ട് ഒന്നുമായി വർദ്ധിപ്പിക്കുന്നു. സമാന്തരമായി, ശേഷിക്കുന്ന വാർപ്പുകൾ എല്ലാ ബ്ലോക്കുകളിലും ഒരേസമയം സിമ്മട്രിക് മെമ്മറിയിലേക്ക് ടോക്കണുകൾ പകർത്താൻ സഹകരിക്കുന്നു, ആക്റ്റിവേഷനുകൾ, സ്കെയിലിംഗ് ഘടകങ്ങൾ, ലോക്കൽ റാങ്കിലെ അവയുടെ സൂചിക എന്നിവ ഒരു തുടർച്ചയായ കഷണമായി പായ്ക്ക് ചെയ്യുന്നു. ടോക്കൺ റൈറ്റ് ചെയ്യപ്പെടും വിലാസം നിർണ്ണയിക്കാൻ കോംബൈൻ സെൻഡർക്ക് സൂചിക ആവശ്യമാണ്. അടുത്തതായി, വാർപ്പ് ഗ്രൂപ്പുകളിലുടനീളമുള്ള ഒരു ബാരിയർ വഴി എല്ലാ ഡാറ്റയും പകർത്തപ്പെട്ടിട്ടുണ്ടെന്ന് ഉറപ്പാക്കിയ ശേഷം, വാർപ്പുകൾ വീണ്ടും സ്വതന്ത്രമായി പ്രവർത്തിക്കുന്നു, ഓരോന്നും ഒരേ ബഫർ വ്യത്യസ്ത എക്സ്പെർട്ടിലേക്ക് സമാന്തരമായി അയക്കുന്നു. nvshmemx_putmem_signal_nbi_warp ഉപയോഗിച്ചാണ് ടോക്കണുകൾ അയക്കുന്നത്, ഇത് റിമോട്ട് ഡിവൈസിലെ ലോക്കൽ റാങ്കിൽ നിന്നുള്ള അയച്ച ടോക്കണുകളുടെ എണ്ണം ആറ്റോമിക്കായി വർദ്ധിപ്പിക്കുകയും ചെയ്യുന്നു. ഒരു DP ഗ്രൂപ്പിനുള്ളിൽ, അയക്കേണ്ട ടോക്കണിന്റെ ഒരു റെപ്ലിക്ക ഓരോ റാങ്കിനും സ്വന്തമായുള്ളതിനാൽ, ഡിസ്പാച്ച് തുല്യമായി ബാലൻസ് ചെയ്യപ്പെടുന്നു, ഓരോ ഉപകരണവും ടോക്കണുകളുടെ ഒരു ഉപസെറ്റ് പുറത്തേക്ക് അയക്കുന്നു.

റിസീവ് വശത്ത്, മൊത്തം ടോക്കൺ കൗണ്ടുകൾ അയക്കുന്നതിനായി കാത്തിരുന്നുകൊണ്ട് എല്ലാ കേർണലുകളും ആദ്യം സെൻഡർ റാങ്കുകളുമായി സിൻക്രണൈസ് ചെയ്യുന്നു. അതിനുശേഷം, സോഴ്സ് റാങ്കുകളിൽ നിന്നുള്ള എല്ലാ പേലോഡുകളും അയച്ചുകഴിഞ്ഞിരിക്കുന്നു എന്ന് സൂചിപ്പിച്ചുകൊണ്ട്, ആറ്റോമിക്കായി വർദ്ധിപ്പിച്ച സെൻഡ് ടോക്കൺ കൗണ്ടുകൾ മൊത്തം കൗണ്ടുകളിലേക്ക് സ്ഥിരത പ്രാപിക്കുന്നതിനായി അവയെല്ലാം കാത്തിരിക്കുന്നു, ഇത് putmem കോളുകളുടെ സെമാന്റിക്സ് കാരണം സാധ്യമാകുന്നു. എല്ലാ ബ്ലോക്കുകളിലും ത്രെഡുകളിലും പ്രവർത്തനം സമാന്തരമാക്കിക്കൊണ്ട് കേർണലുകൾ nvshmem_uint64_wait_until ഉപയോഗിച്ച് കൗണ്ടുകളിൽ പോൾ ചെയ്യുന്നു. തുടർന്ന്, എല്ലാ ഡാറ്റയും ശരിയായി ലഭിച്ചില്ലെങ്കിൽ ഒരു ബ്ലോക്കും ബഫറുകൾ റീഡ് ചെയ്യുന്നില്ലെന്ന് ക്രോസ്-ബ്ലോക്ക് ബാരിയർ ഉറപ്പാക്കുന്നു. ബ്ലോക്കുകളിലുടനീളം വ്യാപിപ്പിച്ച് ആറ്റോമിക്കായി വർദ്ധിപ്പിച്ച കൗണ്ടർ വഴി സിൻക്രണൈസ് ചെയ്ത, ടോക്കണുകൾ സിമ്മട്രിക് മെമ്മറിയിൽ നിന്ന് റെഗുലർ തുടർച്ചയായ ബഫറുകളിലേക്ക് പകർത്തപ്പെടുന്നു, അവ എക്സ്പെർട്ടുകളെ നടപ്പിലാക്കുന്ന ഡൗൺസ്ട്രീം കേർണലുകളിലേക്കുള്ള ഇൻപുട്ടുകളായി മാറുന്നു. സോഴ്സ് റാങ്ക്, എക്സ്പെർട്ട് സൂചിക, ടോക്കൺ സൂചിക എന്നിവ പ്രത്യേകം സംഭരിച്ചിരിക്കുന്നു, കോംബൈൻ കേർണൽ ആക്റ്റിവേഷനുകൾ അയക്കേണ്ട സ്ഥലം ഇത് കൃത്യമായി ചൂണ്ടിക്കാണിക്കുന്നു. ഒരു DP ഗ്രൂപ്പിനുള്ളിൽ നിന്നുള്ള ടോക്കണുകൾ വ്യത്യസ്ത ഉപകരണങ്ങളിൽ നിന്നാണ് അയക്കുന്നത് എങ്കിലും, ബന്ധപ്പെട്ട എക്സ്പെർട്ടിന് കൈമാറുന്നതിനായി അവയെല്ലാം ഒന്നിച്ചു ഗ്രൂപ്പ് ചെയ്യപ്പെടുന്നു.

കോംബൈൻ

കോംബൈൻ കേർണലുകൾ വീണ്ടും സെൻഡ്, റിസീവർ പാതികളായി വിഭജിക്കപ്പെട്ടിരിക്കുന്നു: സെൻഡറുകൾ ക്വാണ്ടിഡൈസ് ചെയ്യാത്ത 16-ബിറ്റ് ആക്റ്റിവേഷനുകൾ റിമോട്ട് ഉപകരണങ്ങളിലേക്ക് പകർത്തുന്നു, റിസീവ് കേർണലുകൾ ഡാറ്റ അയക്കുന്നതിനായി കാത്തിരിക്കുന്നു കൂടാതെ എല്ലാ എക്സ്പെർട്ട് സംഭാവനകളുടെയും വെയ്റ്റഡ് ആവറേജ് ലോക്കലായി കണക്കാക്കുന്നു. കൂടാതെ, ഡിസ്പാച്ച്-കോംബൈൻ സീക്വൻസ് സിൻക്രണൈസ് ചെയ്യാൻ അവ ഒരു ബാരിയറായും പ്രവർത്തിക്കുന്നു: സെൻഡ് കേർണലുകളിൽ പ്രവേശിക്കുമ്പോൾ ഓരോ റാങ്കും ഓരോ പിയറിലും ഒരു ഫ്ലാഗ് സജ്ജമാക്കുന്നു, ഫ്ലാഗ് സജ്ജീകരിച്ചിരിക്കുന്നത് നിരീക്ഷിച്ചില്ലെങ്കിൽ റിസീവ് കേർണലുകൾക്ക് മടങ്ങാൻ അനുവാദമില്ല. യഥാർത്ഥ ആശയവിനിമയവുമായും കമ്പ്യൂട്ടേഷനുമായും ഓവർലാപ്പ് ചെയ്യുന്നതിനാൽ സിൻക്രണൈസേഷന്റെ ലേറ്റൻസി കുറവാണ്.

വലുതാക്കുക

സെൻഡർ വശത്ത്, കേർണലുകൾ എല്ലാ ലോക്കൽ എക്സ്പെർട്ടുകൾക്കും നൽകിയിട്ടുള്ള ടോക്കണുകളുടെ ലിസ്റ്റ് സമാന്തരമായി ട്രാവേഴ്സ് ചെയ്യുന്നു, ഡെസ്റ്റിനേഷൻ റാങ്കിലെ ഒരു ബഫറിലേക്ക് അവയെ റൈറ്റ് ചെയ്യുന്നു. ടോക്കണുകൾ ലഭിച്ചപ്പോൾ സ്കാറ്റർ കേർണലുകൾ പോപ്പുലേറ്റ് ചെയ്ത പെർ-ടോക്കൺ ബഫറുകളിൽ നിന്നാണ് ടാർഗെറ്റ് റാങ്ക്, എക്സ്പെർട്ട് സൂചിക, ടോക്കൺ സൂചിക എന്നിവ റീഡ് ചെയ്യുന്നത്. മുകളിലുള്ള ചിത്രത്തിൽ സൂചിപ്പിച്ചിരിക്കുന്നതുപോലെ, ഓരോ സെൻഡർക്കും റൈറ്റ് ചെയ്യാൻ ഓരോ എക്സ്പെർട്ടിനും സ്വന്തമായി സ്വകാര്യ മെമ്മറി മേഖലയുണ്ട്, ഇത് സിൻക്രണൈസ് ചെയ്യേണ്ട ആവശ്യം ഒഴിവാക്കുന്നു. ഡിസ്പാച്ചിന് സമാനമായി, ഡാറ്റ ലഭിച്ചുവെന്ന് സൂചിപ്പിക്കാൻ കോംബൈൻ ഡെസ്റ്റിനേഷൻ റാങ്കിലെ പെർ ടോക്കൺ കൗണ്ടറുകൾ ആറ്റോമിക്കായി വർദ്ധിപ്പിക്കുന്നു: കൗണ്ടർ ഒരു ടോക്കൺ ഡിസ്പാച്ച് ചെയ്ത എക്സ്പെർട്ടുകളുടെ എണ്ണവുമായി പൊരുത്തപ്പെടുമ്പോൾ, ടോക്കൺ ഉള്ളടക്കങ്ങൾ ആക്സസ് ചെയ്യാൻ കഴിയും.

റിസീവ് കേർണലിൽ, സിഗ്നലിംഗ് പ്രവർത്തനം സജ്ജമാക്കിയ ഫ്ലാഗ് പോൾ ചെയ്തുകൊണ്ട് അവയുടെ ഉള്ളടവുകൾ വരുന്നതിനായി കാത്തിരുന്നുകൊണ്ട്, ഒന്നിലധികം ബ്ലോക്കുകളിലുടനീളം ടോക്കണുകളുടെ ലിസ്റ്റ് സമാന്തരമായി ട്രാവേഴ്സ് ചെയ്യപ്പെടുന്നു. എത്തുമ്പോൾ, റൂട്ടിംഗ് ടേബിൾ ഏത് ബഫറിൽ നിന്ന് റീഡ് ചെയ്യണമെന്നും ഓരോ എക്സ്പെർട്ടിനും എന്ത് ഭാരം നൽകണമെന്നും സൂചിപ്പിച്ചുകൊണ്ട്, സ്വകാര്യ ബഫറുകളിൽ നിന്ന് പേലോഡുകൾ റീഡ് ചെയ്യപ്പെടുന്നു. എല്ലാ ഉപകരണങ്ങളും ബാരിയർ പാസാക്കിക്കഴിഞ്ഞാൽ കേർണൽ എക്സിക്യൂഷൻ പൂർത്തിയാക്കിക്കൊണ്ട് ഫലങ്ങൾ ബാഹ്യമായി അലോക്കേറ്റ് ചെയ്ത ടെൻസറുകളിലേക്ക് റൈറ്റ് ചെയ്യപ്പെടുന്നു. ഒരു DP ഗ്രൂപ്പിലുടനീളം, എല്ലാ റാങ്കുകളും അവയുടെ സ്വന്തം റെപ്ലിക്കകൾ കണക്കാക്കാൻ ഓരോ എക്സ്പെർട്ട് ആക്റ്റിവേഷന്റെയും ഒരു കോപ്പി സ്വീകരിക്കുന്നു.

ബെഞ്ച്മാർക്കുകൾ

InfiniBand, CX-7 NIC-കൾ വഴി കണക്റ്റുചെയ്തിരിക്കുന്ന 16 നോഡുകളിലായി വ്യാപിച്ചുകിടക്കുന്ന 128 GPU-കളുടെ ക്ലസ്റ്ററിലാണ് ഞങ്ങൾ ഞങ്ങളുടെ കേർണലുകൾ വിലയിരുത്തുന്നത്. ഞങ്ങൾ GPUDirect Async (IBGDA), ഒരു CPU പ്രോക്സിയോടുകൂടിയ റിലയബിൾ കണക്ഷൻ (RC) എന്നിവ പരീക്ഷിക്കുന്നു. ഒപ്റ്റിമൈസ് ചെയ്ത ഡീപ്സീക്ക് നടപ്പാക്കലുമായും, PyTorch (NCCL വഴി) അല്ലെങ്കിൽ NVSHMEM നൽകുന്ന ഡെൻസ് പ്രിമിറ്റീവുകളുമായും ഞങ്ങൾ അവയെ താരതമ്യം ചെയ്യുന്നു.

വലുതാക്കുക

ഡെൻസ് NVSHMEM ഓൾ-ടു-ഓൾ പ്രവർത്തനത്തിൽ, IBRC, IBGDA എന്നിവയുടെ പ്രകടനം സമാനമാണെങ്കിലും (6378 µs vs 6180 µs), സ്പാഴ്സ് കേർണലുകളിൽ IBGDA വളരെ വേഗതയുള്ളതാണ്. ഡിസ്പാച്ചും കോംബൈനും കൂട്ടിച്ചേർക്കുമ്പോൾ, IBGDA 902 µs ഉപയോഗിക്കുന്നുവെങ്കിൽ IBRC-ക്ക് 3223 µs എടുക്കുന്നു - ലേറ്റൻസിയിൽ 3.6x മെച്ചപ്പെടുത്തൽ. ഓൾ-ടു-ഓൾ ബാൻഡ്‌വിഡ്ത്ത് ബൗണ്ട് ആയിരിക്കുമ്പോൾ, സ്പാഴ്സ് കേർണലുകൾ വളരെ കുറഞ്ഞ അളവിലുള്ള ഡാറ്റ ബ്രോഡ്കാസ്റ്റ് ചെയ്യുന്നു, ലേറ്റൻസി വഴി ബന്ധിക്കപ്പെട്ടിരിക്കുന്നു. GPU-വിനെ കോർഡിനേറ്റ് ചെയ്യാൻ ഒരു CPU പ്രോക്സിയുടെ ആവശ്യമില്ലാതെ, നേരിട്ട് ഒരു GPU-ൽ നിന്ന് നെറ്റ്‌വർക്ക് കൈമാറ്റങ്ങൾ ട്രിഗർ ചെയ്യുന്നതിലൂടെ, എൻഡ്-ടു-എൻഡ് ലേറ്റൻസി ഗണ്യമായി കുറയുന്നു.

ഞങ്ങളുടെ പോർട്ടബിൾ കേർണലുകൾ ഉയർന്ന ഒപ്റ്റിമൈസ് ചെയ്ത ഡീപ്സീക്ക് കേർണലുകളേക്കാൾ ഏകദേശം 2x മന്ദഗതിയിലാണെങ്കിലും, ഡെൻസ് കേർണലുകളുമായി താരതമ്യം ചെയ്യുമ്പോൾ അവ ലേറ്റൻസി 10x മെച്ചപ്പെടുത്തുന്നു (902 µs vs 9944 µs).

കൂടാതെ, ലൈബ്രറി പ്രിമിറ്റീവുകളിൽ നിന്ന് വ്യത്യസ്തമായി, സെൻഡർ-റിസീവർ ഘടകങ്ങളിലേക്കുള്ള വിഭജനം ലേറ്റൻസിയിൽ കുറച്ചു ഭാഗം മറച്ചുവെക്കാൻ അനുവദിക്കുന്നു.

വലുതാക്കുക

സിംഗിൾ-നോഡിൽ (EP8), NVSHMEM ട്രാൻസ്പോർട്ടേഷനായി NVLINK ഉപയോഗിക്കുന്നു, ഇത് നോഡ്-ഇന്റർ നെറ്റ്‌വർക്കിംഗിനേക്കാൾ കുറഞ്ഞ ലേറ്റൻസിയും ഉയർന്ന ത്രൂപുട്ടും നൽകുന്നു. ഞങ്ങളുടെ പോർട്ടബിൾ കേർണലുകൾ സിംഗിൾ-നോഡിൽ DeepEP-യേക്കാൾ ഏകദേശം 2.5x വേഗതയുള്ളതാണ് (186µs vs 481 µs).

കൂടുതൽ മെച്ചപ്പെടുത്തലുകൾ

ഇവിടെ വിവരിച്ചിരിക്കുന്ന കേർണലുകൾ ML ഫ്രെയിംവർക്കുകളുടെ ബിൽറ്റ്-ഇൻ പ്രിമിറ്റീവുകളെ മറികടക്കുകയും NVLink അല്ലെങ്കിൽ InfiniBand പോലെയുള്ള പ്രത്യേക ഡിവൈസ്-ഇന്റർ ട്രാൻസ്പോർട്ടുകൾക്കായി അമിതമായി പ്രത്യേകവൽക്കരിക്കാതെ മികച്ച പ്രകടനം വാഗ്ദാനം ചെയ്യുകയും ചെയ്യുന്നു. ഇതിനകം സൂചിപ്പിച്ച അവസരങ്ങൾക്ക് പുറമേ, ആശയവിനിമയ പ്രിമിറ്റീവുകൾ കൂടുതൽ പ്രത്യേക പതിപ്പുകൾ ഉപയോഗിച്ച് മാറ്റിസ്ഥാപിക്കുന്നതിലൂടെ കൂടുതൽ പ്രകടന നേട്ടങ്ങൾ കൈവരിക്കാനാകും. ഉദാഹരണത്തിന്, NVLink-ൽ, സിമ്മട്രിക് മെമ്മറിയുടെ ഉപയോഗം ഉപകരണങ്ങളിൽ ഉടനീളം പങ്കിട്ട ബഫറുകൾ ഉപയോഗിച്ച് മാറ്റിസ്ഥാപിക്കാം, ഇത് ചില പകർപ്പുകൾ ഇല്ലാതാക്കുകയും ടോക്കൺ ബാച്ചുകൾക്ക് പകരം വ്യക്തിഗത ടോക്കണുകളിലുടനീളം, സൂക്ഷ്മമായ സിൻക്രണൈസേഷൻ സ്കീമുകൾ അനുവദിക്കുകയും ചെയ്യുന്നു. InfiniBand-ൽ, ഒരു നടപ്പാക്കലിന് അടിസ്ഥാന ക്യൂ ജോഡികൾ നേരിട്ട് ആക്സസ് ചെയ്യാൻ കഴിയും, ഇത് പ്രതീക്ഷിക്കുന്നതും അയച്ചതുമായ ടോക്കൺ കൗണ്ടുകൾ കൈമാറേണ്ടതിന്റെ ആവശ്യകത ഇല്ലാതാക്കുകയും ടോക്കണുകൾ എത്തുമ്പോൾ അവയെ ക്യൂകളിൽ നിന്ന് ഔട്ട്പുട്ട് ബഫറുകളിലേക്ക് ഫ്ലഷ് ചെയ്യുകയും ചെയ്യുന്നു. എന്നിരുന്നാലും, അത്തരം നടപ്പാക്കലുകൾ പോർട്ടബിലിറ്റിയുടെ ചെലവിലാണ് വരുന്നത്, അതേസമയം നിർദ്ദിഷ്ട ഹാർഡ്‌വെയറിനായി ഒപ്റ്റിമൈസ് ചെയ്യാൻ ഡെവലപ്മെന്റ് ശ്രമങ്ങൾ നടത്തുന്നതിന് മുമ്പ് നിർദ്ദിഷ്ട കോൺഫിഗറേഷനുകളുള്ള ക്ലസ്റ്ററുകൾ വിലയിരുത്താൻ ഞങ്ങളുടെ നടപ്പാക്കൽ അനുവദിക്കുന്നു.

ഉപസംഹാരം

വൈവിധ്യമാർന്ന ഹാർഡ്‌വെയർ കോൺഫിഗറേഷനുകളിലുടനീളം അനുയോജ്യത നിലനിർത്തിക്കൊണ്ട്, സാധാരണ ഓൾ-ടു-ഓൾ ആശയവിനിമയത്തേക്കാൾ 10 മടങ്ങ് വേഗതയേറിയ പ്രകടനം കൈവരിക്കുന്ന MoE ആശയവിനിമയത്തിനായി ഉയർന്ന പ്രകടനമുള്ള, പോർട്ടബിൾ ലൈബ്രറി ഞങ്ങൾ അവതരിപ്പിച്ചു. NVLink ഉള്ള സിംഗിൾ-നോഡ് ഡിപ്ലോയ്മെന്റുകളിൽ, ഞങ്ങളുടെ പരിഹാരം മുൻകാല നടപ്പാക്കലുകളേക്കാൾ 2.5x കുറഞ്ഞ ലേറ്റൻസി പ്രകടമാക്കുന്നു.

GPU-ആരംഭിച്ച ആശയവിനിമയ പിന്തുണ, കമ്പ്യൂട്ടേഷൻ-കമ്മ്യൂണിക്കേഷൻ ഓവർലാപ്പ് പ്രാപ്തമാക്കുന്ന സ്പ്ലിറ്റ് കേർണൽ ആർക്കിടെക്ചർ, കുറഞ്ഞ NVSHMEM പ്രിമിറ്റീവുകൾ ഉപയോഗിച്ചുള്ള കാര്യക്ഷമമായ ടോക്കൺ ഡിസ്പാച്ച് എന്നിവയുൾപ്പെടെയുള്ള പ്രധാന നവീകരണങ്ങളിലൂടെ ഞങ്ങളുടെ സമീപനം പ്രകടനവും പോർട്ടബിലിറ്റിയും ബാലൻസ് ചെയ്യുന്നു. മൾട്ടി-നോഡ് സജ്ജീകരണങ്ങളിൽ ഉയർന്ന പ്രത്യേക നടപ്പാക്കലുകളേക്കാൾ ഏകദേശം 2x മന്ദഗതിയിലാണെങ്കിലും, വിവിധ നെറ്റ്‌വർക്ക് പരിതസ്ഥിതികളിൽ (NVLink, CX-7, EFA) ഞങ്ങളുടെ ലൈബ്രറി മികച്ച ഫ്ലെക്സിബിലിറ്റി വാഗ്ദാനം ചെയ്യുന്നു.

MoE മോഡലുകൾ സ്കെയിൽ ചെയ്യുന്നത് തുടരുമ്പോൾ, ഞങ്ങളുടേത് പോലെയുള്ള കാര്യക്ഷമമായ ആശയവിനിമയ തന്ത്രങ്ങൾ പ്രായോഗിക ഡിപ്ലോയ്മെന്റിന് കൂടുതൽ പ്രാധാർഹമുള്ളതായിത്തീരും. ഞങ്ങളുടെ പൂർണ്ണമായും ഓപ്പൺ സോഴ്സ് നടപ്പാക്കൽ https://github.com/ppl-ai/pplx-kernels എന്നതിൽ ലഭ്യമാണ്.

റഫറൻസുകൾ

DeepEP

DeepSeek-V3 സാങ്കേതിക റിപ്പോർട്ട്

NVIDIA NVSHMEM

ഉദ്ധരണി

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}