[feat]: preserve Apple MLX future 5B, I2V, and self-forcing tracks by aryan5v · Pull Request #13 · aryan5v/FastVideo

aryan5v · 2026-07-09T22:06:35Z

Summary

Future-only consolidation of the former stacked Apple MLX PRs.

Preserves causal/self-forcing, hardware-tiering, Wan2.2 5B, latent-only I2V, 48-channel decode/DMD, QAD arming, benchmarks, and parity scaffolds from former fork PRs Mac Day-1: numerics gate + mx.compile unblock (1.4x) + Track C Rung 1 #4 and [feat]: Hardware-adaptive model tiering (auto-select quant+caps by unified memory) #6 through [feat]: MLX VAE/TAEHV decode + validated 5B DMD sampler #11.
Retains the newer MLX 0.31.2 QAT contract and adds a precise PR-by-PR preservation map in tests/local_tests/wan2_2_ti2v_5b/README.md.
This is not a public 5B/I2V release: source revision/checksum, strict real-weight loading, component/pipeline parity, and media-quality gates remain required.

Validation

pytest fastvideo/tests/mlx -q -rs: 110 passed, 4 expected artifact-gated skips.
git diff --check: passed.

Supersedes

Former fork PRs #4, #6, #7, #8, #9, #10, and #11. Each is mapped to its preserved implementation and remaining gate in the port-preparation README.

… Phase A) fastvideo/layers/quantization/mlx_affine_qat.py transcribes MLX's affine quantizer from the v0.31.2 CPU kernel: fp32 group min/max, negative scale when |max| >= |min|, the anchor endpoint re-fit to an exact integer multiple of the scale (q0 = rint(edge/scale); scale = edge/q0; bias = edge), rint rounding, and clamped codes. fake_quantize_mlx_affine simulates the full deploy pipeline (master weight -> fp16 cast -> quantize -> dequantize) with a straight-through estimator, returning fp32 so the fp16 deploy values survive bf16 master dtypes. test_mlx_affine_qat_parity.py pins the transcription against the real mx.quantize/mx.dequantize bitwise -- unpacked codes, scales, biases, and dequantized weights, for int8 and int4 at fp32 and fp16 -- plus STE gradient passthrough and a quantized-matmul tolerance check. This is the roadmap's 'parity test before any GPU spend' gate for the Mac-targeted QAT run; both CI smoke jobs now run it. Verified: pytest fastvideo/tests/mlx/ -> 53 passed (mlx-cpu 0.31.2). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

…hase A complete) MLXQuantizationAwareCallback registers a weight parametrization on the student transformer so every training forward sees weights on the exact MLX affine deploy grid (fake_quantize_mlx_affine; numerics pinned bitwise in test_mlx_affine_qat_parity.py) while gradients pass straight-through to the master weights. Composes with any TrainingMethod via YAML ('mlx_qat' builtin); targets 2-D and conv weights whose grouped dim divides group_size, excluding norms/scale_shift_table, matching what the MLX loader quantizes. examples/train/configs/distribution_matching/wan/dmd2_t2v_mlx_int8.yaml is the ready-to-launch recipe: Wan2.1-T2V-1.3B teacher -> 3-step INT8 student on FastWan timesteps [1000, 757, 522], 4 GPUs, W&B tracking. First multi-GPU smoke run should verify parametrizations shard cleanly under HSDP. Verified: 5 callback tests (targeting, deploy-grid forward incl. conv flattening, STE gradient flow + requantization after optimizer step, error paths) + full MLX suite -> 58 passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

Operator-facing steps for M4 Phase B: preflight, install, dataset pull, mandatory 100-step smoke run with explicit pass/fail checks (QAT callback armed, no parametrization x HSDP crash, finite losses, s/step extrapolation), the full 4-GPU launch, DCP->Diffusers export, and the deliverables to report back for Mac-side evaluation. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

…masters The DGX smoke run trained 50 steps cleanly but crashed in validation: 'Input type (float) and bias type (c10::BFloat16) should be the same' at the Wan patch embedding. Under HSDP mixed precision the master weights (parametrizations.weight.original) are stored fp32 and only cast to bf16 at forward time; the QAT parametrization returned the fake-quantized weight in the master's dtype, so module.weight presented fp32 to everything outside autocast -- validation pipelines and dtype sniffing -- against bf16 biases. The parametrization now takes an explicit compute_dtype and always presents the weight in it (recipe sets bf16); gradients still reach the fp32 master through the cast. Registered with unsafe=True since the parametrized dtype legitimately differs from the master's. New regression test covers the exact scenario: fp32 masters, bf16 presentation, STE gradient flow. Verified: 7 callback tests + full MLX suite -> 65 passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

VideoCaptionMergedDataset subclasses torch.distributed.checkpoint.stateful.Stateful but the module never imported that submodule, so importing the dataset only worked when something else had loaded it first (surfaced by the lightweight DGX training env, which had to bootstrap around it). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

Training-time validation crashed on FSDP/HSDP models -- base DMD2 recipe and QAT alike (A/B-verified on a DGX B200): the stages derived target_dtype from next(parameters()).dtype, which under mixed precision returns the fp32 master storage. That left validation latents fp32 AND disabled autocast (the target_dtype != fp32 gate), while FSDP cast the weights to bf16 per-forward: 'Input type (float) and bias type (c10::BFloat16) should be the same' at the Wan patch embedding. transformer_compute_dtype() now prefers the mixed-precision policy the model loader registers (get_compute_dtype) and falls back to parameter sniffing when no policy is set, preserving the pure-inference and MPS fp16 paths. Both sniff sites in denoising.py use it. Verified: 4 new dtype-resolution tests (policy over fp32 masters, fallback sniffing, fp32-policy deferral, .module unwrap) + QAT callback + MLX suites -> 63 passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

…trize The DGX smoke run hit 'aten.convolution.default: got mixed torch.Tensor and DTensor' at the patch embedding: torch.nn.utils.parametrize restructures the weight into a submodule and computes from the raw master, which under FSDP2 is a sharded DTensor outside the unshard window -- so the parametrized weight mixed sharded DTensors with the unsharded plain tensors FSDP provides inside forwards. The callback now wraps each target module's forward and swaps the weight for its fake-quantized version only for the duration of that call. Outside forwards the module is completely vanilla (FSDP, optimizers, dtype sniffing, DCP checkpointing, and dcp_to_diffusers export all see ordinary parameters); inside forwards the fake-quant operates on exactly the unsharded compute-dtype weight the matmul would have used, which also makes the earlier compute_dtype option unnecessary (removed). Restoration is exception-safe and wrapping is idempotent. Verified: 8 callback tests (targeting, vanilla-outside-forward invariant, deploy-grid forwards incl. conv, STE gradients + requantization, idempotence, exception-safe restore) + dtype-resolution + MLX suites -> 65 passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

Rebuilds the run's callbacks so the checkpoint's 'callbacks' DCP entry loads into the EMA callback's shadow container, then swaps the EMA weights into the role's transformer (ema_context) for the export. Fails with a clear message when the config has no EMA callback or the checkpoint's EMA never started. Runbook now exports both raw and EMA for Mac-side comparison. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

Run 1's EMA was unrecoverable: EMA_FSDP local_shard shadows are plain per-rank tensors, DCP deduplicated them as replicated (only rank 0's quarter shards were saved, incl. a [0,...] empty shard), and the 1-GPU --ema export loaded those into full-shape buffers -- noise weights whose uniform INT8-vs-FP16 SSIM masqueraded as a QAT win (probes on the DGX confirmed: 826 keys = one shard per param, [384,1536] shapes; export values ~1.22 relative distance from both raw and base). EMACallback now checkpoints full tensors with activation-checkpointing wrapper names normalized (DTensor gather on save, re-slice to the current topology on load) and refuses legacy per-shard state instead of loading silently-corrupt weights. CPU round-trip tests cover world-size-1 and name normalization; the multi-rank path uses the same DTensor helpers and gets exercised by the run-2 smoke. dmd2_t2v_mlx_int8_v2.yaml is run 2: student+critic init from the already 3-step-distilled FastWan weights (teacher stays base Wan2.1), gradient_accumulation_steps 4 -- run 1's student moved only ~0.2% from init at global batch 4 and shows motion defects; both levers target that. Verified: 71 tests passed (EMA round-trip, QAT callback, dtype resolution, MLX suites). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

…ctives The run-2 smoke hung at checkpoint-20: the previous portable-EMA fix gathered full tensors in state_dict() via DTensor.full_tensor(), which is a collective (all-gather). Issuing 826 collectives from inside DCP's save path deadlocked and NCCL timed out after 10 minutes. state_dict() now presents each local shard as a DTensor via DTensor.from_local (metadata-only, no communication) under the 'student_ema_sharded' key. DCP saves every rank's shard and reshards natively on load at any world size -- the same mechanism the model weights under roles.* already use. load_state_dict re-slices via to_local; legacy plain-shard state is still refused. Verified: EMA round-trip (world=1), QAT callback, dtype, and MLX suites -> 68 passed. The multi-rank gather/reshard path is exercised by the next DGX smoke. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

coderabbitai · 2026-07-09T22:06:44Z

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: e5b8009c-83cf-4545-b270-b7983503c5cb

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch aryan/future/fastwan-qad-5b-i2v

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

gemini-code-assist

Code Review

This pull request introduces comprehensive support for Apple Silicon MLX runtime inference and quantization-aware training (QAT) for Wan2.1 and Wan2.2 models, including hardware-adaptive tiering, causal streaming attention, and VAE decode helpers. The review feedback highlights a few issues: a potential AttributeError in wan22_i2v.py due to the use of dim() instead of ndim, a contradiction between a comment and the download-blocking logic in wan_vae.py, dead code in wan22_5b_qad_arming.py consisting of an empty if statement, and an unused import in wan22_sampler_reference.py.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-09T22:25:27Z

+            if img.dim() == 5:
+                img = img[:, :, 0]


Using "img.dim()" will raise an "AttributeError" if "image_latent_frame" is passed as a NumPy array or an MLX array, as they do not implement the ".dim()" method. Since "ndim" is universally supported by PyTorch, NumPy, and MLX (and is already used elsewhere in this file on lines 80 and 97), we should use "img.ndim" instead.

if img.ndim == 5:\n img = img[:, :, 0]

gemini-code-assist · 2026-07-09T22:25:27Z

+# Leave empty to skip pin for taew2_2 until a long-lived pin is recorded (still
+# downloads; user-supplied paths are never hash-checked).
+TAEW2_2_SHA256 = ""


The comment states that "taew2_2" still downloads when the SHA256 pin is empty. However, the implementation on line 94 explicitly raises a "RuntimeError" and blocks the download if "expect" is empty (which "TAEW2_2_SHA256" is). Please update the comment to accurately reflect that downloading is blocked until a valid SHA256 pin is recorded, or adjust the download logic if downloading without a pin was intended.

gemini-code-assist · 2026-07-09T22:25:27Z

+    if not isinstance(sample.weight, type(next(transformer.parameters()))):
+        # Parameter check
+        pass
+    if "weight" not in sample._parameters:


This "if" statement is empty and contains only a "pass" statement, making it dead code. If a type check is intended here, please implement the assertion or error handling (e.g., printing a failure and returning an error code). Otherwise, this redundant block should be removed.

if \"weight\" not in sample._parameters:

gemini-code-assist · 2026-07-09T22:25:27Z

+            current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")):
+        for i, t in enumerate(timesteps):
+            t_val = float(t.item())
+            ts = torch.full((1, tokens_n), t_val, device=device, dtype=torch.long)


The import of "pred_noise_to_pred_video as torch_p2v" is unused in this function and can be safely removed to keep the code clean.

aryan5v and others added 20 commits July 9, 2026 14:21

[feat]: add MLX FastWan 1.3B Apple release runtime

0fada8b

[fix]: trace FastWan MLX compile scalars

05f4eab

[docs]: runbook guidance for running QAD on fewer or more than 4 GPUs

5048cd3

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HmszJweXHHu14NY4zQ2uF8

[feat]: stage MLX QAD 5B port foundation

fda8e6e

[bugfix]: consolidate MLX release dependency pin

bcb4658

[feat]: preserve internal causal MLX future track

da208bc

[feat]: preserve internal Wan2.2 5B MLX scaffolding

a25ba98

[feat]: preserve latent-only I2V and MLX decode preparation

bbabc3d

[docs]: record future-stack preservation and hard gates

178f4a0

[style]: normalize preserved MLX files

dd97402

gemini-code-assist Bot reviewed Jul 9, 2026

View reviewed changes

aryan5v force-pushed the aryan/release/fastwan-qad-int8-1.3b-mlx branch from d60d8e1 to 7ed0356 Compare July 23, 2026 22:35

aryan5v force-pushed the aryan/release/fastwan-qad-int8-1.3b-mlx branch from 7ed0356 to 7822aed Compare July 23, 2026 22:59

Navigation Menu

[feat]: preserve Apple MLX future 5B, I2V, and self-forcing tracks - #13

[feat]: preserve Apple MLX future 5B, I2V, and self-forcing tracks#13
aryan5v wants to merge 20 commits into
aryan/release/fastwan-qad-int8-1.3b-mlxfrom
aryan/future/fastwan-qad-5b-i2v

aryan5v commented Jul 9, 2026

Uh oh!

coderabbitai Bot commented Jul 9, 2026

Review skipped

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 9, 2026

Uh oh!

gemini-code-assist Bot Jul 9, 2026

Uh oh!

gemini-code-assist Bot Jul 9, 2026

Uh oh!

gemini-code-assist Bot Jul 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

aryan5v commented Jul 9, 2026

Summary

Validation

Supersedes

Uh oh!

coderabbitai Bot commented Jul 9, 2026

Review skipped

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants