2 సెకన్లలోపు RL పోస్ట్-ట్రైనింగ్ కోసం వెయిట్ బదిలీ

మేము ఇటీవల Kimi-K2 (1T పారామీటర్‌లు) కోసం 1.3-సెకన్ల క్రాస్-మెషిన్ పారామీటర్ అప్‌డేట్‌లను సాధించాము, 256 ట్రైనింగ్ GPUల (BF16) నుండి 128 ఇన్‌ఫరెన్స్ GPUలకు (FP8) వెయిట్‌లను బదిలీ చేశాము.

అసింక్రోనస్ రీన్‌ఫోర్స్‌మెంట్ లెర్నింగ్ ఫైన్-ట్యూనింగ్‌లో, ట్రైనింగ్ మరియు ఇన్‌ఫరెన్స్ వేర్వేరు GPUలపై రన్ అవుతాయి. ప్రతి ట్రైనింగ్ స్టెప్ తర్వాత, కొత్త వెయిట్‌లను ఇన్‌ఫరెన్స్ నోడ్స్‌కు పంపాలి. ట్రిలియన్ పారామీటర్ మోడల్స్ కోసం అనేక ఇప్పటికే ఉన్న ఫ్రేమ్‌వర్క్‌లు అనేక సెకన్లు—లేదా నిమిషాలు కూడా—తీసుకుంటాయి.

RDMA పాయింట్-టు-పాయింట్ కమ్యూనికేషన్‌ను ఉపయోగించడం ద్వారా, ఇన్‌ఫరెన్స్ ఇంజిన్‌ను మార్చకుండా, వెయిట్ బదిలీని చాలా వేగంగా చేయగలము మరియు కోడ్‌ను వ్రాయడం మరియు నిర్వహించడం సులభం చేయగలము.

మా పరిష్కారం 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.
    ...

గమ్యస్థాన వైపు బదిలీ కోసం నోటిఫికేషన్ కూడా పొందదు. ఇది ఇన్‌ఫరెన్స్ నోడ్స్‌పై ఎలాంటి కంట్రోల్ లాజిక్ లేకుండా ట్రైనింగ్ నోడ్స్ ద్వారా నడిచే తక్కువ-లేటెన్సీ, హై-త్రూపుట్, జీరో-కాపీ బదిలీలను అందిస్తుంది.

మెటాడేటా సేకరణ – కంట్రోలర్ అన్ని ట్రైనింగ్ మరియు ఇన్‌ఫరెన్స్ 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 మెమరీ వినియోగాన్ని పరిచయం చేస్తాయి. అవుట్ ఆఫ్ మెమరీ లోపాన్ని నివారించడానికి, ప్రస్తుత ఆన్-ది-ఫ్లై టాస్క్‌లు కాన్ఫిగర్ చేయదగిన వాటర్‌మార్క్ కంటే తక్కువ తాత్కాలిక GPU మెమరీని ఆక్రమిస్తేనే మేము టాస్క్ ఎగ్జిక్యూషన్‌ను ప్రారంభిస్తాము.

అనేక డిజైన్ ఎంపికలు మా సిస్టమ్‌ను సాధారణ ఓపెన్ సోర్స్ పరిష్కారాల కంటే అమలు చేయడానికి చాలా వేగంగా మరియు నిర్వహించడానికి సులభంగా చేస్తాయి.

Rank0-ఆధారిత వెయిట్ ట్రాన్స్‌ఫర్ వాటిని ఇన్‌ఫరెన్స్ నోడ్‌కు పంపిణీ చేసే ముందు వెయిట్‌లను సేకరిస్తుంది.

పాయింట్-టు-పాయింట్ కమ్యూనికేషన్, రెండు ట్రైనింగ్ నోడ్‌ల నుండి వెయిట్‌లను ఇన్‌ఫరెన్స్ నోడ్ యొక్క నెట్‌వర్క్ ఇంటర్‌ఫేస్‌లకు బదిలీ చేస్తుంది.

అన్ని పారామీటర్‌లను ర్యాంక్-0 GPUల ద్వారా ఫన్నెల్ చేయడం ఒక సాధారణ నమూనా: ట్రైనింగ్ ర్యాంక్-0 వద్ద సేకరించండి, ఇన్‌ఫరెన్స్ ర్యాంక్-0కి పంపండి, ఆపై మళ్ళీ స్కాటర్ చేయండి. ఇది త్వరగా ఒక చోక్‌పాయింట్‌గా మారుతుంది, ఒకే GPU యొక్క PCIe బ్యాండ్‌విడ్త్ మరియు NIC (ఉదా., 400 Gbps ≈ 50 GB/s) ద్వారా పరిమితం చేయబడుతుంది.

దీనికి విరుద్ధంగా, మా పాయింట్-to-పాయింట్ సెటప్ ప్రతి ట్రైనింగ్ GPU ప్రతి ఇన్‌ఫరెన్స్ GPUకి నేరుగా పంపడానికి అనుమతిస్తుంది, ఒకే లింక్ కాకుండా పూర్తి నెట్‌వర్క్ ఫ్యా్రిక్‌ను సంతృప్తపరుస్తుంది.

RPC-ఆధారిత ట్రాన్స్‌ఫర్ ఇన్‌ఫరెన్స్-వైపు బఫర్ మరియు వర్కర్ కాల్‌లను ఉపయోగిస్తుంది, అయితే వన్-సైడెడ్ ట్రాన్స్‌ఫర్ ఇన్‌ఫరెన్స్-వర్కర్ పాల్గొనకుండానే వెయిట్‌లను నేరుగా వ్రాస్తుంది.

కొన్ని సిస్టమ్‌లు ప్రతి టెన్సార్ కోసం ఇన్‌ఫరెన్స్ ఇంజిన్ యొక్క update_weight() పద్ధతికి కాల్ చేయడంపై ఆధారపడతాయి. అంటే ఇన్‌ఫరెన్స్ కోడ్‌లో అంతరాయం కలిగించే మార్పులు, మరియు RPCలు, సీరియలైజేషన్, మరియు కంట్రోల్-ప్లేన్ సమన్వయం నుండి ఓవర్‌హెడ్.

RDMA WRITE ప్రిమిటివ్‌తో, మేము అదనపు కాపీలు లేకుండా, ఇన్‌ఫరెన్స్ GPU మెమరీపై వెయిట్‌లను సైలెంట్‌గా అప్‌డేట్ చేస్తాము. కంట్రోల్ ప్లేన్ సందేశం మరియు CPU కంట్రోల్ లాజిక్ ఏదీ ప్రమేయం లేదు. ఇన్‌ఫరెన్స్ ఇంజిన్‌కు ఎలాంటి మార్పు అవసరం లేదు.

వెయిట్ బదిలీ ప్రక్రియ నాలుగు రకాల హార్డ్‌వేర్ వనరులను ఉపయోగించుకోవచ్చు: (1) హోస్ట్-డివైజ్ డేటా కదలిక (2) ప్రొజెక్షన్ ఫ్యూజన్ మరియు క్వాంటైజేషన్ కోసం GPU గణన (3) డేటా ప్లేన్ కోసం RDMA నెట్‌వర్క్ (4) కంట్రోల్ ప్లేన్ కోసం ఈథర్నెట్.

మా డిజైన్ వెయిట్ బదిలీ టాస్క్‌లను పైప్‌లైన్ దశలుగా విభజించింది, వివిధ హార్డ్‌వేర్ వనరుల మధ్య సులభంగా అతివ్యాప్తి చెందడానికి అనుమతిస్తుంది.

కొన్ని ఇంప్లిమెంటేషన్‌లు ప్రతి ట్రైనింగ్ స్టెప్‌లో బదిలీ షెడ్యూల్‌ను మళ్లీ లెక్కిస్తాయి, మెటాడేటాను పదేపదే సేకరిస్తాయి మరియు సూచనలను పంపిణీ చేస్తాయి. ఇది అనవసరమైన కంట్రోల్-ప్లేన్ లేటెన్సీని జోడిస్తుంది.

మా షెడ్యూల్ ప్రారంభ సమయంలో ఒకసారి లెక్కించబడుతుంది. ప్రతి ట్రైనింగ్ పునరావృతం కేవలం ప్రణాళికను రీప్లే చేస్తుంది: కంట్రోలర్ “గో” సిగ్నల్‌ను జారీ చేస్తుంది, మరియు 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}
}

RDMA WRITE: వన్-సైడెడ్ బదిలీలు

హై-లెవెల్ వర్క్‌ఫ్లో

వెయిట్ బదిలీ ఎగ్జిక్యూషన్

DeviceMesh మరియు మెష్ గ్రూపులు

టాస్క్ పైప్‌లైన్

GPU మెమరీ వినియోగ నియంత్రణ

ఇది ఎందుకు వేగంగా మరియు సులభంగా ఉంటుంది

పాయింట్-to-పాయింట్ కమ్యూనికేషన్

వన్-సైడెడ్ డేటా బదిలీ

పైప్‌లైనింగ్

స్టాటిక్ షెడ్యూల్

క్లీన్ విభజన

ముగింపు

ఉదాహరణ