MiniMax-H3-MLX-8bit / PATCHES.md
MrMofer's picture
feat: i2v vision tower + processor + PATCHES (update pipeline)
166a35f verified
|
Raw
History Blame Contribute Delete
15.6 kB

PATCHES.md β€” MiniMax-H3 MLX 8-bit pipeline modifications (complete port)

All modifications relative to upstream PipeNetwork/minimax-h3-mlx base commit b2f7e4d2b7861cefe68b75e4b59ab81cc4e7c318 (2026-08-10). The first committed patch is video-lab-8bit-fixes 7210b93e6df86bf9c7206091c9542b6983c10c30 (2026-08-17, two files); fixes S8/S9/S10 extend the same branch in the working tree (dit.py two-phase, pipeline release + TeaCache hook, text_encoder processor+scatter, teacache.py, generate.py flags, processor/ + vision shards). Everything else is byte-identical to the base commit.

Base commit b2f7e4d2b7861cefe68b75e4b59ab81cc4e7c318
Branch video-lab-8bit-fixes (+ working-tree S8/S9/S10 on video-lab)
Patch commit (committed) 7210b93e6df86bf9c7206091c9542b6983c10c30
Date (committed) 2026-08-17
Working-tree S8/S9/S10 2026-08-18–2026-08-19 (pipeline.py, dit.py extension, teacache.py, text_encoder.py B1/B2, generate.py, processor/ + vision shards)
Validated with Real generations incl. Turbo LoRA merges + I2V encode checks (see README Validation report + tmp/S9_teacache_spike_RESULT.md)
License Port: Apache-2.0 (PipeNetwork). Weights: MiniMax H3 Community License (LICENSE). Turbo LoRA: Apache-2.0 (larryvrh).

(a) minimax_h3_mlx/dit.py β€” DiT attention QKV blocked layout + TeaCache two-phase hooks

File: minimax_h3_mlx/dit.py (~20 lines β†’ ~230 lines added vs base)

QKV blocked layout (S7, committed 7210b93, lines ~12–17 docstring + lines 153–165 in Attention.__call__):

  • Raw-checkpoint attn.qkv_proj rows are blocked [q | k | v] β€” three contiguous thirds of width heads * head_dim β€” not per-head interleaved [h0: q,k,v][h1: q,k,v].... The old code reshaped to (B, S, heads, 3, head_dim) and indexed qkv[...,0] which scrambled heads. Fix splits and reshapes each third independently (same layout mlx-serve consumes via splitEqual(qkv,3); video VAE to_qkv remains per-head interleaved β€” different codebase, different convention):
    # minimax_h3_mlx/dit.py:160-164 (Attention)
    qkv = self.qkv_proj(x)
    q, k, v = mx.split(qkv, 3, axis=-1)
    q = q.reshape(B, S, self.heads, self.head_dim)
    k = k.reshape(B, S, self.heads, self.head_dim)
    v = v.reshape(B, S, self.heads, self.head_dim)
    
  • Module docstring (lines 12–17, 157–158) now documents the blocked layout and the VAE contrast.

TeaCache two-phase forward (S9, working-tree, lines ~335–510):

  • New helpers TYPE_CHECKING import of ModulationCache, _resolve_cache_layer (lines 347–354), _prepare_stages (355–391, shared input projections + token_refiner + packed buffer + temb + adaln indices so the prefix graph is built exactly once), _run_blocks (393–409, iterates blocks[i] with optional ModulationCache.get(i)), _finish_heads (411–423, final_layer.norm_out + video_out/audio_out heads).
  • Public two-phase entry points: forward_to_cache_layer (425–464, partial forward up to and including cache_layer β€” default last block len(blocks)-1 β€” returning (hidden, ctx) where hidden is pre-final-norm residual after that block and ctx carries rotary/temb/adaln/layer bookkeeping) and finish_from_cache_layer (466–487, completes from cached hidden/ctx without recomputing prefix, returns (video_velocity, audio_velocity)). __call__ (489–510) now delegates to _prepare_stages/_run_blocks/_finish_heads and gains return_hidden for probing. Saving per skip = suffix blocks[layer+1..49] + final_layer.norm_out + both heads β€” with default last block only heads+norm are saved (~<5% wall); moving to block 40 saves ~20% blocks per skip (see S9 spike result: 0/7 skips at last block, no wall saving).

(b) minimax_h3_mlx/text_encoder.py β€” 8-bit quantized loader + processor fallback (B1) + positional scatter (B2)

File: minimax_h3_mlx/text_encoder.py (+84 lines vs base, ~6 β†’ ~17 imports)

8-bit quantized loader (S7, committed 7210b93, lines ~64–110 + ~143–180):

  • Detects MLX affine 8-bit pack by presence of quant_config.json (line 71: self.quantized = (model_dir / "quant_config.json").exists()). Before loading, quantizes the module tree so packed U32/scales/biases key 1:1 β€” same recipe load.py replays from quant_config.json:
    # lines 90-92
    nn.quantize(self.language, group_size=64, bits=8,
                class_predicate=lambda _path, m: isinstance(m, nn.Linear))
    
  • Vision tower: if load_vision=True and vision_quant_config.json exists, replays its recorded quantized paths (89 tensors) with same group/bits (lines 98–110):
    with open(model_dir / "vision_quant_config.json") as fh: vq = json.load(fh)
    paths = set(vq.get("quantized", ()))
    nn.quantize(self.vision, group_size=vq.get("group_size",64), bits=vq.get("bits",8),
                class_predicate=lambda _p, m: isinstance(m, nn.Linear) and _p in paths)
    
  • Quantized shards loaded raw (no astype(dtype), lines 166–170: buckets[bucket][path] = tensor when self.quantized else tensor.astype(dtype)), preserving packed integers. Only dense path casts.
  • Serve pack omits model.norm (never evaluated β€” H3 conditions pre-norm); loader fabricates zeros like dense converter (lines 176–180):
    buckets["language"]["norm.weight"] = mx.zeros(tuple(module.norm.weight.shape), dtype=mx.bfloat16)
    

Processor property with torch-free fallback (B1, S10, lines 207–233):

  • @property processor (lines 207–233) tries AutoProcessor.from_pretrained(processor_dir) (full Qwen3VLProcessor with torch video/image sub-processors). On ImportError (no torch/torchvision; Qwen3VLVideoProcessor hard-requires them, and transformers 5.15 AutoImageProcessor is itself gated), falls back to PIL-only Qwen2VLImageProcessorPil:
    try: self._processor = AutoProcessor.from_pretrained(processor_dir)
    except Exception:
        try: image_processor = AutoImageProcessor.from_pretrained(processor_dir)
        except Exception:
            from transformers.models.qwen2_vl.image_processing_pil_qwen2_vl import Qwen2VLImageProcessorPil
            image_processor = Qwen2VLImageProcessorPil.from_pretrained(processor_dir)
        self._processor = SimpleNamespace(image_processor=image_processor)
    
    The facade only ever reads .image_processor (build_request line 251), tokenizer comes from separate tokenizer property (lines 197–205), so image-only processor is sufficient for I2V.

Encode positional scatter fix (B2, S10, lines 324–344):

  • encode() (lines 306–344) previously did mx.where(image_mask[...,None], hidden[None], inputs_embeds) which mis-broadcasts (1,3096,1) vs (1,3072,5120) because vision hidden rows are compact while image pads sit at arbitrary positions in the longer token row β€” a positional scatter, not a broadcast.
  • Fix validates hidden.shape[0] == num_image (lines 330–334), then does positional REPLACE scatter via scatter_add (mlx 0.32 has no nonzero/index_put, only at[].add):
    pos_np = np.nonzero(np.array(image_mask))[0]   # line 339 (mlx has no nonzero, use numpy)
    pos = mx.array(pos_np)                         # line 340
    target = inputs_embeds[0]                      # line 341
    inputs_embeds = target.at[pos].add(hidden.astype(target.dtype) - target[pos])  # line 342
    inputs_embeds = inputs_embeds[None]            # line 343
    
    Validated with tmp/s10_validate_encode.py (1,3096,5120) finite, 3074 video-tagged rows (start+3072 pads+end). build_request (lines 237–277) tags whole vision block as TAG_VIDEO (not text) β€” the DiT AdaLN key.

(c) minimax_h3_mlx/pipeline.py β€” release_text_encoder headroom (S8/S8b) + TeaCache hook (S9) + checkpoint resume

File: minimax_h3_mlx/pipeline.py (+~205 lines vs base)

Release text encoder (S8/S8b, lines 59–73 __init__, 77–118 from_pretrained, 228–331 __call__, 470–532 _release_text_encoder_now):

  • MiniMaxH3Pipeline.__init__ gains release_text_encoder: bool (line 66, stored as self._release_text_encoder line 73). from_pretrained threads it (lines 84, 118).
  • In __call__ (lines 328–331), right after prompt_embeds built (line 1, text conditioning + vision rows already encoded, step-4 noise already drawn β€” trajectory fixed), opt-in release_text_encoder or self._release_text_encoder drops the encoder:
    if (release_text_encoder or self._release_text_encoder) and getattr(self, "text_encoder", None) is not None:
        self._release_text_encoder_now()
    
  • _release_text_encoder_now (lines 470–532) drops self.text_encoder (del + gc.collect() pattern), sizes correctly via mlx.utils.tree_flatten (lines 484,496,510 β€” previous module.parameters().values() summed dicts not arrays β†’ 0.0 GB), calls mx.clear_cache() to return unified Metal memory (line 528), logs released text encoder after conditioning (freeing ~N GB) (line 532). Second call is no-op (line 471 guard). Safe: inner estimator errors just leave freed-size as "(sizing indeterminado)" while delete still happens. Measured: frees ~27.5 GB / ~22 GB resident before denoise loop; pipeline then only touches DiT+VAEs.
  • CLI: scripts/generate.py --release-encoder (line 56) β†’ pipeline(release_text_encoder=args.release_encoder) (line 96).
  • Checkpoint resume (bonus, lines 371–389): resume_from/checkpoint_dir floats inside the loop (not a pipeline-level patch for HF, but present in pipeline.py working tree).

TeaCache hook (S9, lines 43, 228–230, 342–410):

  • Imports TeaCacheConfig/TeaCacheController (line 43), __call__ gains teacache, teacache_config, teacache_layer (lines 228–230 docstring 244–252). Controller instantiated when teacache=True (lines 342–350, resolves tc_layer = num_layers-1 if None else int, validates range). Per-step two-phase: hidden, ctx = dit.forward_to_cache_layer(..., cache_layer=tc_layer) β†’ mx.eval(hidden) β†’ skip = controller.decide(i, hidden) β†’ on skip reuse cached_preds = (video_pred, audio_pred) joint (lines 391–399), else finish_from_cache_layer and cache preds (lines 406–410). Log: step N/T teacache skip (rel_l1=..., thresh=...) (lines 398–401). Default last-block hook measured 0/7 skips at 768x448 turbo8 (see teacache.py below β†’ OFF recommended).

Other: load.py one-line import fix for quantized path (line 1 added), scripts/generate.py adds --teacache*/--release-encoder flags (lines 45–56, 85–96).


(d) minimax_h3_mlx/teacache.py β€” two-phase TeaCache controller (S9, new file)

File: minimax_h3_mlx/teacache.py (11 KB, new β€” 223 lines)

  • Intent: feature-caching (ByteDance/ali-vilab line) β€” each step runs partial forward to a probe hidden, compares against previous step's feature, reuses previous (video, audio) velocity if similar. Joint reuse (both modalities from one forward must be reused together). Noise/keyframe sampling untouched β†’ seed still fixes trajectory.
  • TeaCacheConfig (lines 40–71): rel_l1_thresh=0.2 (default, ref 0.15 for 10s 544x960), metric='rel_l1'|'cosine', start_at_step=3, compute_last_step=True (protects first 3 and last step), cache_type='avg'|'last'. __post_init__ validates.
  • Pure helpers (lines 87–134): relative_l1_distance (mean|curr-prev|/mean|prev|, lines 87–100), cosine_similarity (lines 103–113), teacache_gate (1-min(dist,1) for rel_l1 so higher=more similar, lines 115–124), teacache_should_skip (lines 126–134: rel_l1 <= thresh or cosine >= thresh).
  • TeaCacheController (lines 148–242): stateful per-run (total_steps, cached_feature, computed/skipped, by_step). decide(step_index, feature) (lines 171–210): must_compute if step==0 or step<start or cache is None or last step β†’ prime cache; else compute metric_value (cosine or rel_l1 distance), skip = metric >=/<= thresh, smoothing avg folds current into cache as (prev+curr)/2 on skip, else replaces. Stats: last_gate (similarity-domain), last_metric (raw distance/similarity for honest log), metric_name, to_stats().
  • Pipeline integration: see (c) above; dit.py hooks provide the probe feature. Measurement (M4 Max, 768x448 turbo8, 7 forwards): default last-block thresh 0.2 β†’ 0 skips (+9.7% overhead), 0.35 β†’ 1/7 skip (rel_l1 0.297) still no wall saving (partial forward already 49/50 blocks, only final_layer.norm_out+heads saved). tmp/S9_teacache_spike_RESULT.md β€” verdict: TeaCache OFF by default; aggressive needs earlier hook layer 40 + thresh 0.25–0.35 for double-digit % (trades quality).

(e) processor/ + vision shards (S10 I2V)

Files added to checkpoint (not code patches but data required for I2V):

  • processor/ (7 files, ~11.6 MB): preprocessor_config.json (390 B), video_preprocessor_config.json (385 B), chat_template.json (5.5 KB), tokenizer_config.json (11 KB), tokenizer.json (6.7 MB), vocab.json (2.6 MB), merges.txt (1.6 MB) β€” copied from models/minimax-h3-ckpt/processor/ (Qwen3VLProcessor, Qwen2VLImageProcessorFast + Qwen3VLVideoProcessor configs; runtime uses PIL fallback per (b) B1).
  • text_encoder/model-00005-of-00008-vision.safetensors (291,980,081 B), model-00006-of-00008-vision.safetensors (341,371,441 B), model-00007-of-00008-vision.safetensors (129,267,789 B) β€” 529 model.visual.* tensors (vision tower: depth 27, hidden 1152, patch 16, merge 2, out 5120). Built by scripts/rebuild_h3_text_encoder_vision_8bit.py which transposes patch_embed.proj.weight from torch (O,C,kD,kH,kW) to MLX channels-last and quantizes 89 paths.
  • text_encoder/vision_quant_config.json (2,461 B, 89 quantized paths: blocks.*.attn.qkv/proj, blocks.*.mlp.linear_fc1, merger.*, deepstack_merger_list.*).

Why this matters: without processor/ the pipeline cannot build pixel_values/image_grid_thw (I2V falls back to T2V). Without vision shards text_encoder.encode(prompt, images) raises load_vision=False or missing-tensor error. With both, encode() (b-B2) scatters 3072 vision patches into inputs_embeds and hits the DiT as TAG_VIDEO rows, then pipeline._encode_keyframes patches conditioning latents β€” FL2VA keyframe path (anchors first/last) validated end-to-end (see README ## Image-to-Video).


Other minor port fixes

  • minimax_h3_mlx/load.py: quantization structure replay for DiT (already in committed 7210b93) β€” one-line import guard retained.
  • scripts/generate.py: flags --teacache, --teacache-thresh (default 0.2), --teacache-metric (rel_l1/cosine), --teacache-layer, --teacache-start (default 3), --release-encoder (all S8/S9).
  • tests/test_release_text_encoder.py, tests/test_teacache.py: unit/mocked tests for (c)/(d) (not shipped in HF package, kept in repo).

Reproducing the diff

git -C models/minimax-h3-mlx diff b2f7e4d2 -- minimax_h3_mlx/dit.py minimax_h3_mlx/text_encoder.py minimax_h3_mlx/pipeline.py minimax_h3_mlx/teacache.py scripts/generate.py
# committed two-file patch:
git -C models/minimax-h3-mlx show 7210b93 --stat

License

  • Weights: MiniMax H3 Community License (included as LICENSE) β€” not open source; territorial exclusions apply (EU/UK/KR/US excluded per mlx-serve pack note).
  • Port code: Apache-2.0 (PipeNetwork/minimax-h3-mlx).
  • Turbo LoRA: Apache-2.0 (larryvrh/MiniMax-H3-Turbo-Lora).