torchao

torchao는 PyTorch 모델의 추론과 학습을 위한 양자화 및 희소성 기반의 고성능 dtype과 최적화를 제공합니다. Accelerate로 로딩을 지원하고 torch.nn.Linear 레이어를 포함하는 모든 모달리티의 모든 모델에서 사용할 수 있습니다.

아래 명령어로 Pytorch 2.5+와 torchao가 설치되어 있는지 확인하세요.

uv pip install -U torch torchao

각 양자화 dtype는 AOBaseConfig 클래스의 개별 인스턴스로 사용할 수 있습니다. 이는 더 많은 사용 가능한 인수를 노출시켜 보다 유연한 구성 옵션을 제공합니다.

Pass the AOBaseConfig of a quantization dtype, like Int4WeightOnlyConfig to TorchAoConfig in 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"는 로딩 중 각 레이어를 GPU에서 양자화합니다. 이 방식은 빠르지만, 원래 가중치와 양자화된 가중치를 위해 일시적으로 추가 GPU 메모리가 필요합니다. 모델이 이미 GPU 메모리의 대부분을 사용 중이라면, 메모리 부족 오류로 인해 로딩이 실패할 수 있습니다. 이 경우 device_map을 제거하고 로딩 후 파이프라인을 GPU로 이동하세요:

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

torchao는 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 저장소에서도 확인할 수 있습니다.

torchao의 FP8 사후 학습 양자화 스킴은 컴퓨트 캐피빌리티가 8.9 이상인 GPU(RTX-4090, Hopper 등)에서 효과적입니다. 이미지와 비디오 생성 시 FP8은 속도, 메모리, 품질 간의 최적의 트레이드오프를 제공하는 경우가 많습니다. GPU가 호환되는 경우 FP8과 torch.compile을 함께 사용하는 것을 권장합니다.

지원되는 양자화 유형

torchao는 int8, float3-float8, uint1-uint7에 대해 가중치 전용 양자화 및 가중치와 동적 활성화 양자화를 지원합니다.

Weight-only quantization은 모델 가중치를 특정 저비트 데이터 형식으로 저장하지만, bfloat16와 같은 더 높은 정밀도의 데이터 형식으로 계산을 수행합니다. 이를 통해 모델 가중치에 따른 메모리 요구량을 줄이되, 활성화 계산에 따른 메모리 피크는 유지됩니다.

Dynamic activation quantization은 모델 가중치를 저비트 dtype로 저장하는 동시에, 추론 시점에 활성화 값을 동적으로 양자화하여 추가적인 메모리를 절약합니다. 이를 통해 모델 가중치로 인한 메모리 요구량이 줄어드는 것은 물론, 활성화 계산으로 인한 메모리 오버헤드도 낮아집니다. 다만, 경우에 따라 품질이 저하될 수 있으므로 다양한 모델을 철저히 테스트하는 것이 권장됩니다.

사용 가능한 양자화 방법에 대한 더 자세한 내용은 공식 torchao 문서를 참고하세요. 전체 구성 옵션 목록은 여기에서 확인할 수 있습니다.

다음은 몇 가지 인기 있는 양자화 구성 예시입니다:

카테고리	구성 클래스
정수 양자화	`Int4WeightOnlyConfig`, `Int8WeightOnlyConfig`, `Int8DynamicActivationInt8WeightConfig`
부동 소수점 8비트 양자화	`Float8WeightOnlyConfig`, `Float8DynamicActivationFloat8WeightConfig`
부호 없는 정수 양자화	`IntxWeightOnlyConfig`

양자화 모델 직렬화 및 역직렬화

주어진 dtype로 양자화된 모델을 직렬화하려면, 먼저 원하는 양자화 dtype로 모델을 로드한 다음 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 weight-only와 같은 일부 양자화 방법은 직접 로드할 수 없으며, 모델을 로드하려고 할 때 UnpicklingError가 발생할 수 있지만, 저장할 때는 정상적으로 작동합니다. 이를 우회하려면 state dict를 모델에 수동으로 로드할 수 있습니다. 단, torch.load에서 weights_only=False를 사용해야 하므로, 가중치가 신뢰할 수 있는 출처에서 가져온 경우에만 실행해야 합니다.

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)

아래 예시와 같이 PyTorch >= 2.6에서 AutoModel API가 지원됩니다.

리소스

Update on GitHub