Tatzmori/orthdion-audit-ckpts

Dion assumption-audit checkpoints (checkpoint_v1)

Checkpoints + scalars for the spectral assumption audit (Δgap, τ/η, κr, ρt) and the νt / precision study. Private; do not redistribute.

Runs

The normalization column tells you what produced each pack's stored right_factor_t_minus_1 (= the optimizer's live Q = previous step's V̄); it is also recorded per pack in meta["normalization_method"].

Attribution pair: orthdion_r384_seed0 ↔ dion_colnorm_r384_seed0 run the same orth_dion code path with the same config; the single difference is the ortho_dion_method flag (cholesky vs colnorm). Any ε̂ difference between these two packs is attributable to the V̄ geometry alone. (dion_r384_seed0 is the upstream-module run kept for fidelity with the paper's published Dion trajectory; it is algorithmically identical to dion_colnorm_r384_seed0 but a different code path.) Sanity check: the QR arm's fp64-replay νt = 1.0 to 14 decimals at every checkpoint; the ColNorm arms' νt > 1.

Each mirrors its original's exact module, config, and CLI args (verified by script diff), plus checkpointing. The r=96 original ran on the orth_dion module in ColNorm mode (that's how its ν_t was logged); same math as upstream Dion, different code path — kept identical here for trajectory fidelity.

Common setup: LLaMA3 320M (18 layers, dim 768), c4, GBS 256, LBS 32, seq 2048, seed 0, lr 0.012 (Dion & scalar AdamW), wd 0.1, warmup 610 + cosine→0, steps 6104, bf16 mixed precision, FSDP shard 8 (GH200 ×8), compile. Training live-logs to wandb llm_jp_pp/checkpoint_v1.

Layout

<run>/step-<N>/audit_pack_step-<N>.pt      N ∈ {610, 1220, ..., 6100, 6104}
<run>/step-<N>/derived_scalars_step-<N>.csv
<run>/step-<N>/dcp_step-<N>.tar
<run>/scalars.csv

audit_pack_step-<N>.pt — primary artifact

torch.load(path, map_location="cpu", weights_only=False) → dict with keys meta, params, derived.

params[fqn] for each of the 126 transformer matrices (fqn like layers.3.attention.wq.weight; 72× (768,768) attn, 36× (2048,768) w1/w3, 18× (768,2048) w2). All tensors fp32 — lossless upcasts of the bf16 training state. Notation follows Algorithm 1 (Orth-Dion paper); t = checkpoint step:

The step's full input/output set is pinned down: lines 2–5 (Ut, V̄t, D̂t) are deterministically recomputable from `momentumt and rightfactortminus1` (orthonormalization is unique for full-rank input).

derived[fqn] — fp64 offline replay of Alg.1 L2–L5 (QR for Ut; V̄t by the run's own normalization, see meta["normalization_method"]), computed at export time: inner_M_Dhat = ⟨Mt, D̂t⟩, cos_M_Dhat, nu_t = max(σmax(V̄t), ‖V̄t‖F/√r), M_fro, Dhat_fro, G_fro, R_fro, R_next_fro.

Orientation: FSDP shards dim 0 of every param, so upstream Dion runs every matrix in transposed mode — Ut = orth(Mtᵀ V{t−1}), with V{t−1} of shape (dim0, r). The replay applies this convention unconditionally and asserts ‖D̂t‖F = √r per matrix (an identity under both ColNorm and QR). Derived values regenerated 2026-07-30 after fixing a square-matrix orientation bug (validated against the r=96 run's inline ν_t: ≤0.14% rel. difference).

meta — provenance, field semantics, dtype note.

derived_scalars_step-<N>.csv

The derived dict as CSV: one row per matrix, same columns. For plotting without loading tensors.

scalars.csv — dense trajectory (one per run)

Frozen copy of the run's full wandb history; one row per training step (~6104 rows; loss is logged every step, the other columns populate every 10th row). Columns: train/val loss, lr (ηt), global grad-norm, per-matrix `weights/<fqn>{g,r,m}normF = ‖G_t‖_F, ‖R_t‖_F, ‖M_t‖_F (captured pre-step, exact), aggregates z/optimizer/{grad,efresidual,momentum}normF/{mean,max,min,std}`, and for the r=96 run additionally per-matrix νt (its config logs ν inline).

dcp_step-<N>.tar — raw checkpoint

The sharded torch.distributed.checkpoint dir torchtitan wrote (8 .distcp shards + .metadata): model, full optimizer state (the six buffers above under state.<fqn>.<key>, plus Q, momentum, AdamW states), lrscheduler, dataloader, trainstate. Only needed for exact restore/resume of training; the audit pack is derived from it.

What lives where (summary)

Precision caveat

Training state (W, momentum, Q) is stored in bf16; the optimizer upcasts to fp32 only for Gram accumulation. Audit-pack fp32 tensors are exact copies of the bf16 values — measurements on them reflect what training actually saw (orthogonality error floor ≈ 3e-3 from bf16 rounding; breakdown κ ≈ 1e3). The derived scalars are fp64 replays on those bf16-valued inputs.