SkewAdam Mixture-of-Experts శిక్షణ మెమరీని 60% కంటే ఎక్కువగా తగ్గిస్తుంది మరియు ఖచ్చితత్వంలో (accuracy) గణనీయమైన పెరుగుదలను అందిస్తుంది, దీనివల్ల డెవలపర్లు 6.78 బిలియన్ పారామీటర్ల MoE మోడల్‌ను ఒకే 40 GB GPUలో ఉపయోగించుకోవచ్చు.

MoE శిక్షణ ఎందుకు మెమరీ పరిమితులకు (memory wall) ఎదురవుతుంది

Mixture-of-Experts ఆర్కిటెక్చర్‌లు ఒక డెన్స్ బ్యాక్‌బోన్‌ను (dense backbone) ఉంచుతూనే, ప్రతి ఇన్‌పుట్‌ను "ఎక్స్‌పర్ట్" సబ్-నెట్‌వర్క్‌ల యొక్క చిన్న సబ్‌సెట్ ద్వారా పంపిస్తాయి. దీని వల్ల కంప్యూటేషన్ (compute) గణనీయంగా పెరగకుండానే, మోడల్‌ను బిలియన్ల కొద్దీ పారామీటర్లకు విస్తరించవచ్చు. వాస్తవానికి, ఆప్టిమైజర్—ముఖ్యంగా చాలా టీమ్‌లు ఉపయోగించే AdamW వెరియంట్—GPU RAMలో ఎక్కువ భాగాన్ని తీసుకుంటుంది. 6.78 B-పారామీటర్ల మోడల్‌కు, ఆప్టిమైజర్ యొక్క momentum మరియు variance టెన్సర్లు మాత్రమే సుమారు 50 GB అవసరమవుతాయి. దీనికి activations మరియు మోడల్ వెయిట్స్‌ను కలిపితే, పీక్ మెమరీ 81.4 GB వరకు చేరుతుంది, దీనివల్ల మల్టీ-GPU సెటప్‌లు లేదా నెమ్మదైన శిక్షణ షెడ్యూల్‌లు తప్పనిసరి అవుతాయి.

SkewAdam దేనిని భిన్నంగా చేస్తుంది

SkewAdam కొత్త లెర్నింగ్ రూల్‌ను కనిపెట్టలేదు. ఇది Adam యొక్క moments ఎక్కడ ఉండాలో రీషఫుల్ (reshuffle) చేస్తుంది:

  • డెన్స్ బ్యాక్‌బోన్ ఫుల్-ప్రిసిషన్ momentumను ఉంచుకుంటుంది, దీనివల్ల డెన్స్ లేయర్‌లకు అవసరమైన స్మూత్ అప్‌డేట్‌లు సాధ్యమవుతాయి.
  • ఎక్స్‌పర్ట్ బ్యాంక్ variance యొక్క ఫ్యాక్టరైజ్డ్ అప్రోక్సిమేషన్‌ను (factored approximation) నిల్వ చేస్తుంది, దీనివల్ల ప్రతి ఎక్స్‌పర్ట్ కోసం నిల్వ చేయాల్సిన డేటా తగ్గుతుంది.
  • ఏ ఎక్స్‌పర్ట్‌లను యాక్టివేట్ చేయాలో నిర్ణయించే రూటర్ (router), ఖచ్చితమైన సెకండ్-మోమెంట్ ఎస్టిమేట్‌ను (second-moment estimate) కలిగి ఉంటుంది.

ఈ మూడు భాగాలను విడివిడి "టైర్లు" (tiers) గా పరిగణించడం ద్వారా, ఆప్టిమైజర్ తక్కువ ప్రాముఖ్యత ఉన్న చోట అనవసరమైన ప్రిసిషన్‌ను తగ్గిస్తుంది మరియు అత్యంత ముఖ్యమైన చోట దానిని అలాగే ఉంచుతుంది.

కొలవదగిన ప్రభావం (Measurable impact)

అదే 6.78 B-పారామీటర్ MoE మోడల్‌ను SkewAdamతో నడపడం వల్ల లభించే ఫలితాలు:

  • పీక్ GPU మెమరీ: 31.3 GB (81.4 GB నుండి తగ్గింది)
  • ఆప్టిమైజర్-స్టేట్ ఫుట్‌ప్రింట్: 1.29 GB (50 GB నుండి తగ్గింది)

తగ్గించబడిన ఈ స్టేట్ (state) ఒకే 40 GB యాక్సిలరేటర్‌లో సులభంగా సరిపోతుంది.

కేవలం పొదుపు మాత్రమే కాదు, ఖచ్చితత్వంలో కూడా వృద్ధి

మెమరీ తగ్గింపులు తరచుగా మోడల్ నాణ్యతను దెబ్బతీస్తాయి, కానీ SkewAdam perplexity—ఒక ప్రామాణిక లాంగ్వేజ్-మోడల్ మెట్రిక్—ను 126.8 (AdamW) నుండి 108.4 కి మెరుగుపరుస్తుంది. ఇది అదే టాస్క్‌లో Muon (120.2) మరియు Lion (393.7) కంటే మెరుగైన పనితీరును కనబరుస్తుంది. మూడు టైర్ల అంతటా momentumను కాపాడటం వల్ల ఈ వృద్ధి లభిస్తుందని రచయితలు చెబుతున్నారు; Adafactor వంటి momentumను పూర్తిగా వదిలివేసే పద్ధతులు వెనుకబడి ఉంటాయి.

పెండింగ్ ప్రశ్నలు (Open questions)

SkewAdam ఫలితాలు 6.78 B-పారామీటర్ మోడల్‌పై ప్రదర్శించబడ్డాయి. అంతకంటే చాలా పెద్ద మోడల్‌లకు లేదా లాంగ్వేజ్ మోడలింగ్ వెలుపల ఉన్న టాస్క్‌లకు కూడా ఇదే మెమరీ-స్టేట్ నిష్పత్తులు వర్తిస్తాయా అనేది ఇంకా స్పష్టంగా లేదు.

గమనించవలసిన అంశాలు

  • పెద్ద MoE కాన్ఫిగరేషన్‌లపై (పదుల కొద్దీ బిలియన్ల పారామీటర్లు) బెంచ్‌మార్క్‌లు.
  • ఓపెన్-సోర్స్ ఫ్రేమ్‌వర్క్‌ల ద్వారా దీనిని స్వీకరించడం మరియు ప్రజాదరణ పొందిన శిక్షణ స్క్రిప్ట్‌లలో చేర్చడం.
  • కన్వర్జెన్స్ స్పీడ్ (convergence speed) లేదా వివిధ లెర్నింగ్-రేట్ షెడ్యూల్‌ల కింద స్థిరత్వం వంటి దాగి ఉన్న లాభనష్టాలపై (trade-offs) కమ్యూనిటీ ఫీడ్‌బ్యాక్.

మూలం: ఆప్టిమైజర్ యొక్క డిజైన్ మరియు అనుభవపూర్వక ఫలితాలను వివరించే డెవలపర్ పోస్ట్.