Speed up quantization processing on H3

extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py

@@ -921,10 +921,12 @@ class MinimaxH3Model(BaseModel):
         return ["blocks"]
 
     def get_quantization_exclude_modules(self) -> Optional[List[str]]:
-        # float32 islands, the conditioning projection, and the token refiner
+        # float32 islands, the conditioning projection, the token refiner and
-        # (shipped bf16 in the pre-quantized checkpoints — excluding it makes
+        # the AdaLN projections — all shipped unquantized in the pre-quantized
-        # quantize with the checkpoint's own qtype an exact no-op). The
+        # checkpoints (pruned files carry tiny fp16 adaln linears fed by the
-        # per-block adaln_proj stays quantizable (half the parameter count).
+        # 8-dim time table), so excluding them makes quantize with the
+        # checkpoint's own qtype an exact no-op and keeps the sensitive
+        # modulation path at full precision under any other qtype.
         return [
             "video_patch_proj*",
             "audio_patch_proj*",
@@ -932,6 +934,7 @@ class MinimaxH3Model(BaseModel):
             "final_layer*",
             "condition_proj*",
             "token_refiner*",
+            "*adaln_proj*",
         ]
 
     def convert_lora_weights_before_save(self, state_dict):

toolkit/util/quantize.py

@@ -233,6 +233,34 @@ def quantize(
             # raise e
 
 
+def _has_quantizable_linear(module: torch.nn.Module, weights, exclude=None) -> bool:
+    """Whether quantizing ``module`` with ``weights`` would change anything.
+
+    False when every non-excluded linear is already quantized with the same
+    ostris qtype (or cannot be quantized at all) — the pre-quantized-checkpoint
+    case, where the whole block can be skipped without the device round-trip.
+    ``exclude`` patterns are matched against module names relative to
+    ``module`` (use leading wildcards for patterns aimed at inner layers)."""
+    if not isinstance(weights, ostristype):
+        return True
+    for name, m in module.named_modules():
+        if not isinstance(m, torch.nn.Linear):
+            continue
+        if exclude is not None and any(fnmatch(name, pattern) for pattern in exclude):
+            continue
+        if isinstance(m, OstrisLinear):
+            if getattr(
+                m.ostris_quantizer, "qtype", None
+            ) != weights.quantizer.qtype and weights.quantizer.can_quantize(m):
+                return True
+            continue
+        if m.__class__.__name__ in Q_MODULES:
+            continue
+        if weights.quantizer.can_quantize(m):
+            return True
+    return False
+
+
 def quantize_model(
     base_model: "BaseModel",
     model_to_quantize: torch.nn.Module,
@@ -417,15 +445,28 @@ def quantize_model(
         base_model.print_and_status_update(
             f" - quantizing {len(all_blocks)} transformer blocks"
         )
+        already_quantized = 0
         for block in tqdm(all_blocks):
+            if not _has_quantizable_linear(block, quantization_type, exclude_modules):
+                # pre-quantized checkpoint with a matching qtype: nothing in
+                # this block would change — skip the device round-trip and the
+                # dtype cast entirely so the load stays byte-identical
+                already_quantized += 1
+                continue
             block.to(base_model.device_torch, dtype=base_model.torch_dtype, non_blocking=True)
-            quantize(block, weights=quantization_type)
+            # exclude patterns with a leading wildcard (e.g. "*adaln_proj*")
+            # also apply inside blocks, where names are block-relative
+            quantize(block, weights=quantization_type, exclude=exclude_modules)
             freeze(block)
             # NOT non_blocking: an async D2H allocates the cpu destination in pinned
             # memory, which the caching host allocator keeps forever (with power-of-2
             # bucket rounding on top) — that silently retained a model-sized chunk of
             # host ram after the weights moved back to the gpu for training
             block.to("cpu")
+        if already_quantized:
+            base_model.print_and_status_update(
+                f" - {already_quantized} blocks already quantized with a matching qtype; left untouched"
+            )
 
         # todo, on extras find a universal way to quantize them on device and move them back to their original
         # device without having to move the transformer blocks to the device first