John 40688778c2 SparseStructureFlowModel verified against upstream at correlation 1.00000000

700/700 params, max abs diff 1.2e-5 on the real 1.3B checkpoint.

Key realisation: the flow models have no sparse conv, so upstream RUNS on CPU torch
with flash-attn swapped for SDPA. That gives a real numerical oracle - unavailable for
the sparse path, where spconv cannot be installed at all.

It was needed. Three bugs survived a loader reporting a perfect 700/700 with zero
missing and zero unmapped keys:
- a parameterless final LayerNorm (no params -> no checkpoint trace) that the output
  was 200x too large without
- rope_phases being complex64, so the rotation is a complex multiply
- qk_rms_norm belonging before rope rather than after

Weight-key matching is necessary but nowhere near sufficient for a port.

2026-08-02 12:03:12 +10:00

3.0 KiB

Raw Blame History

pixal3d_mrp_mlx — working notes

MLX port of TencentARC Pixal3D, on top of the shared trellis_sparse_mlx core.

Measured facts (don't re-derive)

Pixal3D inherits TRELLIS.2's sparse module, same as LATO.2. conv_spconv.py is the same code: stride==1 and padding is None -> spconv.SubMConv3d.
EVERY SparseConv3d in the model is (channels, out, 3) -> submanifold only. The shared core's SubMConv3d covers all of them.
attn_mode is 'full' everywhere -> SDPA. No flash-attn/xformers needed.
Container differs from LATO.2: SparseTensor(VarLenTensor). VarLenTensor is feats + explicit slice-per-batch layout, no coords. Both are in the shared core now.
Cache API spelling is get_spatial_cache/register_spatial_cache (LATO.2 uses cache_get/cache_put). Shared core exposes both.
_scale uses Fraction, not int, upstream. Shared core keeps scale as an opaque tuple so either works.
Only sparse-conv-using model file is models/sc_vaes/sparse_unet_vae.py.

Remaining ops to port

SparseUpsample(2) — pixal3d/modules/sparse/spatial/basic.py. Cache-paired: needs the subdivision stored by its matching SparseDownsample, or an explicit subdivision tensor.
SparseSpatial2Channel(2) — pixal3d/modules/sparse/spatial/spatial2channel.py. Sparse pixel-shuffle; also cache-backed.

Conventions

Depends on ../trellis_sparse_mlx (editable install). Do NOT vendor a second copy of the sparse core.
upstream/Pixal3D is vendored read-only reference; never edit it.
Weights live at MODELBEAST/vendor/pixal3d-weights (24GB, outside this repo).

Weights (measured 2026-08-02)

Every checkpoint ships a sibling .json with {name, args} — the authoritative config. Do NOT infer architecture from shapes here (that was needed for LATO.2, not this).
ss_flow + slat_flow x3 (~20GB, the bulk): ZERO 5-D tensors. Pure transformers, no sparse conv. They pass through the converter untouched.
shape_dec / tex_dec: 40 KRSC sparse kernels each. ss_dec: 20 dense Conv3d. classify_5d() distinguishes them correctly on the real weights.
rope_phases ships as a stored tensor — RoPE phases are precomputed, not derived.
Converter preserves dtype (fp16 stays fp16); upcasting doubled 24GB for nothing.

Numerical oracle (important)

The flow models have NO sparse conv, so upstream RUNS ON CPU TORCH here. Patch pixal3d.modules.attention.modules.scaled_dot_product_attention with a torch SDPA wrapper (permute to [B,H,N,D] and back) and it loads the real checkpoint. Use this to diff any flow-model change — do not reason about correctness, measure it.

Bugs found that weight-matching could NOT catch (loader said 700/700, 0 missing)

Parameterless final F.layer_norm between last block and out_layer. No params -> no checkpoint trace. Without it output is ~200x too large.
rope_phases is complex64 (torch.polar). Rotation is a complex multiply; cos/sin are the real/imag parts, NOT cos(phase).
qk_rms_norm goes BEFORE rope, not after. Non-commutative. One block correlates 0.9998 when reversed; compounds to 0.84 over 30 blocks.

3.0 KiB Raw Blame History