Ulysses Sequence Parallelism: ก้าวข้ามข้อจำกัด การเทรนโมเดลด้วยบริบทนับล้านโทเค็น

Ulysses Sequence Parallelism: ก้าวข้ามข้อจำกัด การเทรนโมเดลด้วยบริบทนับล้านโทเค็นการพัฒนาปัญญาประดิษฐ์ (AI) โดยเฉพาะโมเดลภาษาขนาดใหญ่ (LLM) กำลังก้าวไปอย่างรวดเร็ว หนึ่งในความท้าทายสำคัญคือการจัดการกับ "บริบท" หรือจำนวนโทเค็นที่โมเดลสามารถประมวลผลได้ในการตอบสนองหรือสร้างข้อความ การเทรนโมเดลให้เข้าใจบริบทที่ยาวขึ้นเรื่อยๆ เป็นสิ่งจำเป็นสำหรับงานที่ซับซ้อน เช่น การวิเคราะห์เอกสารขนาดยาว การทำความเข้าใจโค้ดโปรแกรมขนาดใหญ่ หรือการทำงานที่ต้องการการคิดวิเคราะห์เป็นขั้นตอน แต่ปัญหาคือ กลไก Attention ในโมเดล Transformer แบบดั้งเดิมนั้นใช้ทรัพยากรจำนวนมหาศาล ทั้งหน่วยความจำและพลังประมวลผล โดยมีสัดส่วนการใช้ทรัพยากรเพิ่มขึ้นเป็น กำลังสอง (quadratic) ตามความยาวของลำดับโทเค็น (Sequence Length)แม้จะมีเทคนิคอย่าง FlashAttention ที่ช่วยลดการใช้หน่วยความจำลงได้มาก แต่ข้อจำกัดด้านการประมวลผลก็ยังคงอยู่ ทำให้การเทรนโมเดลด้วยบริบทที่ยาวมากๆ (เช่น 32k โทเค็นขึ้นไป) ยังคงเป็นเรื่องยากบน GPU เดี่ยวๆUlysses Sequence Parallelism: ทางออกอันชาญฉลาดUlysses Sequence Parallelism (SP) เป็นส่วนหนึ่งของโปรโตคอล Arctic Long Sequence Training (ALST) จาก Snowflake AI Research ที่เข้ามาแก้ปัญหานี้ได้อย่างตรงจุด โดยใช้แนวคิดการกระจายการคำนวณ Attention ไปยัง GPU หลายตัว ผ่านการทำ Attention Head Parallelism หรือการแบ่งส่วนของ Attention Head ออกไปประมวลผลบน GPU ที่แตกต่างกันกลไกการทำงานของ Ulysses SPSequence Sharding: ลำดับอินพุต (Input Sequence) จะถูกแบ่งออกเป็นส่วนๆ ตามมิติของ Sequence และกระจายไปยัง GPU จำนวน P ตัว แต่ละ GPU จะรับผิดชอบส่วนของ Sequence ที่แตกต่างกันQKV Projection: แต่ละ GPU จะคำนวณ Query (Q), Key (K), และ Value (V) สำหรับส่วนของ Sequence ที่ตัวเองรับผิดชอบAll-to-All Communication: ข้อมูล QKV ที่ได้จะถูกส่งผ่านการสื่อสารแบบ All-to-All เพื่อให้แต่ละ GPU ได้รับข้อมูล QKV ของทุกส่วนของ Sequence แต่จะได้รับสำหรับ Attention Head เพียงบางส่วนเท่านั้นLocal Attention: แต่ละ GPU จะทำการคำนวณ Attention สำหรับ Attention Head ที่ตนเองมีอยู่ โดยใช้กลไก Attention มาตรฐาน (เช่น FlashAttention หรือ SDPA)All-to-All Communication (อีกครั้ง): ข้อมูล Attention ที่คำนวณได้จะถูกส่งกลับผ่าน All-to-All เพื่อให้แต่ละ GPU กลับมามีข้อมูลในรูปแบบ Sequence Sharding เหมือนเดิมOutput Projection: แต่ละ GPU จะคำนวณ Output Projection สำหรับส่วนของ Sequence ที่ตนเองรับผิดชอบหัวใจสำคัญของ Ulysses คือการตระหนักว่า Attention Heads สามารถคำนวณแยกจากกันได้ โดยการแลกเปลี่ยนการประมวลผลตามตำแหน่งใน Sequence (Sequence Locality) กับการประมวลผลตาม Attention Head (Head Locality) ทำให้สามารถทำ Parallelism ได้อย่างมีประสิทธิภาพ โดยมีต้นทุนการสื่อสารที่ต่ำเปรียบเทียบ Ulysses SP กับ Ring Attention| คุณสมบัติ | Ulysses Sequence Parallelism | Ring Attention || :----------------- | :-------------------------------------------------------------------------------------------------------------------------- | :---------------------------------------------------------------------------------------------------------------------------------------- || การสื่อสาร | ใช้ All-to-All 2 ครั้งต่อ Layer มีปริมาณการสื่อสารรวม O(n⋅d/P) ต่อ GPU | สื่อสาร O(n⋅d) ต่อ GPU ผ่าน P-1 การส่งข้อมูลแบบ Point-to-Point เรียงกันไป || ประสิทธิภาพ | ใช้ประโยชน์จาก Bandwidth เต็มที่ในการสื่อสารแบบ All-to-All ในครั้งเดียว | มี Latency สูงกว่าเนื่องจากต้องส่งข้อมูลต่อเนื่องกันหลายครั้ง || ข้อจำกัด | จำนวน Attention Heads ต้องมากกว่าหรือเท่ากับจำนวน GPU ที่ใช้ทำ SP (numheads >= spsize) | ไม่มีข้อจำกัดเรื่องจำนวน Attention Heads || การใช้งาน | เหมาะกับการกระจายการคำนวณ Attention Head | เหมาะกับการกระจายการคำนวณในลักษณะวงแหวน |การผสานรวม Ulysses SP เข้ากับ Hugging Face EcosystemUlysses SP ได้รับการผสานรวมเข้ากับเครื่องมือยอดนิยมของ Hugging Face เพื่อให้การเทรนโมเดลด้วยบริบทขนาดยาวเป็นไปได้ง่ายขึ้น1. Hugging Face AccelerateAccelerate เป็นพื้นฐานสำคัญที่ช่วยให้ Ulysses SP ทำงานได้ โดยผ่านคลาส ParallelismConfig และการผสานรวมกับ DeepSpeedเมื่อเรียกใช้ accelerator.prepare() Ulysses SP จะถูกตั้งค่าโดยอัตโนมัติ:Model Registration: ลงทะเบียนโมเดลกับ DeepSpeed's UlyssesSPAttentionHFDataloader Wrapping: ห่อหุ้ม Dataloader ด้วย UlyssesSPDataLoaderAdapter เพื่อจัดการ Sequence ShardingLoss Aggregation: จัดการการคำนวณ Loss ที่ถูกต้อง โดยการรวม Loss จากแต่ละ GPU และถ่วงน้ำหนักตามจำนวนโทเค็นที่ถูกต้อง (Automatic weighted loss aggregation) เพื่อให้ Gradient ถูกต้อง แม้ว่าโทเค็นจะกระจายตัวไม่เท่ากันในแต่ละ Rankหมายเหตุ: การรวม Loss แบบถ่วงน้ำหนักนี้สำคัญมาก โดยเฉพาะเมื่อบาง Rank มีเพียง Padding หรือ Prompt tokens ที่ถูก Mask ไว้ หากใช้ Transformers Trainer หรือ TRL's SFTTrainer จะจัดการให้โดยอัตโนมัติ2. Hugging Face Transformers TrainerTrainer ทำให้การใช้งาน Ulysses SP ง่ายขึ้นไปอีก เพียงแค่ตั้งค่า parallelism_config ใน TrainingArgumentsTrainingArguments จะจัดการรายละเอียดต่างๆ ให้โดยอัตโนมัติ:Dataloader Wrapping: Trainer จะห่อหุ้ม Dataloader ด้วย UlyssesSPDataLoaderAdapter หลังจากการเตรียมโมเดลLoss Computation: ฟังก์ชัน computeloss จะตรวจจับโหมด SP และส่งต่อไปยัง deepspeedspcompute_loss เพื่อรวม Loss และคำนวณน้ำหนักที่ถูกต้องBatch Size Calculation: ขนาด Batch ที่แท้จริงจะถูกคำนวณโดยคำนึงถึง SP (dpworldsize = worldsize // spsize)Dataloader Length Adjustment: จำนวน Step ในการเทรนจะถูกปรับให้สอดคล้องกับผลกระทบของ SP ต่อจำนวน Iteration3. TRL's SFTTrainerSFTTrainer จากไลบรารี TRL (Transformer Reinforcement Learning) สร้างต่อยอดจาก Transformers Trainer โดยมีการปรับแต่งเพิ่มเติมสำหรับการทำ Supervised Fine-Tuning (SFT) ด้วยบริบทขนาดยาวพารามิเตอร์สำคัญใน SFTConfig สำหรับ Ulysses:use_ulysses: ตั้งค่าเป็น True เพื่อเปิดใช้งาน Ulysses SPulyssesconfig: กำหนดค่าการตั้งค่า Ulysses SP เช่น spsize (จำนวน GPU ที่ใช้ทำ Sequence Parallelism)SFTTrainer จะจัดการเรื่องการ Shift Labels โดยอัตโนมัติเมื่อเปิดใช้งาน Ulysses SPข้อควรพิจารณาและเคล็ดลับในการใช้งาน Ulysses SPเพื่อให้การเทรนโมเดลด้วย Ulysses SP ได้ผลลัพธ์ที่ดีที่สุด ควรพิจารณาข้อแนะนำดังนี้:Sequence Length Divisibility: ตรวจสอบให้แน่ใจว่าความยาว Sequence ของคุณ หารลงตัว ด้วย sp_size เสมอใช้ Flash Attention: Flash Attention 2 ให้ผลลัพธ์ที่ดีและมีประสิทธิภาพมากกว่า SDPA โดยเฉพาะบน Hardware รุ่นใหม่ๆผสานรวมกับ DeepSpeed ZeRO: สำหรับโมเดลขนาดใหญ่มาก ควรใช้ Ulysses SP ร่วมกับ DeepSpeed ZeRO Stage 3 เพื่อการจัดการหน่วยความจำที่มีประสิทธิภาพสูงสุดใช้ Memory Fragmentation-Friendly PyTorch Allocator: การตั้งค่า Environment Variable นี้สามารถช่วยให้เทรนด้วย Sequence Length ที่ยาวขึ้นได้การกำหนดค่า 2D Parallelism: ปรับสมดุลระหว่าง Sequence Parallelism (SP) และ Data Parallelism (DP) ให้เหมาะสมกับจำนวน GPU ที่มี ตัวอย่างเช่น dpreplicatesize × dpshardsize × spsize = numprocessesLiger-Kernel (ถ้าใช้ได้): หากโมเดลรองรับ Liger-Kernel ก็จะเข้ากันได้ดีกับ Ulysses SP และสามารถเปิดใช้งานได้ด้วย Flag ง่ายๆFusedLinearCrossEntropy และ TiledMLP: เทคนิคเหล่านี้ช่วยลดการใช้หน่วยความจำโดยหลีกเลี่ยงการสร้าง Tensor ขนาดใหญ่ของ Logits และการทำ Matrix Operations ที่มีขนาดใหญ่Token Distribution: ไม่ต้องกังวลกับการกระจาย Token ด้วยตนเอง เพราะการคำนวณ Loss Aggregation สามารถจัดการกับการกระจายตัวที่ไม่เท่ากันได้อย่างมีประสิทธิภาพการวัดผลและประสิทธิภาพจากการทดสอบเทรนโมเดล Qwen3-4B ด้วย TRL's SFTTrainer บน GPU H100 80GB โดยใช้ DeepSpeed ZeRO-3, CPU optimizer offloading, gradient checkpointing, และ Flash-attn2 พบว่า:Loss Equivalence: เมื่อปรับ Scaling ของ Global Accumulated Steps (GAS) ให้เหมาะสมกับ sp_size แล้ว การเทรนด้วย Ulysses SP ให้ผลลัพธ์ Loss ที่เทียบเท่ากับการเทรนแบบ Data Parallelism แบบดั้งเดิมMemory Usage: ที่ Sequence Length 8K, การใช้หน่วยความจำต่อ GPU ใกล้เคียงกันระหว่าง DP และ SP แต่ Ulysses SP สามารถขยายไปสู่ Sequence Length ที่ยาวกว่ามากได้ (เช่น 96K) โดยยังคงอยู่ในขีดจำกัดของ GPU H100Throughput: ที่ Sequence Length เท่ากัน (8K) Throughput ของ SP ใกล้เคียงกับ Baseline แต่เมื่อ Sequence Length ยาวขึ้น (เช่น 64K) Ulysses SP จะมี Throughput สูงกว่าอย่างมีนัยสำคัญ (3.7x) เนื่องจากต้นทุนการประมวลผล Attention ที่เพิ่มขึ้นตามกำลังสองจะกลายเป็นปัจจัยหลักที่ส่งผลต่อประสิทธิภาพการใช้ Ulysses SP ร่วมกับ GPU จำนวนมากขึ้น หรือการใช้ 2D Parallelism (SP + DP) จะช่วยให้สามารถเทรนโมเดลด้วยบริบทที่ยาวขึ้นไปได้อีก (เช่น 256K+ โทเค็น)ความต้องการของระบบHF Accelerate: deepspeed>=0.18.1, accelerate>=1.12HF Trainer: deepspeed>=0.18.1, accelerate>=1.12, transformers>=5.0Ulysses Sequence Parallelism เป็นเทคนิคที่ทรงพลัง ช่วยปลดล็อกศักยภาพของโมเดลภาษาขนาดใหญ่ในการทำความเข้าใจและประมวลผลบริบทที่ยาวนับล้านโทเค็น ทำให้งาน AI ซับซ้อนต่างๆ มีความเป็นไปได้มากขึ้น#AI #LLM #DeepLearning #HuggingFace #NLPhttps://huggingface.co/blog/ulysses-sp