NVIDIA ModelOpt

NVIDIA-ModelOpt는 양자화, 프루닝, 증류, 추측 디코딩 등 최첨단 모델 최적화 기법을 통합한 라이브러리입니다. TensorRT-LLM 또는 TensorRT와 같은 다운스트림 배포 프레임워크에서 딥러닝 모델을 압축하여 추론 속도를 최적화합니다.

시작하기 전에 nvidia_modelopt가 설치되어 있는지 확인하세요.

pip install -U "nvidia_modelopt[hf]"

NVIDIAModelOptConfig를 from_pretrained()에 전달하여 모델을 양자화합니다(사전 양자화된 모델을 로드할 수도 있습니다). 이는 Accelerate로 로딩을 지원하고 torch.nn.Linear 레이어를 포함하는 모든 모달리티의 모든 모델에 적용됩니다.

…

import torch
from diffusers import AutoModel, SanaPipeline, NVIDIAModelOptConfig

model_id = "Efficient-Large-Model/Sana_600M_1024px_diffusers"
dtype = torch.bfloat16

quantization_config = NVIDIAModelOptConfig(quant_type="FP8", quant_method="modelopt")
transformer = AutoModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=dtype,
)
pipe = SanaPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=dtype,
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

print(f"Pipeline memory usage: {torch.cuda.max_memory_reserved() / 1024**3:.3f} GB")

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

참고:

NVIDIA-ModelOpt의 양자화 방법은 모델 성능을 유지하면서 다양한 QAT(Quantization-Aware Training) 및 PTQ(Post-Training Quantization) 기법을 활용하여 모델 가중치의 메모리 사용량을 줄이는 데 설계되었습니다. 그러나 추론 시 실제 성능 향상은 배포 프레임워크(예: TRT-LLM, TensorRT)와 특정 하드웨어 구성에 따라 달라집니다.

자세한 내용은 여기에서 확인할 수 있습니다.

NVIDIAModelOptConfig

NVIDIAModelOptConfig 클래스는 세 가지 파라미터를 받습니다:

quant_type: 아래 양자화 유형 중 하나를 나타내는 문자열 값.
modules_to_not_convert: 양자화를 수행하지 않을 모듈의 전체/부분 모듈 이름 목록입니다. 예를 들어, SD3Transformer2DModel의 pos_embed 투영 블록에 양자화를 수행하지 않으려면 다음과 같이 지정합니다: modules_to_not_convert=["pos_embed.proj.weight"].
disable_conv_quantization: 모델의 모든 컨볼루션 레이어에 대한 양자화를 비활성화하려면 True로 설정하는 불리언 값입니다. 채널 및 블록 양자화는 일반적으로 컨볼루션 레이어(INT4, NF4, NVFP4와 함께 사용)와 잘 동작하지 않으므로 유용합니다. 특정 컨볼루션 레이어에 대해서만 양자화를 비활성화하려면 대신 modules_to_not_convert를 사용하세요.
algorithm: 스케일을 결정하는 데 사용할 알고리즘입니다. 기본값은 "max"입니다. 추가 알고리즘 및 세부 사항은 modelopt 문서를 확인하세요.
forward_loop: 양자화 과정에서 활성화 값을 보정하는 데 사용할 forward loop 함수입니다. 제공되지 않으면 가중치만 사용하여 계산된 정적 스케일 값에 의존합니다.
kwargs: quant_type에 따라 호출되는 기본 양자화 메서드에 전달할 키워드 인자의 딕셔너리.

지원되는 양자화 유형

ModelOpt는 weight-only, channel 및 block 양자화를 지원하며, int8, fp8, int4, nf4, nvfp4 포맷을 사용할 수 있습니다. 이러한 양자화 방식은 추론 시 모델의 성능을 유지하면서 모델 가중치의 메모리 사용량을 줄이는 것을 목표로 설계되었습니다.

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

지원되는 양자화 방법은 다음과 같습니다:

양자화 유형	지원되는 스키마	필수 Kwargs	추가 참고 사항
INT8	`int8 weight only`, `int8 channel quantization`, `int8 block quantization`	`quant_type`, `quant_type + channel_quantize`, `quant_type + channel_quantize + block_quantize`
FP8	`fp8 weight only`, `fp8 channel quantization`, `fp8 block quantization`	`quant_type`, `quant_type + channel_quantize`, `quant_type + channel_quantize + block_quantize`
INT4	`int4 weight only`, `int4 block quantization`	`quant_type`, `quant_type + channel_quantize + block_quantize`	`channel_quantize = -1 is only supported for now`
NF4	`nf4 weight only`, `nf4 double block quantization`	`quant_type`, `quant_type + channel_quantize + block_quantize + scale_channel_quantize` + `scale_block_quantize`	`channel_quantize = -1 and scale_channel_quantize = -1 are only supported for now`
NVFP4	`nvfp4 weight only`, `nvfp4 block quantization`	`quant_type`, `quant_type + channel_quantize + block_quantize`	`channel_quantize = -1 is only supported for now`

사용 가능한 양자화 방법과 사용 가능한 구성 옵션의 전체 목록에 대한 더 자세한 이해를 위해 공식 modelopt 문서를 참조하세요.

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

주어진 dtype로 양자화된 모델을 직렬화하려면, 먼저 원하는 양자화 dtype로 모델을 로드한 다음 save_pretrained() 메서드를 사용하여 저장합니다.

import torch
from diffusers import AutoModel, NVIDIAModelOptConfig
from modelopt.torch.opt import enable_huggingface_checkpointing

enable_huggingface_checkpointing()

model_id = "Efficient-Large-Model/Sana_600M_1024px_diffusers"
quant_config_fp8 = {"quant_type": "FP8", "quant_method": "modelopt"}
quant_config_fp8 = NVIDIAModelOptConfig(**quant_config_fp8)
model = AutoModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quant_config_fp8,
    dtype=torch.bfloat16,
)
model.save_pretrained('path/to/sana_fp8', safe_serialization=False)

시리얼라이즈된 양자화 모델을 로드하려면 from_pretrained() 메서드를 사용합니다.

import torch
from diffusers import AutoModel, NVIDIAModelOptConfig, SanaPipeline
from modelopt.torch.opt import enable_huggingface_checkpointing

enable_huggingface_checkpointing()

quantization_config = NVIDIAModelOptConfig(quant_type="FP8", quant_method="modelopt")
transformer = AutoModel.from_pretrained(
    "path/to/sana_fp8",
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
pipe = SanaPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_600M_1024px_diffusers",
    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=50, guidance_scale=4.5, max_sequence_length=512
).images[0]
image.save("output.png")

Update on GitHub

Diffusers

NVIDIA ModelOpt

NVIDIAModelOptConfig

지원되는 양자화 유형

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