การใช้งาน Mamba ในไฟล์เดียวด้วย PyTorch
(github.com/johnma2006)- mamba-minimal เป็นโปรเจ็กต์ที่นำ Mamba มาทำเป็นเวอร์ชันเรียบง่ายและมินิมอลในไฟล์ PyTorch ไฟล์เดียว
- ตั้งเป้าให้ได้ผลลัพธ์เชิงตัวเลขแบบเดียวกับอิมพลีเมนเทชันทางการใน forward/backward pass
- โค้ดถูกทำให้ง่ายขึ้น และจัดทำในรูปแบบ มีคอมเมนต์ประกอบ เพื่อให้อ่านง่าย
- ไม่ได้รวมการปรับแต่งประสิทธิภาพหลักของอิมพลีเมนเทชันทางการ จึงไม่ได้ให้ ความเร็ว และไม่ได้รวมการกำหนดค่าเริ่มต้นพารามิเตอร์ที่เหมาะสม
- เดโมใช้
state-spaces/mamba-370mและโทเคไนเซอร์EleutherAI/gpt-neox-20bเพื่อรันตัวอย่างการเติมพรอมป์ต์
ภาพรวมโปรเจ็กต์
- mamba-minimal คืออิมพลีเมนเทชันแบบเรียบง่ายและมินิมอลของ Mamba ในไฟล์ PyTorch ไฟล์เดียว
- เป้าหมายคือแสดงพฤติกรรมเดียวกับอิมพลีเมนเทชันทางการด้วยโค้ดที่อ่านง่ายกว่า
- คุณสมบัติหลัก:
- ให้ผลลัพธ์เชิงตัวเลขเทียบเท่ากับอิมพลีเมนเทชันทางการใน forward pass และ backward pass
- โค้ดที่ถูกทำให้ง่ายขึ้น
- อิมพลีเมนเทชันที่อ่านง่ายและมีคอมเมนต์กำกับ
สิ่งที่ไม่ได้รวมไว้
- ความเร็ว ไม่ใช่เป้าหมาย
- อิมพลีเมนเทชันทางการมีการปรับแต่งอย่างหนัก
- การปรับแต่งดังกล่าวเป็นส่วนสำคัญของผลงานในงานวิจัย Mamba
- อิมพลีเมนเทชันนี้คงโค้ดส่วนใหญ่ให้ง่ายไว้เพื่อให้อ่านได้สะดวก
- ไม่ได้รวม การกำหนดค่าเริ่มต้นพารามิเตอร์ที่เหมาะสม
- ระบุว่าเป็นสิ่งที่สามารถเพิ่มได้โดยไม่ต้องแลกกับความอ่านง่ายมากนัก
ตัวอย่างการใช้งานเดโม
- สามารถดูตัวอย่างการเติมพรอมป์ต์ได้ที่
demo.ipynb - ตัวอย่างใช้
model.MambaและAutoTokenizerจาก Hugging Facetransformers - โมเดลและโทเคไนเซอร์ที่ใช้:
state-spaces/mamba-370mEleutherAI/gpt-neox-20b
- พรอมป์ต์ตัวอย่างคือ
Mamba is theและผลลัพธ์ที่สร้างขึ้นมีประโยคที่อธิบาย Mamba ว่าเป็นงูพิษชนิดหนึ่ง
เอกสารอ้างอิง
- สถาปัตยกรรม Mamba ถูกนำเสนอใน Mamba: Linear-Time Sequence Modeling with Selective State Spaces
- ผู้เขียนงานวิจัยคือ Albert Gu และ Tri Dao
- อิมพลีเมนเทชันทางการอยู่ในรีโพซิทอรี
state-spaces/mamba
1 ความคิดเห็น
ความคิดเห็นจาก Hacker News
ก่อนหน้านี้ผมกับเพื่อนร่วมงานเคยสร้างไลบรารีที่แยกโค้ดโมเดลส่วนที่ใช้ร่วมกันส่วนใหญ่ออกมา และเมื่อใช้มันก็สามารถอิมพลีเมนต์โมเดลจำนวนมากได้ใน ประมาณ 100 บรรทัด หากไม่นับ Python import และคอมเมนต์
BERT: https://github.com/explosion/curated-transformers/blob/main/...
Llama 1/2: https://github.com/explosion/curated-transformers/blob/main/...
MPT: https://github.com/explosion/curated-transformers/blob/main/...
รองรับฟีเจอร์ต่าง ๆ เช่น TorchScript JIT และ PyTorch flash attention ด้วย
xformers ก็จัดการปัญหาคล้ายกัน แต่เน้นมากกว่าในการ提供โมดูล Transformer ประสิทธิภาพสูงโดยใช้ Triton อย่างไรก็ตาม การหยิบเฉพาะคอมโพเนนต์บางส่วนของไลบรารีมาใช้ไม่ค่อยง่าย และเกิด runtime error อยู่เรื่อย ๆ เลยพักไว้ก่อน ตอนนี้กำลังทำอะไรบางอย่างบนพื้นฐานสถาปัตยกรรม BERT เลยจะลองใช้เป็นข้อมูลอ้างอิงดู
คิดว่าจะลองใช้ในโปรเจกต์ถัดไป
เดิมที โค้ด Mamba มีการปรับแต่งความเร็วและองค์ประกอบอื่น ๆ มาก ทำให้เข้าใจได้ยากในทันที แต่อิมพลีเมนเทชันนี้น่าจะช่วยในการเรียนรู้ได้
ตอน inference ทีละโทเค็น ทุกอย่างจะง่ายขึ้นมาก ผมก็มีอิมพลีเมนเทชัน inference ของ Mamba ที่ทำเองด้วย: https://github.com/rbitr/llm.f90/tree/master/ssm
รู้ว่ามันเป็นรากฐานของโค้ดคำนวณเชิงวิทยาศาสตร์ที่ผ่านการพิสูจน์มายาวนาน และมักถูกห่อใช้งานผ่านไลบรารีอย่าง PyTorch หรือ Numpy แต่ทุกวันนี้ก็ไม่ใช่ภาษายอดนิยม เลยอยากรู้เหตุผลที่เลือกใช้
มีส่วนที่อยากให้ช่วยอธิบาย Mamba ให้แม้แต่คนที่ไม่ใช่นักวิจัยแมชชีนเลิร์นนิงก็เข้าใจได้
แก่นของ Mamba อยู่ที่การแก้ปัญหาเก่าแก่ของโมเดลปริภูมิสถานะ โมเดลปริภูมิสถานะเก่งในการบีบอัดบริบทอินพุต แต่ระหว่างการบีบอัดอินพุตให้เป็นสถานะซ่อน ข้อมูลที่จำเป็นต่อการใช้บริบทอย่างมีประสิทธิภาพแบบ Transformer จะถูกลบทิ้งไป
วิธีแก้คือการสร้างสิ่งที่论文เรียกว่า กลไกการเลือก กลไกนี้ขึ้นกับอินพุต ดังนั้นเมื่ออินพุตเปลี่ยน โมเดลจึงปรับเอาต์พุตของแต่ละขั้นได้ เพื่อทำเช่นนี้ จึงทำให้ตัวแปรปริภูมิสถานะบางตัวขึ้นกับอินพุตแทนที่จะคงที่ต่ออินพุต และเพิ่มเลเยอร์เชิงเส้น เป็นต้น เพื่อฉายอินพุตในแต่ละเวลาไปเป็นตัวแปรปริภูมิสถานะ
แต่การทำให้ตัวแปรปริภูมิสถานะขึ้นกับอินพุตทำให้เกิด overhead ด้านการคำนวณ จึงแก้ด้วย อัลกอริทึมที่รับรู้ฮาร์ดแวร์ ซึ่งใช้โครงสร้างหน่วยความจำของ GPU สมัยใหม่ให้เต็มที่ และหลีกเลี่ยงการย้ายข้อมูลเข้าออก HBM ให้มากที่สุด
Tri Dao คือคนที่สร้าง Flash Attention และนั่นก็เป็นวิธีใช้ฮาร์ดแวร์ให้มีประสิทธิภาพมากขึ้นใน Transformer เช่นกัน เรื่องแนวนี้เป็นความเชี่ยวชาญของเขาจริง ๆ
Mamba และ Based ที่ออกมาก่อน NeurIPS 2023 ไม่นาน ใส่การเรียกคืนเชิงเชื่อมโยงแบบหลายคิวรี (MQAR) และการพึ่งพาข้อมูลของ gating/selection ที่ได้แรงบันดาลใจจาก multi-head Attention เข้ามา สองสิ่งนี้เป็นองค์ประกอบสำคัญที่ Hyena และสถาปัตยกรรมปริภูมิสถานะก่อนหน้าไม่มี ทำให้โมเดลใหม่ ๆ ดีพอ ๆ กับ Attention ในงานเรียกคืนเชิงเชื่อมโยง และในงานอื่นที่ไม่ใช่การค้นคืน ก็มีแนวโน้มว่าอาจดีกว่า Attention เล็กน้อยด้วย
แน่นอนว่ารายละเอียดใหญ่ของ Mamba คือ การอิมพลีเมนต์ CUDA ที่มีประสิทธิภาพ ถ้าไม่มีสิ่งนั้น ความหมายของสถาปัตยกรรมนี้ในงานที่ Transformer เหมาะอยู่แล้วอาจลดลง
แม้ไม่ต้องกังวลเรื่องความยาวบริบทมากนัก ก็ยังเปิดพื้นที่ใหม่ได้มากมาย การวิเคราะห์ลำดับ DNA เป็นงานเชิงเส้นที่มี dependency ยาว และยังอาจมองภาพ วิดีโอ และข้อมูลมิติสูงเป็นสตรีมโทเค็นได้ด้วย คล้ายกับการกวาดพิกเซลแบบจอ CRT สมัยก่อน
หนึ่งในความฝันยุคแรกของ AI คือ trajectory การเรียนรู้เดียวของเอเจนต์ที่โต้ตอบกับสภาพแวดล้อมอย่างต่อเนื่องและวิวัฒน์ไปเรื่อย ๆ และ โมเดลความยาวบริบทอนันต์ แบบนี้อาจทำให้ความฝันนั้นง่ายขึ้น
อย่างไรก็ตาม ในตอนนี้ การประยุกต์ downstream ของโมเดลแบบนี้กับงานจริงที่สำคัญ โดยรวมแล้วยังผ่านการตรวจสอบและจูนมาน้อยกว่าแอปพลิเคชันที่ใช้ Attention ซึ่งเติบโตเต็มที่แล้ว การเทียบเคียงกับโครงข่ายประสาทแบบวนซ้ำเก่าช่วยได้ในระดับหนึ่ง แต่ตลอด 5 ปีที่ผ่านมา ผู้คนเชี่ยวชาญเฉพาะทางกับ Attention และ Transformer มากเกินไป จึงมีแรงเฉื่อยฝั่ง Transformer สูง
ใน论文ต้นฉบับอธิบายว่า หลังจากแปลงพารามิเตอร์แล้ว โมเดลสามารถคำนวณได้สองแบบ คือเป็นสมการเวียนเกิดเชิงเส้น หรือเป็นคอนโวลูชันทั่วทั้งลำดับ โดยทั่วไป ในการฝึกที่สามารถเห็นลำดับอินพุตทั้งหมดล่วงหน้า จะใช้โหมดคอนโวลูชันที่ขนานได้ง่าย และในการอนุมานแบบ autoregressive ที่เห็นอินพุตทีละช่วงเวลา จะสลับไปใช้โหมดวนซ้ำที่มีประสิทธิภาพ
ดังนั้นการฝึกจึง ทำแบบขนานได้ เหมือนโหมด forward แบบขนานของ RetNet ส่วนการอนุมานพื้นฐานทำในโหมดวนซ้ำเพื่อให้ได้บริบทที่ยาวที่สุดเท่าที่ทำได้ และเพราะไม่มี chunking จึงยากจะตัดสินว่าระหว่างอนุมานจะกิน RAM กับ VRAM แค่ไหน
นอกจากอธิบาย论文แล้ว ยังให้บริบทค่อนข้างมากว่ามันเข้ากับภาพใหญ่อย่างไร ฟังการปูเรื่องแล้วค่อนข้างน่าสนใจ
https://youtu.be/ouF-H35atOY?si=y2Ckp9MCFd7ulLL3
แทนที่จะใช้ Attention แบบกำลังสองที่คำนวณว่าแต่ละโทเค็นควรสนใจโทเค็นอื่นทั้งหมดมากแค่ไหน ก็หาทางคำนวณเคอร์เนลคอนโวลูชันยาวเท่ากับอินพุต แล้วนำ conv1d ไปใช้
จากความเข้าใจที่จำกัดของผม มันเกี่ยวข้องอยู่บ้างกับการทำ FFT, คูณเมทริกซ์ แล้วแปลงกลับด้วย IFFT รู้ว่ามันทำงานได้ แต่ช้า มีหลายวิธีในการคำนวณ FFT และหนึ่งในนั้นคือ เมทริกซ์ผีเสื้อ อาจเป็นแค่การประมาณ แต่ดูเหมือนว่าดีพอ และเร็วมาก มีประสิทธิภาพมากบนฮาร์ดแวร์ปัจจุบัน
ความซับซ้อนกำลังสองฟังดูแย่ แต่ในทางปฏิบัติ อัลกอริทึมที่ต่ำกว่ากำลังสองมักช้ากว่าเพราะข้อจำกัดของฮาร์ดแวร์ ดังนั้นแม้จะคาดหวังกับโมเดลปริภูมิสถานะมาก ก็ยังพูดได้ไม่ง่ายว่า Llama จบแล้ว เรายังไม่รู้ด้วยว่า Mamba จะยังทำงานได้ดีเมื่อขยายสเกลหรือไม่ และถ้าจะรู้ ก็ต้องใช้เงินจริงหลายล้านดอลลาร์ไปกับการฝึก ถึงอย่างนั้นก็ยังมองในแง่ดี
อีกโมเดลหนึ่งในสายต่ำกว่ากำลังสองที่น่าสนใจคือ RWKV ควรลองดู แต่คิดว่าอาจเคยพูดถึงในพอดแคสต์ไปแล้ว
ผมเรียนเอง และ论文ก็แค่เคยอ่านผ่าน ๆ เมื่อนานมาแล้ว จึงอาจผิดได้มาก อีกอย่าง โดยปกติ Attention มี KV cache ซึ่งช่วยประสิทธิภาพได้มาก แต่ผมมองว่า Mamba ทำแบบนั้นไม่ได้
ผมหัวเราะกับประโยคที่ว่า “Mamba คืองูพิษที่ยาวที่สุดในโลก โดยมีความยาวประมาณกว่า 150 เมตร”
แต่ก็เป็นบทความที่ยอดเยี่ยมจริง ๆ และดีที่อ้างอิงบทความ arXiv ไว้ ทำให้คนอย่างผมที่อ่านบทความแนวนี้มากกว่าจะไปตีความเปเปอร์เอง ได้แอบมองเข้าไปข้างในนิดหน่อย
คงตลกดีถ้าภายหลังต้องออกคำแก้ไขสำหรับประโยคนั้น
ผมเดาไว้ว่าแก่นของอัลกอริทึมน่าจะเป็น parallel prefix scan นั่นแหละดูเหมือนจะเป็นประเด็นสำคัญของ Mamba ไม่ใช่หรือ
for i in range(l):x = deltaA[:, :, i] * x + deltaB_u[:, :, i]y = einsum(x, C[:, i, :], 'b d_in n , b n -> b d_in')ys.append(y)อาจเป็นคำถามโง่ ๆ ก็ได้ แต่ผมสงสัยว่าการเทรน โมเดล Mamba ที่อยู่บน Hugging Face นั้นยากแค่ไหน
โมเดลใหญ่สุดดูเหมือนจะเป็น 2.8B ถ้าจะเทรนด้วยชุดข้อมูลอย่าง The Pile ต้องใช้ GPU กี่ตัวและนานแค่ไหน?
การ inference ก็น่าจะทำได้ เร็วกว่า 3~5 เท่า โดยใช้ RAM เพียงครึ่งเดียว
เคยพยายามแกะ เวอร์ชัน CUDA อย่างเป็นทางการ แต่หลังจากลองครั้งแรกแล้วล้มเหลว สุดท้ายก็ไม่ได้แตะต่อเลย implementation นี้ดูดีกว่ามาก
เป็น implementation PyTorch แบบไฟล์เดียวอีกอันที่ยอดเยี่ยมจริง ๆ หวังว่า hlb-CIFAR10 และโปรเจ็กต์ที่เกี่ยวข้องที่เคยทำมาก่อน รวมถึงอิทธิพลจากงานก่อนหน้าอย่าง minGPT หรือ DawnBench จะช่วยผลักดัน รูปแบบไฟล์เดียวที่เรียบง่าย ได้บ้าง
งานแบบนี้สำคัญต่อการวิจัย machine learning ที่มีประสิทธิภาพ และอาจเป็นหนึ่งในสิ่งสำคัญที่สุดที่เราทำได้เพื่อวงการนี้ในตอนนี้
งานวิจัยเดินหน้าไปตามความเร็วของนวัตกรรม และนวัตกรรมก็เร็วขึ้นตามส่วนกลับของเวลาที่ใช้รันการทดลอง ซึ่งเห็นได้ชัดว่าเกี่ยวข้องกับ Kolmogorov complexity ของโค้ดสำหรับงานวิจัยหรือการแฮ็กแบบง่าย ๆ
ผมย้ำเท่าไรก็ไม่พอว่าเครื่องมือแบบนี้สำคัญต่อการวิจัยแค่ไหน และโดยส่วนตัวมันเร่งกระบวนการค้นพบความรู้ของผมได้มากเพียงใด ความสามารถในการร่างไอเดียอย่างรวดเร็วภายในไม่กี่นาทีแล้วได้ผลลัพธ์ที่มีอัตราสัญญาณต่อสัญญาณรบกวนสูงทันที กลายเป็นองค์ประกอบจำเป็นของความก้าวหน้าในการวิจัยไปแล้ว
ผมมองว่า knowledge distillation และ MDL(https://en.wikipedia.org/wiki/Minimum_description_length) สำคัญมากในการย้อนกลับบรรดาเครื่องตกแต่งที่ไม่จำเป็น ของรก ๆ และการแข่งขันหัวข้อคุณค่าต่ำที่หนาแน่นเกินไปแบบ “กลัวตกกระแส” ซึ่งกระบวนการส่งและรีวิวเปเปอร์ในปัจจุบันดูเหมือนจะส่งเสริมอยู่
ช่วงหลังผมเริ่มหลีกเลี่ยงปัญหานี้และอยากไปสู่แนวทางขยายผลที่ดีกว่าเล็กน้อย จึงเริ่มเผยแพร่โค้ดเป็น “code sketch” ในรูป gist สั้น ๆ แบบ self-contained ไฟล์เดียว ช่วยลดเวลาในการพัฒนา และส่งโค้ดที่ทำงานได้แบบหยาบ ๆ ยังไม่ขัดเกลาแต่บรรจุแนวคิดไปถึงผู้คนได้ทันที จนถึงตอนนี้ดูเหมือนจะได้ผลค่อนข้างดี และอยากทำต่อไป
อยากเห็นโค้ดแบบนี้มากขึ้น ถ้าเป็นนักวิจัยที่เทรนข้อมูลในสเกลใหญ่ ก็ควรจะ data-efficient แม้กระทั่งในวิธีเผยแพร่ข้อมูลด้วย
มนุษยชาติอาจไม่เคยพัฒนาบางสิ่งที่มีความซับซ้อนมากขนาดนี้ได้รวดเร็วเช่นนี้มาก่อน
ที่เห็นความเร็วใกล้เคียงกันคงมีแค่ SpaceX ซึ่งปีนี้ก็ปล่อยจรวดระดับล้ำสมัยไปสองลำแล้ว อยากรู้ว่าในปี 2024 จะมีอะไรออกมา
x_projไม่มี bias จึงดูเหมือนว่าสามารถรวม น้ำหนักของ x_proj และ dt_proj เข้าด้วยกันได้ถ้ามีข้อกำหนดเรื่องการปรับน้ำหนัก ก็อาจทำแบบง่าย ๆ ตอน runtime ได้ และสุดท้าย kernel เดียวกับ bias น่าจะเร็วกว่า แต่ไม่แน่ใจ
สงสัยว่ามีการถกเถียงเกี่ยวกับเปเปอร์ต้นฉบับหรือเปล่า ผมคงพลาดไป แต่น่าสนใจทีเดียว
ผมไม่ค่อยเข้าใจส่วนที่ว่า “เนื่องจากขาด implementation ที่มีประสิทธิภาพ ทำให้เกิดหน่วยความจำไม่พอหรือความต้องการคำนวณที่ไม่สมจริง จึงไม่มีผลลัพธ์เต็มสำหรับ context length 8k ของ baseline อย่าง RWKV และ RetNet ซึ่งเป็นโมเดล recurrent ที่แข็งแกร่งก่อนหน้านี้และสามารถตีความเป็น SSM ได้เช่นกัน”
RetNet ไม่ได้ใช้หน่วยความจำมาก และถ้าใช้ implementation forward pass แบบแบ่ง chunk การใช้ VRAM ก็ถูกจำกัดด้วยขนาด chunk นี่คือหัวใจของการทดสอบ context length
สงสัยว่ามีใครลองทดสอบโมเดล Mamba ต้นฉบับแล้วหรือยัง ความเร็วในการเทรนจะเป็นอย่างไรเมื่อเทียบกับ RetNet ในโหมด forward pass แบบขนาน?
https://openreview.net/forum?id=AL1fq05o7H
implementation ที่กลั่นสิ่งซับซ้อนให้เหลือแต่แก่นนั้นดีเสมอ