mflux-debugging

mflux debugging (MLX parity vs PyTorch/diffusers)

Use this skill when you are porting a model to MLX and need to prove numerical parity (or isolate where it diverges) versus a PyTorch reference implementation (often from diffusers).

This skill defaults to export-then-compare:

• Run the reference once and export deterministic artifacts (tensors + optional images).
• Load those artifacts in MLX and compare with clear thresholds.

When to Use

• You suspect a port mismatch (wrong shapes/layout, RoPE, scheduler math, dtype casting, etc).
• You want a repeatable workflow to narrow down the first layer/block where outputs diverge.
• You need evidence of correctness before refactoring (see mflux-model-porting).

Ground Rules (repo norms)

• Use uv to run Python: uv run python -m ...
• If you run pytest, preserve outputs: MFLUX_PRESERVE_TEST_OUTPUT=1 (see mflux-testing and the Makefile test targets).
• Do not update or replace reference (“golden”) images unless explicitly asked.
• Debug artifacts (tensor dumps) should live in a local folder and must not be committed unless explicitly asked.
• If you need the broader porting workflow (milestones, ordering, when to refactor), follow mflux-model-porting.
• RNG warning: PyTorch and MLX RNGs are different. Matching the same integer seed is not enough for parity—export the exact initial noise/latents from the reference and load them in MLX.
• Practical setup: the PyTorch reference repo (often diffusers/) and mflux/ are frequently next to each other on disk (e.g. both on your Desktop). Use absolute paths when in doubt.

Default Workflow (export-then-compare)

Preferred workflow: two tiny scripts + inline dumps

For day-to-day debugging, prefer a minimal paired repro:

• One simple script in the reference repo (often diffusers/), e.g. diffusers/flux2_klein_edit_debug.py
• One simple script in mflux/, e.g. mflux/flux2_klein_edit_debug.py

Keep them “boring”: hardcoded variables, no cli, no framework, and just a few np.savez(...) / mx.save(...) lines at the right spot.

The key trick for RNG parity:

• In the reference script, compute latents once, save them, then pass them back into the pipeline (latents=...) so the run definitely uses the dumped tensor.
• In rare cases where a tensor needs to be saved from within a loop, make sure its name reflects the loop conditions (e.g the 4th noise predcition in a 10 step loop etc).
• In the MLX script, load that same latents file and feed it into the MLX run (do not rely on matching integer seeds).

1) Pick a single deterministic repro

• Fix seed(s), prompt(s), height/width, steps, guidance, and any input image paths.
• Keep the first repro small if possible (fewer steps, smaller resolution) to iterate quickly.

2) Decide your checkpoints (what to dump)

Start coarse, then narrow:

• VAE: packed latents before decode; optional intermediate activations for one block if needed.
• Transformer: hidden states at entry/exit of the model, then per-block (or every N blocks), then inside attention/MLP.
• Text encoder: token ids + attention mask, embeddings output, then per-layer hidden states if needed.
• Scheduler: timesteps/sigmas/alphas and the predicted noise/velocity per step.

Tip: work “backwards from pixels” like mflux-model-porting suggests: validate VAE decode first with exported latents, then the diffusion/transformer loop, then text encoder.

3) Export artifacts from the PyTorch reference (no logic changes)

Create a run directory like:

• ./debug_artifacts/<run_id>/ref/

Export with one of these patterns:

• NumPy: np.savez(path, **tensors_as_numpy)
• PyTorch: torch.save(dict_of_tensors, path)

4) Run the MLX side with the same inputs and compare

Create a matching run directory:

• ./debug_artifacts/<run_id>/mlx/

Load and compare tensors. For each checkpoint, report:

• Shape + dtype
• max_abs_diff, mean_abs_diff
• max_rel_diff (guarding division by zero)
• Pass/fail with a clearly stated rtol/atol
• It is more important to inspect actual tensor values (e.g., first 10 elements) than rely on summary stats.
• Statistics can mislead; small-looking stats can hide systematic drift or sign flips.
• Prefer runtime tensor dumps over code reading; code can use different conventions yet still represent the same math.

Suggested tolerance starting points (adjust per component):

• fp32 comparisons: atol=1e-5, rtol=1e-5
• fp16/bf16 comparisons: atol=1e-2, rtol=1e-2
• If comparing images: compare both (a) tensor space before final clamp and (b) saved png visually, since tiny numeric diffs can look identical.

If a checkpoint fails:

• Add an earlier checkpoint and repeat (binary search through the forward path).

Common Causes of Divergence (high-signal checklist)

• Layout mistakes: NCHW vs NHWC, transposes around convs/attention, flatten/reshape ordering. Some operations like convolutions can have different conventions between libraries.
• Broadcasting: scale/shift vectors applied on the wrong axis (common in RoPE and modulation).
• Dtype casting: reference silently upcasts to fp32 for norm/softmax; MLX path stays in fp16.
• RoPE details: position ids, reshape order, whether cos/sin are broadcast over heads vs sequence.
• Scheduler math: timestep indexing, sigma/alpha definitions, and off-by-one step order.
• Scheduler config: compare sigma schedules directly.
• Seed/RNG: ensure you aren’t comparing stochastic paths (dropout, noise sampling) without controlling RNG.
• Device dtype: MPS float16 can produce NaNs; prefer bfloat16 for reference dumps if you see NaNs.
• Do not use CPU for comparisons; always keep reference runs on MPS to avoid misleading behavior.

Artifact Hygiene

• Prefer debug_artifacts/<run_id>/... at repo root.
• Do not commit debug_artifacts/ unless explicitly asked.
• If you convert the parity check into a test, follow the repo’s testing conventions and preserve outputs (see mflux-testing).
• Clean up old artifacts when they are no longer needed, only focus on the current problem and avoid confusion with older artifacts that are not relevant for the current task.

Isolate transformer + VAE (inject reference initial latents)

When full end-to-end images differ between mflux and diffusers but you need to trust the core model (transformer, text encoder, VAE decode), export the reference initial noise and run mflux denoising from that tensor. If injected-latent mflux looks good while native-seed runs differ, the forward path is likely sound and any gap is probably RNG, scheduler, or optional reference-only components.

Do not commit these scripts or Desktop .npy files unless asked. Use inline uv run python <<'PY' ... PY heredocs in the reference repo and in mflux/.

Step 1 — Export initial latents from diffusers

Run from the local diffusers/ clone (often next to mflux/ on Desktop). Match mflux settings: height, width, seed, dtype (bfloat16 on MPS).

Load the reference pipeline the same way you would for a normal generation, but disable optional subsystems that mflux does not implement or download (check model_index.json vs mflux get_download_patterns() / weight definition). Pass None for unused components, set pipeline flags to skip them, and use local_files_only=True so missing weights surface immediately instead of downloading extras.

Then export noise with the same shape helper the pipeline uses internally:

from diffusers.utils.torch_utils import randn_tensor

pipe = <ReferencePipeline>.from_pretrained(
    "<org/model>",
    torch_dtype=torch.bfloat16,
    local_files_only=True,
    # plus any kwargs to omit optional reference-only modules
)
latent_h = height // pipe.vae_scale_factor  # or read from pipeline/config
latent_w = width // pipe.vae_scale_factor
ch = pipe.transformer.config.in_channels    # name varies by architecture
generator = torch.Generator(device="mps").manual_seed(seed)
latents = randn_tensor(
    (1, ch, latent_h, latent_w),
    generator=generator,
    device="mps",
    dtype=torch.bfloat16,
)
np.save("/path/to/init_latent.npy", latents.detach().cpu().float().numpy())

Step 2 — Denoise in mflux from the loaded latent

If generate_image() has no latents= argument, inline the denoising loop from the variant class instead of patching production code:

1. latents = mx.array(np.load(path)).astype(mx.bfloat16)
2. Text encode step from the variant (e.g. _encode_prompts, prompt cache, etc.)
3. Any position encoding / conditioning setup the variant does before the loop
4. Denoise loop: noise prediction + config.scheduler.step(...) over config.time_steps
5. VAE decode + ImageUtil.to_image(...) (or the variant’s decode helper) → save PNG

Also run one generate_image(seed=...) baseline on the same prompt for side-by-side.

How to read results

If injected-latent mflux looks good, the port’s core forward path is trustworthy. Golden CI tests should still lock mflux-native sampling (see mflux-testing), not diffusers pixel parity.

diffusers reference runs (side-by-side images)

• mflux does not depend on diffusers; run reference scripts from a sibling diffusers/ clone with cd .../diffusers && uv run python.
• Use HF_HUB_OFFLINE=1 when

Content truncated.