పెద్ద భాషా నమూనాలను (Large Language Models) ఫోన్లో నడపడం అనేది ఇకపై కేవలం ఒక పరిశోధనా ప్రయోగం మాత్రమే కాదు. ఇది వాస్తవ రూపంలో అందుబాటులోకి వస్తోంది. అయితే, మీరు ఒక చిన్న డెమో నుండి అసలైన ఉత్పత్తికి మారిన క్షణమే, లేటెన్సీ (latency) మళ్ళీ సమస్యగా మారుతుంది. మీరు మోడల్ను తగ్గించారు, వెయిట్స్ను క్వాంటైజ్ (quantize) చేశారు, అయినా prefill దశ నెమ్మదిస్తుంది. టోకెన్లు ఆగిపోతాయి. UI ఫ్రీజ్ అవుతుంది. దీనికి నిజమైన కారణం ఎప్పుడూ బయటపడదు.
ఆండ్రాయిడ్ పరికరాలలో, LLM prefill సమయంలో కంప్యూట్ (compute) అనేది దాదాపు ఎప్పుడూ అడ్డంకి (bottleneck) కాదు. మెమరీ బ్యాండ్విడ్త్ (Memory bandwidth) అడ్డంకిగా మారుతుంది. ఆధునిక ఫ్లాగ్షిప్ SoCలు శక్తివంతమైన GPU మరియు NPU కోర్లతో వస్తాయి, ఇవి మెమరీ సబ్సిస్టమ్ డేటాను అందించే వేగం కంటే చాలా వేగంగా గణిత ప్రక్రియలను (arithmetic) పూర్తి చేయగలవు. మీరు ఒక సాధారణ అటెన్షన్ ఇంప్లిమెంటేషన్ను ప్రొఫైల్ చేసినప్పుడు, ఎగ్జిక్యూషన్ యూనిట్లు పూర్తిగా ఉపయోగించబడవు. అవి వేచి ఉంటాయి. DRAM కోసం వేచి ఉంటాయి.
మీ క్వాంటైజ్డ్ (Quantized) మోడల్ ఇంకా నెమ్మదిగా ఎందుకు అనిపిస్తుంది?
ఆన్-డివైస్ ఇన్ఫరెన్స్ (on-device inference) కోసం క్వాంటైజేషన్ (Quantization) అనేది ప్రాథమిక దశగా మారింది. FP16 నుండి INT8కి వెయిట్స్ను తగ్గించడం వల్ల మోడల్ పరిమాణం సగానికి తగ్గుతుంది మరియు స్టోరేజ్ ఆదా అవుతుంది. ఇది సహాయపడుతుంది, కానీ అటెన్షన్ లేయర్ లేటెన్సీని సరిచేయదు. కారణం సరళం: క్వాంటైజేషన్ మీరు నిల్వ చేసే డేటా పరిమాణాన్ని తగ్గిస్తుంది, కానీ అటెన్షన్ మెకానిజం చేసే మెమరీ ట్రాన్సాక్షన్ల సంఖ్యను తగ్గించదు.
టెక్స్ట్బుక్ పద్ధతిలో అమలు చేయబడిన ప్రామాణిక మల్టీ-హెడ్ అటెన్షన్ లేయర్, ప్రతి లేయర్ కోసం DRAMకి మూడు పూర్తి రౌండ్ ట్రిప్స్ చేస్తుంది. Query, Key, మరియు Value మ్యాట్రిసెస్ మెయిన్ మెమరీ నుండి చదవబడతాయి, స్కోర్లు లెక్కించబడతాయి మరియు మధ్యంతర ఫలితాలు తిరిగి వ్రాయబడతాయి. గణితం చాలా సులభం. కానీ డేటా మూవ్మెంట్ మాత్రం చాలా కష్టతరమైనది. పవర్ మరియు థర్మల్ బడ్జెట్లు తక్కువగా ఉండే ఆండ్రాయిడ్లో, ఈ విధానం మెమరీ బస్ను తీవ్రంగా ప్రభావితం చేస్తుంది. ప్రాసెసర్ ఒకే వంతెనను దాటడానికి మూడుసార్లు టోల్ (toll) చెల్లిస్తున్నట్లు అన్నమాట.
మీరు INT8 మోడళ్లను ఉపయోగిస్తూ, ప్రాంప్ట్ పొడవు పెరిగే కొద్దీ prefill దశ ఇంకా ఎందుకు క్వాడ్రాటిక్గా (quadratically) పెరుగుతుందో అని ఆలోచిస్తుంటే, ఇక్కడ సమాధానం ఉంది. వెయిట్స్ చిన్నవిగా ఉన్నప్పటికీ, యాక్టివేషన్ ట్రాఫిక్ మాత్రం భారీగా ఉంటుంది.
అడ్డంకి మెమరీయే తప్ప, గణితం కాదు
పరిష్కారాన్ని అర్థం చేసుకోవడానికి, రూఫ్లైన్ (roofline) వైపు చూడండి. Snapdragon 8 Gen 3 మరియు Dimensity 9300 వంటి చిప్లలోని మొబైల్ GPUలు మరియు NPUలు, వాటి LPDDR5X ఇంటర్ఫేస్లు తట్టుకోగలిగే సామర్థ్యం కంటే చాలా ఎక్కువ థియరిటికల్ కంప్యూట్ త్రూపుట్ను కలిగి ఉంటాయి. ఒక సాధారణ అటెన్షన్ కెర్నల్లో, ప్రతి హెడ్ Q మరియు K యొక్క లబ్ధంపై softmaxను లెక్కిస్తుంది, ఆపై దానిని Vతో గుణిస్తుంది. ప్రతి మధ్యంతర స్కోర్ మ్యాట్రిక్స్ గ్లోబల్ మెమరీలో నిక్షిప్తం చేయబడుతుంది. అంటే మీ పీక్ DRAM రీడ్స్ సీక్వెన్స్ పొడవు యొక్క వర్గంతో, లేదా O(n²) తో పెరుగుతాయి. 1024-టోకెన్ ప్రాంప్ట్ కోసం, మెమరీ ట్రాఫిక్ ఇప్పటికే ఎగ్జిక్యూషన్ సమయాన్ని ఎక్కువగా ఆక్రమించేంత పెద్దదిగా ఉంటుంది.
కోర్లు తక్కువగా ఉపయోగించబడుతున్నాయి ఎందుకంటే అవి లేటెన్సీని దాచలేవు. ఆధునిక ప్రాసెసర్లు పైప్లైన్లను నిరంతరం నింపడానికి క్యాచెస్ (caches) పై ఆధారపడతాయి. ఒక అల్గారిథమ్ నిరంతరం క్యాష్లో దొరకకపోయినా మరియు DRAM నుండి డేటాను తీసుకురావాల్సి వచ్చినా, ఎగ్జిక్యూషన్ యూనిట్లు ఖాళీగా ఉంటాయి. కంప్యూట్ సామర్థ్యం మరియు మెమరీ సరఫరా మధ్య ఉన్న ఈ నిర్మాణాత్మక అసమతుల్యతను ఎంత క్వాంటైజేషన్ చేసినా సరిచేయలేదు.
టైలింగ్ (Tiling) ద్వారా బ్యాండ్విడ్త్ను ఎలా తిరిగి పొందవచ్చు
దీనికి పరిష్కారం టైలింగ్ స్ట్రాటజీ (tiling strategy). ఇది మధ్యంతర స్కోర్లను DRAMకి పంపే బదులు, ఆన్-చిప్ SRAMలో ఉంచుతుంది. Flash Attentionని నడిపించే అదే ఆలోచన ఇది, దీనిని ఆండ్రాయిడ్ కంప్యూట్ స్టాక్ కోసం అనువదించడమైనది. మెమరీలో పూర్తి n × n స్కోర్ మ్యాట్రిస్ను సృష్టించే బదులు, మీరు గణనను L1 క్యాష్లో సరిపోయే చిన్న టైల్స్గా విభజిస్తారు. మీరు లోకల్ సాఫ్ట్మాక్స్ స్టాటిస్టిక్స్ను లెక్కిస్తారు, రన్నింగ్ మాక్స్ విలువలను మరియు నార్మలైజేషన్ సమ్స్ను సేకరిస్తారు మరియు చివరిగా వెయిటెడ్ అవుట్పుట్లను మాత్రమే మెమరీకి తిరిగి వ్రాస్తారు.
ఇది బ్యాండ్విడ్త్ కాంప్లెక్సిటీని మారుస్తుంది. పీక్ DRAM రీడ్స్ O(n²) నుండి O(n)కి తగ్గుతాయి, ఎందుకంటే మీరు ఇకపై పూర్తి స్కోర్ మ్యాట్రిస్లను మెయిన్ మెమరీ ద్వారా తరలించాల్సిన అవసరం లేదు. ప్రధానమైన పని ఎగ్జిక్యూషన్ యూనిట్ల పక్కనే ఉన్న SRAM లోనే జరుగుతుంది.
ఒక స్పష్టమైన ఉదాహరణ కోసం, టైల్ సైజు 64 మరియు హెడ్ డైమెన్షన్ 128 అని అనుకుందాం. స్కోర్ టైల్ 16 KB స్థలాన్ని ఆక్రమిస్తుంది. ఈ ఫుట్ప్రింట్ Snapdragon 8 Gen 3 మరియు Dimensity 9300 వంటి ప్రస్తుత ఫ్లాగ్షిప్ SoCల L1 క్యాష్లో సులభంగా సరిపోతుంది. గణితం లోకల్గా ఉంటుంది. మెమరీ బస్ పై భారం తగ్గుతుంది.