กำลังเพิ่มการรองรับ CUDA ให้กับ Apple MLX

• MLX คือ เฟรมเวิร์กแบบอาร์เรย์ สำหรับนักวิจัยด้านแมชชีนเลิร์นนิงบนพื้นฐาน Apple Silicon โดยมอบประสบการณ์การใช้งานที่คล้ายกับ NumPy และ PyTorch
• ขณะนี้กำลังดำเนินการเพิ่ม CUDA backend ให้กับ MLX
• เป้าหมายหลัก คือการรองรับ unified memory และการใช้งานฮาร์ดแวร์ NVIDIA ได้อย่างกว้างขวาง
• ตอนนี้ยังใช้งานได้เฉพาะ ตัวอย่างจากทิวทอเรียล เท่านั้น และมีการยืนยันการ build และ test บนสภาพแวดล้อม Ubuntu 22.04 + CUDA 11.6 แล้ว
• ขณะนี้มีการปรับปรุง optimization และ refactoring ซ้ำอย่างต่อเนื่อง ระหว่างกระบวนการแก้ปัญหาด้านประสิทธิภาพและคอขวดในช่วงแรก
• CUDA backend ของ MLX ดำเนินงานภายใต้การสนับสนุนจาก Apple โดยมีเป้าหมายเพื่อมอบประสบการณ์ที่สม่ำเสมอระหว่างการพัฒนาบน Mac และการนำไป deploy บนสภาพแวดล้อมคอมพิวต์ขนาดใหญ่ ผ่านการรองรับ unified memory และฮาร์ดแวร์ NVIDIA

---

แนะนำเฟรมเวิร์ก MLX

• MLX เป็น เฟรมเวิร์กอาร์เรย์สำหรับแมชชีนเลิร์นนิง ที่พัฒนาโดยทีมวิจัยแมชชีนเลิร์นนิงของ Apple
• พัฒนาขึ้นเพื่อช่วยให้วิศวกรและนักวิจัยทดลองได้อย่างมีประสิทธิภาพ และตรวจสอบไอเดียได้อย่างรวดเร็ว
• แม้จะมีจุดเด่นที่รองรับ Apple Silicon แต่ตอนนี้กำลังพยายามขยายความสามารถให้ทำงานได้ในสภาพแวดล้อม NVIDIA GPU ด้วยการเพิ่ม CUDA backend

คุณสมบัติหลัก

• API ที่คุ้นเคย: Python API คล้าย NumPy และแพ็กเกจระดับบนอย่าง mlx.nn, mlx.optimizers ให้อินเทอร์เฟซที่คล้าย PyTorch
  • มี API สำหรับ C++, C, Swift ด้วย และคงประสบการณ์การใช้งานที่สอดคล้องกับ Python API
• การแปลงฟังก์ชันที่ประกอบต่อกันได้: มีความสามารถอย่าง automatic differentiation, automatic vectorization และการเพิ่มประสิทธิภาพกราฟการคำนวณ
• การคำนวณแบบหน่วงเวลา: รวมหลายโอเปอเรชันเข้าด้วยกันและคำนวณผลเมื่อจำเป็นจริงเท่านั้น ทำให้ใช้ทรัพยากรได้อย่างมีประสิทธิภาพ
• การสร้างกราฟแบบไดนามิก: แม้ shape ของข้อมูลนำเข้าจะเปลี่ยน ก็สามารถรันได้ทันทีโดยไม่ต้องคอมไพล์ช้า ๆ ทำให้พัฒนาและดีบักได้ง่าย
• โมเดล multi-device และ unified memory:
  • เดิมรองรับเฉพาะ CPU และ Apple GPU แต่ตอนนี้กำลังผลักดันการรองรับสภาพแวดล้อม NVIDIA GPU ผ่าน CUDA backend
  • ด้วยสถาปัตยกรรม unified memory จึงสามารถใช้งานอ็อบเจ็กต์เดียวกันบนอุปกรณ์ที่รองรับทั้งหมดได้โดยไม่ต้องคัดลอกหน่วยความจำ
• การออกแบบที่กระชับและขยายต่อได้: ช่วยให้นักวิจัยขยายและปรับปรุงเฟรมเวิร์กได้ง่าย

ความแตกต่างจากเฟรมเวิร์กอื่นและแรงบันดาลใจ

• ได้รับแรงบันดาลใจด้านการออกแบบจาก NumPy, PyTorch, Jax, ArrayFire เป็นต้น
• จุดเด่นคือการเน้น โมเดล unified memory, อินเทอร์เฟซที่กระชับ และสภาพแวดล้อมแบบ รันทันที/ดีบักได้ทันที

ตัวอย่างเด่นและกรณีการใช้งาน

• ใน MLX Examples repo มีโค้ดตัวอย่างหลากหลาย
  • การฝึก Transformer language model
  • การสร้างข้อความขนาดใหญ่บนพื้นฐาน LLaMA และการ fine-tuning ด้วย LoRA
  • การสร้างภาพด้วย Stable Diffusion
  • รวมถึงตัวอย่างการใช้งานโมเดลล่าสุด เช่น การรู้จำเสียงพูดด้วย OpenAI Whisper

ความหมายของ MLX และการรองรับ CUDA

• MLX ซึ่งเดิมเป็นเฟรมเวิร์กที่เน้น Apple Silicon กำลังขยายศักยภาพให้ใช้งานได้กว้างขึ้นในสภาพแวดล้อม NVIDIA GPU ด้วยการนำ CUDA backend มาใช้
• ด้วยการรองรับ CUDA รุ่นใหม่ MLX มีแนวโน้มจะกลายเป็นเฟรมเวิร์กที่เหมาะทั้งกับ งานวิจัยล่าสุดและการใช้งานในภาคอุตสาหกรรม บนฮาร์ดแวร์ทั้งฝั่ง Apple และ NVIDIA

[gosarinamu]

ถ้าแค่ทำให้ใช้ได้เฉพาะ nvidia gpu บน macos ก็พอแล้ว... เฮ้อๆ

[yangeok]

เป็นข่าวที่ดีเลย 555 หวังว่าจะรองรับ CUDA เร็ว ๆ จะได้เทรนความเร็วสูงบน Mac ได้ด้วยครับ~!

[GN⁺]

ความเห็นจาก Hacker News

• อยากเช็กว่าฉันเข้าใจถูกไหม: ข้อ 1 โปรแกรมที่ build ด้วย MLX สามารถทำงานบนชิปที่รองรับ CUDA ได้ แต่ข้อ 2 โปรแกรม CUDA ไม่สามารถทำงานบน Apple Silicon ได้ ใช่ไหม โดยสงสัยว่าเหตุผลที่ข้อ 2 ทำไม่ได้เป็นเพราะจะละเมิดลิขสิทธิ์หรือเปล่า (โดยเฉพาะกำแพงทางเข้าที่ NVidia สร้างไว้ซึ่งเป็นที่รู้จักกันดี)
  • ข้อ 1 ถูกต้อง สิ่งนี้ทำให้นักพัฒนาสามารถพัฒนาบนอุปกรณ์ Apple ที่สเปกค่อนข้างต่ำกว่า (รวมถึง UMA) แล้วนำโค้ดที่เสร็จแล้วไป deploy บนระบบ Nvidia ที่สเปกค่อนข้างสูงกว่าได้ ซึ่งมีประโยชน์ด้วยหลายเหตุผล
  • ข้อ 2 ไม่ใช่การละเมิดลิขสิทธิ์ การทำ API ขึ้นมาใหม่สามารถทำได้
  • สำหรับข้อ 2 ฉันคิดว่าคงไม่จริงเสียทีเดียว AMD เองก็ทำสิ่งเดียวกันนี้กับ HIP หลังจากเลิกใช้ OpenCL ราวปี 17~18
  • ฉันอยากได้ข้อ 3 มากกว่า คือเอา NVIDIA GPU มาเสียบกับ Apple Silicon แล้วใช้ CUDA ได้ คิดภาพการใช้ Apple Silicon, หน่วยความจำรวม, GPU และ CUDA ร่วมกับ PyTorch, JAX หรือ TensorFlow แต่ฉันยังไม่ได้ลอง MLX อย่างจริงจัง
  • ที่ข้อ 2 ยังทำไม่ได้ ส่วนใหญ่เป็นเพราะมันยากกว่ามาก
• ถ้าคิดว่า "เดี๋ยวก่อน แพลตฟอร์ม Apple ไม่มี CUDA รองรับอย่างเป็นทางการไม่ใช่เหรอ?" ก็ควรดูด้วยว่า patch ชุดนี้ยังรองรับ "แพลตฟอร์ม Linux ที่ใช้ CUDA 12 และ SM 7.0(Volta) ขึ้นไป" ด้วย
  คำแนะนำการติดตั้ง
• โปรเจกต์นี้นำโดย zcbenz ผู้เป็นที่รู้จักจากการสร้าง Electron และอื่น ๆ
  แนะนำ zcbenz
• หนึ่งในคุณสมบัติหลักของ MLX คือการใช้สถาปัตยกรรมหน่วยความจำรวม (UMA) เลยสงสัยว่ามันทำงานอย่างไร
  ดู bullet ใน readme ของรีโป: MLX GitHub
  คิดว่าน่าจะน่าสนใจถ้าขยายไปยัง UMA APU ทั้งหมด (ไม่ว่าผู้ผลิตไหน) แต่ก็สงสัยว่าบน discrete GPU น่าจะต้องใช้วิธีที่ต่างออกไป
  พอไปอ่านคอมเมนต์ใน PR ดูแล้ว เหมือนว่า CUDA ก็รองรับ UMA API โดยตรงเช่นกัน และจะจัดการ copy ให้แบบโปร่งใสเมื่อจำเป็น
  • จากประสบการณ์ของฉัน การขาดฟีเจอร์ prefetch ทำให้เกิด memory bottleneck หนักจากการรอ copy ถ้า dataset ทั้งหมดใส่ลง VRAM ได้ การ prefetch เองด้วยมือก็พอไหว แต่ในแอปพลิเคชันของฉัน (การเทรน ML) ประสิทธิภาพตกมากเกินไป เลยเปลี่ยนไปใช้การโหลดแบบสตรีมด้วยตัวเอง
• พูดอีกเรื่องนิดหนึ่ง แต่ดูเหมือนว่าคนทำ MLX จำนวนมากไม่ได้สังกัด Apple อย่างเป็นทางการ ตัวอย่างเช่นดู ทวิตเตอร์ของ prince_canuma
  เลยสงสัยว่า Apple สนับสนุนอยู่เบื้องหลังโดยไม่ออกหน้าอย่างเป็นทางการหรือเปล่า มีข่าวลือด้วยว่าทีมกำลังย้าย
• ในคำอธิบาย PR เขียนว่าเป็น "ongoing effort" ในการเพิ่ม "CUDA backend"
  กล่าวคือหมายความว่าโค้ด MLX สามารถ build และรันบนฮาร์ดแวร์ x86 + GeForce ได้ ไม่ใช่ในทางกลับกัน
• อย่างที่มีคนพูดว่า "ถ้าเขียน/ทดสอบโค้ดบน Mac แล้ว deploy ไปซูเปอร์คอมพิวเตอร์ได้ ก็คงเป็นประสบการณ์พัฒนาที่ดี" เลยมีคำถามว่าตอนนี้ใช้ MLX บน Linux ได้หรือยัง
  ฉันลองเองแล้ว ตอนนี้มีแค่เวอร์ชัน Python 3.12 เท่านั้นที่ขึ้นบน pypi
  หน้า MLX-CUDA บน PYPI
• สงสัยว่านี่เป็นผลจาก Strix Halo หรือเปล่า เครื่องทำงานของฉันสเปกเหลือ เลยซื้อ EVO-X2 มา ซึ่งมันใกล้เคียงหรือดีกว่าเครื่อง znver5 EPYC ระดับกลางหลายเครื่อง และเหนือกว่า EC2 หรือ GCE ส่วนใหญ่แบบทิ้งห่าง (แม้กระทั่ง NVMe) ฮาร์ดแวร์นี้ราคา 1800
  ช่วงหลังเริ่มมีรีวิวจากยูทูบเบอร์ว่า DGX Spark คุ้มค่าน้อยกว่า Strix Halo ข้อเสียคือ ROCm ยังไม่ค่อยพร้อม (แต่ดูเหมือนจะดีขึ้นเร็ว ๆ นี้) ถ้า CUDA บนอุปกรณ์ Apple ใช้งานได้สมบูรณ์จริง ต่อให้ Strix จะถูกและดีแค่ไหน มันก็เป็นตัวเลือกที่ต้องคิดหนักเหมือนกัน
  • Strix Halo ก็คือ AMD Ryzen AI Max+ 395 นั่นเอง ใช้อยู่ใน Framework Desktop และมินิพีซีหลายรุ่น แบนด์วิดท์หน่วยความจำของชิปนี้อยู่ที่ 200GB/s ซึ่งถือว่าดีมากสำหรับแพลตฟอร์ม x86 ทั่วไป แต่ยังห่างจาก Nvidia GPU (เช่น 5090 ที่ 1792GB/s) หรือ Apple M3 Ultra (800GB/s) มาก ความคุ้มค่าต่อราคายอดเยี่ยม แต่สำหรับงานที่กินหน่วยความจำมากอย่าง LLM ฉันคิดว่ามันเพิ่งแตะระดับประสิทธิภาพขั้นต่ำเท่านั้น
  • จากคำอธิบาย PR จะเห็นชัดว่ากำหนดเป้าหมายไปที่การเทรนบนคลัสเตอร์คลาวด์
  • มีคนพูดกันมา 10 ปีแล้วว่า neural net compute ของ AMD GPU จะกลายเป็นคู่แข่งของ Nvidia ในไม่ช้า แต่จนถึงตอนนี้ก็ยังไม่เกิดขึ้น
  • สงสัยว่าเทียบกับ Mac mini m4 แล้วเป็นอย่างไร
• ฉันประทับใจกับโมเดลของ MLX มาก เปิดให้ทั้งบ้านใช้โมเดลแบบ local ได้โดยไม่ต้องกังวลเรื่องเสี่ยงไฟไหม้แบบคอม Nvidia หวังว่า Apple Silicon จะกลายเป็นคู่แข่งจริงจังของฝั่งชิป Nvidia และก็สงสัยเหมือนกันว่าการรองรับ CUDA นี่อาจเป็นกลยุทธ์แบบ embrace, extend, extinguish (EEE) หรือเปล่า
• Apple มีแผนจะสร้างดาต้าเซ็นเตอร์ที่ใช้ชิปตระกูล m ในอนาคต เพื่อนำไปใช้กับการพัฒนาแอป การทดสอบ และการโฮสต์บริการภายนอก