Hunyuan3D-2.2-mrp-MLX/hy3dpaint/tests/comparison_report.md

MLX vs PyTorch UNet Numerical Comparison Report

Test configuration

• Latent spatial size: 8 x 8
• Input channels: 12 (4 latent + 4 normal + 4 position)
• Encoder text: zeros (1, 77, 1024)
• Timestep: 500
• Both models cast to fp32
• Tolerance: 0.01 (abs)

Setup fixes made to the script

1. MLX UNet built at in_channels=12 with attention_head_dim=(5, 10, 20, 20) (matching the PyTorch config), no _enhance_unet() — we compare the vanilla UNet topology.
2. PyTorch UNet rebuilt from config.json with in_channels=12 (the config says 4 but the actual checkpoint has 12-channel conv_in). Loading via from_pretrained silently drops most weights because the checkpoint's keys are nested inside unet. / unet_dual. with a .transformer. wrapper for BasicTransformerBlock. We now:
   • load the .bin file manually,
   • strip the unet. prefix,
   • collapse transformer_blocks.0.transformer.X to transformer_blocks.0.X,
   • drop 2.5D-only tensors (attn_multiview, attn_refview, attn_dino, .processor.).
   • This yields 686 weights loaded, 0 missing, 0 unexpected.

Weight spot-check

All 18 spot-checked weight tensors (conv_in, time_embedding, conv_out, conv_norm_out, resnet 0 weights, attention.0 norm/proj_in, attn1 to_q/to_k/to_v for both down_blocks[0] and mid_block) are bit-identical between the PT state_dict and the MLX safetensors after layout transposition. Conclusion: the converter is fine; the bug is in the MLX modules themselves.

Forward-pass divergence (plain run, same weights, same input)

layer	shape	max_abs	status
conv_in	(1, 8, 8, 320)	1.5e-07	MATCH
down_block_0	(1, 4, 4, 320)	6.29	MISMATCH
mid_block	(1, 1, 1, 1280)	15.18	MISMATCH
up_block_0	(1, 2, 2, 1280)	6.19	MISMATCH
conv_out / final	(1, 8, 8, 4)	0.68	MISMATCH

First divergent layer: down_blocks[0] (while conv_in still matches).

Drill-down inside down_blocks[0]

PyTorch forward hooks captured each sub-module output; MLX sub-modules were run in the same order.

sub-module	max_abs (plain)	max_abs (GN fix)	max_abs (PT temb)
resnet 0	1.83	0.042	1.85
attention 0	2.33	0.75	2.31
resnet 1	2.51	0.78	2.49
attention 1	2.70	0.99	2.68

Each experimental column changes one variable vs the "plain" column:

• GN fix: every nn.GroupNorm in the model patched to pytorch_compatible=True.
• PT temb: the timestep embedding for the MLX forward is replaced by the exact numpy tensor produced by PyTorch's time_embedding.

The GN fix collapses resnet-0 error by ~44x (from 1.83 to 0.042 — roughly at fp32 tolerance). Substituting the PT timestep embedding alone barely moves the needle at this stage.

---

Bugs identified

Bug 1 (primary, high impact): nn.GroupNorm uses MLX-native mode, not PyTorch-compatible mode

mlx.nn.GroupNorm.__init__ defaults to pytorch_compatible=False, which implements a different normalization scheme. Standalone test with identity affine parameters on a (1, 8, 8, 320) tensor:

MLX default                 vs PyTorch GroupNorm: max |diff| = 3.45e-01
MLX pytorch_compatible=True vs PyTorch GroupNorm: max |diff| = 7.15e-07

Every nn.GroupNorm in the MLX port is affected. Locations:

• hunyuanpaintpbr_mlx/unet/blocks_mlx.py lines 64, 66 (ResnetBlock2D)
• hunyuanpaintpbr_mlx/unet/blocks_mlx.py line 403 (Transformer2DModel.norm)
• hunyuanpaintpbr_mlx/unet/modules_mlx.py line 338 (2.5D Transformer2DModel.norm)
• hunyuanpaintpbr_mlx/unet/unet_mlx.py line 153 (conv_norm_out)
• hunyuanpaintpbr_mlx/vae_mlx.py lines 24, 26, 59, 221, 263

Fix: pass pytorch_compatible=True on every nn.GroupNorm(...) construction site. The drill-down showed that applying this fix at runtime shrinks the resnet-0 max_abs from 1.83 to 0.042 (with the remaining 0.042 attributable to Bug 2).

Bug 2 (secondary, clearly measurable): timestep-embedding frequency denominator

hunyuanpaintpbr_mlx/unet/blocks_mlx.py line 20 uses:

emb = math.log(10000) / (half_dim - 1)

Diffusers' get_timestep_embedding with downscale_freq_shift=0 (the value SD 2.1 uses) divides by half_dim, not half_dim - 1. Direct comparison at dim=320, t=500:

[timestep_sinusoidal] max_abs=1.06  mean_abs=0.21  [MISMATCH]

The first element matches (frequency 0) but all others drift — by up to 1.06 on individual entries. Swapping MLX's timestep embedding for PyTorch's only moved resnet-0 max_abs by ~0.02 at this stage, so Bug 2 is quantitatively smaller than Bug 1 for this particular input, but it is still present and will stack with Bug 1 after many layers / non-trivial text inputs.

Fix: change to emb = math.log(10000) / half_dim (and in principle take a downscale_freq_shift argument, but 0 is what we need for SD 2.1).

Non-bug: the residual attention error (~0.75 after GN fix)

After fixing GroupNorm, attention-0 still shows max_abs ~0.75 at down_blocks[0]. This is explained by:

• Its input already carries ~0.042 error from resnet-0.
• The first attention's softmax and proj_in amplify small input perturbations.
• Fixing Bug 2 removes a further small component.

No obvious structural bug in the attention code was found during the drill-down; the QKV weights compare bit-identically, heads=5 matches PT's per-block head-dim config, and to_out path matches. Recommend re-running this comparison after fixing Bugs 1 and 2 to confirm attention falls within tolerance.

---

How to reproduce

cd /Users/dgrauet/Work/Hunyuan3D-2.1-mlx/hy3dpaint
python3 tests/compare_mlx_pytorch.py

The script prints MATCH/MISMATCH verdicts with max absolute difference, max relative difference, PT-norm and MLX-norm for every comparison, and writes a summary here.

Key files referenced:

• /Users/dgrauet/Work/Hunyuan3D-2.1-mlx/hy3dpaint/tests/compare_mlx_pytorch.py
• /Users/dgrauet/Work/Hunyuan3D-2.1-mlx/hy3dpaint/hunyuanpaintpbr_mlx/unet/blocks_mlx.py
• /Users/dgrauet/Work/Hunyuan3D-2.1-mlx/hy3dpaint/hunyuanpaintpbr_mlx/unet/modules_mlx.py
• /Users/dgrauet/Work/Hunyuan3D-2.1-mlx/hy3dpaint/hunyuanpaintpbr_mlx/unet/unet_mlx.py
• /Users/dgrauet/Work/Hunyuan3D-2.1-mlx/hy3dpaint/hunyuanpaintpbr_mlx/vae_mlx.py