Phase 1: correctness — every audited defect fixed, 16 xfails flipped#11
Open
AshishKumar4 wants to merge 5 commits into
Open
Phase 1: correctness — every audited defect fixed, 16 xfails flipped#11AshishKumar4 wants to merge 5 commits into
AshishKumar4 wants to merge 5 commits into
Conversation
…, DDIM eta, RK4, DPM reuse The sampler step domain was hardcoded to [0, 1000] and rescaled to the schedule's domain internally, so the default start_step=max_timesteps was wrong for every continuous schedule: the inference pipeline collapsed to a single no-op step at sigma_min and returned the initial noise. Steps are now in the schedule's own native domain and the defaults are correct by construction. Linear spacing on a Karras schedule is already the EDM rho spacing, so the broken quadratic/karras/exponential spacing modes (log(0), duplicate int16 steps) are gone. Also: DDPMSampler is now the algebraically exact ancestral posterior (handles batches and arbitrary strides; the posterior-table variant crashed on batch > 1 and ignored the stride), DDIM eta > 0 uses the paper formula with the missing direction-term reduction, RK4 no longer jits over a python callable, MultiStepDPM resets its history per trajectory, Heun falls back to euler when sigma reaches 0, and every sampler reshapes rates against the actual sample rank so 5D video works. All samplers are now verified end to end against an analytic gaussian oracle denoiser; the corresponding strict xfails are flipped to passing tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…, EDM weights, stable hashing - training.py silently ignored --noise_schedule and always trained EDM while inference honored the flag, so a cosine-flagged run trained one convention and sampled another. Both sides now build from a single get_diffusion_preset, with a test pinning their agreement. - UViT applied the hilbert permutation twice on the way in but inverted it once on the way out, training on a spatial scramble (SimpleDiT was correct; UDiT recomputed what patchify already returns). - The EDM loss weight's epsilon guard halved the weight at sigma_min ((sigma*sd)^2 == 1e-6 exactly); rewritten as 1/sd^2 + 1/sigma^2 which needs no guard. - FourierEmbedding frequencies came from jax's PRNG, which changed defaults in 0.5.0 and silently altered every model's time conditioning; now fixed numpy values, identical across jax versions. - arguments_hash used python hash(), randomized per process, so identical runs never mapped to the same experiment directory; now sha256. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ling, full state restore, checkpoint deletion - self.wandb was only assigned when wandb_config was set on process 0, so training without wandb (or on any non-zero process) crashed on the first logging tick; it is now always initialized and the getattr patches at the call sites are gone. - pbar was unbound on non-zero processes (UnboundLocalError at step 0 of any multi-host run since d50040d). - Checkpoint restore rebuilt the TrainState via create(), silently zeroing the adam moments and resetting the optax step counter (restarting lr warmup) on every resume. Restore now goes through orbax with the freshly initialized states as a typed template, so optimizer state, step counter and EMA params all round-trip. The unused param_transforms hook is gone. - The NaN 'recovery' block reset params and optimizer to a possibly epochs-old best_state mid-epoch and patched the loss to 1.0 to keep the average looking sane. A diverged run now fails loudly after 5 consecutive non-finite steps. - save() deleted the checkpoint it just wrote even when the registry push was skipped; deletion now only happens after a successful push. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ntion path, dead flags, bf16 output heads --dropout_rate was threaded into every DiT-family block and applied nowhere (any sweep over it measured noise). nn.Dropout now sits on both residual branches of DiTBlock/SSMDiTBlock/MMDiTBlock with a uniform train flag across every model, and the trainers thread a dedicated dropout rng. UViT's variant was never even forwarded, so that field is gone. --flash_attention is deleted along with EfficientAttention: the pallas call passed no sm_scale (inflating logits 8x at head_dim 64 vs the normal path), crashed on GPU/CPU where the import guard swallowed the dependency but the call remained, and used a different param tree so checkpoints weren't interchangeable. On every DiT-family model the flag was accepted and silently ignored anyway. A proper fused-attention path (cudnn/splash with identical params) comes with the parallelism work. Also: dead norm_groups/activation fields removed from the DiT family (training.py now scopes those args to the unet/uvit which read them), S5Layer validates its features field, parse_config filters config keys against the model's actual fields (old runs logged since-removed flags) and resolves dotted function paths instead of silently dropping them, the activation string is no longer mutated into a live function object inside the logged config, DiT/MMDiT/hybrid output heads are fp32 (the loss is fp32; bf16 heads quantized early-training residuals), the attention-block FFN no longer silently runs fp32 inside bf16 models, and jax.tree_map (removed in jax 0.6) is jax.tree.map. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The old example called generate_images(num_images=...) which never matched the actual signature, and showed a custom sampler with a stale take_next_step contract. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Second PR of the revamp series — stacked on #10 (review that first; this diff shows only Phase 1 commits).
What
Every defect from the audit, grouped into 5 commits:
Samplers (
fix(samplers)): the step domain was hardcoded to [0,1000] and rescaled internally, so the defaultstart_stepwas wrong for every continuous schedule —DiffusionInferencePipeline.generate_samples()collapsed to a single no-op step at sigma_min and returned the initial noise. Samplers now work in the schedule's native domain and the defaults are correct by construction. Also: DDPMSampler is the algebraically exact ancestral posterior (old one crashed on batch>1 and ignored strides), DDIM eta uses the paper formula (old crashed on.sqrt()and inflated variance), RK4 runs (jitted over a python callable before), MultiStepDPM resets history per trajectory (the pipeline caches samplers), Heun guards sigma=0, all samplers are 5D/video-safe, broken spacing modes deleted (linear on Karras is rho-spacing).Schedule presets:
training.pysilently ignored--noise_schedule(always trained EDM) while inference honored it — a cosine-flagged run trained one convention and sampled another, silently. Both sides now build from oneget_diffusion_preset, with a test pinning agreement. Plus: UViT double-Hilbert fix (it trained on a spatial scramble), EDM weight epsilon fix (halved the weight at sigma_min), FourierEmbedding frequencies pinned via numpy (jax 0.5.0's PRNG flip silently changed every model's time conditioning), sha256arguments_hash(was per-process randomized).Trainer: training without wandb crashed; non-zero processes crashed (
pbar); resume silently zeroed adam moments and restarted lr warmup (restore now goes through orbax with the initialized states as a typed template — opt_state, step counter, EMA all round-trip); the NaN "recovery" that reset params mid-epoch and cosmetically patched loss=1.0 is replaced by a loud failure after 5 bad steps;save()no longer deletes checkpoints it didn't push to the registry.Models:
--dropout_rateis now actually implemented (was threaded everywhere, applied nowhere — sweeps over it measured noise);--flash_attentionandEfficientAttentionare deleted (no sm_scale → 8x inflated logits, crashed off-TPU, incompatible param tree; a real cudnn/splash path comes with the parallelism PR); dead fields removed; fp32 output heads for bf16 models;parse_configfilters stale config keys and resolves dotted function paths instead of silently dropping them.Verification
16 of Phase 0's 17 strict xfails are flipped to passing tests (the 17th, non-square unpatchify, is fixed by the Phase 2 backbone rewrite). Every sampler now converges against the analytic gaussian oracle, including video shapes and the no-argument default path. Suite: 56 passed, 1 xfailed. Grain iterator checkpointing (mid-epoch data resume) moved into Phase 2's trainer merge so save/load is rewritten once.
🤖 Generated with Claude Code