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 justfile test recipes). - 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: ```python 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 | Comparison | What it tells you | |---|---| | mflux native seed vs itself (re-run) | Sampling is deterministic on this hardware | | mflux + diffusers latent vs diffusers full run | Transformer + VAE + CFG path quality | | mflux native seed vs diffusers full run | **Not** a parity test — different RNG (`mx.random` vs `torch.Generator`) and often different sigma schedules | | mflux + diffusers latent vs mflux native seed | Large diff is **expected** if schedulers/sigmas differ, even with identical starting noise | 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 weights are cached to catch missing components the reference pipeline expects but mflux omits from its download patterns. - Compare at matched resolution, steps, guidance, seed, and precision when possible (note mflux `-q 8` vs diffusers bf16 is not apples-to-apples for speed or pixels). - Save outputs to explicit paths (e.g. Desktop) and report wall time with `/usr/bin/time -p`. ## See Also - `mflux-model-porting`: correctness-first workflow (validate components and lock behavior before refactor). - `mflux-testing`: how to run tests safely and handle image outputs/goldens.