Phase 5: JEPA + the Objective seam, integrated onto FSDP - #16
Closed
AshishKumar4 wants to merge 11 commits into
Closed
Phase 5: JEPA + the Objective seam, integrated onto FSDP#16AshishKumar4 wants to merge 11 commits into
AshishKumar4 wants to merge 11 commits into
Conversation
…roblem The train step hardcoded diffusion: normalize, VAE encode, CFG dropout, sample timesteps, corrupt, apply, weight. Nothing about sharding, EMA, checkpointing or the loops needs to know any of that, and JEPA needs a different loss over several parameter subtrees with an EMA that tracks only one of them. An Objective now owns what is being learned - init_params, loss, an EMA spec and the validation artifacts - and the trainer owns the mechanics. The loss always runs under has_aux, so an objective can report per-component diagnostics; the train loop forwards them to wandb. The EMA spec carries a step-indexed decay schedule and a subtree path, because I-JEPA ramps its momentum and averages only its context encoder. DiffusionObjective is the existing objective moved behind that interface, verified bit-identical: params, EMA and adam state after five real steps match the inlined implementation exactly, and that fingerprint is now a test. Two fixes fell out of the move. The video check read len(sample_data_shape) == 5 on a shape that never carries the batch axis, so it was always false and video validation sampled with sequence_length=None. And unconditional ModulatedBlocks now exist: conditioning=None gives a plain pre-norm residual block with learned affine norms, which is the ViT block a JEPA encoder needs. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Three changes to the training objective layer, sharing one contract. Flow matching slots into the existing scheduler/predictor seams rather than next to them: FlowMatchingScheduler is the linear path alpha = 1 - t, sigma = t with SD3's logit-normal timesteps and Flux's resolution shift, and FlowMatchPredictionTransform predicts the velocity u = eps - x0. Because the path is linear, the unmodified DDIM and Euler updates already are the flow Euler step; there is a test proving the identity exactly. min-SNR-gamma (Hang et al. 2023) is defined on the x_0 loss, so converting it into the space the trainer actually computes the loss in needs the parameterization. get_weights becomes a template over get_schedule_weights and asks the transform for that conversion, which also makes the schedule/transform pairing explicit instead of implicit. No second weighting system. EDMNoiseScheduler's P_mean/P_std were hardcoded to EDM1's -1.2/1.2; they are now constructor arguments defaulting to EDM2's -0.4/1.0. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Guidance costs diversity at high noise and buys nothing at low noise (Kynkaanniemi et al. 2024). guidance_start/guidance_stop bound the fraction of the trajectory over which it applies; outside the interval the scale drops to 1, which is exactly the plain conditional prediction. The default interval is the full range, so existing samplers are unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The jax-fid port in metrics/inception.py has been sitting unused. Wire it into an EvaluationMetric the same way metrics/images.py builds the CLIP ones, sharing a single cached copy of the weights. Per-batch FID is noisy and only meaningful as a trend across checkpoints, which the docstring says outright. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Latents are normalized as (z - latent_shift) * latent_scale and inverted on decode, the SD3 convention, so a caller can hand in a dataset's own mean and 1/std instead of relying on the VAE's single fixed constant. The SD VAE's scaling_factor now rides on that same seam rather than being applied inside its jitted frame functions, so there is one normalization path and the defaults reproduce the previous behavior exactly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
--noise_schedule gains flow/flow_matching, plus --min_snr_gamma, --flow_shift, and 'fid' as a validation metric. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The encoders are arrangements, not new transformers. JepaEncoder is PatchSequenceEmbed plus a stack of unmodulated ModulatedBlocks; JepaVideoEncoder is the same blocks in VideoDiT's factorized spatial-temporal layout; the predictor is a narrower stack of the same blocks over context embeddings, mask tokens and target positions. mixer patterns carry straight through, so an encoder can be all-attention, all-S5 or any ratio between - the hybrid SSM encoder is a config string, not a second implementation. Masking follows the paper's multi-block strategy, resolved against the static patch grid at construction so every step produces identically-shaped masks under jit. The block area is the value inside the scale range with the most (h, w) factorizations in the aspect range, so shape and position stay random per sample; the context is a random subset of the complement of the target union, sized so it is always available. Video masks are tubelets: one spatial mask applied to every frame, which leaves the factorized layout intact. Position comes from the 2D sincos signal alone, with RoPE at identity: a masked subsequence has no meaningful 1D index, and sincos travels with the token when it is gathered. Time is untouched by masking, so the temporal half of the video stack keeps real RoPE. Collapse telemetry ships with the loss rather than after it. repr_std and repr_cov_offdiag go into the aux dict on every step, and the trainer already forwards aux to wandb, so a run that quietly agrees on a constant shows it in the curves. Both are tested against a degenerate encoder, not just asserted. Also here: frozen-encoder linear and k-NN probes as EvaluationMetrics that the existing validation loop runs over the target encoder's embeddings, the models in the registry, and training_jepa.py. Two fixes fell out. The validation loop's numpy import lived inside an `if i == 0 and wandb` branch while np.mean ran outside it, so any run with eval metrics and no wandb died in validation. And the wandb artifact name was hardcoded to "diffusion-<dataset>-res<n>", which would have collided across objectives; it now takes the objective's tag. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Phase 3 rewrote how training executes; phase 5 rewrote what it optimizes. Both rewrote _define_train_step, so the merged step keeps phase 3's mechanics and gets its loss from phase 5's Objective seam. The step is still jax.jit with in_shardings/out_shardings and donate_argnums over the state, and still returns the in-step is_finite flag the loop uses to detect divergence without a host sync. What changed is the body: instead of inlining the diffusion maths it calls Objective.loss through value_and_grad with has_aux, and returns the objective's aux dict alongside the loss, which the train loop folds into its wandb logging next to the throughput metrics. Signature is now (state, loss, aux, rng_state, is_finite). EMA moves to phase 5's EMASpec: a subtree path plus a step-indexed decay schedule, read at train_state.step. Diffusion declares a constant schedule over the whole tree, so its behaviour is unchanged; JEPA tracks only the context encoder with a momentum ramp. The EMA copy still inherits the params' sharding through the spec tree, so averaging a subtree stays device-local. generate_states routes params through Objective.init_params but still materialises them via _build_state, so the sharding constraints come from jax.eval_shape over the objective's own tree and neither the objective nor the model has to describe a layout. Dropped from phase 5: shard_map, the per-device RNG fold-in, the hand-written pmean over grads and loss, and donation of the batch - all superseded by phase 3's GSPMD path. Dropped from HEAD: the inlined diffusion loss, the sampler-specific validation step and generate_modelname, all now owned by DiffusionObjective. The trainer no longer mirrors the objective's configuration (noise_schedule, autoencoder, ema_decay, conditional_inputs, frames_per_sample were all write-only after the merge). Two latent bugs in HEAD's validation_loop fall out of taking phase 5's version: an undefined global_device_count, and np used without an import. Its is_video test compared len(sample_data_shape) against 5 on a shape that excludes the batch axis, so video runs validated as images; the objective tests the same shape against 4. Tests: the golden fingerprint in test_objectives was captured before phase 3 changed the step's RNG derivation, so it is re-baselined against HEAD's inlined step - the merged objective reproduces it exactly, which is what makes the refactor provably behaviour-preserving. The EMA tracking test kept a live reference to the train state across fit, which donation now deletes; it snapshots to the host instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…sion JEPA over a 4x2 mesh: 21 of the objective's 59 parameter leaves land split over the fsdp axis, the EMA target encoder inherits their specs, and two real steps come back finite with the collapse telemetry intact. Compiling was never the claim - this is what shows the two halves of the merge compose. The mixed-precision branch of the step was rewritten to thread has_aux and had no coverage on either side of the merge, so it gets a test that an objective's aux survives it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Owner
Author
|
Landed on |
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.
Final PR of the revamp series — stacked on #15. Contains Phase 5 plus the semantic merge that reconciles it with Phase 3. (Until #14 merges, GitHub will also show Phase 4's delta here; that collapses as the stack lands.)
The Objective seam
GeneralDiffusionTrainerhardcoded the diffusion objective in its train step. There is now anObjectiveABC —init_params(rng),loss(...) -> (scalar, aux)always withhas_aux,make_validation_step,log_validation_artifacts, plusEMASpec(decay: Callable[[step], value], path). Five members, and it earns its place by having two real implementations:DiffusionObjectiveand JEPA. The trainer owns mechanics (sharding, EMA bookkeeping, checkpointing, loops); the objective owns what is learned.The step-indexed decay schedule isn't speculative generality — I-JEPA's momentum ramps 0.996→1.0, and the subtree
pathis what lets JEPA EMA only its context encoder while diffusion EMAs the whole tree.I-JEPA and V-JEPA
Built by composing
dit_commonblocks — no second attention implementation. Multi-block masking per the paper (4 targets, scale 0.15–0.2, aspect 0.75–1.5, context = complement), L2 on layer-normalized targets, stop-gradient + EMA target encoder. V-JEPA adds tubelet masking on the existing video machinery. Linear and k-NN probes ship asEvaluationMetrics.Collapse telemetry from day one:
repr_stdandrepr_cov_offdiagin the aux dict — the second catches dimensional collapse, which representation-std alone misses. Three tests prove the metrics actually detect collapse rather than merely computing.test_encoder_runs_on_the_ssm_mixercovers the hybrid SSM encoder — the novel research angle.The merge: mechanics from Phase 3, objective from Phase 5
Both branches rewrote
_define_train_stepfor different reasons. The resolution keeps Phase 3's frame —jax.jitwithin_shardings/out_shardings,donate_argnums=(0,), noshard_map(still 0 repo-wide), one partitionable key per step, in-stepis_finite— and takes the body from the objective viavalue_and_grad(..., has_aux=True). EMA goes throughEMASpecwhile the EMA copy still inherits Phase 3's sharding specs, so subtree averaging stays device-local.Behaviour-preservation was proven, not assumed. The golden fingerprint in
test_objectives.pywas stale (captured against a per-devicefold_inthat Phase 3 deleted). Rather than re-record blindly, the fingerprint was captured from HEAD's inlined step before any code changed, then the merged objective-based step was asserted to reproduce it — matching to the original1e-9/1e-12tolerances, with an identical epoch loss (0.7228752374649048).Two real bugs in Phase 3 that this fixes
I verified both against the
phase-3-parallelismbranch:validation_loopreferenced an undefinedglobal_device_count— aNameErroron any validation run._is_video_datacomparedlen(sample_data_shape) == 5, but that shape excludes the batch axis, so it was always False and video runs would have validated as images.If #15 is reviewed standalone, note these are live there and fixed here.
Verification
185 passed (from 152 at the merge base), re-run independently by me. New:
test_jepa_trains_under_fsdp— a 4×2 mesh where 21 of 59 parameter leaves genuinely split over the fsdp axis (each shard exactly half the global param), the EMA target carries identical specs, and two real steps return finite losses with telemetry intact. Compiling was never the claim.Honest gaps
process_count > 1) is untested — CPU simulation is one process with N devices. This is the one thing that needs a real TPU pod slice to validate.train_state.step, which increments per micro-batch, so under--grad_accum_stepsthe momentum ramp advances per micro-step rather than per optimizer update. Inherited behavior, but more visible now that decay is a schedule — worth a decision before running JEPA with accumulation.🤖 Generated with Claude Code