Phân tích hiệu năng trong PyTorch (Phần 1): Hướng dẫn cho người mới bắt đầu về torch.profiler

Những gì bạn không thể phân tích, bạn không thể tối ưu hóa.

Cho dù bạn đang cố gắng tăng số lượng token mỗi giây cho một Mô hình Ngôn ngữ Lớn (LLM), cắt giảm vài mili giây thời gian suy luận, hay chỉ đơn giản là muốn hiểu tại sao vòng lặp huấn luyện của bạn chạy chậm hơn so với thông số kỹ thuật, con đường cuối cùng đều dẫn đến việc phân tích hiệu năng (profiling).

Tuy nhiên, việc phân tích hiệu năng có một “ngưỡng nhập môn” khá dốc. Các bản trace là những bức tường dày đặc các hình chữ nhật đầy màu sắc. Các sự kiện mang những cái tên đáng sợ. Hầu hết các hướng dẫn đều giả định rằng bạn đã biết cách đọc chúng. Vì vậy, ngay cả khi biết mình nên phân tích, việc mở một bản trace có thể mang lại cảm giác như một công việc nặng nề nên để lại sau (hoặc để cho người khác làm). Bài viết này, và loạt bài mà nó khởi đầu, là nỗ lực của chúng tôi nhằm hạ thấp ngưỡng nhập môn đó.

Đây là bài mở đầu của chuỗi bài Profiling in PyTorch, nơi chúng ta sẽ chậm rãi xây dựng kỹ năng đọc các bản trace của profiler và sử dụng nó để thúc đẩy tối ưu hóa. Kế hoạch như sau:

1. Phần 1 (bài này): Bắt đầu với thao tác đơn giản nhất có thể, một phép nhân ma trận theo sau là phép cộng bias, và học cách đọc những gì profiler trả về.
2. Phần 2: Mở rộng lên nn.Linear và một mạng MLP nhỏ, sử dụng các bản trace để thúc đẩy tối ưu hóa và xem xét các kernels bên dưới.
3. Phần 3: Áp dụng tất cả vào các Mô hình Ngôn ngữ Lớn với thư viện transformers.

Chúng tôi ghi lại hành trình này từ góc nhìn của một người mới bắt đầu. Không yêu cầu điều kiện tiên quyết nào ngoài kiến thức cơ bản về PyTorch. Hãy coi đây là một bài đọc thư giãn với một vài khoảnh khắc “Aha!”. Cấu trúc của bài viết được thiết kế theo hướng đặt câu hỏi: chúng ta mở một bản trace, hỏi “đợi đã, tại sao điều đó lại xảy ra?”, và đuổi theo câu trả lời cho đến khi hiểu ra vấn đề. Đến cuối bài, bạn sẽ biết:

• Cách thiết lập torch.profiler và những gì nó thực sự trả về.
• Cách đọc bảng profiler và bản trace (luồng CPU, luồng GPU và những khoảng trống đáng nghi ở giữa).
• Chuỗi sự kiện từ một lời gọi Python cho đến tận một CUDA kernel.
• Những gì thay đổi (và điều thú vị hơn là những gì không thay đổi) khi bạn sử dụng torch.compile.

Trước khi bắt đầu, hai định nghĩa sẽ giúp bạn dễ đọc nội dung bên dưới hơn:

1. Một GPU kernel là một chương trình chạy song song trên nhiều luồng (threads) của GPU.
2. CPU lập lịch và kích hoạt (schedule and launch) các kernel này.

Thông thường, bạn không cần tự viết GPU kernels; khi bạn sử dụng một thao tác trong PyTorch, nó sẽ tự động được dịch sang một hoặc nhiều kernel để thực hiện công việc trên GPU.

Với hai ý tưởng đó, hãy bắt đầu đặt câu hỏi.

Lưu ý: Đây là toàn bộ mã nguồn chúng tôi sử dụng cho bài viết: 01_matmul_add.py. Chúng tôi khuyên bạn nên mở mã nguồn này trong một tab riêng và đi qua từng bước code. Chúng tôi sử dụng GPU NVIDIA A100-SXM4-80GB để chạy các script.

Thao tác nhân ma trận và cộng ma trận

Như Tiến sĩ Sara Hooker đã nhận định chính xác, giống như con người chủ yếu được cấu tạo từ nước, các Mạng Thần kinh Sâu chủ yếu được cấu tạo từ các phép nhân ma trận. Vì chúng cơ bản như vậy, sẽ thật đáng tiếc nếu chúng ta bắt đầu hành trình phân tích hiệu năng với bất kỳ thứ gì khác.

def fn(x, w, b):
  return torch.add(torch.matmul(x, w), b)

Phép cộng ma trận cùng với phép nhân ma trận mô phỏng cách trọng số (weights) và bias tương tác trong một neuron. Phép cộng này sẽ giúp chúng ta hiểu cách nó mở đường cho việc biên dịch (compilation) ở phần sau của bài viết.

Để phân tích, chúng ta sẽ sử dụng mô-đun torch.profiler. Các bước thực hiện bao gồm:

1. Chuẩn bị mã cần phân tích (ở đây là def fn, bao bọc phép nhân và phép cộng ma trận).
2. Chú thích (Annotate) thuật toán. Điều này hoàn toàn tùy chọn, nhưng chúng tôi khuyên dùng. record_function chú thích hàm của chúng ta là matmul_add, giúp dễ dàng điều hướng trong các bản trace.

def step():
  with torch.profiler.record_function("matmul_add"):
    return fn(x, w, b)

3. Bao bọc mã bằng context manager torch.profiler.profile.

  with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,  # các hoạt động của cpu
        torch.profiler.ProfilerActivity.CUDA, # các hoạt động của gpu
    ],
  ) as prof:
    # Khuyến nghị chạy sự kiện nhiều lần để "làm nóng" (warm up) GPU
    for _ in range(5):
      step()
      prof.step()

4. Xuất kết quả profile.

# bảng profiler
prof.key_averages().table(sort_by="cuda_time_total", row_limit=15)

# bản trace profiler
prof.export_chrome_trace(trace_path)

Profiler xuất ra hai sản phẩm riêng biệt:

1. Bảng profiler (Profiler table): Cung cấp tóm tắt thống kê của thuật toán. Nó trả lời câu hỏi “Cái gì đang tốn nhiều thời gian nhất”. Điều này rất hữu ích để tìm ra các “điểm nóng” (hotspots) — những sự kiện tốn nhiều thời gian nhất, có thể là nút thắt cổ chai của pipeline, hoặc một sự kiện bị kích hoạt quá nhiều lần.
2. Bản trace profiler (Profiler trace): Cung cấp góc nhìn thực thi theo thời gian. Trả lời câu hỏi “Khi nào và Tại sao một thao tác xảy ra”, mô tả các hoạt động diễn ra trên CPU và GPU. Điều này hữu ích khi chúng ta muốn điều tra các kernel đã được kích hoạt, bất kỳ sự chậm trễ nào trong việc kích hoạt chúng, bất kỳ sự chồng lấp (overlap) nào giữa hoạt động của CPU và GPU, v.v.

Hãy xem cả hai hoạt động trong lần thực thi đầu tiên.

Khuyến nghị chạy script này trên máy có GPU.

uv run 01_matmul_add.py --size 64

Nếu bạn chạy script trên (trên máy có GPU), bạn sẽ tìm thấy thư mục traces/01_matmul_add với hai sản phẩm: 64_bf16_cold_eager.json 64_bf16_cold_eager.txt

File .txt chứa bảng profiler. Khi mở file, như trong Hình 1, bạn sẽ thấy một bảng lớn với cột đầu tiên là các sự kiện được kích hoạt trong phạm vi profile.

Các cột khác liên quan đến thời gian sự kiện tốn trên CPU hoặc GPU. Hãy xem sự kiện nào tốn nhiều thời gian nhất và thử suy luận xem sự kiện đó có thực sự nên tốn nhiều thời gian như vậy không. Điều quan trọng là phải nhìn vào cột “# of Calls” để biết sự kiện đó được kích hoạt bao nhiêu lần.

Nhân tiện, hãy nói về “Self CPU/CUDA” và “CPU/CUDA total”. Các cột “Self” đo thời gian chỉ dành cho chính sự kiện đó, không bao gồm các sự kiện con. Các cột “total” bao gồm cả sự kiện đó và tất cả các sự kiện con của nó. Vì vậy, nếu bạn nhìn vào “CPU total” của matmul_add, nó bao gồm thời gian thực hiện chính nó cộng với các sự kiện con mà nó kích hoạt. Đây là một sắc thái quan trọng cần lưu ý.

Nếu bạn nhìn vào hai dòng cuối của bảng, bạn sẽ thấy profiler cho biết: Self CPU time total: 2.314ms Self CUDA time total: 23.104us

Thời gian CPU tính bằng ms trong khi thời gian GPU tính bằng us. Để dễ hình dung, thời gian tiêu tốn trên GPU (kernel ampere_bf16_s16816gemm...) chiếm chưa đến 1% thời gian tiêu tốn trên CPU (thao tác matmul_add). GPU ở trạng thái nhàn rỗi hầu hết thời gian, đây là một dấu hiệu cảnh báo đỏ ngay lập tức. Lý do là GPU có thể tính toán một phép nhân ma trận nhỏ rất nhanh, vì vậy mã của chúng ta dành phần lớn thời gian để chuẩn bị các kernel, kích hoạt chúng trên GPU, gửi dữ liệu và thu thập kết quả. Khái niệm này được gọi là thuật toán bị giới hạn bởi chi phí quản lý (overhead-bound).

Cách dễ nhất để thoát khỏi trạng thái này là sử dụng các phép nhân ma trận lớn hơn.

uv run 01_matmul_add.py --size 4096 

Hai dòng cuối trong Hình 2 là: Self CPU time total: 4.908ms Self CUDA time total: 4.495ms

Cả hai thời gian đều tính bằng ms, nghĩa là chúng ta đã tạo ra nhiều thời gian chạy GPU hơn chỉ bằng cách tăng kích thước ma trận. Trong Hình 2, bạn cũng sẽ thấy thời gian CUDA nhiều nhất hiện giờ là do GPU kernel (ampere_bf16_s16816gemm_..) chứ không phải do thao tác CPU kích hoạt nó (matmul_add). Điều này có nghĩa là chúng ta thực sự đã chuyển từ trạng thái overhead-bound sang bị giới hạn bởi tính toán (compute-bound).

Bây giờ chúng ta chuyển sang trực quan hóa chuỗi điều phối (dispatch chain), nằm trong các sản phẩm .json. Bạn có thể tải chúng lên Perfetto UI để xem trace.

Bản trace 64x64

Trong Hình 3, chúng ta thấy bản trace cho phép nhân và cộng ma trận. Ở đây, chiều rộng của thanh biểu thị thời gian diễn ra sự kiện, phân cấp theo chiều dọc là phân cấp gọi hàm, luồng CPU biểu thị các sự kiện xảy ra trên CPU, trong khi luồng GPU hiển thị các lần thực thi kernel thực tế. Bạn cũng có thể nhận thấy những khoảng trống, đó là thời gian chờ hoặc thời gian nhàn rỗi.

Script được chạy với các cấu hình mặc định:

• size 64: Đầu vào, trọng số và bias có kích thước (64, 64).
• dtype bf16: Kiểu dữ liệu là bfloat16.
• no compile: Chúng ta không biên dịch các thao tác torch.
• no warmup: Chúng ta không làm nóng GPU trước khi phân tích.

Với Perfetto, chúng tôi khuyên bạn nên sử dụng bàn phím để truy cập trace nhanh hơn. Có thể sử dụng “W A S D” để điều hướng trace.

Có hai luồng trong Hình 4, một cho hoạt động CPU và một cho hoạt động GPU. Trong luồng CPU, bạn sẽ thấy ba bước profile (bắt đầu từ ProfilerStep#2). Điều này đến từ schedule.

schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)

wait bỏ qua các khởi tạo gây nhiễu (ProfilerStep#0), warmup chạy qua profiler mà không ghi lại (ProfilerStep#1), và active là những gì hiển thị trong trace.

Hãy cùng đóng vai thám tử để điều tra bản trace và đặt một số câu hỏi.

Tại sao ProfilerStep#2 lại tốn nhiều thời gian đến vậy?

Trong Hình 5, chúng ta thấy ProfileStep#2 tốn nhiều thời gian hơn các bước khác, và khi nhìn kỹ, bạn sẽ thấy một mô hình tương tự với chú thích matmul_add. Manh mối nằm bên trong chú thích, chứ không phải chính chú thích đó:

Khoảng ~228 µs trong Hình 6 là “cửa sổ chết” giữa lúc bắt đầu record_function("matmul_add") và lúc PyTorch thực sự điều phối aten::matmul. Điều này có thể xảy ra vì nhiều lý do, bao gồm phân bổ vùng làm việc (workspace allocations), các heuristic của cuBLAS (thư viện tăng tốc GPU độc quyền của NVIDIA để thực hiện các thao tác đại số tuyến tính cơ bản), hoặc tải mô-đun chậm (lazy module loading). Chúng ta có thể bỏ qua hoặc chạy thêm một số bước làm nóng trước khi phân tích (đây là cách làm tiêu chuẩn).

Trong phân tích hiệu năng, làm nóng (warmup) là khi bạn chạy các sự kiện một vài lần trước khi thực sự phân tích. Các công việc chuẩn bị của GPU là nỗ lực một lần mà chúng ta không muốn đưa vào bản profile. Trong ví dụ của chúng ta, chúng ta có hai giai đoạn làm nóng: một là lặp qua hàm trước khi vào profiler, và hai là bên trong profiler thông qua đối số warmup.

uv run 01_matmul_add.py --warmup

Bản Trace Perfetto cho 64x64 có Warmup

Trong Hình 7, chúng ta thấy mỗi bước profile tốn thời gian tương tự nhau, nhưng điều này không có nghĩa là chúng ta đã tối ưu hóa được các chi phí một lần đó. Chúng ta chỉ làm nóng để các chi phí đó không bị ghi lại trong bản profile.

Tại sao có một khoảng lệch ~2.5 ms giữa luồng CPU và luồng GPU?

Trong Hình 8, chúng ta thấy luồng CPU và GPU có một khoảng lệch khoảng 2.5 ms: đây là độ trễ sau khi CPU gửi các CUDA kernel và thời điểm chúng thực sự bắt đầu thực thi. Có người sẽ nghĩ giai đoạn làm nóng kết hợp với wait và warmup của schedule sẽ giữ cho GPU bận rộn và làm giảm khoảng lệch này.

Để tìm ra điều gì đang thực sự xảy ra, hãy thay đổi schedule một chút: - schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1) + schedule = torch.profiler.schedule(wait=0, warmup=0, active=3, repeat=1)

Hình 9 cho thấy có một Activity Buffer Request trong luồng GPU trước bất kỳ thao tác nào. Hãy phóng to thêm một chút.

Khi phóng to vào trace GPU, chúng ta nhận thấy các kernel matmul và add cho ProfileStep#0 diễn ra liên tiếp nhau, trong khi các kernel cho ProfileStep#1 có một khoảng trống ở giữa. Giải thích hợp lý nhất là đã xảy ra hiện tượng tràn bộ đệm (overflow of buffers), và một yêu cầu bộ đệm khác (yêu cầu cấp phát bộ nhớ trên VRAM của GPU) đã được phát ra trong quá trình thực thi kernel.

Cách tốt nhất để loại trừ các khả năng khác là phân tích nhiều lần lặp hơn. Chúng ta chạy với active=20.

Như Hình 11 cho thấy, chúng ta thấy xu hướng tương tự ở ProfileStep#1. Điều này khớp với những phát hiện trước đó, và chúng ta có thể kết luận an toàn rằng đó thực sự là một yêu cầu bộ đệm.

Chuỗi sự kiện

Trong Hình 12, chúng ta thấy các lời gọi CPU lồng nhau. Đây là một hình ảnh quan trọng để hiểu một chuỗi điều phối thực sự trông như thế nào.

Chúng ta bắt đầu với ProfileStep#<id>, bao bọc bước profiling. Nhờ việc chú thích, chúng ta thấy hàng matmul_add. matmul_add bao gồm hai lời gọi aten, một cho nhân ma trận và một cho cộng ma trận.

aten::matmul là mức điều phối ATen mà các lời gọi matmul của người dùng trong PyTorch sẽ dẫn đến. aten::mm là backend nhân ma trận 2D.

Một điểm thú vị là PyTorch sẽ gọi aten::bmm (nhân ma trận theo lô - batched matrix multiplication) nếu chúng ta thêm trục batch vào ma trận.

# thêm batch size là 8
x = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
w = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
b = torch.randn(8, args.size, args.size, device=device, dtype=dtype)

Trong Hình 13, khi thêm trục batch, aten::matmul hiện bao bọc một loạt các lời gọi CUDA runtime tiên quyết cùng với aten::bmm (thay vì aten::mm). Điều này cũng gợi ý về các heuristic mà cuBLAS cần thực hiện để điều phối kernel phù hợp nhất cho chương trình.

Tại sao matmul có một lời gọi CUDA runtime bổ sung?

Chúng ta nhận thấy đối với aten::mm có hai lời gọi CUDA Runtime, cụ thể là cudaOccupancyMaxActiveBlocksPerMultiprocessor (được đóng khung trong Hình 14) và cudaLaunchKernel, trong khi đối với aten::add chỉ có cudaLaunchKernel.

cudaOccupancyMaxActiveBlocksPerMultiprocessor là một lời gọi lập kế hoạch và hoàn toàn nằm ở phía CPU. Nó hỏi: “với một hàm kernel, một kích thước block đã chọn và một kích thước bộ nhớ chia sẻ động đã chọn, có bao nhiêu block của kernel này có thể cùng tồn tại trên một SM (Streaming Multiprocessor)?”

Câu hỏi đặt ra là, tại sao chúng ta cần lập kế hoạch cho matmul mà không cần cho add?

Để hiểu điều này, chúng ta phải nhìn vào dấu chân tài nguyên (resource footprint) của kernel.

Trong Hình 15, ta thấy đối với phép nhân ma trận, registers per thread (thanh ghi mỗi luồng) và shared memory (bộ nhớ chia sẻ) là động (dựa trên kích thước ma trận). cuBLAS cung cấp hàng trăm biến thể kernel, và mỗi biến thể có một đường dẫn kích hoạt dựa trên heuristic cần thông tin runtime về khả năng của phần cứng. Truy vấn occupancy là một phần của heuristic đó.

Từ Hình 16, chúng ta thấy dấu chân của phép cộng là 32 thanh ghi và 0 bộ nhớ chia sẻ. Điều này hiển nhiên là vừa vặn. Không có gì để truy vấn vì không có tài nguyên phần cứng nào sẽ giới hạn occupancy. Kernel này, theo thiết kế, tiêu tốn rất ít tài nguyên.

Bạn có thể sử dụng điều này như một chẩn đoán nhanh khi đọc bất kỳ bản trace nào. Quét luồng CPU tìm cudaOccupancyMaxActiveBlocksPerMultiprocessor. Mỗi lần xuất hiện đánh dấu một kernel “nặng, được kích hoạt thích ứng”, thường là GEMM, conv hoặc tương tự. Các kernel không có truy vấn occupancy đi trước là nhóm elementwise/reduction mà PyTorch kích hoạt một cách máy móc.

Tại sao cudaDeviceSynchronize lại lớn như vậy (~1.78 ms)?

cudaDeviceSynchronize chặn CPU cho đến khi tất cả công việc GPU trên thiết bị này kết thúc. Profiler phát ra lệnh đồng bộ này ở cuối cửa sổ active để xả (flush) các sự kiện. Nếu không có nó, thời gian chạy của kernel sẽ bị thiếu.

Một lệnh đồng bộ 1.78 ms bao phủ 26 µs công việc GPU thực tế cho bạn biết rằng lần chạy này nhàn rỗi 98%. Đó là triệu chứng điển hình của trạng thái overhead-bound.

Bản trace 4096x4096

Chúng ta đã biết từ phân tích bảng profiler rằng việc cung cấp ma trận lớn hơn sẽ đưa thuật toán ra khỏi vùng overhead-bound và trở thành compute-bound. Hãy cùng xem chi tiết hơn.

uv run 01_matmul_add.py --size 4096 --warmup

Tại sao cùng một kernel lại tốn thời gian khác nhau?

Trong Hình 17, chúng ta thấy kernel matmul cho ProfileStep#3 tốn nhiều thời gian hơn trên GPU so với các bước khác. Điều này đặc biệt thú vị vì các kernel được kích hoạt là hoàn toàn giống nhau, nghĩa là không có heuristic cuBLAS nào can thiệp.

Bản trace trong Hình 17 chỉ ra một điểm hữu ích: thời gian chạy của kernel không phải là hằng số, ngay cả trên cùng một môi trường phần cứng chạy mã giống hệt nhau trên dữ liệu giống hệt nhau.

Hãy chạy lặp lại 20 lần để quan sát.

- schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=20, repeat=1)

- for _ in range(5):
+ for _ in range(20):

Hình 18 tiết lộ kết quả tương tự. Thời gian tính toán khác nhau có thể do nhiều nguyên nhân:

• Xung nhịp GPU khi nhàn rỗi và khi boost.
• Nhiệt độ GPU.
• Quản lý năng lượng GPU.
• Các công việc dọn dẹp (housekeeping) phía driver.

Một người chỉ nhìn vào giá trị trung bình sẽ kết luận matmul tốn ~1 ms; một người nhìn vào bản trace sẽ thấy matmul tốn ~580 µs trừ khi GPU “nổi nóng”. Đó là hai mô hình tư duy rất khác nhau, và chỉ có một cái là đúng.

Xem torch.compile hoạt động

Làm việc với torch.compile luôn gây kinh ngạc. Bạn viết mã PyTorch eager bình thường, nhưng PyTorch cố gắng nắm bắt các vùng nặng về tensor, biến chúng thành đồ thị (graphs), tối ưu hóa chúng và chạy mã được tạo ra. Backend mặc định thường là TorchInductor, và pipeline tổng quát là:

1. TorchDynamo nắm bắt thực thi Python thành một đồ thị FX.
2. AOTAutograd chuẩn bị đồ thị forward/backward khi có gradient.
3. Inductor hạ cấp (lowers) đồ thị thành mã CPU hoặc GPU đã được tối ưu hóa.

uv run 01_matmul_add.py --size 4096 --warmup --compile

Mã kích hoạt:

def fn(x, w, b):
  return torch.add(torch.matmul(x, w), b)

fn = torch.compile(fn) if args.compile else fn

Trong Hình 19, chúng ta thấy các hàng CPU mới tên là Torch-Compiled Region: 0/0, chỉ ra các hàm đã biên dịch đang được sử dụng.

Chúng ta có hợp nhất (fuse) các kernel matmul và add thành một không?

Nhìn vào Hình 20, chúng ta tự hỏi: liệu chúng ta có thực sự hợp nhất các thao tác nhân và cộng lại thành một không?

Đây là sự hợp nhất toán tử ở mức đồ thị. Inductor đã lấy torch.add(torch.matmul(x, w), b) và viết lại thành một lời gọi aten::addmm(b, x, w) duy nhất. Điều quan trọng cần lưu ý là nó không tạo ra một CUDA kernel hợp nhất mới. Công việc GPU thực tế vẫn là ampere_bf16_s16816gemm..., cùng một kernel cuBLAS mà chế độ eager đã sử dụng. Vì vậy, sự “hợp nhất” ở đây nằm ở mức điều phối (dispatcher), không phải ở mức kernel.

PyTorch cung cấp hàm torch.addmm thực hiện việc nhân rồi cộng trong một bước. Chúng tôi khuyến khích bạn xem bản trace của hàm này và chia sẻ quan sát bên dưới!

Kiến trúc runtime của torch.compile

Hãy nhìn vào phân cấp phía CPU phản ánh kiến trúc runtime của torch.compile.

TorchDynamo Cache Lookup là nơi Dynamo kiểm tra xem lời gọi hiện tại có còn khớp với những gì đã được biên dịch với cùng hình dạng (shapes), kiểu dữ liệu (dtypes), thiết bị và siêu dữ liệu tensor hay không. Nếu có bất kỳ điểm không khớp nào, Dynamo sẽ biên dịch lại. Chi phí này phải trả cho mọi lời gọi, ngay cả sau khi biên dịch.

Torch-Compiled Region là wrapper để “đi vào” phiên bản biên dịch. AOTDispatcher Runtime Wrapper Prologue là runtime wrapper của AOT Autograd. Ngay cả khi chúng ta không cần gradient, AOTDispatcher luôn nằm trong stack để xử lý siêu dữ liệu tensor, theo dõi view, và sẽ thiết lập lượt chạy backward nếu requires_grad là true.

## Call CompiledFxGraph là nơi mã được tạo ra thực sự chạy. Chuỗi ký tự sau “CompiledFxGraph” là mã băm (hash) nội dung của đồ thị FX. Nó giống nhau trong cả ba bước active, xác nhận là đã hit cache.

Bạn có thể tìm thấy mã được tạo ra trên đĩa tại /tmp/torchinductor_<user>/fxgraph theo mã băm này, hữu ích khi bạn muốn đọc mã Triton/C++ mà Inductor thực sự tạo ra.

Các lượt kích hoạt CUDA có giảm đi một nửa không?

Nhìn vào trace trong Hình 21, chúng ta rất vui khi thấy chỉ có một cudaLaunchKernel mỗi bước. Tuy nhiên, quan sát này mâu thuẫn với những gì chúng ta thấy trong trace GPU. Vẫn có hai kernel được kích hoạt mỗi bước: Memcpy DtoD (Device -> Device) và GEMM. Quay lại trace CPU, chúng ta thấy mình đã bỏ sót lệnh điều phối cudaMemcpyAsync.

addmm tính toán out = α·A·B + β·C, và epilogue (phần kết) GEMM-with-bias-add của cuBLAS ghi vào một bộ đệm đích vốn đã phải chứa bias. Epilogue có thể hiểu là tất cả các thao tác xảy ra sau một GEMM.

Vì vậy, mã được Inductor tạo ra thực hiện:

• out = copy(C) $\rightarrow$ đó là DtoD memcpy (32 MB, tốn ~33 µs).
• out = α·(A·B) + β·out $\rightarrow$ GEMM với $\alpha=\beta=1$, hợp nhất việc cộng bias vào quá trình ghi kết quả (writeback).

Kết quả toán học vẫn như vậy. Việc cộng bias không hề miễn phí, vì chúng ta phải trả chi phí memcpy trước cộng với một epilogue GEMM đắt hơn một chút.

Sự hợp nhất mà người ta kỳ vọng, nơi x·w + b thu gọn thành một kernel duy nhất không có lưu lượng bộ nhớ thừa, đã không xảy ra. Inductor vẫn giữ hai thao tác chạm bộ nhớ, nó chỉ dán nhãn lại việc sao chép bias thành memcpy và phép cộng thành GEMM epilogue.

Một bản thực thi hợp nhất thực sự sẽ bỏ qua memcpy. Đó là những gì các kernel viết tay kiểu FlashAttention làm, và là điều Inductor có thể làm thông qua Triton codegen, nhưng đối với một phép nhân ma trận 4096×4096 bf16, Inductor rõ ràng quyết định rằng “sử dụng cuBLAS, thực hiện bias qua epilogue” là con đường tốt nhất.

Chi phí CPU tăng lên chứ không giảm đi

Đây là điều dễ bị bỏ qua nhất khi so sánh giữa lần chạy eager và biên dịch:

Biên dịch tốn kém hơn khoảng 2 lần trên CPU mỗi bước. Đó là vì mọi lời gọi đều phải đi qua toàn bộ stack Dynamo $\rightarrow$ AOTAutograd $\rightarrow$ Inductor, trên nền lệnh điều phối aten::addmm vốn đã có sẵn. Pipeline biên dịch được xây dựng cho các mô hình ML với hàng chục toán tử, nơi chi phí cho mỗi lời gọi được phân bổ (amortize). Đối với một toán tử đơn lẻ, đây giống như một loại “thuế”.

torch.compile có đối số mode. Bạn hãy đọc tài liệu và tìm xem mode nào có thể làm giảm chi phí CPU này nhé. 🤗

Bảng tra cứu nhanh đọc trace (Cheatsheet)

Một tham chiếu nhanh cho các mô hình chúng ta đã đi qua. Ý tưởng là: nếu bạn thấy điều này trong trace, thì thông thường nó có nghĩa là…

Bảng Profiler

Luồng CPU (CPU lane)

Luồng GPU (GPU lane)

Chuỗi điều phối (Dispatch chain)

torch.compile

Kết luận

Chúng ta bắt đầu với một phép matmul + add nhỏ và sử dụng nó như một cái cớ để học cách đọc PyTorch profiler. Qua đó, chúng ta đã thu thập được một vài mô hình tư duy có thể áp dụng cho các khối lượng công việc lớn hơn. Đây là điểm dừng đầu tiên trong chuỗi bài Profiling PyTorch. Trong các bài tiếp theo, chúng ta sẽ dần rời xa “món đồ chơi” hai thao tác này để leo lên nấc thang phức tạp hơn, xem xét các khối xây dựng lớn hơn và cuối cùng là các mô hình thực tế.

Link bài viết gốc

• Tags:
• Ai
• May 29, 2026
• Huggingface.co