mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
[2/3] Implicit Gemm NVFP4 (#1227)
### What does this PR do? Type of change: new feature <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> - Add Conv3D implicit GEMM kernel with BF16 WMMA tensor cores and fused NVFP4 activation quantization for video diffusion VAE layers - Integrate into _QuantConv3d via QuantModuleRegistry — automatically dispatched when NVFP4 quantization is applied to nn.Conv3d - Move kernel from `experimental/conv/ to modelopt/torch/kernels/conv/`; move tests to `tests/gpu/torch/quantization/kernels/` ### Testing <!-- Mention how have you tested your change if applicable. --> - Added test cases to measure the difference between cuDNN and our CUDA implicit GEMM kernel - Added an NVFP4 fake quantization test using CUDA code ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ <!--- If ❌, explain why. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ <!--- Mandatory --> - Did you write any new necessary tests?: ✅ <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Per-backbone quantization/export in a single run with per-backbone checkpoints and backbone-aware quant filters * Configurable NVFP4 block-size via CLI/config; improved NVFP4 Conv3D inference path and Wan 2.2 quantization support * **Bug Fixes** * Video-model calibration now respects extra params and forces video decoding during calibration * **Documentation** * Added comprehensive Conv3D implicit‑GEMM kernel documentation; removed experimental Conv3D prototype docs/benchmark * **Tests** * New Wan 2.2 quantization/export tests and expanded Conv3D/FP4 kernel test coverage <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
This commit is contained in:
jingyu-ml
@@ -11,15 +11,3 @@ coverage:
|
||||
target: auto
|
||||
threshold: 1% # Allow atmost 1% coverage drop from main branch.
|
||||
patch: false
|
||||
|
||||
# Exclude GPU-only Triton kernel files from ALL codecov calculations (project
|
||||
# and patch checks, all flags). Rationale: these files are dominated by
|
||||
# @triton.jit kernel bodies that CPU unit tests cannot exercise. GPU tests
|
||||
# cover them end-to-end (see tests/gpu/torch/sparsity/attention_sparsity/) but
|
||||
# the `gpu`-flag upload may race with the PR status check, so relying on flag
|
||||
# combination alone leaves the project check flaky. Dropping these files here
|
||||
# makes the check deterministic — local `pytest --cov` and GPU runs still
|
||||
# measure them; only the codecov PR status ignores them.
|
||||
ignore:
|
||||
- "modelopt/torch/kernels/triton_fa.py"
|
||||
- "modelopt/torch/kernels/hf_triton_attention.py"
|
||||
|
||||
@@ -16,6 +16,7 @@ Changelog
|
||||
- Add support for vLLM fakequant reload using ModelOpt state for HF models. See `examples/vllm_serve/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/vllm_serve#load-qatptq-model-and-serve-in-vllm-wip>`_ for more details.
|
||||
- [Early Testing] Add Claude Code PTQ skill (``.claude/skills/ptq/``) for agent-assisted post-training quantization. The skill guides the agent through environment detection, model support checking, format selection, and execution via the launcher or manual SLURM/Docker/bare GPU paths. Includes handling for unlisted models with custom module patching. This feature is in early testing — use with caution.
|
||||
- Add performant layerwise calibration for large models that don't fit on GPU (e.g. DeepSeek-R1, Kimi-K2). See `modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml <https://github.com/NVIDIA/Model-Optimizer/blob/main/modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml>`_ for usage. Layerwise calibration also supports PTQ with intermediate progress saving — useful when long PTQ runs get hit with Slurm timeouts. See `modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yaml <https://github.com/NVIDIA/Model-Optimizer/blob/main/modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yaml>`_ for usage.
|
||||
- Add implicit GEMM CUDA kernel for Conv3D with fused NVFP4 fake quantization (``modelopt.torch.quantization.src.conv``). When NVFP4 quantization is applied to an ``nn.Conv3d`` layer via ModelOpt PTQ, the implicit GEMM path is used automatically instead of cuDNN. Uses BF16 WMMA tensor cores (SM80+) with FP32 accumulation and in-kernel FP4 (E2M1) activation quantization. Grouped convolution (``groups > 1``) falls back to the default cuDNN path. Inference only — training mode falls back to cuDNN with a warning.
|
||||
|
||||
**Backward Breaking Changes**
|
||||
|
||||
|
||||
@@ -117,6 +117,20 @@ python quantize.py \
|
||||
--hf-ckpt-dir ./hf_ckpt
|
||||
```
|
||||
|
||||
#### Wan 2.2 VAE NVFP4 (Conv3D Implicit GEMM)
|
||||
|
||||
The Wan 2.2 VAE (`AutoencoderKLWan`, shared between the 5B and 14B pipelines) is built from 3D convolutions. When quantizing the VAE with NVFP4, the `Conv3d` layers are automatically dispatched through a custom BF16 WMMA implicit-GEMM kernel with fused FP4 activation quantization. Requires SM80+ (Ampere or newer). See [`modelopt/torch/quantization/src/conv/README.md`](../../modelopt/torch/quantization/src/conv/README.md) for kernel details.
|
||||
|
||||
```sh
|
||||
python quantize.py \
|
||||
--model {wan2.2-t2v-14b|wan2.2-t2v-5b} \
|
||||
--backbone vae \
|
||||
--format fp4 --quant-algo max --collect-method default \
|
||||
--model-dtype BFloat16 --trt-high-precision-dtype BFloat16 \
|
||||
--batch-size 1 --calib-size 32 --n-steps 30 \
|
||||
--quantized-torch-ckpt-save-path ./wan22_vae_fp4.pt
|
||||
```
|
||||
|
||||
#### [LTX-2](https://github.com/Lightricks/LTX-2) FP4
|
||||
|
||||
> [!WARNING]
|
||||
|
||||
@@ -108,11 +108,12 @@ class Calibrator:
|
||||
def _run_wan_video_calibration(
|
||||
self, prompt_batch: list[str], extra_args: dict[str, Any]
|
||||
) -> None:
|
||||
extra_params = self.pipeline_manager.config.extra_params
|
||||
kwargs = {}
|
||||
kwargs["negative_prompt"] = extra_args["negative_prompt"]
|
||||
kwargs["height"] = extra_args["height"]
|
||||
kwargs["width"] = extra_args["width"]
|
||||
kwargs["num_frames"] = extra_args["num_frames"]
|
||||
kwargs["height"] = extra_params.get("height", extra_args["height"])
|
||||
kwargs["width"] = extra_params.get("width", extra_args["width"])
|
||||
kwargs["num_frames"] = extra_params.get("num_frames", extra_args["num_frames"])
|
||||
kwargs["guidance_scale"] = extra_args["guidance_scale"]
|
||||
if "guidance_scale_2" in extra_args:
|
||||
kwargs["guidance_scale_2"] = extra_args["guidance_scale_2"]
|
||||
@@ -154,7 +155,11 @@ class Calibrator:
|
||||
"images": extra_params.get("images", []),
|
||||
"tiling_config": extra_params.get("tiling_config", TilingConfig.default()),
|
||||
}
|
||||
self.pipe(prompt=prompt, **kwargs)
|
||||
decoded_video, decoded_audio = self.pipe(prompt=prompt, **kwargs)
|
||||
# vae_decode_video returns a lazy generator — consume it so the
|
||||
# video decoder's forward() actually runs during calibration.
|
||||
for _ in decoded_video:
|
||||
pass
|
||||
|
||||
def _run_ltx_video_calibration(
|
||||
self, prompt_batch: list[str], extra_args: dict[str, Any]
|
||||
|
||||
@@ -33,7 +33,9 @@ except ImportError:
|
||||
from utils import (
|
||||
filter_func_default,
|
||||
filter_func_flux_dev,
|
||||
filter_func_ltx2_vae,
|
||||
filter_func_ltx_video,
|
||||
filter_func_wan_vae,
|
||||
filter_func_wan_video,
|
||||
)
|
||||
|
||||
@@ -54,31 +56,30 @@ class ModelType(str, Enum):
|
||||
WAN22_T2V_5b = "wan2.2-t2v-5b"
|
||||
|
||||
|
||||
def get_model_filter_func(model_type: ModelType) -> Callable[[str], bool]:
|
||||
"""
|
||||
Get the appropriate filter function for a given model type.
|
||||
_FILTER_FUNC_MAP: dict[ModelType, Callable[[str], bool]] = {
|
||||
ModelType.FLUX_DEV: filter_func_flux_dev,
|
||||
ModelType.FLUX2_DEV: filter_func_flux_dev,
|
||||
ModelType.LTX_VIDEO_DEV: filter_func_ltx_video,
|
||||
ModelType.LTX2: filter_func_ltx_video,
|
||||
ModelType.WAN22_T2V_14b: filter_func_wan_video,
|
||||
ModelType.WAN22_T2V_5b: filter_func_wan_video,
|
||||
}
|
||||
|
||||
Args:
|
||||
model_type: The model type enum
|
||||
_VAE_FILTER_FUNC_MAP: dict[tuple[ModelType, str], Callable[[str], bool]] = {
|
||||
(ModelType.LTX2, "video_decoder"): filter_func_ltx2_vae,
|
||||
(ModelType.WAN22_T2V_14b, "vae"): filter_func_wan_vae,
|
||||
(ModelType.WAN22_T2V_5b, "vae"): filter_func_wan_vae,
|
||||
}
|
||||
|
||||
Returns:
|
||||
A filter function appropriate for the model type
|
||||
"""
|
||||
filter_func_map = {
|
||||
ModelType.FLUX_DEV: filter_func_flux_dev,
|
||||
ModelType.FLUX_SCHNELL: filter_func_default,
|
||||
ModelType.FLUX2_DEV: filter_func_flux_dev,
|
||||
ModelType.SDXL_BASE: filter_func_default,
|
||||
ModelType.SDXL_TURBO: filter_func_default,
|
||||
ModelType.SD3_MEDIUM: filter_func_default,
|
||||
ModelType.SD35_MEDIUM: filter_func_default,
|
||||
ModelType.LTX_VIDEO_DEV: filter_func_ltx_video,
|
||||
ModelType.LTX2: filter_func_ltx_video,
|
||||
ModelType.WAN22_T2V_14b: filter_func_wan_video,
|
||||
ModelType.WAN22_T2V_5b: filter_func_wan_video,
|
||||
}
|
||||
|
||||
return filter_func_map.get(model_type, filter_func_default)
|
||||
def get_model_filter_func(
|
||||
model_type: ModelType, backbone_name: str = "transformer"
|
||||
) -> Callable[[str], bool]:
|
||||
"""Get the appropriate filter function for a given model type and backbone."""
|
||||
vae_func = _VAE_FILTER_FUNC_MAP.get((model_type, backbone_name))
|
||||
if vae_func is not None:
|
||||
return vae_func
|
||||
return _FILTER_FUNC_MAP.get(model_type, filter_func_default)
|
||||
|
||||
|
||||
# Model registry with HuggingFace model IDs
|
||||
|
||||
@@ -42,6 +42,7 @@ class PipelineManager:
|
||||
self.pipe: Any | None = None
|
||||
self.pipe_upsample: LTXLatentUpsamplePipeline | None = None # For LTX-Video upsampling
|
||||
self._transformer: torch.nn.Module | None = None
|
||||
self._video_decoder: torch.nn.Module | None = None
|
||||
|
||||
@staticmethod
|
||||
def create_pipeline_from(
|
||||
@@ -58,23 +59,20 @@ class PipelineManager:
|
||||
Raises:
|
||||
ValueError: If model type is unsupported
|
||||
"""
|
||||
try:
|
||||
pipeline_cls = MODEL_PIPELINE[model_type]
|
||||
if pipeline_cls is None:
|
||||
raise ValueError(f"Model type {model_type.value} does not use diffusers pipelines.")
|
||||
model_id = (
|
||||
MODEL_REGISTRY[model_type] if override_model_path is None else override_model_path
|
||||
)
|
||||
pipe = pipeline_cls.from_pretrained(
|
||||
model_id,
|
||||
torch_dtype=torch_dtype,
|
||||
use_safetensors=True,
|
||||
**MODEL_DEFAULTS[model_type].get("from_pretrained_extra_args", {}),
|
||||
)
|
||||
pipe.set_progress_bar_config(disable=True)
|
||||
return pipe
|
||||
except Exception as e:
|
||||
raise e
|
||||
pipeline_cls = MODEL_PIPELINE[model_type]
|
||||
if pipeline_cls is None:
|
||||
raise ValueError(f"Model type {model_type.value} does not use diffusers pipelines.")
|
||||
model_id = (
|
||||
MODEL_REGISTRY[model_type] if override_model_path is None else override_model_path
|
||||
)
|
||||
pipe = pipeline_cls.from_pretrained(
|
||||
model_id,
|
||||
torch_dtype=torch_dtype,
|
||||
use_safetensors=True,
|
||||
**MODEL_DEFAULTS[model_type].get("from_pretrained_extra_args", {}),
|
||||
)
|
||||
pipe.set_progress_bar_config(disable=True)
|
||||
return pipe
|
||||
|
||||
def create_pipeline(self) -> Any:
|
||||
"""
|
||||
@@ -157,42 +155,32 @@ class PipelineManager:
|
||||
self.logger.info("Enabling VAE tiling for LTX-Video")
|
||||
self.pipe.vae.enable_tiling()
|
||||
|
||||
def get_backbone(self) -> torch.nn.Module:
|
||||
"""
|
||||
Get the backbone model (transformer or UNet).
|
||||
|
||||
Returns:
|
||||
Backbone model module
|
||||
"""
|
||||
if not self.pipe:
|
||||
raise RuntimeError("Pipeline not created. Call create_pipeline() first.")
|
||||
|
||||
backbone_pairs = list(self.iter_backbones())
|
||||
if len(backbone_pairs) == 1:
|
||||
return backbone_pairs[0][1]
|
||||
return torch.nn.ModuleList([module for _, module in backbone_pairs])
|
||||
|
||||
def iter_backbones(self) -> Iterator[tuple[str, torch.nn.Module]]:
|
||||
"""
|
||||
Yield backbone modules by name, based on a backbone spec.
|
||||
|
||||
Yields:
|
||||
(backbone_name, module) pairs
|
||||
Yield (backbone_name, module) pairs.
|
||||
"""
|
||||
if not self.pipe:
|
||||
raise RuntimeError("Pipeline not created. Call create_pipeline() first.")
|
||||
|
||||
names = list(self.config.backbone)
|
||||
|
||||
if self.config.model_type == ModelType.LTX2:
|
||||
self._ensure_ltx2_transformer_cached()
|
||||
name = names[0] if names else "transformer"
|
||||
yield name, self._transformer
|
||||
return
|
||||
|
||||
if not names:
|
||||
raise RuntimeError("No backbone names provided.")
|
||||
|
||||
if self.config.model_type == ModelType.LTX2:
|
||||
for name in names:
|
||||
if name == "video_decoder":
|
||||
self._ensure_ltx2_video_decoder_cached()
|
||||
yield name, self._video_decoder
|
||||
elif name == "transformer":
|
||||
self._ensure_ltx2_transformer_cached()
|
||||
yield name, self._transformer
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported LTX-2 backbone name '{name}'. "
|
||||
"Expected 'transformer' or 'video_decoder'."
|
||||
)
|
||||
return
|
||||
|
||||
for name in names:
|
||||
module = getattr(self.pipe, name, None)
|
||||
if module is None:
|
||||
@@ -207,6 +195,16 @@ class PipelineManager:
|
||||
self.pipe.stage_1_model_ledger.transformer = lambda: transformer
|
||||
self._transformer = transformer
|
||||
|
||||
def _ensure_ltx2_video_decoder_cached(self) -> None:
|
||||
if not self.pipe:
|
||||
raise RuntimeError("Pipeline not created. Call create_pipeline() first.")
|
||||
if self._video_decoder is None:
|
||||
video_decoder = self.pipe.stage_1_model_ledger.video_decoder()
|
||||
# Cache it so subsequent calls return the same (quantized) instance
|
||||
self.pipe.stage_1_model_ledger.video_decoder = lambda: video_decoder
|
||||
self.pipe.stage_2_model_ledger.video_decoder = lambda: video_decoder
|
||||
self._video_decoder = video_decoder
|
||||
|
||||
def _create_ltx2_pipeline(self) -> Any:
|
||||
params = dict(self.config.extra_params)
|
||||
checkpoint_path = params.pop("checkpoint_path", None)
|
||||
@@ -261,7 +259,6 @@ class PipelineManager:
|
||||
return TI2VidTwoStagesPipeline(**pipeline_kwargs)
|
||||
|
||||
def print_quant_summary(self):
|
||||
backbone_pairs = list(self.iter_backbones())
|
||||
for name, backbone in backbone_pairs:
|
||||
for name, backbone in self.iter_backbones():
|
||||
self.logger.info(f"{name} quantization info:")
|
||||
mtq.print_quant_summary(backbone)
|
||||
|
||||
@@ -116,33 +116,48 @@ class Quantizer:
|
||||
|
||||
if self.config.format == QuantFormat.INT8:
|
||||
if self.config.algo == QuantAlgo.SMOOTHQUANT:
|
||||
quant_config = mtq.INT8_SMOOTHQUANT_CFG
|
||||
base_cfg = mtq.INT8_SMOOTHQUANT_CFG
|
||||
else:
|
||||
quant_config = INT8_DEFAULT_CONFIG
|
||||
base_cfg = INT8_DEFAULT_CONFIG
|
||||
if self.config.collect_method != CollectMethod.DEFAULT:
|
||||
reset_set_int8_config(
|
||||
quant_config,
|
||||
base_cfg,
|
||||
self.config.percentile,
|
||||
n_steps,
|
||||
collect_method=self.config.collect_method.value,
|
||||
backbone=backbone,
|
||||
)
|
||||
elif self.config.format == QuantFormat.FP8:
|
||||
quant_config = FP8_DEFAULT_CONFIG
|
||||
base_cfg = FP8_DEFAULT_CONFIG
|
||||
elif self.config.format == QuantFormat.FP4:
|
||||
if self.model_config.model_type.value.startswith("flux"):
|
||||
quant_config = NVFP4_FP8_MHA_CONFIG
|
||||
base_cfg = NVFP4_FP8_MHA_CONFIG
|
||||
else:
|
||||
quant_config = NVFP4_DEFAULT_CONFIG
|
||||
base_cfg = NVFP4_DEFAULT_CONFIG
|
||||
else:
|
||||
raise NotImplementedError(f"Unknown format {self.config.format}")
|
||||
|
||||
# Build a fresh config dict so we never mutate the global constants.
|
||||
quant_cfg_list = list(base_cfg["quant_cfg"])
|
||||
|
||||
if self.config.format == QuantFormat.FP4:
|
||||
for i, entry in enumerate(quant_cfg_list):
|
||||
if isinstance(entry, dict) and "block_sizes" in entry.get("cfg", {}):
|
||||
new_block_sizes = {**entry["cfg"]["block_sizes"], -1: self.config.block_size}
|
||||
quant_cfg_list[i] = {
|
||||
**entry,
|
||||
"cfg": {**entry["cfg"], "block_sizes": new_block_sizes},
|
||||
}
|
||||
|
||||
if self.config.quantize_mha:
|
||||
quant_config["quant_cfg"].append(
|
||||
quant_cfg_list.append(
|
||||
{
|
||||
"quantizer_name": "*[qkv]_bmm_quantizer",
|
||||
"cfg": {"num_bits": (4, 3), "axis": None},
|
||||
}
|
||||
)
|
||||
|
||||
quant_config = {**base_cfg, "quant_cfg": quant_cfg_list}
|
||||
set_quant_config_attr(
|
||||
quant_config,
|
||||
self.model_config.trt_high_precision_dtype.value,
|
||||
@@ -158,6 +173,7 @@ class Quantizer:
|
||||
backbone: torch.nn.Module,
|
||||
quant_config: Any,
|
||||
forward_loop: callable, # type: ignore[valid-type]
|
||||
backbone_name: str = "transformer",
|
||||
) -> torch.nn.Module:
|
||||
"""
|
||||
Apply quantization to the model.
|
||||
@@ -166,15 +182,18 @@ class Quantizer:
|
||||
backbone: Model backbone to quantize
|
||||
quant_config: Quantization configuration
|
||||
forward_loop: Forward pass function for calibration
|
||||
backbone_name: Name of the backbone being quantized
|
||||
"""
|
||||
self.logger.info("Checking for LoRA layers...")
|
||||
check_lora(backbone)
|
||||
|
||||
self.logger.info("Starting model quantization...")
|
||||
self.logger.info(f"Starting model quantization for {backbone_name}...")
|
||||
mtq.quantize(backbone, quant_config, forward_loop)
|
||||
# Get model-specific filter function
|
||||
model_filter_func = get_model_filter_func(self.model_config.model_type)
|
||||
self.logger.info(f"Using filter function for {self.model_config.model_type.value}")
|
||||
model_filter_func = get_model_filter_func(self.model_config.model_type, backbone_name)
|
||||
self.logger.info(
|
||||
f"Using filter function for {self.model_config.model_type.value}/{backbone_name}"
|
||||
)
|
||||
|
||||
self.logger.info("Disabling specific quantizers...")
|
||||
mtq.disable_quantizer(backbone, model_filter_func)
|
||||
@@ -221,20 +240,27 @@ class ExportManager:
|
||||
return True
|
||||
return False
|
||||
|
||||
def save_checkpoint(self, backbone: torch.nn.Module) -> None:
|
||||
def save_checkpoint(
|
||||
self,
|
||||
backbone: torch.nn.Module,
|
||||
backbone_name: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Save quantized model checkpoint.
|
||||
|
||||
Args:
|
||||
backbone: The quantized backbone module to save (must be the same instance
|
||||
that was passed to mtq.quantize, as it carries the _modelopt_state).
|
||||
backbone_name: Optional name for the backbone file (defaults to "backbone").
|
||||
"""
|
||||
if not self.config.quantized_torch_ckpt_path:
|
||||
return
|
||||
|
||||
ckpt_path = self.config.quantized_torch_ckpt_path
|
||||
ckpt_path.mkdir(parents=True, exist_ok=True)
|
||||
target_path = ckpt_path / "backbone.pt"
|
||||
filename = f"{backbone_name}.pt" if backbone_name else "backbone.pt"
|
||||
target_path = ckpt_path / filename
|
||||
|
||||
self.logger.info(f"Saving backbone to {target_path}")
|
||||
mto.save(backbone, str(target_path))
|
||||
|
||||
@@ -292,14 +318,19 @@ class ExportManager:
|
||||
if self.pipeline_manager is None:
|
||||
raise RuntimeError("Pipeline manager is required for per-backbone checkpoints.")
|
||||
|
||||
backbone = self.pipeline_manager.get_backbone()
|
||||
if restore_path.exists() and restore_path.is_dir():
|
||||
source_path = restore_path / "backbone.pt"
|
||||
if not restore_path.exists() or not restore_path.is_dir():
|
||||
raise FileNotFoundError(f"Checkpoint directory not found: {restore_path}")
|
||||
|
||||
for backbone_name, backbone in self.pipeline_manager.iter_backbones():
|
||||
source_path = restore_path / f"{backbone_name}.pt"
|
||||
if not source_path.exists():
|
||||
raise FileNotFoundError(f"Backbone checkpoint not found: {source_path}")
|
||||
self.logger.info(f"Restoring backbone from {source_path}")
|
||||
raise FileNotFoundError(
|
||||
f"Checkpoint not found for '{backbone_name}' in {restore_path}"
|
||||
)
|
||||
self.logger.info(f"Restoring {backbone_name} from {source_path}")
|
||||
mto.restore(backbone, str(source_path))
|
||||
self.logger.info("Backbone checkpoints restored successfully")
|
||||
|
||||
self.logger.info("Checkpoints restored successfully")
|
||||
|
||||
# TODO: should not do the any data type
|
||||
def export_hf_ckpt(self, pipe: Any, model_config: ModelConfig | None = None) -> None:
|
||||
@@ -368,9 +399,9 @@ def create_argument_parser() -> argparse.ArgumentParser:
|
||||
nargs="+",
|
||||
default=None,
|
||||
help=(
|
||||
"Model backbone(s) in the DiffusionPipeline to work on. "
|
||||
"Provide one name or multiple names separated by space or comma. "
|
||||
"If not provided use default based on model type."
|
||||
"Model backbone(s) in the DiffusionPipeline to quantize. "
|
||||
"Provide one or more names (e.g., 'transformer', 'video_decoder'). "
|
||||
"If not provided, uses default based on model type."
|
||||
),
|
||||
)
|
||||
model_group.add_argument(
|
||||
@@ -448,6 +479,12 @@ def create_argument_parser() -> argparse.ArgumentParser:
|
||||
action="store_true",
|
||||
help="Compress quantized weights to reduce memory footprint (FP8/FP4 only)",
|
||||
)
|
||||
quant_group.add_argument(
|
||||
"--block-size",
|
||||
type=int,
|
||||
default=16,
|
||||
help="Block size for NVFP4 quantization (default: 16)",
|
||||
)
|
||||
|
||||
calib_group = parser.add_argument_group("Calibration Configuration")
|
||||
calib_group.add_argument("--batch-size", type=int, default=2, help="Batch size for calibration")
|
||||
@@ -535,6 +572,7 @@ def main() -> None:
|
||||
lowrank=args.lowrank,
|
||||
quantize_mha=args.quantize_mha,
|
||||
compress=args.compress,
|
||||
block_size=args.block_size,
|
||||
)
|
||||
|
||||
if args.prompts_file is not None:
|
||||
@@ -571,7 +609,6 @@ def main() -> None:
|
||||
pipe = pipeline_manager.create_pipeline()
|
||||
pipeline_manager.setup_device()
|
||||
|
||||
backbone = pipeline_manager.get_backbone()
|
||||
export_manager = ExportManager(export_config, logger, pipeline_manager)
|
||||
|
||||
if export_config.restore_from and export_config.restore_from.exists():
|
||||
@@ -581,37 +618,48 @@ def main() -> None:
|
||||
logger.info("Initializing calibration...")
|
||||
calibrator = Calibrator(pipeline_manager, calib_config, model_config.model_type, logger)
|
||||
batched_prompts = calibrator.load_and_batch_prompts()
|
||||
|
||||
quantizer = Quantizer(quant_config, model_config, logger)
|
||||
backbone_quant_config = quantizer.get_quant_config(calib_config.n_steps, backbone)
|
||||
|
||||
# Pipe loads the ckpt just before the inference.
|
||||
def forward_loop(mod):
|
||||
calibrator.run_calibration(batched_prompts)
|
||||
for backbone_name, backbone in pipeline_manager.iter_backbones():
|
||||
logger.info(f"Quantizing backbone: {backbone_name}")
|
||||
backbone_quant_config = quantizer.get_quant_config(calib_config.n_steps, backbone)
|
||||
|
||||
quantizer.quantize_model(backbone, backbone_quant_config, forward_loop)
|
||||
# Calibration runs the full pipeline (not just `mod`), so the
|
||||
# closure intentionally ignores the backbone argument.
|
||||
def forward_loop(mod):
|
||||
calibrator.run_calibration(batched_prompts)
|
||||
|
||||
# Compress model weights if requested (only for FP8/FP4)
|
||||
if quant_config.compress:
|
||||
logger.info("Compressing model weights to reduce memory footprint...")
|
||||
mtq.compress(backbone)
|
||||
logger.info("Model compression completed")
|
||||
quantizer.quantize_model(
|
||||
backbone,
|
||||
backbone_quant_config,
|
||||
forward_loop,
|
||||
backbone_name=backbone_name,
|
||||
)
|
||||
|
||||
export_manager.save_checkpoint(backbone)
|
||||
# Compress model weights if requested (only for FP8/FP4)
|
||||
if quant_config.compress:
|
||||
logger.info(f"Compressing {backbone_name} weights...")
|
||||
mtq.compress(backbone)
|
||||
logger.info(f"{backbone_name} compression completed")
|
||||
|
||||
# TODO (Jingyu): To update this function, as we are focusing more on the torch deployment side.
|
||||
check_conv_and_mha(
|
||||
backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha
|
||||
)
|
||||
# For VAE backbones, skip check_conv_and_mha — the whole point
|
||||
# of VAE quantization is to quantize Conv layers.
|
||||
if backbone_name not in ("video_decoder", "vae"):
|
||||
check_conv_and_mha(
|
||||
backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha
|
||||
)
|
||||
|
||||
export_manager.save_checkpoint(backbone, backbone_name)
|
||||
|
||||
pipeline_manager.print_quant_summary()
|
||||
|
||||
export_manager.export_onnx(
|
||||
pipe,
|
||||
backbone,
|
||||
model_config.model_type,
|
||||
quant_config.format,
|
||||
)
|
||||
for backbone_name, backbone in pipeline_manager.iter_backbones():
|
||||
export_manager.export_onnx(
|
||||
pipe,
|
||||
backbone,
|
||||
model_config.model_type,
|
||||
quant_config.format,
|
||||
)
|
||||
|
||||
export_manager.export_hf_ckpt(pipe, model_config)
|
||||
|
||||
|
||||
@@ -79,6 +79,7 @@ class QuantizationConfig:
|
||||
lowrank: int = 32 # SVDQuant lowrank
|
||||
quantize_mha: bool = False
|
||||
compress: bool = False
|
||||
block_size: int = 16 # NVFP4 block size
|
||||
|
||||
def validate(self) -> None:
|
||||
"""Validate configuration consistency."""
|
||||
@@ -120,7 +121,7 @@ class ModelConfig:
|
||||
|
||||
model_type: ModelType = ModelType.FLUX_DEV
|
||||
model_dtype: dict[str, torch.dtype] = field(default_factory=lambda: {"default": torch.float16})
|
||||
backbone: str = ""
|
||||
backbone: list[str] = field(default_factory=list)
|
||||
trt_high_precision_dtype: DataType = DataType.HALF
|
||||
override_model_path: Path | None = None
|
||||
cpu_offloading: bool = False
|
||||
|
||||
@@ -46,7 +46,7 @@ def filter_func_default(name: str) -> bool:
|
||||
|
||||
def check_conv_and_mha(backbone, if_fp4, quantize_mha):
|
||||
for name, module in backbone.named_modules():
|
||||
if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)) and if_fp4:
|
||||
if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d)) and if_fp4:
|
||||
module.weight_quantizer.disable()
|
||||
module.input_quantizer.disable()
|
||||
|
||||
@@ -87,6 +87,22 @@ def filter_func_flux_dev(name: str) -> bool:
|
||||
return pattern.match(name) is not None
|
||||
|
||||
|
||||
def filter_func_ltx2_vae(name: str) -> bool:
|
||||
"""Filter for LTX-2 VAE: keeps only conv1/conv2 in up_blocks resnets."""
|
||||
keep = re.compile(r".*up_blocks\.\d+\.resnets\.\d+\.conv[12](?:\.|$)")
|
||||
return not keep.match(name)
|
||||
|
||||
|
||||
def filter_func_wan_vae(name: str) -> bool:
|
||||
"""Filter for Wan 2.2 VAE: keeps only conv1/conv2 in resnet blocks."""
|
||||
keep = re.compile(
|
||||
r".*(down_blocks\.\d+\.(?:resnets\.\d+\.)?conv[12]"
|
||||
r"|mid_block\.resnets\.\d+\.conv[12]"
|
||||
r"|up_blocks\.\d+\.resnets\.\d+\.conv[12])(?:\.|$)"
|
||||
)
|
||||
return not keep.match(name)
|
||||
|
||||
|
||||
def filter_func_wan_video(name: str) -> bool:
|
||||
"""Filter function specifically for WAN-Video models."""
|
||||
pattern = re.compile(
|
||||
|
||||
@@ -1,105 +0,0 @@
|
||||
# Conv3D Implicit GEMM (Experimental)
|
||||
|
||||
Experimental Conv3D kernel prototype using implicit GEMM, with optional fused FP4 fake quantization for activations.
|
||||
|
||||
This code is kept under `experimental/` by design and is **not** part of the stable `modelopt.torch.quantization` API.
|
||||
|
||||
## Model Support
|
||||
|
||||
| Model/Framework | Supported | Notes |
|
||||
|-----------------|-----------|-------|
|
||||
| Video diffusion VAE Conv3D layers | Tested | Validated on VAE encoder/decoder Conv3D layers in video diffusion models |
|
||||
| Generic LLM backbones | No | Conv3D path is not relevant |
|
||||
| End-to-end ModelOpt PTQ/QAT pipeline | No | Not wired into formal quantization/export/compress flows |
|
||||
|
||||
## Deployment
|
||||
|
||||
| Framework | Supported | Notes |
|
||||
|-----------|-----------|-------|
|
||||
| TensorRT-LLM | No | No formal export integration for this kernel path |
|
||||
| vLLM | No | No integration |
|
||||
| SGLang | No | No integration |
|
||||
| PyTorch runtime (CUDA) | Yes (experimental) | JIT-compiles CUDA extension on first use |
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
from experimental.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda
|
||||
from modelopt.torch.quantization.tensor_quant import dynamic_block_quantize_op
|
||||
|
||||
x = torch.randn(1, 128, 21, 60, 106, device="cuda")
|
||||
w = torch.randn(512, 128, 3, 3, 3, device="cuda")
|
||||
block_size = 128
|
||||
|
||||
# Without FP4 activation quantization (drop-in-style Conv3D call)
|
||||
out = conv3d_implicit_gemm_cuda(x, w, stride=(1, 1, 1), padding=(1, 1, 1))
|
||||
|
||||
# Optional FP4 block quantization of weights along the GEMM K dimension.
|
||||
# The kernel's A-tile (activations) is quantized along K = Cin*kD*kH*kW,
|
||||
# so weights must be flattened to [Cout, K] before quantizing to match.
|
||||
Cout, Cin = w.shape[:2]
|
||||
K = Cin * w.shape[2] * w.shape[3] * w.shape[4]
|
||||
w_flat = w.reshape(Cout, K)
|
||||
w_q_flat = dynamic_block_quantize_op(
|
||||
w_flat,
|
||||
block_size,
|
||||
w_flat.abs().max().unsqueeze(0),
|
||||
4, # num_bits
|
||||
2, # exponent_bits
|
||||
8, # scale_num_bits
|
||||
4, # scale_exponent_bits
|
||||
)
|
||||
w_q = w_q_flat.reshape_as(w)
|
||||
|
||||
# With FP4 activation fake quantization
|
||||
out_q = conv3d_implicit_gemm_cuda(
|
||||
x,
|
||||
w_q,
|
||||
stride=(1, 1, 1),
|
||||
padding=(1, 1, 1),
|
||||
act_amax=x.abs().max().unsqueeze(0),
|
||||
quant_act=True,
|
||||
fp4_block_size=block_size, # 16, 32, 64, 128, or 256
|
||||
)
|
||||
```
|
||||
|
||||
## API
|
||||
|
||||
Function: `conv3d_implicit_gemm_cuda(...)` from `experimental/conv/implicit_gemm_cuda.py`
|
||||
|
||||
| Parameter | Description |
|
||||
|-----------|-------------|
|
||||
| `x` | Input tensor `[N, Cin, D, H, W]` |
|
||||
| `w` | Weight tensor `[Cout, Cin, kD, kH, kW]` |
|
||||
| `bias` | Optional bias `[Cout]` |
|
||||
| `stride` | Convolution stride `(D, H, W)` |
|
||||
| `padding` | Convolution padding `(D, H, W)` |
|
||||
| `dilation` | Convolution dilation `(D, H, W)` |
|
||||
| `act_amax` | Activation abs-max scalar tensor (required when `quant_act=True`) |
|
||||
| `quant_act` | Enable FP4 fake quantization on activations |
|
||||
| `fp4_block_size` | FP4 quantization block size (`16`, `32`, `64`, `128`, or `256`) |
|
||||
|
||||
## Status
|
||||
|
||||
Current state: **Prototype**
|
||||
|
||||
Known limitations:
|
||||
|
||||
- API is unstable and may change without notice.
|
||||
- Not registered in core quantization module registries.
|
||||
- Not covered by formal export/compress integration.
|
||||
- CUDA extension compile latency on first invocation.
|
||||
- Validation and performance coverage are limited to local experiments.
|
||||
|
||||
## Notes
|
||||
|
||||
- The CUDA kernel is JIT-compiled on first call (can take several seconds).
|
||||
- Output shape matches `torch.nn.functional.conv3d`.
|
||||
- FP4 path applies quantize-dequantize in-kernel for activation tiles.
|
||||
|
||||
## References
|
||||
|
||||
- Implicit GEMM-based convolution design patterns in GPU kernels.
|
||||
- ModelOpt FP4-related quantization utilities in `modelopt.torch.quantization.tensor_quant`.
|
||||
@@ -1,208 +0,0 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Latency benchmark: implicit GEMM (quant / non-quant) vs cuDNN conv3d.
|
||||
|
||||
Usage:
|
||||
python -m experimental.conv.bench_implicit_gemm
|
||||
python -m experimental.conv.bench_implicit_gemm --shapes wan22
|
||||
python -m experimental.conv.bench_implicit_gemm --shapes all --warmup 20 --iters 100
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark shapes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# (name, N, Cin, D, H, W, Cout, kD, kH, kW, stride, padding, dilation)
|
||||
SHAPES = {
|
||||
"small": [
|
||||
("small_16x32_3x3x3", 1, 16, 8, 8, 8, 32, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)),
|
||||
],
|
||||
"medium": [
|
||||
("med_64x128_3x3x3", 1, 64, 16, 32, 32, 128, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)),
|
||||
("med_128x256_3x3x3", 1, 128, 8, 16, 16, 256, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)),
|
||||
("med_128x128_1x3x3", 1, 128, 16, 32, 32, 128, 1, 3, 3, (1, 1, 1), (0, 1, 1), (1, 1, 1)),
|
||||
],
|
||||
"wan22": [
|
||||
("wan22_128x512", 1, 128, 21, 60, 106, 512, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)),
|
||||
("wan22_512x512", 1, 512, 21, 60, 106, 512, 1, 1, 1, (1, 1, 1), (0, 0, 0), (1, 1, 1)),
|
||||
("wan22_512x128", 1, 512, 21, 60, 106, 128, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)),
|
||||
],
|
||||
"stride": [
|
||||
("stride2_64x128", 1, 64, 16, 32, 32, 128, 3, 3, 3, (2, 2, 2), (1, 1, 1), (1, 1, 1)),
|
||||
("stride2_128x256", 1, 128, 16, 32, 32, 256, 3, 3, 3, (2, 2, 2), (1, 1, 1), (1, 1, 1)),
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def get_shapes(name: str):
|
||||
"""Return list of benchmark shapes by name or all shapes."""
|
||||
if name == "all":
|
||||
result = []
|
||||
for v in SHAPES.values():
|
||||
result.extend(v)
|
||||
return result
|
||||
return SHAPES[name]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Timing utility
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def bench_fn(fn, warmup: int, iters: int) -> float:
|
||||
"""Benchmark a callable, return median time in ms."""
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
times = []
|
||||
for _ in range(iters):
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
fn()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
times.append(start.elapsed_time(end))
|
||||
|
||||
times.sort()
|
||||
return times[len(times) // 2] # median
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_benchmark(shapes_name: str, warmup: int, iters: int, fp4_block_size: int):
|
||||
"""Run latency benchmark for the given shapes."""
|
||||
from experimental.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda
|
||||
|
||||
shapes = get_shapes(shapes_name)
|
||||
|
||||
# Header
|
||||
print(f"\n{'=' * 100}")
|
||||
print(
|
||||
f"Conv3D Latency Benchmark | warmup={warmup} iters={iters} fp4_block_size={fp4_block_size}"
|
||||
)
|
||||
print(f"GPU: {torch.cuda.get_device_name()}")
|
||||
print(f"{'=' * 100}")
|
||||
print(
|
||||
f"{'Shape':<25} {'M':>10} {'K':>8} {'N':>6} "
|
||||
f"{'cuDNN':>9} {'GEMM':>9} {'GEMM+FP4':>9} "
|
||||
f"{'GEMM/cuDNN':>11} {'FP4/cuDNN':>10}"
|
||||
)
|
||||
print("-" * 100)
|
||||
|
||||
for name, n, cin, d, h, w, cout, kd, kh, kw, stride, padding, dilation in shapes:
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(n, cin, d, h, w, device="cuda", dtype=torch.float32)
|
||||
weight = torch.randn(cout, cin, kd, kh, kw, device="cuda", dtype=torch.float32)
|
||||
act_amax = x.abs().max().unsqueeze(0)
|
||||
|
||||
# Compute GEMM dimensions for display
|
||||
sd, sh, sw = stride
|
||||
dd, dh, dw = dilation
|
||||
pd, ph, pw = padding
|
||||
od = (d + 2 * pd - dd * (kd - 1) - 1) // sd + 1
|
||||
oh = (h + 2 * ph - dh * (kh - 1) - 1) // sh + 1
|
||||
ow = (w + 2 * pw - dw * (kw - 1) - 1) // sw + 1
|
||||
gemm_m = n * od * oh * ow
|
||||
gemm_k = cin * kd * kh * kw
|
||||
gemm_n = cout
|
||||
|
||||
# cuDNN (torch.nn.functional.conv3d)
|
||||
t_cudnn = bench_fn(
|
||||
lambda: F.conv3d(x, weight, stride=stride, padding=padding, dilation=dilation),
|
||||
warmup,
|
||||
iters,
|
||||
)
|
||||
|
||||
# Implicit GEMM (non-quantized)
|
||||
t_gemm = bench_fn(
|
||||
lambda: conv3d_implicit_gemm_cuda(
|
||||
x,
|
||||
weight,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
quant_act=False,
|
||||
fp4_block_size=fp4_block_size,
|
||||
),
|
||||
warmup,
|
||||
iters,
|
||||
)
|
||||
|
||||
# Implicit GEMM (FP4 quantized)
|
||||
t_fp4 = bench_fn(
|
||||
lambda: conv3d_implicit_gemm_cuda(
|
||||
x,
|
||||
weight,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
act_amax=act_amax,
|
||||
quant_act=True,
|
||||
fp4_block_size=fp4_block_size,
|
||||
),
|
||||
warmup,
|
||||
iters,
|
||||
)
|
||||
|
||||
ratio_gemm = t_gemm / t_cudnn
|
||||
ratio_fp4 = t_fp4 / t_cudnn
|
||||
|
||||
print(
|
||||
f"{name:<25} {gemm_m:>10,} {gemm_k:>8,} {gemm_n:>6,} "
|
||||
f"{t_cudnn:>8.3f}ms {t_gemm:>8.3f}ms {t_fp4:>8.3f}ms "
|
||||
f"{ratio_gemm:>10.2f}x {ratio_fp4:>9.2f}x"
|
||||
)
|
||||
|
||||
print(f"{'=' * 100}")
|
||||
print("Ratios > 1.0x mean slower than cuDNN; < 1.0x mean faster.")
|
||||
print()
|
||||
|
||||
|
||||
def main():
|
||||
"""Entry point for the benchmark CLI."""
|
||||
parser = argparse.ArgumentParser(description="Conv3D latency benchmark")
|
||||
parser.add_argument(
|
||||
"--shapes",
|
||||
default="all",
|
||||
choices=[*list(SHAPES.keys()), "all"],
|
||||
help="Which shape set to benchmark (default: all)",
|
||||
)
|
||||
parser.add_argument("--warmup", type=int, default=20, help="Warmup iterations")
|
||||
parser.add_argument("--iters", type=int, default=100, help="Benchmark iterations")
|
||||
parser.add_argument(
|
||||
"--fp4-block-size",
|
||||
type=int,
|
||||
default=128,
|
||||
choices=[128, 256],
|
||||
help="FP4 block size (default: 128)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
run_benchmark(args.shapes, args.warmup, args.iters, args.fp4_block_size)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -15,8 +15,12 @@
|
||||
|
||||
"""Quantized convolution."""
|
||||
|
||||
import warnings
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda
|
||||
|
||||
from ... import tensor_quant
|
||||
from .quant_module import QuantLinearConvBase, QuantModuleRegistry, _LegacyQuantLinearConvBaseMixin
|
||||
|
||||
@@ -62,12 +66,85 @@ class QuantConv2d(_LegacyQuantLinearConvBaseMixin, nn.Conv2d):
|
||||
default_quant_desc_weight = _QuantConv2d.default_quant_desc_weight
|
||||
|
||||
|
||||
def _is_nvfp4_quantizer(quantizer) -> bool:
|
||||
"""Check if a TensorQuantizer is configured for NVFP4 dynamic block quantization."""
|
||||
return (
|
||||
quantizer.num_bits == (2, 1)
|
||||
and quantizer.block_sizes is not None
|
||||
and quantizer.block_sizes.get("scale_bits") == (4, 3)
|
||||
and quantizer.block_sizes.get("type") == "dynamic"
|
||||
)
|
||||
|
||||
|
||||
def _nvfp4_quantize_weight_along_k(weight, weight_quantizer):
|
||||
"""Apply NVFP4 fake quantization to Conv3D weight along the GEMM K dimension."""
|
||||
w_flat = weight.reshape(weight.shape[0], -1)
|
||||
return weight_quantizer(w_flat).reshape_as(weight)
|
||||
|
||||
|
||||
@QuantModuleRegistry.register({nn.Conv3d: "nn.Conv3d"})
|
||||
class _QuantConv3d(QuantLinearConvBase):
|
||||
"""Quantized 3D convolution."""
|
||||
"""Quantized 3D convolution.
|
||||
|
||||
For NVFP4, uses a fused implicit GEMM kernel with activation FP4 quantization
|
||||
inside the kernel. For all other configs, the default cuDNN path is used.
|
||||
"""
|
||||
|
||||
default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL
|
||||
|
||||
def _should_use_implicit_gemm(self):
|
||||
"""Check if both quantizers are NVFP4 and the implicit GEMM kernel is available."""
|
||||
return (
|
||||
hasattr(self, "input_quantizer")
|
||||
and hasattr(self, "weight_quantizer")
|
||||
and _is_nvfp4_quantizer(self.input_quantizer)
|
||||
and _is_nvfp4_quantizer(self.weight_quantizer)
|
||||
and self.groups == 1
|
||||
)
|
||||
|
||||
def _implicit_gemm_forward(self, input):
|
||||
"""Run NVFP4 implicit GEMM kernel. Input may already be padded."""
|
||||
# _get_amax is an internal TensorQuantizer method with no public equivalent;
|
||||
# block_sizes is a public property.
|
||||
act_amax = self.input_quantizer._get_amax(input)
|
||||
weight = _nvfp4_quantize_weight_along_k(self.weight, self.weight_quantizer)
|
||||
fp4_block_size = self.input_quantizer.block_sizes.get(-1, 16)
|
||||
|
||||
output = conv3d_implicit_gemm_cuda(
|
||||
input,
|
||||
weight,
|
||||
bias=self.bias,
|
||||
stride=self.stride,
|
||||
padding=self.padding,
|
||||
dilation=self.dilation,
|
||||
act_amax=act_amax,
|
||||
quant_act=self.input_quantizer.is_enabled,
|
||||
fp4_block_size=fp4_block_size,
|
||||
)
|
||||
return self.output_quantizer(output)
|
||||
|
||||
def forward(self, input, *args, **kwargs):
|
||||
"""Forward with implicit GEMM for NVFP4, default path otherwise."""
|
||||
if not self._should_use_implicit_gemm():
|
||||
return super().forward(input, *args, **kwargs)
|
||||
|
||||
if self.training:
|
||||
warnings.warn(
|
||||
"Implicit GEMM Conv3D kernel is inference-only and does not support training. "
|
||||
"Falling back to the default cuDNN quantization path, which could produce "
|
||||
"different numerics.",
|
||||
stacklevel=2,
|
||||
)
|
||||
return super().forward(input, *args, **kwargs)
|
||||
|
||||
# During calibration, only collect amax — use the faster cuDNN path.
|
||||
# _if_calib/_if_quant are internal TensorQuantizer state with no public property;
|
||||
# toggled via enable_calib()/disable_calib()/enable_quant()/disable_quant().
|
||||
if self.input_quantizer._if_calib and not self.input_quantizer._if_quant:
|
||||
return super().forward(input, *args, **kwargs)
|
||||
|
||||
return self._implicit_gemm_forward(input)
|
||||
|
||||
|
||||
class QuantConv3d(_LegacyQuantLinearConvBaseMixin, nn.Conv3d):
|
||||
"""Quantized 3D convolution."""
|
||||
|
||||
@@ -63,6 +63,7 @@ from ...nn import (
|
||||
QuantModuleRegistry,
|
||||
TensorQuantizer,
|
||||
)
|
||||
from ...nn.modules.quant_conv import _QuantConv3d
|
||||
from ..custom import _QuantFunctionalMixin
|
||||
|
||||
onnx_dtype_map = {
|
||||
@@ -278,3 +279,37 @@ class FP8SDPA(Function):
|
||||
high_precision_flag,
|
||||
disable_fp8_mha,
|
||||
)
|
||||
|
||||
|
||||
# WanCausalConv3d quantization support (diffusers VAE)
|
||||
try:
|
||||
from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d
|
||||
|
||||
@QuantModuleRegistry.register({WanCausalConv3d: "WanCausalConv3d"})
|
||||
class _QuantDiffusersWanCausalConv3d(_QuantConv3d):
|
||||
"""Quantized WanCausalConv3d — applies causal padding before quantized conv."""
|
||||
|
||||
def forward(self, x, cache_x=None):
|
||||
# Apply WanCausalConv3d-specific causal padding
|
||||
padding = list(self._padding)
|
||||
if cache_x is not None and self._padding[4] > 0:
|
||||
cache_x = cache_x.to(x.device)
|
||||
x = torch.cat([cache_x, x], dim=2)
|
||||
padding[4] -= cache_x.shape[2]
|
||||
x = F.pad(x, padding)
|
||||
|
||||
# NVFP4 implicit GEMM path (self.padding is (0,0,0) since padding already applied)
|
||||
if self._should_use_implicit_gemm():
|
||||
if not (self.input_quantizer._if_calib and not self.input_quantizer._if_quant):
|
||||
return self._implicit_gemm_forward(x)
|
||||
|
||||
# Default quantized conv path (skip WanCausalConv3d.forward to avoid double-padding)
|
||||
with self.quantize_weight():
|
||||
input = self.input_quantizer(x)
|
||||
output = torch.nn.Conv3d.forward(self, input)
|
||||
if isinstance(output, tuple):
|
||||
return (self.output_quantizer(output[0]), *output[1:])
|
||||
return self.output_quantizer(output)
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
# Conv3D Implicit GEMM
|
||||
|
||||
Conv3D kernel using implicit GEMM with BF16 WMMA tensor cores and optional fused FP4 (E2M1) fake quantization.
|
||||
|
||||
This kernel is integrated into `modelopt.torch.quantization` via `_QuantConv3d` — when NVFP4 quantization is applied to an `nn.Conv3d` layer through ModelOpt PTQ, the implicit GEMM path is used automatically. We have only tested it on VAE Conv3D layers from video generation models (e.g. Wan2.2).
|
||||
|
||||
## Requirements
|
||||
|
||||
- **GPU:** SM80+ (Ampere or newer) for BF16 WMMA tensor cores
|
||||
- **PyTorch:** CUDA toolkit with JIT C++ extension support (`torch.utils.cpp_extension`)
|
||||
- **Grouped convolution is not supported** (groups must be 1)
|
||||
|
||||
## Data Types
|
||||
|
||||
| Stage | Precision |
|
||||
|-------|-----------|
|
||||
| Input / output tensors | FP32, FP16, or BF16 (dtype is preserved) |
|
||||
| Internal compute | BF16 via WMMA m16n16k16 tensor cores |
|
||||
| Accumulation | FP32 |
|
||||
| FP4 activation quantization | E2M1 values, FP8 E4M3 scales |
|
||||
|
||||
## Integration with ModelOpt Quantization
|
||||
|
||||
When NVFP4 quantization is configured on a `Conv3d` layer via ModelOpt PTQ, the implicit GEMM kernel is used automatically during quantized inference. The integration is in `_QuantConv3d` (`modelopt/torch/quantization/nn/modules/quant_conv.py`):
|
||||
|
||||
- During **calibration**, the standard cuDNN path is used (faster).
|
||||
- During **quantized inference** with NVFP4 input and weight quantizers, the kernel fuses activation FP4 quantization inside the GEMM.
|
||||
- For all other quantization configs, the default cuDNN path is used as fallback.
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda
|
||||
from modelopt.torch.quantization.tensor_quant import dynamic_block_quantize_op
|
||||
|
||||
x = torch.randn(1, 128, 21, 60, 106, device="cuda")
|
||||
w = torch.randn(512, 128, 3, 3, 3, device="cuda")
|
||||
block_size = 128
|
||||
|
||||
# Without FP4 activation quantization (drop-in-style Conv3D call)
|
||||
out = conv3d_implicit_gemm_cuda(x, w, stride=(1, 1, 1), padding=(1, 1, 1))
|
||||
|
||||
# Optional FP4 block quantization of weights along the GEMM K dimension.
|
||||
# The kernel's A-tile (activations) is quantized along K = Cin*kD*kH*kW,
|
||||
# so weights must be flattened to [Cout, K] before quantizing to match.
|
||||
Cout, Cin = w.shape[:2]
|
||||
K = Cin * w.shape[2] * w.shape[3] * w.shape[4]
|
||||
w_flat = w.reshape(Cout, K)
|
||||
w_q_flat = dynamic_block_quantize_op(
|
||||
w_flat,
|
||||
block_size,
|
||||
w_flat.abs().max().unsqueeze(0),
|
||||
4, # num_bits
|
||||
2, # exponent_bits
|
||||
8, # scale_num_bits
|
||||
4, # scale_exponent_bits
|
||||
)
|
||||
w_q = w_q_flat.reshape_as(w)
|
||||
|
||||
# With FP4 activation fake quantization
|
||||
out_q = conv3d_implicit_gemm_cuda(
|
||||
x,
|
||||
w_q,
|
||||
stride=(1, 1, 1),
|
||||
padding=(1, 1, 1),
|
||||
act_amax=x.abs().max().unsqueeze(0),
|
||||
quant_act=True,
|
||||
fp4_block_size=block_size, # 16, 32, 64, 128, or 256
|
||||
)
|
||||
```
|
||||
|
||||
## API
|
||||
|
||||
### `conv3d_implicit_gemm_cuda`
|
||||
|
||||
`from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda`
|
||||
|
||||
| Parameter | Description |
|
||||
|-----------|-------------|
|
||||
| `x` | Input tensor `[N, Cin, D, H, W]` |
|
||||
| `w` | Weight tensor `[Cout, Cin, kD, kH, kW]` |
|
||||
| `bias` | Optional bias `[Cout]` |
|
||||
| `stride` | Convolution stride `(D, H, W)` |
|
||||
| `padding` | Convolution padding `(D, H, W)` |
|
||||
| `dilation` | Convolution dilation `(D, H, W)` |
|
||||
| `act_amax` | Activation abs-max scalar tensor (required when `quant_act=True`) |
|
||||
| `quant_act` | Enable FP4 fake quantization on activations |
|
||||
| `fp4_block_size` | FP4 quantization block size (`16`, `32`, `64`, `128`, or `256`) |
|
||||
|
||||
### `fp4_fake_quant`
|
||||
|
||||
`from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import fp4_fake_quant`
|
||||
|
||||
Standalone FP4 (E2M1) blockwise fake quantization with FP8 E4M3 scale quantization. Uses the same CUDA device functions as the fused path inside the GEMM kernel.
|
||||
|
||||
| Parameter | Description |
|
||||
|-----------|-------------|
|
||||
| `x` | Input tensor (any shape; `numel` must be divisible by `block_size`) |
|
||||
| `global_amax` | Scalar tensor — global abs max for scale computation |
|
||||
| `block_size` | Number of elements per FP4 quantization block (default `16`) |
|
||||
|
||||
## Testing
|
||||
|
||||
```bash
|
||||
# Run tests (requires GPU)
|
||||
python -m pytest tests/gpu/torch/quantization/kernels/test_implicit_gemm.py -v
|
||||
```
|
||||
|
||||
## Status
|
||||
|
||||
Current state: **Integrated** (registered in `QuantModuleRegistry`, auto-dispatched for NVFP4 Conv3D)
|
||||
|
||||
Known limitations:
|
||||
|
||||
- CUDA extension compile latency on first invocation (~seconds).
|
||||
- Grouped convolution (`groups > 1`) is not supported. In the ModelOpt E2E flow, `_QuantConv3d` automatically falls back to the default cuDNN path for grouped convolutions.
|
||||
- BF16 rounding error accumulates with the K dimension — expect max abs diff scaling roughly as `sqrt(K)` compared to cuDNN FP32.
|
||||
- Inference only (`@torch.no_grad`) — not suitable for QAT backward pass.
|
||||
|
||||
## Notes
|
||||
|
||||
- The CUDA kernel is JIT-compiled on first call via `torch.utils.cpp_extension.load()`.
|
||||
- Output shape matches `torch.nn.functional.conv3d`.
|
||||
- FP4 path applies quantize-dequantize in-kernel for activation tiles (no extra global memory pass).
|
||||
- Tile config: BLOCK_M=64, BLOCK_N=64, BLOCK_K=256, 8 warps (256 threads), ~70 KB shared memory per block.
|
||||
- The kernel body is guarded by `#if __CUDA_ARCH__ >= 800` so it compiles as an empty stub when nvcc targets pre-Ampere archs (PyTorch's default `-gencode` list can include sm_75, which lacks BF16 WMMA fragments). Dispatch is enforced at runtime by `_get_cuda_module()` via `_MIN_SM_MAJOR = 8`.
|
||||
|
||||
## Files
|
||||
|
||||
| File | Role |
|
||||
|------|------|
|
||||
| `implicit_gemm_cuda.py` | Python API and JIT compilation |
|
||||
| `implicit_gemm_kernel.cu` | CUDA kernel (BF16 WMMA + FP4 quantization) |
|
||||
| `implicit_gemm_binding.cpp` | PyTorch C++ extension binding |
|
||||
|
||||
## References
|
||||
|
||||
- Implicit GEMM-based convolution design patterns in GPU kernels.
|
||||
- ModelOpt FP4-related quantization utilities in `modelopt.torch.quantization.tensor_quant`.
|
||||
+5
@@ -145,6 +145,10 @@ __global__ void __launch_bounds__(WARPS_M * WARPS_N * 32, 2)
|
||||
const float *__restrict__ act_amax, int Cin, int Dp, int Hp, int Wp,
|
||||
int Cout, int OD, int OH, int OW, int kD, int kH, int kW, int sd,
|
||||
int sh, int sw, int dd, int dh, int dw, int M, int K) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
// BF16 WMMA fragments are only available on sm_80+. On earlier archs this
|
||||
// kernel compiles as an empty stub; the runtime gate in _get_cuda_module()
|
||||
// prevents dispatching it on unsupported hardware.
|
||||
// Derived constants
|
||||
constexpr int NUM_WARPS = WARPS_M * WARPS_N;
|
||||
constexpr int NUM_THREADS = NUM_WARPS * 32;
|
||||
@@ -469,6 +473,7 @@ __global__ void __launch_bounds__(WARPS_M * WARPS_N * 32, 2)
|
||||
y[m_idx * Cout + n_idx] = result;
|
||||
}
|
||||
}
|
||||
#endif // __CUDA_ARCH__ >= 800
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
@@ -0,0 +1,31 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def tiny_wan22_path(tmp_path_factory):
|
||||
"""Create a tiny Wan 2.2 (14B-style) pipeline and return its path.
|
||||
|
||||
Built once per session and shared across all tests that need it.
|
||||
"""
|
||||
try:
|
||||
from _test_utils.torch.diffusers_models import create_tiny_wan22_pipeline_dir
|
||||
except ImportError:
|
||||
pytest.skip("Wan 2.2 diffusers models not available (requires diffusers with WanPipeline)")
|
||||
|
||||
tmp_path = tmp_path_factory.mktemp("wan22")
|
||||
return str(create_tiny_wan22_pipeline_dir(tmp_path))
|
||||
@@ -151,6 +151,119 @@ def test_diffusers_quantization(
|
||||
model.inference(tmp_path)
|
||||
|
||||
|
||||
class Wan22Model(NamedTuple):
|
||||
model: str
|
||||
backbone: str | None
|
||||
format_type: str
|
||||
quant_algo: str
|
||||
collect_method: str
|
||||
|
||||
def _ckpt_path(self, tmp_path: Path) -> str:
|
||||
stem = self.model.replace("wan2.2-t2v-", "")
|
||||
parts = [stem, *([self.backbone] if self.backbone else []), self.format_type]
|
||||
return str(tmp_path / f"wan22_{'_'.join(parts)}.pt")
|
||||
|
||||
def _common_args(self, tiny_wan22_path: str) -> list[str]:
|
||||
cmd_args = [
|
||||
"python",
|
||||
"quantize.py",
|
||||
"--model",
|
||||
self.model,
|
||||
"--override-model-path",
|
||||
tiny_wan22_path,
|
||||
"--format",
|
||||
self.format_type,
|
||||
"--quant-algo",
|
||||
self.quant_algo,
|
||||
"--collect-method",
|
||||
self.collect_method,
|
||||
"--model-dtype",
|
||||
"BFloat16",
|
||||
"--trt-high-precision-dtype",
|
||||
"BFloat16",
|
||||
"--calib-size",
|
||||
"2",
|
||||
"--batch-size",
|
||||
"1",
|
||||
"--n-steps",
|
||||
"2",
|
||||
# Tiny video dims — override MODEL_DEFAULTS for fast CI.
|
||||
"--extra-param",
|
||||
"height=16",
|
||||
"--extra-param",
|
||||
"width=16",
|
||||
"--extra-param",
|
||||
"num_frames=5",
|
||||
]
|
||||
if self.backbone is not None:
|
||||
cmd_args.extend(["--backbone", self.backbone])
|
||||
return cmd_args
|
||||
|
||||
def quantize(self, tiny_wan22_path: str, tmp_path: Path) -> None:
|
||||
run_example_command(
|
||||
[
|
||||
*self._common_args(tiny_wan22_path),
|
||||
"--quantized-torch-ckpt-save-path",
|
||||
self._ckpt_path(tmp_path),
|
||||
],
|
||||
"diffusers/quantization",
|
||||
)
|
||||
|
||||
def restore(self, tiny_wan22_path: str, tmp_path: Path) -> None:
|
||||
run_example_command(
|
||||
[*self._common_args(tiny_wan22_path), "--restore-from", self._ckpt_path(tmp_path)],
|
||||
"diffusers/quantization",
|
||||
)
|
||||
|
||||
|
||||
# The VAE (``AutoencoderKLWan``) is shared between Wan 2.2 14B and 5B, so the
|
||||
# Conv3D NVFP4 implicit-GEMM dispatch exercises the same kernel either way; we
|
||||
# parametrize both ``--model`` values to also cover the ``quantize.py`` dispatch
|
||||
# for each.
|
||||
@pytest.mark.parametrize(
|
||||
"wan_model",
|
||||
[
|
||||
Wan22Model("wan2.2-t2v-14b", None, "int8", "smoothquant", "min-mean"),
|
||||
pytest.param(
|
||||
Wan22Model("wan2.2-t2v-14b", None, "fp8", "max", "default"),
|
||||
marks=minimum_sm(89),
|
||||
),
|
||||
pytest.param(
|
||||
Wan22Model("wan2.2-t2v-14b", None, "fp4", "max", "default"),
|
||||
marks=minimum_sm(89),
|
||||
),
|
||||
pytest.param(
|
||||
Wan22Model("wan2.2-t2v-14b", "vae", "fp8", "max", "default"),
|
||||
marks=minimum_sm(89),
|
||||
),
|
||||
pytest.param(
|
||||
Wan22Model("wan2.2-t2v-14b", "vae", "fp4", "max", "default"),
|
||||
marks=minimum_sm(89),
|
||||
),
|
||||
pytest.param(
|
||||
Wan22Model("wan2.2-t2v-5b", "vae", "fp8", "max", "default"),
|
||||
marks=minimum_sm(89),
|
||||
),
|
||||
pytest.param(
|
||||
Wan22Model("wan2.2-t2v-5b", "vae", "fp4", "max", "default"),
|
||||
marks=minimum_sm(89),
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"wan22_14b_transformer_int8_smoothquant",
|
||||
"wan22_14b_transformer_fp8_max",
|
||||
"wan22_14b_transformer_fp4_max",
|
||||
"wan22_14b_vae_fp8_max",
|
||||
"wan22_14b_vae_fp4_max",
|
||||
"wan22_5b_vae_fp8_max",
|
||||
"wan22_5b_vae_fp4_max",
|
||||
],
|
||||
)
|
||||
def test_wan22_quantization(wan_model: Wan22Model, tiny_wan22_path: str, tmp_path: Path) -> None:
|
||||
wan_model.quantize(tiny_wan22_path, tmp_path)
|
||||
wan_model.restore(tiny_wan22_path, tmp_path)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_name", "model_path", "torch_compile"),
|
||||
[
|
||||
|
||||
@@ -128,3 +128,84 @@ def test_diffusers_hf_ckpt_export(model: DiffuserHfExportModel, tmp_path: Path)
|
||||
|
||||
weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin"))
|
||||
assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}"
|
||||
|
||||
|
||||
class Wan22HfExportModel(NamedTuple):
|
||||
model: str
|
||||
backbone: str | None
|
||||
format_type: str
|
||||
quant_algo: str
|
||||
collect_method: str
|
||||
|
||||
def _suffix(self) -> str:
|
||||
stem = self.model.replace("wan2.2-t2v-", "")
|
||||
parts = [stem, *([self.backbone] if self.backbone else []), self.format_type]
|
||||
return "_".join(parts)
|
||||
|
||||
def quantize_and_export_hf(self, tiny_wan22_path: str, tmp_path: Path) -> Path:
|
||||
hf_ckpt_dir = tmp_path / f"wan22_{self._suffix()}_hf_ckpt"
|
||||
cmd_args = [
|
||||
"python",
|
||||
"quantize.py",
|
||||
"--model",
|
||||
self.model,
|
||||
"--override-model-path",
|
||||
tiny_wan22_path,
|
||||
"--format",
|
||||
self.format_type,
|
||||
"--quant-algo",
|
||||
self.quant_algo,
|
||||
"--collect-method",
|
||||
self.collect_method,
|
||||
"--model-dtype",
|
||||
"BFloat16",
|
||||
"--trt-high-precision-dtype",
|
||||
"BFloat16",
|
||||
"--calib-size",
|
||||
"2",
|
||||
"--batch-size",
|
||||
"1",
|
||||
"--n-steps",
|
||||
"2",
|
||||
# Tiny video dims — override MODEL_DEFAULTS for fast CI.
|
||||
"--extra-param",
|
||||
"height=16",
|
||||
"--extra-param",
|
||||
"width=16",
|
||||
"--extra-param",
|
||||
"num_frames=5",
|
||||
"--hf-ckpt-dir",
|
||||
str(hf_ckpt_dir),
|
||||
]
|
||||
if self.backbone is not None:
|
||||
cmd_args.extend(["--backbone", self.backbone])
|
||||
run_example_command(cmd_args, "diffusers/quantization")
|
||||
return hf_ckpt_dir
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"wan_model",
|
||||
[
|
||||
Wan22HfExportModel("wan2.2-t2v-14b", None, "int8", "smoothquant", "min-mean"),
|
||||
pytest.param(
|
||||
Wan22HfExportModel("wan2.2-t2v-14b", None, "fp8", "max", "default"),
|
||||
marks=minimum_sm(89),
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"wan22_14b_transformer_int8_smoothquant",
|
||||
"wan22_14b_transformer_fp8_max",
|
||||
],
|
||||
)
|
||||
def test_wan22_hf_ckpt_export(
|
||||
wan_model: Wan22HfExportModel, tiny_wan22_path: str, tmp_path: Path
|
||||
) -> None:
|
||||
hf_ckpt_dir = wan_model.quantize_and_export_hf(tiny_wan22_path, tmp_path)
|
||||
|
||||
assert hf_ckpt_dir.exists(), f"HF checkpoint directory was not created: {hf_ckpt_dir}"
|
||||
|
||||
config_files = list(hf_ckpt_dir.rglob("config.json"))
|
||||
assert len(config_files) > 0, f"No config.json found in {hf_ckpt_dir}"
|
||||
|
||||
weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin"))
|
||||
assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}"
|
||||
|
||||
+505
-80
@@ -28,7 +28,7 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not
|
||||
@pytest.fixture(scope="module")
|
||||
def cuda_conv3d():
|
||||
"""Import and return the CUDA implicit GEMM conv3d function."""
|
||||
from experimental.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda
|
||||
from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda
|
||||
|
||||
return conv3d_implicit_gemm_cuda
|
||||
|
||||
@@ -303,9 +303,9 @@ class TestConv3dDeterminism:
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def cuda_fp4_quant():
|
||||
"""Import FP4 fake quant for reference comparisons."""
|
||||
from experimental.conv.implicit_gemm_cuda import fp4_fake_quant
|
||||
def cuda_fp4():
|
||||
"""Import and return the CUDA FP4 fake quant function."""
|
||||
from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import fp4_fake_quant
|
||||
|
||||
return fp4_fake_quant
|
||||
|
||||
@@ -361,7 +361,7 @@ class TestConv3dFP4QuantBlockSizes:
|
||||
assert torch.equal(out1, out2), f"Non-deterministic for fp4_block_size={fp4_block_size}"
|
||||
|
||||
@pytest.mark.parametrize("fp4_block_size", [16, 32, 64, 128, 256])
|
||||
def test_quant_vs_unfused_reference(self, cuda_conv3d, cuda_fp4_quant, fp4_block_size):
|
||||
def test_quant_vs_unfused_reference(self, cuda_conv3d, cuda_fp4, fp4_block_size):
|
||||
"""Compare fused kernel vs unfused: fp4(im2col) @ fp4(weight).
|
||||
|
||||
Uses a shape where K is a multiple of 256 so all K-tiles are full
|
||||
@@ -382,9 +382,9 @@ class TestConv3dFP4QuantBlockSizes:
|
||||
im2col = x.permute(0, 2, 3, 4, 1).reshape(-1, cin) # [M, K]
|
||||
|
||||
# 2. FP4 fake-quant both matrices along K with the same block_size
|
||||
im2col_q = cuda_fp4_quant(im2col, act_amax, fp4_block_size)
|
||||
im2col_q = cuda_fp4(im2col, act_amax, fp4_block_size)
|
||||
w_flat = w.reshape(cout, cin).transpose(0, 1).contiguous() # [K, Cout]
|
||||
w_flat_q = cuda_fp4_quant(w_flat, w_amax, fp4_block_size)
|
||||
w_flat_q = cuda_fp4(w_flat, w_amax, fp4_block_size)
|
||||
|
||||
# 3. Matmul (in BF16 to match kernel's WMMA path)
|
||||
ref_out = (im2col_q.bfloat16() @ w_flat_q.bfloat16()).float()
|
||||
@@ -421,81 +421,33 @@ class TestConv3dFP4QuantBlockSizes:
|
||||
"""Smaller FP4 block sizes should generally produce lower quantization error.
|
||||
|
||||
Finer-grained blocks capture local ranges better, reducing quant error vs cuDNN.
|
||||
Test monotonicity: error(16) <= error(32) <= ... <= error(256) (with some tolerance).
|
||||
Reports detailed accuracy metrics for each block size vs cuDNN baseline.
|
||||
Test monotonicity on a medium config: error(16) <= error(64) <= error(256) (with 1.2x slack).
|
||||
"""
|
||||
torch.manual_seed(42)
|
||||
|
||||
# Test multiple shapes to get a comprehensive picture
|
||||
configs = [
|
||||
("Small K=432", 1, 16, 8, 8, 8, 32, 3, 3, 3),
|
||||
("Medium K=1728", 1, 64, 8, 8, 8, 64, 3, 3, 3),
|
||||
("Large K=3456", 1, 128, 5, 8, 8, 256, 3, 3, 3),
|
||||
]
|
||||
# Medium K=1728: Cin=64, 3x3x3 kernel
|
||||
cin, cout = 64, 64
|
||||
x = torch.randn(1, cin, 8, 8, 8, device="cuda", dtype=torch.float32)
|
||||
w = torch.randn(cout, cin, 3, 3, 3, device="cuda", dtype=torch.float32)
|
||||
act_amax = x.abs().max().unsqueeze(0)
|
||||
ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), dilation=(1, 1, 1))
|
||||
|
||||
block_sizes = [16, 32, 64, 128, 256]
|
||||
all_errors = {}
|
||||
|
||||
for desc, n, cin, d, h, w_s, cout, kd, kh, kw in configs:
|
||||
x = torch.randn(n, cin, d, h, w_s, device="cuda", dtype=torch.float32)
|
||||
w = torch.randn(cout, cin, kd, kh, kw, device="cuda", dtype=torch.float32)
|
||||
act_amax = x.abs().max().unsqueeze(0)
|
||||
k_size = cin * kd * kh * kw
|
||||
|
||||
ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), dilation=(1, 1, 1))
|
||||
ref_abs_mean = ref.abs().mean().item()
|
||||
|
||||
# Also compute no-quant baseline (BF16 rounding only)
|
||||
out_nq = cuda_conv3d(
|
||||
errors = {}
|
||||
for bs in block_sizes:
|
||||
out = cuda_conv3d(
|
||||
x,
|
||||
w,
|
||||
stride=(1, 1, 1),
|
||||
padding=(1, 1, 1),
|
||||
dilation=(1, 1, 1),
|
||||
quant_act=False,
|
||||
act_amax=act_amax,
|
||||
quant_act=True,
|
||||
fp4_block_size=bs,
|
||||
)
|
||||
nq_diff = (out_nq - ref).abs()
|
||||
errors[bs] = (out - ref).abs().mean().item()
|
||||
|
||||
print(
|
||||
f"\n {desc} (K={k_size}), output range [{ref.min().item():.1f}, {ref.max().item():.1f}]"
|
||||
)
|
||||
print(
|
||||
f" {'Block Size':>10} | {'Max Diff':>10} | {'Mean Diff':>10} | {'RMSE':>10} | {'Rel Err%':>8}"
|
||||
)
|
||||
print(f" {'-' * 10}-+-{'-' * 10}-+-{'-' * 10}-+-{'-' * 10}-+-{'-' * 8}")
|
||||
print(
|
||||
f" {'no-quant':>10} | {nq_diff.max().item():>10.4f} | "
|
||||
f"{nq_diff.mean().item():>10.6f} | "
|
||||
f"{((out_nq - ref) ** 2).mean().sqrt().item():>10.4f} | "
|
||||
f"{nq_diff.mean().item() / ref_abs_mean * 100:>7.3f}%"
|
||||
)
|
||||
|
||||
errors = {}
|
||||
for bs in block_sizes:
|
||||
out = cuda_conv3d(
|
||||
x,
|
||||
w,
|
||||
stride=(1, 1, 1),
|
||||
padding=(1, 1, 1),
|
||||
dilation=(1, 1, 1),
|
||||
act_amax=act_amax,
|
||||
quant_act=True,
|
||||
fp4_block_size=bs,
|
||||
)
|
||||
diff = (out - ref).abs()
|
||||
max_d = diff.max().item()
|
||||
mean_d = diff.mean().item()
|
||||
rmse = ((out - ref) ** 2).mean().sqrt().item()
|
||||
rel_err = mean_d / ref_abs_mean * 100
|
||||
errors[bs] = mean_d
|
||||
print(
|
||||
f" {bs:>10} | {max_d:>10.4f} | {mean_d:>10.6f} | "
|
||||
f"{rmse:>10.4f} | {rel_err:>7.3f}%"
|
||||
)
|
||||
all_errors[desc] = errors
|
||||
|
||||
# Monotonicity check on the medium config
|
||||
errors = all_errors["Medium K=1728"]
|
||||
# Monotonicity: smaller blocks should have equal or lower error
|
||||
for smaller, larger in [(16, 64), (16, 256), (32, 256), (64, 256)]:
|
||||
assert errors[smaller] <= errors[larger] * 1.2, (
|
||||
f"Expected error({smaller})={errors[smaller]:.6f} <= "
|
||||
@@ -590,14 +542,6 @@ class TestConv3dFP4QuantBlockSizes:
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def cuda_fp4():
|
||||
"""Import and return the CUDA FP4 fake quant function."""
|
||||
from experimental.conv.implicit_gemm_cuda import fp4_fake_quant
|
||||
|
||||
return fp4_fake_quant
|
||||
|
||||
|
||||
def _py_fp4_fake_quant_ref(x_flat, global_amax, block_size):
|
||||
"""Pure Python reference for FP4 fake quant (no BF16 rounding).
|
||||
|
||||
@@ -909,8 +853,8 @@ requires_dynamic_block_quantize = pytest.mark.skipif(
|
||||
class TestFP4FakeQuantVsModelopt:
|
||||
"""Compare experimental CUDA FP4 fake quant against all modelopt FP4 implementations.
|
||||
|
||||
This ensures the standalone FP4 kernel in experimental/conv produces the same
|
||||
results as the official modelopt quantization paths:
|
||||
This ensures the standalone FP4 kernel produces the same results as the
|
||||
other modelopt quantization paths:
|
||||
1. Triton fp4_fake_quant_block (Hopper+ dynamic blockwise)
|
||||
2. cuda_ext_mx.fused_amax_convert (CUDA extension fallback)
|
||||
3. dynamic_block_quantize_op (high-level API that dispatches to either)
|
||||
@@ -1065,3 +1009,484 @@ class TestFP4FakeQuantVsModelopt:
|
||||
# BF16/FP16 input rounding may cause small diffs
|
||||
tol = 1e-2 if dtype != torch.float32 else 1e-5
|
||||
assert max_diff < tol, f"dtype={dtype}: experimental vs Triton max diff: {max_diff:.6e}"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Input Validation / Error Path Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestConv3dInputValidation:
|
||||
"""Verify error paths raise appropriate exceptions."""
|
||||
|
||||
def test_invalid_fp4_block_size(self, cuda_conv3d):
|
||||
"""fp4_block_size not in {16, 32, 64, 128, 256} should raise ValueError."""
|
||||
x = torch.randn(1, 4, 4, 4, 4, device="cuda")
|
||||
w = torch.randn(8, 4, 3, 3, 3, device="cuda")
|
||||
with pytest.raises(ValueError, match="fp4_block_size"):
|
||||
cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), fp4_block_size=7)
|
||||
|
||||
def test_non_5d_input(self, cuda_conv3d):
|
||||
"""Non-5D tensors should raise ValueError."""
|
||||
x = torch.randn(1, 4, 4, 4, device="cuda") # 4D
|
||||
w = torch.randn(8, 4, 3, 3, 3, device="cuda")
|
||||
with pytest.raises(ValueError, match="5D"):
|
||||
cuda_conv3d(x, w)
|
||||
|
||||
def test_non_5d_weight(self, cuda_conv3d):
|
||||
"""Non-5D weight should raise ValueError."""
|
||||
x = torch.randn(1, 4, 4, 4, 4, device="cuda")
|
||||
w = torch.randn(8, 4, 3, 3, device="cuda") # 4D
|
||||
with pytest.raises(ValueError, match="5D"):
|
||||
cuda_conv3d(x, w)
|
||||
|
||||
def test_grouped_conv_error(self, cuda_conv3d):
|
||||
"""Mismatched Cin (groups > 1) should raise ValueError."""
|
||||
x = torch.randn(1, 8, 4, 4, 4, device="cuda")
|
||||
w = torch.randn(8, 4, 3, 3, 3, device="cuda") # Cin=4 != x.Cin=8
|
||||
with pytest.raises(ValueError, match="Grouped convolution"):
|
||||
cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1))
|
||||
|
||||
def test_quant_act_without_amax(self, cuda_conv3d):
|
||||
"""quant_act=True without act_amax should raise ValueError."""
|
||||
x = torch.randn(1, 4, 4, 4, 4, device="cuda")
|
||||
w = torch.randn(8, 4, 3, 3, 3, device="cuda")
|
||||
with pytest.raises(ValueError, match="act_amax"):
|
||||
cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=True, act_amax=None)
|
||||
|
||||
def test_fp4_numel_not_divisible(self, cuda_fp4):
|
||||
"""fp4_fake_quant should error when numel is not divisible by block_size."""
|
||||
inp = torch.randn(17, device="cuda")
|
||||
amax = torch.tensor([1.0], device="cuda")
|
||||
with pytest.raises(AssertionError, match="divisible"):
|
||||
cuda_fp4(inp, amax, block_size=16)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Input Dtype Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestConv3dInputDtypes:
|
||||
"""Verify conv3d works with non-float32 inputs and preserves output dtype."""
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_dtype_preservation(self, cuda_conv3d, dtype):
|
||||
"""Output dtype should match input dtype."""
|
||||
x = torch.randn(1, 16, 8, 8, 8, device="cuda", dtype=dtype)
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda", dtype=dtype)
|
||||
out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
assert out.dtype == dtype, f"Expected {dtype}, got {out.dtype}"
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_dtype_correctness(self, cuda_conv3d, dtype):
|
||||
"""Non-float32 inputs should produce correct results (vs F.conv3d in float32)."""
|
||||
torch.manual_seed(42)
|
||||
x_fp32 = torch.randn(1, 16, 8, 8, 8, device="cuda")
|
||||
w_fp32 = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
x = x_fp32.to(dtype)
|
||||
w = w_fp32.to(dtype)
|
||||
|
||||
out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
ref = F.conv3d(x_fp32, w_fp32, stride=(1, 1, 1), padding=(1, 1, 1))
|
||||
|
||||
# Both BF16 input rounding and internal BF16 WMMA contribute to error
|
||||
max_diff = (out.float() - ref).abs().max().item()
|
||||
k_size = 16 * 27
|
||||
scaled_atol = ATOL * (k_size / 1000.0) ** 0.5 * 2 # extra slack for input rounding
|
||||
assert max_diff < scaled_atol, f"dtype={dtype}: max diff {max_diff:.4f} > {scaled_atol:.4f}"
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_dtype_quant_path(self, cuda_conv3d, dtype):
|
||||
"""FP4 quantized path should also work with non-float32 inputs."""
|
||||
x = torch.randn(1, 16, 8, 8, 8, device="cuda", dtype=dtype)
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda", dtype=dtype)
|
||||
act_amax = x.float().abs().max().unsqueeze(0)
|
||||
|
||||
out = cuda_conv3d(
|
||||
x,
|
||||
w,
|
||||
stride=(1, 1, 1),
|
||||
padding=(1, 1, 1),
|
||||
act_amax=act_amax,
|
||||
quant_act=True,
|
||||
fp4_block_size=16,
|
||||
)
|
||||
assert out.dtype == dtype
|
||||
assert not torch.isnan(out).any()
|
||||
assert out.abs().max() > 0
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Non-Contiguous Input Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestConv3dNonContiguous:
|
||||
"""Verify kernel handles non-contiguous tensors (via internal .contiguous() calls)."""
|
||||
|
||||
def test_non_contiguous_input(self, cuda_conv3d):
|
||||
"""Permuted (non-contiguous) input should produce correct results."""
|
||||
torch.manual_seed(42)
|
||||
# Create non-contiguous tensor via permute + permute back
|
||||
x_base = torch.randn(1, 8, 8, 8, 16, device="cuda")
|
||||
x = x_base.permute(0, 4, 1, 2, 3) # [1, 16, 8, 8, 8] but non-contiguous
|
||||
assert not x.is_contiguous()
|
||||
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
x_contig = x.contiguous()
|
||||
|
||||
out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
ref = cuda_conv3d(x_contig, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
assert torch.equal(out, ref), "Non-contiguous input produced different results"
|
||||
|
||||
def test_non_contiguous_weight(self, cuda_conv3d):
|
||||
"""Transposed (non-contiguous) weight should produce correct results."""
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(1, 16, 8, 8, 8, device="cuda")
|
||||
# Create non-contiguous weight
|
||||
w_base = torch.randn(16, 32, 3, 3, 3, device="cuda")
|
||||
w = w_base.transpose(0, 1) # [32, 16, ...] but non-contiguous
|
||||
assert not w.is_contiguous()
|
||||
|
||||
w_contig = w.contiguous()
|
||||
out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
ref = cuda_conv3d(x, w_contig, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
assert torch.equal(out, ref), "Non-contiguous weight produced different results"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Combined Conv Parameter Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestConv3dCombinedParams:
|
||||
"""Test combinations of stride + dilation + padding that were never combined."""
|
||||
|
||||
def test_stride_and_dilation(self, cuda_conv3d):
|
||||
"""Stride > 1 and dilation > 1 simultaneously."""
|
||||
x = torch.randn(1, 16, 16, 16, 16, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
_run_conv3d_test(cuda_conv3d, x, w, None, (2, 2, 2), (1, 1, 1), (2, 2, 2))
|
||||
|
||||
def test_asymmetric_stride_and_padding(self, cuda_conv3d):
|
||||
"""Asymmetric stride with asymmetric padding."""
|
||||
x = torch.randn(1, 16, 16, 16, 16, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
_run_conv3d_test(cuda_conv3d, x, w, None, (1, 2, 2), (0, 1, 2), (1, 1, 1))
|
||||
|
||||
def test_all_non_default(self, cuda_conv3d):
|
||||
"""Non-default stride + padding + dilation all at once."""
|
||||
x = torch.randn(1, 16, 16, 16, 16, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
_run_conv3d_test(cuda_conv3d, x, w, None, (1, 2, 1), (1, 0, 1), (1, 2, 1))
|
||||
|
||||
def test_bias_with_stride(self, cuda_conv3d):
|
||||
"""Bias with non-default stride."""
|
||||
x = torch.randn(1, 16, 16, 16, 16, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
b = torch.randn(32, device="cuda")
|
||||
_run_conv3d_test(cuda_conv3d, x, w, b, (2, 2, 2), (1, 1, 1), (1, 1, 1))
|
||||
|
||||
def test_bias_with_dilation(self, cuda_conv3d):
|
||||
"""Bias with non-default dilation."""
|
||||
x = torch.randn(1, 16, 16, 16, 16, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
b = torch.randn(32, device="cuda")
|
||||
_run_conv3d_test(cuda_conv3d, x, w, b, (1, 1, 1), (2, 2, 2), (2, 2, 2))
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# FP4 Quantized Path: Advanced Conv Params
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _run_quant_smoke_test(cuda_conv3d, x, w, bias, stride, padding, dilation, fp4_block_size=16):
|
||||
"""Helper: run FP4-quantized conv and verify basic sanity."""
|
||||
act_amax = x.abs().max().unsqueeze(0)
|
||||
out = cuda_conv3d(
|
||||
x,
|
||||
w,
|
||||
bias=bias,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
act_amax=act_amax,
|
||||
quant_act=True,
|
||||
fp4_block_size=fp4_block_size,
|
||||
)
|
||||
ref = F.conv3d(x, w, bias=bias, stride=stride, padding=padding, dilation=dilation)
|
||||
assert out.shape == ref.shape, f"Shape mismatch: {out.shape} vs {ref.shape}"
|
||||
assert not torch.isnan(out).any(), "Output contains NaN"
|
||||
assert not torch.isinf(out).any(), "Output contains Inf"
|
||||
# Quantized output should be in a reasonable range relative to reference
|
||||
if ref.abs().max() > 0:
|
||||
ratio = out.abs().max().item() / ref.abs().max().item()
|
||||
assert 0.01 < ratio < 100, f"Output magnitude ratio {ratio:.2f} is unreasonable"
|
||||
return out
|
||||
|
||||
|
||||
class TestConv3dFP4QuantAdvanced:
|
||||
"""FP4 quantized path with non-trivial stride, dilation, and kernel shapes."""
|
||||
|
||||
def test_quant_with_stride(self, cuda_conv3d):
|
||||
"""FP4 quant with stride=(2,2,2)."""
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(1, 16, 16, 16, 16, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
_run_quant_smoke_test(cuda_conv3d, x, w, None, (2, 2, 2), (1, 1, 1), (1, 1, 1))
|
||||
|
||||
def test_quant_with_dilation(self, cuda_conv3d):
|
||||
"""FP4 quant with dilation=(2,2,2)."""
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(1, 16, 16, 16, 16, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
_run_quant_smoke_test(cuda_conv3d, x, w, None, (1, 1, 1), (2, 2, 2), (2, 2, 2))
|
||||
|
||||
def test_quant_with_asymmetric_kernel(self, cuda_conv3d):
|
||||
"""FP4 quant with 1x3x3 kernel."""
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(1, 16, 8, 16, 16, device="cuda")
|
||||
w = torch.randn(32, 16, 1, 3, 3, device="cuda")
|
||||
_run_quant_smoke_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 1, 1), (1, 1, 1))
|
||||
|
||||
def test_quant_with_stride_and_dilation(self, cuda_conv3d):
|
||||
"""FP4 quant with both stride>1 and dilation>1."""
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(1, 16, 16, 16, 16, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
_run_quant_smoke_test(cuda_conv3d, x, w, None, (2, 2, 2), (1, 1, 1), (2, 2, 2))
|
||||
|
||||
def test_quant_with_no_padding(self, cuda_conv3d):
|
||||
"""FP4 quant with padding=(0,0,0)."""
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(1, 16, 8, 8, 8, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
_run_quant_smoke_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1))
|
||||
|
||||
def test_quant_bias_reference(self, cuda_conv3d, cuda_fp4):
|
||||
"""FP4 quant + bias: verify bias is added correctly by comparing with/without.
|
||||
|
||||
The difference between bias and no-bias output should equal the bias broadcast.
|
||||
"""
|
||||
torch.manual_seed(42)
|
||||
cin, cout = 256, 64
|
||||
x = torch.randn(1, cin, 4, 4, 4, device="cuda")
|
||||
w = torch.randn(cout, cin, 1, 1, 1, device="cuda")
|
||||
b = torch.randn(cout, device="cuda")
|
||||
act_amax = x.abs().max().unsqueeze(0)
|
||||
|
||||
kwargs = {
|
||||
"stride": (1, 1, 1),
|
||||
"padding": (0, 0, 0),
|
||||
"dilation": (1, 1, 1),
|
||||
"act_amax": act_amax,
|
||||
"quant_act": True,
|
||||
"fp4_block_size": 16,
|
||||
}
|
||||
out_bias = cuda_conv3d(x, w, bias=b, **kwargs)
|
||||
out_no_bias = cuda_conv3d(x, w, bias=None, **kwargs)
|
||||
|
||||
# Difference should be the bias broadcast over spatial dims
|
||||
diff = out_bias - out_no_bias # [1, Cout, D, H, W]
|
||||
expected_bias = b.view(1, -1, 1, 1, 1).expand_as(diff)
|
||||
assert torch.allclose(diff, expected_bias, atol=1e-5), (
|
||||
f"Bias diff mismatch: max {(diff - expected_bias).abs().max().item():.6e}"
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Zero / Degenerate Input Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestConv3dZeroInputs:
|
||||
"""Tests with zero and degenerate inputs."""
|
||||
|
||||
def test_zero_input(self, cuda_conv3d):
|
||||
"""Zero activation tensor should produce zero (or bias-only) output."""
|
||||
x = torch.zeros(1, 16, 8, 8, 8, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1))
|
||||
out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
assert torch.allclose(out, ref, atol=1e-5), f"Max diff: {(out - ref).abs().max().item()}"
|
||||
|
||||
def test_zero_weight(self, cuda_conv3d):
|
||||
"""Zero weight tensor should produce zero output."""
|
||||
x = torch.randn(1, 16, 8, 8, 8, device="cuda")
|
||||
w = torch.zeros(32, 16, 3, 3, 3, device="cuda")
|
||||
out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
assert torch.allclose(out, torch.zeros_like(out), atol=1e-5)
|
||||
|
||||
def test_zero_input_quant(self, cuda_conv3d):
|
||||
"""Zero input with FP4 quant should not produce NaN."""
|
||||
x = torch.zeros(1, 16, 8, 8, 8, device="cuda")
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda")
|
||||
# act_amax=0 is a tricky edge case — the kernel's scale guard should handle it
|
||||
act_amax = torch.tensor([1e-10], device="cuda") # near-zero but not exactly 0
|
||||
out = cuda_conv3d(
|
||||
x,
|
||||
w,
|
||||
stride=(1, 1, 1),
|
||||
padding=(1, 1, 1),
|
||||
act_amax=act_amax,
|
||||
quant_act=True,
|
||||
fp4_block_size=16,
|
||||
)
|
||||
assert not torch.isnan(out).any(), "Zero input with quant produced NaN"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Numerical Stability Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestConv3dNumericalStability:
|
||||
"""Test with extreme value ranges."""
|
||||
|
||||
def test_large_values(self, cuda_conv3d):
|
||||
"""Large input values (randn * 100)."""
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(1, 16, 8, 8, 8, device="cuda") * 100
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda") * 100
|
||||
ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1))
|
||||
out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
# With large values, BF16 rounding error scales proportionally
|
||||
rel_err = (out - ref).abs().max().item() / ref.abs().max().item()
|
||||
assert rel_err < 0.05, f"Relative error {rel_err:.4f} too high for large values"
|
||||
|
||||
def test_small_values(self, cuda_conv3d):
|
||||
"""Small input values (randn * 1e-3)."""
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(1, 16, 8, 8, 8, device="cuda") * 1e-3
|
||||
w = torch.randn(32, 16, 3, 3, 3, device="cuda") * 1e-3
|
||||
ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1))
|
||||
out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
assert out.shape == ref.shape
|
||||
# Small values: absolute error is small, relative error may be larger due to BF16
|
||||
max_diff = (out - ref).abs().max().item()
|
||||
assert max_diff < 1e-5, f"Max diff {max_diff:.6e} for small values"
|
||||
|
||||
def test_uniform_input(self, cuda_conv3d):
|
||||
"""Uniform input (all ones) — exposes accumulation patterns."""
|
||||
x = torch.ones(1, 16, 8, 8, 8, device="cuda")
|
||||
w = torch.ones(32, 16, 3, 3, 3, device="cuda")
|
||||
ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1))
|
||||
out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False)
|
||||
k_size = 16 * 27
|
||||
scaled_atol = ATOL * (k_size / 1000.0) ** 0.5
|
||||
max_diff = (out - ref).abs().max().item()
|
||||
assert max_diff < scaled_atol, f"Uniform input: max diff {max_diff:.4f}"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Exact Block Boundary Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestConv3dExactBoundaries:
|
||||
"""Shapes that land exactly on BLOCK_M=64, BLOCK_N=64, BLOCK_K=256 boundaries."""
|
||||
|
||||
def test_m_exact_128(self, cuda_conv3d):
|
||||
"""M = 128 = 2 * BLOCK_M (exactly 2 M-tiles, no remainder)."""
|
||||
# batch=1, output 4x4x8 = 128 with kernel 1x1x1
|
||||
x = torch.randn(1, 32, 4, 4, 8, device="cuda")
|
||||
w = torch.randn(64, 32, 1, 1, 1, device="cuda")
|
||||
_run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1))
|
||||
|
||||
def test_k_exact_512(self, cuda_conv3d):
|
||||
"""K = 512 = 2 * BLOCK_K (exactly 2 K-tiles, no remainder)."""
|
||||
# Cin=512, kernel 1x1x1 -> K=512
|
||||
x = torch.randn(1, 512, 4, 4, 4, device="cuda")
|
||||
w = torch.randn(64, 512, 1, 1, 1, device="cuda")
|
||||
_run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1))
|
||||
|
||||
def test_cout_exact_64(self, cuda_conv3d):
|
||||
"""Cout = 64 = 1 * BLOCK_N (exactly 1 N-tile)."""
|
||||
x = torch.randn(1, 16, 8, 8, 8, device="cuda")
|
||||
w = torch.randn(64, 16, 3, 3, 3, device="cuda")
|
||||
_run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (1, 1, 1), (1, 1, 1))
|
||||
|
||||
def test_all_exact_multiples(self, cuda_conv3d):
|
||||
"""M=64, N=64, K=256 — single tile in each dimension."""
|
||||
# batch=1, Cin=256, kernel 1x1x1 -> K=256; output 4x4x4=64; Cout=64
|
||||
x = torch.randn(1, 256, 4, 4, 4, device="cuda")
|
||||
w = torch.randn(64, 256, 1, 1, 1, device="cuda")
|
||||
_run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1))
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# FP4 Fake Quant: Shape and Edge Case Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestFP4FakeQuantShapes:
|
||||
"""Test fp4_fake_quant with multi-dimensional inputs."""
|
||||
|
||||
def test_3d_shape_preservation(self, cuda_fp4):
|
||||
"""3D input should preserve shape after quantization."""
|
||||
inp = torch.randn(4, 8, 32, device="cuda") # numel=1024
|
||||
amax = inp.abs().max().unsqueeze(0)
|
||||
out = cuda_fp4(inp, amax, block_size=16)
|
||||
assert out.shape == (4, 8, 32)
|
||||
|
||||
def test_4d_shape_preservation(self, cuda_fp4):
|
||||
"""4D input should preserve shape."""
|
||||
inp = torch.randn(2, 4, 8, 16, device="cuda") # numel=1024
|
||||
amax = inp.abs().max().unsqueeze(0)
|
||||
out = cuda_fp4(inp, amax, block_size=16)
|
||||
assert out.shape == (2, 4, 8, 16)
|
||||
|
||||
def test_5d_shape_preservation(self, cuda_fp4):
|
||||
"""5D input (like a Conv3D activation) should preserve shape."""
|
||||
inp = torch.randn(1, 4, 4, 4, 16, device="cuda") # numel=1024
|
||||
amax = inp.abs().max().unsqueeze(0)
|
||||
out = cuda_fp4(inp, amax, block_size=16)
|
||||
assert out.shape == (1, 4, 4, 4, 16)
|
||||
|
||||
def test_multidim_correctness(self, cuda_fp4):
|
||||
"""Multi-dim quantization should equal flatten -> quant -> reshape."""
|
||||
torch.manual_seed(42)
|
||||
inp = torch.randn(4, 8, 32, device="cuda")
|
||||
amax = inp.abs().max().unsqueeze(0)
|
||||
|
||||
out_3d = cuda_fp4(inp, amax, block_size=16)
|
||||
out_flat = cuda_fp4(inp.reshape(-1), amax, block_size=16).reshape(4, 8, 32)
|
||||
assert torch.equal(out_3d, out_flat)
|
||||
|
||||
|
||||
class TestFP4FakeQuantEdgeCases:
|
||||
"""Edge cases for fp4_fake_quant."""
|
||||
|
||||
def test_very_large_values(self, cuda_fp4):
|
||||
"""Very large input values should saturate to max E2M1 level, not produce NaN."""
|
||||
inp = torch.tensor([1e6, -1e6, 5e5, -5e5, 1e4, -1e4, 100, -100], device="cuda")
|
||||
amax = inp.abs().max().unsqueeze(0)
|
||||
out = cuda_fp4(inp, amax, block_size=8)
|
||||
assert not torch.isnan(out).any()
|
||||
assert not torch.isinf(out).any()
|
||||
|
||||
def test_very_small_values(self, cuda_fp4):
|
||||
"""Very small input values should quantize to zero or near-zero."""
|
||||
inp = torch.tensor([1e-8, -1e-8, 1e-10, -1e-10, 1e-6, -1e-6, 0, 0], device="cuda")
|
||||
amax = torch.tensor([1.0], device="cuda")
|
||||
out = cuda_fp4(inp, amax, block_size=8)
|
||||
assert not torch.isnan(out).any()
|
||||
# Very small values relative to amax should quantize to ~0
|
||||
assert out.abs().max() < 1e-3
|
||||
|
||||
def test_uniform_block(self, cuda_fp4):
|
||||
"""All-same-value block."""
|
||||
inp = torch.full((16,), 3.0, device="cuda")
|
||||
amax = inp.abs().max().unsqueeze(0)
|
||||
out = cuda_fp4(inp, amax, block_size=16)
|
||||
# All elements are the same, so they should all quantize to the same E2M1 level
|
||||
assert (out == out[0]).all(), f"Uniform block produced non-uniform output: {out}"
|
||||
|
||||
def test_near_zero_amax(self, cuda_fp4):
|
||||
"""Very small global_amax should not produce NaN/Inf."""
|
||||
inp = torch.randn(16, device="cuda") * 1e-8
|
||||
amax = torch.tensor([1e-10], device="cuda")
|
||||
out = cuda_fp4(inp, amax, block_size=16)
|
||||
assert not torch.isnan(out).any(), "Near-zero amax produced NaN"
|
||||
assert not torch.isinf(out).any(), "Near-zero amax produced Inf"
|
||||
@@ -0,0 +1,129 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Tests for the quantized WanCausalConv3d wrapper in the diffusers plugin.
|
||||
|
||||
WanCausalConv3d applies asymmetric causal padding before the underlying Conv3D,
|
||||
so the quantized subclass has to replicate that padding logic around the
|
||||
quantized forward. These tests run on CPU with the NVFP4 dispatch disabled (by
|
||||
only exercising the default quantized path); the NVFP4 implicit-GEMM path is
|
||||
covered by the GPU tests under ``tests/gpu/torch/quantization/kernels``.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytest.importorskip("onnx")
|
||||
|
||||
from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d
|
||||
|
||||
# Triggers registration of _QuantDiffusersWanCausalConv3d.
|
||||
import modelopt.torch.quantization.plugins.diffusion.diffusers # noqa: F401
|
||||
from modelopt.torch.quantization.nn import QuantModuleRegistry
|
||||
from modelopt.torch.quantization.plugins.diffusion.diffusers import _QuantDiffusersWanCausalConv3d
|
||||
|
||||
|
||||
def _make_quantized(in_ch: int = 4, out_ch: int = 6, padding=(1, 1, 1)) -> torch.nn.Module:
|
||||
m = WanCausalConv3d(in_ch, out_ch, kernel_size=3, padding=padding)
|
||||
m.eval()
|
||||
# Convert via the registry so the generated class picks up WanCausalConv3d
|
||||
# in its MRO (needed for ``nn.Conv3d._conv_forward`` to be reachable).
|
||||
return QuantModuleRegistry.convert(m)
|
||||
|
||||
|
||||
def _strip_quant_state(state_dict: dict) -> dict:
|
||||
return {
|
||||
k: v
|
||||
for k, v in state_dict.items()
|
||||
if not any(
|
||||
k.startswith(p) for p in ("input_quantizer", "weight_quantizer", "output_quantizer")
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
class TestQuantWanCausalConv3dRegistration:
|
||||
def test_registered(self):
|
||||
assert WanCausalConv3d in QuantModuleRegistry
|
||||
# Registry-returned class is a generated subclass of our dm class.
|
||||
assert issubclass(QuantModuleRegistry[WanCausalConv3d], _QuantDiffusersWanCausalConv3d)
|
||||
|
||||
def test_convert_preserves_type_identity(self):
|
||||
m = _make_quantized()
|
||||
mro = [c.__name__ for c in type(m).__mro__]
|
||||
assert "WanCausalConv3d" in mro
|
||||
assert "_QuantDiffusersWanCausalConv3d" in mro
|
||||
|
||||
def test_no_implicit_gemm_without_nvfp4(self):
|
||||
# Default quantizer config is INT8 — must NOT route through implicit GEMM.
|
||||
m = _make_quantized()
|
||||
assert not m._should_use_implicit_gemm()
|
||||
|
||||
|
||||
class TestQuantWanCausalConv3dForward:
|
||||
"""Exercise the default quantized path (NVFP4 kernel is GPU-only).
|
||||
|
||||
We disable the quantizers and assert output matches an unquantized
|
||||
``WanCausalConv3d`` with the same weights — this verifies the causal-padding
|
||||
logic in the overridden ``forward`` is preserved after conversion.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize("padding", [(1, 1, 1), (2, 0, 0), (0, 1, 1)])
|
||||
def test_matches_unquantized_no_cache(self, padding):
|
||||
torch.manual_seed(0)
|
||||
m = _make_quantized(in_ch=4, out_ch=6, padding=padding)
|
||||
m.input_quantizer.disable()
|
||||
m.weight_quantizer.disable()
|
||||
|
||||
m_ref = WanCausalConv3d(4, 6, kernel_size=3, padding=padding)
|
||||
m_ref.eval()
|
||||
m_ref.load_state_dict(_strip_quant_state(m.state_dict()), strict=False)
|
||||
|
||||
x = torch.randn(1, 4, 5, 6, 6)
|
||||
out = m(x)
|
||||
out_ref = m_ref(x)
|
||||
assert torch.allclose(out, out_ref, atol=1e-5), (
|
||||
f"Max diff: {(out - out_ref).abs().max().item()}"
|
||||
)
|
||||
|
||||
def test_matches_unquantized_with_cache_x(self):
|
||||
"""cache_x is the temporal-cache branch used during causal decoding."""
|
||||
torch.manual_seed(0)
|
||||
m = _make_quantized(padding=(1, 1, 1)) # _padding[4] == 2 > 0 → cache path active
|
||||
m.input_quantizer.disable()
|
||||
m.weight_quantizer.disable()
|
||||
|
||||
m_ref = WanCausalConv3d(4, 6, kernel_size=3, padding=(1, 1, 1))
|
||||
m_ref.eval()
|
||||
m_ref.load_state_dict(_strip_quant_state(m.state_dict()), strict=False)
|
||||
|
||||
x = torch.randn(1, 4, 5, 6, 6)
|
||||
cache_x = torch.randn(1, 4, 1, 6, 6)
|
||||
out = m(x, cache_x=cache_x)
|
||||
out_ref = m_ref(x, cache_x=cache_x)
|
||||
assert torch.allclose(out, out_ref, atol=1e-5)
|
||||
|
||||
def test_output_quantizer_applied(self):
|
||||
"""Enabling the output quantizer must change the forward output."""
|
||||
torch.manual_seed(0)
|
||||
m = _make_quantized()
|
||||
m.input_quantizer.disable()
|
||||
m.weight_quantizer.disable()
|
||||
# Output quantizer is disabled by default; enable it and check it takes effect.
|
||||
x = torch.randn(1, 4, 3, 5, 5)
|
||||
out_disabled = m(x)
|
||||
m.output_quantizer.enable()
|
||||
out_enabled = m(x)
|
||||
# INT8 default config clamps and rounds; at least some elements differ.
|
||||
assert not torch.allclose(out_disabled, out_enabled)
|
||||
@@ -15,6 +15,8 @@
|
||||
|
||||
"""Tests of QuantConv module."""
|
||||
|
||||
import warnings
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
@@ -23,11 +25,19 @@ from torch import nn
|
||||
from modelopt.torch.quantization import tensor_quant
|
||||
from modelopt.torch.quantization.config import QuantizerAttributeConfig
|
||||
from modelopt.torch.quantization.nn.modules import quant_conv
|
||||
from modelopt.torch.quantization.nn.modules.quant_conv import (
|
||||
_is_nvfp4_quantizer,
|
||||
_nvfp4_quantize_weight_along_k,
|
||||
)
|
||||
from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
|
||||
|
||||
NUM_IN_CHANNELS = 3
|
||||
NUM_OUT_CHANNELS = 5
|
||||
|
||||
_NVFP4_CFG = QuantizerAttributeConfig(
|
||||
num_bits=(2, 1), block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)}
|
||||
)
|
||||
|
||||
|
||||
class TestQuantConvND:
|
||||
@pytest.mark.parametrize(
|
||||
@@ -316,3 +326,134 @@ class TestQuantConvND:
|
||||
output = conv(test_input)
|
||||
|
||||
assert torch.allclose(quant_output, output)
|
||||
|
||||
|
||||
class TestQuantConv3dNVFP4:
|
||||
"""Tests for the NVFP4 implicit GEMM dispatch path in ``_QuantConv3d``.
|
||||
|
||||
The CUDA kernel itself is GPU-only (covered in ``tests/gpu``); here we only
|
||||
exercise the CPU-side predicate and fallback branches that currently lack
|
||||
coverage.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _make_nvfp4_conv3d(groups: int = 1, bias: bool = False) -> quant_conv.QuantConv3d:
|
||||
return quant_conv.QuantConv3d(
|
||||
NUM_IN_CHANNELS * groups,
|
||||
NUM_OUT_CHANNELS * groups,
|
||||
kernel_size=3,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
quant_desc_input=_NVFP4_CFG,
|
||||
quant_desc_weight=_NVFP4_CFG,
|
||||
)
|
||||
|
||||
def test_is_nvfp4_quantizer_true(self):
|
||||
q = TensorQuantizer(_NVFP4_CFG)
|
||||
assert _is_nvfp4_quantizer(q)
|
||||
|
||||
def test_is_nvfp4_quantizer_false_for_int8(self):
|
||||
q = TensorQuantizer(QuantizerAttributeConfig(num_bits=8))
|
||||
assert not _is_nvfp4_quantizer(q)
|
||||
|
||||
def test_is_nvfp4_quantizer_false_for_fp8(self):
|
||||
# FP8 E4M3: num_bits == (4, 3), block_sizes is None
|
||||
q = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3)))
|
||||
assert not _is_nvfp4_quantizer(q)
|
||||
|
||||
def test_is_nvfp4_quantizer_false_for_static_block(self):
|
||||
static_cfg = QuantizerAttributeConfig(
|
||||
num_bits=(2, 1),
|
||||
block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)},
|
||||
)
|
||||
q = TensorQuantizer(static_cfg)
|
||||
assert not _is_nvfp4_quantizer(q)
|
||||
|
||||
def test_should_use_implicit_gemm_true(self):
|
||||
m = self._make_nvfp4_conv3d()
|
||||
assert m._should_use_implicit_gemm()
|
||||
|
||||
def test_should_use_implicit_gemm_false_groups_gt_1(self):
|
||||
m = self._make_nvfp4_conv3d(groups=NUM_IN_CHANNELS)
|
||||
# Groups > 1 disqualifies the fused kernel even with NVFP4 quantizers.
|
||||
assert not m._should_use_implicit_gemm()
|
||||
|
||||
def test_should_use_implicit_gemm_false_non_nvfp4(self):
|
||||
# INT8 per-tensor config — default Conv3D quantizers, no NVFP4.
|
||||
m = quant_conv.QuantConv3d(NUM_IN_CHANNELS, NUM_OUT_CHANNELS, kernel_size=3, bias=False)
|
||||
assert not m._should_use_implicit_gemm()
|
||||
|
||||
def test_nvfp4_quantize_weight_along_k_reshape(self):
|
||||
"""Verify weight is flattened/restored along the K (input) dimension.
|
||||
|
||||
Uses a stub quantizer so this test stays CPU-only (the real NVFP4 dynamic
|
||||
quantizer requires a CUDA tensor).
|
||||
"""
|
||||
|
||||
def identity_quantizer(x):
|
||||
# K-dim must be the last axis when passed to the quantizer so NVFP4's
|
||||
# block-wise scaling aligns with the GEMM reduction axis.
|
||||
assert x.dim() == 2 and x.shape[0] == NUM_OUT_CHANNELS
|
||||
return x
|
||||
|
||||
w = torch.randn(NUM_OUT_CHANNELS, NUM_IN_CHANNELS, 3, 3, 3)
|
||||
qw = _nvfp4_quantize_weight_along_k(w, identity_quantizer)
|
||||
assert qw.shape == w.shape
|
||||
assert torch.equal(qw, w)
|
||||
|
||||
def test_forward_non_nvfp4_matches_unquantized(self):
|
||||
# Disabled quantizers: forward must match plain conv3d exactly.
|
||||
m = quant_conv.QuantConv3d(NUM_IN_CHANNELS, NUM_OUT_CHANNELS, kernel_size=3, bias=False)
|
||||
m.input_quantizer.disable()
|
||||
m.weight_quantizer.disable()
|
||||
x = torch.randn(1, NUM_IN_CHANNELS, 4, 4, 4)
|
||||
out = m(x)
|
||||
ref = F.conv3d(x, m.weight)
|
||||
assert torch.allclose(out, ref)
|
||||
|
||||
def test_forward_nvfp4_training_warns_and_falls_back(self):
|
||||
"""Training mode must fall back to the default (cuDNN) path with a warning.
|
||||
|
||||
The implicit-GEMM kernel is inference-only; this exercises the CPU-visible
|
||||
training-fallback branch. We disable the quantizers so the default-path
|
||||
``super().forward()`` does not try to run NVFP4 dynamic quantization
|
||||
(which requires CUDA).
|
||||
"""
|
||||
m = self._make_nvfp4_conv3d()
|
||||
# NVFP4 predicate reads configuration (num_bits/block_sizes), not enable
|
||||
# state, so disabling the quantizers still routes through the NVFP4 branch.
|
||||
m.input_quantizer.disable()
|
||||
m.weight_quantizer.disable()
|
||||
assert m._should_use_implicit_gemm()
|
||||
m.train()
|
||||
x = torch.randn(1, NUM_IN_CHANNELS, 4, 4, 4)
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
out = m(x)
|
||||
assert out.shape == (1, NUM_OUT_CHANNELS, 2, 2, 2)
|
||||
assert any("inference-only" in str(w.message) for w in caught), (
|
||||
f"Expected an 'inference-only' warning, got: {[str(w.message) for w in caught]}"
|
||||
)
|
||||
|
||||
def test_forward_nvfp4_calib_only_uses_default_path(self):
|
||||
"""Calibration-only mode must NOT try to invoke the CUDA kernel.
|
||||
|
||||
When the input quantizer is in calibration mode without quant enabled,
|
||||
the forward must fall back to the default path. This exercises the
|
||||
calib-only early-return branch; we assert that the output matches the
|
||||
default-path output exactly (i.e. the implicit-GEMM path wasn't used).
|
||||
"""
|
||||
m = self._make_nvfp4_conv3d()
|
||||
m.eval()
|
||||
# Match the state toggled by TensorQuantizer.disable_quant()/enable_calib().
|
||||
m.input_quantizer.disable_quant()
|
||||
m.input_quantizer.enable_calib()
|
||||
m.weight_quantizer.disable_quant()
|
||||
m.weight_quantizer.enable_calib()
|
||||
assert m.input_quantizer._if_calib and not m.input_quantizer._if_quant
|
||||
|
||||
x = torch.randn(1, NUM_IN_CHANNELS, 4, 4, 4)
|
||||
out = m(x)
|
||||
# Default path with quant disabled should equal plain conv3d.
|
||||
ref = F.conv3d(x, m.weight.detach())
|
||||
assert torch.allclose(out, ref, atol=1e-5)
|
||||
|
||||
Reference in New Issue
Block a user