เจาะลึก PyTorch Profiling: จาก nn.Linear สู่ Fused MLP ที่เร็วขึ้น

เจาะลึก PyTorch Profiling: จาก nn.Linear สู่ Fused MLP ที่เร็วขึ้นการทำความเข้าใจเบื้องหลังการทำงานของ PyTorch เป็นกุญแจสำคัญในการเพิ่มประสิทธิภาพโมเดล AI ให้ถึงขีดสุด ในซีรีส์ "Profiling in PyTorch" นี้ เราจะค่อยๆ สร้างทักษะการอ่าน PyTorch profiler traces เพื่อนำไปสู่การปรับปรุงประสิทธิภาพ ในตอนที่ 1 เราได้เรียนรู้วิธีอ่าน profiler traces เบื้องต้น และในตอนนี้ เราจะก้าวไปอีกขั้นด้วยการเปลี่ยนจากการใช้ torch.add(torch.matmul(x, w), b) มาเป็น nn.Linear ซึ่งเป็นส่วนประกอบพื้นฐานที่ทุกโมเดล Deep Learning ต้องใช้จาก matmul-add สู่ nn.Linear: การทำงานเบื้องหลังnn.Linear คือ wrapper ที่ครอบการทำงานของการคูณเมทริกซ์และการบวกที่เราได้เห็นในตอนที่ 1 โดยมีน้ำหนัก (weight) และไบแอส (bias) เป็นพารามิเตอร์ของตัวเอง และมีเมธอด forward ที่ผู้ใช้ PyTorch คุ้นเคย การดำเนินการนี้สามารถเขียนได้ในรูป:y = x @ w + bเมื่อ x คืออินพุต, w คือน้ำหนัก และ b คือไบแอสเมื่อเราใช้ profiler กับ nn.Linear เราจะเห็นการทำงานที่ซับซ้อนกว่าที่คิดเล็กน้อยการทำงานของ Transpose (aten::t)เมื่อเจาะลึก profiler trace เราจะสังเกตเห็น aten::t (transpose) ก่อน aten::addmm (การคูณและบวก) สิ่งนี้บ่งบอกว่า nn.Linear ทำการ transpose น้ำหนัก (w) ก่อนนำไปคูณกับอินพุต (x)สิ่งสำคัญที่ควรทราบ: aten::t ไม่ได้ทำการคัดลอกหรือจัดเรียงข้อมูลใหม่จริง ๆ แต่เป็นการแก้ไขเมทาดาทา (shape และ stride) ของเทนเซอร์บน CPU เพื่อให้มองเห็นเมทริกซ์ที่ถูก transpose แล้ว โดย ไม่ได้เรียกใช้ GPU kernel คุณสามารถตรวจสอบได้จากการดูเลน GPU ใน trace หรือเวลาที่ใช้บน CUDA สำหรับ aten::tทำไมไม่มี Kernel แยกสำหรับ mul และ add?คุณจะสังเกตเห็นว่าไม่มี aten::add แยกออกมาใน dispatch chain ของ nn.Linear ทั้งนี้เป็นเพราะการบวกไบแอสได้ถูก รวมเข้ากับการคูณเมทริกซ์ (GEMM) แล้ว ผ่านสิ่งที่เรียกว่า epilogueEpilogue คือการคำนวณเล็กๆ น้อยๆ ที่ GEMM kernel ทำก่อนที่จะเขียนผลลัพธ์กลับไปยัง HBM (High Bandwidth Memory) การบวกไบแอส การใช้ activation function หรือการ scale เป็นตัวอย่างของ epilogue การทำเช่นนี้ช่วยหลีกเลี่ยงการโหลดหรือเขียนข้อมูลไปยัง HBM ซ้ำสองครั้ง ซึ่งการรับส่งข้อมูลผ่านหน่วยความจำเป็นส่วนที่ทำให้การทำงานช้าลงnn.Linear จะเรียกใช้ torch.nn.functional.linear ซึ่งจะเรียก aten::linear ต่อไป aten::linear จะตรวจสอบอินพุต หากพบว่ามีการส่ง bias เข้ามา จะเลือกใช้ aten::addmm(bias, x, weight) แทนการทำ matmul และ add แยกกันaddmm สามารถคำนวณได้ดังนี้: y = x @ w + bKernel ของ cuBLAS GEMM ที่ทำงานบน GPU มีฟังก์ชันสำหรับ bias-add ในตัว และนั่นคือ kernel ที่ aten::addmm เลือกใช้ การบวก (add) จึงไม่ปรากฏเป็น kernel แยก เพราะเป็นส่วนหนึ่งของการเขียนผลลัพธ์ของ matmul kernel ซึ่งก็คือ epilogue นั่นเองtorch.compile ช่วยอะไรกับ Linear Layer เดียวได้หรือไม่?แม้ว่า nn.Linear จะใช้งาน kernel ที่มีประสิทธิภาพอยู่แล้ว การลองคอมไพล์ forward call ด้วย torch.compile และดู profiler trace จะพบว่า:GPU Kernel เดียวกัน: ยังคงใช้ cuBLAS GEMM kernel เดิมCPU Op เดียวกัน: ยังคงเป็น aten::addmmมีแถวเพิ่มเติมบน CPU: ที่เกี่ยวข้องกับการทำงานของ compileนี่เป็นจุดที่สำคัญที่ควรทำความเข้าใจ: หลายครั้งที่เรารู้สึกว่าโมเดลทำงานช้า สิ่งแรกที่เรามักจะนึกถึงคือ torch.compile แต่สำหรับ GEMM เพียงหนึ่งครั้งที่มี bias, compile แทบจะไม่มีอะไรให้ทำ การที่ compile จะแสดงประสิทธิภาพได้ ต้องมี การดำเนินการมากกว่าหนึ่งอย่าง เพื่อให้เกิดการ fuse (หลอมรวม) ได้การซ้อน Linear Layer สามชั้น: สร้าง Multilayer Perceptron (MLP)เมื่อเรานำ nn.Linear มาซ้อนกันสามชั้นพร้อม activation function ตรงกลาง เราจะได้โครงสร้าง Multilayer Perceptron (MLP) ซึ่งเป็นส่วนประกอบสำคัญของโมเดล Deep Learning จำนวนมากสิ่งที่คาดหวังก่อนเปิด Traceเมื่อพิจารณา MLP ที่มีการใช้งาน GeGLU activation เราคาดว่า:จะมีการเรียก aten::linear สามครั้ง (สำหรับ nn.Linear แต่ละชั้น)จะมีการเรียก kernel สำหรับ activation (GeLU) และการคูณ (mul) สองครั้งการตั้งสมมติฐานก่อนดู trace เป็นนิสัยที่มีประโยชน์อย่างยิ่งในการทำ profiling เพราะเราจะใช้ trace เพื่อยืนยันหรือหักล้างสมมติฐานของเราผลลัพธ์จาก Profilerจาก profiler trace เราจะเห็นว่า:ต่อ 1 forward pass ของ MLP, GPU จะรัน 5 kernelsสำหรับ nn.Linear จะมีการเรียก cudaOccupancyMaxActiveBlocksPerMultiprocessor ก่อนการ launch kernel (เป็นการคำนวณขนาด grid ของ cuBLAS)สำหรับ pointwise ops (GeLU, mul) จะเป็นการ launch kernel โดยตรงOps ที่ไม่มีการ launch kernel: aten::t, aten::transpose, aten::reshape, aten::view, aten::asstrided และ aten::unsafe_view จะไม่ launch kernel ใดๆ เลย โดยจะแสดงเวลา CUDA เป็น 0.000us ในตาราง เพราะทำเพียงแค่การแก้ไขเมทาดาทาของเทนเซอร์ (shape และ stride) บน CPU เท่านั้น แม้ในตารางจะเห็นชื่อ op จำนวนมากต่อหนึ่ง linear layer แต่มีเพียง mm เท่านั้นที่ไปทำงานบน GPU จริงๆทำไม GEMM Kernel ถึงมีสองแบบ?MLP จะทำการ flatten ข้อมูลจาก [batch, seq, dim] เป็น [batch seq, dim] ก่อนทำการ matmul ตัวอย่างเช่น หาก batch = 64 และ seq = 128 จะได้ 8192 (64 128)แม้ว่า GEMM ทั้งสามตัวจะมีการคำนวณ FLOPs เท่ากัน (ประมาณ 38.7 GFLOP แต่ละอัน) แต่ downproj จะเร็วกว่าประมาณ 10% สาเหตุมาจาก shape ที่แตกต่างกัน (N=768 เทียบกับ 3072) ทำให้ cuBLAS เลือกใช้ tiling ที่แตกต่างกัน (128×256 กับ pipeline stages64x3 ที่ลึกกว่า) ซึ่งให้ผลลัพธ์ที่ดีกว่าสำหรับการ reuse ข้อมูลใน shape นั้นๆนี่คือเหตุผลที่ตาราง profiler แสดง GEMM สองแถว: แถว 128x128 สำหรับ gate+up และแถว 128x256 สำหรับ downtorch.compile ทำอะไรได้บ้าง?เมื่อเราคอมไพล์ forward method และดู trace:ในโหมด eager, nn.Linear แต่ละตัวจะถูกขยายเป็น chain ของ dispatcher ops (aten::linear → aten::t → aten::transpose → aten::matmul → aten::reshape → aten::mm) ซึ่งเป็น high-level wrappers ที่ ATen ใช้ก่อนจะถึง GEMM จริงtorch.compile ขจัด chain นี้ออกไป เมื่อกราฟที่คอมไพล์ทำงาน จะไม่มี linear, matmul, transpose หรือ reshape อีกต่อไป โดย ops เหล่านี้จะถูกรวมเข้ากับการเรียก mm โดยตรงเราจะเห็นการเรียก aten::mm ภายนอกเพียงสามครั้ง และหลักฐานที่พิสูจน์ว่าเป็น GEMM เดียวกันคือ ชื่อ kernel ที่เหมือนกันทุกประการ กับโหมด eager (...128x128...stages32x5tn สำหรับ gate และ up, และ ...128x256...stages64x3tn สำหรับ down)Fused Triton Kernel: หัวใจของการเพิ่มประสิทธิภาพtorch.compile สามารถ fuse การดำเนินการหลายอย่างเข้าด้วยกัน สร้างเป็น kernel ที่มีประสิทธิภาพสูงขึ้น ซึ่งมักจะใช้ Triton kernel โดยเฉพาะอย่างยิ่งเมื่อมีการรวมการดำเนินการที่ซับซ้อนเข้าด้วยกัน เช่น การบวก bias, การใช้ activation function หรือการ transpose เข้าไปใน GEMM kernel ตัวเดียวการเข้าใจการทำงานเบื้องหลังของ PyTorch ผ่าน profiler จะช่วยให้นักพัฒนาสามารถระบุคอขวด (bottleneck) และปรับปรุงประสิทธิภาพของโมเดลได้อย่างตรงจุด โดยเฉพาะอย่างยิ่งการใช้ประโยชน์จาก torch.compile เพื่อสร้าง fused kernel ที่มีประสิทธิภาพสูงสุด#PyTorch #Profiling #DeepLearning #MLP #AIhttps://huggingface.co/blog/torch-mlp-fusion