Profiling trong PyTorch (Phần 3): Tối ưu hóa hiệu năng cho cơ chế Attention

Quay lại danh sách bài viết

Loạt bài "Profiling in PyTorch" (Phân tích hiệu năng trong PyTorch) nhằm giúp bạn làm quen với việc đọc các dấu vết (traces) và bảng dữ liệu từ trình phân tích hiệu năng (profiler). Trong Phần 1, chúng ta đã phân tích các phép toán cơ bản như cộng và nhân. Chúng ta đã thấy cách bảng dữ liệu của profiler làm lộ ra các "điểm nóng" (hotspots), và cách dấu vết của profiler hiển thị thứ tự thực thi của một thuật toán theo thời gian.

Trong Phần 2, chúng ta đã gói các phép cộng và nhân đó vào một lớp linear của torch. Sau đó, chúng ta xếp chồng nhiều lớp linear lên nhau (mạng perceptron đa lớp) và phân tích hiệu năng của nó. Trong quá trình đó, chúng ta cũng đã phân tích các kernel được hợp nhất (fused) và tinh chỉnh thủ công.

Từ góc độ kiến trúc Transformer, bước logic tiếp theo để chúng ta phân tích là một thuật toán nền tảng khác: attention. Mặc dù nổi tiếng với độ phức tạp thời gian bậc hai, nhưng có rất nhiều thủ thuật thông minh tồn tại để giảm thiểu vấn đề đó và tăng tốc độ. Mục tiêu của chúng ta ở đây không phải là trình bày chi tiết mọi thủ thuật. Thay vào đó, chúng ta muốn xem mỗi thủ thuật trông như thế nào dưới góc nhìn của profiler.

Các tập lệnh cho bài blog này nằm tại đây: 04_a_naive_attention.py, 04_b_inplace_ops_attention.py, 04_c_sdpa_attention.py, và 04_d_kernels_attention.py. Giống như trước đây, bạn nên mở chúng trong một tab riêng và xem qua mã nguồn khi đọc. Chúng tôi sử dụng GPU NVIDIA A100-SXM4-80GB để chạy các tập lệnh này. Việc thiết lập GPU trên cơ sở hạ tầng của Hugging Face và thử nghiệm với các tập lệnh bằng Dev Mode với Spaces thực sự rất dễ dàng. Bạn cũng có thể chạy các tập lệnh này với pipeline Hugging Face Jobs.

Naive attention (Attention cơ bản)

Attention hoạt động với các Queries (q), Keys (k), và Values (v). Sự tương tác giữa chúng có thể được viết thành một chuỗi các bước ngắn gọn:

Vậy attention thực chất là một tập hợp các phép toán nguyên thủy. Một số trong đó chúng ta đã biết (các phép nhân ma trận - matmuls), và những phần còn lại rất dễ nhận ra. Hãy viết một module attention cơ bản trong PyTorch và phân tích hiệu năng của nó.

Trước khi mở dấu vết, hãy thực hiện bài tập quen thuộc của chúng ta là dự đoán những gì sẽ thấy. Khi theo dõi (trace) quá trình forward của module này, chúng ta kỳ vọng:

Hình 1 hiển thị làn CPU của profile (làn GPU được thu gọn để không làm chúng ta bị choáng ngợp). Bên trong attn_fwd (lệnh gọi forward đã được chú thích của chúng ta), chúng ta có thể thấy chính xác các phép toán mà chúng ta đã dự đoán. Phép nhân ma trận (matmul) giờ đã là một người bạn cũ, và các phép toán mới rất dễ nhận ra:

Bây giờ, hãy mở rộng làn GPU để xem những kernel nào thực sự đã được khởi chạy.

Hình 2 hiển thị làn GPU bên cạnh làn CPU. Hãy phóng to vào một khối attn_fwd duy nhất trên làn GPU để xem xét từng kernel một.

Hình 3 cho phép chúng ta đọc các kernel riêng lẻ cho một bước phân tích:

Năm trong số này là những gì chúng ta mong đợi. Phép sao chép bộ nhớ (memory copy) là phần khác biệt, vậy nó đến từ đâu? Manh mối nằm ở chỗ PyTorch có các phép toán thực hiện tại chỗ (in-place operations). Khi bạn thao tác trên một tensor theo cách thông thường (out-of-place), PyTorch thường tạo ra một bản sao, áp dụng phép toán yêu cầu lên đó và trả về bản sao đó. Theo chuỗi các phép toán, "thủ phạm" ở đây chính là masked_fill của chúng ta.

Điều gì sẽ xảy ra nếu chúng ta thay thế nó bằng một phép toán in-place?

Naive attention với in-place causal masking

Tất cả những gì chúng ta thay đổi là masked_fill thành masked_fill_ (lưu ý dấu gạch dưới ở cuối, quy ước của PyTorch cho các phép toán in-place), và chúng ta chạy cùng một tập lệnh.

Hãy xem dấu vết và kiểm tra xem có gì thay đổi không.

Phiên bản in-place (Hình 5) gói ít phép toán CPU hơn nhiều trong bước masking so với phiên bản out-of-place (Hình 4). Đây là một tín hiệu đáng mừng. Hãy mở rộng làn GPU để xác nhận những gì đã xảy ra ở đó.

Trên làn GPU, kernel Memcpy đã biến mất hoàn toàn (Hình 6 và 7). Chỉ với một dòng thay đổi, chúng ta đã loại bỏ được một kernel khỏi mỗi lượt forward pass. Điều này có vẻ không đáng kể, nhưng hãy nhớ đây chỉ là một phép toán attention đơn lẻ. Trong bối cảnh của một mô hình lớn dựa trên transformer (LLMs, Diffusion models, v.v.), nó lặp lại một lần cho mỗi lớp, và có rất nhiều lớp, vì vậy mức tiết kiệm sẽ tăng lên nhanh chóng (và nếu nó giúp bạn được tăng lương, việc chia sẻ ít nhất 10% với chúng tôi là điều công bằng).

Out-of-place là mặc định của PyTorch vì một lý do. Để tính toán gradient, autograd phải ghi nhớ các giá trị tensor mà nó đã thấy trong lượt forward pass, vì nhiều công thức backward cần sử dụng lại chúng. Một phép toán in-place sẽ ghi đè lên các giá trị đó trong bộ nhớ, vì vậy backward pass sẽ đọc sai số liệu. Do chúng ta chạy forward dưới torch.no_grad, in-place là an toàn cho chúng ta, vì không có backward pass và không có gì bị hỏng. Cũng cần lưu ý rằng các phép toán in-place không chỉ tiết kiệm thời gian (như chúng ta thấy trong trường hợp này) mà còn tiết kiệm bộ nhớ (do không có bản sao thừa), điều này rất tuyệt vời cho các tensor lớn như logits!

Scaled Dot Product Attention

Chúng ta vừa xây dựng attention từ các thành phần nguyên thủy, và thậm chí đã loại bỏ được một Memcpy. Tin tốt là đội ngũ PyTorch đã làm tất cả những điều này cho chúng ta, và đóng gói toàn bộ pipeline vào một hàm duy nhất:

Dòng này thay thế cho module viết tay của chúng ta, và is_causal=True thậm chí còn giúp chúng ta không cần phải tự xây dựng mask. Rất đáng để dừng lại một chút để trân trọng những gì mà lệnh gọi này che giấu. Và nó che giấu nhiều thứ hơn là chỉ các dòng mã. Scaled Dot Product Attention (SDPA) không có một cách triển khai duy nhất. Bên dưới, nó điều phối đến một trong nhiều backend và chọn backend nhanh nhất hỗ trợ các đầu vào của chúng ta (dtype, head dimension, mask, phần cứng, v.v.).

Hướng dẫn chính thức về SDPA hướng dẫn chúng ta qua quá trình lựa chọn này, và bản thân các backend được liệt kê trong enum torch.nn.attention.SDPBackend:

Thông thường SDPA tự chọn cho chúng ta, nhưng chúng ta có thể ghim một backend cụ thể bằng trình quản lý ngữ cảnh torch.nn.attention.sdpa_kernel. Đây là những gì chúng ta làm trong các tập lệnh của mình. Điều này cho phép chúng ta phân tích từng backend riêng lẻ và đọc xem chúng xuất hiện khác biệt như thế nào trong dấu vết. Hãy thực hiện từng cái một.

Math backend

Trước khi mở bất cứ thứ gì, hãy dự đoán. Chúng ta đã thay thế attention viết tay (matmul, mul, mask, softmax, matmul) bằng một dòng lệnh duy nhất, vì vậy chúng ta nên kỳ vọng dấu vết sẽ đơn giản và nhanh hơn. Ít kernel hơn, ít điều phối CPU hơn, thậm chí có thể là một kernel được hợp nhất. Hãy kiểm tra bảng profiler trước.

Đây là bất ngờ đầu tiên của chúng ta, dòng lệnh đơn đó chậm hơn 3,7 lần.

Mở dấu vết (Hình 9) cho thấy lý do tại sao chuông báo động vang lên, math backend khởi chạy 20 GPU kernel cho mỗi lượt forward thay vì 5 kernel như với cách triển khai attention cơ bản của chúng ta (Hình 8). Điều này trái ngược với những gì chúng ta dự đoán. Hãy tìm hiểu lý do tại sao điều này xảy ra.

Tensor cores bị bỏ trống

Trong Phần 2, chúng ta đã học cách đọc tên kernel như một dấu vân tay. Hãy sử dụng thói quen đó ở đây:

Các GPU A100 mà chúng ta sử dụng để ghi lại các dấu vết này được trang bị Tensor Cores, phần cứng chuyên dụng cho các phép nhân ma trận tăng tốc, được biết là nhanh hơn nhiều so với các nhân CUDA thông thường. Để hiểu tại sao điều đó quan trọng ở đây, bạn cần biết những gì nằm bên trong một GPU. Streaming Multiprocessor (SM) là đơn vị tính toán của GPU, và mỗi SM có hai loại đơn vị số học: nhân CUDA và Tensor Cores. Nhân CUDA là mục đích chung và xử lý một vài phần tử tại một thời điểm, trong khi Tensor Cores nhân và tích lũy toàn bộ một ma trận nhỏ trong một lệnh duy nhất. Vì vậy, câu hỏi rất đơn giản: "Mỗi backend có thực sự đang sử dụng đường dẫn nhanh (fast path) không?"

Tên các kernel đã trả lời câu hỏi đó. s16816 trong kernel cơ bản (Hình 10) là chữ ký của một phép nhân ma trận Tensor Core bfloat16 (lệnh Tensor Core 16x8x16), vì vậy phiên bản cơ bản đang đi trên đường dẫn nhanh. sgemm (Hình 11) là phép nhân ma trận độ chính xác đơn (FP32) cổ điển chạy trên các nhân CUDA thông thường. Nói cách khác, math backend không bao giờ chạm vào Tensor Cores: để đánh đổi tốc độ lấy độ chính xác số học, nó nâng cấp (upcast) các tensor lên FP32 (làm tăng gấp đôi dữ liệu được di chuyển, ngay cả khi đầu vào là bf16) và quay lại sử dụng các nhân CUDA chậm hơn.