Google के Tensor Processing Unit को समझना

• TPU एक Google domain-specific accelerator है, जो deep learning के लिए सचमुच ज़रूरी गणनाओं को ही रखता है और बाकी को साहसपूर्वक हटा देता है
• इसकी शुरुआत 2013 में data center विस्तार की सीमाओं के कारण हुई, और 12 साल में यह 7वीं पीढ़ी ‘Ironwood’ तक विकसित हो गया
• semiconductor performance के अपने-आप बेहतर होते जाने का दौर खत्म होने के साथ, इंतज़ार करने के बजाय खुद design करने की रणनीति एक अहम विकल्प बन गई
• TPU कोई एक chip नहीं, बल्कि hardware, compiler, network और operating software का co-design system है
• हर नई पीढ़ी के साथ performance competition से ज़्यादा power, deployment और operating cost (TCO) design के केंद्र में आ गए
• TPU की प्रतिस्पर्धात्मक बढ़त किसी एक तकनीक में नहीं, बल्कि 10 साल से अधिक समय में जमा हुए design और operations experience के योग में है

---

Something New

• TPU कोई गुप्त हथियार नहीं, बल्कि लंबे समय तक चली public research और internal iteration से निखरा हुआ परिणाम है
• Google ने data center को लगातार बड़ा करने के बजाय, गणना करने के तरीके को मूल रूप से बदलने का रास्ता चुना
• 2013 के आसपास data center capacity को दोगुना करने का दबाव और समय की पाबंदी, 15 महीनों में TPU के जन्म तक जा पहुँची
• अप्रैल 2025 में Google Cloud Next पर 7वीं पीढ़ी के TPU Ironwood की घोषणा हुई, जिसमें 9,216 chips/pod, 42.5 Exaflops और 10MW के आँकड़े दिए गए
• GPU deep learning के लिए बनाए गए डिवाइस नहीं थे, और TPU को शुरुआत से ही neural network computation को ध्यान में रखकर design किया गया
• इस विकल्प से केवल compute performance ही नहीं, बल्कि power efficiency और operational stability में भी संरचनात्मक बढ़त मिली
• यह “संयोग” नहीं, बल्कि constraints, trade-offs और co-design की बार-बार की गई प्रक्रिया से जमा हुआ परिणाम है

Slowing Down

• Moore’s Law और Dennard Scaling के कमज़ोर पड़ने से पहले केवल नए CPU का इंतज़ार करने भर से program तेज़ हो जाते थे, लेकिन अब वह आधार टूट चुका है
• transistor की संख्या बढ़ती रही, लेकिन power और heat की सीमाएँ performance improvement को रोकने लगीं
• साथ ही neural network को बड़े data और बड़े model चाहिए, जिससे compute demand विस्फोटक रूप से बढ़ी
• इसी वजह से “हर काम थोड़ा-थोड़ा अच्छा करने वाली chip” के बजाय “एक काम बेहद अच्छी तरह करने वाली chip” की ज़रूरत बढ़ी
• neural network computation का केंद्र matrix multiplication आधारित repetitive operations होना, specialized design को संभव बनाता है

The Inference Chip

• पहली TPU ने training नहीं, बल्कि inference, यानी पहले से train किए गए model को चलाने पर ध्यान दिया
• TPUv1 ने cache, branch prediction और multithreading को हटाकर control cost को न्यूनतम किया
• इसके बजाय बड़े matrix multiplication को लगातार प्रोसेस करने वाले Systolic Array (MXU) पर सारे resources केंद्रित किए गए
• execution के दौरान निर्णय लेने के बजाय, यह compile समय पर तय किए गए execution order का पालन करती है
• नतीजतन, समान power पर यह GPU और CPU की तुलना में कहीं अधिक inference work संभाल सकती थी

The Training Chip

• training को inference की तुलना में कहीं अधिक compute और विस्तृत numeric representation range की ज़रूरत होती है
• TPUv2 से inference-only architecture से आगे बढ़कर training के लिए flexibility जोड़ी गई
• मुख्य बदलाव matrix (MXU), vector (VPU) और control (Scalar Unit) की भूमिकाओं का विभाजन था
• execution flow को XLA compiler पहले से compute और decide करता है, और chip उसे वैसे ही execute करती है
• कई TPU को एक डिवाइस की तरह काम कराने के लिए dedicated high-speed interconnect (ICI) भी साथ में design किया गया

Scaling Up

• system बड़ा होने के साथ सवाल “यह कितना तेज़ है” से बदलकर “इसे कितनी देर तक और कितनी कम लागत में चलाया जा सकता है?” हो गया
• इसके लिए compute units के पास बड़ी on-chip memory (CMEM) रखी गई, ताकि धीमी DRAM access कम हो
• recommendation systems जैसे sparse data वाले workloads के लिए SparseCore जैसे dedicated units भी जोड़े गए
• chip के अंदर और chip के बीच communication को अलग करके wiring complexity और bottleneck को संरचनात्मक रूप से कम किया गया
• performance numbers से ज़्यादा operational efficiency पूरे design को निर्धारित करने लगी

Island Hopping

• हज़ारों TPU वाले environment में failure कोई अपवाद नहीं, बल्कि एक मूल धारणा है
• लक्ष्य ऐसा system है जो रुके नहीं, यानी partial failure को absorb कर सके
• काम कई TPU में बाँटकर चलाया जाता है, लेकिन उसे एक ही program की तरह दिखने और manage होने दिया जाता है
• समस्या आने पर पूरे system को रोकने के बजाय तेज़ reallocation और restart चुना जाता है
• यह जटिल प्रक्रिया ज़्यादातर operating software अपने-आप संभालता है

Datacenter Network विस्तार

• जब TPU का एक group पर्याप्त नहीं रहता, तो कई groups को जोड़ना पड़ता है
• सामान्य network की सीमाओं के कारण optical switching (OCS) अपनाया गया
• इससे पूरे data center को एक विशाल compute resource की तरह संगठित करना संभव हुआ
• मौजूदा execution model को बढ़ाने वाला तरीका और पूरी तरह नया asynchronous execution model (Pathways) साथ-साथ मौजूद हैं
• अब बड़े model और ज़्यादा जटिल communication patterns को संभालना संभव हो गया

Ceci n’est pas une TPU

• नवीनतम TPU संख्याओं के हिसाब से भले ही अत्यंत प्रभावशाली हों, लेकिन मुख्य सिद्धांत शुरुआती दौर जैसे ही हैं
• ज़रूरी computation पर ध्यान देना और अनावश्यक complexity को हटाना अब भी दिशा का आधार है
• केवल hardware specs के आधार पर इस system की नकल नहीं की जा सकती
• compiler (XLA), dedicated interconnect (ICI), optical switching (OCS) और operations scheduler को साथ काम करना पड़ता है
• TPU एक बार की invention नहीं, बल्कि सैकड़ों साधारण निर्णयों के जमा होते जाने का परिणाम है

ज़रूर याद रखने योग्य कुछ मुख्य तकनीकें

• Systolic Array (MXU): matrix multiplication को उच्च दक्षता से संभालने वाला TPU का दिल
• XLA compiler: execution order को पहले से calculate करके control cost हटाता है
• BF16: training के लिए आवश्यक range बनाए रखते हुए hardware cost घटाने वाला numeric format
• ICI / OCS: chip, rack और data center को एक साथ बाँधने वाली dedicated communication architecture
• TCO-केंद्रित design: instant performance के बजाय long-term operating cost को optimize करने वाली सोच

[GN⁺]

Hacker News टिप्पणियाँ

• Scaling ML किताब में भी TPU पर एक शानदार सेक्शन है - How to Think About TPUs
  • मैंने भी इस लेख के साथ Henry Ko का TPU गहन विश्लेषण दिलचस्पी से पढ़ा। XLA और scheduler जो काम करते हैं, वह सच में प्रभावशाली है। यह Itanium से कहीं अधिक जटिल संरचना है, फिर भी software इस राक्षसी chip का वास्तव में उपयोग कर पाता है, यह हैरान करने वाली बात है। अच्छा होता अगर XLA को और व्यापक रूप से अपनाया जाता। open source होने के बावजूद industry में रुचि की कमी खलती है। लगता है Nvidia अब Tiles के साथ कुछ वैसी ही दिशा में बढ़ रही है। हालांकि, मेरी जानकारी में XLA अभी कई machines के बीच scheduling में बहुत उपयोगी नहीं है
• इस लेख की संरचनात्मक व्याख्या अच्छी लगी। TPU पर ज्यादातर लेख व्यावहारिक हिस्सों को छोड़ देते हैं, लेकिन इस बार लेख ने concepts को वास्तविक चीजों से जोड़ा, इसलिए समझ एकदम साफ हो गई
• TPU architecture को उद्देश्य के हिसाब से design करने का स्तर सिर्फ एक generation के design पर खत्म नहीं होता। Ironwood 7वीं पीढ़ी का TPU है, और उसका evolution बहुत महत्वपूर्ण है
• मुझे अब भी लगता है कि Moore's Law मरा नहीं है। अगर 1965 से 2025 तक 60 सालों में हर 2 साल पर दोगुना मानें, तो 30 बार doubling होती है। सैद्धांतिक रूप से 2025 में लगभग 107 अरब transistors होने चाहिए, और वास्तव में Apple M1 Ultra में 114 अरब transistors हैं
  • कुछ लोग Moore's Law को इस मजबूत अर्थ में लेते हैं कि 'doubling की गति स्थिर रहती है', लेकिन वह तो बहुत पहले ही खत्म हो चुका। फिर भी अगर उस गति को धीरे-धीरे बदलने वाला स्थिरांक माना जाए, तो यह अब भी मान्य है। समस्या यह है कि केवल boundary values देखकर average निकालने से हालिया बदलाव के trend का सही प्रतिबिंब नहीं मिलता
• यह हैरानी की बात है कि कुछ ही वर्षों में चीन TPU का बड़े पैमाने पर उत्पादन कर सकता है, इस दृष्टिकोण को बड़ी खबर नहीं माना गया। अगर ऐसा हुआ, तो Google, NVIDIA जैसी कंपनियों को बड़ा झटका लग सकता है। 2022~2023 में चीनियों द्वारा TPUv4 और v6 दस्तावेज़ों के लीक का मामला भी था। पहले से ही एक चीनी startup अपना TPU cluster बनाकर revenue कमा रहा है
  • लेकिन वास्तविक निर्माण ही सबसे कठिन हिस्सा है। चीन के पास design knowledge पर्याप्त है, लेकिन असली chip बनाने की क्षमता कम है। semiconductor manufacturing के लिए TSMC जैसी 'जादुई तकनीकी क्षमता' चाहिए। Intel और Samsung भी कुछ हद तक सक्षम हैं, लेकिन अंतर अब भी बड़ा है
  • लेख का आधा हिस्सा TPU, Borb, lilpunet, optical switching network आदि के बीच software dependency पर था। ऐसे जटिल systems को केवल manufacturing technology से कॉपी करना मुश्किल है
  • Google TPU को अपनी services के लिए इस्तेमाल करता है, इसलिए कोई दूसरी कंपनी वैसी chip बना भी ले, तो उसे बहुत बड़ा नुकसान नहीं होगा। बल्कि NVIDIA का बाज़ार एकाधिकार टूटना ज्यादा यथार्थवादी है। FMA/MAC units का design सरल है, इसलिए Apple, Qualcomm, AMD, Amazon, Huawei जैसी लगभग हर कंपनी अपनी 'TPU' पहले से डाल रही है। अमेरिका 6 लाख चीनी छात्रों को train भी कर दे, तब भी असली कुंजी manufacturing और process technology ही है
  • “परमाणु अड्डे और श्रमशक्ति” वाली बात समझ नहीं आती। परमाणु अड्डों का semiconductor manufacturing से क्या संबंध है, समझ नहीं आता। और 6 लाख छात्र सभी chip design नहीं सीख रहे होते
  • यह कहना कि TPU और सस्ते हो जाने का डर है, थोड़ा व्यंग्य जैसा लगता है
• GCP में TPU इस्तेमाल करते समय क्या अब भी वही अजीब Google bucket system से बंधे रहना पड़ता है? पहले वह सच में बहुत असुविधाजनक था

[crawler]

मेरी एक निजी इच्छा है
कृपया TPU इतना व्यापक रूप से commercial हो जाए कि कंपनियों को GPU इस्तेमाल करने की ज़रूरत ही न पड़े।
ताकि Nvidia फिर से personal consumer GPUs पर ध्यान दे.....

[xguru]

मूल लेख तकनीकी रूप से बहुत गहरा है, इसलिए उसे ज्यों का त्यों संक्षेप में बताने पर वह सिर्फ तकनीकी शब्दों की सूची बनकर रह जाएगा। इसलिए मैंने इसे थोड़ा इस तरह व्यवस्थित किया है कि विकास की स्थिति को समझना आसान हो। जिन लोगों को तकनीकी विवरण जानने हैं, वे चित्रों के साथ समझाने वाले मूल लेख को देखें।