PyTorch torch.compile ประสิทธิภาพก้าวกระโดด: การปรับปรุงเคอร์เนล LayerNorm และ RMSNorm เพิ่มประสิทธิภาพ GPU ถึงระดับ SOTA

คำสำคัญ: torch.compile, ตัวดำเนินการปรับมาตรฐาน, LayerNorm, RMSNorm, การปรับปรุงประสิทธิภาพ GPU

LayerNorm และ RMSNorm เป็นตัวดำเนินการปรับมาตรฐานพื้นฐานในโมเดลการเรียนรู้เชิงลึก ใช้สำหรับประมวลผลมาตรฐานข้อมูลอินพุต พวกมันเป็นโมดูลสำคัญที่รับประกันการลู่เข้าอย่างราบรื่นในการฝึกโมเดลขนาดใหญ่และเพิ่มประสิทธิภาพการอนุมาน ประสิทธิภาพของเคอร์เนลของพวกมันบนแพลตฟอร์ม GPU ประสิทธิภาพสูงกำหนดปริมาณงานการฝึกโดยรวมโดยตรง

ปัจจุบัน เคอร์เนลปรับมาตรฐานระดับแนวหน้าของอุตสาหกรรมส่วนใหญ่พึ่งพาการปรับปรุงเชิงลึกด้วยมือ ในทางตรงกันข้าม torch.compile ดั้งเดิมของ PyTorch มีช่องว่างประสิทธิภาพในตัวดำเนินการประเภทนี้มาเป็นเวลานาน ซึ่งแสดงให้เห็นหลักๆ จากการใช้ประโยชน์แบนด์วิดท์หน่วยความจำต่ำ ตรรกะการรีดิวซ์และกลยุทธ์แบ่งส่วนที่ไม่สมเหตุสมผล และอื่นๆ ทำให้ยากที่จะใช้ศักยภาพฮาร์ดแวร์ของ GPU รุ่นใหม่ เช่น H100, B200 ได้อย่างเต็มที่

บทความนี้ใช้เคอร์เนล Quack ที่ปรับปรุงด้วยมือเป็นเกณฑ์มาตรฐานประสิทธิภาพ และมุ่งเน้นการปรับปรุงแบ็กเอนด์ Inductor ของ torch.compile ผ่านการปรับปรุงพารามิเตอร์หลัก เช่น การแบ่งส่วนรีดิวซ์ภายใน การกำหนดค่ากลุ่มเธรด (warp) เสนอแผนการรีดิวซ์ลำดับผสม และผสานรวมกลยุทธ์การแยกรีดิวซ์ การวางท่อซอฟต์แวร์ และการปรับแต่งอัตโนมัติ เราได้เพิ่มประสิทธิภาพของเคอร์เนลการส่งต่อและแพร่ย้อนกลับอย่างมาก ในที่สุดทำให้ตัวดำเนินการปรับมาตรฐานที่สร้างจากการคอมไพล์บรรลุระดับแนวหน้าของอุตสาหกรรม ในขณะที่ยังคงข้อได้เปรียบของการผสานตัวดำเนินการโดยรอบอัตโนมัติ ซึ่งเป็นโซลูชันที่เป็นประโยชน์สำหรับการปรับปรุงการคอมไพล์โมเดลขนาดใหญ่

เราได้ประเมินและปรับปรุงประสิทธิภาพของ torch.compile สำหรับ LayerNorm/RMSNorm บน GPU H100 และ B200 ของ NVIDIA ทำให้ได้ประสิทธิภาพในระดับเคอร์เนลใกล้เคียงกับระดับแนวหน้าของอุตสาหกรรม (SOTA) ในขณะเดียวกันยังสามารถเพิ่มความเร็วได้อีกผ่านความสามารถในการผสานอัตโนมัติ

สารบัญบทความ

• หนึ่ง. การส่งต่อ
  • 1.1 การปรับมาตรฐานตามชั้น (LayerNorm)
  • 1.2 การปรับมาตรฐานรากที่สองของค่าเฉลี่ยกำลังสอง (RMSNorm)
  • 1.3 เคอร์เนลมาตรฐาน Quack
  • 1.4 ตรรกะการสร้างของ torch.compile
  • 1.5 ผลการทดสอบมาตรฐาน
• สอง. การแพร่ย้อนกลับ
  • 2.1 การรีดิวซ์แบบผสาน
  • 2.2 เคอร์เนลการแพร่ย้อนกลับแบบผสานที่สร้างโดย Inductor
  • 2.3 การปรับแต่งอัตโนมัติขนาดการแยก
  • 2.4 การปรับปรุงการวางท่อซอฟต์แวร์
  • 2.5 ผลการทดสอบมาตรฐาน
• สรุป

หนึ่ง. การส่งต่อ

1.1 การปรับมาตรฐานตามชั้น (LayerNorm)

การปรับมาตรฐานตามชั้นถูกเสนอครั้งแรกในบทความวิจัย “Layer Normalization” มันทำการปรับมาตรฐานโดยการคำนวณค่าเฉลี่ยและความแปรปรวนของเทนเซอร์อินพุต และแนะนำพารามิเตอร์น้ำหนัก (γ) และอคติ (β) ที่เรียนรู้ได้สำหรับการปรับขนาดและการเลื่อน

1.2 การปรับมาตรฐานรากที่สองของค่าเฉลี่ยกำลังสอง (RMSNorm)

การปรับมาตรฐานรากที่สองของค่าเฉลี่ยกำลังสองเป็นแผนการปรับปรุงของ LayerNorm ถูกเสนอในบทความวิจัย “Root Mean Square Layer Normalization” วิธีนี้ไม่ทำการปรับให้อยู่กึ่งกลางอีกต่อไป แต่ใช้ค่ารากที่สองของค่าเฉลี่ยกำลังสอง (RMS) ของผลรวมกำลังสองของค่าอินพุตเพื่อการปรับมาตรฐาน มันยังใช้พารามิเตอร์น้ำหนัก (γ) ที่เรียนรู้ได้สำหรับการปรับขนาด แต่ลบพารามิเตอร์อคติออก

ตรรกะการส่งต่อของ LayerNorm และ RMSNorm มีความคล้ายคลึงกัน โดยทั่วไปจะดำเนินการรีดิวซ์บนมิติที่ต่อเนื่องกัน และมาพร้อมกับการดำเนินการแบบองค์ประกอบต่อองค์ประกอบหลายครั้ง ปริมาณการคำนวณของ RMSNorm โดยทั่วไปต่ำกว่า เพราะมันลดการดำเนินการจุดลอยตัวและไม่จำเป็นต้องจัดการกับพารามิเตอร์อคติ เนื่องจากความคล้ายคลึงกันของเคอร์เนลของทั้งสอง บทความนี้จะสลับแสดงผลการทดสอบมาตรฐานของพวกมัน

1.3 Quack

Quack เป็นไลบรารีเคอร์เนล GPU ที่ปรับปรุงอย่างสุดขีดซึ่งพัฒนาโดยทีม Tri Dao และใช้ CuteDSL เอกสาร README แสดงให้เห็นว่า บน GPU H100 ประสิทธิภาพของ Quack สำหรับเคอร์เนลรีดิวซ์ประเภทนี้ดีกว่า torch.compile อย่างมีนัยสำคัญ ดังนั้น เราจึงใช้ Quack เป็นเกณฑ์มาตรฐานประสิทธิภาพระดับแนวหน้าของอุตสาหกรรม (SOTA) เพื่อประเมินระดับการปรับปรุงของ torch.compile

(คำอธิบายภาพ: ที่ความแม่นยำ bf16 ประสิทธิภาพของ Quack ในตัวดำเนินการหลัก เช่น RMSNorm, Softmax, Cross-Entropy นำหน้า torch.compile และไลบรารีเคอร์เนลอื่นๆ อย่างครอบคลุม)

ตามการทดสอบมาตรฐานของ Quack ประสิทธิภาพของ torch.compile ก่อนการปรับปรุงโดยทั่วไปสามารถทำได้เพียงประมาณ 50% ของ Quack

1.4 torch.compile

ต่อไปนี้จะอธิบายตรรกะทั่วไปที่ torch.compile สร้างเคอร์เนลสำหรับการส่งต่อของ LayerNorm (แนวคิดการใช้งานของ RMSNorm เหมือนกันทุกประการ) เราสมมติว่าจำนวนองค์ประกอบของมิติรีดิวซ์ของอินพุตเป็นแบบต่อเนื่อง ซึ่งใน Inductor เรียกว่าการรีดิวซ์ภายใน

(คำอธิบายภาพ: กระบวนการคำนวณแบบขนานแบบแบ่งส่วนสำหรับการส่งต่อของ RMSNorm/LayerNorm ผ่านกลยุทธ์แบ่งส่วนที่เหมาะสมเพื่อปรับให้เข้ากับสถาปัตยกรรมแบบขนานของ GPU และเพิ่มประสิทธิภาพการเข้าถึงหน่วยความจำ)

แม้ว่าตรรกะการสร้างเคอร์เนลจะเกี่ยวข้องกับการแบ่งส่วนและการจัดตารางเวลาแบบขนาน แต่หลักการพื้นฐานของมันเป็นไปตามสัญชาตญาณ

1.4 การใช้งานเคอร์เนลและกลยุทธ์การปรับปรุง

เคอร์เนลการส่งต่อของ LayerNorm/RMSNorm ที่สร้างโดย Inductor ปฏิบัติตามขั้นตอนการคำนวณต่อไปนี้:
1. รักษาผลรวมบางส่วนขนาด R_BLOCK สำหรับแต่ละแถวของอินพุต
2. ใช้ผลรวมบางส่วนเพื่อคำนวณค่าเฉลี่ยและความแปรปรวน
3. ดำเนินการดำเนินการแบบองค์ประกอบต่อองค์ประกอบกับอินพุตตามสูตรการปรับมาตรฐานตามชั้น
4. เก็บผลลัพธ์ของการดำเนินการแบบองค์ประกอบต่อองค์ประกอบ:
* หากเปิดการแปลงแบบแอฟฟินองค์ประกอบต่อองค์ประกอบ (elementwise_affine=True) และจำเป็นต้องคำนวณเกรเดียนต์ (requires_grad=True) จะเก็บค่าเฉลี่ยและความแปรปรวนเพิ่มเติมสำหรับใช้ในการแพร่ย้อนกลับ

คำอธิบายเพิ่มเติม: หากขนาดมิติรีดิวซ์น้อยกว่าเกณฑ์เชิงประจักษ์หนึ่ง (1024) Inductor จะสร้างเคอร์เนลการรีดิวซ์แบบคงทน (persistent reduction) ในกรณีนี้ไม่จำเป็นต้องวนซ้ำผ่านมิติรีดิวซ์ สามารถคำนวณค่าเฉลี่ยได้โดยตรง

เมื่อเปรียบเทียบการส่งต่อของ RMSNorm ระหว่าง torch.compile และการใช้งาน Quack บนการ์ด H100 และ B200 เราได้ทำซ้ำปัญหาที่ประสิทธิภาพของ torch.compile ต่ำกว่า Quack อย่างมาก หลังจากปรับแต่งอัตโนมัติและปรับปรุงพารามิเตอร์เริ่มต้นของ Inductor แล้ว torch.compile บน H100 และ B200 บรรลุประสิทธิภาพระดับแนวหน้าของอุตสาหกรรม (SOTA)

วิธีการปรับปรุงหลักที่ทำให้เกิดผลลัพธ์นี้มีดังนี้:
* แทรก torch._dynamo.reset() ในการทดสอบมาตรฐาน: เพื่อให้แน่ใจว่า torch.compile ไม่เปิดใช้งานการอนุมานรูปร่างไดนามิกอัตโนมัติ ก่อนหน้านี้การเรียก torch.compile แยกกันสำหรับแต่ละรูปร่าง ทำให้คอมไพเลอร์ตัดสินใจใช้รูปร่างไดนามิก
* แก้ไขการตัดสินใจกำหนดค่าการปรับแต่งอัตโนมัติที่ไม่สมเหตุสมผล: การกำหนดค่าการปรับแต่งอัตโนมัติเริ่มต้นบน H100 และ B200 จะทำให้เลือกทางเลือกที่ด้อยกว่า ซึ่งนำไปสู่ประสิทธิภาพที่ไม่ดี การตั้งค่า mode='max-autotune' สามารถบรรเทาปัญหานี้ได้ เราทำการปรับปรุงหลายประการสำหรับกฎฮิวริสติกเริ่มต้น:
1. ขยายขนาด RBLOCK ของการรีดิวซ์ภายใน
2. ขยายขนาด XBLOCK ในการรีดิวซ์แบบคงทน สำหรับสถานการณ์รีดิวซ์ขนาดเล็กที่มีจำนวนองค์ประกอบรีดิวซ์ ≤2048
3. ลดจำนวนกลุ่มเธรด (num_warps) ตามมิติรีดิวซ์เฉพาะ เมื่อจำนวนกลุ่มเธรดมากเกินไป มักจะขัดขวางการเวกเตอร์ไอซ์พีค การเวกเตอร์ไอซ์พีคเป็นกุญแจสำคัญในการเพิ่มจำนวนไบต์ที่กำลังบิน (in-flight bytes) และทำให้เวิร์กโหลดที่ถูกจำกัดด้วยหน่วยความจำใช้แบนด์วิดท์หน่วยความจำสูงสุดได้เต็มที่ และสถาปัตยกรรม Blackwell (B200) มีความไวต่อแบนด์วิดท์หน่วยความจำมากขึ้น ปัญหานี้จึงมีผลกระทบที่เด่นชัดกว่า

1.5 ผลการทดสอบมาตรฐาน

ด้านล่างแสดงข้อมูลเปรียบเทียบประสิทธิภาพระหว่าง torch.compile 2.11 และ Quack (เวอร์ชัน主干 24 มีนาคม 2026) ในรูปร่างการทดสอบมาตรฐานของ Quack และรูปร่าง M ใหญ่ N เล็กที่พบบ่อยในธุรกิจจริง ผลลัพธ์แสดงให้เห็นว่าประสิทธิภาพของ torch.compile โดยพื้นฐานแล้วเทียบเท่ากับ Quack

มีสถานการณ์ประสิทธิภาพถดถอยสองประเภทในการทดสอบ:
1. การถดถอยประสิทธิภาพเล็กน้อยเมื่อ N=384 เนื่องจาก Triton ไม่สามารถแสดงขนาดบล็อกที่ไม่ใช่กำลังสองของจำนวนเต็มได้อย่างกระชับ
2. การถดถอยประสิทธิภาพอย่างมากบน H100 เมื่อขนาด N ใหญ่มาก เนื่องจาก Triton ยังไม่รองรับการแสดงหน่วยความจำแชร์แบบกระจาย

ภาพ: การทดสอบประสิทธิภาพ RMSNorm ที่ความแม่นยำ bf16 บน GPU B200 เปรียบเทียบระหว่าง Quack และโซลูชัน Static torch.compile ที่ปรับปรุงแล้ว

ภาพ: การทดสอบประสิทธิภาพ RMSNorm ที่ความแม่นยำ bf16 บน GPU H100 เปรียบเทียบระหว่าง Quack และโซลูชัน Static torch.compile ที่ปรับปรุงแล้ว

สอง. การแพร่ย้อนกลับ

กระบวนการแพร่ย้อนกลับของ LayerNorm/RMSNorm ซับซ้อนกว่าการส่งต่อ อย่างน้อยจำเป็นต้องคำนวณเกรเดียนต์สองประเภท: เกรเดียนต์ของอินพุต เกรเดียนต์ของน้ำหนัก ใน LayerNorm ยังต้องคำนวณเกรเดียนต์อคติเพิ่มเติม

ภาพ: แผนผังเค้าโครงการรีดิวซ์เกรเดียนต์สำหรับการแพร่ย้อนกลับ ใช้สำหรับคำนวณเกรเดียนต์สามประเภทอย่างมีประสิทธิภาพ

เพื่อลดความซับซ้อนของการใช้งานและหลีกเลี่ยงสูตรคณิตศาสตร์ที่ซับซ้อน จากมุมมองประสิทธิภาพ การคำนวณเกรเดียนต์เหล่านี้จำเป็นต้องดำเนินการรีดิวซ์บนสองมิติของเกรเดียนต์อินพุตของการแพร่ย้อนกลับ (ซึ่งก็คือเกรเดียนต์ของเอาต์พุตขั้นตอนก่อนหน้าของการส่งต่อ) แยกกัน

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

2.1 การรีดิวซ์แบบผสาน (Fused Reductions)

สำหรับรูปร่างที่สมเหตุสมผลซึ่งจำนวนองค์ประกอบโดยทั่วไปไม่มาก และข้อมูลแถวเดียวสามารถใส่ลงในบล็อกเธรดได้อย่างเพียงพอ (โดยทั่วไป ≤16384) สามารถใช้เคอร์เนลแบบผสานที่มีประสิทธิภาพดีกว่าได้ โดยไม่ทำให้การใช้งานหน่วยความจำแชร์/เรจิสเตอร์ระเบิด

พูดง่ายๆ เคอร์เนลนี้จะดำเนินการคำนวณรีดิวซ์สำหรับเกรเดียนต์น้ำหนักและอคติตามปกติ ในขณะเดียวกันก็ดำเนินการรีดิวซ์มิติคอลัมน์สำหรับเกรเดียนต์อินพุตให้เสร็จแบบขนานสำหรับแต่ละแถว มีเอกสารหลายฉบับที่ศึกษาวิธีการผสานประเภทนี้

ใน Inductor เราใช้ประเภทที่แตกต่างกันเพื่อระบุการดำเนินการรีดิวซ์:
* การรีดิวซ์ภายใน (INNER reduction): ดำเนินการรีดิวซ์บนมิติต่อเนื่องที่มีสตรายด์เป็น 1
* การรีดิวซ์ภายนอก (OUTER reduction):

⚠️ หมายเหตุ: เนื้อหาได้รับการแปลโดย AI และตรวจสอบโดยมนุษย์ หากมีข้อผิดพลาดโปรดแจ้ง

☕ สนับสนุนค่ากาแฟทีมงาน

หากคุณชอบบทความนี้ สามารถสนับสนุนเราได้ผ่าน PromptPay

SCAN TO PAY WITH ANY BANK