From d496e5977731dd619cea7e207b8aafb9c2cfc0fa Mon Sep 17 00:00:00 2001 From: Uxito-Ada <414416158@qq.com> Date: Thu, 23 Jul 2026 02:45:54 +0000 Subject: [PATCH 1/8] Support Online FP8 & NF4 Weight Quantization --- docs/quantization/telefuser_fp8_deployment.md | 25 +++ docs/quantization/telefuser_nf4_deployment.md | 30 ++++ .../qwen_image_t2i_telefuser_fp8_h100.py | 151 ++++++++++++++++++ .../qwen_image_t2i_telefuser_nf4_h100.py | 151 ++++++++++++++++++ telefuser/core/base_stage.py | 21 +++ telefuser/core/config.py | 10 ++ telefuser/core/module_manager.py | 32 ++-- telefuser/models/ltx_dit.py | 27 +++- telefuser/models/qwen_image_dit.py | 34 ++++ telefuser/models/wan_video_dit.py | 33 ++++ telefuser/ops/bnb_nf4_linear.py | 99 ++++++++++++ telefuser/ops/torchao_fp8_linear.py | 116 ++++++++++++++ 12 files changed, 717 insertions(+), 12 deletions(-) create mode 100644 docs/quantization/telefuser_fp8_deployment.md create mode 100644 docs/quantization/telefuser_nf4_deployment.md create mode 100644 examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py create mode 100644 examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py create mode 100644 telefuser/ops/bnb_nf4_linear.py create mode 100644 telefuser/ops/torchao_fp8_linear.py diff --git a/docs/quantization/telefuser_fp8_deployment.md b/docs/quantization/telefuser_fp8_deployment.md new file mode 100644 index 0000000..3348b90 --- /dev/null +++ b/docs/quantization/telefuser_fp8_deployment.md @@ -0,0 +1,25 @@ +# TeleFuser FP8 deployment for Qwen-Image + +Telefuser users torchao as its backend of FP8 weight-only linear quantization. + +First, install Telefuser as [here](https://github.com/Tele-AI/TeleFuser#install). + +Then, install torchao: + +```bash +pip install torchao +``` + +Next, download `Qwen-Image-2512` model to your `TF_MODEL_ZOO_PATH` (or specify the model path to `--model_root`), and run the NF4 `Qwen-Image-2512` example: + +```bash +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ +python examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py \ + --prompt "A cat playing piano" \ + --aspect_ratio 1:1 \ + --num-inference-steps 16 \ + --seed 42 \ + --output qwen_image_fp8.png +``` + +Then generated image is saved as `qwen_image_fp8.png` at current directory. diff --git a/docs/quantization/telefuser_nf4_deployment.md b/docs/quantization/telefuser_nf4_deployment.md new file mode 100644 index 0000000..1d7d558 --- /dev/null +++ b/docs/quantization/telefuser_nf4_deployment.md @@ -0,0 +1,30 @@ +# TeleFuser NF4 deployment for Qwen-Image + +Telefuser users bitsandbytes as its backend of NF4 weight-only linear quantization. + +First, install Telefuser as [here](https://github.com/Tele-AI/TeleFuser#install). + +Then, install bitsandbytes: + +```bash +pip install -U bitsandbytes +``` + +Next, download `Qwen-Image-2512` model to your `TF_MODEL_ZOO_PATH` (or specify the model path to `--model_root`), and run the NF4 `Qwen-Image-2512` example: + +```bash +# set the below cuda versions according to your environment +export BNB_CUDA_VERSION=128 +export CUDA_HOME=/usr/local/cuda-12.8 +export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:${LD_LIBRARY_PATH:-} + +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ +python examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py \ + --prompt "A cat playing piano" \ + --aspect_ratio 1:1 \ + --num-inference-steps 16 \ + --seed 42 \ + --output qwen_image_nf4.png +``` + +Then generated image is saved as `qwen_image_nf4.png` at current directory. diff --git a/examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py b/examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py new file mode 100644 index 0000000..6617446 --- /dev/null +++ b/examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py @@ -0,0 +1,151 @@ +import os +import time + +import click +import torch + +from telefuser.core.config import ( + AttentionConfig, + AttnImplType, + QuantConfig, + QuantKernelBackend, + QuantType, + WeightOffloadType, +) +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.qwen_image import QwenImagePipeline, QwenImagePipelineConfig +from telefuser.pipelines.qwen_image.qwen_image import ASPECT_RATIO_TO_SIZE +from telefuser.utils.utils import get_example_name + + +def configure_attention_backends(): + """Avoid cuDNN SDPA plan failures in the Qwen2.5-VL text encoder.""" + if hasattr(torch.backends.cuda, "enable_cudnn_sdp"): + torch.backends.cuda.enable_cudnn_sdp(False) + if hasattr(torch.backends.cuda, "enable_flash_sdp"): + torch.backends.cuda.enable_flash_sdp(True) + if hasattr(torch.backends.cuda, "enable_math_sdp"): + torch.backends.cuda.enable_math_sdp(True) + if hasattr(torch.backends.cuda, "enable_mem_efficient_sdp"): + torch.backends.cuda.enable_mem_efficient_sdp(True) + +TF_MODEL_ZOO_PATH = os.environ.get("TF_MODEL_ZOO_PATH", "model_zoo") +PPL_CONFIG = dict( + name="qwen_image_t2i_torchao_fp8", + model_root=TF_MODEL_ZOO_PATH + "/Qwen-Image-2512", + dit_path_list=[ + "transformer/diffusion_pytorch_model-00001-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00002-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00003-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00004-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00005-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00006-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00007-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00008-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00009-of-00009.safetensors", + ], + vae_path_list=["vae/diffusion_pytorch_model.safetensors"], + text_encoder_path_list=[ + "text_encoder/model-00001-of-00004.safetensors", + "text_encoder/model-00002-of-00004.safetensors", + "text_encoder/model-00003-of-00004.safetensors", + "text_encoder/model-00004-of-00004.safetensors", + ], + tokenizer_path="tokenizer", + negative_prompt=( + "low resolution, low quality, distorted anatomy, malformed fingers, over saturated, " + "waxy skin, missing facial details, over-smoothed, AI artifacts, messy composition, " + "blurry text, distorted text" + ), + attn_impl=AttnImplType.TORCH_SDPA, + seed=42, + sample_solver="euler", + cfg_scale=1.0, + num_inference_steps=16, +) + + +def get_pipeline(model_root=PPL_CONFIG["model_root"]): + dit_paths = [os.path.join(model_root, p) for p in PPL_CONFIG["dit_path_list"]] + vae_paths = [os.path.join(model_root, p) for p in PPL_CONFIG["vae_path_list"]] + text_encoder_paths = [os.path.join(model_root, p) for p in PPL_CONFIG["text_encoder_path_list"]] + tokenizer_path = os.path.join(model_root, PPL_CONFIG["tokenizer_path"]) + + mm = ModuleManager(torch_dtype=torch.bfloat16, device="cpu") + mm.load_model(dit_paths, device="cpu", torch_dtype=torch.bfloat16) + mm.load_model(vae_paths, device="cpu", torch_dtype=torch.bfloat16) + mm.load_model(text_encoder_paths, device="cpu", torch_dtype=torch.bfloat16) + mm.load_from_huggingface(tokenizer_path, "transformers", module_name="tokenizer") + + pipeline = QwenImagePipeline(device="cuda", torch_dtype=torch.bfloat16) + pipe_config = QwenImagePipelineConfig() + pipe_config.dit_config.attention_config = AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]) + pipe_config.dit_config.offload_config.offload_type = WeightOffloadType.NO_CPU_OFFLOAD + pipe_config.dit_config.quant_config = QuantConfig( + enabled=True, + quant_type=QuantType.TORCHAO_FP8, + kernel_backend=QuantKernelBackend.TORCHAO, + ) + pipe_config.sample_solver = PPL_CONFIG["sample_solver"] + pipeline.init(mm, pipe_config) + return pipeline + + +def run( + pipeline: QwenImagePipeline, + prompt: str, + aspect_ratio: str = "1:1", + negative_prompt: str = PPL_CONFIG["negative_prompt"], + seed: int = PPL_CONFIG["seed"], + num_inference_steps: int = PPL_CONFIG["num_inference_steps"], +): + height, width = ASPECT_RATIO_TO_SIZE[aspect_ratio] + image = pipeline( + prompt, + height=height, + width=width, + negative_prompt=negative_prompt, + seed=seed, + num_inference_steps=num_inference_steps, + rand_device="cpu", + cfg_scale=PPL_CONFIG["cfg_scale"], + ) + return image + + +@click.command() +@click.option("--aspect_ratio", "-ar", default="1:1", help="Image ratio such as 1:1, 16:9", type=str) +@click.option("--prompt", default="A cat playing piano", help="Custom prompt text") +@click.option("--negative_prompt", default=PPL_CONFIG["negative_prompt"], help="Negative prompt") +@click.option("--output", default=get_example_name(__file__, "png"), help="Output image filename") +@click.option("--model_root", default=PPL_CONFIG["model_root"], help="Model root directory") +@click.option( + "--num-inference-steps", + "--num_inference_steps", + default=PPL_CONFIG["num_inference_steps"], + type=int, + help="Number of denoising steps", +) +@click.option("--seed", default=PPL_CONFIG["seed"], type=int, help="Random seed") +def main(aspect_ratio, prompt, negative_prompt, output, model_root, num_inference_steps, seed): + configure_attention_backends() + pipeline = get_pipeline(model_root) + s = time.time() + images = run( + pipeline, + prompt, + aspect_ratio, + negative_prompt=negative_prompt, + seed=seed, + num_inference_steps=num_inference_steps, + ) + print(f"pipe cost {time.time() - s} s") + for i, image in enumerate(images): + image.save(output.replace(".png", f"_{i}.png")) + del pipeline + + +if __name__ == "__main__": + main() + + diff --git a/examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py b/examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py new file mode 100644 index 0000000..06bd5ae --- /dev/null +++ b/examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py @@ -0,0 +1,151 @@ +import os +import time + +import click +import torch + +from telefuser.core.config import ( + AttentionConfig, + AttnImplType, + QuantConfig, + QuantKernelBackend, + QuantType, + WeightOffloadType, +) +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.qwen_image import QwenImagePipeline, QwenImagePipelineConfig +from telefuser.pipelines.qwen_image.qwen_image import ASPECT_RATIO_TO_SIZE +from telefuser.utils.utils import get_example_name + + +def configure_attention_backends(): + """Avoid cuDNN SDPA plan failures in the Qwen2.5-VL text encoder.""" + if hasattr(torch.backends.cuda, "enable_cudnn_sdp"): + torch.backends.cuda.enable_cudnn_sdp(False) + if hasattr(torch.backends.cuda, "enable_flash_sdp"): + torch.backends.cuda.enable_flash_sdp(True) + if hasattr(torch.backends.cuda, "enable_math_sdp"): + torch.backends.cuda.enable_math_sdp(True) + if hasattr(torch.backends.cuda, "enable_mem_efficient_sdp"): + torch.backends.cuda.enable_mem_efficient_sdp(True) + +TF_MODEL_ZOO_PATH = os.environ.get("TF_MODEL_ZOO_PATH", "model_zoo") +PPL_CONFIG = dict( + name="qwen_image_t2i_bnb_nf4", + model_root=TF_MODEL_ZOO_PATH + "/Qwen-Image-2512", + dit_path_list=[ + "transformer/diffusion_pytorch_model-00001-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00002-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00003-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00004-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00005-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00006-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00007-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00008-of-00009.safetensors", + "transformer/diffusion_pytorch_model-00009-of-00009.safetensors", + ], + vae_path_list=["vae/diffusion_pytorch_model.safetensors"], + text_encoder_path_list=[ + "text_encoder/model-00001-of-00004.safetensors", + "text_encoder/model-00002-of-00004.safetensors", + "text_encoder/model-00003-of-00004.safetensors", + "text_encoder/model-00004-of-00004.safetensors", + ], + tokenizer_path="tokenizer", + negative_prompt=( + "low resolution, low quality, distorted anatomy, malformed fingers, over saturated, " + "waxy skin, missing facial details, over-smoothed, AI artifacts, messy composition, " + "blurry text, distorted text" + ), + attn_impl=AttnImplType.TORCH_SDPA, + seed=42, + sample_solver="euler", + cfg_scale=1.0, + num_inference_steps=16, +) + + +def get_pipeline(model_root=PPL_CONFIG["model_root"]): + dit_paths = [os.path.join(model_root, p) for p in PPL_CONFIG["dit_path_list"]] + vae_paths = [os.path.join(model_root, p) for p in PPL_CONFIG["vae_path_list"]] + text_encoder_paths = [os.path.join(model_root, p) for p in PPL_CONFIG["text_encoder_path_list"]] + tokenizer_path = os.path.join(model_root, PPL_CONFIG["tokenizer_path"]) + + mm = ModuleManager(torch_dtype=torch.bfloat16, device="cpu") + mm.load_model(dit_paths, device="cpu", torch_dtype=torch.bfloat16) + mm.load_model(vae_paths, device="cpu", torch_dtype=torch.bfloat16) + mm.load_model(text_encoder_paths, device="cpu", torch_dtype=torch.bfloat16) + mm.load_from_huggingface(tokenizer_path, "transformers", module_name="tokenizer") + + pipeline = QwenImagePipeline(device="cuda", torch_dtype=torch.bfloat16) + pipe_config = QwenImagePipelineConfig() + pipe_config.dit_config.attention_config = AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]) + pipe_config.dit_config.offload_config.offload_type = WeightOffloadType.NO_CPU_OFFLOAD + pipe_config.dit_config.quant_config = QuantConfig( + enabled=True, + quant_type=QuantType.BNB_NF4, + kernel_backend=QuantKernelBackend.BITSANDBYTES, + ) + pipe_config.sample_solver = PPL_CONFIG["sample_solver"] + pipeline.init(mm, pipe_config) + return pipeline + + +def run( + pipeline: QwenImagePipeline, + prompt: str, + aspect_ratio: str = "1:1", + negative_prompt: str = PPL_CONFIG["negative_prompt"], + seed: int = PPL_CONFIG["seed"], + num_inference_steps: int = PPL_CONFIG["num_inference_steps"], +): + height, width = ASPECT_RATIO_TO_SIZE[aspect_ratio] + image = pipeline( + prompt, + height=height, + width=width, + negative_prompt=negative_prompt, + seed=seed, + num_inference_steps=num_inference_steps, + rand_device="cpu", + cfg_scale=PPL_CONFIG["cfg_scale"], + ) + return image + + +@click.command() +@click.option("--aspect_ratio", "-ar", default="1:1", help="Image ratio such as 1:1, 16:9", type=str) +@click.option("--prompt", default="A cat playing piano", help="Custom prompt text") +@click.option("--negative_prompt", default=PPL_CONFIG["negative_prompt"], help="Negative prompt") +@click.option("--output", default=get_example_name(__file__, "png"), help="Output image filename") +@click.option("--model_root", default=PPL_CONFIG["model_root"], help="Model root directory") +@click.option( + "--num-inference-steps", + "--num_inference_steps", + default=PPL_CONFIG["num_inference_steps"], + type=int, + help="Number of denoising steps", +) +@click.option("--seed", default=PPL_CONFIG["seed"], type=int, help="Random seed") +def main(aspect_ratio, prompt, negative_prompt, output, model_root, num_inference_steps, seed): + configure_attention_backends() + pipeline = get_pipeline(model_root) + s = time.time() + images = run( + pipeline, + prompt, + aspect_ratio, + negative_prompt=negative_prompt, + seed=seed, + num_inference_steps=num_inference_steps, + ) + print(f"pipe cost {time.time() - s} s") + for i, image in enumerate(images): + image.save(output.replace(".png", f"_{i}.png")) + del pipeline + + +if __name__ == "__main__": + main() + + diff --git a/telefuser/core/base_stage.py b/telefuser/core/base_stage.py index cea4470..6a578ed 100644 --- a/telefuser/core/base_stage.py +++ b/telefuser/core/base_stage.py @@ -142,6 +142,27 @@ def parallel_models(self) -> None: """Apply parallelization to models. Override in subclass.""" pass + def apply_model_quantization(self, model: Any, runtime_config: ModelRuntimeConfig | None = None) -> None: + """Apply explicit runtime quantization config to a model if requested.""" + runtime_config = runtime_config or self.model_runtime_config + quant_config = runtime_config.quant_config + if not quant_config.enabled: + return + if not hasattr(model, "enable_quant"): + logger.warning(f"Quantization requested for {self.name}, but model has no enable_quant()") + return + + from .config import QuantType + + if quant_config.quant_type in (QuantType.BNB_NF4, QuantType.TORCHAO_FP8): + if runtime_config.offload_config.offload_type != WeightOffloadType.NO_CPU_OFFLOAD: + raise ValueError("Runtime quantization requires NO_CPU_OFFLOAD because tensor must stay on device") + logger.info(f"onload {self.name} model to {self.device} before runtime quantization") + model.to(self.device) + self.onload_models_flag = True + + model.enable_quant(quant_config) + def setup_feature_cache( self, model: Any, diff --git a/telefuser/core/config.py b/telefuser/core/config.py index 04f6e99..fd5137e 100644 --- a/telefuser/core/config.py +++ b/telefuser/core/config.py @@ -330,6 +330,8 @@ class QuantType(Enum): MXFP6: Microscaling FP6 (OCP standard). MXFP4: Microscaling FP4 (OCP standard). NVFP4: NVIDIA FP4 format (Blackwell+). + BNB_NF4: bitsandbytes weight-only NF4 linear path. + TORCHAO_FP8: TorchAO dynamic-activation FP8 linear path. """ FP8 = auto() @@ -338,6 +340,8 @@ class QuantType(Enum): MXFP6 = auto() MXFP4 = auto() NVFP4 = auto() + BNB_NF4 = auto() + TORCHAO_FP8 = auto() class QuantKernelBackend(Enum): @@ -347,6 +351,8 @@ class QuantKernelBackend(Enum): TF_KERNEL = auto() # TeleFuser custom kernel VLLM = auto() # vLLM kernel CUTLASS = auto() # NVIDIA CUTLASS + TORCHAO = auto() # TorchAO quantization backends + BITSANDBYTES = auto() # bitsandbytes quantization backends @dataclass @@ -388,6 +394,10 @@ class QuantConfig: quant_type: QuantType = QuantType.FP8 kernel_backend: QuantKernelBackend = QuantKernelBackend.AUTO weight_block_size: tuple[int, int] | None = None + group_size: int = 16 + quantize_modules: tuple[str, ...] | None = None + skip_modules: tuple[str, ...] = ("head", "time_embedding", "time_projection", "patch_embedding") + keep_fp16_weight: bool = False @dataclass diff --git a/telefuser/core/module_manager.py b/telefuser/core/module_manager.py index 2bb39d4..5427a0e 100755 --- a/telefuser/core/module_manager.py +++ b/telefuser/core/module_manager.py @@ -14,6 +14,7 @@ from telefuser.utils.logging import logger from telefuser.utils.model_weight import hash_state_dict_keys, init_weights_on_device, load_state_dict +from .config import QuantConfig, QuantType from .model_registry import ModelRegistry @@ -26,7 +27,7 @@ def load_model_from_single_file( device: str | torch.device, low_cpu_mem_usage: bool = False, converter_kwargs: dict[str, Any] | None = None, - strict: bool = True, + quant_config: QuantConfig | None = None, ) -> tuple[list[str], list[nn.Module]]: """Load models from state dict with format conversion. @@ -39,7 +40,6 @@ def load_model_from_single_file( device: Target device for model low_cpu_mem_usage: If True, keep weights in CPU memory until moved to device converter_kwargs: Extra kwargs passed to state_dict_converter() - strict: Whether checkpoint keys must exactly match the model. Returns: Tuple of (model_names, loaded_models) @@ -66,8 +66,14 @@ def load_model_from_single_file( with init_weights_on_device("meta"): model = model_class(**extra_kwargs) - # Enable quantization if needed - if torch_dtype == torch.float8_e4m3fn: + # Enable quantization if needed. Keep the legacy FP8 dtype path for + # existing examples, but prefer explicit QuantConfig for new modes. + if quant_config is not None and quant_config.enabled: + if quant_config.quant_type in (QuantType.BNB_NF4, QuantType.TORCHAO_FP8): + logger.info("Deferring TorchAO quantization until runtime stage initialization") + else: + model.enable_quant(quant_config) + elif torch_dtype == torch.float8_e4m3fn: model.enable_quant(torch_dtype) if hasattr(model, "eval"): model = model.eval() @@ -78,7 +84,7 @@ def load_model_from_single_file( model_state_dict = {k: v.to("cpu").clone() for k, v in model_state_dict.items()} # Load weights and move to target device/dtype - model.load_state_dict(model_state_dict, strict=strict, assign=True) + model.load_state_dict(model_state_dict, assign=True) if torch_dtype != torch.float8_e4m3fn: model = model.to(dtype=torch_dtype) model = model.to(device) @@ -134,6 +140,7 @@ def load( torch_dtype: torch.dtype = torch.float16, low_cpu_mem_usage: bool = False, converter_kwargs: dict[str, Any] | None = None, + quant_config: QuantConfig | None = None, ) -> tuple[list[str], list[nn.Module]]: """Load model from file with automatic type detection.""" if not state_dict: @@ -152,6 +159,7 @@ def load( device, low_cpu_mem_usage, converter_kwargs, + quant_config, ) # Fall back to key-only hash @@ -167,6 +175,7 @@ def load( device, low_cpu_mem_usage, converter_kwargs, + quant_config, ) return [], [] @@ -189,9 +198,10 @@ def load_models( device: str | None = None, torch_dtype: torch.dtype | None = None, low_cpu_mem_usage: bool = False, + quant_config: QuantConfig | None = None, ) -> None: for file_path in file_path_list: - self.load_model(file_path, device, torch_dtype, low_cpu_mem_usage) + self.load_model(file_path, device, torch_dtype, low_cpu_mem_usage, quant_config=quant_config) def load_model( self, @@ -203,7 +213,7 @@ def load_model( model_class: type[nn.Module] | None = None, model_resource: str = "official", converter_kwargs: dict[str, Any] | None = None, - strict: bool = True, + quant_config: QuantConfig | None = None, ) -> None: """Load model from file path with automatic type detection. @@ -220,8 +230,6 @@ def load_model( converter_kwargs: Extra kwargs passed to model_class.state_dict_converter(). Used to provide explicit config overrides when hash-based config detection may not work (e.g., MoE distilled checkpoints). - strict: Whether checkpoint keys must exactly match the model when an - explicit ``model_class`` is supplied. """ device = device or self.device torch_dtype = torch_dtype or self.torch_dtype @@ -259,7 +267,7 @@ def load_model( device, low_cpu_mem_usage, converter_kwargs, - strict, + quant_config, ) for mn, model in zip(model_names, models): self.modules.append(model) @@ -270,7 +278,9 @@ def load_model( for detector in self.model_detectors: if detector.match(file_path, state_dict): - model_names, models = detector.load(file_path, state_dict, device, torch_dtype, low_cpu_mem_usage) + model_names, models = detector.load( + file_path, state_dict, device, torch_dtype, low_cpu_mem_usage, quant_config=quant_config + ) # Override model name if specified if name is not None: model_names = [name] * len(models) diff --git a/telefuser/models/ltx_dit.py b/telefuser/models/ltx_dit.py index a588312..443c811 100644 --- a/telefuser/models/ltx_dit.py +++ b/telefuser/models/ltx_dit.py @@ -1500,7 +1500,7 @@ def apply_cross_attention_adaln( class PixArtAlphaTextProjection(torch.nn.Module): """ Projects caption embeddings using dual linear layers. - Flow: linear_1 → activation → linear_2 + Flow: linear_1 -> activation -> linear_2 Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py """ @@ -2070,6 +2070,31 @@ def get_tp_plan(self) -> dict: raise NotImplementedError("Tensor parallelism plan is not implemented for LTXVideoTransformer yet.") def enable_quant(self, quant_type: str | torch.dtype) -> None: + from telefuser.core.config import QuantConfig, QuantType + + if isinstance(quant_type, QuantConfig) and quant_type.quant_type == QuantType.BNB_NF4: + from telefuser.ops.bnb_nf4_linear import replace_linear_layers_with_bnb_nf4 + + replaced = replace_linear_layers_with_bnb_nf4( + self.velocity_model.transformer_blocks, + compute_dtype=torch.bfloat16, + include_names=quant_type.quantize_modules, + exclude_names=quant_type.skip_modules, + ) + logger.info(f"BNB NF4 converted {replaced} Linear layers in LTX transformer blocks") + self.quant_type = quant_type.quant_type + return + if isinstance(quant_type, QuantConfig) and quant_type.quant_type == QuantType.TORCHAO_FP8: + from telefuser.ops.torchao_fp8_linear import replace_linear_layers_with_torchao_fp8 + + replaced = replace_linear_layers_with_torchao_fp8( + self.velocity_model.transformer_blocks, + include_names=quant_type.quantize_modules, + exclude_names=quant_type.skip_modules, + ) + logger.info(f"TorchAO FP8 converted {replaced} Linear layers in LTX transformer blocks") + self.quant_type = quant_type.quant_type + return # Keep a flag for downstream pipeline logic (e.g., FSDP buffer conversion). self.quant_type = quant_type diff --git a/telefuser/models/qwen_image_dit.py b/telefuser/models/qwen_image_dit.py index 4560e8f..35ba463 100644 --- a/telefuser/models/qwen_image_dit.py +++ b/telefuser/models/qwen_image_dit.py @@ -600,6 +600,40 @@ def __init__(self, num_layers: int = 60): def enable_quant(self, quant_type: str | torch.dtype): """Enable quantization for transformer blocks.""" + from telefuser.core.config import QuantConfig, QuantType + + if isinstance(quant_type, QuantConfig): + if quant_type.quant_type == QuantType.BNB_NF4: + logger.info("loading weights with BNB NF4, start quantize linear layers") + from telefuser.ops.bnb_nf4_linear import replace_linear_layers_with_bnb_nf4 + + target = self.transformer_blocks if hasattr(self, "transformer_blocks") else self.blocks + replaced = replace_linear_layers_with_bnb_nf4( + target, + compute_dtype=torch.bfloat16, + include_names=quant_type.quantize_modules, + exclude_names=quant_type.skip_modules, + ) + self.bnb_nf4_replaced_linear = replaced + logger.info(f"BNB NF4 converted {replaced} Linear layers") + self.quant_type = quant_type.quant_type + return + if quant_type.quant_type == QuantType.TORCHAO_FP8: + logger.info("loading weights with TorchAO FP8, start quantize linear layers") + from telefuser.ops.torchao_fp8_linear import replace_linear_layers_with_torchao_fp8 + + target = self.transformer_blocks if hasattr(self, "transformer_blocks") else self.blocks + replaced = replace_linear_layers_with_torchao_fp8( + target, + include_names=quant_type.quantize_modules, + exclude_names=quant_type.skip_modules, + ) + self.torchao_fp8_replaced_linear = replaced + logger.info(f"TorchAO FP8 converted {replaced} Linear layers") + self.quant_type = quant_type.quant_type + return + quant_type = torch.float8_e4m3fn if quant_type.quant_type == QuantType.FP8 else quant_type.quant_type + if quant_type in [torch.float8_e4m3fn]: logger.info(f"loading weights with {quant_type}, start convert linear layer to {quant_type}") from telefuser.ops.quantized_linear import replace_linear_layers diff --git a/telefuser/models/wan_video_dit.py b/telefuser/models/wan_video_dit.py index d116469..c87b085 100755 --- a/telefuser/models/wan_video_dit.py +++ b/telefuser/models/wan_video_dit.py @@ -506,6 +506,39 @@ def reset_y_camera_status(self): def enable_quant(self, quant_type: str | torch.dtype): """Enable quantization for transformer blocks.""" + from telefuser.core.config import QuantConfig, QuantType + + if isinstance(quant_type, QuantConfig): + if quant_type.quant_type == QuantType.BNB_NF4: + logger.info("loading weights with BNB NF4, start quantize linear layers") + from telefuser.ops.bnb_nf4_linear import replace_linear_layers_with_bnb_nf4 + + replaced = replace_linear_layers_with_bnb_nf4( + self.blocks, + compute_dtype=torch.bfloat16, + include_names=quant_type.quantize_modules, + exclude_names=quant_type.skip_modules, + ) + self.bnb_nf4_replaced_linear = replaced + logger.info(f"BNB NF4 converted {replaced} Linear layers") + self.quant_type = quant_type.quant_type + return + if quant_type.quant_type == QuantType.TORCHAO_FP8: + logger.info("loading weights with TorchAO FP8, start quantize linear layers") + from telefuser.ops.torchao_fp8_linear import replace_linear_layers_with_torchao_fp8 + + target = self.transformer_blocks if hasattr(self, "transformer_blocks") else self.blocks + replaced = replace_linear_layers_with_torchao_fp8( + target, + include_names=quant_type.quantize_modules, + exclude_names=quant_type.skip_modules, + ) + self.torchao_fp8_replaced_linear = replaced + logger.info(f"TorchAO FP8 converted {replaced} Linear layers") + self.quant_type = quant_type.quant_type + return + quant_type = torch.float8_e4m3fn if quant_type.quant_type == QuantType.FP8 else quant_type.quant_type + if quant_type in [torch.float8_e4m3fn]: logger.info(f"loading weights with {quant_type}, start convert linear layer to {quant_type}") from telefuser.ops.quantized_linear import replace_linear_layers diff --git a/telefuser/ops/bnb_nf4_linear.py b/telefuser/ops/bnb_nf4_linear.py new file mode 100644 index 0000000..1537e3a --- /dev/null +++ b/telefuser/ops/bnb_nf4_linear.py @@ -0,0 +1,99 @@ +"""bitsandbytes NF4 helpers for TeleFuser DiT linear layers.""" + +from __future__ import annotations + +from importlib import metadata +from typing import Iterable + +import torch +import torch.nn as nn + +from telefuser.utils.logging import logger + + +def _check_bnb_available() -> None: + try: + metadata.version("bitsandbytes") + except metadata.PackageNotFoundError as exc: + raise RuntimeError("BNB NF4 requires bitsandbytes to be installed") from exc + try: + import bitsandbytes as bnb # noqa: F401 + except OSError as exc: + raise RuntimeError( + "bitsandbytes failed to load its CUDA library. Check LD_LIBRARY_PATH/CUDA_HOME; " + "for CUDA 13 PyTorch wheels, libnvJitLink.so.13 must be discoverable." + ) from exc + except Exception as exc: + raise RuntimeError(f"bitsandbytes import failed: {exc}") from exc + + +def _matches_filter(name: str, include_names: Iterable[str] | None, exclude_names: Iterable[str]) -> bool: + if include_names is not None and not any(token in name for token in include_names): + return False + return not any(token and token in name for token in exclude_names) + + +def _make_bnb_nf4_linear(linear: nn.Linear, *, compute_dtype: torch.dtype, compress_statistics: bool) -> nn.Module: + import bitsandbytes as bnb + + device = linear.weight.device + new_linear = bnb.nn.Linear4bit( + linear.in_features, + linear.out_features, + bias=linear.bias is not None, + compute_dtype=compute_dtype, + compress_statistics=compress_statistics, + quant_type="nf4", + device=device, + ) + new_linear.weight = bnb.nn.Params4bit( + linear.weight.detach().contiguous(), + requires_grad=False, + compress_statistics=compress_statistics, + quant_type="nf4", + ) + if linear.bias is not None: + new_linear.bias = nn.Parameter(linear.bias.detach().to(device=device, dtype=compute_dtype), requires_grad=False) + new_linear = new_linear.to(device=device) + new_linear.requires_grad_(False) + return new_linear + + +def replace_linear_layers_with_bnb_nf4( + module: nn.Module, + *, + compute_dtype: torch.dtype = torch.bfloat16, + include_names: Iterable[str] | None = None, + exclude_names: Iterable[str] = ("head", "time_embedding", "time_projection", "patch_embedding"), + compress_statistics: bool = True, + _prefix: str = "", +) -> int: + """Replace selected ``nn.Linear`` modules with bitsandbytes NF4 Linear4bit.""" + _check_bnb_available() + replaced = 0 + for child_name, child in list(module.named_children()): + full_name = f"{_prefix}.{child_name}" if _prefix else child_name + if isinstance(child, nn.Linear): + if _matches_filter(full_name, include_names, exclude_names): + setattr( + module, + child_name, + _make_bnb_nf4_linear( + child, + compute_dtype=compute_dtype, + compress_statistics=compress_statistics, + ), + ) + replaced += 1 + continue + replaced += replace_linear_layers_with_bnb_nf4( + child, + compute_dtype=compute_dtype, + include_names=include_names, + exclude_names=exclude_names, + compress_statistics=compress_statistics, + _prefix=full_name, + ) + if not _prefix: + logger.info(f"BNB NF4 replaced {replaced} Linear layers") + return replaced diff --git a/telefuser/ops/torchao_fp8_linear.py b/telefuser/ops/torchao_fp8_linear.py new file mode 100644 index 0000000..4cf93d4 --- /dev/null +++ b/telefuser/ops/torchao_fp8_linear.py @@ -0,0 +1,116 @@ +"""TorchAO FP8 helpers for TeleFuser DiT linear layers. + +This backend applies TorchAO dynamic-activation FP8 + FP8 weight quantization +to selected ``nn.Linear`` modules. It targets W8A8 inference on Hopper/H100 +and keeps the integration close to TorchAO's native ``quantize_`` API. +""" + +from __future__ import annotations + +import inspect +from importlib import metadata +from typing import Iterable + +import torch.nn as nn + +from telefuser.utils.logging import logger + + +def _import_first_attr(module_names: tuple[str, ...], attr_names: tuple[str, ...]): + errors: list[str] = [] + for module_name in module_names: + try: + module = __import__(module_name, fromlist=list(attr_names)) + except Exception as exc: # pragma: no cover - diagnostic path + errors.append(f"{module_name}: {type(exc).__name__}: {exc}") + continue + for attr_name in attr_names: + if hasattr(module, attr_name): + return getattr(module, attr_name) + raise ImportError("; ".join(errors) if errors else f"none of {attr_names} found") + + +def _instantiate_config(config_cls, **kwargs): + try: + signature = inspect.signature(config_cls) + accepted = {k: v for k, v in kwargs.items() if k in signature.parameters} + except (TypeError, ValueError): + accepted = kwargs + try: + return config_cls(**accepted) + except TypeError: + return config_cls() + + +def _check_torchao_fp8_available() -> None: + try: + metadata.version("torchao") + except metadata.PackageNotFoundError as exc: + raise RuntimeError("TorchAO FP8 requires torchao to be installed") from exc + + +def _matches_filter(name: str, include_names: Iterable[str] | None, exclude_names: Iterable[str]) -> bool: + if include_names is not None and not any(token in name for token in include_names): + return False + return not any(token and token in name for token in exclude_names) + + +def _count_matching_linear_layers( + module: nn.Module, + *, + include_names: Iterable[str] | None, + exclude_names: Iterable[str], + _prefix: str = "", +) -> int: + count = 0 + for child_name, child in module.named_children(): + full_name = f"{_prefix}.{child_name}" if _prefix else child_name + if isinstance(child, nn.Linear): + if _matches_filter(full_name, include_names, exclude_names): + count += 1 + continue + count += _count_matching_linear_layers( + child, + include_names=include_names, + exclude_names=exclude_names, + _prefix=full_name, + ) + return count + + +def replace_linear_layers_with_torchao_fp8( + module: nn.Module, + *, + include_names: Iterable[str] | None = None, + exclude_names: Iterable[str] = ("head", "time_embedding", "time_projection", "patch_embedding"), +) -> int: + """Quantize selected ``nn.Linear`` modules with TorchAO FP8. + + Returns the number of selected Linear layers. TorchAO performs in-place + conversion through ``quantize_``. + """ + _check_torchao_fp8_available() + + try: + from torchao.quantization import quantize_ + except ImportError as exc: + raise RuntimeError("TorchAO FP8 requires torchao.quantization.quantize_") from exc + + fp8_api = _import_first_attr( + ("torchao.quantization", "torchao.quantization.quant_api", "torchao.float8"), + ( + "float8_dynamic_activation_float8_weight", + "Float8DynamicActivationFloat8WeightConfig", + "Float8WeightOnlyConfig", + ), + ) + quant_config = fp8_api() if not inspect.isclass(fp8_api) else _instantiate_config(fp8_api) + selected = _count_matching_linear_layers(module, include_names=include_names, exclude_names=exclude_names) + + def filter_fn(target: nn.Module, *args) -> bool: + name = next((str(arg) for arg in args if isinstance(arg, str)), "") + return isinstance(target, nn.Linear) and _matches_filter(name, include_names, exclude_names) + + quantize_(module, quant_config, filter_fn=filter_fn) + logger.info(f"TorchAO FP8 quantized {selected} Linear layers") + return selected From 48b855a5f22d71fc461a8b5d5ff903f9d57d0558 Mon Sep 17 00:00:00 2001 From: Uxito-Ada <414416158@qq.com> Date: Thu, 23 Jul 2026 05:54:32 +0000 Subject: [PATCH 2/8] refactor denoise --- telefuser/pipelines/ltx_video/dit_denoising.py | 2 ++ telefuser/pipelines/qwen_image/dit_denoising.py | 1 + telefuser/pipelines/wan_video/single_dit_denoising.py | 1 + 3 files changed, 4 insertions(+) diff --git a/telefuser/pipelines/ltx_video/dit_denoising.py b/telefuser/pipelines/ltx_video/dit_denoising.py index 4b306dc..016eed8 100644 --- a/telefuser/pipelines/ltx_video/dit_denoising.py +++ b/telefuser/pipelines/ltx_video/dit_denoising.py @@ -560,6 +560,7 @@ def __init__( self.dit: LTXVideoTransformer = module_manager.fetch_module(model_name) if self.dit is not None and hasattr(self.dit, "set_attention_config"): self.dit.set_attention_config(model_runtime_config.attention_config) + self.apply_model_quantization(self.dit) self.model_names = ["dit"] self.scheduler = scheduler self.flow_match_scheduler = FlowMatchScheduler(template="LTX.2") @@ -1008,3 +1009,4 @@ def process( def stage2_sigmas(device: torch.device) -> torch.Tensor: """Return the fixed distilled sigma schedule used by stage 2.""" return torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES, device=device, dtype=torch.float32) + diff --git a/telefuser/pipelines/qwen_image/dit_denoising.py b/telefuser/pipelines/qwen_image/dit_denoising.py index da7d854..bc6e435 100644 --- a/telefuser/pipelines/qwen_image/dit_denoising.py +++ b/telefuser/pipelines/qwen_image/dit_denoising.py @@ -37,6 +37,7 @@ def __init__( super().__init__(name, model_runtime_config) self.dit: QwenImageDiT = module_manager.fetch_module("qwen_image_dit") self.dit.set_attention_config(model_runtime_config.attention_config) + self.apply_model_quantization(self.dit) if self.model_runtime_config.offload_config.offload_type == WeightOffloadType.ASYNC_CPU_OFFLOAD: self.dit.enable_async_offload(self.device, offload_config=self.model_runtime_config.offload_config) diff --git a/telefuser/pipelines/wan_video/single_dit_denoising.py b/telefuser/pipelines/wan_video/single_dit_denoising.py index 829fcb9..108c03b 100644 --- a/telefuser/pipelines/wan_video/single_dit_denoising.py +++ b/telefuser/pipelines/wan_video/single_dit_denoising.py @@ -39,6 +39,7 @@ def __init__( self.dit: WanModel = module_manager.fetch_module("wan_video_dit") self.dit.set_attention_config(model_runtime_config.attention_config) self.load_loras() + self.apply_model_quantization(self.dit) self.model_names = ["dit"] self.scheduler = scheduler if model_runtime_config.offload_config.offload_type == WeightOffloadType.SEQUENTIAL_CPU_OFFLOAD: From bd34b813608ce84b2b93aa56e1482d9f0ecfcd36 Mon Sep 17 00:00:00 2001 From: Uxito-Ada <414416158@qq.com> Date: Thu, 23 Jul 2026 06:23:06 +0000 Subject: [PATCH 3/8] refactor quant happen place --- telefuser/core/base_stage.py | 21 ----------------- telefuser/core/module_manager.py | 23 ++++++++++--------- .../pipelines/ltx_video/dit_denoising.py | 1 - .../pipelines/qwen_image/dit_denoising.py | 1 - .../wan_video/single_dit_denoising.py | 1 - 5 files changed, 12 insertions(+), 35 deletions(-) mode change 100755 => 100644 telefuser/core/module_manager.py diff --git a/telefuser/core/base_stage.py b/telefuser/core/base_stage.py index 6a578ed..cea4470 100644 --- a/telefuser/core/base_stage.py +++ b/telefuser/core/base_stage.py @@ -142,27 +142,6 @@ def parallel_models(self) -> None: """Apply parallelization to models. Override in subclass.""" pass - def apply_model_quantization(self, model: Any, runtime_config: ModelRuntimeConfig | None = None) -> None: - """Apply explicit runtime quantization config to a model if requested.""" - runtime_config = runtime_config or self.model_runtime_config - quant_config = runtime_config.quant_config - if not quant_config.enabled: - return - if not hasattr(model, "enable_quant"): - logger.warning(f"Quantization requested for {self.name}, but model has no enable_quant()") - return - - from .config import QuantType - - if quant_config.quant_type in (QuantType.BNB_NF4, QuantType.TORCHAO_FP8): - if runtime_config.offload_config.offload_type != WeightOffloadType.NO_CPU_OFFLOAD: - raise ValueError("Runtime quantization requires NO_CPU_OFFLOAD because tensor must stay on device") - logger.info(f"onload {self.name} model to {self.device} before runtime quantization") - model.to(self.device) - self.onload_models_flag = True - - model.enable_quant(quant_config) - def setup_feature_cache( self, model: Any, diff --git a/telefuser/core/module_manager.py b/telefuser/core/module_manager.py old mode 100755 new mode 100644 index 5427a0e..829c848 --- a/telefuser/core/module_manager.py +++ b/telefuser/core/module_manager.py @@ -14,7 +14,7 @@ from telefuser.utils.logging import logger from telefuser.utils.model_weight import hash_state_dict_keys, init_weights_on_device, load_state_dict -from .config import QuantConfig, QuantType +from .config import QuantConfig from .model_registry import ModelRegistry @@ -66,15 +66,6 @@ def load_model_from_single_file( with init_weights_on_device("meta"): model = model_class(**extra_kwargs) - # Enable quantization if needed. Keep the legacy FP8 dtype path for - # existing examples, but prefer explicit QuantConfig for new modes. - if quant_config is not None and quant_config.enabled: - if quant_config.quant_type in (QuantType.BNB_NF4, QuantType.TORCHAO_FP8): - logger.info("Deferring TorchAO quantization until runtime stage initialization") - else: - model.enable_quant(quant_config) - elif torch_dtype == torch.float8_e4m3fn: - model.enable_quant(torch_dtype) if hasattr(model, "eval"): model = model.eval() model.requires_grad_(False) @@ -83,12 +74,22 @@ def load_model_from_single_file( if not low_cpu_mem_usage: model_state_dict = {k: v.to("cpu").clone() for k, v in model_state_dict.items()} - # Load weights and move to target device/dtype + # Load weights and move to target device/dtype first. Runtime quantization + # backends such as TorchAO FP8 and BNB NF4 need real assigned weights on + # the target device before replacing Linear layers. model.load_state_dict(model_state_dict, assign=True) if torch_dtype != torch.float8_e4m3fn: model = model.to(dtype=torch_dtype) model = model.to(device) + # Enable quantization after weights are assigned and the model is on device. + # This preserves the original ModuleManager ownership of quantization while + # allowing online backends to inspect/convert real parameters. + if quant_config is not None and quant_config.enabled: + model.enable_quant(quant_config) + elif torch_dtype == torch.float8_e4m3fn: + model.enable_quant(torch_dtype) + loaded_model_names.append(model_name) loaded_models.append(model) diff --git a/telefuser/pipelines/ltx_video/dit_denoising.py b/telefuser/pipelines/ltx_video/dit_denoising.py index 016eed8..ec2ee88 100644 --- a/telefuser/pipelines/ltx_video/dit_denoising.py +++ b/telefuser/pipelines/ltx_video/dit_denoising.py @@ -560,7 +560,6 @@ def __init__( self.dit: LTXVideoTransformer = module_manager.fetch_module(model_name) if self.dit is not None and hasattr(self.dit, "set_attention_config"): self.dit.set_attention_config(model_runtime_config.attention_config) - self.apply_model_quantization(self.dit) self.model_names = ["dit"] self.scheduler = scheduler self.flow_match_scheduler = FlowMatchScheduler(template="LTX.2") diff --git a/telefuser/pipelines/qwen_image/dit_denoising.py b/telefuser/pipelines/qwen_image/dit_denoising.py index bc6e435..da7d854 100644 --- a/telefuser/pipelines/qwen_image/dit_denoising.py +++ b/telefuser/pipelines/qwen_image/dit_denoising.py @@ -37,7 +37,6 @@ def __init__( super().__init__(name, model_runtime_config) self.dit: QwenImageDiT = module_manager.fetch_module("qwen_image_dit") self.dit.set_attention_config(model_runtime_config.attention_config) - self.apply_model_quantization(self.dit) if self.model_runtime_config.offload_config.offload_type == WeightOffloadType.ASYNC_CPU_OFFLOAD: self.dit.enable_async_offload(self.device, offload_config=self.model_runtime_config.offload_config) diff --git a/telefuser/pipelines/wan_video/single_dit_denoising.py b/telefuser/pipelines/wan_video/single_dit_denoising.py index 108c03b..829fcb9 100644 --- a/telefuser/pipelines/wan_video/single_dit_denoising.py +++ b/telefuser/pipelines/wan_video/single_dit_denoising.py @@ -39,7 +39,6 @@ def __init__( self.dit: WanModel = module_manager.fetch_module("wan_video_dit") self.dit.set_attention_config(model_runtime_config.attention_config) self.load_loras() - self.apply_model_quantization(self.dit) self.model_names = ["dit"] self.scheduler = scheduler if model_runtime_config.offload_config.offload_type == WeightOffloadType.SEQUENTIAL_CPU_OFFLOAD: From 858c974ccc85635ae6bebca0e02bea3bd5bda999 Mon Sep 17 00:00:00 2001 From: Uxito-Ada <414416158@qq.com> Date: Thu, 23 Jul 2026 07:55:28 +0000 Subject: [PATCH 4/8] refactor quant entrypoint --- .../qwen_image/qwen_image_t2i_telefuser_fp8_h100.py | 12 ++++++------ .../qwen_image/qwen_image_t2i_telefuser_nf4_h100.py | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py b/examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py index 6617446..b779377 100644 --- a/examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py +++ b/examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py @@ -71,8 +71,13 @@ def get_pipeline(model_root=PPL_CONFIG["model_root"]): text_encoder_paths = [os.path.join(model_root, p) for p in PPL_CONFIG["text_encoder_path_list"]] tokenizer_path = os.path.join(model_root, PPL_CONFIG["tokenizer_path"]) + quant_config = QuantConfig( + enabled=True, + quant_type=QuantType.TORCHAO_FP8, + kernel_backend=QuantKernelBackend.TORCHAO, + ) mm = ModuleManager(torch_dtype=torch.bfloat16, device="cpu") - mm.load_model(dit_paths, device="cpu", torch_dtype=torch.bfloat16) + mm.load_model(dit_paths, device="cuda", torch_dtype=torch.bfloat16, quant_config=quant_config) mm.load_model(vae_paths, device="cpu", torch_dtype=torch.bfloat16) mm.load_model(text_encoder_paths, device="cpu", torch_dtype=torch.bfloat16) mm.load_from_huggingface(tokenizer_path, "transformers", module_name="tokenizer") @@ -81,11 +86,6 @@ def get_pipeline(model_root=PPL_CONFIG["model_root"]): pipe_config = QwenImagePipelineConfig() pipe_config.dit_config.attention_config = AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]) pipe_config.dit_config.offload_config.offload_type = WeightOffloadType.NO_CPU_OFFLOAD - pipe_config.dit_config.quant_config = QuantConfig( - enabled=True, - quant_type=QuantType.TORCHAO_FP8, - kernel_backend=QuantKernelBackend.TORCHAO, - ) pipe_config.sample_solver = PPL_CONFIG["sample_solver"] pipeline.init(mm, pipe_config) return pipeline diff --git a/examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py b/examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py index 06bd5ae..f50c475 100644 --- a/examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py +++ b/examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py @@ -71,8 +71,13 @@ def get_pipeline(model_root=PPL_CONFIG["model_root"]): text_encoder_paths = [os.path.join(model_root, p) for p in PPL_CONFIG["text_encoder_path_list"]] tokenizer_path = os.path.join(model_root, PPL_CONFIG["tokenizer_path"]) + quant_config = QuantConfig( + enabled=True, + quant_type=QuantType.BNB_NF4, + kernel_backend=QuantKernelBackend.BITSANDBYTES, + ) mm = ModuleManager(torch_dtype=torch.bfloat16, device="cpu") - mm.load_model(dit_paths, device="cpu", torch_dtype=torch.bfloat16) + mm.load_model(dit_paths, device="cuda", torch_dtype=torch.bfloat16, quant_config=quant_config) mm.load_model(vae_paths, device="cpu", torch_dtype=torch.bfloat16) mm.load_model(text_encoder_paths, device="cpu", torch_dtype=torch.bfloat16) mm.load_from_huggingface(tokenizer_path, "transformers", module_name="tokenizer") @@ -81,11 +86,6 @@ def get_pipeline(model_root=PPL_CONFIG["model_root"]): pipe_config = QwenImagePipelineConfig() pipe_config.dit_config.attention_config = AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]) pipe_config.dit_config.offload_config.offload_type = WeightOffloadType.NO_CPU_OFFLOAD - pipe_config.dit_config.quant_config = QuantConfig( - enabled=True, - quant_type=QuantType.BNB_NF4, - kernel_backend=QuantKernelBackend.BITSANDBYTES, - ) pipe_config.sample_solver = PPL_CONFIG["sample_solver"] pipeline.init(mm, pipe_config) return pipeline From da1bff165fba3f4d03c3e8b9eff9701412560dba Mon Sep 17 00:00:00 2001 From: Uxito-Ada <414416158@qq.com> Date: Thu, 23 Jul 2026 08:04:07 +0000 Subject: [PATCH 5/8] format --- telefuser/pipelines/ltx_video/dit_denoising.py | 1 - 1 file changed, 1 deletion(-) diff --git a/telefuser/pipelines/ltx_video/dit_denoising.py b/telefuser/pipelines/ltx_video/dit_denoising.py index ec2ee88..4b306dc 100644 --- a/telefuser/pipelines/ltx_video/dit_denoising.py +++ b/telefuser/pipelines/ltx_video/dit_denoising.py @@ -1008,4 +1008,3 @@ def process( def stage2_sigmas(device: torch.device) -> torch.Tensor: """Return the fixed distilled sigma schedule used by stage 2.""" return torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES, device=device, dtype=torch.float32) - From 3f0ba6561fe0d2bea3542fcd8b66ae291963ca07 Mon Sep 17 00:00:00 2001 From: Heyang Sun <60865256+Uxito-Ada@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:05:37 +0800 Subject: [PATCH 6/8] format From 5d1744c157a5ed6ba4e7b24c5e59767214d9f4b9 Mon Sep 17 00:00:00 2001 From: Uxito-Ada <414416158@qq.com> Date: Thu, 23 Jul 2026 08:35:11 +0000 Subject: [PATCH 7/8] move dependencies to install --- docs/quantization/telefuser_fp8_deployment.md | 6 ------ docs/quantization/telefuser_nf4_deployment.md | 6 ------ pyproject.toml | 2 ++ 3 files changed, 2 insertions(+), 12 deletions(-) diff --git a/docs/quantization/telefuser_fp8_deployment.md b/docs/quantization/telefuser_fp8_deployment.md index 3348b90..210eee1 100644 --- a/docs/quantization/telefuser_fp8_deployment.md +++ b/docs/quantization/telefuser_fp8_deployment.md @@ -4,12 +4,6 @@ Telefuser users torchao as its backend of FP8 weight-only linear quantization. First, install Telefuser as [here](https://github.com/Tele-AI/TeleFuser#install). -Then, install torchao: - -```bash -pip install torchao -``` - Next, download `Qwen-Image-2512` model to your `TF_MODEL_ZOO_PATH` (or specify the model path to `--model_root`), and run the NF4 `Qwen-Image-2512` example: ```bash diff --git a/docs/quantization/telefuser_nf4_deployment.md b/docs/quantization/telefuser_nf4_deployment.md index 1d7d558..aa670dc 100644 --- a/docs/quantization/telefuser_nf4_deployment.md +++ b/docs/quantization/telefuser_nf4_deployment.md @@ -4,12 +4,6 @@ Telefuser users bitsandbytes as its backend of NF4 weight-only linear quantizati First, install Telefuser as [here](https://github.com/Tele-AI/TeleFuser#install). -Then, install bitsandbytes: - -```bash -pip install -U bitsandbytes -``` - Next, download `Qwen-Image-2512` model to your `TF_MODEL_ZOO_PATH` (or specify the model path to `--model_root`), and run the NF4 `Qwen-Image-2512` example: ```bash diff --git a/pyproject.toml b/pyproject.toml index 5bb520d..cde900a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,8 @@ dependencies = [ "transformers==4.57.3", "triton", "uvicorn>=0.30,<1.0", + "torchao", + "bitsandbytes", ] [project.optional-dependencies] From 1858ead617c00dad1bf895c55de8522b8acc053b Mon Sep 17 00:00:00 2001 From: Uxito-Ada <414416158@qq.com> Date: Fri, 24 Jul 2026 02:48:45 +0000 Subject: [PATCH 8/8] refactor doc --- .../quantization.md} | 24 +++++++++- docs/quantization/telefuser_fp8_deployment.md | 19 -------- docs/zh/quantization.md | 47 +++++++++++++++++++ 3 files changed, 70 insertions(+), 20 deletions(-) rename docs/{quantization/telefuser_nf4_deployment.md => en/quantization.md} (53%) delete mode 100644 docs/quantization/telefuser_fp8_deployment.md create mode 100644 docs/zh/quantization.md diff --git a/docs/quantization/telefuser_nf4_deployment.md b/docs/en/quantization.md similarity index 53% rename from docs/quantization/telefuser_nf4_deployment.md rename to docs/en/quantization.md index aa670dc..3da4b2f 100644 --- a/docs/quantization/telefuser_nf4_deployment.md +++ b/docs/en/quantization.md @@ -1,4 +1,26 @@ -# TeleFuser NF4 deployment for Qwen-Image +# Quantizaiton + +## TeleFuser FP8 deployment for Qwen-Image + +Telefuser users torchao as its backend of FP8 weight-only linear quantization. + +First, install Telefuser as [here](https://github.com/Tele-AI/TeleFuser#install). + +Next, download `Qwen-Image-2512` model to your `TF_MODEL_ZOO_PATH` (or specify the model path to `--model_root`), and run the NF4 `Qwen-Image-2512` example: + +```bash +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ +python examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py \ + --prompt "A cat playing piano" \ + --aspect_ratio 1:1 \ + --num-inference-steps 16 \ + --seed 42 \ + --output qwen_image_fp8.png +``` + +Then generated image is saved as `qwen_image_fp8.png` at current directory. + +## TeleFuser NF4 deployment for Qwen-Image Telefuser users bitsandbytes as its backend of NF4 weight-only linear quantization. diff --git a/docs/quantization/telefuser_fp8_deployment.md b/docs/quantization/telefuser_fp8_deployment.md deleted file mode 100644 index 210eee1..0000000 --- a/docs/quantization/telefuser_fp8_deployment.md +++ /dev/null @@ -1,19 +0,0 @@ -# TeleFuser FP8 deployment for Qwen-Image - -Telefuser users torchao as its backend of FP8 weight-only linear quantization. - -First, install Telefuser as [here](https://github.com/Tele-AI/TeleFuser#install). - -Next, download `Qwen-Image-2512` model to your `TF_MODEL_ZOO_PATH` (or specify the model path to `--model_root`), and run the NF4 `Qwen-Image-2512` example: - -```bash -PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ -python examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py \ - --prompt "A cat playing piano" \ - --aspect_ratio 1:1 \ - --num-inference-steps 16 \ - --seed 42 \ - --output qwen_image_fp8.png -``` - -Then generated image is saved as `qwen_image_fp8.png` at current directory. diff --git a/docs/zh/quantization.md b/docs/zh/quantization.md new file mode 100644 index 0000000..5f1591e --- /dev/null +++ b/docs/zh/quantization.md @@ -0,0 +1,47 @@ +# 量化 + +## 使用 TeleFuser 为 Qwen-Image 部署 FP8 量化 + +TeleFuser 使用 TorchAO 作为 FP8 仅权重线性量化的后端。 + +首先，按照[此处说明](https://github.com/Tele-AI/TeleFuser#install)安装 TeleFuser。 + +接下来，将 `Qwen-Image-2512` 模型下载到 `TF_MODEL_ZOO_PATH` 目录中（也可以通过 `--model_root` 指定模型路径），然后运行 FP8 版本的 `Qwen-Image-2512` 示例： + +```bash +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ +python examples/qwen_image/qwen_image_t2i_telefuser_fp8_h100.py \ + --prompt "A cat playing piano" \ + --aspect_ratio 1:1 \ + --num-inference-steps 16 \ + --seed 42 \ + --output qwen_image_fp8.png +``` + +生成的图像将以 `qwen_image_fp8.png` 为文件名保存在当前目录中。 + +## 使用 TeleFuser 为 Qwen-Image 部署 NF4 量化 + +TeleFuser 使用 bitsandbytes 作为 NF4 仅权重线性量化的后端。 + +首先，按照[此处说明](https://github.com/Tele-AI/TeleFuser#install)安装 TeleFuser。 + +接下来，将 `Qwen-Image-2512` 模型下载到 `TF_MODEL_ZOO_PATH` 目录中（也可以通过 `--model_root` 指定模型路径），然后运行 NF4 版本的 `Qwen-Image-2512` 示例： + +```bash +# 请根据你的运行环境设置以下 CUDA 版本 +export BNB_CUDA_VERSION=128 +export CUDA_HOME=/usr/local/cuda-12.8 +export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:${LD_LIBRARY_PATH:-} + +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ +python examples/qwen_image/qwen_image_t2i_telefuser_nf4_h100.py \ + --prompt "A cat playing piano" \ + --aspect_ratio 1:1 \ + --num-inference-steps 16 \ + --seed 42 \ + --output qwen_image_nf4.png +``` + +生成的图像将以 `qwen_image_nf4.png` 为文件名保存在当前目录中。 +