เทรนโมเดลให้เร็วทันใจด้วย Distributed Training บน Vertex AI

เคยไหมครับที่ต้องเทรนโมเดลใหญ่ ๆ อย่าง Large Language Model หรือโมเดลเกี่ยวกับรูปภาพที่ใช้ข้อมูลมหาศาล แล้วรู้สึกว่าต้องรอนานจนดริปกาแฟเสร็จไปหลายแก้วก็ยังไม่เสร็จสักที นั่นแหละครับคือปัญหาคลาสสิกของ AI Engineer ทั่วโลก เพราะการเทรนโมเดลบนเครื่องเดียวอาจจะช้าเกินไป นั่นแหละครับที่ Distributed Training เกิดมาเพื่อแก้ปัญหานี้

Distributed Training คืออะไร ? 

อธิบายง่ายๆ Distributed Training คือการ แบ่งงานเทรนโมเดลออกไปให้คอมพิวเตอร์หลาย ๆ เครื่องช่วยกันทำพร้อม ๆ กัน แทนที่จะให้เครื่องเดียวทำงานหนักคนเดียว เหมือนการจ้างทีมงานหลายคนมาช่วยกันทำโปรเจกต์ใหญ่ ๆ ให้เสร็จเร็วขึ้นนั่นเองครับ

พอเราเข้าใจแนวคิดของ Distributed Training แล้ว คำถามถัดมาคือ “แล้วจะเริ่มโค้ดยังไง ?” ซึ่งคนที่ใช้ TensorFlow ก็ไม่ต้องกังวลเลย เพราะ TensorFlow มีฟีเจอร์สำหรับการเทรนแบบกระจายตัวมาให้อยู่แล้ว เรียกว่า Strategy ซึ่งก็มีหลัก ๆ 3 แบบ แล้วแต่การใช้งานดังนี้ครับ

1. MirroredStrategy: เหมาะสำหรับการเทรนบนเครื่องเดียว แต่มีหลาย GPU
2. MultiWorkerMirroredStrategy: เหมาะกับเทรนกับหลายเครื่อง แต่ละเครื่องก็มี GPU หลายตัว ซึ่งเป็น Strategy ยอดนิยมที่เราจะพูดถึงครับ
3. TPUStrategy: สำหรับการเทรนบนโดยใช้ TPU บน Cloud TPU โดยเฉพาะ

แต่ TPU อาจไม่ได้เหมาะกับทุกโจทย์เสมอไปครับ

เหมือนคนบางคนแม้จะดีเพียบพร้อมขนาดไหน ก็อาจไม่ได้เป็นคนรักของเราครับ เพราะคนดีที่กับคนที่รักอาจจะเป็นคนละคนกัน ดังนั้น TPU หรือ Tensor Processing Unit ที่ออกแบบมาเพื่อเทรน Machine Learning โดยเฉพาะและขึ้นชื่อเรื่องเร็วมาก แต่ก็มีบางเคสที่อาจจะไม่ใช่คนรักที่ดีที่สุดครับ

โดยทั่วไปแล้ว TPU จะใช้การคำนวณแบบ bfloat16 ซึ่งย่อมาจาก Brain Floating Point เป็นรูปแบบตัวเลขแบบ 16-bit ที่จุดเด่นคือช่วงตัวเลข หรือ range กว้างเท่ากับ 32-bit แต่ใช้พื้นที่และพลังงานน้อยกว่า ทำให้สามารถคำนวณแบบ Matrix ได้อย่างรวดเร็วและประหยัด ซึ่งแลกมากับความละเอียดทศนิยมที่น้อยกว่า fp64 หรือ Floating Point 64 ที่ช่วงกว้าง 64-bit ที่ใช้กันใน GPU ทั่วไป

ดังนั้น หากงานของเราต้องอาศัยการคำนวณที่มีความแม่นยำสูงมาก ๆ หรือ high-precision arithmetic ที่ต้องการทศนิยมละเอียด ๆ อย่างพวกงานทางการแพทย์ หรือ วิทยาศาสตร์ต่าง ๆ การใช้ TPUStrategy อาจทำให้ได้ผลลัพธ์ที่ไม่แม่นยำเท่าที่ควร บางกรณีการใช้ MultiWorkerMirroredStrategy ที่รันบน GPU อาจรองรับการคำนวณที่แม่นยำกว่าอย่าง fp64 ทำให้ผลลัพธ์ที่ดีกว่า หรือแม้กระทั่งเร็วกว่าครับ

Distributed Training บน Vertex AI เป็นยังไง ?

สำหรับบน Vertex AI นั้น การทำ Distributed Training จะมีลักษณะเป็น Cluster หรือกลุ่มของคอมพิวเตอร์หลาย ๆ เครื่อง โดยแต่ละเครื่องเรียกว่า Worker แต่จะมี Worker ตัวหนึ่งที่เป็น Chief (หรือ Master เป็นหัวหน้าทีม) คอยดูแลและควบคุมการเทรน

Worker แต่ละตัวจะได้รับข้อมูลส่วนหนึ่งไปประมวลผลและคำนวณ Gradient (ค่าที่บอกว่าต้องปรับเปลี่ยนพารามิเตอร์โมเดลไปในทิศทางไหน) จากนั้นข้อมูล Gradient เหล่านี้จะถูกส่งกลับมาเพื่อรวมเข้าด้วยกันและอัปเดตน้ำหนักของโมเดล ซึ่งกระบวนการนี้เรียกว่า Collective Operation

หนึ่งในเทคนิคที่นิยมใช้กันมากคือ All-reduce ซึ่งเป็นกลไกที่ช่วยให้ Worker ทุกตัวสามารถแลกเปลี่ยนและรวม Gradient เข้าด้วยกันได้อย่างมีประสิทธิภาพ จนกระทั่งทุก Worker มีข้อมูล Gradient ที่อัปเดตแล้วเหมือนกันทั้งหมด แต่การสื่อสารระหว่าง Worker จำนวนมาก ๆ ก็อาจจะกลายเป็นคอขวด ทำให้การเทรนช้าลงได้เหมือนกันครับ

สมมติเรามี GPU 4 ตัวเป็น 4 Worker ซึ่งแต่ละตัวก็ได้มอบหมายงานคำนวณ Gradient จากข้อมูลที่ตัวเองดูแลอยู่ แต่ไม่รู้ Gradient ของคนอื่น จากนั้นก็เข้าสู่ขั้นตอน All-reduce ที่ Woker แต่ละตัวเริ่มคุยกันแลกเปลี่ยน Gradient กับ Worker ตัวอื่น ๆ จนทุกตัวก็จะมี Gradient ที่รวมกันแล้วเหมือนกันหมด นั่นก็คือ A+B+C+D พูดง่ายๆ ก็เหมือนกับทุกตัวอัพเดทค่า Weight ครับ

เมื่อการคุยกันกลายเป็นคอขวด Reduction Server จึงเข้ามา

อย่างที่บอกไปว่าการที่ Worker ทุกตัวต้องคุยกันมากไปเพื่อแลกเปลี่ยน Gradient มันเกิด Latency (ความหน่วง) และ Bottleneck (คอขวด) โดยเฉพาะอย่างยิ่งถ้าใช้ GPUs จำนวนมาก ๆ ๆ ๆ ในการเทรน

นี่คือจุดที่ผมจะมาเล่าให้ฟัง Reduction Server เข้ามาช่วยอย่างไร ?

Reduction Server คือเซิร์ฟเวอร์เฉพาะทางที่ออกแบบมาเพื่อจัดการงาน รวม Gradient แทนที่ Worker ทุกตัวจะต้องคุยกันเองเพื่อทำ All-reduce เปลี่ยนเป็น Worker ส่งข้อมูล Gradient ไปที่ Reduction Server แทน และให้ Reduction Server ทำหน้าที่รวมข้อมูลทั้งหมด แล้วจึงส่งข้อมูลที่รวมแล้วกลับไปให้แต่ละ Worker อีกครั้ง

ถ้าเทียบ All-reduce แบบดั้งเดิมอย่าง Ring All-reduce และ Tree All-reduce จะเห็นว่า Data Transfer ที่จะส่งกว่าจะครอบทุกตัวก็ 2(n-1)/n ใช่ไหมครับ ขณะที่พอใช้ Reduction Server ก็ส่งแค่ครั้งเดียว

ในมุมของ Latency ก็ได้ O(1) ใครเรียนวิชา Algorithm มาจะรู้ว่านี่คือ BigO ที่ดีที่สุด ทำงานแค่ 1 ครั้ง ซึ่งหมายความว่า Latency ไม่ได้เพิ่มขึ้นเลยแม้จะมี Worker เพิ่มขึ้นมากี่ node ก็ตาม แต่ถ้า O(n) แปลว่า ยิ่ง Woker มากยิ่งหน่วยสูง สูงแบบ linear น่ะครับ

Reduction Server ทำหน้าที่เป็นศูนย์กลางในการรวม Gradient แต่สำหรับโมเดลขนาดใหญ่มาก ๆ การใช้ Reducer แค่ตัวเดียวก็อาจทำงานไม่ทัน การออกแบบระบบของ Google Cloud จึงทำมาเพื่อรองรับการใช้งาน Reducer หลายตัวพร้อมกัน ลองดูภาพนี้ครับ

ภาพนี้แสดงให้เห็นว่า เราไม่ได้มี Reducer แค่ตัวเดียวอีกต่อไป แต่มีหลายตัว (ในภาพมี 5 ตัว) ซึ่งแต่ละตัวไม่ได้ทำงานร่วมกันแบบสุ่ม ๆ แต่มีการแบ่งงานกันอย่างเป็นระบบครับ

หลักการทำงานคือ

1. Partitioning: Gradient Array ที่อยู่บน Worker แต่ละตัวจะถูกแบ่งออกเป็นหลายส่วน (ในภาพแบ่งเป็นส่วน a, b, c, d, e) จำนวนส่วนจะเท่ากับจำนวนของ Reducer ที่เรามี
2. ส่ง Gradient แบบเจาะจง: Worker แต่ละตัวจะส่งข้อมูล Gradient ส่วนของตัวเองไปให้ Reducer ที่รับผิดชอบส่วนนั้นโดยเฉพาะ เช่น Reducer ตัวที่ 1 จะรับผิดชอบข้อมูลส่วน a จากทุก Worker เท่านั้น เช่นเดียวกับ Reducer ตัวที่ 2 ก็จะรับข้อมูลจากส่วน b ของแต่ละ Worker เช่นกัน
3. รวมข้อมูลกันแบบขนาน: Reducer แต่ละตัวจะทำหน้าที่รวมข้อมูล Gradient ของส่วนที่ตัวเองได้รับมาพร้อม ๆ กัน
4. ส่งผลลัพธ์กลับไป: เมื่อรวมเสร็จแล้ว Reducer แต่ละตัวก็จะส่งข้อมูลที่รวมแล้วกลับไปให้ Worker ทุกตัว

วิธีนี้ทำให้การรวม Gradient ทั้งก้อนสามารถทำได้อย่างรวดเร็วแบบกันครับ ทำให้ไม่ต้องรอ Reducer ตัวเดียวทำงานจนเสร็จ

แล้ว Reduction Server ช่วยอะไรได้บ้าง ?

• ลด Latency และ Bandwidth: เมื่อ Worker แต่ละตัวไม่ต้องคุยกันเองทั้งหมด ก็จะทำงานเร็วขึ้น ใช้ Bandwidth เช่น Network คุ้มค่าขึ้น
• ลดการสื่อสารที่ซับซ้อน: All-reduce ในระบบที่มีหลาย Workers อาจจะยุ่งยาก ต้อง Setup แต่ละ Node ให้รู้จักกัน แต่ Reduction Server จะช่วยจัดการส่วนนี้แทนเรา ไม่ต้องเขียนโค้ดเยอะ
• เพิ่มประสิทธิภาพการเทรน: พอไม่มีคอขวด โมเดลก็เทรนได้เร็วขึ้นอย่างมาก ซึ่งหมายความว่าเราจะประหยัดเวลาและค่าใช้จ่ายในการเทรนโมเดลที่ใหญ่และซับซ้อนครับ

ผลลัพธ์ที่พิสูจน์แล้ว: เร็วขึ้น 45% กับ Reduction Server!

ทีมงานของ Google ได้ทำการทดลองเพื่อเปรียบเทียบประสิทธิภาพการเทรนของ PyTorch Distributed Training ระหว่างการใช้ Reduction Server กับการใช้เทคนิค NCCL All-reduce (NVIDIA Collective Communications Library) โดยการทดลองนี้เป็นการ Fine-tune โมเดลภาษาขนาดใหญ่อย่าง BERT large บน Dataset สำหรับทำ Sentiment Classification อย่าง IMDb ซึ่งเป็นโมเดลและ Dataset ที่นิยมใช้กันอย่างแพร่หลายครับ ได้ผลลัพธ์ดังนี้

• เวลาในการเทรนลดลงมากกว่า 45%: จากที่เคยต้องใช้เวลาเทรนเกือบ 5,000 วินาที ก็ลดลงเหลือเพียงประมาณ 2,600 วินาทีเท่านั้น ซึ่งหมายถึงการประหยัดเวลาและค่าใช้จ่ายในการเทรนไปได้เกือบครึ่งเลยทีเดียว
• อัตรา Throughput เพิ่มขึ้นถึง 80%: Throughput หรืออัตราความเร็วในการเทรน (จำนวน Train steps ต่อวินาที) เพิ่มขึ้นจาก 0.54 เป็น 0.989 อย่างเห็นได้ชัด ซึ่งสะท้อนให้เห็นว่าการเทรนมีความเร็วและมีประสิทธิภาพสูงขึ้นมาก

อีกผลการทดลองที่ทาง Google ทำไว้เพื่อ Fine-tune โมเดลภาษาอย่าง BERT บน Dataset MNLI ก็ยืนยันข้อนี้อย่างชัดเจนครับ

ทางทีมได้ทำการทดลองด้วยเครื่องมือที่เหมือนกันทุกอย่าง ไม่ว่าจะเป็นการใช้ Worker node 8 ตัว แต่ละตัวมี 8 NVIDIA A100 GPUs เหมือนกันหมด โดยมีเพียงสิ่งเดียวที่แตกต่างคือสีน้ำเงินใช้ NCCL All-reduce ส่วนสีแดงใช้ Reduction Server โดยเพิ่ม Reducer node เข้าไปอีก 20 ตัวครับ

แล้วผลที่ได้ … ก็คือ ฝั่งที่ใช้ Reduction Server จะมีค่าใช้จ่ายส่วนของ Reducer เพิ่มขึ้นมาอีกประมาณ $11.34 ต่อชั่วโมง แต่ผมอยากให้ดู cost ฝั่ง Training Step ครับ

• NCCL All-reduce มีค่าใช้จ่ายอยู่ที่ $0.17 ต่อ Training Step
• ขณะที่การใช้ Reduction Server สามารถลดค่าใช้จ่ายลงเหลือเพียง $0.10 ต่อ Training Step เท่านั้น

การลดลงเกือบ 40% ของต้นทุนต่อ Step นี้ ถือเป็นตัวเลขที่น่าประทับใจมาก ๆ ครับ เพราะมันหมายถึงการประหยัดค่าใช้จ่ายรวมในการเทรนโมเดลขนาดใหญ่ได้อย่างมหาศาล

อย่างไรก็ตาม ในโลกธุรกิจจริงแล้ว การลดต้นทุนอาจไม่ใช่สิ่งสำคัญที่สุดเสมอไปครับ เวลาก็สำคัญมาก การที่ Reduction Server ช่วยให้เราสามารถเทรนโมเดลให้เสร็จได้เร็วขึ้น อาจดีกับธุรกิจมากกว่า เช่น ช่วยให้เราสามารถนำโมเดลออกสู่ตลาดได้เร็วขึ้น ขายเร็วขึ้น หรือช่วยทำโมเดลได้ตามแผน ซึ่งสิ่งเหล่านี้ประเมินมูลค่าเป็นเงินไม่ได้เลยครับ

สรุป

การเทรนโมเดลขนาดใหญ่ให้สำเร็จได้อย่างรวดเร็วและมีประสิทธิภาพไม่ใช่เรื่องง่าย แต่ด้วย Distributed Training บน Vertex AI และเทคโนโลยีอย่าง Reduction Server ที่เข้ามาช่วยลดภาระด้านการสื่อสาร ทำให้เราสามารถโฟกัสไปที่การสร้าง AI เจ๋ง ๆ ครับ