torchao

يوفر torchao أنواع بيانات عالية الأداء وتحسينات مبنية على التكميم والتخلخل للاستدلال والتدريب على نماذج PyTorch. وهو مدعوم لأي نموذج في أي نمط، ما دام يدعم التحميل عبر Accelerate ويحتوي على طبقات torch.nn.Linear.

تأكد من تثبيت Pytorch 2.5+ وtorchao باستخدام الأمر أدناه.

uv pip install -U torch torchao

يتوفر كل نوع بيانات (dtype) للتكميم كحالة منفصلة من فئة AOBaseConfig. وهذا يوفر خيارات تكوين أكثر مرونة من خلال كشف المزيد من الحجج المتاحة.

مرّر AOBaseConfig لنوع بيانات التكميم، مثل Int4WeightOnlyConfig إلى TorchAoConfig في from_pretrained().

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" يقوم بتكميم كل طبقة على وحدة معالجة الرسوميات أثناء تحميلها. هذه العملية سريعة، لكنها تتطلب مؤقتًا ذاكرة إضافية لوحدة معالجة الرسوميات للأوزان الأصلية والمكمّمة. إذا كان النموذج يستخدم بالفعل معظم ذاكرة وحدة معالجة الرسوميات لديك، فقد يفشل التحميل بسبب خطأ نفاد الذاكرة. في هذه الحالة، احذف device_map وانقل خط المعالجة إلى وحدة معالجة الرسوميات بعد التحميل:

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، يقوم Diffusers بتكميم الطبقات على وحدة المعالجة المركزية (CPU). هذا الأسلوب أبطأ، لكنه يتجنب الزيادة المؤقتة في ذاكرة وحدة معالجة الرسوميات (GPU) أثناء التكميم. لتقليل استخدام ذاكرة وحدة معالجة الرسوميات (GPU) بشكل أكبر، استخدم enable_model_cpu_offload() بدلاً من ذلك. يمكنك أيضاً تكميم مكونات إضافية، مثل مُشفّر النص.

torch.compile

يدعم torch.compile تسريع الاستدلال بسطر واحد من الكود.

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 وCogVideoX. كما تتوفر معايير أداء إضافية على أجهزة مختلفة في مستودع torchao المستودع.

تُعدّ مخططات الكمّية بعد التدريب بدقة FP8 في torchao فعّالة للوحدات المعالجة الرسومية التي تتمتع بقابلية حسابية لا تقل عن 8.9 (مثل RTX-4090 وHopper وغيرها). غالباً ما توفر FP8 أفضل توازن بين السرعة والذاكرة والجودة عند توليد الصور والفيديوهات. ننصح بدمج FP8 مع torch.compile إذا كانت وحدة المعالجة الرسومية لديك متوافقة.

أنواع التكميم المدعومة

يدعم torchao تكميم الأوزان فقط، وتكميم الأوزان والتفعيلات الديناميكية، لكل من int8 وfloat3-float8 وuint1-uint7.

تخزين التكميم على الأوزان فقط يحفظ أوزان النموذج في نوع بيانات منخفض البتات محدد، لكنه يؤدي العمليات الحسابية بنوع بيانات أعلى دقة، مثل bfloat16. وهذا يقلل من متطلبات الذاكرة الناتجة عن أوزان النموذج، مع الحفاظ على ذروات الذاكرة اللازمة لحساب التنشيطات.

تخزّن الكمية الديناميكية للتفعيلات أوزان النموذج في نوع بيانات منخفض البت، مع تكميم التفعيلات أثناء التشغيل لتوفير ذاكرة إضافية. يقلل هذا من متطلبات الذاكرة الناتجة عن أوزان النموذج، كما يقلل من العبء الذاكرتي الناتج عن حسابات التفعيلات. ومع ذلك، قد يأتي ذلك على حساب الجودة في بعض الأحيان، لذا يُوصى باختبار النماذج المختلفة بشكل شامل.

راجع التوثيق الرسمي لـ torchao لفهم أفضل لطرق التكميم المتاحة. تتوفر قائمة شاملة لخيارات التكوين هنا.

بعض تهيئةات التكميم الشائعة هي على النحو التالي:

تسلسل وفك تسلسل النماذج المُكمّاة

لتسلسل نموذج مُكمَّم في نوع بيانات محدد، حمّل النموذج أولاً بنوع بيانات الكمّ المطلوب ثم احفظه باستخدام طريقة save_pretrained().

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)

لتحميل نموذج كمّي مُسلسَل، استخدم طريقة from_pretrained().

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، فإن بعض طرق التكميم، مثل uint4 للأوزان فقط، لا يمكن تحميلها مباشرة وقد يؤدي ذلك إلى ظهور UnpicklingError عند محاولة تحميل النماذج، لكنها تعمل كما هو متوقع عند حفظها. للتحايل على هذه المشكلة، يمكن تحميل قاموس الحالة يدويًا إلى النموذج. لاحظ، ومع ذلك، أن هذا يتطلب استخدام weights_only=False في torch.load، لذا يجب تشغيله فقط إذا كانت الأوزان قد تم الحصول عليها من مصدر موثوق.

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)

تُدعم واجهة AutoModel لـ PyTorch >= 2.6 كما هو موضح في الأمثلة أدناه.

الموارد

• واجهة TorchAO للكمّية
• أمثلة Diffusers-TorchAO