Giảm mức sử dụng bộ nhớ

Các mô hình khuếch tán hiện đại như Flux và Wan có hàng tỷ tham số, chiếm nhiều bộ nhớ trên phần cứng của bạn khi suy luận. Điều này gây khó khăn vì các GPU phổ biến thường không có đủ bộ nhớ. Để khắc phục giới hạn về bộ nhớ, bạn có thể sử dụng nhiều hơn một GPU (nếu có), chuyển một số thành phần của pipeline sang CPU, và nhiều phương pháp khác.

Hướng dẫn này sẽ cho bạn thấy cách giảm mức sử dụng bộ nhớ của bạn.

Hãy lưu ý rằng các kỹ thuật này có thể cần được điều chỉnh tùy thuộc vào mô hình. Ví dụ, một mô hình khuếch tán dựa trên transformer có thể không hưởng lợi từ các tối ưu hóa bộ nhớ này một cách đồng đều như một mô hình dựa trên UNet.

Nhiều GPU

Nếu bạn có quyền truy cập vào nhiều hơn một GPU, có một số tùy chọn để tải và phân phối một mô hình lớn một cách hiệu quả trên phần cứng của bạn. Các tính năng này được hỗ trợ bởi thư viện Accelerate, vì vậy hãy đảm bảo cài đặt nó trước.

pip install -U accelerate

Điểm kiểm tra phân mảnh

Việc nạp các checkpoint lớn được chia thành nhiều shard rất hữu ích vì các shard được nạp lần lượt một một. Điều này giúp giữ mức sử dụng bộ nhớ ở mức thấp, chỉ cần đủ bộ nhớ cho kích thước mô hình và kích thước shard lớn nhất. Chúng tôi khuyến nghị chia shard khi checkpoint fp32 lớn hơn 5GB. Kích thước shard mặc định là 5GB.

Tạo các shard cho checkpoint trong save_pretrained() bằng tham số max_shard_size.

from diffusers import AutoModel

unet = AutoModel.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet"
)
unet.save_pretrained("sdxl-unet-sharded", max_shard_size="5GB")

Bây giờ bạn có thể sử dụng checkpoint phân mảnh thay vì checkpoint thông thường để tiết kiệm bộ nhớ.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

unet = AutoModel.from_pretrained(
    "username/sdxl-unet-sharded", dtype=torch.float16
)
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    unet=unet,
    dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

Đặt thiết bị

Việc phân bổ thiết bị là một tính năng thử nghiệm và API có thể thay đổi. Hiện tại, chỉ có chiến lược balanced được hỗ trợ. Chúng tôi dự định hỗ trợ các chiến lược ánh xạ bổ sung trong tương lai.

Tham số device_map kiểm soát cách các thành phần của mô hình trong một pipeline hoặc các lớp trong một mô hình riêng lẻ được phân bổ trên các thiết bị.

mức pipeline

mức mô hình

Khi thiết kế device_map của riêng bạn, nó phải là một từ điển gồm tên module hoặc lớp cụ thể của mô hình và một định danh thiết bị (một số nguyên cho GPU, cpu cho CPU, và disk cho ổ đĩa).

Gọi hf_device_map trên một model để xem các lớp của model được phân bổ như thế nào, sau đó thiết kế cấu hình của riêng bạn.

print(transformer.hf_device_map)
{'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 'cpu', 'single_transformer_blocks.11': 'cpu', 'single_transformer_blocks.12': 'cpu', 'single_transformer_blocks.13': 'cpu', 'single_transformer_blocks.14': 'cpu', 'single_transformer_blocks.15': 'cpu', 'single_transformer_blocks.16': 'cpu', 'single_transformer_blocks.17': 'cpu', 'single_transformer_blocks.18': 'cpu', 'single_transformer_blocks.19': 'cpu', 'single_transformer_blocks.20': 'cpu', 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'}

Ví dụ, device_map bên dưới đặt single_transformer_blocks.10 đến single_transformer_blocks.20 trên GPU thứ hai (1).

import torch
from diffusers import AutoModel

device_map = {
    'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 1, 'single_transformer_blocks.11': 1, 'single_transformer_blocks.12': 1, 'single_transformer_blocks.13': 1, 'single_transformer_blocks.14': 1, 'single_transformer_blocks.15': 1, 'single_transformer_blocks.16': 1, 'single_transformer_blocks.17': 1, 'single_transformer_blocks.18': 1, 'single_transformer_blocks.19': 1, 'single_transformer_blocks.20': 1, 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'
}

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    subfolder="transformer",
    device_map=device_map,
    dtype=torch.bfloat16
)

Truyền một từ điển ánh xạ mức sử dụng bộ nhớ tối đa tới từng thiết bị để áp đặt giới hạn. Nếu một thiết bị không có trong max_memory, nó sẽ bị bỏ qua và các thành phần của pipeline sẽ không được phân phối tới thiết bị đó.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

max_memory = {0:"1GB", 1:"1GB"}
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="balanced",
    max_memory=max_memory
)

Diffusers sử dụng bộ nhớ tối đa của tất cả các thiết bị theo mặc định, nhưng nếu chúng không vừa trên các GPU, bạn sẽ cần sử dụng một GPU duy nhất và offload sang CPU bằng các phương pháp bên dưới.

enable_model_cpu_offload() chỉ hoạt động trên một GPU duy nhất nhưng một mô hình rất lớn có thể không vừa trên GPU đó
enable_sequential_cpu_offload() có thể hoạt động nhưng cực kỳ chậm và cũng chỉ giới hạn ở một GPU duy nhất

Sử dụng phương thức reset_device_map() để đặt lại device_map. Điều này là cần thiết nếu bạn muốn sử dụng các phương thức như .to(), enable_sequential_cpu_offload() và enable_model_cpu_offload() trên một pipeline đã được ánh xạ thiết bị.

pipeline.reset_device_map()

VAE slicing

VAE slicing tiết kiệm bộ nhớ bằng cách chia các lô đầu vào lớn thành một lô dữ liệu duy nhất và xử lý chúng riêng biệt. Phương pháp này hoạt động tốt nhất khi tạo ra nhiều hơn một hình ảnh cùng lúc.

Ví dụ, nếu bạn đang tạo 4 hình ảnh cùng lúc, quá trình giải mã sẽ làm tăng bộ nhớ kích hoạt đỉnh lên 4 lần. VAE slicing giảm điều này bằng cách chỉ giải mã 1 hình ảnh tại một thời điểm thay vì tất cả 4 hình ảnh cùng lúc.

Gọi enable_slicing() để bật VAE dạng lát cắt. Bạn có thể kỳ vọng một mức tăng hiệu năng nhỏ khi giải mã các lô nhiều ảnh và không có tác động đến hiệu năng với các lô đơn ảnh.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_slicing()
pipeline(["An astronaut riding a horse on Mars"]*32).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

Các lớp AutoencoderKLWan và AsymmetricAutoencoderKL không hỗ trợ cắt lát (slicing).

VAE tiling

VAE tiling tiết kiệm bộ nhớ bằng cách chia hình ảnh thành các ô nhỏ có phần chồng lấn thay vì xử lý toàn bộ hình ảnh cùng lúc. Điều này cũng làm giảm mức sử dụng bộ nhớ đỉnh vì GPU chỉ xử lý một ô tại một thời điểm.

Gọi enable_tiling() để bật tính năng VAE tiling. Hình ảnh được tạo ra có thể có một số biến đổi về tông màu giữa các tile do chúng được giải mã riêng biệt, nhưng sẽ không có đường nối rõ ràng giữa các tile. Tiling sẽ bị tắt đối với độ phân giải thấp hơn một ngưỡng được quy định trước (nhưng có thể cấu hình). Ví dụ, ngưỡng này là 512x512 đối với VAE trong StableDiffusionPipeline.

import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image

pipeline = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_tiling()

init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, image=init_image, strength=0.5).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

AutoencoderKLWan và AsymmetricAutoencoderKL không hỗ trợ tiling.

Offloading

Các chiến lược offload di chuyển các lớp hoặc mô hình hiện không hoạt động sang CPU để tránh làm tăng bộ nhớ GPU. Các chiến lược này có thể kết hợp với lượng tử hóa và torch.compile để cân bằng giữa tốc độ suy luận và mức sử dụng bộ nhớ.

Tham khảo hướng dẫn Compile and offloading quantized models để biết thêm chi tiết.

Tách bộ nhớ ra CPU

CPU offloading di chuyển trọng số từ GPU sang CPU một cách có chọn lọc. Khi một thành phần được yêu cầu, nó sẽ được chuyển đến GPU và khi không còn cần, nó sẽ được chuyển về CPU. Phương pháp này hoạt động trên các submodule thay vì toàn bộ mô hình. Nó tiết kiệm bộ nhớ bằng cách tránh việc lưu trữ toàn bộ mô hình trên GPU.

Việc offload lên CPU làm giảm đáng kể lượng bộ nhớ sử dụng, nhưng cũng cực kỳ chậm vì các submodule được truyền đi truyền lại nhiều lần giữa các thiết bị. Do tốc độ chậm, phương pháp này thường không thực tế.

Không chuyển pipeline sang CUDA trước khi gọi enable_sequential_cpu_offload(), nếu không lượng bộ nhớ tiết kiệm được sẽ rất ít (xem issue này để biết thêm chi tiết). Đây là một thao tác có trạng thái, cài đặt các hook lên mô hình.

Gọi enable_sequential_cpu_offload() để bật tính năng này trên một pipeline.

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_sequential_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

Tải xuống mô hình

Model offloading di chuyển toàn bộ mô hình sang GPU thay vì chỉ di chuyển một số lớp hoặc thành phần mô hình một cách chọn lọc. Một trong các mô hình pipeline chính, thường là text encoder, UNet và VAE, được đặt trên GPU trong khi các thành phần khác được giữ trên CPU. Các thành phần như UNet, vốn được chạy nhiều lần, sẽ ở lại trên GPU cho đến khi hoàn tất hoàn toàn và không còn cần thiết nữa. Điều này loại bỏ chi phí truyền thông của CPU offloading và biến model offloading thành một phương án thay thế nhanh hơn. Đổi lại, mức tiết kiệm bộ nhớ sẽ không lớn bằng.

Lưu ý rằng nếu các model được tái sử dụng bên ngoài pipeline sau khi các hook đã được cài đặt (xem Removing Hooks để biết thêm chi tiết), bạn cần chạy toàn bộ pipeline và các model theo thứ tự dự kiến để offload chúng một cách chính xác. Đây là một thao tác có trạng thái, cài đặt các hook lên model.

Gọi enable_model_cpu_offload() để bật tính năng này trên một pipeline.

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_model_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

enable_model_cpu_offload() cũng hữu ích khi bạn sử dụng phương thức encode_prompt() một mình để tạo ra trạng thái ẩn của các text encoder.

Tải xuống nhóm

Group offloading di chuyển các nhóm lớp nội bộ (torch.nn.ModuleList hoặc torch.nn.Sequential) sang CPU. Phương pháp này sử dụng ít bộ nhớ hơn model offloading và nhanh hơn CPU offloading vì nó giảm thiểu chi phí truyền thông.

Group offloading có thể không hoạt động với tất cả các mô hình nếu cài đặt forward chứa việc chuyển đổi thiết bị (device casting) cho các đầu vào phụ thuộc vào trọng số, vì điều này có thể xung đột với cơ chế chuyển đổi thiết bị của group offloading.

Bật offloading nhóm bằng cách cấu hình tham số offload_type thành block_level hoặc leaf_level.

block_level offload các nhóm lớp dựa trên tham số num_blocks_per_group. Ví dụ, nếu num_blocks_per_group=2 trên một mô hình có 40 lớp, thì 2 lớp sẽ được onload và offload cùng một lúc (tổng cộng 20 lần onload/offload). Điều này giảm đáng kể yêu cầu về bộ nhớ.
leaf_level offload từng lớp ở mức thấp nhất và tương đương với CPU offloading. Tuy nhiên, bạn có thể tăng tốc độ bằng cách sử dụng streams mà không phải đánh đổi tốc độ suy luận.

Group offloading được hỗ trợ cho toàn bộ pipeline hoặc các model riêng lẻ. Áp dụng group offloading cho toàn bộ pipeline là lựa chọn đơn giản nhất, trong khi áp dụng có chọn lọc cho các model riêng lẻ mang lại cho người dùng nhiều linh hoạt hơn để sử dụng các kỹ thuật offloading khác nhau cho các model khác nhau.

pipeline

model

CUDA stream

Tham số use_stream có thể được kích hoạt cho các thiết bị CUDA hỗ trợ luồng chuyển dữ liệu bất đồng bộ để giảm thời gian thực thi tổng thể so với CPU offloading. Nó chồng lấp việc chuyển dữ liệu và tính toán bằng cách sử dụng prefetching lớp. Lớp tiếp theo cần thực thi sẽ được nạp lên GPU trong khi lớp hiện tại vẫn đang được thực thi. Nó có thể làm tăng đáng kể bộ nhớ CPU, vì vậy hãy đảm bảo bạn có gấp đôi dung lượng bộ nhớ so với kích thước mô hình.

Đặt record_stream=True để tăng tốc độ hơn, với chi phí là mức sử dụng bộ nhớ tăng lên một chút. Tham khảo tài liệu torch.Tensor.record_stream để biết thêm.

Khi use_stream=True trên các VAE có bật tiling, hãy đảm bảo thực hiện một bước forward giả lập (có thể dùng cả các input giả lập) trước khi suy luận để tránh các lỗi không khớp thiết bị. Điều này có thể không hoạt động trên tất cả các cài đặt, vì vậy hãy thoải mái mở một issue nếu bạn gặp bất kỳ vấn đề nào.

Nếu bạn đang sử dụng group offloading ở mức block_level với use_stream được bật, tham số num_blocks_per_group nên được đặt thành 1, nếu không một cảnh báo sẽ được phát ra.

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True, record_stream=True)

Tham số low_cpu_mem_usage có thể được đặt thành True để giảm mức sử dụng bộ nhớ CPU khi dùng streams trong quá trình offloading nhóm. Cách này phù hợp nhất với offloading ở mức leaf_level và khi bộ nhớ CPU là điểm nghẽn. Bộ nhớ được tiết kiệm bằng cách tạo các tensor pinned theo thời gian thực thay vì cố định chúng trước. Tuy nhiên, điều này có thể làm tăng tổng thời gian thực thi.

Chuyển xuống đĩa

Group offloading có thể tiêu thụ một lượng lớn bộ nhớ hệ thống tùy thuộc vào kích thước mô hình. Trên các hệ thống có bộ nhớ hạn chế, hãy thử group offloading ra đĩa như một bộ nhớ phụ.

Đặt đối số offload_to_disk_path trong enable_group_offload() hoặc apply_group_offloading() để offload mô hình xuống đĩa.

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", offload_to_disk_path="path/to/disk")

apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2, offload_to_disk_path="path/to/disk")

Tham khảo hai bảng này để so sánh sự đánh đổi giữa tốc độ và bộ nhớ.

Casting theo từng lớp

Kết hợp layerwise casting với group offloading để tiết kiệm thêm bộ nhớ.

Layerwise casting lưu trữ trọng số ở định dạng dữ liệu nhỏ hơn (ví dụ, torch.float8_e4m3fn và torch.float8_e5m2) để sử dụng ít bộ nhớ hơn và chuyển đổi các trọng số đó lên độ chính xác cao hơn như torch.float16 hoặc torch.bfloat16 cho phép tính toán. Một số lớp (các trọng số liên quan đến chuẩn hóa và điều biến) được bỏ qua vì việc lưu trữ chúng ở fp8 có thể làm giảm chất lượng sinh.

Việc chuyển đổi kiểu theo từng lớp có thể không hoạt động với tất cả các mô hình nếu phần triển khai forward chứa việc chuyển đổi kiểu nội bộ cho các trọng số. Phần triển khai hiện tại của việc chuyển đổi kiểu theo từng lớp giả định rằng quá trình forward độc lập với độ chính xác của trọng số và kiểu dữ liệu đầu vào luôn được chỉ định trong compute_dtype (xem ở đây để biết một phần triển khai không tương thích).

Việc chuyển đổi kiểu dữ liệu theo từng lớp (layerwise casting) cũng có thể thất bại trên các cài đặt mô hình tùy chỉnh có chứa các lớp PEFT. Có một số kiểm tra khả dụng nhưng chúng không được kiểm tra rộng rãi và không được đảm bảo hoạt động trong mọi trường hợp.

Gọi enable_layerwise_casting() để đặt kiểu dữ liệu lưu trữ và tính toán.

import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)
transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)

pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b",
    transformer=transformer,
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"
prompt = (
    "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "
    "The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "
    "pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "
    "casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "
    "The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "
    "atmosphere of this unique musical performance."
)
video = pipeline(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
export_to_video(video, "output.mp4", fps=8)

Phương thức apply_layerwise_casting() cũng có thể được sử dụng nếu bạn cần nhiều quyền kiểm soát và tính linh hoạt hơn. Nó có thể được áp dụng một phần vào các lớp của mô hình bằng cách gọi nó trên các module nội bộ cụ thể. Sử dụng các tham số skip_modules_pattern hoặc skip_modules_classes để chỉ định các module cần tránh, chẳng hạn như các lớp chuẩn hóa và điều biến.

import torch
from diffusers import CogVideoXTransformer3DModel
from diffusers.hooks import apply_layerwise_casting

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)

# skip the normalization layer
apply_layerwise_casting(
    transformer,
    storage_dtype=torch.float8_e4m3fn,
    compute_dtype=torch.bfloat16,
    skip_modules_classes=["norm"],
    non_blocking=True,
)

torch.channels_last

torch.channels_last đảo ngược cách lưu trữ tensor từ (batch size, channels, height, width) sang (batch size, height, width, channels). Điều này giúp tensor khớp với cách phần cứng truy cập tuần tự các tensor được lưu trong bộ nhớ và tránh việc nhảy lung tung trong bộ nhớ để truy cập giá trị pixel.

Không phải tất cả các toán tử hiện tại đều hỗ trợ định dạng channels-last và có thể dẫn đến hiệu suất kém hơn, nhưng vẫn đáng để thử.

print(pipeline.unet.conv_out.state_dict()["weight"].stride())  # (2880, 9, 3, 1)
pipeline.unet.to(memory_format=torch.channels_last)  # in-place operation
print(
    pipeline.unet.conv_out.state_dict()["weight"].stride()
)  # (2880, 1, 960, 320) having a stride of 1 for the 2nd dimension proves that it works

Chú ý tiết kiệm bộ nhớ

Diffusers hỗ trợ nhiều backend attention tiết kiệm bộ nhớ (FlashAttention, xFormers, SageAttention, và nhiều hơn) thông qua set_attention_backend(). Tham khảo hướng dẫn Attention backends để tìm hiểu cách chuyển đổi giữa chúng.

Update on GitHub