torch.compile cho suy luận

torch.compile biên dịch mã PyTorch thành các kernel tối ưu, giúp tăng tốc độ suy luận đáng kể. Tính năng này dựa vào TorchDynamo để biên dịch mã thành các đồ thị và TorchInductor để tiếp tục biên dịch các đồ thị đó thành các kernel tối ưu. Đây là một công cụ tối ưu mạnh mẽ và trong nhiều trường hợp, chỉ cần thêm một dòng mã duy nhất.

Bọc một model bằng torch.compile để biên dịch và trả về một model đã được tối ưu.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model)

Gọi torch.compile lần đầu tiên sẽ chậm vì mô hình cần được biên dịch. Các lần gọi mô hình đã biên dịch sau đó sẽ nhanh hơn nhiều vì không cần biên dịch lại.

Có một số tham số để tùy chỉnh quá trình biên dịch. Hai tham số quan trọng hơn được liệt kê bên dưới. Để xem danh sách đầy đủ các tham số, hãy tham khảo tài liệu của torch.compile.

Các chế độ

Tham số mode cung cấp một số tùy chọn hiệu năng cho việc biên dịch. Hãy thử các chế độ khác nhau để xem chế độ nào hoạt động tốt nhất cho trường hợp sử dụng của bạn.

• default là tùy chọn cân bằng giữa tốc độ và bộ nhớ.
• reduce-overhead giảm bớt chi phí xử lý của Python với cái giá là tốn thêm một chút bộ nhớ, nhưng nó có thể nhanh hơn.
• max-autotune mang lại tốc độ nhanh nhất, nhưng quá trình biên dịch mất nhiều thời gian hơn.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead")

Fullgraph

Fullgraph cố gắng biên dịch toàn bộ mô hình thành một đồ thị duy nhất để tối đa hóa hiệu suất. torch.compile sẽ gây ra lỗi nếu gặp một điểm ngắt đồ thị (graph break), điều này có nghĩa là nó không thể biên dịch mô hình thành một đồ thị duy nhất.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

Benchmarks

Tham khảo bảng bên dưới để xem các chỉ số hiệu năng so sánh thời gian suy luận trung bình tính bằng mili giây khi bật và tắt torch.compile trên các GPU và kích thước lô khác nhau, sử dụng cùng một hình ảnh cho các tác vụ thị giác khác nhau.

Chọn Subset trong bảng bên dưới để chuyển đổi giữa các GPU khác nhau, cũng như các benchmark trên PyTorch nightly 2.1.0dev và torch.compile với chế độ reduce-overhead được bật.

Các bước tiếp theo

• Xem hướng dẫn torch.compile cho huấn luyện để tìm hiểu cách sử dụng torch.compile cho huấn luyện.