Automated PR - 2026-01-13

packages/ltx-trainer/src/ltx_trainer/quantization.py

@@ -7,7 +7,6 @@ from optimum.quanto import qtype
 from ltx_trainer import logger
 
 QuantizationOptions = Literal[
-    "no_change",
     "int8-quanto",
     "int4-quanto",
     "int2-quanto",
@@ -30,37 +29,40 @@ def quantize_model(
     Returns:
         The quantized model, or the original model if no quantization is performed.
     """
-    if precision is None or precision == "no_change":
-        return model
-
     from optimum.quanto import freeze, quantize  # noqa: PLC0415
 
     weight_quant = _quanto_type_map(precision)
     extra_quanto_args = {
         "exclude": [
-            "proj_in",
-            "time_embed.*",
-            "caption_projection.*",
-            "rope",
-            "*norm*",
+            # Input/output projection layers
+            "patchify_proj",
+            "audio_patchify_proj",
             "proj_out",
+            "audio_proj_out",
+            # Timestep embedding layers - int4 tinygemm requires strict bfloat16 input
+            # and these receive float32 sinusoidal embeddings that are cast to bfloat16
+            "*adaln*",
+            "time_proj",
+            "timestep_embedder*",
+            # Caption/text projection layers
+            "caption_projection*",
+            "audio_caption_projection*",
+            # Normalization layers (usually excluded from quantization)
+            "*norm*",
         ]
     }
     if quantize_activations:
-        logger.info("Freezing model weights and activations")
+        logger.debug("Quantizing model weights and activations")
         extra_quanto_args["activations"] = weight_quant
     else:
-        logger.info("Freezing model weights only")
+        logger.debug("Quantizing model weights only")
 
     quantize(model, weights=weight_quant, **extra_quanto_args)
     freeze(model)
     return model
 
 
-def _quanto_type_map(precision: QuantizationOptions) -> torch.dtype | qtype | None:  # noqa: PLR0911
-    if precision == "no_change":
-        return None
-
+def _quanto_type_map(precision: QuantizationOptions) -> torch.dtype | qtype | None:
     from optimum.quanto import (  # noqa: PLC0415
         qfloat8,
         qfloat8_e4m3fnuz,

packages/ltx-trainer/src/ltx_trainer/trainer.py

@@ -344,7 +344,7 @@ class LtxvTrainer:
         logger.debug("Loading text encoder...")
         if self._config.acceleration.load_text_encoder_in_8bit:
             logger.warning(
-                "⚠️  load_text_encoder_in_8bit is set to True but 8-bit text encoder loading "
+                "⚠️ load_text_encoder_in_8bit is set to True but 8-bit text encoder loading "
                 "is not currently implemented. The text encoder will be loaded in bfloat16 precision."
             )
 
@@ -428,7 +428,7 @@ class LtxvTrainer:
             if self._config.model.training_mode == "full":
                 raise ValueError("Quantization is not supported in full training mode.")
 
-            logger.warning(f"Quantizing model with precision: {self._config.acceleration.quantization}")
+            logger.info(f'Quantizing model with "{self._config.acceleration.quantization}". This may take a while...')
             self._transformer = quantize_model(
                 self._transformer,
                 precision=self._config.acceleration.quantization,