ความก้าวหน้าของเทคโนโลยีแมชชีนเลิร์นนิงเพื่อแมชชีนเลิร์นนิง

• ยิ่งโมเดล ML มีขนาดใหญ่ขึ้น ประสิทธิภาพการรันก็ไม่ได้ขึ้นอยู่กับฮาร์ดแวร์เท่านั้น แต่ยังขึ้นอยู่กับ การปรับแต่งคอมไพเลอร์ อย่างมากด้วย โดย Google Research และ Google DeepMind เปิดเผยแนวทางที่นำ ML กลับมาปรับปรุงด้านนี้
• ภายใต้เฟรมเวิร์กอย่าง TensorFlow, JAX, PyTorch การตัดสินใจของคอมไพเลอร์ สามารถเปลี่ยนเวลาในการรันและการใช้ทรัพยากรของโมเดลเดียวกันได้อย่างมาก
• TpuGraphs คือชุดข้อมูลที่รวบรวมกราฟการคำนวณของโปรแกรม ML สำหรับ TPU, การตั้งค่าคอมไพเลอร์ และเวลาในการรัน เพื่อใช้ในการวิจัยโมเดลต้นทุนแบบอิงการเรียนรู้
• เพื่อลดข้อจำกัดที่ทำให้การเรียนรู้จากกราฟการคำนวณขนาดใหญ่ในคราวเดียวเป็นเรื่องยาก Graph Segment Training จะแบ่งกราฟออกเป็นเซกเมนต์ ลดการใช้หน่วยความจำ และย่นเวลาการฝึกลง 3 เท่า
• ในการแข่งขัน Kaggle มีผู้เข้าร่วม 792 คนจาก 616 ทีมใน 66 ประเทศ ซึ่งช่วยตรวจสอบเทคนิคปรับปรุงโมเดลทำนายต้นทุนจริง เช่น การบีบอัดกราฟ การปรับค่า padding การเพิ่มคุณลักษณะของโหนด และ attention ระหว่างการตั้งค่า

---

เหตุผลที่คอมไพเลอร์ ML เป็นตัวกำหนดประสิทธิภาพการรัน

• โมเดล ML สมัยใหม่ทำงานอย่างการเข้าใจภาษาธรรมชาติ การสนทนา การสร้างภาพ และการสร้างวิดีโอ และถูกเขียน·ฝึกด้วย เฟรมเวิร์กการเขียนโปรแกรม ML เช่น TensorFlow, JAX, PyTorch
• เฟรมเวิร์กมีการดำเนินการพีชคณิตเชิงเส้น เช่น การคูณเมทริกซ์และคอนโวลูชัน รวมถึงเลเยอร์โครงข่ายประสาท เช่น เลเยอร์คอนโวลูชัน 2D และเลเยอร์ทรานส์ฟอร์เมอร์
• ผู้ใช้ไม่จำเป็นต้องจัดการรายละเอียดวิธีรันโมเดลบนฮาร์ดแวร์อย่างมีประสิทธิภาพด้วยตนเอง เพราะ คอมไพเลอร์ ที่อยู่ใต้เฟรมเวิร์กจะปรับแต่งโมเดลโดยอัตโนมัติ
• อย่างไรก็ตาม คอมไพเลอร์มักแก้ปัญหาการปรับแต่งที่ซับซ้อนด้วยฮิวริสติก จึงอาจไม่ให้ประสิทธิภาพที่เหมาะสมที่สุดเสมอไป

กราฟการคำนวณและการปรับแต่งสองระดับ

• คอมไพเลอร์ ML แปลงคำสั่งทางคณิตศาสตร์ที่ผู้ใช้เขียนให้เป็นคำสั่งที่รันบนฮาร์ดแวร์จริงได้
• โปรแกรม ML สามารถแทนได้ด้วย กราฟการคำนวณ
  • โหนดแทนการดำเนินการกับเทนเซอร์ เช่น matrix multiplication
  • เอดจ์แทนเทนเซอร์ที่ไหลจากโหนดหนึ่งไปยังอีกโหนดหนึ่ง
• การปรับแต่งคอมไพเลอร์แบ่งได้กว้าง ๆ เป็นสองประเภท
  • การปรับแต่งระดับกราฟ: ตัดสินใจโดยคำนึงถึงบริบทของกราฟทั้งก้อน และแปลงกราฟทั้งหมด
  • การปรับแต่งระดับเคอร์เนล: แปลงเคอร์เนลหนึ่งตัวซึ่งเป็น fused subgraph โดยเป็นอิสระจากเคอร์เนลอื่น

trade-off ด้านประสิทธิภาพของ memory layout

• เทนเซอร์ 2D เช่น เมทริกซ์ สามารถจัดเก็บในหน่วยความจำในรูปแบบ [A B C a b c] หรือ [A a B b C c] ซึ่งสอดคล้องกับเลย์เอาต์แบบ row-major และ column-major ตามลำดับ
• หนึ่งในการปรับแต่งสำคัญของคอมไพเลอร์ ML คือการกำหนด memory layout ให้กับเทนเซอร์กลางทั้งหมดของโปรแกรม
• เลย์เอาต์บางแบบอาจมีประสิทธิภาพที่สุดสำหรับการดำเนินการแต่ละรายการ แต่หากเลย์เอาต์ไม่ตรงกันระหว่าง add กับ convolution คอมไพเลอร์ต้องแทรก การดำเนินการ copy เพิ่มเติม
• ในทางกลับกัน แม้ประสิทธิภาพของการดำเนินการแต่ละรายการจะต่ำลงเล็กน้อย แต่การจัดวางที่ไม่ต้องแปลงเลย์เอาต์อาจดีกว่าสำหรับการรันโดยรวม
• ใน XLA benchmark suite พบว่าเมื่อเลือกการกำหนดค่าเลย์เอาต์ที่เหมาะสมที่สุดแทนการตั้งค่าคอมไพเลอร์พื้นฐาน ความเร็วเพิ่มขึ้นสูงสุด 32%

ชุดข้อมูล TpuGraphs

• TpuGraphs คือชุดข้อมูลสำหรับ โมเดลต้นทุนแบบอิงการเรียนรู้ สำหรับโปรแกรมที่รันบน TPU แบบปรับแต่งเองของ Google
• เป้าหมายคือการฝึกโมเดลต้นทุนที่รับโปรแกรมอินพุตและการตั้งค่าคอมไพเลอร์ แล้วทำนายเวลาในการรันของโปรแกรม
• ชุดข้อมูลนี้มุ่งเป้าไปที่การตั้งค่าคอมไพเลอร์ XLA สองชนิด
  • layout: การตั้งค่าที่ขยายแนวคิด row-major·column-major ของเมทริกซ์ไปสู่เทนเซอร์มิติสูง
  • tiling: การตั้งค่าขนาดไทล์
• แต่ละตัวอย่างประกอบด้วยกราฟการคำนวณของเวิร์กโหลด ML, การตั้งค่าคอมไพเลอร์ และเวลาในการรันเมื่อคอมไพล์ด้วยการตั้งค่านั้น
• กราฟถูกรวบรวมจากโปรแกรม ML แบบโอเพนซอร์ส และรวมสถาปัตยกรรมโมเดลอย่าง ResNet, EfficientNet, Mask R-CNN, Transformer
• วิธีดาวน์โหลดและโค้ดเริ่มต้นมีให้ที่ TpuGraphs GitHub
• TpuGraphs มีจำนวนกราฟมากกว่าชุดข้อมูล graph property prediction ที่ใหญ่ที่สุดเดิมซึ่งมีขนาดกราฟใกล้เคียงกัน 25 เท่า และมีขนาดกราฟเฉลี่ยใหญ่กว่าชุดข้อมูลทำนายประสิทธิภาพโปรแกรม ML เดิม 770 เท่า

โมเดลต้นทุนอ้างอิงและโครงสร้าง GNN

• TpuGraphs มีโมเดลต้นทุนแบบอิงการเรียนรู้อ้างอิงให้ด้วย และเนื่องจากโปรแกรมอินพุตถูกแทนเป็นกราฟ จึงใช้ GNN
• คุณลักษณะของโหนดประกอบด้วยสองส่วน
  • opcode id: ข้อมูลโหนดที่สำคัญที่สุด ซึ่งแสดงชนิดของการดำเนินการกับเทนเซอร์
  • คุณลักษณะโหนดอื่น ๆ
• โมเดลอ้างอิงแปลง opcode id เป็น opcode embedding ผ่าน embedding lookup table
• รวม opcode embedding กับคุณลักษณะโหนดที่เหลือ แล้วใช้เป็นอินพุตของ GNN
• node embedding ที่ GNN สร้างขึ้นจะถูกรวมเป็น graph embedding ขนาดคงที่ด้วย graph pooling reduction แบบง่าย ๆ เช่น sum และ mean
• graph embedding สุดท้ายถูกส่งผ่าน feedforward layer แล้วแปลงเป็นเอาต์พุต scalar หนึ่งค่า

การฝึกกราฟขนาดใหญ่ด้วย Graph Segment Training

• Graph Segment Training คือ เทคนิคขยายการฝึก GNN สำหรับจัดการกราฟขนาดใหญ่บนอุปกรณ์ที่มีหน่วยความจำจำกัด
• วิธีนี้มุ่งไปที่สถานการณ์ graph-level prediction ซึ่งเป้าหมายการทำนายไม่ใช่โหนดหรือเอดจ์ แต่เป็นกราฟทั้งก้อน
• กราฟการคำนวณอาจมีโหนดหลายแสนโหนด ทำให้ Full Graph Training ที่ใช้กราฟทั้งหมดในคราวเดียวอาจเป็นไปไม่ได้ในเชิงการคำนวณ
• GST แบ่งกราฟขนาดใหญ่ออกเป็นเซกเมนต์เล็ก ๆ และเลือกเพียง subset แบบสุ่มของเซกเมนต์เพื่ออัปเดตโมเดล
• เซกเมนต์ที่เหลือจะสร้าง embedding โดยไม่เก็บ activation กลาง เพื่อลดการใช้หน่วยความจำ
• รวม embedding ของทุกเซกเมนต์เพื่อสร้าง embedding ของกราฟขนาดใหญ่เดิม และใช้สำหรับการทำนาย
• มีการนำ historical embedding table และ segment dropout มาใช้ร่วมกัน เพื่อลด staleness ของ historical embedding
• วิธีทั้งหมดช่วยลดเวลาฝึกแบบ end-to-end ลง 3 เท่า

เทคนิคปรับปรุงที่ตรวจสอบแล้วในการแข่งขัน Kaggle

• การแข่งขัน Kaggle Fast or Slow? Predict AI Model Runtime จัดขึ้นบนพื้นฐานชุดข้อมูล TpuGraphs และมีผู้เข้าร่วม 792 คนจาก 616 ทีมใน 66 ประเทศ
• มีการส่งผลงาน 10,507 รายการ โดย 153 คนเป็นการแข่งขัน Kaggle ครั้งแรก และในจำนวนนี้ 47 คนติด 100 อันดับแรก
• ทีมที่เข้าร่วมทดลองหลายเทคนิค
  • การ pruning·compression กราฟ: ทดลองวิธีบีบอัดกราฟขนาดใหญ่แทน GST โดยใช้วิธีคงไว้เฉพาะ subgraph ที่มีโหนดที่กำหนดค่าได้และเพื่อนบ้านทันทีของโหนดเหล่านั้น
  • การเปลี่ยนค่า padding: เนื่องจากค่า padding พื้นฐาน 0 ชนกับค่า feature ที่ถูกต้อง การใช้ -1 จึงช่วยปรับปรุงความแม่นยำของโมเดลได้อย่างมาก
  • การเพิ่มคุณลักษณะโหนด·การเปลี่ยนการเข้ารหัส: คุณลักษณะโหนดเพิ่มเติม เช่น contracting dimensions ของ dot general มีความสำคัญ และวิธีเข้ารหัสคุณลักษณะโหนดก็อาจส่งผลต่อผลลัพธ์
  • cross-configuration attention: ทีมชนะเลิศออกแบบเลเยอร์แบบง่ายที่ทำให้โมเดลเปรียบเทียบการตั้งค่าต่าง ๆ ได้อย่างชัดเจน และให้ผลลัพธ์ดีกว่าวิธีให้อนุมานแต่ละการตั้งค่าแยกกันอย่างมาก
• ผลการแข่งขันและโซลูชันของผู้ชนะมีกำหนดจะถูกนำเสนอในเซสชันการแข่งขันของ ML for Systems workshop ที่ NeurIPS วันที่ 16 ธันวาคม 2023

เซสชันที่เกี่ยวข้องใน NeurIPS Expo

• สำหรับผู้อ่านที่สนใจข้อมูลมีโครงสร้างและการวิจัยปัญญาประดิษฐ์ พาเนล NeurIPS Expo Graph Learning Meets Artificial Intelligence จะจัดขึ้นในวันที่ 9 ธันวาคม 2023
• พาเนลนี้จะครอบคลุมเรื่องต่าง ๆ เช่น ความก้าวหน้าของโมเดลต้นทุนแบบอิงการเรียนรู้

[GN⁺]

ความคิดเห็นใน Hacker News

• คอมไพเลอร์ ML ถูกอวยเกินจริง มันเป็นการแลกเปลี่ยนแบบเดียวกับคอมไพเลอร์ดั้งเดิม: ได้ throughput มากกว่าการจ้างโปรแกรมเมอร์สาย performance เฉพาะทางมาก แต่ฝ่ายหลังมักจะเร็วกว่าเยอะ และในบางกรณีอาจนำอยู่หลายลำดับขั้น
  ยังขาดตกบกพร่องในหลายระดับ ในระดับอัลกอริทึม มันไม่สามารถป้อนกลับเคล็ดลับที่จะทำให้เครือข่ายเร็วขึ้นให้มนุษย์ได้ ให้ได้แค่สัญญาณพื้นฐานมาก ๆ เท่านั้น เจตนาก็สูญหายไปด้วย นักออกแบบเครือข่าย ML กำหนดโครงสร้างด้วย Python แต่พอผ่านการแปลงลดระดับหลายขั้น อาจได้ผลลัพธ์ที่ผิดไปคนละเรื่อง ช่วงหลังผมเห็นคอมไพเลอร์ตัวหนึ่งทำ slice update โดยสร้างช่วงดัชนีที่เป็นไปได้ทั้งหมดของอาร์เรย์ แล้วตัดออกมาเพื่อให้ได้ดัชนีที่จะอัปเดต จากนั้นค่อย scatter ซึ่งผมแทนที่ได้ด้วยการเรียก memcpy ครั้งเดียว เคอร์เนลก็ไม่มีประสิทธิภาพด้วย ทุกครั้งที่เอา output ของคอมไพเลอร์แบบนี้ไปชนกับโปรแกรมเมอร์ assembly ที่ชำนาญ คอมไพเลอร์จะแพ้ และมักต่างกัน 30% ขึ้นไป ดูเหมือนจะเป็นปัญหาที่แก้ง่าย แต่ถ้าตลอด 50 ปีที่ผ่านมาไม่มีใครแก้ได้ดีจริง ๆ ก็ชัดเจนว่ามันไม่ได้ง่ายอย่างที่พูด

  • ถ้าดูเอนจินหมากรุก Stockfish จะเห็นว่ามันเลิกใช้ heuristic ที่มนุษย์เขียนสะสมมาหลายปีในการประเมินกระดาน แล้วแทนด้วย neural network ขนาดเล็ก ซึ่งทำได้ดีกว่า
    ในคอมไพเลอร์ก็มี heuristic เยอะ เช่น inlining, loop unrolling, vectorization ดังนั้น neural network อาจช่วยได้ และอาจดูแลรักษาง่ายกว่า heuristic จำนวนมากที่มนุษย์เขียน
  • ที่บอกว่าเป็นการแลกเปลี่ยนแบบเดียวกับคอมไพเลอร์ดั้งเดิมก็น่าสนใจ เพราะคอมไพเลอร์ดั้งเดิมเหล่านั้นมีประโยชน์มหาศาล
  • ฟังดูฟันธงและปิดกั้นเกินไป
  • ใช่เลย จ้างคนมาเขียน assembly ด้วยมือก็ได้ แล้วทำไมใครจะใช้ gcc/clang กันล่ะ?
  • Throughput นั่นแหละคือประเด็นหลัก เราไม่สามารถเอาผู้เชี่ยวชาญ performance ไปประกบทุกงาน ML ได้
    optimization แบบนี้ยังดีกว่าไม่มีเลยมาก

• ช่วยอธิบายให้เป็นภาพจริงหน่อยได้ไหม? อยากรู้ว่าสถานะจริงของ คอมไพเลอร์ ML ตอนนี้เป็นอย่างไร และในอนาคตอันใกล้เราควรคาดหวังอะไรได้บ้าง

  • หนึ่งในวิธีที่ง่ายที่สุดคือ torch.compile ซึ่งเป็น iteration ล่าสุดของคอมไพเลอร์ PyTorch ก่อนหน้านี้มีวิธีอย่าง TorchScript และ FX Tracing
    แค่เขียนว่า model = torch.compile(model) ก็ได้แล้ว “ในบรรดาโมเดลโอเพนซอร์ส 163 ตัวนี้ torch.compile ใช้งานได้ใน 93% ของกรณี และทำให้การเทรนบน NVIDIA A100 GPU เร็วขึ้น 43% โดยเฉลี่ยเร็วขึ้น 21% ที่ความแม่นยำ Float32 และ 51% ที่ความแม่นยำ AMP”[1] ดูเหมือน Google ต้องการให้มีคนเข้ามามีส่วนร่วมกับ R&D ของวิธีแบบนี้มากขึ้น
    [1] https://pytorch.org/get-started/pytorch-2.0/
  • สิ่งที่คาดหวังได้ในอนาคตอันใกล้คือการใช้ AMD, CUDA, TPU, CPU ฯลฯ ได้ แม้จะไม่มีการสนับสนุนจาก vendor อย่างชัดเจนสำหรับ framework ที่ใช้พัฒนาโมเดลนั้น
    ความจริงซับซ้อนกว่านี้ แต่ถ้าพูดแบบย่อมาก ๆ คือคอมไพล์ computation graph ให้เป็น intermediate representation บางแบบ แล้ว implement backend ที่รองรับสิ่งนั้น โปรเจกต์ที่เกี่ยวข้องดูได้ที่ stableHLO, IREE, openXLA คอมไพเลอร์ jit ของ Jax ก็ถือเป็นคอมไพเลอร์ลักษณะนี้รูปแบบหนึ่งได้ มัน lowering operations ที่ trace ได้ไปเป็น XLA แล้ว XLA ก็ใช้เวทมนตร์สารพัดให้รันบน backend ได้ สุดท้ายยิ่งลงไปชั้นล่าง ๆ ก็เป็น การแปลงและการทำ abstraction ต่อเนื่องไปเรื่อย ๆ
  • ดู torch.compile ได้เลย

• สรุปคือ งานนี้คือการปรับปรุงการทำนาย performance runtime ของ computation graph ด้วย graph neural network (GNN) โดยใช้ embedding dictionary สำหรับ opcode ของแต่ละ node ร่วมกับคุณลักษณะอื่นของ node เช่น shape, bits, window size ([1])
  เขาเผยแพร่ dataset กราฟขนาดใหญ่ที่มีการตั้งค่าคอมไพล์ XLA หลากหลายและผล performance บน TPU ไว้ใน [2] และปรับปรุงการทำนายกราฟที่ใหญ่กว่าก่อนหน้าใน [3] ด้วยวิธีแบ่งกราฟเป็นส่วน ๆ (เพิ่งเคยเห็น METIS graph partition เป็นครั้งแรก) และเทคนิคการเรียนรู้หลายแบบ เรื่องนี้เป็นการทำนาย performance ของกราฟที่กำหนดมา ไม่ใช่การปรับปรุง เสนอ หรือแก้ไขกราฟใหม่ที่เทียบเท่ากัน โมเดลที่ทำนายได้ดีพอสมควร เช่น FunSearch สามารถใช้ร่วมกับ evolutionary search ได้
  [1] https://github.com/google-research-datasets/tpu_graphs#featu...
  [2] TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs https://arxiv.org/abs/2308.13490
  [3] Learning Large Graph Property Prediction via Graph Segment Training https://arxiv.org/abs/2305.12322

• ใครช่วยอธิบายได้ไหมว่า convolution ในกราฟนั้นทำงานอย่างไร? เอา tensor shape [2,4,16] ไป convolution กับ kernel shape [4,16,8] แล้วได้ tensor [2,8] ออกมา เป็นไปได้อย่างไร?

  • ไม่รู้ว่าจะช่วยได้ไหม แต่ใน input tensor [2,4,16] ค่า 2 มองได้ว่าเป็น batch size, 4 เป็นมิติของ input features และ 16 เป็นมิติของ input channels
    ใน kernel [4,16,8] ค่า 4 คือขนาดของ filter window, 16 เป็นค่าที่ตรงกับมิติ input channels และ 8 คือมิติ output channels ใน output [2,8] ค่า 2 คงไว้เป็น batch size และ 8 ตรงกับมิติ output channels ของ kernel แม้มองเผิน ๆ เหมือนมิติจะไม่ตรงกัน แต่ convolution บนกราฟใช้ โครงสร้างเพื่อนบ้าน เคอร์เนลจะเลื่อนไปบนกราฟ ใช้น้ำหนักกับ feature ของโหนดปัจจุบันและเพื่อนบ้านในรัศมีที่กำหนด แล้วรวมผลรวมถ่วงน้ำหนักนั้นเพื่อสร้าง feature ใหม่ของแต่ละ output channel รายละเอียดการ implement เช่น โครงสร้างกราฟ น้ำหนัก edge, padding และ stride ก็อาจมีผลต่อ output shape ได้ด้วย

• Gemini อยู่ในสถานะไหนแล้ว?

  • น่าสนใจที่ GPT-4 ยังคงเหนือกว่า: https://twitter.com/lmsysorg/status/1735729398672716114
    เท่าที่นึกออกก็มี foundation model อย่างน้อยห้าตัว เช่น Llama, Claude, Gemini, Falcon, Mistral ที่ผลัดกันขึ้นลง แต่ GPT ก็ยังอยู่เหนือไปอีกขั้น และเป็นแบบนั้นมาหนึ่งปีแล้ว โมเดลภาษาขนาดใหญ่บนฐาน Transformer ดูเหมือนจะเรียบง่ายพอที่ใครก็ตามที่ใช้เวลา GPU มูลค่าราวหนึ่งล้านดอลลาร์ก็สร้างได้ แต่ก็ยังตาม OpenAI ไม่ทันเสียที เคล็ดลับพิเศษของพวกเขาคืออะไร?

• แล้ว Transformer เองล่ะ? มีเบาะแสไหมว่ามัน optimal ในความหมายใดบ้าง?

• รู้สึกว่าเอาประเด็นหลักไปฝังไว้ในย่อหน้าแรก แต่เนื้อหาที่เหลือก็ยอดเยี่ยม

• ความเร็วของ พัฒนาการ ML ตอนนี้น่าทึ่งมาก ผมไม่เชื่อเรื่อง singularity แต่สิ่งนี้กำลังเปลี่ยนซอฟต์แวร์และสังคมในแบบที่ไม่มีใครคาดเดาได้

  • พอดูสิ่งนี้กับ FunSearch แล้วเหมือน singularity ใกล้เข้ามาแล้ว
    https://deepmind.google/discover/blog/funsearch-making-new-d...
  • สำหรับผม มันดูเหมือน ยุคตื่นทอง อีกครั้งต่อจาก dot-com, mobile, cloud, VR
  • อีก 5 ปี ผมคิดว่าผู้คนจะไม่ได้เขียนโปรแกรมกันแบบทุกวันนี้แล้ว
  • ก่อนอื่นอยากเห็นมันเสนอ วิธีรักษา โรคที่รักษายากเสียก่อน singularity เองไม่มีความหมายถ้าไม่ก่อประโยชน์ต่อมนุษย์ และประโยชน์นั้นควรอยู่ที่การทำให้สุขภาพดีขึ้นและลดความทุกข์เป็นหลัก