This repository was archived by the owner on Jul 17, 2026. It is now read-only.
refactor(components): owner-partitioned stacked persistence for V/U (flagship muon sharding)#989
Draft
ocg-goodfire wants to merge 2 commits into
Draft
refactor(components): owner-partitioned stacked persistence for V/U (flagship muon sharding)#989ocg-goodfire wants to merge 2 commits into
ocg-goodfire wants to merge 2 commits into
Conversation
…PEC D4) The check predated the S18 amendment that made the model input token ids (Int[B,T] -> embed_tokens indexing); its float32 'resid' batch crashed clean_output with 'Indexer must have integer or boolean type'. Draw random int tokens like the llama8b tests do. Same failure family as the slow-eval hidden_acts crash fixed by #919. Verified: XLA_FLAGS=--xla_force_host_platform_device_count=4 passes (worst rel 1.02e-04 within tol). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…/U masters DecompVU persists as same-shape stacks — one (Vs [g,d_in,C], Us [g,C,d_out]) pair per shape group + a static site->(shape,slot) index — instead of 2*n_sites per-site leaves. Placement (hybrid HSDP rule): stack axis /replicate (whole matrices owned per node-group: zero cross-node weight collectives per step, muon NS node-local), matrix d dims /fsdp, C /tp — total /N, same memory as the retired intra-matrix ZeRO-1; per-group fallback to intra-matrix sharding when a stack doesn't tile replicate (site subsets). SPEC D4 amended 2026-07-15 (Oli-requested). - llama8b _stack_per_kind_vu: contiguous-slot fast path = a static slice of the resting stack (the per-step 32-way restack disappears for full-model runs); _per_kind_dims reads shapes statically off site_slots. - init_decomp_vu_placed: one jit, 2*n_shapes outputs (the stack/unstack fan-out and its transient copy are gone). - muon: the V/U tree is all-3D now, so BOTH optimizer groups label leaves via stacked_muon_dimension_numbers (renamed from chunk_stacked_*; optax's default 2D rule would silently Adam every V/U leaf). - decomp_vu_from_sites: explicit per-site-dict constructor (toys/tests). - cherry-pick f63dcd5 (invariance_check int tokens — the fix was stranded on its bridge branch; harness needed here). Validated: full suite 535 passed + multidevice 5 passed; invariance_check at 4 sim devices OK (worst rel 9.6e-6, reassociation-only — the D4 anchor for the layout change). Known follow-ups: one-off checkpoint layout migration for pre-change runs; grad_norms per-leaf wandb keys now path through stacks. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SwEM9ZetQw9W4yjH8okwaH
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 subscribe to this conversation on GitHub.
Already have an account?
Sign in.
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.
Design B from the flagship muon-sharding discussion (bridge: component-norm-optimizer): the V/U masters (+ optimizer moments) persist as same-shape stacks —
(Vs [g, d_in, C], Us [g, C, d_out])per shape group + a staticsite_slotsindex — replacing 2·n_sites per-site leaves.Placement (hybrid HSDP rule): stack axis ÷
replicate(whole matrices owned per node-group → zero cross-node weight collectives per step; muon NS node-local), matrix dims ÷fsdp(NVLink), C ÷tp. Total ÷N — same memory as intra-matrix ZeRO-1. Per-group fallback to intra-matrix sharding when a stack doesn't tilereplicate(site subsets, e.g. L18-only). Atreplicate=1the layouts coincide (zero change single-node).Why: at 32L/dp32 the per-kind stacks the scan already consumes become static slices of the resting representation (the per-step 32-way restack disappears); ENTRY's cross-
replicategather and the backward's grad reduce land in collectives that must exist anyway; muon's NS runs on whole owned matrices with no cross-node reshard — the Kimi owner-partitioned recipe expressed GSPMD-natively, with the trainer's existing structures.Muon interplay: V/U leaves are now 3D, so both optimizer groups use
stacked_muon_dimension_numbers(renamed fromchunk_stacked_*) — optax's default 2D rule would silently Adam every V/U leaf.Validation: full suite 535 passed / 11 xfail; multidevice 5 passed (placement + bit-parity init pinned to the new rule);
invariance_checkat 4 sim devices OK (worst rel 9.6e-6 — the SPEC D4 anchor: same math, different layout). SPEC D4 amended (2026-07-15, Oli-requested). Includes cherry-pick of f63dcd5 (invariance_check int-tokens fix that was stranded on its bridge branch).Known follow-ups before landing: one-off checkpoint layout migration for pre-change runs (restore-onto-reference means old per-site checkpoints don't open at tip;
migrate_c49k_checkpoint.pyis precedent);grad_norms/components*per-leaf wandb keys now path through stacks (summaries unchanged); CI-fn masters deliberately keep intra-matrix ZeRO-1 (their owner-partitioning is a separate decision).Stacked on #982 (base:
feature/muon).Bridge thread: component-norm-optimizer
🤖 Generated with Claude Code
https://claude.ai/code/session_01SwEM9ZetQw9W4yjH8okwaH