Hai con số trên một GPU
Bạn đã biết cách thực hiện c = a + b trong Python. Bạn đã làm điều đó hàng nghìn lần. Ở đây, chúng ta sẽ nói về nó trong tensor PyTorch. Tensor chỉ là một mảng các số. Đặt nó lên GPU có nghĩa là mảng đó sống trong bộ nhớ của GPU thay vì bộ nhớ CPU thông thường. Khi a và b là hai tensor trên GPU, một dòng lệnh đó hoàn thành đủ nhanh đến mức bạn không bao giờ nghĩ về nó.
Bây giờ hãy thu nhỏ nó lại. Giả sử a và b là hai số float đơn lẻ, cả hai đều nằm trên GPU. Cùng một dòng lệnh. c = a + b. Điều gì thực sự chạy trên chip?
Câu trả lời là một kernel. Trong thế giới này, kernel là một chương trình nhỏ mà GPU chạy trên một phần dữ liệu nào đó. Không phải kernel hệ điều hành mà máy tính xách tay của bạn khởi động. Không phải kernel toán học từ sách giáo khoa đại số tuyến tính. Từ này được sử dụng lại nhiều và đó không phải lỗi của bạn. Trong thế giới GPU, kernel chỉ đơn giản có nghĩa là: một hàm nhỏ mà GPU được yêu cầu chạy, ngay lập tức, song song, trên dữ liệu bạn đưa cho nó.
Đến cuối bài viết này, bạn sẽ có thể nhìn vào một đoạn mã PyTorch và đếm xem GPU sẽ chạy bao nhiêu kernel. Nghe có vẻ như một mẹo nhỏ, và đúng là như vậy, nhưng đó cũng là bước đầu tiên để thoát khỏi cảm giác "GPU như một hộp đen", cảm giác khi mô hình của bạn chạy chậm và bạn không biết tại sao. Mọi câu hỏi bạn có thể đặt ra về hiệu suất GPU cuối cùng đều quay lại "kernel nào đã chạy và chúng đã làm gì". Vì vậy, đây là nơi chúng ta bắt đầu.
Kernel đầu tiên của bạn
Hãy làm cho a và b lớn hơn một chút. Lần này là tensor độ dài 8. Vẫn một dòng PyTorch: c = a + b.
Khi bạn chạy dòng này, CPU của bạn (máy thực hiện Python thực tế) bảo GPU: này, hãy chạy chương trình này trên dữ liệu này. Chỉ thị đó được gọi là một **lần khởi chạy. Những gì được khởi chạy là một kernel: một chương trình, sẵn sàng chạy. Bản thân các lần khởi chạy rất rẻ, mỗi lần vài micro giây. Những gì xung quanh* lần khởi chạy (dữ liệu đi đến GPU, kết quả quay về) là nơi chi phí thực sự nằm, và đó là thứ chúng ta sẽ đếm khi tiến hành.
Bên trong kernel, công việc thực tế được thực hiện bởi các công nhân nhỏ gọi là thread. Một GPU có hàng nghìn thread có sẵn. Đối với phép cộng độ dài 8 của chúng ta, 8 thread nhận công việc: thread 0 xử lý phần tử 0, thread 1 xử lý phần tử 1, v.v. cho đến thread 7. Mỗi thread chạy cùng một chương trình nhỏ: đọc một phần tử của a, đọc phần tử tương ứng của b, cộng chúng, ghi kết quả vào c.
(Trong thực tế, GPU khởi chạy các thread theo nhóm kích thước cố định gọi là warp, luôn là 32 thread trên card NVIDIA, và mặt nạ các thread thừa khi mảng của bạn không chia đều. Có thể bỏ qua an toàn cho bây giờ.)
Vậy chúng ta có một dòng PyTorch, một lần khởi chạy, một kernel, 8 thread thực hiện 8 phép cộng. Bây giờ hãy đếm những gì thực sự đi qua chip. Để thực hiện phép cộng, mỗi thread cần phần tử của nó từ a và phần tử của nó từ b. Đó là 8 lần đọc từ a và 8 lần đọc từ b. Sau đó, mỗi thread ghi kết quả của nó vào c. Đó là 8 lần ghi.
Những lần đọc và ghi đó đi đến bộ nhớ lớn nằm ngay cạnh chip GPU. Trên các card trung tâm dữ liệu (A100, H100), bộ nhớ đó được gọi là HBM (bộ nhớ băng thông cao). Trên các card tiêu dùng (RTX 4090) và các T4 kiểu Colab mà mọi người thực sự có thể thử điều này, nó được gọi là GDDR. Dù bằng cách nào, đó là bộ nhớ nhanh nằm cạnh chip, và chúng ta sẽ chỉ gọi nó là bộ nhớ GPU. Nó nhanh, nhưng không miễn phí, và mỗi lần truy cập vào nó đều tốn kém.
Một kernel = một lần khởi chạy = một lần truyền dữ liệu. Bất kể kernel làm gì bên trong, các lần đọc và ghi ở các cạnh của nó (các lần truy cập ra bộ nhớ GPU để lấy đầu vào, lần quay lại để ghi đầu ra) là phần tốn kém. Đó là toàn bộ hình dạng.
Không có gì thay đổi khi các tensor lớn hơn. Cùng một dòng PyTorch, cùng một kernel, chỉ có nhiều thread hơn. Nếu a và b mỗi cái có một triệu phần tử, GPU khởi chạy cùng một kernel với một bầy thread lớn hơn. Phép toán mở rộng, số byte di chuyển mở rộng, mô hình tinh thần không thay đổi. Một dòng, một kernel.
Điều gì xảy ra giữa hai op
1c = (a + b).relu()
Bạn biết Python đủ để hiểu đây là hai thao tác, một phép cộng và sau đó là một relu, được xâu chuỗi. Trong trình thông dịch, đó là hai lệnh gọi hàm. Trên GPU, trong PyTorch eager, đó là hai lần khởi chạy kernel: một cho phép cộng, một cho relu. Cho đến nay, vẫn chưa có gì đáng ngạc nhiên.
Điều thực sự thú vị là những gì xảy ra giữa hai kernel.
Khi phép cộng kết thúc, nó phải đặt kết quả của nó vào một nơi nào đó. Nơi đó là bộ nhớ GPU. Phép cộng ghi toàn bộ một mảng trung gian (gọi nó là tmp) ra bộ nhớ. Sau đó, một lát sau, relu khởi chạy, và công việc đầu tiên của nó là đọc lại cùng một mảng tmp đó từ bộ nhớ. Nó đọc toàn bộ, áp dụng relu cho từng phần tử, ghi kết quả ra c.
Đếm lưu lượng bộ nhớ cho hai kernel đó:
- Phép cộng: đọc
a, đọcb, ghitmp. Ba lần truyền dữ liệu kích thước mảng. - Relu: đọc
tmp, ghic. Thêm hai lần nữa.
Tổng cộng năm lần truyền dữ liệu kích thước mảng. So sánh với phép cộng độ dài 8 một mình từ phần trước, chỉ có ba lần. Thêm .relu() vào chuỗi không chỉ tốn kém cho việc tính toán của relu. Nó tốn kém cho một vòng đi-về đầy đủ của mảng qua bộ nhớ GPU, bởi vì tmp phải được ghi ra chỉ để kernel tiếp theo có thể đọc lại nó.
Không có gì được lưu trong bộ nhớ đệm. tmp không được giữ lại trong thanh ghi hoặc bộ nhớ đệm cục bộ nhanh. Nó đi ra bộ nhớ GPU (loại chậm và xa) và quay lại ngay lập tức. Hai kernel là những người xa lạ với nhau. Chúng phải trao đổi dữ liệu thông qua phương tiện duy nhất mà cả hai đều biết cách nói chuyện: bộ nhớ GPU.
Tại sao PyTorch làm theo cách này? Bởi vì trong chế độ eager, khi bạn viết a + b, PyTorch chạy nó
gay lập tức**. Nó không biết bạn sắp gọi .relu() tiếp theo. Mỗi op được gửi đi ngay khi dòng Python của nó được thực thi. Không có kế hoạch, không có dự đoán trước. Mỗi op đứng một mình, tạo ra một mảng thực tế, và giao nó cho bất kỳ thứ gì đến tiếp theo thông qua bộ nhớ.
Đây là mẫu hình cần nắm giữ. Mọi giá trị trung gian trong mã PyTorch của bạn đều được ghi vật lý ra bộ nhớ GPU và được đọc lại bởi op tiếp theo. Mọi giá trị. Đó là những gì "số lượng kernel" thực sự đo lường. Mỗi kernel thừa là một vòng đi-về khác mà dữ liệu của bạn phải thực hiện qua bộ nhớ GPU.
Fusion: hai op, một kernel
Hãy tưởng tượng một kernel duy nhất thực hiện tất cả điều này trong một lần: đọc phần tử của nó từ a, đọc phần tử của nó từ b, cộng chúng, áp dụng relu cho kết quả (tất cả bên trong kernel, trên không gian làm việc tạm thời nhỏ cho mỗi thread không bao giờ rời khỏi chip), và chỉ sau đó ghi giá trị cuối cùng vào c. Giá trị trung gian (a + b) vẫn tồn tại, nhưng chỉ bên trong kernel, trong không gian làm việc riêng tư của mỗi thread. Nó không bao giờ được ghi ra bộ nhớ GPU. tmp, như một mảng thực tế, hoàn toàn không tồn tại.
Đếm các lần truyền dữ liệu bây giờ. Đọc a: 1 lần mỗi phần tử. Đọc b: 1 lần mỗi phần tử. Ghi c: 1 lần mỗi phần tử. Ba lần truyền dữ liệu kích thước mảng. Cùng một phép toán như phiên bản hai kernel, nhưng ít hơn hai vòng đi-về.
Ở độ dài 8, đây là một sai số làm tròn. Không ai quan tâm. Ở độ dài 1 triệu, hoặc 100 triệu, những vòng đi-về thừa đó trở thành một phần lớn của thời gian chạy, và đồng hồ thực tế phản ánh điều đó. \Tại sao** lưu lượng bộ nhớ lại chiếm ưu thế như vậy là toàn bộ chủ đề của Bài viết 2 trong loạt bài này, vì vậy tôi sẽ để "tại sao" ở đây. Điểm chính ngay bây giờ chỉ là: cùng một phép toán, ít lần truyền hơn, nhanh hơn trong thực tế.
Thủ thuật đó, kết hợp các op mà lẽ ra sẽ là các kernel riêng biệt thành một kernel duy nhất để giá trị trung gian không bao giờ phải đến bộ nhớ GPU, có một cái tên. Nó được gọi là fusion. Đó là toàn bộ từ. Đó là toàn bộ ý tưởng.
Bây giờ đến phần khó xử. Viết kernel kết hợp đó bằng tay trông có vẻ dễ đối với add + relu. Hai op. Một dòng "tính toán" ở giữa. Nhưng mã PyTorc thực tế có hàng chục op được xâu chuỗi với nhau, mỗi op có hình dạng, kiểu dữ liệu và quy tắc broadcasting riêng. Viết một kernel kết hợp xử lý tất cả điều đó một cách chính xác là một công việc kỹ thuật thực sự. Bạn thường sẽ không viết các kernel elementwise thông thường này bằng tay.
Tin tốt: PyTorch đi kèm với một công cụ thực hiện việc viết lại cho bạn, chính xác cho loại trường hợp này, một cách tự động. Nó được gọi là torch.compile.
Bạn gần như chắc chắn đã thấy torch.compile được sử dụng như thế này:
1model = torch.compile(model)
Một dòng. Ai đó trên internet đã nói với bạn rằng nó làm cho mọi thứ nhanh hơn. Đây là những gì nó thực sự làm, bằng tiếng Việt đơn giản: thay vì chạy các thao tác của bạn từng cái một theo cách chế độ eager, torch.compile nắm bắt các thao tác tensor mà hàm của bạn thực hiện, tìm kiếm cơ hội để kết hợp chúng, và tạo ra mã được tối ưu hóa. Các lệnh gọi sau đó khớp với cùng các giả định có thể tái sử dụng mã đó.
Sự kết hợp mà chúng ta đã thực hiện trên giấy ở trên (add và relu, chia sẻ một kernel, tmp không bao giờ chạm vào bộ nhớ) chính xác là loại điều mà torch.compile sẽ làm với mã của bạn một cách tự động, miễn là các op đủ đơn giản. Khi mọi người nói torch.compile "làm cho PyTorch nhanh hơn", đây là một phần lớn của những gì họ muốn nói.
Đối với các trường hợp mà torch.compile không thể tự kết hợp (các op tùy chỉnh mà nó không nhận ra, các phép reduction không điển hình, các bố trí bộ nhớ kỳ lạ), ai đó vẫn phải viết kernel bằng tay. Đó là những gì các công cụ như Triton và CUDA dành cho. Một bài viết riêng.
Tự mình nhìn thấy nó
Mọi thứ ở trên đã đếm kernel trên giấy. Đã đến lúc đếm chúng trên một GPU thực tế. Nếu bạn có bất kỳ máy nào có GPU CUDA (một máy trạm, một notebook Colab, một phiên bản đám mây), bạn có thể tự chạy điều này trong vài phút.
Công cụ là torch.profiler. Nó được tích hợp trong PyTorch. Tất cả những gì nó làm là ghi lại những gì GPU thực sự đã làm trong khi mã của bạn chạy, và trả lại cho bạn một bảng bạn có thể đọc.
Bước 1: phiên bản eager.
Bọc dòng hai op trong một hàm để có thứ gọi:
1import torch2from torch.profiler import profile, ProfilerActivity34def add_relu(a, b):5 return (a + b).relu()67a = torch.randn(1_000_000, device="cuda")8b = torch.randn(1_000_000, device="cuda")910with profile(activities=[ProfilerActivity.CUDA]) as prof:11 add_relu(a, b)12 torch.cuda.synchronize()1314print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
torch.cuda.synchronize() chỉ ở đó để đảm bảo GPU đã hoàn thành trước khi chúng ta đọc thời gian. Công việc GPU chạy không đồng bộ, và nếu không có đồng bộ, đôi khi bạn sẽ đo chi phí khởi chạy thay vì công việc kernel thực tế.
Bước 2: đọc đầu ra.
Đầu ra profiler thực tế của bạn sẽ có nhiều hàng hơn bạn có thể mong đợi. Một loạt các dòng cấp phát bộ nhớ và sổ sách kế toán PyTorch sẽ được trộn vào. Các hàng chúng ta quan tâm là các kernel CUDA, các hàm thực tế mà GPU đã chạy. Tìm các hàng có kernel trong tên. Hai trong số chúng sẽ trông đại khái như thế này:
1vectorized_elementwise_kernel<...BinaryFunctor_add...> 1 12us2vectorized_elementwise_kernel<...threshold_kernel...> 1 9us
Tên template chính xác thay đổi giữa các phiên bản PyTorch (relu thường hiển thị là threshold vì đó là op cơ bản, và add đôi khi là CUDAFunctor_add). Đừng cố gắng phân tích toàn bộ. Chỉ cần đếm số hàng. **Hai hàng. Hai kernel.*\* Một cho phép cộng, một cho relu. Chính xác như những gì chúng ta đã nói sẽ xảy ra ở một phần trước.
Bước 3: phiên bản đã biên dịch.
Thay đổi một dòng. Bọc hàm trong torch.compile:
1compiled = torch.compile(add_relu)
Bây giờ có một cái bẫy đáng biết trước khi bạn chạy nó. Lần gọi \đầu tiên đến một hàm `torch.compile`d rất chậm, đôi khi chậm một cách kỳ lạ, bởi vì đó là lúc `torch.compile` thực hiện công việc của nó: phân tích mã của bạn, tìm ra những gì để kết hợp, tạo ra kernel kết hợp. Nếu bạn profile lệnh gọi đầu tiên, bạn đang đo bước biên dịch, không phải kernel. Vì vậy, mẫu hình là: gọi nó một lần để làm nóng, bỏ kết quả đi, sau đó** profile.
1compiled(a, b) # warm-up, throw away2torch.cuda.synchronize()34with profile(activities=[ProfilerActivity.CUDA]) as prof:5 compiled(a, b)6 torch.cuda.synchronize()78print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
Bước 4: đọc đầu ra một lần nữa.
Lần này:
1triton_poi_fused_add_relu_0 1 14us
Một hàng. Một kernel. Tên thậm chí còn cho bạn biết nó đã làm gì: fused add và relu. Cùng một phép toán như trước, một lần khởi chạy thay vì hai.
Bạn vừa làm điều mà bài viết này đã nói đến, trong một câu: bạn yêu cầu PyTorch kết hợp hai op thành một kernel, xem profiler, và xác nhận số lượng giảm từ hai xuống một. Fusion, ngoài đời thực, trên máy của bạn.
Nếu bạn muốn thấy nó ấn tượng hơn, hãy thử nó với một vài kích thước tensor khác nhau. Ở độ dài 100, cả hai phiên bản chạy nhanh đến mức sự khác biệt bị mất trong nhiễu. Ở độ dài 10 triệu hoặc 100 triệu, phiên bản đã biên dịch bắt đầu vượt trội rõ rệt, bởi vì vòng đi-về mà chúng ta cắt giảm là một phần thực tế của công việc ở quy mô đó.
Đếm kernel không còn là lời khuyên trừu tượng nữa. Bạn có một cách để kiểm tra.
Tổng kết
Đây là toàn bộ vấn đề trong một lần.
Mã PyTorch của bạn, khi chạy trên GPU, biến thành một chuỗi các kernel. Mỗi kernel là một lần khởi chạy, một lần truyền dữ liệu, một vòng đi-về qua bộ nhớ GPU để lấy đầu vào và ghi đầu ra. Các op đơn giản trở thành một kernel. Các chuỗi op trở thành một kernel \mỗi op** theo mặc định, với các giá trị trung gian đi vòng qua bộ nhớ ở giữa. torch.compile có thể kết hợp các chuỗi đơn giản cho bạn để các giá trị trung gian đó không bao giờ chạm vào bộ nhớ. Ít kernel hơn thường có nghĩa là ít lưu lượng bộ nhớ hơn. Và ít lưu lượng bộ nhớ hơn thường có nghĩa là nhanh hơn.
