torchao

torchao, PyTorch modellerinin çıkarım ve eğitiminde kullanılmak üzere, kuantizasyon ve seyreklik tabanlı yüksek performanslı dtypes ve optimizasyonlar sağlar. Herhangi bir modalitedeki herhangi bir model için desteklenir; yeter ki Accelerate ile yüklenmeyi desteklesin ve torch.nn.Linear katmanları içersin.

Aşağıdaki komutla Pytorch 2.5+ ve torchao'nun yüklendiğinden emin olun.

uv pip install -U torch torchao

Her kuantizasyon dtype'ı, AOBaseConfig sınıfının ayrı bir örneği olarak kullanılabilir. Bu, daha fazla kullanılabilir argüman açığa çıkararak daha esnek yapılandırma seçenekleri sunar.

AOBaseConfig'ini, Int4WeightOnlyConfig gibi bir kuantizasyon veri türü olarak, from_pretrained() içindeki TorchAoConfig'ye iletin.

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(group_size=128, version=2))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)

device_map="cuda", her katmanı yüklenirken GPU üzerinde kuantize eder. Bu hızlıdır, ancak orijinal ve kuantize ağırlıklar için geçici olarak ek GPU belleği gerektirir. Model zaten GPU belleğinizin çoğunu kullanıyorsa, yükleme bellek yetersizliği (out-of-memory) hatasıyla başarısız olabilir. Bu durumda, device_map parametresini kaldırın ve pipeline'ı yükleme işleminden sonra GPU'ya taşıyın:

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    torch_dtype=torch.bfloat16,
)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"

device_map olmadan Diffusers, katmanları CPU üzerinde kuantize eder. Bu daha yavaştır ancak kuantizasyon sırasında geçici GPU bellek artışını önler. GPU bellek kullanımını daha da azaltmak için bunun yerine enable_model_cpu_offload() kullanın. Metin kodlayıcı gibi ek bileşenleri de kuantize edebilirsiniz.

torch.compile

torchao, tek satır kodla çıkarımı hızlandırabilen torch.compile özelliğini destekler.

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int4WeightOnlyConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)

pipeline.transformer.compile(transformer, mode="max-autotune", fullgraph=True)

Flux ve CogVideoX için çıkarım hızı ve bellek kullanımı ölçümlerine ilişkin bu tabloya bakın. Çeşitli donanımlara ait daha fazla ölçüm, torchao deposunda da mevcuttur.

torchao içindeki FP8 eğitim sonrası nicelleştirme şemaları, hesaplama yeteneği en az 8.9 olan GPU'lar (RTX-4090, Hopper vb.) için etkilidir. Görsel ve video üretirken FP8 genellikle hız, bellek ve kalite arasında en iyi dengeyi sağlar. GPU'nuz uyumluysa FP8'i torch.compile ile birleştirmenizi öneririz.

Desteklenen kuantizasyon türleri

torchao, int8, float3-float8 ve uint1-uint7 için yalnızca ağırlık kuantizasyonu ile ağırlık ve dinamik aktivasyon kuantizasyonunu destekler.

Sadece ağırlık kuantizasyonu, model ağırlıklarını belirli bir düşük bit veri türünde saklar ancak bfloat16 gibi daha yüksek hassasiyetli bir veri türüyle hesaplama yapar. Bu, model ağırlıklarından kaynaklanan bellek gereksinimini azaltır ancak aktivasyon hesaplama için bellek zirvelerini korur.

Dinamik aktivasyon kuantizasyonu, model ağırlıklarını düşük bitlik bir veri türünde (dtype) saklarken, ek bellek tasarrufu sağlamak için aktivasyonları da anlık olarak kuantize eder. Bu, model ağırlıklarından kaynaklanan bellek gereksinimini azaltırken, aktivasyon hesaplamalarından kaynaklanan bellek yükünü de düşürür. Ancak bu, bazen kalite açısından bir ödün gerektirebilir; bu nedenle farklı modelleri kapsamlı bir şekilde test etmeniz önerilir.

Mevcut kuantizasyon yöntemlerini daha iyi anlamak için resmi torchao dokümantasyonuna başvurun. Yapılandırma seçeneklerinin kapsamlı bir listesi burada mevcuttur.

Bazı popüler örnek kuantizasyon yapılandırmaları aşağıdaki gibidir:

Kantize modellerin seri hale getirilmesi ve dizide çıkarılması

Belirtilen bir dtype ile kuantize edilmiş bir modeli serialize etmek için, önce modeli istenen kuantizasyon dtype'i ile yükleyin ve ardından save_pretrained() yöntemini kullanarak kaydedin.

import torch
from diffusers import AutoModel, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig

quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_int8wo", safe_serialization=False)

Serileştirilmiş bir kuantize modeli yüklemek için from_pretrained() yöntemini kullanın.

import torch
from diffusers import FluxPipeline, AutoModel

transformer = AutoModel.from_pretrained("/path/to/flux_int8wo", dtype=torch.bfloat16, use_safetensors=False)
pipe = FluxPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", transformer=transformer, dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.0).images[0]
image.save("output.png")

torch<=2.6.0 kullanıyorsanız, uint4 weight-only gibi bazı kuantizasyon yöntemleri doğrudan yüklenemez ve modelleri yüklemeye çalışırken UnpicklingError ile sonuçlanabilir, ancak kaydetme sırasında beklenen şekilde çalışır. Bu durumu aşmak için state dict'i modele manuel olarak yükleyebilirsiniz. Bununla birlikte, bunun torch.load içinde weights_only=False kullanılmasını gerektirdiğini unutmayın; bu nedenle, ağırlıkların güvenilir bir kaynaktan elde edildiğinden emin olmanız durumunda çalıştırılmalıdır.

import torch
from accelerate import init_empty_weights
from diffusers import FluxPipeline, AutoModel, TorchAoConfig
from torchao.quantization import IntxWeightOnlyConfig

# Serialize the model
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=TorchAoConfig(IntxWeightOnlyConfig(dtype=torch.uint4)),
    dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_uint4wo", safe_serialization=False, max_shard_size="50GB")
# ...

# Load the model
state_dict = torch.load("/path/to/flux_uint4wo/diffusion_pytorch_model.bin", weights_only=False, map_location="cpu")
with init_empty_weights():
    transformer = AutoModel.from_config("/path/to/flux_uint4wo/config.json")
transformer.load_state_dict(state_dict, strict=True, assign=True)

Aşağıdaki örneklerde gösterildiği gibi AutoModel API'si PyTorch >= 2.6 için desteklenmektedir.

Kaynaklar

• TorchAO Niceliklendirme API'si
• Diffusers-TorchAO örnekleri