మీరు ఎప్పుడైనా ఒక LLM సుదీర్ఘమైన సమాధానాన్ని జనరేట్ చేయడం చూసి, ప్రారంభంలో వేగంగా ఉన్నా, ఆ తర్వాత ఎందుకు నెమ్మదించిపోతుందో అని ఆలోచించి ఉంటే, మీరు నిజానికి ఒక హార్డ్వేర్ అడ్డంకిని (hardware bottleneck) ప్రత్యక్షంగా చూస్తున్నారని అర్థం. చాలా మంది డెవలపర్లు తమ Python కోడ్, ఫ్రేమ్వర్క్ లేదా మోడల్ యొక్క పరిమాణం వల్ల ఇలా జరుగుతుందని భావిస్తారు. వారు ఫంక్షన్లను ప్రొఫైల్ చేస్తారు, ఆప్టిమైజర్లను మారుస్తారు మరియు ప్రీప్రాసెసింగ్లో మిల్లీసెకన్లను తగ్గించడానికి ప్రయత్నిస్తారు. కానీ వాటి వల్ల అసలు సమస్య పరిష్కారం కాదు. వేగ పరిమితి మీ సాఫ్ట్వేర్లో లేదు, అది సిలికాన్లో ఉంది.
ప్రతి లార్జ్ లాంగ్వేజ్ మోడల్ ఇన్ఫరెన్స్ (inference) పని మీ సర్వర్లో ఉన్న GPU యొక్క రెండు భౌతిక లక్షణాలపై ఆధారపడి ఉంటుంది: అది ఎంత వేగంగా గణనలు (numbers) చేయగలదు, మరియు ఆ గణనల కోసం ఆ నంబర్లను ఎంత వేగంగా సరైన స్థానానికి తరలించగలదు.
గణన సులభం. డేటాను తరలించడం కష్టం
GPU మార్కెటింగ్ కంప్యూట్ (compute) గురించి మాట్లాడటానికి ఇష్టపడుతుంది. సెకనుకు ట్రిలియన్ల కొద్దీ ఫ్లోటింగ్-పాయింట్ ఆపరేషన్లు. ఆ సంఖ్యలు అద్భుతంగా ఉంటాయి. కానీ కంప్యూట్ అనేది కథలో సగం మాత్రమే. మిగిలిన సగం మెమరీ బ్యాండ్విడ్త్ (memory bandwidth), అంటే డేటా హై-బ్యాండ్విడ్త్ మెమరీ నుండి అసలు గణిత ప్రక్రియలు జరిగే కంప్యూట్ కోర్లకు ఎంత వేగంగా ప్రయాణిస్తుంది అనే అంశం.
ఈ రెండు లింకులలో ఏది బలహీనంగా ఉంటే, ఒక LLM అంత వేగంతో మాత్రమే నడవగలదు. ఇరవై మంది మాస్టర్ చెఫ్లు ఉన్న ఒక వాణిజ్య వంటగదిని ఊహించుకోండి. ఓవెన్లు వేడిగా ఉన్నాయి, చాకులు పదునుగా ఉన్నాయి, ప్రతి వంట మనిషి సిద్ధంగా ఉన్నాడు. కానీ కూరగాయల డెలివరీ మాత్రం ఒక సైకిల్పై, ఒక్కో బుట్ట చొప్పున వస్తోంది. దీనివల్ల వంటగది పని ఆగిపోతుంది. మరిన్ని వంట మనుషులను చేర్చడం వల్ల ఇది పరిష్కారం కాదు. వేగవంతమైన ఓవెన్లను కొనడం వల్ల కూడా ఇది పరిష్కారం కాదు. ఇక్కడ అడ్డంకి ఆ రోడ్డు మాత్రమే.
ఆధునిక డేటాసెంటర్ GPUలలో, ఆర్థమెటిక్ యూనిట్లు ఎంత శక్తివంతమైనవంటే, అవి తరచుగా తమ గణనలను పూర్తి చేసి, మెమరీ నుండి వెయిట్స్ (weights) మరియు యాక్టివేషన్స్ (activations) స్ట్రీమ్ అయ్యే వరకు ఖాళీగా వేచి ఉంటాయి. ఈ అసమతుల్యత మీ కోడ్లో ఉన్న బగ్ కాదు. చిప్లు నిర్మించబడిన భౌతిక వాస్తవికత ఇది. మెమరీ బ్యాండ్విడ్త్, ముడి కంప్యూట్ వేగంతో సమానంగా పెరగలేదు. LLMలు ఈ అసమతుల్యత వల్ల మరింత ఇబ్బంది పడతాయి, ఎందుకంటే వాటి ఫార్వర్డ్ పాస్లు ప్రతి అవుట్పుట్ టోకెన్ కోసం ప్రతి ఒక్క పారామీటర్ను తాకాల్సి ఉంటుంది.
ప్రాంప్ట్లు వేగంగా, జనరేషన్ నెమ్మదిగా ఎందుకు అనిపిస్తాయి?
LLM ఇన్ఫరెన్స్ రెండు విభిన్న దశలుగా విభజించబడింది, మరియు అవి హార్డ్వేర్పై పూర్తిగా భిన్నమైన రీతిలో ఒత్తిడిని కలిగిస్తాయి.
Prefill అనేది మీ ప్రాంప్ట్ మోడల్కు చేరుకున్నప్పుడు జరుగుతుంది. అన్ని టోకెన్లు ఒకేసారి వస్తాయి. GPU వాటిని పెద్ద మ్యాట్రిక్స్-మ్యాట్రిక్స్ మల్టిప్లికేషన్స్ (matrix-matrix multiplications) ఉపయోగించి ప్యారలల్గా ప్రాసెస్ చేయగలదు. వేలకొద్దీ ఆర్థమెటిక్ యూనిట్లు ఒకేసారి పనిచేస్తాయి, మరియు వర్క్లోడ్ సాంద్రంగా ఉంటుంది. ఈ దశ కంప్యూట్-బౌండ్ (compute-bound). ప్రారంభంలో మీరు చూసే ఆ అకస్మాత్తుగా పెరిగే వేగం? అది GPU దేని కోసం నిర్మించబడిందో సరిగ్గా అదే పని చేస్తోంది.
Decode అనేది సమస్యలు మొదలయ్యే చోటు. మోడల్ తదుపరి టోకెన్ను జనరేట్ చేసేటప్పుడు, అది ఒక్కో టోకెన్ను మాత్రమే చేస్తుంది. ఈ దశ మ్యాట్రిక్స్-వెక్టర్ ఆపరేషన్లపై (matrix-vector operations) ఆధారపడి ఉంటుంది, ఇవి GPU యొక్క ప్యారలల్ సామర్థ్యంలో చాలా స్వల్ప భాగాన్ని మాత్రమే ఉపయోగిస్తాయి. అంతకంటే దారుణమైన విషయం ఏమిటంటే, ప్రతి కొత్త టోకెన్ GPUని మెమరీ నుండి మొత్తం మోడల్ వెయిట్స్ను మళ్ళీ లోడ్ చేయమని బలవంతం చేస్తుంది. ఆర్థమెటిక్ యూనిట్లు పని కోసం ఎదురుచూస్తాయి, కానీ వాటికి పని దొరకదు. కాబట్టి అవి వేచి ఉంటాయి. Decode అనేది మెమరీ-బౌండ్ (memory-bound). గణిత ఇంజన్లు విశ్రాంతి తీసుకుంటున్న సమయంలో, GPU కేవలం ఒక ఖరీదైన ట్రాఫిక్ కంట్రోలర్లా పనిచేస్తూ, పారామీటర్లను మెమరీ బస్ ద్వారా అటు ఇటు తరలిస్తూ ఉంటుంది. అందుకే ప్రారంభ ప్రాంప్ట్ విశ్లేషణ తక్షణమే జరిగినట్లు అనిపించినప్పటికీ, వంద పదాల సమాధానం రావడానికి పది సెకన్ల సమయం పట్టవచ్చు.
KV cache దీనిని మరింత ఆసక్తికరంగా మారుస్తుంది. Decode సమయంలో, మోడల్ ప్రతి మునుపటి టోకెన్ కోసం కీ (key) మరియు వ్యాల్యూ (value) టెన్సర్లను నిల్వ చేస్తుంది, తద్వారా అది అటెన్షన్ (attention) ప్రక్రియను మొదటి నుండి మళ్ళీ లెక్కించాల్సిన అవసరం ఉండదు. ఈ క్యాష్ సీక్వెన్స్ పొడవుతో పెరుగుతుంది. ఇది కూడా మెమరీలోనే ఉంటుంది. కాబట్టి ఇప్పుడు GPU కేవలం వెయిట్స్ను మాత్రమే లోడ్ చేయడం లేదు; ప్రతి ఫార్వర్డ్ పాస్లో నిరంతరం విస్తరిస్తున్న క్యాష్ను చదువుతూ మరియు రాస్తూ ఉంటుంది. మెమరీ బస్ రెండింటి కోసం శ్రమిస్తుంటే, కంప్యూట్ కోర్లు మాత్రం పెద్దగా కష్టపడవు.
మెమరీ వాల్ను అధిగమించడం
డేటా ఎంత తక్కువగా తరలించాలో తగ్గించడానికి లేదా కనీసం ఆ తరలింపు ఖర్చును పంచుకోవడానికి ఇంజనీర్లు కొన్ని పద్ధతులను అభివృద్ధి చేశారు.
Batching అనేది అత్యంత సరళమైన పద్ధతి. ఒక వినియోగదారుని రిక్వెస్ట్ వల్ల మెమరీ నుండి పూర్తి వెయిట్స్ను లోడ్ చేయాల్సి వస్తే, ఒకేసారి ఎనిమిది లేదా పదహారు రిక్వెస్ట్లను ప్రాసెస్ చేయడం వల్ల GPU ఆ లోడ్ను అన్నింటికీ పంచుకోగలదు (amortize). వెయిట్స్ ఒక్కసారి మాత్రమే చదవబడతాయి మరియు బ్యాచ్లోని ప్రతి సీక్వెన్స్ కోసం మళ్ళీ ఉపయోగించబడతాయి. ప్రొడక్షన్లో, అధునాతన షెడ్యూలింగ్ సిస్టమ్లు రిక్వెస్ట్లను డైనమిక్గా సమూహపరుస్తాయి, దీనిని కొన్నిసార్లు కంటిన్యూయస్ లేదా ఇన్-ఫ్లైట్ బ్యాచింగ్ (continuous or in-flight batching) అని పిలుస్తారు, తద్వారా GPU అరుదుగా ఆగిపోతుంది. ఇది ఒకే మార్గంలో వెళ్లే ఒక బస్సు మరియు పదహారు వేర్వేరు కార్ల మధ్య ఉన్న తేడా వంటిది.
Quantization నేరుగా బ్యాండ్విడ్త్ సమస్యను ఎదుర్కొంటుంది. మోడల్ వెయిట్స్ (weights) సాధారణంగా పదహారు-బిట్ ఫ్లోటింగ్-పాయింట్ ఫార్మాట్లలో నిల్వ చేయబడతాయి. వాటిని ఎనిమిది-బిట్ లేదా నాలుగు-బిట్ ఇంటిజర్లుగా కుదించడం ద్వారా, బస్ (bus) ద్వారా ప్రయాణించే డేటా పరిమాణాన్ని మీరు సగానికి లేదా అంతకంటే ఎక్కువ తగ్గించవచ్చు. మోడల్ అర్థవంతమైన అవుట్పుట్ను అందించడానికి తగినంత ఖచ్చితత్వం (precision) అవసరం, కానీ ఆధునిక పోస్ట్-ట్రైనింగ్ క్వాంటైజేషన్ పద్ధతులు నాణ్యతను దెబ్బతీయకుండా మోడల్ యొక్క మెమరీ ఫుట్ప్రింట్ను గణనీయంగా తగ్గించగలవు. డేటా ప్రయాణించే పరిమాణం తగ్గితే, మెమరీ కంట్రోలర్ వద్ద వేచి ఉండాల్సిన సమయం కూడా తగ్గుతుంది.
FlashAttention ఇంటర్మీడియట్ ఫలితాలను GPU యొక్క వేగవంతమైన ఆన్-చిప్ మెమరీలోనే ఉంచడానికి అటెన్షన్ మెకానిజంను పునర్నిర్మిస్తుంది. స్టాండర్డ్ అటెన్షన్ పెద్ద అటెన్షన్ మ్యాట్రిక్స్లను నెమ్మదిగా ఉండే ఎక్స్టర్నల్ మెమరీలోకి వ్రాయాల్సి వచ్చేది మరియు మళ్ళీ వాటిని చదవాల్సి వచ్చేది. FlashAttention గణనను SRAMలో సరిపోయే చిన్న టైల్స్గా విభజిస్తుంది, సాఫ్ట్మాక్స్ (softmax) మరియు స్కేలింగ్ దశలను ఆన్-చిప్లోనే నిర్వహిస్తుంది, మరియు కేవలం తుది అవుట్పుట్లను మాత్రమే హై-బ్యాండ్విడ్త్ మెమరీలోకి వ్రాస్తుంది. ఇది మెయిన్ మెమరీకి చేసే అనవసరమైన ప్రయాణాలను తగ్గించడం కోసం కొంచెం అదనపు కంప్యూట్ను ఉపయోగిస్తుంది, ఇది దాదాపు ఎల్లప్పుడూ లాభదాయకమైన నిర్ణయం.
PagedAttention మెమరీ వృధా అయ్యే మరొక రకమైన సమస్యను పరిష్కరిస్తుంది. డీకోడింగ్ సమయంలో, KV cache ఊహించని విధంగా పెరుగుతుంది. సాంప్రదాయ వ్యవస్థలు ప్రతి సీక్వెన్స్ కోసం నిర్ణీత, నిరంతర (contiguous) మెమరీ భాగాలుగా కేటాయిస్తాయి, దీనివల్ల కొన్ని సీక్వెన్స్లు త్వరగా ముగియడం మరియు మరికొన్ని విస్తరించడం వల్ల మెమరీలో పెద్ద ఖాళీలు ఏర్పడతాయి. PagedAttention ఆపరేటింగ్ సిస్టమ్స్ నుండి వర్చువల్ మెమరీ భావనను తీసుకుంటుంది. ఇది KV cache ఎంట్రీలను నిర్ణీత పరిమాణ బ్లాక్లలో నిల్వ చేస్తుంది, వీటిని నాన్-కంటిగ్యువస్ (non-contiguously) గా కేటాయించవచ్చు మరియు ఇండైరెక్షన్ టేబుల్ ద్వారా మ్యాప్ చేయవచ్చు. ఇది రిజర్వ్ చేయబడిన కానీ సగం ఖాళీగా ఉన్న బఫర్లలో మెమరీ వృధా కాకుండా నిరోధిస్తుంది మరియు పెద్ద బ్యాచ్ సైజులను అనుమతిస్తుంది, తద్వారా ఫ్రాగ్మెంటేషన్ ఓవర్హెడ్ కంటే ఉపయోగకరమైన పనులతో మెమరీ బస్ను బిజీగా ఉంచుతూ మొత్తం త్రూపుట్ను మెరుగుపరుస్తుంది.
ప్రశ్నను మార్చండి
లేటెన్సీ (latency) పెరిగినప్పుడు, చాలా బృందాలు తాము చిన్న మోడల్కు మారాలా లేదా తమ ఇన్ఫరెన్స్ సర్వర్ను తిరిగి వ్రాయాలా అని అడుగుతాయి. ఆ ప్రశ్నలు ముఖ్యమే, కానీ అవి ద్వితీయమైనవి. మొదటి ప్రశ్న హార్డ్వేర్ గురించి ఉండాలి. మీ GPU నిజంగా కంప్యూటింగ్ చేయడంలో బిజీగా ఉందా, లేదా డేటా కోసం ఎదురుచూస్తూ ఖాళీగా ఉందా?
మీ యుటిలైజేషన్ మెట్రిక్స్ను చూడండి. GPU కంప్యూట్ ఆక్యుపెన్సీతో పాటు మెమరీ బ్యాండ్విడ్త్ సాచురేషన్ను ప్రొఫైల్ చేయండి. డీకోడింగ్ సమయంలో మీరు అధిక మెమరీ కాంటెన్షన్ మరియు తక్కువ అర్థమెటిక్ ఇంటెన్సిటీని చూస్తే, అది మోడల్ ఆర్కిటెక్చర్ సమస్య కాదు. అది ఫిజిక్స్ సమస్య. దీనికి పరిష్కారం మరింత క్లీన్ పైథాన్ (Python) కోడ్ రాయడం వల్ల రాదు. మరింత అగ్రెసివ్గా బ్యాచింగ్ చేయడం, పైపు ద్వారా వేగంగా వెళ్లడానికి మీ వెయిట్స్ను క్వాంటైజ్ చేయడం, ఆన్-చిప్లోనే ఉండేలా అటెన్షన్ను పునర్నిర్మించడం మరియు మెమరీ తక్కువ కాకుండా పెద్ద బ్యాచ్లను సరిపోయేలా KV క్యాచ్ను నిర్వహించడం ద్వారా మాత్రమే పరిష్కారం లభిస్తుంది.
ఒకసారి మీరు ఇన్ఫరెన్స్ను ఈ కోణంలో చూడటం ప్రారంభిస్తే, ఆప్టిమైజేషన్ అనేది ఒక యాంత్రిక ప్రక్రియలా మారుతుంది. మోడల్ ఇంటెలిజెన్స్ వల్ల వేగం తగ్గుతుందనే అపోహల వెంట పడటం మానేసి, హార్డ్వేర్ వాస్తవంగా ఏమి అందించగలదనే దానిపై ఆధారపడి ఇంజనీరింగ్ నిర్ణయాలు తీసుకోవడం ప్రారంభిస్తారు. స్కేల్ అయ్యే ప్రొడక్షన్ సిస్టమ్స్ను, కేవలం పనిచేసే సిస్టమ్స్ నుండి వేరు చేసేది ఇదే మార్పు.