diff --git a/docs/en/quantization.md b/docs/en/quantization.md new file mode 100644 index 0000000..3da4b2f --- /dev/null +++ b/docs/en/quantization.md @@ -0,0 +1,46 @@ +# 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. + +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 +# 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/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` 为文件名保存在当前目录中。 + 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..b779377 --- /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"]) + + 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="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") + + 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.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..f50c475 --- /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"]) + + 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="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") + + 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.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/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] 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 old mode 100755 new mode 100644 index 2bb39d4..829c848 --- 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 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,9 +66,6 @@ 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: - model.enable_quant(torch_dtype) if hasattr(model, "eval"): model = model.eval() model.requires_grad_(False) @@ -77,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 - model.load_state_dict(model_state_dict, strict=strict, assign=True) + # 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) @@ -134,6 +141,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 +160,7 @@ def load( device, low_cpu_mem_usage, converter_kwargs, + quant_config, ) # Fall back to key-only hash @@ -167,6 +176,7 @@ def load( device, low_cpu_mem_usage, converter_kwargs, + quant_config, ) return [], [] @@ -189,9 +199,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 +214,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 +231,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 +268,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 +279,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