Phase 3: real FSDP, jit+NamedSharding, measurement, mid-epoch resume

[AshishKumar4]

Sixth PR of the revamp series — stacked on #14. Merges cleanly onto it (verified).

FSDP is real now

Before this PR the mesh was a single 'data' axis and the train step passed P() for the train state — parameters, EMA and both Adam moments were replicated on every device. I verified that empirically before and after; same inspection, opposite result:

fsdp_size	params ≥256 elems sharded	optimizer-state leaves sharded
1 (default)	0 / 13 — PartitionSpec()	0 / 26
4	13 / 13 — PartitionSpec('fsdp',)	26 / 26

The mesh is now ('data','fsdp'), and the parameter spec tree is derived once in generate_states from jax.eval_shape plus a size heuristic (--fsdp_min_param_size); optimizer state and EMA inherit the param specs through tx.init. Zero model-file changes — exactly as designed. fsdp_size=1 reproduces the old replicated behavior.

shard_map is gone

grep -rn shard_map flaxdiff/ returns 0. The train step is jax.jit with in_shardings/out_shardings and donation; GSPMD derives the collectives, so the hand-written pmean, the local_device_index argument and the fold_in device-index RNG hack are all deleted (jax_threefry_partitionable already defaults True on jax 0.9).

Two parity gates guard the migration, both runnable on CPU with simulated devices: single-device vs multi-device loss trajectories agree, and FSDP vs replicated agree.

Measurement, which did not exist at all before

train/step_time_ms, train/samples_per_sec and an MFU estimate now go to wandb; --profile_steps writes a jax.profiler trace; --compilation_cache_dir enables the persistent XLA cache (the biggest single win for restart-heavy TPU workflows).

Async dispatch unblocked

The old loop did if not jnp.isfinite(loss) every step — a Python branch on a device array, forcing a host sync that defeated JAX's async dispatch. is_finite is now an output of the jitted step and the consecutive-bad-step counter stays on device (jnp.where), checked on host only at the logging cadence. Batch H2D moved into a prefetch iterator so transfer overlaps compute, and donate_argnums now donates the state (the large buffer) rather than the batch, which required breaking the best_state aliasing first.

Checkpointing caught up with the sharding

Sharded orbax save/restore with restore_args built from the target sharding tree — replacing a full host gather (np.array on every leaf) that immediately called wait_until_finished(), defeating its own AsyncCheckpointer. Restore onto a different mesh is tested. Grain iterator state is checkpointed too, so resume is finally mid-epoch-correct — previously the data iterator restarted from the top of the epoch on every resume.

Capacity

--remat rematerializes transformer blocks (dots_with_no_batch_dims_saveable) — this sets your maximum trainable model size. --grad_accum_steps adds optax.MultiSteps, tested to hold params still until the accumulation boundary and to survive the sharding heuristic.

Verification

113 passed (from 73), re-run by me independently. tests/test_parallelism.py adds 19 tests covering the sharding heuristic, mesh construction, prefetch ordering/error propagation, both parity gates, FSDP sharding proof, checkpoint roundtrip, cross-mesh restore, gradient accumulation and mid-epoch resume.

Not done: the unified precision policy (deliberately scoped last and left out; the ad-hoc dtype= threading remains).

🤖 Generated with Claude Code

[AshishKumar4]

Landed on main as part of the full revamp stack (main fast-forwarded to bc9d06b). GitHub reports no new commits between this head and main, so every commit here is merged; closing as superseded by the stack merge rather than leaving it open.