wehub-resource-sync
79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# You may not use this file except in compliance with the License.
|
|
# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
"""TorchAO FP8 post-training quantization helpers."""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
# Small conditioning/output projections are both numerically sensitive and too
|
|
# small to amortize dynamic activation quantization on H100.
|
|
_BF16_MODULES = {
|
|
"text_embedding.0",
|
|
"text_embedding.2",
|
|
"time_embedding.0",
|
|
"time_embedding.2",
|
|
"time_projection.1",
|
|
"head.head",
|
|
}
|
|
|
|
|
|
def quantize_model_fp8(model: nn.Module, *, verbose: bool = False) -> int:
|
|
"""Quantize compatible BF16 linear layers to row-wise dynamic FP8 in-place."""
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("TorchAO FP8 inference requires a CUDA GPU.")
|
|
|
|
device = next(model.parameters()).device
|
|
if device.type != "cuda":
|
|
raise ValueError("Move the BF16 model to CUDA before applying FP8 quantization.")
|
|
if torch.cuda.get_device_capability(device) < (8, 9):
|
|
raise RuntimeError("TorchAO FP8 inference requires compute capability 8.9 or newer.")
|
|
|
|
try:
|
|
from torchao.quantization import (
|
|
Float8DynamicActivationFloat8WeightConfig,
|
|
PerRow,
|
|
quantize_,
|
|
)
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"FP8 inference requires a TorchAO version compatible with the installed PyTorch."
|
|
) from exc
|
|
|
|
quantized_names = []
|
|
skipped_names = []
|
|
|
|
def filter_fn(module: nn.Module, fqn: str) -> bool:
|
|
if not isinstance(module, nn.Linear):
|
|
return False
|
|
if fqn in _BF16_MODULES:
|
|
skipped_names.append(fqn)
|
|
return False
|
|
if module.weight.dtype != torch.bfloat16:
|
|
raise TypeError(f"FP8 layer {fqn!r} must be BF16, got {module.weight.dtype}.")
|
|
out_features, in_features = module.weight.shape
|
|
if in_features % 16 or out_features % 16:
|
|
skipped_names.append(fqn)
|
|
return False
|
|
quantized_names.append(fqn)
|
|
return True
|
|
|
|
quantize_(
|
|
model,
|
|
Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()),
|
|
filter_fn=filter_fn,
|
|
)
|
|
if verbose:
|
|
print(
|
|
f"[FP8] TorchAO W8A8 quantized {len(quantized_names)} linear layers "
|
|
f"with row-wise scaling; kept {len(skipped_names)} layers in BF16"
|
|
)
|
|
return len(quantized_names)
|