วิเคราะห์ TPU เชิงลึก

 (henryhmko.github.io)
13 คะแนน โดย GN⁺ 2025-06-23 | ยังไม่มีความคิดเห็น | แชร์ทาง WhatsApp
  • TPU คือชิปแบบกำหนดเองสำหรับการฝึกและอนุมาน AI ขนาดใหญ่ที่พัฒนาโดย Google และมี แนวคิดการออกแบบ ที่แตกต่างจาก GPU
  • เน้น ความสามารถในการขยายระบบ และ ประสิทธิภาพด้านพลังงาน โดยออกแบบทั้งฮาร์ดแวร์ (เช่น โครงสร้าง system-on-chip, หน่วยความจำบนชิปขนาดใหญ่) และซอฟต์แวร์ (XLA compiler) ร่วมกัน
  • โครงสร้างหลักคือ systolic array ร่วมกับการทำ pipelining และการคอมไพล์ล่วงหน้า ซึ่งเหมาะอย่างยิ่งกับงาน deep learning ส่วนใหญ่ (โดยเฉพาะการคูณเมทริกซ์)
  • ด้วยเทคโนโลยี OCI และ OCS จึงสามารถจัดโหนดได้อย่างยืดหยุ่น ประมวลผลแบบขนานประสิทธิภาพสูง และเลือก topology ได้หลากหลาย
  • ระบบ TPU สามารถขยายแบบลำดับชั้นได้ตั้งแต่ชิปเดี่ยวไปจนถึง multi-pod ขนาดมหึมา ทำให้ฝึกโมเดล AI ขนาดใหญ่และใช้ทรัพยากรได้อย่างมีประสิทธิภาพ

ภาพรวมและภูมิหลัง

  • TPU เป็นชิปที่ใช้ ASIC ซึ่ง Google พัฒนาขึ้นเพื่อรองรับการขยายตัวของบริการ AI ในระดับมหาศาล
  • ในช่วงแรกมีการชั่งน้ำหนักระหว่าง GPU, FPGA และ ASIC แต่เมื่อการใช้ deep learning เช่นการค้นหาด้วยเสียงขยายตัวมากขึ้นในปี 2013 ความจำเป็นของฮาร์ดแวร์แบบเฉพาะทางก็เพิ่มขึ้น
  • ปัจจุบันมีบทบาทสำคัญต่อบริการ AI หลักอย่าง Gemini, Veo และยังถูกใช้อย่างกว้างขวางกับโมเดลแนะนำ (DLRM) เป็นต้น

โครงสร้างชิปเดี่ยวของ TPU

องค์ประกอบพื้นฐาน

  • ใน TPUv4 ชิปเดี่ยวมี TensorCore 2 ชุด (TPU ที่เน้นงานอนุมานมี 1 ชุด)
  • TensorCore แต่ละชุดเชื่อมต่อกับหน่วยความจำ CMEM (128MiB) และ HBM (32GiB)

โครงสร้างภายใน TensorCore

  • Matrix Multiply Unit (MXU) : ใช้ systolic array ขนาด 128x128 รับผิดชอบการคูณเมทริกซ์โดยเฉพาะ
  • Vector Unit (VPU) : ประมวลผลการคำนวณแบบองค์ประกอบทั่วไป
  • Vector Memory (VMEM; 32MiB) : คัดลอกข้อมูลจาก HBM มาใช้เป็นหน่วยความจำเตรียมงาน
  • Scalar Unit + Scalar Memory (SMEM; 10MiB) : ดูแล control flow การคำนวณแบบ scalar และการจัดการที่อยู่หน่วยความจำ

ความแตกต่างเชิงโครงสร้างจาก GPU

  • TPU มีหน่วยความจำบนชิป (CMEM, VMEM, SMEM) ใหญ่กว่า GPU มาก
  • ความจุ HBM นั้น GPU มากกว่า และจำนวนคอร์ประมวลผลก็มีมากกว่ามากเช่นกัน
  • อิงตาม TPUv5p สามารถให้สมรรถนะ 500 TFLOPs/sec ต่อชิป และทั้ง pod (8960 ชิป) อยู่ที่ระดับ 4.45 ExaFLOPs/sec
โฆษณา

แนวคิดการออกแบบ TPU

1. Systolic array และ pipelining

  • systolic array คือการจัดเรียง processing element (PE) เป็นอาร์เรย์ เพื่อส่งต่อผลลัพธ์การคำนวณไปยังองค์ประกอบที่อยู่ติดกัน
  • หลังป้อนข้อมูลเข้าแล้วสามารถคำนวณต่อเนื่องได้โดยแทบไม่ต้องมีการควบคุมเพิ่มเติม และการอ่าน/เขียนหน่วยความจำจะเกิดขึ้นเฉพาะตอนรับเข้าและส่งออกเท่านั้น
  • เหมาะอย่างยิ่งกับงาน การคูณเมทริกซ์ และ convolution
  • ใช้การประมวลผลแบบ pipeline เพื่อให้การคำนวณและการเคลื่อนย้ายข้อมูลเกิดขึ้นพร้อมกัน จึงปรับ throughput ได้ดี

ข้อเสียของ systolic array - sparsity

  • หน่วยคำนวณทุกตัวจะทำงานอยู่เสมอ จึงไม่เหมาะกับ เมทริกซ์แบบ sparse
  • หากในอนาคตโมเดล DL มุ่งไปสู่ความ sparse แบบไม่เป็นระเบียบมากขึ้น จุดนี้อาจกลายเป็นข้อจำกัด

2. การคอมไพล์ล่วงหน้า (AoT) และการลดการพึ่งพาแคช

  • การออกแบบร่วมกันระหว่าง TPU-XLA ช่วยลดการพึ่งพาแคชที่ต้องใช้กับการเข้าถึงหน่วยความจำแบบไม่สม่ำเสมอ จึงลดการใช้พลังงาน
  • XLA compiler จะวิเคราะห์ computation graph เพื่อคำนวณรูปแบบการเข้าถึงหน่วยความจำล่วงหน้า และใช้งานแบบยึด scratchpad memory เป็นหลักแทนแคช
  • @jit ของ JAX เป็นรูปแบบกึ่งกลางระหว่าง JIT และ AoT โดยจะสร้าง static graph เมื่อรันครั้งแรกแล้วจึงคอมไพล์แบบ AoT ผ่าน XLA
  • หากรูปแบบอินพุตเปลี่ยนไปจะต้องคอมไพล์ใหม่ ทำให้ไม่มีประสิทธิภาพกับ dynamic padding/loop บางแบบ
  • แม้จะได้ประสิทธิภาพด้านพลังงานสูง แต่ก็มีข้อเสียคือความยืดหยุ่นน้อย

ประสิทธิภาพด้านพลังงานของ TPUv4

  • ชิปสมัยใหม่ใช้หน่วยความจำ HBM3 เพื่อลดการใช้พลังงาน
  • งานหน่วยความจำใช้พลังงานมากกว่างานคำนวณหลายสิบถึงหลายร้อยเท่า ดังนั้นการลดการเข้าถึงหน่วยความจำจึงช่วยเพิ่มประสิทธิภาพได้อย่างมาก

โครงสร้างหลายชิปของ TPU

ระดับถาด (Tray/Board; 4 ชิป)

  • 1 ถาดประกอบด้วยชิป TPU 4 ตัว (TensorCore 8 ชุด) และ CPU Host
  • การเชื่อมต่อ Host↔Chip ใช้ PCIe ส่วน Chip↔Chip ใช้ Inter-Core Interconnect (ICI) ซึ่งให้แบนด์วิดท์สูงกว่า
โฆษณา

ระดับแร็ก (Rack; 4x4x4=64 ชิป)

  • 1 แร็กมีชิป TPU 64 ตัว เชื่อมต่อกันด้วย 3D torus (4x4x4) ผ่าน ICI และ OCS (Optical Circuit Switching)
  • ที่ Google มีการแยกความหมายของ rack, pod และ slice
    • rack: หน่วยกายภาพขนาด 64 ชิป (=ลูกบาศก์)
    • pod: หน่วยสูงสุดที่เชื่อมต่อได้ด้วย ICI และ OCS (เช่น TPUv4=4096 ชิป=64 rack)
    • slice: หน่วยนามธรรมที่จัดรูปแบบได้ตามต้องการตั้งแต่ 4 ชิปไปจนถึง Superpod

ข้อดีของ OCS

  1. wraparound : ทำให้แต่ละแกนเป็นวงแหวน (1D torus) เพื่อลดจำนวน hop สูงสุดระหว่างโหนด
  2. slice แบบ non-contiguous ที่ยืดหยุ่น: ด้วยโครงสร้างการสวิตช์ของ OCS แม้โหนดจะอยู่ห่างกันทางกายภาพก็รวมเป็น slice เดียวกันได้ ทำให้ใช้ทรัพยากรและบำรุงรักษาได้ง่าย
  3. topology แบบบิด (twisted topology): แม้มีจำนวนชิปเท่าเดิม (x,y,z คงที่) ก็สามารถเปลี่ยนโครงสร้างการเชื่อมต่อเพื่อปรับความเร็วให้เหมาะกับรูปแบบการคำนวณเฉพาะได้ (เช่น twisted torus)

ตัวอย่างการใช้งาน topology

  • แบบลูกบาศก์: เหมาะกับ data/tensor parallel (แบนด์วิดท์สูงสุด)
  • แบบเส้นตรง (ทรงซิการ์) : เหมาะกับ pipeline parallel
  • twisted torus: ช่วยเพิ่มความเร็วเมื่อต้องสื่อสารแบบ all-to-all (เช่น tensor parallel)
โฆษณา

ระดับซูเปอร์พอด (Full Pod/Superpod; TPUv4: 4096 ชิป/64 แร็ก)

  • เชื่อมต่อหลายแร็กเข้าด้วยกันผ่าน ICI และ OCS เพื่อสร้างระบบขนาดมหึมา
  • สมรรถนะด้านแบนด์วิดท์การสื่อสาร/การประมวลผลขนานจะแตกต่างกันตามชนิดของ slice topology
  • ด้วย OCS จึงได้ความยืดหยุ่นทั้งในด้าน slice แบบ non-contiguous และ topology แบบ twisted

ระดับ multi-pod (Multi-pod/Multislice; TPUv4: มากกว่า 4096 ชิป)

  • เชื่อมต่อหลาย pod ผ่าน เครือข่ายดาต้าเซ็นเตอร์ (DCN) เพื่อสร้างโครงสร้างพื้นฐานการฝึกขนาดใหญ่ แต่แบนด์วิดท์ต่ำกว่า ICI
  • การฝึก PaLM ใช้ 2 pod (TPUv4 จำนวน 6144 ตัว) และบริหารทรัพยากรรวม 6 pod
  • ในการพัฒนาโมเดลขนาดใหญ่ XLA compiler จะปรับรูปแบบการสื่อสารให้เหมาะสมตามค่าที่นักวิจัยตั้งไว้ (เช่น มิติของ parallelism)
  • XLA จะแทรกงานสื่อสารระหว่าง slice และ pod แต่ละส่วน ทำให้สร้างการฝึกแบบกระจายขนาดใหญ่ได้โดยแก้โค้ดน้อยที่สุด

ตัวอย่างฮาร์ดแวร์จริงและแผนภาพ

  • TPU rack: 3D torus แบบ 4x4x4 เป็นหนึ่งยูนิต โดยแต่ละแถวมี 2 tray (8 ชิป)
  • TPUv4 tray: ในฮาร์ดแวร์จริงมี PCIe 4 พอร์ต (หนึ่งพอร์ตต่อ TPU หนึ่งตัว)
  • TPUv4 chip: มี ASIC อยู่ตรงกลาง และมี HBM stack 4 ชุดล้อมรอบ (อิงตาม 2 TensorCore)
  • TPUv4i (สำหรับงานอนุมาน) floorplan ของชิป: มี 1 TensorCore และใช้พื้นที่ CMEM ขนาดใหญ่

ปิดท้าย

  • ขอขอบคุณ Google TPU Research Cloud(TRC) สำหรับการสนับสนุนงานวิจัย

เอกสารอ้างอิง

บทความที่เกี่ยวข้อง

ยังไม่มีความคิดเห็น

ยังไม่มีความคิดเห็น