เมื่อโลกของ AI Inference เปลี่ยนจาก GPU สู่ ASIC

ช่วงไม่กี่ปีที่ผ่านมา เราปฏิเสธไม่ได้เลยว่า AI มันมาจริง ทุกคนใช้งาน AI กันฉ่ำเกิน จนทำให้แสงส่องไปที่สถาปัตยกรรมของ GPU ซึ่งเปรียบได้กับราชาในของของ AI เพราะความสามารถในการทำ Parallel Processing ได้มหาศาล แต่เมื่อ Model นั้นถูก Train จนเสร็จและพร้อมใช้งานในสเกลใหญ่มาก ๆ กฏเกณฑ์ของเกมการประมวลผลก็เริ่มเปลี่ยนไป ผู้ให้บริการหลายเจ้า เริ่มย้ายจากการใช้ GPU แบบเดิม ๆ และหันมาทุ่มเทกับการสร้าง ASIC วันนี้เราจะมาเข้าใจ และการจัดการปัญหานี้ของผู้ให้บริการกันว่า เขาทำอย่างไร

ปัญหาของ GPU กับคอขวดในโลกของการทำ Inference

ในการ Train Model การทำ Parallel Processing ที่ตัวเลือกที่ดีที่สุดในการทำงาน แต่เมื่อเราเข้าถึงการนำมา Inference เกมเปลี่ยนทันที โดยเฉพาะกับการใช้งาน LLM ที่มีขนาดใหญ่กว่า Model ที่เราเคยใช้งานกันมาก่อนหน้านี้

การทำงานของ LLM มันจะแบ่งเป็น 2 ช่วงหลัก ช่วงแรกคือการทำความเข้าใจคำสั่งเริ่มต้น (Prefill) ซึ่งเป็นงานที่เน้นพลังในการประมวลผล (Compute-Bound) เรียกได้ว่าเป็นงานถนัดของ GPU ที่มีจำนวน Core ให้ทำ Parallel Processing ฉ่ำ ๆ และช่วงที่ 2 คือ ช่วงของการสร้างคำตอบ (Auto-Regressive Decode) งานนี้มันจะไปเน้นที่ความเร็ว และความกว้างของช่องทางการส่งข้อมูลแล้ว เพราะช่วงนี้เป็นจุดบอดสำคัญที่ ALU จะต้องนั่งรอคิว เพื่อดึงข้อมูล Model Weight จาก GPU Memory ตลอดเวลา อาการนี้ในทาง Computer Architecture เราเรียกว่า Memory Wall ส่งผลให้ ประสิทธิภาพการคำนวณตกลงไปอย่างน่าตกใจ

นอกจาก Memory Wall แล้ว อีกเรื่องที่เป็นเรื่องใหญ่คือ การออกแบบ GPU มาเพื่อให้มันเก่งในงาน SIMT (Single Instruction, Multiple Threads) ทำให้พื้นที่บน Silicon ต้องเสียไปกับชุดควบคุม และ Register ขนาดใหญ่ ซึ่งโครงสร้างพวกนี้แทบไม่มีความจำเป็นเลยสำหรับงานในการทำ Model Inference แต่จริง ๆ แล้วมันแค่การเอา Matrix มาบวก หรือคูณกันแค่นั้นเลย ผลท่ีเกิดขึ้นคือ เราเสียพลังงานไปกับ Function ที่ไม่ได้ใช้งาน และสร้างคอขวดเพิ่มเติมในการส่งข้อมูลผ่าน Bus ในการเชื่อมต่อ ระหว่าง Main Memory กับ GPU Memory

Tensor Core ปะทะ MXU

เพื่อให้เห็นภาพความแตกต่างที่ชัดเจน และเข้าใจว่าทำไม Cloud Provider ถึงยอมลงทุนสร้าง Chip ใหม่ เราจะต้องเจาะลงลึกเข้าไปดูที่ Microarchitecture ในฝั่งของ GPU สมัยใหม่ เช่นของ Nvidia ตั้งแต่ตระกูล RTX ที่จะมีหน่วยประมวลผลพิเศษที่เรียกว่า Tensor Core หน้าที่หลักของมันคือ คูณและเก็บค่า หรือที่เราเรียกว่า Matrix Multiply-Accumulate หรือ MMA หรือ ๆๆๆ อธิบายเป็นสมการง่าย ๆ ว่า D = A x B + C ให้จบใน 1 Clock Cycle เทียบกับ CPU ที่ไม่มีวงจรแยกก็จะต้องใช้หลาย Clock Cycle ในการทำงาน ทำให้ Tensor Core มันเร็ว และประสิทธิภาพสูงกว่า การใช้ CPU และ GPU ทำงานมหาศาล ปัญหาคือ มันยังต้องทำงานด้วยสถาปัตยกรรมเดิม แบบที่ GPU เป็น ข้อมูลทุกอย่างยังต้องวิ่งผ่านกระบวนการ Instruction Fetch และอ่านเขียนข้อมูลเข้าออก Register ตลอดเวลา ซึ่งกระบวนการพวกนี้สร้าง Overhead ให้กับระบบ อย่างหลีกเลี่ยงไม่ได้

ในทางกลับกัน การใช้ ASIC ที่มี MXU (Matrix Multiply Unit) นั้นมีวิธีการคิดที่แตกต่างออกไปอย่างชัดเจน โดยมีหัวใจสำคัญคือการวางโครงสร้างแบบ Systolic Array และ Dataflow Architecture แทนที่จะดึงข้อมูลเข้าออกหน่วยความจำตลอดเวลา เราก็โหลดข้อมูลพวกนี้ไว้ที่ MAC (Multiply-Accumulate) แต่ละตัวไว้ล่วงหน้า จากนั้น Input ก็จะป้อนเข้าไป และเริ่มคำนวณได้เลย โดยไม่ต้องเสียเวลาโหลด Weight ใหม่ทุกครั้ง ข้อมูลที่ทำงานเสร็จแล้ว ก็จะถูกส่งไปที่ Block ต่อไปที่มันเชื่อมต่อกันด้วย Hardware โดยตรง โดยไม่ต้องวิ่งกลับไปที่ Cache Memory การออกแบบที่ตัดกระบวนการส่วนเกินนี้จะทำให้เราได้ประสิทธิภาพในการทำงานที่สูงขึ้นแบบก้าวกระโดด โคตร ๆ

ทำไม Cloud Provider ถึงเริ่มหันไปใช้ ASIC กันมากขึ้น

เพราะประสิทธิภาพในการทำงานนี่เอง เป็นหนึ่งในแรงพลักดันให้ผู้ให้บริการ Cloud ลงมาออกแบบ Silicon กันเอง เช่น Google ที่เป็นเหมือนผู้บุกเบิกด้วยการสร้าง TPU (Tensor Processing Unit) ขึ้นมาเอง เพื่อรองรับบริการทั้งหมดของบริษัท โดยรุ่นล่าสุดอยาง V5e ไปเน้นเรื่องความคุ้มค่าในการทำงาน พร้อมทั้งพัฒนา Interconnect เพื่อให้รองรับการรัน Model ขนาดใหญ่ กระจายไปตามเครื่องได้มีประสิทธิภาพมากขึ้น ลดเวลาหน่วง และเพิ่มความเสถียรให้กับการทำงาน

หรือในฝั่ง AWS ก็มี Chip ในตระกูล Inferentia ที่เดินทางมาถึงรุ่นที่ 2 แล้ว ซึ่งรุ่นนี้ออกแบบมาเพื่อแก้ปัญหาคอขวดของ Memory Bandwidth โดยเฉพาะ มีการใช้ Ring Topology เพื่อเชื่อมแต่ละ Chip เข้าด้วยกัน ช่วยให้การทำ Pipeline Parallelism ทำได้เร็ว และที่สำคัญคือ หั่นต้นทุนได้มหาศาล

หรือแม้แต่ Microsoft เอง ก็มีการสร้าง Azure Maia 100 และ Mata ที่เป็น MTIA (Meta Training and Inference Accelerator) เพื่อนำมารันระบบ Recommendation System ของตัวเองโดยไม่ต้องจ่ายค่า Hardware ส่วนเกินที่ไม่ได้ใช้งาน

เราจะเห็นว่า พอ AI Inference มันเยอะมากขึ้นเรื่อย ๆ ความต้องการในการ Optimise ก็มีเพิ่มมากขึ้นเรื่อย ๆ เช่นเดียวกัน จากเดิม แค่ Optimise วิธีการทำงานบน Chip เดิม สู่การออกแบบ Custom Silicon สำหรับการทำงานเฉพาะด้านมากขึ้นเรื่อย ๆ เพิ่มความหนาแน่นใน Data Centre อย่างก้าวกระโดด

อนาคตของการทำ AI Inference

การเติบโตอย่างรวดเร็วของ ASIC บน AI Inference ไม่ได้ทำให้ GPU ตายจากไป แต่เราคิดว่า มันน่าจะทำให้เราเข้าสู่ยุคของการใช้อุปกรณ์แบบผสม (Heterogeneous Computation) ซึ่งจะเป็นการรวมกันของ Cluster ขนาดใหญ่ และ Edge Computing โดยมี NPU ทำงานอยู่เบื้องหลัง

ในปัจจุบัน เวลาเราใช้งาน Computer รุ่นใหม่ ๆ ไม่ว่าจะเป็น Apple Silicon หรือฝั่ง PC เราเริ่มจะเห็นประโยชน์ของ Unified Memory Architecture ที่ให้ทั้ง CPU และ NPU สามารถเข้าถึง Model Weight ได้จาก Memory เพียงก้อนเดียวแบบไม่ต้องคัดลอก (Zero-Copy) งานพื้นฐานอย่างการรัน Model ภาษาขนาดเล็ก อาจจะทำงานจบได้ในเครื่องของเราเอง ส่วนงานหนักที่ต้องการความซับซ้อนในการคิดวิเคราะห์ ก็อาจจะตัดเส้นทางไปใช้ ASIC บน Cloud

ซึ่งส่วนตัวเราคิดว่า มันน่าจะเป็นอนาคตของการทำ AI Inference ในเวลาอันใกล้นี้แล้ว