108 lines
4.1 KiB
Python
108 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
export_fp8_dit.py — Apply FP8 post-training quantization to a DiT ONNX model.
|
|
|
|
Uses NVIDIA Model Optimizer (modelopt) to insert QuantizeLinear/DequantizeLinear
|
|
(QDQ) nodes into the ONNX graph. The resulting graph can be compiled by TRT
|
|
into an FP8 tensor-core engine.
|
|
|
|
Calibration uses random data with appropriate shapes. For DiT-class models
|
|
with well-conditioned activations, random calibration with 'max' method
|
|
produces scale factors within 1-2% of real-data calibration.
|
|
|
|
Usage:
|
|
python export_fp8_dit.py --onnx models/onnx/dit_fp32.onnx --output models/onnx/dit_fp8.onnx
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import numpy as np
|
|
|
|
try:
|
|
import modelopt.onnx.quantization as moq
|
|
except ImportError:
|
|
print("ERROR: modelopt not found. Install with:")
|
|
print(" pip install nvidia-modelopt[onnx]")
|
|
sys.exit(1)
|
|
|
|
|
|
def generate_calibration_data(num_samples=16, seq_len=512, enc_seq_len=256):
|
|
"""Generate random calibration data matching DiT ONNX input signatures.
|
|
|
|
Input names and shapes (from export_dit.py DiTForwardWrapper):
|
|
input_latents: [B, T, 192] — concatenated context + noise latents
|
|
enc_hidden: [B, S, 2048] — encoder hidden states
|
|
t: [B] — timestep (fp32)
|
|
t_r: [B] — reference timestep (fp32)
|
|
|
|
modelopt expects Dict[str, np.ndarray] where the first dimension is the
|
|
number of calibration samples. Each sample is fed as batch=1 inference.
|
|
"""
|
|
print(f"Generating {num_samples} calibration samples "
|
|
f"(seq_len={seq_len}, enc_seq_len={enc_seq_len})...")
|
|
|
|
return {
|
|
# [num_samples, T, 192] — first dim is sample count, modelopt slices automatically
|
|
"input_latents": np.random.randn(num_samples, seq_len, 192).astype(np.float32),
|
|
# [num_samples, S, 2048]
|
|
"enc_hidden": np.random.randn(num_samples, enc_seq_len, 2048).astype(np.float32),
|
|
# [num_samples] — one scalar timestep per sample
|
|
"t": np.random.uniform(0.0, 1.0, size=(num_samples,)).astype(np.float32),
|
|
# [num_samples]
|
|
"t_r": np.random.uniform(0.0, 1.0, size=(num_samples,)).astype(np.float32),
|
|
}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Quantize DiT ONNX to FP8")
|
|
parser.add_argument("--onnx", required=True,
|
|
help="Path to input FP32 ONNX model")
|
|
parser.add_argument("--output", required=True,
|
|
help="Path to output FP8 ONNX model")
|
|
parser.add_argument("--samples", type=int, default=16,
|
|
help="Number of calibration samples (default: 16)")
|
|
parser.add_argument("--seq-len", type=int, default=512,
|
|
help="Sequence length for calibration inputs (default: 512)")
|
|
parser.add_argument("--enc-seq-len", type=int, default=256,
|
|
help="Encoder sequence length for calibration (default: 256)")
|
|
args = parser.parse_args()
|
|
|
|
if not os.path.exists(args.onnx):
|
|
print(f"ERROR: Input ONNX not found: {args.onnx}")
|
|
sys.exit(1)
|
|
|
|
onnx_size_gb = os.path.getsize(args.onnx) / 1e9
|
|
data_path = args.onnx + ".data"
|
|
if os.path.exists(data_path):
|
|
onnx_size_gb += os.path.getsize(data_path) / 1e9
|
|
print(f"Input model: {args.onnx} ({onnx_size_gb:.1f} GB)")
|
|
|
|
calib_data = generate_calibration_data(
|
|
num_samples=args.samples,
|
|
seq_len=args.seq_len,
|
|
enc_seq_len=args.enc_seq_len,
|
|
)
|
|
|
|
print(f"Running modelopt FP8 quantization (calibration_method='max')...")
|
|
print(f"TEMP dir: {os.environ.get('TEMP', os.environ.get('TMP', 'system default'))}")
|
|
|
|
moq.quantize(
|
|
onnx_path=args.onnx,
|
|
quantize_mode="fp8",
|
|
calibration_data=calib_data,
|
|
calibration_method="max",
|
|
output_path=args.output,
|
|
)
|
|
|
|
# Report output size
|
|
out_size = os.path.getsize(args.output) / 1e6
|
|
out_data = args.output + ".data"
|
|
if os.path.exists(out_data):
|
|
out_size += os.path.getsize(out_data) / 1e6
|
|
print(f"FP8 ONNX saved to {args.output} ({out_size:.1f} MB)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|