Phase 2: consolidation — one registry, one backbone, real MM-DiT, one trainer#12
Open
AshishKumar4 wants to merge 6 commits into
Open
Phase 2: consolidation — one registry, one backbone, real MM-DiT, one trainer#12AshishKumar4 wants to merge 6 commits into
AshishKumar4 wants to merge 6 commits into
Conversation
- favor_fastattn.py (722 lines, zero importers, hard-depends on undeclared gin/absl, contains a live typo) - unet_3d.py + unet_3d_blocks.py (950 lines, unreachable since the models/__init__ import was commented out; video support will be rebuilt from the shared backbone) - vit_common: AdaLNZero and PositionalEncoding (nothing constructs them) - common: TimeEmbedding (superseded by FourierEmbedding, NaN for features=2) and PixelShuffle (only reference was a commented line) - utils: RMSNorm (nn.RMSNorm is used everywhere), the never-called map_nested_config that silently dropped every non-string value, and the string maps that fed it - trainer: the duplicated form_global_array/convert_to_global_tree copy (the flaxdiff.utils one is the survivor), dead tensorboard stubs, write-only last_save_time, a 39-line commented-out block - simple_vit: an if/else with identical branches and the dead use_remat flag Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…uction training.py and the inference pipeline each carried their own architecture table, their own dtype/precision/activation string maps (three copies total, already drifted: a 'highes' typo on one side, extra keys on another), and their own +2d/+hilbert/+zigzag suffix parser. A model was therefore constructed through two different code paths at train and inference time. flaxdiff.models.registry now owns the architecture table, the string mapping, suffix canonicalization, and field-filtered construction. Model configs stay as plain strings all the way into the logged wandb config, and build_model resolves them identically on both sides - the logged config IS the construction input. The hand-maintained --architecture choices list (which had already drifted from the inference table) is gone; the registry validates. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The patchify/scan-order/sincos/conditioning/final-head skeleton was repeated near-verbatim in SimpleDiT, HybridSSMAttentionDiT, SimpleUDiT and SimpleMMDiT, and the adaLN-Zero block body was copy-pasted three times (attention, SSM, and dual-cond variants). dit_common.py now owns the sandwich (PatchSequenceEmbed / ConditioningEmbed / PatchSequenceOutput) and a single ModulatedBlock with a pluggable mixer (attention | S5 ssm); the S5 layers move to s5.py so the block can use them without importing the full hybrid model. SimpleDiT (291 -> 90 lines) and HybridSSMAttentionDiT (597 -> 142) are now just block arrangements; SimpleUDiT builds its U-skips from the same block. Also: unpatchify takes an explicit patch grid so non-square inputs work (the last remaining xfail from the audit - the whole suite is now xfail-free), and attention gains optional SD3-style qk-norm (RMSNorm on q/k) for bf16 logit safety, exposed on every DiT-family model. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
SimpleMMDiT was MM-DiT in name only: text was mean-pooled into the adaLN vector and never entered the token sequence, attention and MLP ran in parallel off a single norm (verified by ablation - zeroing attention changed the MLP contribution by ~2e-6), so it was a text-modulated DiT at roughly half its nominal depth. The block is now a true SD3-style dual-stream MM-DiT: per-modality qkv/mlp/modulation weights, joint attention over concat([txt, img]) with RoPE on the image tokens only, sequential attn -> mlp, and an adaLN-modulated final layer. HierarchicalMMDiT is ported onto the same block (raster only - its merge/expand grid reshapes assume row-major order, so the old hilbert option was scrambling the neighborhoods being merged). AdaLNParams gains the missing SiLU before its zero-init projection, per the DiT paper. A behavior test pins the dual-stream property: changing text tokens must reach the image stream through joint attention. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…iner GeneralDiffusionTrainer overrode every method the legacy DiffusionTrainer middle class added except TrainState/generate_states/fit - the middle class was a fork whose train/validation code was dead (its validation step could not even construct a sampler post the API change) yet still shipped. Those three live pieces now sit directly on GeneralDiffusionTrainer, the middle class and the broken AutoEncoderTrainer (a copy-paste of the diffusion step referencing attributes that were never set) are deleted, and SimpleTrainer keeps only generic infra with abstract step definitions instead of toy classification steps no real path used. The trainer smoke tests now run the real GeneralDiffusionTrainer end to end (tiny DiT, unconditional synthetic data): fit without wandb, and a save/restore round-trip asserting step counter, adam moments and EMA params all survive. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Remaps the top-level module paths the backbone consolidation moved (embed/, conditioning/, output/) so pre-consolidation DiT-family checkpoints load against the new models. UNet/UViT need no conversion; MMDiT checkpoints are refused implicitly since the architecture itself changed. Documents the FourierEmbedding frequency provenance caveat for pre-jax-0.5 checkpoints. 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.
Third PR of the revamp series — stacked on #11 (which stacks on #10). Review in order; this diff shows only Phase 2 commits.
What
Deletions first (
chore: delete dead code, −1,984 lines): favor_fastattn (722 lines, unimportable), unet_3d* (950 lines, unreachable), dead layers (AdaLNZero, PositionalEncoding, TimeEmbedding, PixelShuffle, utils.RMSNorm), the never-called brokenmap_nested_config, duplicated global-array helpers, dead tensorboard stubs, an if/else with identical branches.One model registry: training.py and inference each had their own architecture table, three drifted copies of the dtype/precision/activation maps (one had a
'highes'typo), and two independent+2d/+hilbert/+zigzagparsers — a model was constructed through different code paths at train and inference time.flaxdiff.models.registrynow owns all of it; configs stay plain strings into the wandb log andbuild_modelresolves them identically on both sides. The hand-drifted--architecturechoices list is gone.One DiT backbone: the patchify/scan/sincos/conditioning/output sandwich was repeated near-verbatim across 4 model files, and the adaLN-Zero block was copy-pasted 3×. Now:
dit_common.pyowns the sandwich + a singleModulatedBlock(mixer = attention | S5),s5.pyowns the SSM layers. SimpleDiT went 291→88 lines, HybridSSMAttentionDiT 597→139. Non-square patch grids now work (the audit's last xfail — the suite is now xfail-free), and attention gains optional SD3-style qk-norm.Real MM-DiT: the old SimpleMMDiT mean-pooled text into the adaLN vector (text never entered the token sequence) and ran attn+MLP in parallel off one norm — a text-modulated DiT at half its nominal depth, verified by ablation. Now a true dual-stream block: per-modality qkv/mlp/modulation, joint attention over concat([txt, img]), RoPE on image tokens only, adaLN-modulated final layer. HierarchicalMMDiT is ported onto it (raster-only: its merge/expand reshapes assume row-major order, so the old hilbert option was scrambling the merged neighborhoods).
AdaLNParamsgains the SiLU the DiT paper specifies. A behavior test pins the dual-stream property.One concrete trainer: the legacy
DiffusionTrainermiddle class was a fork whose train/validation overrides were dead (its validation couldn't even construct a sampler) — only TrainState/generate_states/fit were live, and those now sit onGeneralDiffusionTrainerdirectly. The brokenAutoEncoderTrainer(copy-pasted diffusion step referencing attributes that were never set) is deleted. Trainer smoke tests now run the real trainer end to end and assert step counter, adam moments and EMA all survive a save/restore round-trip.Conversion script (
scripts/convert_legacy_checkpoint.py) remaps old DiT-family checkpoints to the new param layout.Scope notes (deviations from the plan, deliberate)
Objectiveseam is deferred to the JEPA PR where its second implementation exists — an interface with one implementer fails the deletion test today.Verification
pytest -m "not network": 61 passed, 0 xfailed — every bug documented in Phase 0's strict xfails is now fixed or the feature deleted. Net diff for the phase: +1,331 −4,096.🤖 Generated with Claude Code