- TPU Google द्वारा विकसित बड़े पैमाने के AI training और inference के लिए बनाया गया एक custom chip है, जिसकी design philosophy GPU से अलग है
- यह scalability और energy efficiency पर ज़ोर देता है, और hardware (जैसे system-on-chip configuration, large on-chip memory) तथा software (XLA compiler) को साथ में design करता है
- इसकी मुख्य संरचना systolic array, pipelining, और ahead-of-time compilation पर आधारित है, और यह अधिकांश deep learning operations (खासकर matrix multiplication) के लिए optimized है
- OCI और OCS तकनीकों की वजह से flexible node configuration, high-performance parallel processing, और विभिन्न topology विकल्प संभव होते हैं
- TPU system एक single chip से लेकर बहुत बड़े multipod तक hierarchical तरीके से scale होता है, जिससे बड़े AI model training और efficient resource utilization संभव होता है
अवलोकन और पृष्ठभूमि
- TPU Google द्वारा AI services के बड़े पैमाने पर विस्तार की जरूरतों को पूरा करने के लिए विकसित किया गया ASIC आधारित chip है
- शुरुआती दौर में GPU, FPGA और ASIC अपनाने के बीच विचार किया गया था, लेकिन 2013 में voice search आदि में deep learning के बढ़ते उपयोग के साथ custom hardware की जरूरत बढ़ी
- आज यह Gemini, Veo जैसी प्रमुख AI services में अनिवार्य भूमिका निभाता है, और recommendation models (DLRM) आदि में व्यापक रूप से इस्तेमाल हो रहा है
TPU single chip architecture
मूल संरचना
- TPUv4 के आधार पर एक single chip में 2 TensorCore होते हैं (inference-focused TPU में 1)
- प्रत्येक TensorCore CMEM(128MiB), HBM(32GiB) memory units से जुड़ा होता है
TensorCore की आंतरिक संरचना
- Matrix Multiply Unit (MXU) : 128x128 systolic array आधारित, matrix multiplication के लिए समर्पित
- Vector Unit (VPU) : सामान्य element-wise operations करता है
- Vector Memory (VMEM; 32MiB) : HBM से data कॉपी करके working memory के रूप में उपयोग होता है
- Scalar Unit + Scalar Memory (SMEM; 10MiB) : control flow, scalar operations, और memory address management संभालता है
GPU से संरचनात्मक अंतर
- TPU में on-chip memory (CMEM, VMEM, SMEM) GPU की तुलना में काफी बड़ी होती है
- HBM capacity GPU में अधिक होती है, और compute core की संख्या भी GPU में कहीं ज्यादा होती है
- TPUv5p के आधार पर, प्रति chip 500 TFLOPs/sec, और पूरे pod (8960 chips) के आधार पर 4.45 ExaFLOPs/sec स्तर का उच्च प्रदर्शन मिलता है
TPU design philosophy
1. Systolic array और pipelining
- Systolic array में processing elements (PE) एक array में व्यवस्थित होते हैं और computation result को पड़ोसी elements तक पहुंचाते हैं
- data input के बाद अलग control के बिना लगातार computation संभव होती है, और memory read/write केवल input और output पर होता है
- यह मुख्य रूप से matrix multiplication, convolution के लिए optimized है
- pipeline processing के जरिए computation और data movement एक साथ चलाकर throughput optimize किया जाता है
Systolic array की कमी - sparsity
- सभी compute elements हमेशा सक्रिय रहते हैं, इसलिए यह sparse matrices के लिए उपयुक्त नहीं है
- अगर आगे DL models irregular sparsity की ओर बढ़ते हैं, तो यह एक सीमा बन सकती है
2. Ahead-of-time (AoT) compilation और minimal cache
- TPU-XLA co-design के जरिए irregular memory access की जरूरत वाले cache dependency को कम किया जाता है, जिससे energy usage घटता है
- XLA compiler operation graph का विश्लेषण करके memory access pattern पहले से निकालता है, और cache की जगह scratchpad memory केंद्रित संचालन करता है
- JAX का
@jit JIT और AoT के बीच का एक रूप है, जो पहली execution पर static graph बनाकर फिर XLA से AoT compile करता है
- input shape बदलने पर फिर से compilation की जरूरत पड़ती है, और dynamic padding/iteration में inefficiency रहती है
- इससे उच्च energy efficiency मिलती है, लेकिन flexibility की कमी एक downside है
TPUv4 energy efficiency
- आधुनिक chips HBM3 memory का उपयोग करते हैं, जिससे energy consumption कम होता है
- computation की तुलना में memory operations की energy cost दर्जनों से सैकड़ों गुना तक अधिक हो सकती है, इसलिए memory access कम करके efficiency में बड़ा सुधार संभव है
TPU multi-chip architecture
Tray स्तर (Tray/Board; 4 chips)
- 1 tray में 4 TPU chips (8 TensorCore) और CPU Host शामिल होते हैं
- Host↔Chip connection PCIe से होता है, जबकि Chip↔Chip connection Inter-Core Interconnect(ICI) से होता है, जो ज्यादा bandwidth देता है
Rack स्तर (Rack; 4x4x4=64 chips)
- 1 rack में 64 TPU chips होते हैं, जो 3D torus(4x4x4) में ICI और OCS(Optical Circuit Switching) से जुड़े होते हैं
- Google में rack, pod, और slice को अलग-अलग परिभाषित किया जाता है
- Rack: 64-chip physical unit (=cube)
- Pod: ICI और OCS से जुड़ सकने वाली अधिकतम unit (उदाहरण: TPUv4=4096 chips=64 racks)
- Slice: 4 chips से Superpod के बीच की कोई भी arbitrary abstract configuration unit
OCS के फायदे
- Wraparound : हर axis को ring (1D torus) बनाकर nodes के बीच worst-case hop count कम किया जाता है
- Flexible non-contiguous slices: OCS switching structure की वजह से भौतिक रूप से दूर nodes को भी एक slice में जोड़ा जा सकता है, जिससे resource utilization और maintenance आसान होते हैं
- Twisted topology: समान chip count (x,y,z fixed) पर connection structure बदलकर खास operation patterns के लिए speed optimize की जा सकती है (उदाहरण: twisted torus)
Topology उपयोग के उदाहरण
- Cube shape: data/tensor parallelism के लिए उपयुक्त (maximum bandwidth)
- Linear shape (cigar-like) : pipeline parallelism के लिए फायदेमंद
- Twisted torus: all-to-all communication की जरूरत होने पर (जैसे tensor parallel) speed बढ़ती है
Superpod (Full Pod/Superpod; TPUv4: 4096 chips/64 racks)
- कई racks को ICI और OCS से जोड़कर बहुत बड़ा system बनाया जाता है
- Slice topology के प्रकार के अनुसार communication bandwidth और parallel processing performance में अंतर आता है
- OCS की वजह से non-contiguous slices, twisted topology जैसी flexibility मिलती है
Multipod स्तर (Multi-pod/Multislice; TPUv4: 4096 से अधिक)
- कई pods को data center network (DCN) से जोड़ा जाता है, जिससे large-scale training infrastructure मिलता है, लेकिन bandwidth ICI की तुलना में कम होती है
- PaLM training में 2 pods (6144 TPUv4) का उपयोग किया गया, और कुल 6 pods के resource management के साथ इसे चलाया गया
- बड़े models के development में, researcher settings (जैसे parallel dimensions) के अनुसार XLA compiler communication pattern optimize करता है
- XLA हर slice और pod के बीच communication operations insert करता है, जिससे न्यूनतम code changes के साथ large-scale distributed training संभव होती है
वास्तविक hardware और diagram उदाहरण
- TPU rack: 4x4x4 3D torus एक unit होता है, और हर row में 2 trays (8 chips) होते हैं
- TPUv4 tray: वास्तविक रूप में 4 PCIe ports होते हैं (प्रत्येक में 1 TPU)
- TPUv4 chip: बीच में ASIC, और आसपास 4 HBM stacks (2 TensorCore के आधार पर)
- TPUv4i(inference type) chip floorplan: 1 TensorCore, जिसमें CMEM का area काफी बड़ा होता है
समापन
- Google TPU Research Cloud(TRC) के research support के प्रति आभार व्यक्त किया गया है
संदर्भ सामग्री
अभी कोई टिप्पणी नहीं है.