Tối ưu hóa không gian ZeroGPU của bạn với việc biên dịch trước thời hạn PyTorch
Bài viết này hướng dẫn cách tối ưu hóa không gian ZeroGPU bằng cách sử dụng việc biên dịch trước thời hạn của PyTorch.
- 16 min read
Tăng tốc độ Spaces ZeroGPU của bạn với biên dịch trước thời gian (Ahead-of-Time)
Đăng ngày: 2 tháng 9, 2025
ZeroGPU cho phép bất cứ ai khởi chạy phần cứng Nvidia H200 mạnh mẽ trong Hugging Face Spaces mà không cần khóa GPU cho lưu lượng truy cập nhàn rỗi. Nó hiệu quả, linh hoạt và lý tưởng cho các bản demo, nhưng nó không phải lúc nào cũng tận dụng tối đa mọi thứ mà GPU và bộ CUDA có thể cung cấp. Việc tạo ảnh hoặc video có thể mất nhiều thời gian đáng kể. Khả năng khai thác hiệu suất cao hơn, tận dụng phần cứng H200, rất quan trọng trong trường hợp này.
Đây là lúc biên dịch trước thời gian (AoT) của PyTorch phát huy tác dụng. Thay vì biên dịch mô hình khi cần (điều này không hoạt động tốt với các tiến trình ngắn hạn của ZeroGPU), AoT cho phép bạn tối ưu hóa một lần và tải lại ngay lập tức.
Kết quả: các bản demo nhanh hơn và trải nghiệm mượt mà hơn, với tốc độ tăng từ 1.3×–1.8× trên các mô hình như Flux, Wan và LTX 🔥
Trong bài viết này, chúng ta sẽ chỉ cách kết nối biên dịch trước thời gian (AoT) trong Spaces ZeroGPU. Chúng ta sẽ khám phá các thủ thuật nâng cao như lượng tử hóa FP8 và hình dạng động, và chia sẻ các bản demo hoạt động mà bạn có thể thử ngay lập tức. Nếu bạn không thể chờ đợi, chúng tôi mời bạn xem một số bản demo được hỗ trợ bởi ZeroGPU trên tổ chức zerogpu-aoti.
Người dùng Pro và thành viên tổ chức Team/Enterprise có thể tạo Spaces ZeroGPU, trong khi bất kỳ ai cũng có thể sử dụng chúng một cách tự do (người dùng Pro, Team và Enterprise có hạn ngạch ZeroGPU nhiều hơn 8x)
Mục lục
- ZeroGPU là gì
- Biên dịch PyTorch
- Biên dịch trước thời gian trên ZeroGPU
- Những điểm cần lưu ý
- Các bản demo Spaces ZeroGPU đã được biên dịch AoT
- Kết luận
- Tài nguyên
ZeroGPU là gì
Spaces là một nền tảng do Hugging Face cung cấp, cho phép các chuyên gia ML dễ dàng xuất bản các ứng dụng demo.
Các ứng dụng demo điển hình trên Spaces trông như thế này:
import gradio as gr
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(...).to('cuda')
def generate(prompt):
return pipe(prompt).images
gr.Interface(generate, "text", "gallery").launch()
Điều này hoạt động rất tốt, nhưng kết thúc việc dành riêng một GPU cho Space trong suốt vòng đời của nó - ngay cả khi không có hoạt động của người dùng.
Khi thực thi .to('cuda') trên dòng này:
pipe = DiffusionPipeline.from_pretrained(...).to('cuda')
PyTorch khởi tạo trình điều khiển NVIDIA, thiết lập quy trình trên CUDA mãi mãi. Điều này không hiệu quả về tài nguyên, vì lưu lượng truy cập ứng dụng không hoàn toàn mượt mà, mà khá thưa thớt và đột biến.
ZeroGPU sử dụng phương pháp đúng lúc để khởi tạo GPU. Thay vì thiết lập quy trình chính trên CUDA, nó tự động phân nhánh quy trình, thiết lập nó trên CUDA, chạy các tác vụ GPU và cuối cùng xóa nhánh khi cần giải phóng GPU.
Điều này có nghĩa là:
- Khi ứng dụng không nhận được lưu lượng truy cập, nó không sử dụng bất kỳ GPU nào
- Khi nó thực sự đang thực hiện một tác vụ, nó sẽ sử dụng một GPU
- Nó có thể sử dụng nhiều GPU khi cần để thực hiện các tác vụ đồng thời
Nhờ gói spaces của Python, thay đổi mã duy nhất cần thiết để có được hành vi này như sau:
import gradio as gr
+ import spaces
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(...).to('cuda')
+ @spaces.GPU
def generate(prompt):
return pipe(prompt).images
gr.Interface(generate, "text", "gallery").launch()
Bằng cách nhập spaces và thêm trình trang trí @spaces.GPU, chúng ta:
- Chặn các cuộc gọi API PyTorch để hoãn các thao tác CUDA
- Làm cho hàm được trang trí chạy trong một nhánh khi được gọi sau đó
- (Gọi API bên trong để làm cho thiết bị phù hợp hiển thị với nhánh nhưng điều này không nằm trong phạm vi của bài đăng trên blog này)
ZeroGPU hiện đang phân bổ một lát cắt MIG của H200 (mục tiêu
3g.71gb). Các kích thước MIG bổ sung bao gồm lát cắt đầy đủ (mục tiêu7g.141gb) sẽ có vào cuối năm 2025.
Biên dịch PyTorch
Các framework ML hiện đại như PyTorch và JAX có khái niệm biên dịch có thể được sử dụng để tối ưu hóa độ trễ mô hình hoặc thời gian suy luận. Bên dưới, biên dịch áp dụng một loạt các bước tối ưu hóa (thường phụ thuộc vào phần cứng) như hợp nhất toán tử, gập hằng số, v.v.
PyTorch (từ phiên bản 2.0 trở đi) hiện có hai giao diện chính để biên dịch:
- Ngay lập tức với
torch.compile - Trước thời gian với
torch.export+AOTInductor
torch.compile hoạt động tốt trong các môi trường tiêu chuẩn: nó biên dịch mô hình của bạn lần đầu tiên nó chạy và sử dụng lại phiên bản được tối ưu hóa cho các cuộc gọi tiếp theo.
Tuy nhiên, trên ZeroGPU, vì quy trình được khởi chạy mới cho (gần như) mọi tác vụ GPU, điều đó có nghĩa là torch.compile không thể sử dụng lại hiệu quả việc biên dịch và do đó buộc phải dựa vào bộ nhớ cache hệ thống tệp để khôi phục các mô hình đã biên dịch. Tùy thuộc vào mô hình đang được biên dịch, quá trình này mất từ vài chục giây đến vài phút, điều này quá nhiều đối với các tác vụ GPU thực tế trong Spaces.
Trong ví dụ trên, chúng ta chỉ biên dịch bộ chuyển đổi vì trong các mô hình tạo này, bộ chuyển đổi (hoặc tổng quát hơn là bộ khử nhiễu) là thành phần nặng về tính toán nhất.
Đây là lúc biên dịch trước thời gian (AoT) tỏa sáng.
Với AoT, chúng ta có thể xuất một mô hình đã biên dịch một lần, lưu nó và sau đó tải lại ngay lập tức trong bất kỳ quy trình nào, đó chính xác là những gì chúng ta cần cho ZeroGPU. Điều này giúp chúng ta giảm chi phí sử dụng framework và cũng loại bỏ thời gian khởi động lạnh thường xảy ra trong biên dịch đúng lúc.
Nhưng làm thế nào chúng ta có thể biên dịch trước thời gian trên ZeroGPU? Hãy cùng tìm hiểu.
Biên dịch trước thời gian trên ZeroGPU
Hãy quay lại ví dụ cơ sở ZeroGPU của chúng ta và giải nén những gì chúng ta cần để đưa biên dịch AoT vào. Vì mục đích của bản demo này, chúng ta sẽ sử dụng mô hình black-forest-labs/FLUX.1-dev:
import gradio as gr
import spaces
import torch
from diffusers import DiffusionPipeline
MODEL_ID = 'black-forest-labs/FLUX.1-dev'
pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe.to('cuda')
@spaces.GPU
def generate(prompt):
return pipe(prompt).images
gr.Interface(generate, "text", "gallery").launch()
Việc biên dịch mô hình trước thời gian với PyTorch liên quan đến nhiều bước:
1. Lấy ví dụ đầu vào
Hãy nhớ rằng chúng ta sẽ biên dịch mô hình trước thời gian. Do đó, chúng ta cần lấy ví dụ đầu vào cho mô hình. Lưu ý rằng đây là cùng loại đầu vào mà chúng ta mong đợi thấy trong các lần chạy thực tế. Để ghi lại các đầu vào đó, chúng ta sẽ sử dụng trình trợ giúp spaces.aoti_capture từ gói spaces:
with spaces.aoti_capture(pipe.transformer) as call:
pipe("arbitrary example prompt")
Khi được sử dụng như một trình quản lý ngữ cảnh như vậy, nó cho phép ghi lại cuộc gọi của bất kỳ hàm gọi được nào ( pipe.transformer, trong trường hợp của chúng ta), hủy bỏ nó và đặt các giá trị đầu vào đã ghi lại bên trong call.args và call.kwargs.
2. Xuất mô hình
Bây giờ chúng ta đã có các đối số và từ khóa ví dụ cho thành phần bộ chuyển đổi của mình, chúng ta có thể xuất nó thành một ExportedProgram của PyTorch bằng cách sử dụng tiện ích torch.export.export:
exported_transformer = torch.export.export(
pipe.transformer,
args=call.args,
kwargs=call.kwargs,
)
Một chương trình PyTorch được xuất là một biểu đồ tính toán thể hiện các phép tính tensor cùng với các giá trị tham số mô hình ban đầu.
3. Biên dịch mô hình đã xuất
Sau khi mô hình được xuất, việc biên dịch nó khá đơn giản.
Việc biên dịch AoT truyền thống trong PyTorch thường yêu cầu lưu mô hình trên đĩa để nó có thể được tải lại sau này. Trong trường hợp của chúng ta, chúng ta sẽ sử dụng một hàm trợ giúp thuộc gói spaces: spaces.aoti_compile. Nó là một trình bao bọc nhỏ quanh torch._inductor.aot_compile quản lý việc lưu và tải mô hình lười biếng khi cần. Nó được thiết kế để được sử dụng như sau:
compiled_transformer = spaces.aoti_compile(exported_transformer)
compiled_transformer này giờ đây là một nhị phân đã được biên dịch AoT, sẵn sàng được sử dụng để suy luận.
4. Sử dụng mô hình đã biên dịch trong pipeline
Bây giờ chúng ta cần liên kết bộ chuyển đổi đã biên dịch của mình với pipeline gốc, tức là pipeline.
Một phương pháp ngây thơ và gần như hoạt động là chỉ cần vá pipeline của chúng ta như pipe.transformer = compiled_transformer. Thật không may, phương pháp này không hoạt động vì nó xóa các thuộc tính quan trọng như dtype, config, v.v. Chỉ vá phương pháp forward cũng không hoạt động tốt vì sau đó chúng ta đang giữ các tham số mô hình gốc trong bộ nhớ, thường dẫn đến lỗi OOM khi chạy.
Gói spaces cũng cung cấp một tiện ích cho điều này - spaces.aoti_apply:
spaces.aoti_apply(compiled_transformer, pipe.transformer)
Et voilà! Nó sẽ đảm nhiệm việc vá pipe.transformer.forward bằng mô hình đã biên dịch của chúng ta, cũng như xóa các tham số mô hình cũ khỏi bộ nhớ.
5. Kết hợp tất cả lại với nhau
Để thực hiện ba bước đầu tiên (chặn các ví dụ đầu vào, xuất mô hình và biên dịch nó bằng PyTorch inductor), chúng ta cần một GPU thực tế. Mô phỏng CUDA mà bạn nhận được bên ngoài hàm @spaces.GPU là không đủ vì việc biên dịch thực sự phụ thuộc vào phần cứng, ví dụ, dựa trên các lần chạy điểm chuẩn vi mô để điều chỉnh mã được tạo. Đó là lý do tại sao chúng ta cần gói gọn tất cả mọi thứ bên trong một hàm @spaces.GPU và sau đó lấy mô hình đã biên dịch của mình trở lại gốc của ứng dụng. Bắt đầu từ mã demo ban đầu của chúng ta, điều này cho ra:
import gradio as gr
import spaces
import torch
from diffusers import DiffusionPipeline
MODEL_ID = 'black-forest-labs/FLUX.1-dev'
pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe.to('cuda')
+ @spaces.GPU(duration=1500) # thời gian tối đa được phép trong khi khởi động
+ def compile_transformer():
+ with spaces.aoti_capture(pipe.transformer) as call:
+ pipe("arbitrary example prompt")
+
+ exported = torch.export.export(
+ pipe.transformer,
+ args=call.args,
+ kwargs=call.kwargs,
+ )
+ return spaces.aoti_compile(exported)
+
+ compiled_transformer = compile_transformer()
+ spaces.aoti_apply(compiled_transformer, pipe.transformer)
@spaces.GPU
def generate(prompt):
return pipe(prompt).images
gr.Interface(generate, "text", "gallery").launch()
Chỉ với một tá dòng mã bổ sung, chúng ta đã thành công trong việc làm cho bản demo của mình nhanh hơn đáng kể (nhanh hơn 1.7x trong trường hợp của FLUX.1-dev).
Nếu bạn muốn tìm hiểu thêm về biên dịch AoT, bạn có thể đọc hướng dẫn AOTInductor của PyTorch
Những điểm cần lưu ý
Bây giờ chúng ta đã chứng minh được tốc độ tăng mà người ta có thể nhận ra trong các ràng buộc hoạt động với ZeroGPU, chúng ta sẽ thảo luận về một vài điểm cần lưu ý đã xuất hiện khi làm việc với thiết lập này.
Lượng tử hóa
AoT có thể được kết hợp với lượng tử hóa để mang lại tốc độ tăng thậm chí còn lớn hơn.
Đối với việc tạo ảnh và video, các lược đồ lượng tử hóa động sau đào tạo FP8 mang lại sự cân bằng tốt giữa tốc độ và chất lượng.
Tuy nhiên, FP8 yêu cầu khả năng tính toán CUDA ít nhất là 9.0 để hoạt động.
May mắn thay, đối với ZeroGPU, vì chúng dựa trên H200, chúng ta đã có thể tận dụng các lược đồ lượng tử hóa FP8.
Để bật lượng tử hóa FP8 trong quy trình làm việc biên dịch AoT của chúng ta, chúng ta có thể sử dụng các API do torchao cung cấp như sau:
+ from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig
+ # Lượng tử hóa bộ chuyển đổi ngay trước bước xuất.
+ quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
exported_transformer = torch.export.export(
pipe.transformer,
args=call.args,
kwargs=call.kwargs,
)
(Bạn có thể tìm thêm chi tiết về TorchAO ở đây.)
Và sau đó chúng ta có thể tiếp tục với các bước còn lại như đã nêu ở trên. Việc sử dụng lượng tử hóa mang lại thêm 1.2x tốc độ tăng.
Hình dạng động
Hình ảnh và video có thể có nhiều hình dạng và kích thước khác nhau. Do đó, điều quan trọng là cũng phải tính đến tính năng động của hình dạng khi thực hiện biên dịch AoT. Các nguyên hàm do torch.export.export cung cấp giúp dễ dàng cấu hình để cung cấp đầu vào nào nên được xử lý tương ứng cho các hình dạng động, như được hiển thị bên dưới.
Đối với trường hợp của bộ chuyển đổi Flux.1-Dev, sự thay đổi độ phân giải hình ảnh sẽ ảnh hưởng đến hai đối số forward của nó:
hidden_states: Các latent đầu vào nhiễu, mà bộ chuyển đổi được cho là sẽ khử nhiễu. Đó là một tensor 3D, đại diện chobatch_size, flattened_latent_dim, embed_dim. Khi kích thước batch cố định, đó làflattened_latent_dimsẽ thay đổi đối với bất kỳ thay đổi nào được thực hiện đối với độ phân giải hình ảnh.img_ids: Một mảng 2D của các tọa độ pixel được mã hóa có hình dạngheight * width, 3. Trong trường hợp này, chúng ta muốn làm choheight * widthtrở nên động.
Chúng ta bắt đầu bằng cách định nghĩa một phạm vi mà chúng ta muốn cho phép độ phân giải hình ảnh (ẩn) thay đổi.
Để lấy các phạm vi giá trị này, chúng ta đã kiểm tra hình dạng của hidden_states trong pipeline đối với các độ phân giải hình ảnh khác nhau. Các giá trị chính xác phụ thuộc vào mô hình và yêu cầu kiểm tra thủ công và một số trực giác. Đối với Flux.1-Dev, chúng ta đã kết thúc với:
transformer_hidden_dim = torch.export.Dim('hidden', min=4096, max=8212)
Sau đó, chúng ta định nghĩa một bản đồ tên đối số và các chiều nào trong các giá trị đầu vào của chúng ta dự kiến sẽ có tính năng động:
transformer_dynamic_shapes = {
"hidden_dim": {1: transformer_hidden_dim},
"img_ids", {0: transformer_hidden_dim},
}
Sau đó, chúng ta cần làm cho đối tượng hình dạng động của mình sao chép cấu trúc của các đầu vào ví dụ của chúng ta. Các đầu vào không cần hình dạng động phải được đặt thành None. Điều này có thể được thực hiện rất dễ dàng với tiện ích tree_map của PyTorch:
from pytorch.utils._pytree import tree_map
dynamic_shapes = tree_map(lambda v: None, call.kwargs)
dynamic_shapes |= transformer_dynamic_shapes
Bây giờ, khi thực hiện bước xuất, chúng ta chỉ cần cung cấp transformer_dynamic_shapes cho torch.export.export:
exported_transformer = torch.export.export(
pipe.transformer,
args=call.args,
kwargs=call.kwargs,
dynamic_shapes=dynamic_shapes,
)
Hãy xem Space này cho thấy cách sử dụng cả lượng tử hóa và hình dạng động trong bước xuất.
Biên dịch nhiều/chia sẻ trọng số
Hình dạng động đôi khi là không đủ khi tính năng động quá quan trọng.
Ví dụ, đây là trường hợp của họ mô hình tạo video Wan nếu bạn muốn mô hình đã biên dịch của bạn tạo ra các độ phân giải khác nhau.
Một điều có thể được thực hiện trong trường hợp này: biên dịch một mô hình cho mỗi độ phân giải trong khi giữ cho các tham số mô hình được chia sẻ và phân phối đúng mô hình khi chạy.
Đây là một ví dụ tối thiểu về phương pháp này: zerogpu-aoti-multi.py. Bạn cũng có thể thấy một triển khai đầy đủ hoạt động của mô hình này trong Space Wan 2.2.
FlashAttention-3
Vì phần cứng ZeroGPU và trình điều khiển CUDA hoàn toàn tương thích với Flash-Attention 3 (FA3), nên chúng ta có thể sử dụng nó trong Spaces ZeroGPU của mình để tăng tốc mọi thứ hơn nữa. FA3 hoạt động với biên dịch trước thời gian. Vì vậy, điều này rất lý tưởng cho trường hợp của chúng ta.
Việc biên dịch và xây dựng FA3 từ nguồn có thể mất vài phút và quá trình này phụ thuộc vào phần cứng. Là người dùng, chúng ta sẽ không muốn mất những giờ tính toán ZeroGPU quý báu. Đây là lúc thư viện kernels của Hugging Face đến để giải cứu. Nó cung cấp quyền truy cập vào các kernel được xây dựng sẵn, tương thích với phần cứng đã cho. Ví dụ: khi chúng ta cố gắng chạy:
from kernels import get_kernel
vllm_flash_attn3 = get_kernel("kernels-community/vllm-flash-attn3")
Nó cố gắng tải một kernel từ kho lưu trữ kernels-community/vllm-flash-attn3, tương thích với thiết lập hiện tại. Nếu không, nó sẽ báo lỗi do sự cố không tương thích. May mắn thay, điều này hoạt động liền mạch trên Spaces ZeroGPU. Điều này có nghĩa là chúng ta có thể tận dụng sức mạnh của FA3 trên ZeroGPU, sử dụng thư viện kernels.
Đây là ví dụ đầy đủ hoạt động của bộ xử lý chú ý FA3 cho mô hình Qwen-Image.
Các bản demo Spaces ZeroGPU đã được biên dịch AoT
So sánh tốc độ
- FLUX.1-dev không có AoTI
- FLUX.1-dev có AoTI và FA3 (tốc độ tăng 1.75x)
Các Spaces AoTI nổi bật
Kết luận
ZeroGPU trong Hugging Face Spaces là một tính năng mạnh mẽ cho phép các nhà xây dựng AI bằng cách cung cấp quyền truy cập vào khả năng tính toán mạnh mẽ. Trong bài đăng này, chúng tôi đã chỉ ra cách người dùng có thể hưởng lợi từ các kỹ thuật biên dịch trước thời gian của PyTorch để tăng tốc các ứng dụng của họ sử dụng ZeroGPU.
Chúng tôi chứng minh tốc độ tăng với Flux.1-Dev, nhưng các kỹ thuật này không chỉ giới hạn ở mô hình này. Do đó, chúng tôi khuyến khích bạn hãy thử các kỹ thuật này và cung cấp cho chúng tôi phản hồi trong cuộc thảo luận cộng đồng này.
Tài nguyên
- Truy cập tổ chức ZeroGPU-AOTI của chúng tôi trên Hub để tham khảo bộ sưu tập các bản demo sử dụng các kỹ thuật được thảo luận trong bài đăng này.
- Duyệt qua các API
spaces.aoti_*mã nguồn để tìm hiểu thêm về giao diện. - Kiểm tra tổ chức Kernels Community trên hub
- Nâng cấp lên Pro trên Hugging Face để tạo Spaces ZeroGPU của riêng bạn (và nhận được 25 phút sử dụng H200 mỗi ngày)
Lời cảm ơn: Cảm ơn ChunTe Lee đã tạo hình thu nhỏ tuyệt vời cho bài đăng này. Cảm ơn Pedro và Vaibhav đã cung cấp phản hồi về bài đăng.