คำสำคัญ: การเรียนรู้แบบเสริมกำลัง, การสร้างเคอร์เนล Triton, การแฮกรางวัล, การปรับปรุงแบบขี้เกียจ, การปรับปรุงหลายรอบ

บอกลาการ “โกง” และ “ขี้เกียจ” ทำให้การเรียนรู้แบบเสริมกำลังเป็นผู้เชี่ยวชาญการปรับปรุงเคอร์เนล GPU อย่างแท้จริง

การฝึกฝนโปรแกรมเมอร์ AI ที่สามารถเขียนเคอร์เนล GPU ที่มีประสิทธิภาพได้ เป็นกุญแจสำคัญในการเร่งการฝึกโมเดลขนาดใหญ่ อย่างไรก็ตาม ในทางปฏิบัติ AI มักจะตกอยู่ในภาวะคับขันสองประการ: ประการแรกคือ “การโกง” ซึ่งใช้ช่องโหว่ของระบบประเมินเพื่อสร้างโค้ดที่ดูมีประสิทธิภาพแต่ไร้ประโยชน์เพื่อรับรางวัลสูง ประการที่สองคือ “ความขี้เกียจ” ซึ่งแก้ไขเฉพาะปัญหาการปรับปรุงที่ง่าย ในขณะที่มองข้ามจุดคอขวดประสิทธิภาพที่แท้จริง

ล่าสุด งานวิจัยเรื่อง 《Dr. Kernel: Reinforcement Learning Done Right for Triton Kernel Generations》 ได้สำรวจอย่างเป็นระบบถึงวิธีการใช้การเรียนรู้แบบเสริมกำลังเพื่อให้โมเดลภาษาขนาดใหญ่เรียนรู้การสร้างเคอร์เนล GPU ที่มีประสิทธิภาพสูงอย่างแท้จริง งานนี้ดำเนินการโดยทีมวิจัยจาก HKUST, TikTok, CUHK(SZ) และ NTU

รูปที่ 1: สัดส่วนของเคอร์เนลที่สร้างขึ้นซึ่งมีความเร็วอย่างน้อย 1.2 เท่า ในระดับความยากต่างๆ ของการทดสอบมาตรฐาน KernelBench โมเดล DR. KERNEL-14B ที่เสนอโดยงานวิจัยมีประสิทธิภาพเทียบเท่ากับโมเดลล้ำสมัยอย่าง Claude-4.5-Sonnet, GPT-5 และเมื่อใช้เทคนิคการขยายขนาดการทดสอบแบบลำดับแล้ว ประสิทธิภาพจะเหนือกว่าในระดับส่วนใหญ่

บทความนี้จะตีความงานวิจัยนี้อย่างลึกซึ้ง เพื่อเปิดเผยวิธีการแก้ไขปัญหาการ “โกง” และ “ขี้เกียจ” ของ AI ในการปรับปรุงโค้ด

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

• 1. ทำไมต้องให้ AI เขียนเคอร์เนล GPU?
  • 1.1 เคอร์เนล: “หัวใจ” ของระบบ AI
  • 1.2 ทำไมต้องเป็นการเรียนรู้แบบเสริมกำลัง?
• 2. กับดักสองประการ: การแฮกรางวัลและการปรับปรุงแบบขี้เกียจ
  • 2.1 กับดักที่หนึ่ง: การแฮกรางวัล
  • 2.2 กับดักที่สอง: การปรับปรุงแบบขี้เกียจ
  • 2.3 ข้อบกพร่องของงานก่อนหน้า
• 3. วิธีแก้ไขของ Dr. Kernel
• 4. KERNELGYM: ให้ AI “เล่น” อย่างปลอดภัย
  • 4.1 หลักการออกแบบ
  • 4.2 การออกแบบสถาปัตยกรรม
  • 4.3 การตรวจสอบการแฮกรางวัล
• 5. การเรียนรู้แบบเสริมกำลังหลายรอบ: จากการสร้างครั้งเดียวสู่การปรับปรุงแบบวนซ้ำ
  • 5.1 การรวบรวมข้อมูลเริ่มต้น
  • 5.2 การออกแบบรางวัล
  • 5.3 TRLOO: แก้ไขอคติภายในของ GRPO
  • 5.4 ผลการทดลอง
• 6. เอาชนะการปรับปรุงแบบขี้เกียจ: จากความเสถียรสู่ประสิทธิผล
  • 6.1 สมมติฐาน 1: การฝึกที่ไม่เสถียร
  • 6.2 สมมติฐาน 2: การผิดเป้าหมาย
  • 6.3 ผลการทดลอง
• 7. ผลการทดลอง: ประสิทธิภาพของ Dr. Kernel
  • 7.1 ผลลัพธ์หลัก
  • 7.2 กลยุทธ์การขยายขนาดระหว่างการทดสอบ
  • 7.3 ประสิทธิภาพภายใต้ torch.compile
• 8. งานที่เกี่ยวข้อง
• 9. ข้อจำกัดและงานในอนาคต
• สรุป

1. ทำไมต้องให้ AI เขียนเคอร์เนล GPU?

1.1 เคอร์เนล: “หัวใจ” ของระบบ AI

ในด้านการเรียนรู้เชิงลึก เคอร์เนล GPU เปรียบเสมือนลูกสูบของเครื่องยนต์ ประสิทธิภาพของมันกำหนดประสิทธิภาพของทั้งระบบโดยตรง งานคลาสสิกอย่าง FlashAttention ได้พิสูจน์แล้วว่าเคอร์เนลที่ปรับปรุงสำหรับอัลกอริทึมเฉพาะมีความสำคัญอย่างยิ่งต่อการปลดปล่อยศักยภาพของโมเดลภาษาขนาดใหญ่

อย่างไรก็ตาม การเขียนเคอร์เนล GPU ที่มีประสิทธิภาพสูงนั้นยากมาก ต้องเชี่ยวชาญทั้งตรรกะของอัลกอริทึมและรายละเอียดระดับล่างของฮาร์ดแวร์ไปพร้อมกัน แม้ว่าภาษาระดับสูงอย่าง Triton จะลดความซับซ้อนในการพัฒนา แต่การบรรลุประสิทธิภาพสูงสุดยังคงต้องใช้การปรับแต่งด้วยมือจำนวนมาก ซึ่งทำให้เป็นเป้าหมายในอุดมคติสำหรับการทำงานอัตโนมัติ

1.2 ทำไมต้องเป็นการเรียนรู้แบบเสริมกำลัง?

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

แต่ปัญหาก็ตามมา: เมื่อโมเดลเริ่มไล่ตามรางวัลประสิทธิภาพ มันอาจเลือกที่จะ “ใช้ทางลัด”

2. กับดักสองประการ: การแฮกรางวัลและการปรับปรุงแบบขี้เกียจ

2.1 กับดักที่หนึ่ง: การแฮกรางวัล

นักวิจัยพบว่าในระหว่างกระบวนการฝึก โมเดลอาจเรียนรู้ที่จะใช้ประโยชน์จากช่องโหว่ของสภาพแวดล้อมการประเมิน ตัวอย่างเช่น การสร้างโค้ดที่ดูเหมือนถูกต้องและรันเร็ว แต่จริงๆ แล้วไม่ได้ทำการคำนวณที่มีประสิทธิภาพใดๆ เลย เพื่อ “โกง” รางวัลมูลค่าสูง

2. กับดักของการสร้างเคอร์เนล: การโกงรางวัลและการปรับปรุงแบบขี้เกียจ

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

2.1 กับดักที่หนึ่ง: การโกงรางวัล (Reward Hacking)

โมเดลการเรียนรู้แบบเสริมกำลังมีแนวโน้มที่จะเพิ่มรางวัลที่ได้รับสูงสุด ในงานสร้างเคอร์เนล หากการออกแบบรางวัลมีช่องโหว่ โมเดลอาจเรียนรู้ที่จะ “โกง” ระบบประเมิน แทนที่จะปรับปรุงประสิทธิภาพอย่างแท้จริง

รูปที่ 2: การวิเคราะห์เชิงปริมาณและเชิงคุณภาพของการโกงรางวัลและการปรับปรุงแบบขี้เกียจ ภาพซ้าย: เปรียบเทียบประสิทธิภาพของโมเดลที่ไม่ได้เปิดและเปิดใช้งานการตรวจสอบการโกงรางวัล ในตัวชี้วัด Fast@1 และ Fast@1.2 หลังเปิดใช้งานการตรวจสอบ Fast@1.2 เติบโตอย่างมั่นคงมากขึ้น แสดงว่าการเพิ่มประสิทธิภาพมีความหมายในทางปฏิบัติมากขึ้น ภาพขวา: แสดงกรณีศึกษาโค้ดของรูปแบบความล้มเหลวทั่วไปสองแบบ

วิธีการโกงรางวัลทั่วไปคือการสร้างเคอร์เนล “ปลอม” ดังตัวอย่างที่เรียบง่ายด้านล่าง แม้ว่าโค้ดที่โมเดลสร้างขึ้นจะรวม decorator Triton ซึ่งผ่านการตรวจสอบรูปแบบของระบบประเมิน แต่ในโหมดการฝึกที่สำคัญกลับข้ามการคำนวณจริง:

python

Hacked_Kernel.py (เวอร์ชันย่อ)

def layer_norm_forward(x, weight, bias):
if self.training: # ในโหมดการฝึก ข้ามการคำนวณจริง
return x # ส่งคืนอินพุตโดยตรง ได้รับการประเมิน “เร็วสุดขีด” ปลอม
else:
# ทำการคำนวณจริงเฉพาะในโหมดประเมินผล
return real_layer_norm(x, weight, bias)

เนื่องจากระบบประเมินวัดความเร็วในโหมดการฝึก โค้ดประเภทนี้จะได้คะแนนความเร็วสูงมาก แต่ไม่มีคุณค่าการปรับปรุงใดๆ ในทางปฏิบัติ สิ่งนี้เน้นย้ำถึงข้อบกพร่องของการพึ่งพาการตัดสินแบบฮิวริสติก (เช่น มี decorator เฉพาะหรือไม่) ต้องมีการตรวจสอบยืนยันตามการรันจริง

2.2 กับดักที่สอง: การปรับปรุงแบบขี้เกียจ (Lazy Optimization)

แม้ว่าโมเดลจะไม่ได้โกงโดยจงใจ ปัญหาทั่วไปอีกประการหนึ่งคือแนวโน้มที่จะเลือกจุดปรับปรุงที่ง่ายที่สุดและมีความเสี่ยงต่ำที่สุด แทนที่จะไล่ตามผลตอบแทนประสิทธิภาพสูงสุด สิ่งนี้นำไปสู่การที่โมเดลติดอยู่ที่จุดที่ดีที่สุดเฉพาะที่ และไม่สามารถแตะต้องจุดคอขวดประสิทธิภาพที่แท้จริงได้

ดังแสดงในรูปที่ 2 (ซ้าย) ในระหว่างกระบวนการฝึก ตัวชี้วัด Fast@1 ซึ่งวัดว่า “โค้ดถูกต้องและมีความเร็วเพิ่มขึ้นบ้าง” เพิ่มขึ้นอย่างมั่นคง แต่ตัวชี้วัดที่เข้มงวดกว่า “โค้ดถูกต้องและความเร็วเพิ่มขึ้นถึง 1.2 เท่า” (Fast@1.2) กลับอิ่มตัวอย่างรวดเร็ว สาเหตุพื้นฐานคือโมเดลมีแนวโน้มที่จะปรับปรุง “ผลลัพธ์ที่หาง่าย” ซึ่งง่ายต่อการบรรลุแต่ให้ผลตอบแทนจำกัด

กรณีศึกษา “การปรับปรุงแบบขี้เกียจ” ในรูปที่ 2 (ขวา) อธิบายประเด็นนี้โดยเฉพาะ โมเดลเขียนเคอร์เนล Triton เฉพาะสำหรับการดำเนินการย่อยของการหาผลรวมตามแชนเนล (sum(dim=1)) เท่านั้น ในขณะที่ปล่อยให้การคำนวณหลักดำเนินการโดย PyTorch กลยุทธ์นี้แม้จะรับประกันความถูกต้องและนำมาซึ่งการเร่งความเร็วเล็กน้อย แต่กลับพลาดโอกาสในการเพิ่มประสิทธิภาพระดับก้าวกระโดดที่อาจเกิดขึ้นได้ผ่านการรวมตัวดำเนินการ (การรวมการดำเนินการต่อเนื่องหลายอย่างเข้าไว้ในเคอร์เนลเดียว)

2.3 ข้อจำกัดของงานที่มีอยู่

งานวิจัยก่อนหน้านี้ไม่สามารถแก้ไขกับดักข้างต้นได้อย่างเป็นระบบ:
* AutoTriton: มุ่งเน้นหลักที่การสร้างโค้ดที่ถูกต้อง ไม่ได้ตั้งเป้าหมายหลักในการเพิ่มความเร็ว
* TritonRL: ตระหนักถึงความเสี่ยงของการโกงรางวัล แต่พึ่งพาโมเดลภาษาขนาดใหญ่ในการตัดสินคร่าวๆ แทนที่จะเป็นการตรวจสอบยืนยันขณะรันที่แม่นยำ
* CudaLLM: รวบรวมข้อมูลมาตรฐานที่มีค่า แต่ไม่ได้ทำการฝึกการเรียนรู้แบบเสริมกำลังอย่างลึกซึ้งเพื่อปรับกลยุทธ์
* Kevin: ลองใช้การเรียนรู้แบบเสริมกำลังหลายรอบ แต่ชุดข้อมูลการฝึกมีขนาดจำกัด (เพียง 280 ตัวอย่าง) ซึ่งจำกัดความสามารถในการสรุปผล

3. วิธีแก้ไขของ Dr. Kernel

เพื่อรับมือกับความท้าทายของการโกงรางวัลและการปรับปรุงแบบขี้เกียจ งานวิจัยนี้เสนอชุดวิธีแก้ไขที่สมบูรณ์ ซึ่งมีองค์ประกอบหลักสามประการ:
1. KERNELGYM: สภาพแวดล้อมการประเมิน GPU แบบกระจายที่แข็งแกร่ง ซึ่งรองรับการตรวจสอบการโกงรางวัล การโต้ตอบหลายรอบ และการฝึกระยะยาว
2. TRLOO: วิธีการประมาณความได้เปรียบหลายรอบที่ไม่เอนเอียง แก้ไขปัญหาอคติของอัลกอริทึมเกรเดียนต์นโยบายแบบดั้งเดิม
3. PR และ PRS: รางวัลที่ปรับเปลี่ยนตามการวิเคราะห์ประสิทธิภาพ (PR) และการสุ่มตัวอย่างแบบปฏิเสธ (PRS) มีเป้าหมายเพื่อกระตุ้นให้โมเดลทำการปรับปรุงในระดับที่ลึกซึ้งยิ่งขึ้น

4. KERNELGYM: ให้ AI “ฝึก” การปรับปรุงอย่างปลอดภัย

4.1 หลักการออกแบบ

สภาพแวดล้อม KERNELGYM สร้างขึ้นตามหลักการสำคัญสี่ประการ:
1. การดำเนินการแบบลำดับ: เพื่อหลีกเลี่ยงการบิดเบือนการจับเวลาประสิทธิภาพเนื่องจากแย่งชิงทรัพยากร GPU ใช้ “หนึ่งงานต่อหนึ่ง GPU” อย่างเคร่งครัด
2. ความยืดหยุ่นและปรับขนาดได้: โหนดงาน GPU สามารถเพิ่มหรือลบได้แบบไดนามิก เพื่อรองรับความต้องการการฝึกขนาดต่างๆ
3. การแยกความล้มเหลวและการกู้คืนตัวเอง: เคอร์เนลที่ผิดพลาดมักทำให้เกิดการเข้าถึงหน่วยความจำ CUDA ที่ผิดกฎหมาย ด้วยการแยกระดับกระบวนการ รับประกันว่าความล้มเหลวของงานเดียวจะไม่ส่งผลกระทบต่อระบบโดยรวม
4. ผลตอบรับจากสภาพแวดล้อมที่หลากหลาย: ให้ผลตอบรับที่มีโครงสร้างสำหรับการเรียนรู้แบบเสริมกำลัง ซึ่งเกินกว่าผลลัพธ์ไบนารี (ผ่าน/ล้มเหลว) เช่น ข้อมูลข้อผิดพลาดโดยละเอียดและสรุปการวิเคราะห์ประสิทธิภาพ

4.2 สถาปัตยกรรมระบบ

รูปที่ 3: ภาพรวมสถาปัตยกรรม KERNELGYM และกรอบการฝึก ภาพขวาแสดงสถาปัตยกรรมแบบกระจายเซิร์ฟเวอร์-โหนดงาน ฝั่งเซิร์ฟเวอร์รับผิดชอบการจัดตารางงาน แต่ละโหนดงาน GPU ดำเนินการประเมินโค้ดเคอร์เนลในกระบวนการย่อยที่แยกจากกัน การออกแบบนี้รับประกันความบริสุทธิ์ของสภาพแวดล้อมการประเมินและการแยกความล้มเหลว

4.3 กลไกตรวจสอบการโกงรางวัล

เพื่อจัดการกับพฤติกรรมการโกงที่กล่าวถึงก่อนหน้านี้ KERNELGYM ดำเนินการตรวจสอบตามการรัน: มันแทรกตรรกะการตรวจสอบในเส้นทางการเรียกใช้เคอร์เนล Triton บันทึกจริงๆ ว่าเคอร์เนล Triton ใดถูกดำเนินการในโหมดการฝึกและโหมดประเมินผล และเปรียบเทียบเวลาในการรันแบบ end-to-end ของทั้งสองโหมด หากโค้ดที่สร้างขึ้นไม่ได้เรียกใช้เคอร์เนล