SweetSpot: An Analytical Model for Predicting Energy Efficiency of LLM Inference

यहाँ शोध पत्र "SweetSpot: LLM इन्फरेंस की ऊर्जा दक्षता की भविष्यवाणी करने के लिए एक विश्लेषणात्मक मॉडल" का विस्तृत तकनीकी सारांश दिया गया है।

1. समस्या विवरण (Problem Statement)

लार्ज लैंग्वेज मॉडल्स (LLMs) आधुनिक AI के केंद्र बन गए हैं, लेकिन उनके इन्फरेंस चरण (inference phase) में मॉडल के जीवनचक्र की अधिकांश (90% तक) ऊर्जा खर्च होती है। इस ऊर्जा पदचिह्न (energy footprint) का अनुमान लगाने के मौजूदा दृष्टिकोण अत्यधिक सरलीकृत धारणाओं पर निर्भर करते हैं:

• रैखिक धारणा (Linear Assumption): अधिकांश मौजूदा मॉडल और व्यावसायिक मूल्य निर्धारण रणनीतियाँ यह मानती हैं कि ऊर्जा खपत इनपुट ($n_{in}$) और आउटपुट ($n_{out}$) टोकन गणना का एक सरल रैखिक फलन (linear function) है।
• जटिलता की कमी (Missing Complexity): ये मॉडल ट्रांसफॉर्मर की ऑटोरेग्रेसिव प्रकृति (autoregressive nature) को समझाने में विफल रहते हैं, जहाँ प्रीफिल चरण (prefill phase) (इनपुट को समानांतर रूप से प्रोसेस करना) और डिकोड चरण (decode phase) (क्रमवार रूप से टोकन जेनरेट करना) के बीच कम्प्यूटेशनल और मेमोरी लागत मौलिक रूप से भिन्न होती है।
• परिणाम: इससे ऊर्जा भविष्यवाणियों में अशुद्धि आती है और अनुकूलन (optimization) के अवसर छूट जाते हैं, क्योंकि अनुक्रम लंबाई (sequence lengths) और ऊर्जा दक्षता के बीच का संबंध वास्तव में गैर-रैखिक (non-linear) होता है।

2. कार्यप्रणाली (Methodology)

लेखक SweetSpot का प्रस्ताव करते हैं, जो ट्रांसफॉर्मर आर्किटेक्चर की कम्प्यूटेशनल (FLOPs) और मेमोरी-एक्सेस जटिलता से प्राप्त एक विश्लेषणात्मक मॉडल है।

A. सैद्धांतिक सूत्रीकरण (Theoretical Formulation)

मॉडल इन्फरेंस ऊर्जा को दो अलग-अलग चरणों में विभाजित करता है:

1. प्रीफिल चरण (Prefill Phase): पूरे इनपुट अनुक्रम को समानांतर में प्रोसेस किया जाता है। इसकी लागत अटेंशन मैकेनिज्म ( $QK^T$ की गणना) के कारण द्विघातीय (quadratic) जटिलता ($O(n_{in}^2)$) द्वारा संचालित होती है।
2. डिकोड चरण (Decode Phase): टोकन ऑटोरेग्रेसिव तरीके से जेनरेट किए जाते हैं। प्रति टोकन लागत अनुक्रम लंबाई के संबंध में रैखिक (linear) होती है ($O(n_{in} + t)$) क्योंकि इसमें की-वैल्यू (KV) स्टेट्स का पुन: उपयोग शामिल है, लेकिन यह मेमोरी बैंडविड्थ द्वारा सीमित होती है।

लेखक मॉडल के दो संस्करण निकालते हैं:

• SweetSpot (केवल FLOPs): शुद्ध फ्लोटिंग-पॉइंट ऑपरेशंस (floating-point operations) पर आधारित।
• SweetSpot (पूर्ण/Complete): FLOPs मॉडल को मेमोरी-एक्सेस जटिलता ($M$) को शामिल करने के लिए विस्तारित करता है, यह स्वीकार करते हुए कि आधुनिक GPU अक्सर कंप्यूट-बाउंड के बजाय मेमोरी-बाउंड होते हैं, विशेष रूप से डिकोडिंग के दौरान।

अंतिम ऊर्जा-प्रति-टोकन ($E_{tok}$) समीकरण (समीकरण 5) है:
$$E_{tok}(n_{in}, n_{out}) = \theta_0 + \frac{\theta_1 n_{in}^2}{n_{out}} + \theta_2 n_{in} + \frac{\theta_3 n_{in}}{n_{out}} + \theta_4 n_{out} + \frac{\theta_5}{n_{out}}$$

• $\frac{n_{in}^2}{n_{out}}$ वाला पद जेनरेट किए गए टोकनों पर द्विघातीय प्रीफिल लागत के परिशोधन (amortization) को दर्शाता है।
• $n_{out}$ वाला पद डिकोडिंग लागतों के रैखिक संचय (linear accumulation) को दर्शाता है।

B. प्रयोगात्मक सेटअप (Experimental Setup)

• हार्डवेयर: NVIDIA H100 GPUs (4x) जो NVLink के माध्यम से जुड़े हुए हैं।
• इन्फरेंस इंजन: TensorRT-LLM (इसकी अत्याधुनिक ऊर्जा दक्षता और हार्डवेयर-जागरूक अनुकूलन जैसे इन-फ़्लाइट बैचिंग के लिए चुना गया)।
• मॉडेल्स: विभिन्न प्रकार के 12 LLMs, जो 1B से 9B पैरामीटर्स तक हैं (जिसमें OPT, LLaMA, Gemma, Falcon, Qwen2, और Granite शामिल हैं), जो विभिन्न अटेंशन मैकेनिज्म (MHA, GQA) को कवर करते हैं।
• वर्कलोड: इनपुट और आउटपुट लंबाई 64 से 4096 टोकन तक भिन्न थी। प्रयोग 1,000 समवर्ती अनुरोधों (concurrent requests) के साथ चलाए गए थे ताकि संतृप्ति (saturation) सुनिश्चित हो सके और माप शोर (measurement noise) को कम किया जा सके।
• मेट्रिक्स: कुल ऊर्जा ($E_{tot}$) और ऊर्जा-प्रति-टोकन ($E_{tok}$), जिसे 500ms के अंतराल पर NVIDIA NVML के माध्यम से मापा गया।

3. मुख्य योगदान (Key Contributions)

1. SweetSpot मॉडल: पहला विश्लेषणात्मक मॉडल जो LLM इन्फरेंस ऊर्जा को इनपुट और आउटपुट लंबाई के गैर-रैखिक फलन के रूप में सटीक रूप से वर्णित करता है, जो स्पष्ट रूप से द्विघातीय प्रीफिल लागत और रैखिक डिकोड लागत के बीच के अंतर्संबंध को मॉडल करता है।
2. "स्वीट स्पॉट" की खोज: यह शोध पत्र सिद्ध करता है कि ऊर्जा दक्षता एकदिष्ट (monotonic) नहीं है। इनपुट और आउटपुट लंबाई का एक विशिष्ट विन्यास मौजूद है जहाँ दक्षता शिखर (peak) पर होती है।
   • शिखर दक्षता (Peak Efficiency): लघु-से-मध्यम इनपुट और मध्यम-लंबाई के आउटपुट के साथ होती है।
   • अकुशल क्षेत्र (Inefficiency Zones): बहुत लंबे इनपुट (उच्च प्रीफिल लागत के कारण) या बहुत छोटे आउटपुट (प्रीफिल लागत के खराब परिशोधन के कारण) के लिए दक्षता तेजी से गिरती है।
3. वैधीकरण और सटीकता: मॉडल को व्यापक अनुभवजन्य डेटा (empirical data) के विरुद्ध मान्य किया गया था, जिसने 1.79% का मीन एब्सोल्यूट परसेंटेज एरर (MAPE) प्राप्त किया, जो मौजूदा रैखिक बेसलाइन (जिनका MAPE >30% था) से काफी बेहतर प्रदर्शन करता है।
4. वास्तुकला संबंधी अंतर्दृष्टि (Architectural Insights): यह अध्ययन मात्रात्मक रूप से बताता है कि अटेंशन मैकेनिज्म (MHA बनाम GQA) ऊर्जा को कैसे प्रभावित करते हैं। GQA (Grouped-Query Attention) का उपयोग करने वाले मॉडल KV कैश मेमोरी ट्रैफिक को कम करने के कारण MHA (Multi-Head Attention) की तुलना में लगातार अधिक ऊर्जा-कुशल होते हैं।

4. मुख्य परिणाम (Key Results)

• सटीकता: SweetSpot ने भविष्यवाणी त्रुटि को ~31% (पिछले सर्वश्रेष्ठ बेसलाइन) से घटाकर 1.79% (पूर्ण मॉडल) और 2.27% (केवल FLOPs) कर दिया।
• दक्षता लाभ: अनुक्रम लंबाई को "स्वीट स्पॉट" के साथ संरेखित करने से सबसे खराब स्थिति (लंबे इनपुट + छोटा आउटपुट) की तुलना में ऊर्जा उपयोग को 33.41× तक कम किया जा सकता है।
• गैर-रैखिकता (Non-Linearity): अध्ययन ने पुष्टि की कि दक्षता तब चरम पर होती है जब इनपुट लागत आउटपुट लंबाई द्वारा पर्याप्त रूप से परिशोधित (amortized) हो जाती है, लेकिन इससे पहले कि रैखिक डिकोडिंग ओवरहेड हावी हो जाए।
  • उदाहरण: 64 टोकन के स्थिर इनपुट के लिए, दक्षता 128–256 आउटपुट टोकन के आसपास शिखर पर होती है।
• बैचिंग का प्रभाव: अध्ययन में पाया गया कि जबकि बैच आकार पूर्ण थ्रूपुट को प्रभावित करता है, सापेक्ष दक्षता रुझान (वक्र का आकार) सुसंगत रहता है। हालांकि, बैच आकार को कम संतृप्त (under-saturate) करने से समग्र दक्षता काफी कम हो जाती है।

5. महत्व और निहितार्थ (Significance and Implications)

• उत्पादन अनुकूलन (Production Optimization): यह मॉडल अनुकूली पीढ़ी रणनीतियों (adaptive generation strategies) के लिए एक सैद्धांतिक आधार प्रदान करता है। सिस्टम गतिशील रूप से अत्यधिक लंबे इनपुट को ट्रंकेट (truncate) कर सकते हैं या जनरेशन सीमाओं को समायोजित कर सकते हैं ताकि वे "स्वीट स्पॉट" के भीतर काम कर सकें, जिससे उपयोगिता से समझौता किए बिना ऊर्जा लागत को नाटकीय रूप से कम किया जा सकता है।
• सतत AI (Sustainable AI): अनुभवजन्य अनुमान (guesswork) से विश्लेषणात्मक भविष्यवाणी की ओर बढ़कर, यह कार्य अधिक टिकाऊ AI सिस्टम के डिजाइन को सक्षम बनाता है, जो वैश्विक स्तर पर LLM तैनाती के विस्तार के साथ अत्यंत महत्वपूर्ण है।
• हार्डवेयर-सॉफ्टवेयर सह-डिज़ाइन (Hardware-Software Co-design): परिणाम इस बात पर प्रकाश डालते हैं कि सॉफ्टवेयर-स्तरीय अनुकूलन (जैसे बैचिंग और ट्रंकेशन) को अंतर्निहित हार्डवेयर बाधाओं (मेमोरी बैंडविड्थ बनाम कंप्यूट) की समझ द्वारा निर्देशित किया जाना चाहिए।
• भविष्य की दिशाएं: लेखक मॉडल के पैमाने (10B पैरामीटर्स से कम पर ध्यान केंद्रित करना) और हार्डवेयर विविधता के संबंध में सीमाओं का उल्लेख करते हैं, जो बड़े LLMs और विषम (heterogeneous) हार्डवेयर क्लस्टरों तक मॉडल को सामान्य बनाने के लिए भविष्य के कार्य का सुझाव देते हैं।

संक्षेप में, SweetSpot ऊर्जा दक्षता को एक अस्पष्ट, अनुभवजन्य मीट्रिक से एक पूर्वानुमेय, डिजाइन योग्य पैरामीटर में बदल देता है, जो बड़े पैमाने पर AI इन्फरेंस के कार्बन फुटप्रिंट को कम करने के लिए एक रोडमैप प्रदान करता है।