Google की पहली Tensor Processing Unit (TPU): आर्किटेक्चर

 (thechipletter.substack.com)
1 पॉइंट द्वारा GN⁺ 2024-03-26 | 1 टिप्पणियां | WhatsApp पर शेयर करें
  • Google TPU v1 एक ASIC था, जिसे deep learning services की inference cost और scale संभालने के लिए GPU की तुलना में 10 गुना cost-performance के लक्ष्य के साथ 2013 के अंत से 15 महीनों में विकसित किया गया था
  • इसका मुख्य बिंदु 1978 में H.T. Kung और Charles E. Leiserson के systolic array को 256×256 MAC संरचना के रूप में लागू करना था, जिससे matrix multiplication के intermediate results को memory तक आना-जाना कम हुआ
  • TPU v1 PCIe के जरिए host से communication करता है और DDR3-2133 में weights store करता है; inference flow Read_Host_Memory, Read_Weights, Matrix_Multiply / Convolve, Activate, Write_Host_Memory जैसे कुछ commands से बनता है
  • 8-bit×8-bit integer multiplication और quantization के जरिए floating-point computation के die area burden से बचा गया, और TensorFlow से linked User Space Driver·Kernel Driver execution control करते हैं
  • Intel Haswell CPU और Nvidia K80 GPU उस समय comparison targets थे; TPU v1 inference में लगभग 15~30 गुना तेज था और GPU की तुलना में relative incremental performance/Watt 25~29 गुना अधिक थी, लेकिन यह training device नहीं था

Inference cost घटाने के लिए TPU v1 की शुरुआत

  • TPU v1 project का लक्ष्य deep learning-based services के लिए जरूरी बड़े पैमाने के hardware cost को घटाने हेतु inference ASIC को तेजी से बनाना था
    • लक्ष्य GPU की तुलना में inference में 10 गुना cost-performance advantage था
    • साथ में fast development, high performance, large-scale deployment, नए workloads का immediate support और cost efficiency की जरूरत थी
  • TPU नाम इस बात से आया कि यह tensor operations को accelerate करने वाला device है
    • वास्तविक TPU v1 hardware द्वारा संभाली जाने वाली core computation मुख्य रूप से vector और matrix operations हैं
    • neural network की hidden layers और output layers को input value vector और weight matrix के product पर activation function लागू करने के result के रूप में व्यक्त किया जा सकता है
    • जब कई input data आते हैं, तो यह matrix multiplication result के हर item पर activation function लागू करने के रूप में होता है

Systolic array से matrix multiplication processing

  • TPU v1 H.T. Kung और Charles E. Leiserson के 1978 paper Systolic Arrays (for VLSI) में दिए गए systolic system concept का उपयोग करता है
    • यह ऐसी structure है जिसमें कई processors regular rhythm में data compute और pass करते हैं
    • हर processor छोटी computation करते हुए data को लगातार अंदर-बाहर move करता है
  • सरल 2×2 matrix multiplication में, अगर input values को ऊपर और बाईं ओर से सही क्रम में डाला जाए, तो result 2×2 MAC array से स्वाभाविक रूप से निकलता है
    • हर MAC multiplication और accumulation करता है
    • partial sums array के भीतर store होते हैं, और final result चलती हुई diagonal shape में दिखाई देता है
    • 2×2 example में 4 steps लगते हैं, लेकिन वास्तव में जैसे ही upper-left MAC खाली होता है, अगला matrix multiplication शुरू किया जा सकता है, इसलिए हर 2 cycles में नया matrix multiplication संभव है
  • मुख्य बात यह है कि data को सही order में systolic array में feed करने पर values और results का flow ही required computation order बना देता है
    • intermediate results को main memory में store करके फिर वापस लाने की जरूरत नहीं होती
    • Matrix Multiply Unit की structure और input order की वजह से intermediate results जरूरत के समय अपने-आप use हो जाते हैं

TPU v1 system configuration

  • TPU v1 host computer से PCIe high-speed serial bus के जरिए communicate करता है और अपने DDR3 DRAM को directly access करता है
  • प्रमुख components इस प्रकार हैं
    • DDR3 DRAM / Weight FIFO
      • weights DDR3-2133 interface से जुड़े DDR3 RAM chips में store होते हैं
      • host memory से PCIe के जरिए पहले से load किए जाते हैं, फिर Matrix Multiply Unit में उपयोग के लिए Weight FIFO में move होते हैं
    • Matrix Multiply Unit
      • यह 256×256 MAC से बना systolic array है
      • ऊपर से 256 weights और बाईं ओर से 256 data inputs लेता है
    • Accumulators
      • systolic matrix unit के नीचे से निकले results को store करते हैं
    • Activation
      • neural network के activation function को लागू करने वाला stage है
    • Unified Buffer / Systolic Data Setup
      • activation function apply करने के results को store करता है, और next layer computation के लिए Matrix Multiply Unit के input के रूप में फिर feed करने की तैयारी करता है

Computation format और instruction set

  • TPU v1 का Matrix Multiply Unit 8-bit×8-bit integer multiplication करता है
    • quantization का उपयोग करके बड़े die area की जरूरत वाली floating-point computation से बचा जाता है
  • instruction set लगभग 20 commands वाला CISC design है
    • instructions memory से fetch नहीं होते, बल्कि host computer उन्हें PCIe के जरिए भेजता है
  • inference flow का अधिकांश हिस्सा 5 प्रमुख commands से बना है
    • Read_Host_Memory
      • input values को host memory से PCIe के जरिए Unified Buffer में पढ़ता है
    • Read_Weights
      • weight memory से Weight FIFO में weights पढ़ता है
    • Matrix_Multiply / Convolve
      • Unified Buffer के input को Accumulators में भेजते हुए matrix multiplication या convolution करता है
      • B×256 input और 256×256 constant weight input को multiply करके B×256 output बनाता है, और इसमें B pipeline cycles लगते हैं
    • Activate
      • Accumulators के input पर ReLU, Sigmoid आदि artificial neurons के non-linear functions लागू करता है और result को Unified Buffer में output करता है
    • Write_Host_Memory
      • Unified Buffer के result को PCIe के जरिए host memory में लिखता है
  • इस flow को मोटे तौर पर इस रूप में देखा जा सकता है
Read_Host_Memory
Read_Weights
Loop_Start
    Matrix_Multiply
    Activate
Loop_End
Write_Host_Memory
  • matrix unit Unified Buffer read/write को घटाकर energy बचाने के लिए systolic execution का उपयोग करता है
    • data बाईं ओर से आता है, और weights ऊपर से load होते हैं
    • 256 elements की MAC operation diagonal wavefront की तरह matrix से होकर गुजरती है

TensorFlow और driver stack

  • TPU v1 hardware को real services में इस्तेमाल करने के लिए इसे support करने वाले software stack की जरूरत थी
    • Google TensorFlow को develop और use करता था, इसलिए TensorFlow को TPU v1 के साथ काम कराने के लिए drivers बनाना core step था
  • TPU software stack को CPU और GPU stacks के साथ compatible होना था
    • applications को TPU पर तेजी से port किया जा सकना चाहिए था
    • TPU पर चलने वाला application का हिस्सा आम तौर पर TensorFlow में लिखा जाता है, और GPU या TPU पर executable API में compile होता है
  • GPU की तरह TPU stack भी User Space Driver और Kernel Driver में बंटा है
    • Kernel Driver को lightweight रखा जाता है; यह केवल memory management और interrupts handle करता है और long-term stability को target करता है
    • User Space Driver अक्सर बदलता है और TPU execution setup/control, TPU order के अनुसार data reformatting, API calls को TPU commands में बदलना, और application binary generation संभालता है

28nm process और die layout

  • TPU v1 TSMC के अपेक्षाकृत mature 28nm process से manufactured था
    • उस समय Google data centers में इस्तेमाल होने वाले Intel Haswell CPU और Nvidia K80 GPU chips अधिक advanced processes पर बने थे
    • Google के अनुसार TPU v1 का die area इन chips के die area के आधे से भी कम था
  • simple ISA ने decoding और related work के लिए जरूरी die overhead को कम किया
    • control area die area का केवल 2% था
    • Matrix Multiply Unit 24%, और Unified Buffer 29% था

Performance comparison और स्पष्ट सीमाएं

  • TPU v1 पहले से trained models को Google scale की real services में अधिक efficiently use करने के लिए inference device है
    • यह training speed या efficiency improve करने के लिए design किया गया device नहीं है
    • specialized hardware development में inference और training अलग-अलग challenges पैदा करते हैं
  • 2013 के आधार पर core comparison targets Intel Haswell CPU और Nvidia K80 GPU थे
    • TPU v1 में K80 GPU की तुलना में MAC count 25 गुना अधिक है
    • TPU v1 में K80 GPU की तुलना में on-chip memory 3.5 गुना अधिक है
    • TPU v1 inference में K80 GPU और Haswell CPU की तुलना में लगभग 15~30 गुना तेज है
    • GPU की तुलना में relative incremental performance/Watt 25~29 गुना है
  • custom architecture की वजह से TPU v1 ने उस समय के CPU और GPU से अधिक inference performance और कम energy use हासिल किया
  • fast inference और power efficiency के single goal पर केंद्रित first-generation design होने के कारण, यह सीमा बनी रही कि इसे training के लिए design नहीं किया गया था

संबंधित पढ़ाई

1 टिप्पणियां

 
GN⁺ 2024-03-26
Hacker News की राय

  • Groq के CEO Jonathan Ross ने हाल के एक podcast interview में Google में शुरुआती TPU बनाने की कहानी बताई। उनके मुताबिक, वह मूल रूप से inference speed की समस्या झेल रही टीम के पास बैठे थे और 20% time में एक FPGA बनाया था।
    काम करने वाला prototype बन जाने के बाद Jeff Dean ने हिसाब लगाया और ASIC पर जाने का फैसला किया गया।
    आज के हिसाब से मुझे लगता है कि Google को TPU टीम को एक अलग कंपनी के रूप में spin off कर देना चाहिए। Nvidia को चुनौती देने वाला यह इकलौता भरोसेमंद competitor है, और software support भी Nvidia के बाद वाले स्तर का है।
    https://open.spotify.com/episode/0V9kRgNS7Ds6zh3GjdXUAQ?si=q...

    • Nvidia की बढ़त को महत्व के क्रम में देखें तो मैं कहूंगा: reserved foundry capacity, बेहद integrated software, पहले से मौजूद hardware architecture, और customer relationships।
      लेकिन हर एक में कमजोरी है। Foundry capacity तंग है, लेकिन अगर Nvidia ज्यादा महंगे AI chips बेच सकती है तो वह consumer GPU market की बलि दे सकती है। अगर कोई competitor कुछ साल पहले से बड़ा दांव लगा दे, या Intel जैसी ज्यादा production capacity वाली कंपनी अपनी priorities बदल दे, तो यह बढ़त खत्म हो जाएगी।
      Proprietary software अगर industry standard हो तो सुविधाजनक होता है, लेकिन असली महत्व use case पर बहुत निर्भर करता है। TPU के लिए hardware design GPU की तुलना में मूलतः काफी सरल लगता है; ray tracing, texture samplers, rasterization की जरूरत नहीं, और ज्यादातर matrix multiplication व बहुत सारी memory ही चाहिए।
      Customer relationships बातचीत में लगातार बने रहने के लिए उपयोगी हैं, लेकिन ऐसे market में जहां थोड़ा-सा भी फायदा खोजा जाता है, dollar प्रति FLOPS सबसे ज्यादा देने वाला hardware supplier अपनी production capacity भरने लायक customers हासिल कर ही लेगा। इसलिए लगता है कि कुछ सालों में competition काफी तेजी से वास्तविक हो जाएगा।
    • Google को TPU टीम को अलग कंपनी के रूप में spin off करना चाहिए—इस बात पर, market size और लगभग monopoly जैसी स्थिति देखते हुए, मेरा मानना है कि यह Pixel hardware business को लगभग तुरंत पीछे छोड़ सकती है।
      हालांकि TPU Google के भीतर भी अपेक्षाकृत scarce compute resource है, और internal demand पूरी करना भी मुश्किल हो सकता है।
    • Amazon ने इसी तरह का काम कर रही Annapurna Labs को acquire किया था, इसलिए उसके पास अपना Trainium/Inferentia silicon है, और support के मामले में वह Google से निश्चित रूप से काफी आगे है।
    • यह कहना गलत है कि Nvidia का इकलौता भरोसेमंद competitor TPU है। AMD और Intel के पास भी Habana के जरिए H100-स्तर के performance वाले GPU हैं।
    • Groq सच में कमाल है। बहुत से startups सिर्फ hype और वादों के साथ आते हैं, लेकिन Groq पहले ही एक शानदार working product के साथ आया, और सिर्फ यही उसे पसंद करने की पर्याप्त वजह है।
      मैं बहुत कम कंपनियों के लिए कहता हूं कि मैं उनका इतना सम्मान करता हूं, लेकिन Groq का सच में सम्मान करता हूं।

  • Google ने TPU invent किया और Google Research ने LLM paper तक निकाला, फिर भी NVDA और AI startups ने value का लगभग 100% कैसे ले लिया, यह समझ नहीं आता।

    • Xerox और PARC पर एक पुराना मजाक है, जिसमें समझाया जाता है कि “copier company को paperless office बेचना मुश्किल है।”
      Google के मामले में तुलना यह हो सकती है कि अगर ChatGPT जैसी चीज व्यापक रूप से उपलब्ध कराने का प्रस्ताव रखा जाए, तो वह search engine की paid placements और उन sites के ad revenue को खा सकती है जहां लोगों को अब जाना ही नहीं पड़ेगा। इसलिए शायद यह निर्णय हुआ हो कि competitive जरूरत पड़ने पर ही, कम झटकेदार तरीके से सावधानी से इसे लागू किया जाए।
      असल में बात इतनी सरल नहीं होगी, लेकिन अगर यही वजह हो तो काफी मजेदार होगा।
    • Google 18 महीनों से ज्यादा उन products पर focus नहीं कर पाता जो अरबों डॉलर का profit नहीं बनाते। वह ads के नशे में है।
    • यह कहना अभी बहुत जल्दी है कि Google AI से value नहीं ले पाएगा। उसके पास अपने products में AI integrate करने के पर्याप्त मौके हैं।
    • ऐतिहासिक precedent के लिए Xerox PARC देख लीजिए।
    • OpenAI ने बहुत ज्यादा compensation देकर Google के talent को खींच लिया।
      https://www.linkedin.com/posts/eolver_googles-defense-agains...

  • मैं Google employee हूं; अगर आपने कुछ समय से TPU नहीं देखा है, तो v5 जरूर देखें। अब यह PyTorch/JAX support करता है, इसलिए TensorFlow-only वाले दौर की तुलना में इस्तेमाल करना काफी आसान हो गया है।

    • अपने server में लगाने के लिए TPU v5 कहां खरीद सकता हूं? अगर जवाब “cloud” है, तो यही वजह है कि Nvidia इतना dominate करता है।

  • इस लेख ने abstract रूप में बिखरे हुए कई pieces को अच्छी तरह जोड़कर दिखाया कि वे असल में silicon के अंदर कैसे flow करते हैं।
    यह देखना खास तौर पर अच्छा लगा कि simple CISC instructions LLM inference steps से लगभग सीधे-सीधे match करते हैं।

  • यह शायद मेरी अनजानगी दिखाने वाला बेवकूफी भरा सवाल हो, लेकिन consumer side पर मैं लगातार सुन रहा हूं कि M1~M4 chips कुछ AI tasks के लिए अच्छे हैं।
    आजकल मेरे लिए सबसे जरूरी tools Photoshop, Resolve जैसे हैं, और मैंने देखा है कि Apple के नए custom chips पर वे मेरी पुरानी machine की तुलना में कहीं ज्यादा तेजी से चलते हैं।
    हो सकता है कि यह उन कामों से ठीक से न जुड़ता हो जो यह chip या H100 कर सकते हैं, लेकिन मैं जानना चाहता हूं कि क्या इनमें कुछ हद तक संबंध है। बेशक Apple अपने custom chips अलग से नहीं बेचता, इसलिए practical बनाने के लिए उसे external server जैसी किसी form factor में ढेर सारे GPU और AI chips वाला product निकालना होगा।

    • मैं खुद को expert तो नहीं कहूंगा, लेकिन मैंने M1 और कई GPUs benchmark किए हैं।
      M* chips unified memory इस्तेमाल करते हैं, और खासकर Pro/Max/Ultra में 1080 जैसे GPU की तुलना में भी memory bandwidth बहुत ज्यादा है। M1 Ultra की memory bandwidth 2080 और 3090 के बीच के आसपास है।
      छोटे batch size में, खासकर batch 1 पर—जैसा कि ज्यादातर local workloads में होता है—inference compute power की बजाय memory bandwidth से bottleneck होता है। इसी वजह से कहा जाता है कि M* chips machine learning के लिए अच्छे हैं।
      लेकिन H100 मुख्य रूप से बहुत बड़े batch size वाली training में इस्तेमाल होता है, और बड़े models train करने के लिए बहुत सारे interconnects चाहिए होते हैं। उस scale पर arithmetic intensity बहुत ऊंची होती है, इसलिए भले ही M* chips को network से जोड़ा जा सके, वे बहुत competitive नहीं होंगे। यह H100 जैसे ज्यादा power खाने वाले chip की तुलना में power/efficiency Pareto curve का एक अलग point चुनने जैसा है।

  • Google को सच में जो करना चाहिए, वह है 2nm EUV क्षेत्र में उतरना और 2nm से नीचे जाना
    चाहे वह electron lithography हो या ASML जैसी कोई तकनीक जो चिप पर पैटर्न छापती है—अगर उनके पास वैसा कुछ आ गया, तो वे सचमुच खतरनाक खिलाड़ी बन जाएंगे। शायद Google X-स्टाइल का कोई hardcore moonshot project चाहिए
    या फिर शायद उनके पास कोई एक उपकरण खरीदने के लिए करीब 500 मिलियन डॉलर हों। अगर TPU वाकई इतना अच्छा है, तो अपनी तकनीक और अपनी fab तक vertical integration करना अच्छा business बन सकता है

    • ईमानदारी से कहें तो यह लगभग असंभव है। पहले तो दशकों की trade secrets समझनी होंगी, पहली advanced fab बनाने में लगने वाली दसियों अरब या सैकड़ों अरब डॉलर की पूंजी चाहिए होगी, उसे ठीक से काम करने वाले business में परिपक्व होने में 10–20 साल लगेंगे, और वे जो volume बनाएंगे वह बहुत ही मामूली होगा—इन सबको देखते हुए, संभावना यही है कि वे मौजूदा advanced process से कई साल पीछे की जगह पर लगभग 10 साल बाद पहुंचने के लिए 500 अरब डॉलर जला देंगे
      आज की advanced fabs इसलिए पैसा कमाती हैं क्योंकि वे अलग-अलग ग्राहकों और use cases के लिए general-purpose computing devices बनाती हैं, और इसके पीछे दशकों से जमा talent और engineering है। ऊपर से customers भी core areas में स्वतंत्र रूप से innovation को आगे बढ़ाते हैं, जैसे Micron का chip-on-chip HDI yield improvement, Xilinx का die-to-die communication fabric और multichip substrate design
      TPU कभी भी उतना volume नहीं बना पाएगा जितना चाहिए, और वह ऐसे customers भी नहीं ला पाएगा जो profitable economies of scale बना सकें। Google को competitors की तुलना में attractive pricing भी देनी होगी
      अगर business case पर्याप्त रूप से compelling होगा, तो मौजूदा fabs खुशी से capacity allocate कर देंगी। TPU अभी दूर-दूर तक इतना compelling नहीं है

  • TensTorrent के Jim Keller की एक presentation सुनी थी, जिसमें उन्होंने AI cores बनाने का अलग approach समझाया। इसमें 5 RISC-V cores होते हैं: एक data load करने के लिए, एक data upload करने के लिए, और बाकी matrix operations के लिए dedicated
    उन्होंने Google TPU का भी जिक्र किया था; कहा कि इसकी programming VLIW से निपटने जैसी लगती है और compiler पर करीब 500 लोग लगे हुए थे

  • मूल लेख में लिखा है कि “TPU v1 लगभग 20 instructions वाला CISC design है”, और यह देखकर हंसी आती है कि CISC/RISC एक तीखे observation से research program, revolutionary technology, marketing buzzword होते हुए आखिरकार पूरी तरह meaningless बात तक पहुंच गया लगता है
    शायद इसे terminology का lifecycle कहना चाहिए

    • पक्का नहीं, लेकिन computer architecture में जो सीखा था उसके हिसाब से CISC और RISC का फर्क instructions की संख्या से ज्यादा instructions की complexity से जुड़ा था
      इसलिए अगर TPU में instructions कम हैं, लेकिन हर instruction काफी complex है, तो वह CISC हो सकता है। हालांकि computer architecture आखिरी बार 15 साल पहले grad school में पढ़ा था, इसलिए याद धुंधली है। उस semester का ज्यादातर हिस्सा अब बेकार हो चुके Itanium से जुड़े काम में भी निकल गया था
    • ऐसा लगता है कि यह imply कर रहा है कि available instructions की संख्या CISC को अलग करती है, लेकिन असल में वह कभी criteria नहीं था

  • Foundry capacity की demand बहुत ज्यादा लगती है, और मुझे हैरानी है कि जब Microsoft या Google अपने chips बनाते हैं और production की जरूरत होती है, तो वे queue में सबसे आगे कैसे पहुंचते हैं
    क्या वे इतने simple हैं कि “पुरानी और कम-demand वाली” fabs में भी बनाए जा सकें? मेरी समझ है कि Apple और Nvidia ने काफी foundry capacity पहले से ले रखी है

    • मोटे तौर पर यह cutting edge से एक generation पीछे वाली पुरानी fab में चल जाता है
      https://en.wikipedia.org/wiki/Tensor_Processing_Unit#Product...
      HBM जैसे क्षेत्रों में भी उनकी काफी presence और spending है, और SemiAnalysis पर इससे जुड़े अच्छे लेख हैं

  • अगर LLMs सच में -1, 0, 1 quantization के साथ बड़े पैमाने पर चलन में आ जाते हैं, तो hardware कैसे बदलेगा, यह देखने की उत्सुकता है