[feat] add MXFP8 support on H3 (#1796)
This commit is contained in:
Junda Su
@@ -0,0 +1,147 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Generate 5-second FastH3 videos with a supported Preview adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
FASTVIDEO_ROOT = Path(__file__).resolve().parents[3]
|
||||
sys.path.insert(0, str(FASTVIDEO_ROOT))
|
||||
|
||||
from fastvideo import VideoGenerator # noqa: E402
|
||||
from fastvideo.api import ( # noqa: E402
|
||||
CompileConfig,
|
||||
ComponentConfig,
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
ParallelismConfig,
|
||||
PipelineSelection,
|
||||
SamplingConfig,
|
||||
)
|
||||
|
||||
# FastVideo exposes these three backend switches only through environment variables.
|
||||
os.environ.update({
|
||||
"FASTVIDEO_FA4": "1",
|
||||
"FASTVIDEO_MINIMAX_H3_FUSIONS": "all",
|
||||
"FASTVIDEO_VSA_SM100A": "0",
|
||||
})
|
||||
os.environ.pop("FASTVIDEO_INFERENCE_TORCH_COMPILE", None)
|
||||
|
||||
VARIANT_BACKENDS = {
|
||||
"dense-datafree": "FLASH_ATTN",
|
||||
"vsa-datafree": "VIDEO_SPARSE_ATTN_H3",
|
||||
"vsa-synthetic-step1300": "VIDEO_SPARSE_ATTN_H3",
|
||||
"vsa-synthetic-step1900": "VIDEO_SPARSE_ATTN_H3",
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("variant", choices=VARIANT_BACKENDS)
|
||||
parser.add_argument("--prompt", required=True)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Run one compile warmup and three measured generations with one fixed recipe."""
|
||||
args = parse_args()
|
||||
|
||||
attention_backend = VARIANT_BACKENDS[args.variant]
|
||||
experimental = {
|
||||
"attention_backend": attention_backend,
|
||||
"inference_torch_compile": attention_backend == "FLASH_ATTN",
|
||||
"vae_parallel_decode": True,
|
||||
"vae_parallel_decode_strategy": "gather",
|
||||
}
|
||||
if attention_backend == "VIDEO_SPARSE_ATTN_H3":
|
||||
experimental.update({
|
||||
"VSA_sparsity": 0.9,
|
||||
"VSA_tile_size": 64,
|
||||
})
|
||||
|
||||
adapter_path = hf_hub_download(
|
||||
repo_id="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA",
|
||||
filename=f"{args.variant}/adapter_model.safetensors",
|
||||
)
|
||||
output_dir = FASTVIDEO_ROOT / "outputs/fasth3_lora_preview" / args.variant
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="MiniMaxAI/MiniMax-H3",
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(lora_path=adapter_path, lora_strength=1.0),
|
||||
experimental=experimental,
|
||||
),
|
||||
engine=EngineConfig(
|
||||
num_gpus=4,
|
||||
parallelism=ParallelismConfig(tp_size=1, sp_size=4),
|
||||
offload=OffloadConfig(dit=False, dit_layerwise=False),
|
||||
compile=CompileConfig(vae_enabled=True),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
generation_count = 4
|
||||
warmup_count = 1
|
||||
measured_count = generation_count - warmup_count
|
||||
measured_seconds: list[float] = []
|
||||
|
||||
print(f"Variant: {args.variant} ({attention_backend})")
|
||||
print(f"Output directory: {output_dir}")
|
||||
try:
|
||||
for generation_index in range(generation_count):
|
||||
measured = generation_index >= warmup_count
|
||||
if measured:
|
||||
measured_index = generation_index - warmup_count + 1
|
||||
label = f"measured {measured_index}/{measured_count}"
|
||||
output_path = output_dir / f"fasth3_all_run_{measured_index:02d}.mp4"
|
||||
else:
|
||||
warmup_index = generation_index + 1
|
||||
label = f"warmup {warmup_index}/{warmup_count}"
|
||||
output_path = output_dir / f"_fasth3_warmup_{warmup_index:02d}.mp4"
|
||||
|
||||
request = GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
negative_prompt="",
|
||||
sampling=SamplingConfig(
|
||||
height=768,
|
||||
width=1344,
|
||||
num_frames=345, # <-- Change video length: 5sec: 124; 10sec: 243; 15sec: 345.
|
||||
fps=24,
|
||||
num_inference_steps=5,
|
||||
guidance_scale=1.0,
|
||||
batch_cfg=False,
|
||||
seed=1000,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(output_path),
|
||||
save_video=True,
|
||||
return_frames=False,
|
||||
),
|
||||
)
|
||||
started = time.perf_counter()
|
||||
result = generator.generate(request)
|
||||
elapsed = time.perf_counter() - started
|
||||
if measured:
|
||||
measured_seconds.append(elapsed)
|
||||
suffix = "" if measured else " (excluded from median)"
|
||||
print(f"[{label}] {result.video_path or output_path}: {elapsed:.3f}s{suffix}")
|
||||
finally:
|
||||
generator.shutdown()
|
||||
|
||||
print(f"Median E2E wall time: {statistics.median(measured_seconds):.3f}s")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,10 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
PROMPT="integrated_multimodal_description: A red fox runs through fresh snow at dawn. overall_soundscape: Fast pawsteps in snow, winter wind, and distant birds."
|
||||
|
||||
# Options: dense-datafree, vsa-datafree, vsa-synthetic-step1300, vsa-synthetic-step1900
|
||||
VARIANT="vsa-datafree"
|
||||
|
||||
exec python "$(dirname "$0")/basic_fasth3_simplified.py" \
|
||||
"$VARIANT" \
|
||||
--prompt "$PROMPT"
|
||||
@@ -0,0 +1,126 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Profile one FastH3 Preview generation with NVTX after warmup."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
from fastvideo import VideoGenerator
|
||||
from fastvideo.api import (
|
||||
CompileConfig,
|
||||
ComponentConfig,
|
||||
EngineConfig,
|
||||
GenerationRequest,
|
||||
GeneratorConfig,
|
||||
OffloadConfig,
|
||||
OutputConfig,
|
||||
ParallelismConfig,
|
||||
PipelineSelection,
|
||||
QuantizationConfig,
|
||||
SamplingConfig,
|
||||
)
|
||||
from fastvideo.profiler import nvtx_range
|
||||
|
||||
VARIANT_BACKENDS = {
|
||||
"dense-datafree": "FLASH_ATTN",
|
||||
"vsa-datafree": "VIDEO_SPARSE_ATTN_H3",
|
||||
"vsa-synthetic-step1300": "VIDEO_SPARSE_ATTN_H3",
|
||||
"vsa-synthetic-step1900": "VIDEO_SPARSE_ATTN_H3",
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("variant", choices=VARIANT_BACKENDS)
|
||||
parser.add_argument("--prompt", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--num-frames", type=int, default=345)
|
||||
parser.add_argument("--warmup-runs", type=int, default=3)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Warm the selected FastH3 recipe, then profile one identical generation."""
|
||||
args = parse_args()
|
||||
output_dir = Path(args.output)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
attention_backend = VARIANT_BACKENDS[args.variant]
|
||||
experimental = {
|
||||
"attention_backend": attention_backend,
|
||||
"inference_torch_compile": attention_backend == "FLASH_ATTN",
|
||||
"vae_parallel_decode": True,
|
||||
"vae_parallel_decode_strategy": "gather",
|
||||
}
|
||||
if attention_backend == "VIDEO_SPARSE_ATTN_H3":
|
||||
experimental.update({
|
||||
"VSA_sparsity": 0.9,
|
||||
"VSA_tile_size": 64,
|
||||
})
|
||||
|
||||
adapter_path = hf_hub_download(
|
||||
repo_id="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA",
|
||||
filename=f"{args.variant}/adapter_model.safetensors",
|
||||
)
|
||||
generator = VideoGenerator.from_config(
|
||||
GeneratorConfig(
|
||||
model_path="MiniMaxAI/MiniMax-H3",
|
||||
pipeline=PipelineSelection(
|
||||
components=ComponentConfig(lora_path=adapter_path, lora_strength=1.0),
|
||||
experimental=experimental,
|
||||
),
|
||||
engine=EngineConfig(
|
||||
num_gpus=4,
|
||||
parallelism=ParallelismConfig(tp_size=1, sp_size=4),
|
||||
offload=OffloadConfig(dit=False, dit_layerwise=False),
|
||||
compile=CompileConfig(vae_enabled=True),
|
||||
quantization=QuantizationConfig(transformer_quant="MXFP8"),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
request = GenerationRequest(
|
||||
prompt=args.prompt,
|
||||
negative_prompt="",
|
||||
sampling=SamplingConfig(
|
||||
height=768,
|
||||
width=1344,
|
||||
num_frames=args.num_frames,
|
||||
fps=24,
|
||||
num_inference_steps=5,
|
||||
guidance_scale=1.0,
|
||||
batch_cfg=False,
|
||||
seed=1000,
|
||||
),
|
||||
output=OutputConfig(
|
||||
output_path=str(output_dir / "fasth3_profile.mp4"),
|
||||
save_video=True,
|
||||
return_frames=False,
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
for warmup_index in range(args.warmup_runs):
|
||||
warmup_result = generator.generate(request)
|
||||
print(f"Warmup {warmup_index + 1}/{args.warmup_runs}: {warmup_result.video_path}")
|
||||
|
||||
torch.cuda.profiler.start()
|
||||
try:
|
||||
with nvtx_range("fasth3.profiled_generation"):
|
||||
measured_result = generator.generate(request)
|
||||
finally:
|
||||
torch.cuda.profiler.stop()
|
||||
|
||||
print(f"Profiled output: {measured_result.video_path}")
|
||||
if measured_result.generation_time is not None:
|
||||
print(f"Generation time: {measured_result.generation_time:.2f}s")
|
||||
finally:
|
||||
generator.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,53 @@
|
||||
#!/usr/bin/env bash
|
||||
# Capture one FastH3 Preview generation with CUDA and FastVideo NVTX ranges.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
WORKTREE_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
WORKSPACE_ROOT="$(cd -- "${WORKTREE_ROOT}/.." && pwd)"
|
||||
PYTHON_BIN="${WORKSPACE_ROOT}/.venv-fv/bin/python"
|
||||
NSYS_BIN=/usr/local/cuda/bin/nsys
|
||||
PROFILE_SCRIPT="${SCRIPT_DIR}/basic_fasth3_simplified_profile.py"
|
||||
|
||||
export CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||
export FASTVIDEO_FA4=1
|
||||
export FASTVIDEO_INFERENCE_TORCH_COMPILE=0
|
||||
export FASTVIDEO_MINIMAX_H3_FUSIONS=all
|
||||
export FASTVIDEO_NVTX_PROFILE=1
|
||||
export FASTVIDEO_VSA_SM100A=0
|
||||
export PYTHONUNBUFFERED=1
|
||||
export PYTHONPATH="${WORKTREE_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
|
||||
|
||||
PROMPT="integrated_multimodal_description: A red fox runs through fresh snow at dawn. overall_soundscape: Fast pawsteps in snow, winter wind, and distant birds."
|
||||
|
||||
# Options: dense-datafree, vsa-datafree, vsa-synthetic-step1300, vsa-synthetic-step1900
|
||||
VARIANT="vsa-datafree"
|
||||
NUM_FRAMES=345
|
||||
WARMUP_RUNS=3
|
||||
|
||||
PROFILE_ID="${VARIANT}_4gpu_sp4_${NUM_FRAMES}frames"
|
||||
RUN_ID="$(date -u +%Y%m%dT%H%M%SZ)-$$"
|
||||
RESULT_DIR="${WORKTREE_ROOT}/runs/fasth3_lora_preview_profile/${PROFILE_ID}/${RUN_ID}"
|
||||
MEDIA_DIR="${RESULT_DIR}/media"
|
||||
RUN_LOG="${RESULT_DIR}/run.log"
|
||||
mkdir -p "${MEDIA_DIR}"
|
||||
|
||||
"${NSYS_BIN}" profile \
|
||||
--trace=cuda,nvtx \
|
||||
--sample=none \
|
||||
--cpuctxsw=none \
|
||||
--capture-range=cudaProfilerApi \
|
||||
--capture-range-end=stop \
|
||||
--output="${RESULT_DIR}/${PROFILE_ID}" \
|
||||
"${PYTHON_BIN}" "${PROFILE_SCRIPT}" \
|
||||
"${VARIANT}" \
|
||||
--num-frames "${NUM_FRAMES}" \
|
||||
--warmup-runs "${WARMUP_RUNS}" \
|
||||
--prompt "${PROMPT}" \
|
||||
--output "${MEDIA_DIR}" \
|
||||
2>&1 | tee "${RUN_LOG}"
|
||||
|
||||
printf 'RESULT_DIR=%s\n' "${RESULT_DIR}"
|
||||
printf 'NSYS_REPORT=%s\n' "${RESULT_DIR}/${PROFILE_ID}.nsys-rep"
|
||||
printf 'MEDIA_DIR=%s\n' "${MEDIA_DIR}"
|
||||
@@ -0,0 +1,276 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MXFP8 block quantization and linear operations for Blackwell inference."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from quack.mx_utils import to_blocked, to_mx
|
||||
|
||||
MXFP8_BLOCK_SIZE = 32
|
||||
_BLOCKS_PER_PROGRAM = 16
|
||||
|
||||
|
||||
def _validate_mxfp8_matrix(matrix: torch.Tensor) -> None:
|
||||
"""Validate one row-major matrix before MXFP8 quantization."""
|
||||
if matrix.ndim != 2:
|
||||
raise ValueError(f"MXFP8 quantization requires a 2D tensor, got shape {tuple(matrix.shape)}.")
|
||||
if matrix.dtype not in (torch.bfloat16, torch.float32):
|
||||
raise TypeError(f"MXFP8 quantization requires BF16 or FP32 input, got {matrix.dtype}.")
|
||||
if matrix.shape[1] % MXFP8_BLOCK_SIZE:
|
||||
raise ValueError(f"MXFP8 reduction dimension must be divisible by {MXFP8_BLOCK_SIZE}, got {matrix.shape[1]}.")
|
||||
|
||||
|
||||
@torch.compile(dynamic=True, fullgraph=True)
|
||||
def _quantize_mxfp8_weight_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
values, natural_scales = to_mx(matrix.contiguous(), MXFP8_BLOCK_SIZE)
|
||||
return values, to_blocked(natural_scales)
|
||||
|
||||
|
||||
def quantize_mxfp8_weight_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Prequantize a weight with Quack and return hardware-blocked scales."""
|
||||
_validate_mxfp8_matrix(matrix)
|
||||
return _quantize_mxfp8_weight_blockwise(matrix)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _store_mxfp8_blocks(
|
||||
matrix,
|
||||
quantized_ptr,
|
||||
scale_ptr,
|
||||
row,
|
||||
first_block,
|
||||
row_count,
|
||||
column_count,
|
||||
scale_column_count,
|
||||
BLOCKS_PER_PROGRAM: tl.constexpr,
|
||||
):
|
||||
"""Quantize row blocks and store scales in cuBLAS's 32x4x4 layout."""
|
||||
element_offsets = tl.arange(0, BLOCKS_PER_PROGRAM * 32)
|
||||
columns = first_block * 32 + element_offsets
|
||||
valid_elements = (row < row_count) & (columns < column_count)
|
||||
matrix = tl.reshape(matrix, (BLOCKS_PER_PROGRAM, 32))
|
||||
|
||||
# FLOOR scaling encodes exponent(max_abs) - 8 as one E8M0 byte per 32 values.
|
||||
max_abs = tl.max(tl.abs(matrix), axis=1)
|
||||
exponent = (max_abs.to(tl.int32, bitcast=True) >> 23) & 0xFF
|
||||
scale_biased = tl.maximum(tl.minimum(exponent - 8, 255), 0)
|
||||
scale_biased = tl.where(max_abs != max_abs, 255, scale_biased)
|
||||
scale_bits = scale_biased.to(tl.int32) << 23
|
||||
scale = tl.maximum(scale_bits.to(tl.float32, bitcast=True), 1.1754943508222875e-38)
|
||||
|
||||
scaled_matrix = tl.maximum(tl.minimum(matrix / scale[:, None], 448.0), -448.0)
|
||||
tl.store(
|
||||
quantized_ptr + row * column_count + columns,
|
||||
tl.reshape(scaled_matrix.to(tl.float8e4nv), (BLOCKS_PER_PROGRAM * 32, )),
|
||||
mask=valid_elements,
|
||||
)
|
||||
|
||||
blocks = first_block + tl.arange(0, BLOCKS_PER_PROGRAM)
|
||||
# Map natural [row, K / 32] coordinates into cuBLAS's padded 128x4 tiles.
|
||||
row_block = row // 128
|
||||
row_in_block = row % 128
|
||||
column_block = blocks // 4
|
||||
column_in_block = blocks % 4
|
||||
scale_offset = ((row_block * (scale_column_count // 4) + column_block) * 512 + (row_in_block % 32) * 16 +
|
||||
(row_in_block // 32) * 4 + column_in_block)
|
||||
tl.store(scale_ptr + scale_offset, scale_biased.to(tl.uint8), mask=blocks < scale_column_count)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _quantize_mxfp8_kernel(
|
||||
matrix_ptr,
|
||||
quantized_ptr,
|
||||
scale_ptr,
|
||||
row_count,
|
||||
column_count,
|
||||
scale_column_count,
|
||||
BLOCKS_PER_PROGRAM: tl.constexpr,
|
||||
):
|
||||
"""Load BF16 activation blocks for direct MXFP8 quantization."""
|
||||
row = tl.program_id(0)
|
||||
first_block = tl.program_id(1) * BLOCKS_PER_PROGRAM
|
||||
element_offsets = tl.arange(0, BLOCKS_PER_PROGRAM * 32)
|
||||
columns = first_block * 32 + element_offsets
|
||||
matrix = tl.load(
|
||||
matrix_ptr + row * column_count + columns,
|
||||
mask=(row < row_count) & (columns < column_count),
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
_store_mxfp8_blocks(
|
||||
matrix,
|
||||
quantized_ptr,
|
||||
scale_ptr,
|
||||
row,
|
||||
first_block,
|
||||
row_count,
|
||||
column_count,
|
||||
scale_column_count,
|
||||
BLOCKS_PER_PROGRAM,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _swiglu_quantize_mxfp8_kernel(
|
||||
preactivation_ptr,
|
||||
quantized_ptr,
|
||||
scale_ptr,
|
||||
row_count,
|
||||
column_count,
|
||||
scale_column_count,
|
||||
BLOCKS_PER_PROGRAM: tl.constexpr,
|
||||
):
|
||||
"""Apply value-first SwiGLU and directly quantize its BF16 result."""
|
||||
row = tl.program_id(0)
|
||||
first_block = tl.program_id(1) * BLOCKS_PER_PROGRAM
|
||||
element_offsets = tl.arange(0, BLOCKS_PER_PROGRAM * 32)
|
||||
columns = first_block * 32 + element_offsets
|
||||
valid_elements = (row < row_count) & (columns < column_count)
|
||||
packed_row_offset = row * (2 * column_count)
|
||||
values = tl.load(
|
||||
preactivation_ptr + packed_row_offset + columns,
|
||||
mask=valid_elements,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
gates = tl.load(
|
||||
preactivation_ptr + packed_row_offset + column_count + columns,
|
||||
mask=valid_elements,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
postactivation = (values * gates * tl.sigmoid(gates)).to(tl.bfloat16).to(tl.float32)
|
||||
_store_mxfp8_blocks(
|
||||
postactivation,
|
||||
quantized_ptr,
|
||||
scale_ptr,
|
||||
row,
|
||||
first_block,
|
||||
row_count,
|
||||
column_count,
|
||||
scale_column_count,
|
||||
BLOCKS_PER_PROGRAM,
|
||||
)
|
||||
|
||||
|
||||
def _allocate_mxfp8_outputs(matrix: torch.Tensor, column_count: int) -> tuple[torch.Tensor, torch.Tensor, int]:
|
||||
"""Allocate contiguous values and a complete padded blocked-scale buffer."""
|
||||
row_count = matrix.shape[0]
|
||||
scale_column_count = triton.cdiv(column_count // MXFP8_BLOCK_SIZE, 4) * 4
|
||||
blocked_scale_size = triton.cdiv(row_count, 128) * 128 * scale_column_count
|
||||
quantized = torch.empty((row_count, column_count), dtype=torch.float8_e4m3fn, device=matrix.device)
|
||||
scale_storage = torch.empty(blocked_scale_size, dtype=torch.uint8, device=matrix.device)
|
||||
return quantized, scale_storage.view(torch.float8_e8m0fnu), scale_column_count
|
||||
|
||||
|
||||
def quantize_mxfp8_blockwise(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Quantize activation rows and write hardware-blocked scales directly."""
|
||||
_validate_mxfp8_matrix(matrix)
|
||||
matrix = matrix.contiguous()
|
||||
row_count, column_count = matrix.shape
|
||||
quantized, blocked_scales, scale_column_count = _allocate_mxfp8_outputs(matrix, column_count)
|
||||
grid = (triton.cdiv(row_count, 128) * 128, triton.cdiv(scale_column_count, _BLOCKS_PER_PROGRAM))
|
||||
_quantize_mxfp8_kernel[grid](
|
||||
matrix,
|
||||
quantized,
|
||||
blocked_scales.view(torch.uint8),
|
||||
row_count,
|
||||
column_count,
|
||||
scale_column_count,
|
||||
BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
|
||||
num_warps=8,
|
||||
)
|
||||
return quantized, blocked_scales
|
||||
|
||||
|
||||
def swiglu_quantize_mxfp8_blockwise(preactivation: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Apply H3 value-first SwiGLU and quantize the BF16 result to MXFP8."""
|
||||
if preactivation.ndim != 2:
|
||||
raise ValueError(f"MXFP8 SwiGLU requires a 2D tensor, got shape {tuple(preactivation.shape)}.")
|
||||
if preactivation.dtype != torch.bfloat16:
|
||||
raise TypeError(f"MXFP8 SwiGLU requires BF16 input, got {preactivation.dtype}.")
|
||||
if preactivation.shape[1] % (2 * MXFP8_BLOCK_SIZE):
|
||||
raise ValueError("MXFP8 SwiGLU requires each packed half to be divisible by "
|
||||
f"{MXFP8_BLOCK_SIZE}, got packed width {preactivation.shape[1]}.")
|
||||
preactivation = preactivation.contiguous()
|
||||
row_count = preactivation.shape[0]
|
||||
column_count = preactivation.shape[1] // 2
|
||||
quantized, blocked_scales, scale_column_count = _allocate_mxfp8_outputs(preactivation, column_count)
|
||||
grid = (triton.cdiv(row_count, 128) * 128, triton.cdiv(scale_column_count, _BLOCKS_PER_PROGRAM))
|
||||
_swiglu_quantize_mxfp8_kernel[grid](
|
||||
preactivation,
|
||||
quantized,
|
||||
blocked_scales.view(torch.uint8),
|
||||
row_count,
|
||||
column_count,
|
||||
scale_column_count,
|
||||
BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
|
||||
num_warps=8,
|
||||
)
|
||||
return quantized, blocked_scales
|
||||
|
||||
|
||||
def mxfp8_scaled_mm(
|
||||
activation_values: torch.Tensor,
|
||||
activation_scales: torch.Tensor,
|
||||
weight_values: torch.Tensor,
|
||||
weight_scales: torch.Tensor,
|
||||
bias: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
"""Multiply two MXFP8 matrices and return a BF16 matrix."""
|
||||
return F.scaled_mm(
|
||||
mat_a=activation_values,
|
||||
mat_b=weight_values.mT,
|
||||
scale_a=activation_scales,
|
||||
scale_recipe_a=F.ScalingType.BlockWise1x32,
|
||||
scale_b=weight_scales,
|
||||
scale_recipe_b=F.ScalingType.BlockWise1x32,
|
||||
swizzle_a=F.SwizzleType.SWIZZLE_32_4_4,
|
||||
swizzle_b=F.SwizzleType.SWIZZLE_32_4_4,
|
||||
bias=bias,
|
||||
output_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_merged_linear(linear: torch.nn.Module) -> torch.nn.Module:
|
||||
"""Return a linear whose weight includes every active inference adapter."""
|
||||
base_layer = getattr(linear, "base_layer", linear)
|
||||
if base_layer is linear:
|
||||
return base_layer
|
||||
if not getattr(linear, "merged", False) and not getattr(linear, "disable_lora", False):
|
||||
raise RuntimeError("MXFP8 feed-forward requires active LoRA weights to be merged before inference.")
|
||||
return base_layer
|
||||
|
||||
|
||||
def mxfp8_swiglu_feed_forward(
|
||||
hidden_states: torch.Tensor,
|
||||
fc_in: torch.nn.Module,
|
||||
fc_out: torch.nn.Module,
|
||||
) -> torch.Tensor:
|
||||
"""Run the MiniMax-H3 feed-forward network with MXFP8 GEMMs."""
|
||||
fc_in_base = _resolve_merged_linear(fc_in)
|
||||
fc_out_base = _resolve_merged_linear(fc_out)
|
||||
fc_in_method = fc_in_base.quant_method
|
||||
fc_out_method = fc_out_base.quant_method
|
||||
|
||||
preactivation = fc_in_method.apply(fc_in_base, hidden_states, fc_in_base.bias)
|
||||
output_shape = (*hidden_states.shape[:-1], fc_out_base.output_size)
|
||||
preactivation_2d = preactivation.reshape(-1, preactivation.shape[-1])
|
||||
activation_values, activation_scales = swiglu_quantize_mxfp8_blockwise(preactivation_2d)
|
||||
output_2d = fc_out_method.apply_quantized(
|
||||
fc_out_base,
|
||||
activation_values,
|
||||
activation_scales,
|
||||
fc_out_base.bias,
|
||||
)
|
||||
return output_2d.reshape(output_shape)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MXFP8_BLOCK_SIZE",
|
||||
"mxfp8_scaled_mm",
|
||||
"mxfp8_swiglu_feed_forward",
|
||||
"quantize_mxfp8_blockwise",
|
||||
"quantize_mxfp8_weight_blockwise",
|
||||
"swiglu_quantize_mxfp8_blockwise",
|
||||
]
|
||||
@@ -2,7 +2,16 @@ from typing import Literal, get_args
|
||||
|
||||
from fastvideo.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
QuantizationMethods = Literal[None, "AbsMaxFP8", "FP8", "NVFP4", "nvfp4_qat", "nvfp4_qat_train", "fp8_qat_train"]
|
||||
QuantizationMethods = Literal[
|
||||
None,
|
||||
"AbsMaxFP8",
|
||||
"FP8",
|
||||
"MXFP8",
|
||||
"NVFP4",
|
||||
"nvfp4_qat",
|
||||
"nvfp4_qat_train",
|
||||
"fp8_qat_train",
|
||||
]
|
||||
|
||||
QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
|
||||
|
||||
@@ -52,6 +61,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
|
||||
# lazy import to avoid triggering `torch.compile` too early
|
||||
from .absmax_fp8 import AbsMaxFP8Config
|
||||
from .fp8_config import FP8Config
|
||||
from .mxfp8_config import MXFP8Config
|
||||
from .nvfp4_config import NVFP4Config
|
||||
from .nvfp4_qat_config import NVFP4QATConfig
|
||||
from .nvfp4_qat_train_config import NVFP4QATTrainConfig
|
||||
@@ -60,6 +70,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
|
||||
method_to_config: dict[str, type[QuantizationConfig]] = {
|
||||
"AbsMaxFP8": AbsMaxFP8Config,
|
||||
"FP8": FP8Config,
|
||||
"MXFP8": MXFP8Config,
|
||||
"NVFP4": NVFP4Config,
|
||||
"nvfp4_qat": NVFP4QATConfig,
|
||||
"nvfp4_qat_train": NVFP4QATTrainConfig,
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MXFP8 quantization for MiniMax-H3 transformer-block feed-forward layers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.models.utils import set_weight_attrs
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_MINIMAX_H3_FF_PREFIX = re.compile(r"(?:^|\.)transformer_blocks\.\d+\.ff\.(?:fc_in|fc_out)$")
|
||||
|
||||
|
||||
class MXFP8QuantizeMethod(QuantizeMethodBase):
|
||||
"""Dynamically quantize activations against prequantized MXFP8 weights."""
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs: Any,
|
||||
) -> None:
|
||||
"""Create the BF16 checkpoint weight and non-persistent MXFP8 buffers."""
|
||||
del input_size, output_size
|
||||
weight = Parameter(
|
||||
torch.empty(
|
||||
sum(output_partition_sizes),
|
||||
input_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
|
||||
layer.register_parameter("weight", weight)
|
||||
set_weight_attrs(weight, extra_weight_attrs)
|
||||
layer.register_buffer("_mxfp8_weight", None, persistent=False)
|
||||
layer.register_buffer("_mxfp8_weight_scale", None, persistent=False)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
"""Prequantize one adapter-merged BF16 linear weight."""
|
||||
from fastvideo.layers.mxfp8linear import quantize_mxfp8_weight_blockwise
|
||||
|
||||
quantized_weight, blocked_scales = quantize_mxfp8_weight_blockwise(layer.weight.detach())
|
||||
layer._mxfp8_weight = quantized_weight
|
||||
layer._mxfp8_weight_scale = blocked_scales
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
hidden_states: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Quantize one BF16 activation and apply the prequantized linear."""
|
||||
from fastvideo.layers.mxfp8linear import quantize_mxfp8_blockwise
|
||||
|
||||
original_shape = hidden_states.shape
|
||||
hidden_states_2d = hidden_states.reshape(-1, hidden_states.shape[-1])
|
||||
activation_values, activation_scales = quantize_mxfp8_blockwise(hidden_states_2d)
|
||||
output_2d = self.apply_quantized(layer, activation_values, activation_scales, bias)
|
||||
return output_2d.reshape(*original_shape[:-1], layer.output_size)
|
||||
|
||||
def apply_quantized(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
activation_values: torch.Tensor,
|
||||
activation_scales: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Apply one linear to an activation that is already in MXFP8."""
|
||||
from fastvideo.layers.mxfp8linear import mxfp8_scaled_mm
|
||||
|
||||
if layer._mxfp8_weight is None or layer._mxfp8_weight_scale is None:
|
||||
raise RuntimeError(f"MXFP8 weight buffers are not initialized for {layer.prefix}.")
|
||||
return mxfp8_scaled_mm(
|
||||
activation_values,
|
||||
activation_scales,
|
||||
layer._mxfp8_weight,
|
||||
layer._mxfp8_weight_scale,
|
||||
bias,
|
||||
)
|
||||
|
||||
|
||||
class MXFP8Config(QuantizationConfig):
|
||||
"""Select MXFP8 for the main MiniMax-H3 transformer-block FFN linears."""
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "MXFP8"
|
||||
|
||||
def get_supported_act_dtypes(self) -> list[torch.dtype]:
|
||||
return [torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 100
|
||||
|
||||
@staticmethod
|
||||
def get_config_filenames() -> list[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> MXFP8Config:
|
||||
del config
|
||||
return cls()
|
||||
|
||||
def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
|
||||
from fastvideo.layers.linear import LinearBase
|
||||
|
||||
if isinstance(layer, LinearBase) and _MINIMAX_H3_FF_PREFIX.search(prefix):
|
||||
return MXFP8QuantizeMethod()
|
||||
return None
|
||||
|
||||
|
||||
def convert_model_to_mxfp8(model: torch.nn.Module) -> int:
|
||||
"""Prequantize every MXFP8-tagged weight and return the converted count."""
|
||||
converted_count = 0
|
||||
with torch.no_grad():
|
||||
for module in model.modules():
|
||||
quant_method = getattr(module, "quant_method", None)
|
||||
if not isinstance(quant_method, MXFP8QuantizeMethod):
|
||||
continue
|
||||
if converted_count == 0 and (not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10):
|
||||
raise RuntimeError("MXFP8 inference requires an NVIDIA Blackwell GPU with compute capability 10.0+.")
|
||||
quant_method.process_weights_after_loading(module)
|
||||
converted_count += 1
|
||||
|
||||
if converted_count:
|
||||
logger.info("Prequantized %d MiniMax-H3 feed-forward linear weights to MXFP8", converted_count)
|
||||
torch.cuda.empty_cache()
|
||||
return converted_count
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MXFP8Config",
|
||||
"MXFP8QuantizeMethod",
|
||||
"convert_model_to_mxfp8",
|
||||
]
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""LTX-2 NVFP4 quantization (FlashInfer-backed).
|
||||
"""NVFP4 linear quantization for supported transformer layer sets.
|
||||
|
||||
NVFP4 is NVIDIA's block-scaled FP4 format (e2m1 mantissa, fp32 alpha,
|
||||
``layout_128x4`` scale layout, group size 16) — distinct from
|
||||
@@ -8,9 +8,8 @@ explicitly so downstream callers don't conflate it with other FP4
|
||||
variants that may land later (e.g. AMD's MX-FP4 or vendor-neutral
|
||||
e3m0).
|
||||
|
||||
Upstreamed from ``FastVideo-internal`` so consumers that load LTX-2
|
||||
weights with NVFP4 quantization can drive the public package
|
||||
end-to-end.
|
||||
The registered config targets the curated LTX-2 deployment set and the
|
||||
main MiniMax-H3 transformer-block FFN linears.
|
||||
|
||||
`flashinfer` is imported lazily inside the call paths that need it.
|
||||
This keeps ``import fastvideo`` cheap on hosts where flashinfer is
|
||||
@@ -20,6 +19,7 @@ use time, with a clear error.
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
@@ -75,6 +75,7 @@ _LTX2_NVFP4_BLOCK_LINEAR_SUFFIXES = (
|
||||
_LTX2_NVFP4_LINEAR_PREFIXES = frozenset(f"ltx2.blocks.{block_idx}.{suffix}" for block_idx in range(48)
|
||||
for suffix in _LTX2_NVFP4_BLOCK_LINEAR_SUFFIXES) | frozenset(
|
||||
("ltx2.adaln_single.linear", ))
|
||||
_MINIMAX_H3_NVFP4_FF_PREFIX = re.compile(r"(?:^|\.)transformer_blocks\.\d+\.ff\.(?:fc_in|fc_out)$")
|
||||
|
||||
|
||||
def is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool:
|
||||
@@ -82,6 +83,11 @@ def is_ltx2_nvfp4_linear_prefix(prefix: str) -> bool:
|
||||
return prefix in _LTX2_NVFP4_LINEAR_PREFIXES
|
||||
|
||||
|
||||
def is_minimax_h3_nvfp4_linear_prefix(prefix: str) -> bool:
|
||||
"""Return whether *prefix* is a main MiniMax-H3 transformer-block FFN linear."""
|
||||
return _MINIMAX_H3_NVFP4_FF_PREFIX.search(prefix) is not None
|
||||
|
||||
|
||||
def _is_ltx2_refine_only_prefix(prefix: str) -> bool:
|
||||
return any(prefix.endswith(suffix) for suffix in _LTX2_REFINE_ONLY_SUFFIXES)
|
||||
|
||||
@@ -421,13 +427,12 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
|
||||
|
||||
|
||||
class NVFP4Config(QuantizationConfig):
|
||||
"""LTX-2-specific NVFP4 quantization configuration.
|
||||
"""Select NVFP4 for the supported LTX-2 and MiniMax-H3 linear sets.
|
||||
|
||||
NVFP4 is NVIDIA's block-scaled FP4 (e2m1 mantissa, fp32 alpha,
|
||||
``layout_128x4`` scale layout, group size 16). Today this class
|
||||
hardcodes the LTX-2 layer paths it covers. When a second model
|
||||
wants NVFP4, lift the layer-path list into a config field
|
||||
instead of hardcoding it here.
|
||||
``layout_128x4`` scale layout, group size 16). LTX-2 uses its curated
|
||||
attention and FFN deployment set. MiniMax-H3 uses only ``fc_in`` and
|
||||
``fc_out`` in each main transformer-block FFN.
|
||||
"""
|
||||
|
||||
def __init__(self, layer_profile: str = "refine", retain_original_weights: bool | None = None):
|
||||
@@ -468,9 +473,10 @@ class NVFP4Config(QuantizationConfig):
|
||||
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
|
||||
from fastvideo.layers.linear import LinearBase
|
||||
|
||||
# Use the superset at build/load time, then switch active subset
|
||||
# dynamically in NVFP4QuantizeMethod.apply based on stage profile.
|
||||
if isinstance(layer, LinearBase) and is_ltx2_nvfp4_linear_prefix(prefix):
|
||||
# LTX-2 switches its active subset by stage at runtime. MiniMax-H3
|
||||
# uses the fixed main-transformer FFN set selected by its prefix.
|
||||
if isinstance(layer, LinearBase) and (is_ltx2_nvfp4_linear_prefix(prefix)
|
||||
or is_minimax_h3_nvfp4_linear_prefix(prefix)):
|
||||
method = NVFP4QuantizeMethod(layer_prefix=prefix)
|
||||
method._retain_original_weights = self.retain_original_weights
|
||||
return method
|
||||
@@ -552,4 +558,5 @@ __all__ = [
|
||||
"NVFP4QuantizeMethod",
|
||||
"convert_model_to_nvfp4",
|
||||
"is_ltx2_nvfp4_linear_prefix",
|
||||
"is_minimax_h3_nvfp4_linear_prefix",
|
||||
]
|
||||
|
||||
@@ -21,6 +21,7 @@ from fastvideo.distributed.communication_op import (
|
||||
)
|
||||
from fastvideo.distributed.parallel_state import get_sp_world_size, model_parallel_is_initialized
|
||||
from fastvideo.layers.linear import ReplicatedLinear
|
||||
from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod
|
||||
from fastvideo.layers.mlp import MLP
|
||||
from fastvideo.layers.quantization import QuantizationConfig
|
||||
from fastvideo.layers.visual_embedding import Timesteps
|
||||
@@ -119,8 +120,14 @@ class MiniMaxH3FeedForward(nn.Module):
|
||||
prefix=f"{prefix}.fc_out",
|
||||
)
|
||||
self.fuse_swiglu = fuse_swiglu
|
||||
self.use_mxfp8 = isinstance(self.fc_in.quant_method, MXFP8QuantizeMethod) and isinstance(
|
||||
self.fc_out.quant_method, MXFP8QuantizeMethod)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
if self.use_mxfp8:
|
||||
from fastvideo.layers.mxfp8linear import mxfp8_swiglu_feed_forward
|
||||
|
||||
return mxfp8_swiglu_feed_forward(hidden_states, self.fc_in, self.fc_out)
|
||||
hidden_states, _ = self.fc_in(hidden_states)
|
||||
if self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states):
|
||||
hidden_states = minimax_h3_swiglu(hidden_states)
|
||||
@@ -514,41 +521,48 @@ class MiniMaxH3TransformerBlock(nn.Module):
|
||||
rotary_emb: tuple[torch.Tensor, torch.Tensor],
|
||||
original_seq_len: int,
|
||||
) -> torch.Tensor:
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
t.to(hidden_states.dtype) for t in self.adaln_proj(temb))
|
||||
with nvtx_range("minimax_h3.transformer_block.adaln_projection"):
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
t.to(hidden_states.dtype) for t in self.adaln_proj(temb))
|
||||
|
||||
use_modulate_fusion = self.fuse_modulate and _can_run_minimax_h3_fusion(hidden_states)
|
||||
if use_modulate_fusion:
|
||||
norm_hidden_states = fused_rmsnorm_modulate(
|
||||
hidden_states,
|
||||
self.norm1.weight,
|
||||
scale_msa,
|
||||
shift_msa,
|
||||
adaln_indices,
|
||||
self.norm1.eps,
|
||||
)
|
||||
with nvtx_range("minimax_h3.transformer_block.modulate_fusion"):
|
||||
norm_hidden_states = fused_rmsnorm_modulate(
|
||||
hidden_states,
|
||||
self.norm1.weight,
|
||||
scale_msa,
|
||||
shift_msa,
|
||||
adaln_indices,
|
||||
self.norm1.eps,
|
||||
)
|
||||
else:
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
norm_hidden_states = norm_hidden_states * (
|
||||
1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices)
|
||||
attention_output = self.attn(norm_hidden_states, rotary_emb, original_seq_len)
|
||||
with nvtx_range("minimax_h3.transformer_block.no_modulate_fusion"):
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
norm_hidden_states = norm_hidden_states * (
|
||||
1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices)
|
||||
with nvtx_range("minimax_h3.transformer_block.self_attention"):
|
||||
attention_output = self.attn(norm_hidden_states, rotary_emb, original_seq_len)
|
||||
if use_modulate_fusion:
|
||||
hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate(
|
||||
hidden_states,
|
||||
attention_output,
|
||||
gate_msa,
|
||||
self.norm2.weight,
|
||||
scale_mlp,
|
||||
shift_mlp,
|
||||
adaln_indices,
|
||||
self.norm2.eps,
|
||||
)
|
||||
with nvtx_range("minimax_h3.transformer_block.modulate_fusion"):
|
||||
hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate(
|
||||
hidden_states,
|
||||
attention_output,
|
||||
gate_msa,
|
||||
self.norm2.weight,
|
||||
scale_mlp,
|
||||
shift_mlp,
|
||||
adaln_indices,
|
||||
self.norm2.eps,
|
||||
)
|
||||
else:
|
||||
hidden_states = hidden_states + gate_msa.index_select(0, adaln_indices) * attention_output
|
||||
norm_hidden_states = self.norm2(hidden_states)
|
||||
norm_hidden_states = norm_hidden_states * (
|
||||
1.0 + scale_mlp.index_select(0, adaln_indices)) + shift_mlp.index_select(0, adaln_indices)
|
||||
feed_forward_output = self.ff(norm_hidden_states)
|
||||
with nvtx_range("minimax_h3.transformer_block.no_modulate_fusion"):
|
||||
hidden_states = hidden_states + gate_msa.index_select(0, adaln_indices) * attention_output
|
||||
norm_hidden_states = self.norm2(hidden_states)
|
||||
norm_hidden_states = norm_hidden_states * (
|
||||
1.0 + scale_mlp.index_select(0, adaln_indices)) + shift_mlp.index_select(0, adaln_indices)
|
||||
with nvtx_range("minimax_h3.transformer_block.feed_forward"):
|
||||
feed_forward_output = self.ff(norm_hidden_states)
|
||||
return hidden_states + gate_mlp.index_select(0, adaln_indices) * feed_forward_output
|
||||
|
||||
|
||||
|
||||
@@ -36,14 +36,14 @@ def _summarize_param_names(names: set[str]) -> str:
|
||||
return ", ".join(f"{family} x{count}" if count > 1 else family for family, count in sorted(families.items()))
|
||||
|
||||
|
||||
def _maybe_quantize_model(model: nn.Module) -> None:
|
||||
"""Quantize NVFP4- or FP8-tagged linear layers in-place after weights are loaded.
|
||||
def _maybe_quantize_model(model: nn.Module, *, defer_weight_conversion_until_lora_merge: bool = False) -> None:
|
||||
"""Quantize inference linear weights after checkpoint loading.
|
||||
|
||||
Walks the module tree once, looking for layers whose ``quant_method``
|
||||
is an :class:`NVFP4QuantizeMethod` or :class:`FP8QuantizeMethod` (attached
|
||||
at construction time by the respective ``get_quant_method``). When at least
|
||||
Walks the module tree once, looking for layers whose ``quant_method`` is an
|
||||
inference quantization method attached at construction time. When at least
|
||||
one such layer exists, calls the matching conversion function to register
|
||||
quantized weight buffers on each targeted layer.
|
||||
quantized weight buffers on each targeted layer. NVFP4 and MXFP8 conversion
|
||||
can be deferred until a configured inference LoRA is merged.
|
||||
|
||||
The walk returns on the first quantized layer found so unquantized callers
|
||||
pay only an ``isinstance`` check per module. Both imports are deferred so
|
||||
@@ -70,12 +70,19 @@ def _maybe_quantize_model(model: nn.Module) -> None:
|
||||
FP8QuantizeMethod,
|
||||
convert_model_to_fp8,
|
||||
)
|
||||
from fastvideo.layers.quantization.mxfp8_config import (
|
||||
MXFP8QuantizeMethod,
|
||||
convert_model_to_mxfp8,
|
||||
)
|
||||
|
||||
qat_train_attached = 0
|
||||
qat_train_skipped = 0
|
||||
for mod in model.modules():
|
||||
qm = getattr(mod, "quant_method", None)
|
||||
if isinstance(qm, NVFP4QuantizeMethod):
|
||||
if defer_weight_conversion_until_lora_merge:
|
||||
logger.info("Deferring NVFP4 weight conversion until the inference LoRA merge completes")
|
||||
return
|
||||
logger.info("Converting loaded model weights for NVFP4 linear layers")
|
||||
convert_model_to_nvfp4(model)
|
||||
return
|
||||
@@ -87,6 +94,13 @@ def _maybe_quantize_model(model: nn.Module) -> None:
|
||||
logger.info("Converting loaded model weights for FP8 linear layers")
|
||||
convert_model_to_fp8(model)
|
||||
return
|
||||
if isinstance(qm, MXFP8QuantizeMethod):
|
||||
if defer_weight_conversion_until_lora_merge:
|
||||
logger.info("Deferring MXFP8 weight conversion until the inference LoRA merge completes")
|
||||
return
|
||||
logger.info("Converting loaded model weights for MXFP8 linear layers")
|
||||
convert_model_to_mxfp8(model)
|
||||
return
|
||||
# QAT-train configs are mutually exclusive with the inference schemes
|
||||
# above (one quant_config per model), so when they're active the loop
|
||||
# always runs to completion and the counts below are model-wide.
|
||||
@@ -144,6 +158,26 @@ def _prepare_model_for_compile(model: nn.Module, *, regional: bool) -> str | Non
|
||||
return unsupported if isinstance(unsupported, str) and unsupported else None
|
||||
|
||||
|
||||
def _validate_fsdp_inference_quantization(init_params: dict[str, Any], fsdp_inference: bool) -> None:
|
||||
"""Enforce the transformer-quantization allowlist for FSDP inference."""
|
||||
transformer_config = init_params.get("config")
|
||||
quant_config = getattr(transformer_config, "quant_config", None)
|
||||
if not fsdp_inference or quant_config is None:
|
||||
return
|
||||
|
||||
from fastvideo.layers.quantization.nvfp4_qat_train_config import (
|
||||
NVFP4QATTrainConfig, )
|
||||
|
||||
if isinstance(quant_config, NVFP4QATTrainConfig):
|
||||
return
|
||||
|
||||
# TODO: (David) Currently reject every FSDP inference quantization config except NVFP4QATTrainConfig.
|
||||
# Support FSDP inference with precomputed quantized weights.
|
||||
raise NotImplementedError(
|
||||
"FSDP inference supports unquantized transformers and NVFP4QATTrainConfig only; "
|
||||
f"got {type(quant_config).__name__}.")
|
||||
|
||||
|
||||
# Supports optional torch.compile for FSDP-wrapped models during training
|
||||
def maybe_load_fsdp_model(
|
||||
model_cls: type[nn.Module],
|
||||
@@ -178,6 +212,8 @@ def maybe_load_fsdp_model(
|
||||
base checkpoint does not contain, which has to happen while the tensor is still
|
||||
unsharded.
|
||||
"""
|
||||
_validate_fsdp_inference_quantization(init_params, fsdp_inference)
|
||||
|
||||
# NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
|
||||
# manually casting the inputs to the model
|
||||
mp_policy = MixedPrecisionPolicy(param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False)
|
||||
@@ -287,7 +323,7 @@ def maybe_load_fsdp_model(
|
||||
# responsibility is just to materialize the quantized weight buffers
|
||||
# from the freshly-loaded bf16 weights. No-op when no quantized layers
|
||||
# are present (lazy imports inside the helper).
|
||||
_maybe_quantize_model(model)
|
||||
_maybe_quantize_model(model, defer_weight_conversion_until_lora_merge=lora_path is not None)
|
||||
|
||||
compile_in_loader = enable_torch_compile and training_mode
|
||||
if compile_in_loader:
|
||||
|
||||
@@ -32,6 +32,36 @@ from fastvideo.utils import maybe_download_lora
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _has_quantized_mxfp8_weights(transformer_modules: dict[str, nn.Module]) -> bool:
|
||||
"""Return whether any transformer contains an initialized MXFP8 weight."""
|
||||
from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod
|
||||
|
||||
return any(
|
||||
isinstance(getattr(module, "quant_method", None), MXFP8QuantizeMethod)
|
||||
and getattr(module, "_mxfp8_weight", None) is not None for transformer in transformer_modules.values()
|
||||
for module in transformer.modules())
|
||||
|
||||
|
||||
def _has_nvfp4_weights_without_bf16(transformer_modules: dict[str, nn.Module]) -> bool:
|
||||
"""Return whether NVFP4 quantization removed any transformer's BF16 weight."""
|
||||
from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
|
||||
|
||||
return any(
|
||||
isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod)
|
||||
and getattr(module, "_nvfp4_weight", None) is not None and getattr(module, "weight", None) is None
|
||||
for transformer in transformer_modules.values() for module in transformer.modules())
|
||||
|
||||
|
||||
def _has_quantized_nvfp4_weights(transformer_modules: dict[str, nn.Module]) -> bool:
|
||||
"""Return whether any transformer contains an initialized NVFP4 weight."""
|
||||
from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
|
||||
|
||||
return any(
|
||||
isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod)
|
||||
and getattr(module, "_nvfp4_weight", None) is not None for transformer in transformer_modules.values()
|
||||
for module in transformer.modules())
|
||||
|
||||
|
||||
def _get_hook_ctx(module: nn.Module | None):
|
||||
if module is None:
|
||||
return nullcontext()
|
||||
@@ -43,6 +73,24 @@ def _get_hook_ctx(module: nn.Module | None):
|
||||
return nullcontext()
|
||||
|
||||
|
||||
def _convert_quantized_weights_after_lora_merge(transformer_modules: dict[str, nn.Module]) -> None:
|
||||
"""Pack deferred NVFP4 or MXFP8 weights after all LoRA layers are merged."""
|
||||
from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod
|
||||
from fastvideo.layers.quantization.mxfp8_config import convert_model_to_mxfp8
|
||||
from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
|
||||
from fastvideo.layers.quantization.nvfp4_config import convert_model_to_nvfp4
|
||||
|
||||
for transformer_module in transformer_modules.values():
|
||||
for module in transformer_module.modules():
|
||||
quant_method = getattr(module, "quant_method", None)
|
||||
if isinstance(quant_method, NVFP4QuantizeMethod):
|
||||
convert_model_to_nvfp4(transformer_module)
|
||||
break
|
||||
if isinstance(quant_method, MXFP8QuantizeMethod):
|
||||
convert_model_to_mxfp8(transformer_module)
|
||||
break
|
||||
|
||||
|
||||
def _named_module_by_prefix(module: nn.Module,
|
||||
prefixes: list[str]) -> list[tuple[str | None, list[tuple[str, nn.Module]]]]:
|
||||
none_list: list[tuple[str, nn.Module]] = []
|
||||
@@ -389,6 +437,13 @@ class LoRAPipeline(ComposedPipelineBase):
|
||||
if exact_current_adapter:
|
||||
return
|
||||
|
||||
if not self._setting_constructor_adapter and _has_nvfp4_weights_without_bf16(
|
||||
self.trainable_transformer_modules):
|
||||
# TODO(David): Restore the BF16 weights and requantize them after an NVFP4 LoRA adapter change.
|
||||
raise RuntimeError(
|
||||
"Runtime LoRA adapter changes are unsupported after NVFP4 quantization removed the BF16 weights. "
|
||||
"Create a new VideoGenerator with the desired LoRA adapter.")
|
||||
|
||||
if not self._setting_constructor_adapter:
|
||||
if self._constructor_dense_lora_path is not None:
|
||||
raise RuntimeError(
|
||||
@@ -532,6 +587,7 @@ class LoRAPipeline(ComposedPipelineBase):
|
||||
lora_path)
|
||||
if unmatched:
|
||||
logger.warning("LoRA adapter %s: %d weights did not reach a layer", lora_path, len(unmatched))
|
||||
_convert_quantized_weights_after_lora_merge(self.trainable_transformer_modules)
|
||||
|
||||
def merge_lora_weights(self) -> None:
|
||||
for (
|
||||
@@ -547,6 +603,17 @@ class LoRAPipeline(ComposedPipelineBase):
|
||||
layer.merge_lora_weights()
|
||||
|
||||
def unmerge_lora_weights(self) -> None:
|
||||
"""Unmerge LoRA weights when the transformer's quantized weights remain valid."""
|
||||
if _has_quantized_mxfp8_weights(self.trainable_transformer_modules):
|
||||
# TODO(David): Requantize MXFP8 weights after LoRA unmerge before enabling this operation.
|
||||
raise RuntimeError(
|
||||
"LoRA unmerge is unsupported after MXFP8 weight quantization because the quantized weights still "
|
||||
"contain the merged LoRA adapter.")
|
||||
if _has_quantized_nvfp4_weights(self.trainable_transformer_modules):
|
||||
# TODO(David): Preserve BF16 weights and requantize NVFP4 weights after LoRA unmerge.
|
||||
raise RuntimeError(
|
||||
"LoRA unmerge is unsupported after NVFP4 weight quantization because the quantized weights would "
|
||||
"not reflect the unmerged LoRA adapter.")
|
||||
for (
|
||||
transformer_name,
|
||||
transformer_lora_layers,
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""FSDP inference accepts unquantized transformers and NVFP4QATTrainConfig."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from fastvideo.layers.quantization.fp8_qat_train_config import FP8QATTrainConfig
|
||||
from fastvideo.layers.quantization.mxfp8_config import MXFP8Config
|
||||
from fastvideo.layers.quantization.nvfp4_config import NVFP4Config
|
||||
from fastvideo.layers.quantization.nvfp4_qat_train_config import (
|
||||
NVFP4QATTrainConfig, )
|
||||
from fastvideo.models.loader.fsdp_load import (
|
||||
_validate_fsdp_inference_quantization,
|
||||
maybe_load_fsdp_model,
|
||||
)
|
||||
|
||||
|
||||
def _init_params(quant_config: object | None) -> dict[str, object]:
|
||||
return {"config": SimpleNamespace(quant_config=quant_config)}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("quant_config", [MXFP8Config(), NVFP4Config(), FP8QATTrainConfig()])
|
||||
def test_maybe_load_fsdp_model_unsupported_quantization_rejected_before_construction(quant_config: object) -> None:
|
||||
"""Reject unsupported quantization before model construction or FSDP sharding."""
|
||||
with pytest.raises(NotImplementedError, match=type(quant_config).__name__):
|
||||
maybe_load_fsdp_model(
|
||||
model_cls=torch.nn.Module,
|
||||
init_params=_init_params(quant_config),
|
||||
weight_dir_list=[],
|
||||
device=torch.device("cpu"),
|
||||
hsdp_replicate_dim=1,
|
||||
hsdp_shard_dim=1,
|
||||
default_dtype=torch.bfloat16,
|
||||
param_dtype=torch.bfloat16,
|
||||
reduce_dtype=torch.float32,
|
||||
fsdp_inference=True,
|
||||
training_mode=False,
|
||||
)
|
||||
|
||||
|
||||
def test_validate_fsdp_inference_quantization_nvfp4_qat_training_allowed() -> None:
|
||||
_validate_fsdp_inference_quantization(_init_params(NVFP4QATTrainConfig()), fsdp_inference=True)
|
||||
|
||||
|
||||
def test_validate_fsdp_inference_quantization_unquantized_allowed() -> None:
|
||||
_validate_fsdp_inference_quantization(_init_params(None), fsdp_inference=True)
|
||||
|
||||
|
||||
def test_validate_fsdp_inference_quantization_without_fsdp_allowed() -> None:
|
||||
_validate_fsdp_inference_quantization(_init_params(MXFP8Config()), fsdp_inference=False)
|
||||
@@ -0,0 +1,120 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Blackwell numerical parity tests for FastVideo MXFP8 operations."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from quack.mx_utils import to_blocked, to_mx
|
||||
|
||||
from fastvideo.layers.mxfp8linear import (
|
||||
quantize_mxfp8_blockwise,
|
||||
quantize_mxfp8_weight_blockwise,
|
||||
swiglu_quantize_mxfp8_blockwise,
|
||||
)
|
||||
|
||||
|
||||
def _require_blackwell() -> None:
|
||||
"""Require hardware that can execute MXFP8 block-scaled GEMMs."""
|
||||
if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10:
|
||||
pytest.skip("MXFP8 numerical parity requires an NVIDIA Blackwell GPU")
|
||||
|
||||
|
||||
def _quack_mxfp8(matrix: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Quantize one matrix with the independent Quack reference operations."""
|
||||
values, natural_scales = to_mx(matrix.contiguous(), 32)
|
||||
return values, to_blocked(natural_scales)
|
||||
|
||||
|
||||
def _assert_same_mxfp8(
|
||||
actual: tuple[torch.Tensor, torch.Tensor],
|
||||
expected: tuple[torch.Tensor, torch.Tensor],
|
||||
) -> None:
|
||||
"""Require identical FP8 values and E8M0 scale bytes."""
|
||||
actual_values, actual_scales = actual
|
||||
expected_values, expected_scales = expected
|
||||
torch.testing.assert_close(actual_values.float(), expected_values.float(), rtol=0, atol=0)
|
||||
torch.testing.assert_close(
|
||||
actual_scales.view(torch.uint8),
|
||||
expected_scales.view(torch.uint8),
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
|
||||
|
||||
def test_mxfp8_weight_quantization_matches_quack() -> None:
|
||||
"""Compare prequantized weight values and scales with Quack."""
|
||||
_require_blackwell()
|
||||
generator = torch.Generator().manual_seed(20260905)
|
||||
weight = torch.randn(129, 256, generator=generator, dtype=torch.bfloat16).cuda()
|
||||
|
||||
_assert_same_mxfp8(quantize_mxfp8_weight_blockwise(weight), _quack_mxfp8(weight))
|
||||
|
||||
|
||||
def test_mxfp8_activation_quantization_matches_quack() -> None:
|
||||
"""Compare fused activation quantization and scale swizzling with Quack."""
|
||||
_require_blackwell()
|
||||
generator = torch.Generator().manual_seed(20260905)
|
||||
activation = torch.randn(129, 256, generator=generator, dtype=torch.bfloat16).cuda()
|
||||
|
||||
_assert_same_mxfp8(quantize_mxfp8_blockwise(activation), _quack_mxfp8(activation))
|
||||
|
||||
|
||||
def test_mxfp8_swiglu_quantization_matches_bf16_quack() -> None:
|
||||
"""Compare fused SwiGLU quantization with BF16 SwiGLU followed by Quack."""
|
||||
_require_blackwell()
|
||||
generator = torch.Generator().manual_seed(20260905)
|
||||
preactivation = torch.randn(129, 512, generator=generator, dtype=torch.bfloat16).cuda()
|
||||
values, gates = preactivation.chunk(2, dim=-1)
|
||||
bf16_swiglu = (values.float() * gates.float() * torch.sigmoid(gates.float())).to(torch.bfloat16)
|
||||
|
||||
_assert_same_mxfp8(
|
||||
swiglu_quantize_mxfp8_blockwise(preactivation),
|
||||
_quack_mxfp8(bf16_swiglu),
|
||||
)
|
||||
|
||||
|
||||
def test_minimax_h3_mxfp8_feed_forward_matches_bf16() -> None:
|
||||
"""Compare the complete MXFP8 H3 feed-forward output with BF16."""
|
||||
_require_blackwell()
|
||||
from fastvideo.layers.quantization.mxfp8_config import MXFP8Config, convert_model_to_mxfp8
|
||||
from fastvideo.models.dits.minimax_h3 import MiniMaxH3FeedForward
|
||||
|
||||
previous_default_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
try:
|
||||
dense = MiniMaxH3FeedForward(128, 256, prefix="minimax_h3.transformer_blocks.0.ff")
|
||||
quantized = MiniMaxH3FeedForward(
|
||||
128,
|
||||
256,
|
||||
quant_config=MXFP8Config(),
|
||||
prefix="minimax_h3.transformer_blocks.0.ff",
|
||||
)
|
||||
finally:
|
||||
torch.set_default_dtype(previous_default_dtype)
|
||||
|
||||
generator = torch.Generator().manual_seed(20260905)
|
||||
with torch.no_grad():
|
||||
for parameter in dense.parameters():
|
||||
parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02)
|
||||
quantized.load_state_dict(dense.state_dict(), strict=True)
|
||||
dense = dense.cuda().eval()
|
||||
quantized = quantized.cuda().eval()
|
||||
assert convert_model_to_mxfp8(quantized) == 2
|
||||
|
||||
hidden_states = torch.randn(2, 128, 128, generator=generator, dtype=torch.bfloat16).cuda()
|
||||
with torch.inference_mode():
|
||||
dense_output = dense(hidden_states)
|
||||
quantized_output = quantized(hidden_states)
|
||||
|
||||
dense_output_flat = dense_output.float().flatten()
|
||||
quantized_output_flat = quantized_output.float().flatten()
|
||||
similarity = F.cosine_similarity(dense_output_flat, quantized_output_flat, dim=0)
|
||||
relative_l2_error = torch.linalg.vector_norm(dense_output_flat - quantized_output_flat) / torch.linalg.vector_norm(
|
||||
dense_output_flat
|
||||
)
|
||||
assert dense_output.dtype == torch.bfloat16
|
||||
assert quantized_output.dtype == torch.bfloat16
|
||||
assert similarity > 0.995
|
||||
assert relative_l2_error < 0.10
|
||||
@@ -0,0 +1,234 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""MiniMax-H3 NVFP4 targeting and inference-LoRA lifecycle tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import fastvideo.layers.quantization.mxfp8_config as mxfp8
|
||||
import fastvideo.layers.quantization.nvfp4_config as nvfp4
|
||||
|
||||
|
||||
class _TaggedModel(nn.Module):
|
||||
"""Minimal model whose child carries one quantization method."""
|
||||
|
||||
def __init__(self, quant_method: object) -> None:
|
||||
super().__init__()
|
||||
self.linear = nn.Module()
|
||||
self.linear.quant_method = quant_method
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prefix",
|
||||
[
|
||||
"transformer_blocks.0.ff.fc_in",
|
||||
"transformer_blocks.49.ff.fc_out",
|
||||
"minimax_h3.transformer_blocks.12.ff.fc_in",
|
||||
],
|
||||
)
|
||||
def test_is_minimax_h3_nvfp4_linear_prefix_main_ffn_linear(prefix: str) -> None:
|
||||
assert nvfp4.is_minimax_h3_nvfp4_linear_prefix(prefix)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prefix",
|
||||
[
|
||||
"transformer_blocks.0.attn.to_q",
|
||||
"transformer_blocks.0.adaln_proj.linear",
|
||||
"token_refiner.refiner_blocks.0.ff.fc_in",
|
||||
"transformer_blocks.0.ff.extra",
|
||||
],
|
||||
)
|
||||
def test_is_minimax_h3_nvfp4_linear_prefix_non_main_ffn_linear(prefix: str) -> None:
|
||||
assert not nvfp4.is_minimax_h3_nvfp4_linear_prefix(prefix)
|
||||
|
||||
|
||||
def test_nvfp4config_get_quant_method_minimax_h3_feed_forward() -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("NVFP4QuantizeMethod construction requires CUDA")
|
||||
|
||||
from fastvideo.models.dits.minimax_h3 import MiniMaxH3FeedForward
|
||||
|
||||
feed_forward = MiniMaxH3FeedForward(
|
||||
hidden_size=16,
|
||||
ffn_dim=32,
|
||||
quant_config=nvfp4.NVFP4Config(),
|
||||
prefix="minimax_h3.transformer_blocks.0.ff",
|
||||
)
|
||||
|
||||
assert isinstance(feed_forward.fc_in.quant_method, nvfp4.NVFP4QuantizeMethod)
|
||||
assert isinstance(feed_forward.fc_out.quant_method, nvfp4.NVFP4QuantizeMethod)
|
||||
|
||||
|
||||
def test_minimax_h3_feed_forward_nvfp4_cuda_similarity() -> None:
|
||||
"""Run both H3 FFN GEMMs through NVFP4 and compare with BF16."""
|
||||
if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10:
|
||||
pytest.skip("MiniMax-H3 NVFP4 inference requires an NVIDIA Blackwell GPU")
|
||||
pytest.importorskip("flashinfer")
|
||||
|
||||
from fastvideo.models.dits.minimax_h3 import MiniMaxH3FeedForward
|
||||
|
||||
previous_default_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
try:
|
||||
dense = MiniMaxH3FeedForward(128, 256, prefix="minimax_h3.transformer_blocks.0.ff")
|
||||
quantized = MiniMaxH3FeedForward(
|
||||
128,
|
||||
256,
|
||||
quant_config=nvfp4.NVFP4Config(),
|
||||
prefix="minimax_h3.transformer_blocks.0.ff",
|
||||
)
|
||||
finally:
|
||||
torch.set_default_dtype(previous_default_dtype)
|
||||
|
||||
generator = torch.Generator().manual_seed(20260901)
|
||||
with torch.no_grad():
|
||||
for parameter in dense.parameters():
|
||||
parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02)
|
||||
quantized.load_state_dict(dense.state_dict(), strict=True)
|
||||
dense = dense.cuda().eval()
|
||||
quantized = quantized.cuda().eval()
|
||||
nvfp4.convert_model_to_nvfp4(quantized)
|
||||
|
||||
inputs = torch.randn(2, 128, 128, device="cuda", dtype=torch.bfloat16)
|
||||
with torch.inference_mode():
|
||||
dense_output = dense(inputs)
|
||||
quantized_output = quantized(inputs)
|
||||
|
||||
similarity = torch.nn.functional.cosine_similarity(
|
||||
dense_output.float().flatten(),
|
||||
quantized_output.float().flatten(),
|
||||
dim=0,
|
||||
)
|
||||
assert dense_output.dtype == torch.bfloat16
|
||||
assert quantized_output.dtype == torch.bfloat16
|
||||
assert similarity > 0.97
|
||||
assert quantized.fc_in.weight is None
|
||||
assert quantized.fc_out.weight is None
|
||||
|
||||
|
||||
def test_maybe_quantize_model_lora_defers_nvfp4(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from fastvideo.models.loader.fsdp_load import _maybe_quantize_model
|
||||
|
||||
model = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod))
|
||||
converted: list[nn.Module] = []
|
||||
monkeypatch.setattr(nvfp4, "convert_model_to_nvfp4", converted.append)
|
||||
|
||||
_maybe_quantize_model(model, defer_weight_conversion_until_lora_merge=True)
|
||||
assert converted == []
|
||||
|
||||
_maybe_quantize_model(model)
|
||||
assert converted == [model]
|
||||
|
||||
|
||||
def test_convert_quantized_weights_after_lora_merge_dispatches_nvfp4(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from fastvideo.pipelines.lora_pipeline import _convert_quantized_weights_after_lora_merge
|
||||
|
||||
model = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod))
|
||||
nvfp4_converted: list[nn.Module] = []
|
||||
mxfp8_converted: list[nn.Module] = []
|
||||
monkeypatch.setattr(nvfp4, "convert_model_to_nvfp4", nvfp4_converted.append)
|
||||
monkeypatch.setattr(mxfp8, "convert_model_to_mxfp8", mxfp8_converted.append)
|
||||
|
||||
_convert_quantized_weights_after_lora_merge({"transformer": model})
|
||||
|
||||
assert nvfp4_converted == [model]
|
||||
assert mxfp8_converted == []
|
||||
|
||||
|
||||
def test_unmerge_lora_weights_rejects_quantized_mxfp8() -> None:
|
||||
"""Reject LoRA unmerge while the MXFP8 weight still contains the merged adapter."""
|
||||
from fastvideo.pipelines.lora_pipeline import LoRAPipeline
|
||||
|
||||
transformer = _TaggedModel(object.__new__(mxfp8.MXFP8QuantizeMethod))
|
||||
transformer.linear.register_buffer("_mxfp8_weight", torch.empty(1))
|
||||
pipeline = SimpleNamespace(trainable_transformer_modules={"transformer": transformer})
|
||||
|
||||
with pytest.raises(RuntimeError, match="LoRA unmerge is unsupported after MXFP8 weight quantization"):
|
||||
LoRAPipeline.unmerge_lora_weights(pipeline)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("retain_bf16_weight", [False, True])
|
||||
def test_unmerge_lora_weights_rejects_quantized_nvfp4(retain_bf16_weight: bool) -> None:
|
||||
"""Reject NVFP4 LoRA unmerge whether the BF16 weight was removed or retained."""
|
||||
from fastvideo.pipelines.lora_pipeline import LoRAPipeline
|
||||
|
||||
transformer = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod))
|
||||
transformer.linear.register_parameter(
|
||||
"weight",
|
||||
nn.Parameter(torch.empty(1), requires_grad=False) if retain_bf16_weight else None,
|
||||
)
|
||||
transformer.linear.register_buffer("_nvfp4_weight", torch.empty(1))
|
||||
pipeline = SimpleNamespace(trainable_transformer_modules={"transformer": transformer})
|
||||
|
||||
with pytest.raises(RuntimeError, match="LoRA unmerge is unsupported after NVFP4 weight quantization"):
|
||||
LoRAPipeline.unmerge_lora_weights(pipeline)
|
||||
|
||||
|
||||
def test_set_lora_adapter_rejects_nvfp4_without_bf16_weight() -> None:
|
||||
"""Reject runtime adapter changes after NVFP4 removes the BF16 weight."""
|
||||
from fastvideo.pipelines.lora_pipeline import LoRAPipeline
|
||||
|
||||
transformer = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod))
|
||||
transformer.linear.register_parameter("weight", None)
|
||||
transformer.linear.register_buffer("_nvfp4_weight", torch.empty(1))
|
||||
pipeline = SimpleNamespace(
|
||||
_setting_constructor_adapter=False,
|
||||
cur_adapter_name="first-adapter",
|
||||
cur_adapter_path="/first-adapter",
|
||||
cur_adapter_strength=1.0,
|
||||
lora_adapter_paths={},
|
||||
trainable_transformer_modules={"transformer": transformer},
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Runtime LoRA adapter changes are unsupported after NVFP4 quantization"):
|
||||
LoRAPipeline.set_lora_adapter(pipeline, "second-adapter", "/second-adapter")
|
||||
|
||||
|
||||
def test_set_lora_adapter_mxfp8_requantizes_after_adapter_merge(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Requantize MXFP8 weights after merging a replacement LoRA adapter."""
|
||||
from fastvideo.pipelines.lora_pipeline import LoRAPipeline
|
||||
|
||||
events: list[str] = []
|
||||
|
||||
class _RecordingLoRALayer:
|
||||
|
||||
def set_lora_weights(self, lora_a: torch.Tensor, lora_b: torch.Tensor, **kwargs) -> None:
|
||||
del lora_a, lora_b, kwargs
|
||||
events.append("adapter merged")
|
||||
|
||||
transformer = _TaggedModel(object.__new__(mxfp8.MXFP8QuantizeMethod))
|
||||
transformer.linear.register_buffer("_mxfp8_weight", torch.empty(1))
|
||||
transformer_lora_layers = SimpleNamespace(
|
||||
lora_layers_by_block=lambda: iter([(None, {"linear": _RecordingLoRALayer()})])
|
||||
)
|
||||
pipeline = SimpleNamespace(
|
||||
_setting_constructor_adapter=False,
|
||||
_constructor_dense_lora_path=None,
|
||||
cur_adapter_name="first-adapter",
|
||||
cur_adapter_path=None,
|
||||
cur_adapter_strength=1.0,
|
||||
lora_adapter_paths={},
|
||||
lora_adapters={
|
||||
"second-adapter": {
|
||||
"linear.lora_A": torch.empty(1, 1),
|
||||
"linear.lora_B": torch.empty(1, 1),
|
||||
}
|
||||
},
|
||||
lora_initialized=True,
|
||||
lora_layers={"transformer": transformer_lora_layers},
|
||||
fastvideo_args=SimpleNamespace(training_mode=False),
|
||||
trainable_transformer_modules={"transformer": transformer},
|
||||
)
|
||||
monkeypatch.setattr("fastvideo.pipelines.lora_pipeline.dist.get_rank", lambda: 0)
|
||||
monkeypatch.setattr(mxfp8, "convert_model_to_mxfp8", lambda model: events.append("weights requantized"))
|
||||
|
||||
LoRAPipeline.set_lora_adapter(pipeline, "second-adapter")
|
||||
|
||||
assert events == ["adapter merged", "weights requantized"]
|
||||
@@ -188,6 +188,7 @@ dev = [ "fastvideo[lint]", "fastvideo[test]", ]
|
||||
fasth3 = [
|
||||
"flash-attn-4",
|
||||
"fastvideo-kernel==0.3.5; sys_platform == 'linux'",
|
||||
"quack-kernels==0.5.3",
|
||||
]
|
||||
|
||||
prompt-safety = [
|
||||
|
||||
Reference in New Issue
Block a user