perf(compile): stacked V/U init — 13x faster jit_init_decomp_vu compile#951
Closed
danbraunai-goodfire wants to merge 1 commit into
Closed
perf(compile): stacked V/U init — 13x faster jit_init_decomp_vu compile#951danbraunai-goodfire wants to merge 1 commit into
danbraunai-goodfire wants to merge 1 commit into
Conversation
…omp_vu init_decomp_vu_placed compiled a single graph with 2 x n_sites sharded outputs (448 at 32L); the SPMD partitioning/layout pass over that many distinct outputs was a 3m31s compile at 32L/dp128 (210.5s -> 16.3s measured on GPU at the full 224-site shape set). It now compiles two cheap stages: the RNG runs vmap-STACKED per V/U shape (init_decomp_vu_stacked: 2 x n_shapes sharded outputs), then a trivial slice jit (unstack_decomp_vu, stacked input donated) fans out to the per-site ZeRO-1 layout. BIT-identical to the previous jitted per-site init (vmap over the same per-site keys); test_sharding now pins exact equality against the jitted per-site reference (the unjitted eager reference differs from ANY jitted path by ~1 ULP of XLA fusion — pre-existing, previously hidden by allclose). Extracted from the bridge task branch (PR #941). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01E34cK5BaVr6ZFAwRidNp6n
Collaborator
Author
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
2 participants
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.
Description
Targeted extraction of the V/U-init compile fix from #941 (bridge task
reduce-jax-compilation-time), leaving the rest of that branch behind.init_decomp_vu_placedcompiled a single graph with 2×n_sites sharded outputs (448 at 32L); the SPMD partitioning/layout pass over that many distinct outputs was the 3m31sjit_init_decomp_vucompile on every new-config launch. It now compiles two cheap stages:init_decomp_vu_stacked— RNG vmap-STACKED per V/U shape (2×n_shapes sharded outputs, 14 at 32L), sharded per-site behind a leading unsharded stack axis.unstack_decomp_vu— a trivial slice jit (stacked input donated) fanning out to the per-site ZeRO-1 layout.Bit-identical to the previous jitted per-site init (vmap over the same per-site keys);
test_shardingnow pins exactarray_equalagainst the jitted per-site reference. (Found in passing: the unjitted eagerinit_decomp_vudiffers from ANY jitted init by ~1 ULP of XLA fusion — pre-existing, previously hidden by the test'sallclose.)Related Issue
Extracted from #941; measurements from bridge task
reduce-jax-compilation-time.Motivation and Context
Every new-config launch of the 32L Llama-8B decomposition paid a 3m31s init compile before the train-step compile even started. Measured on GPU at the full 224-site/32L shape set: 210.5s → 16.3s (12.9×).
How Has This Been Tested?
make checkclean; sharding tests pass at 1 and 4 simulated CPU devices (pytest param_decomp/tests/test_sharding.py --runmultidevice).array_equal, sharded + unsharded).Does this PR introduce a breaking change?
No. Init values are bit-identical.
🤖 Generated with Claude Code
https://claude.ai/code/session_01E34cK5BaVr6ZFAwRidNp6n