Skip to content

perf: Tune ix shard counts and recalibrate for proving via --ram-budget - #601

Closed
samuelburnham wants to merge 1 commit into
mainfrom
ap/aiur-cost-dump
Closed

perf: Tune ix shard counts and recalibrate for proving via --ram-budget#601
samuelburnham wants to merge 1 commit into
mainfrom
ap/aiur-cost-dump

Conversation

@samuelburnham

Copy link
Copy Markdown
Member

Aiur cost profiling and budget-aware sharding

Summary

This branch adds execution-derived per-constant Aiur cost profiles and uses
measured prover-RAM estimates to make sharding self-correcting. Over-budget
shards are split before proving, audits can emit the partition they actually
validated (including measured peaks), and the aiur-sharded-env benchmark now
audits the same 100-GiB budget used to seed its manifests.

The CLI distinguishes the local prove limit (ix prove --max-ram, defaulting
to 85% of available RAM) from the target machine used by an execution-only
audit (ix check --ram-budget). Static ix shard --max-ram G sizing now uses
the environment's block-size distribution instead of raw .ixe file size.
Explicit --shards N behavior is unchanged.

Updated shard model

For each atomic block b, let s_b be its serialized owned bytes (not .ixe
file bytes):

S(env) = sum_b [28,201 * s_b + 681.08 * s_b^(3/2)]

N_seed = clamp(
  round(233 * (S(env) / 1.254e14)^1.10 * (400 / G)^1.20),
  1,
  max(1, number_of_blocks)
)

G is the requested per-shard RAM budget in GiB. The 400 is only the
calibration reference, so its factor is exactly one at 400 GiB. The per-block
s_b^(3/2) term captures block shape, the 1.10 exponent scales across
environment sizes, and the 1.20 budget exponent is fitted from the 100-GiB
runs below. This remains a seed: measured execution splits any over-budget tail
before proving.

400 GiB seed counts

At the reference budget, the budget exponent has no effect; this table shows
the size-axis correction relative to the previous raw-file-size model.

Environment Previous seed Score-model seed
Init 14 8
InitStd 24 16
Batteries 24 14
Lean 36 22
ISLB 38 23
Mathlib 233 233
FLT 241 246

The eight-shard Init partition measured 266–323 GiB with no split waves, and
Batteries needed one heavy-tail correction (14 → 15). Mathlib's unchanged
233-shard seed produced 242 measured leaves: eight initial shards were over
budget, all final shards landed at 161–397 GiB, and the run completed in 17:20
with zero failing constants. An older 241-shard FLT audit found six shards at
437–648 GiB; each fit after one model-selected halving, with the final maximum
at 99.7% of budget. Mathlib is the normalization point; FLT's new 246 count is
an extrapolation.

100 GiB calibration

Environment Model seed Tested start Initially over budget Split waves Final shards Check time End-to-end
Init 43 42* 0 0 42 44.9 s 49 s
ISLB 122 122 4 3 128 100.9 s ~155 s
Batteries 74 74 1 3 79 75.2 s ~130 s

* Init used the adjacent 42-shard calibration probe; it was clean, and the
final implementation was verified to emit 43. ISLB and Batteries' remaining
splits were localized heavy tails rather than broad under-seeding. For
comparison, the linear-scaling probes (Init at 32 and ISLB at 93) caused 16 and
28 initial misses respectively; a tested 1.36 exponent instead over-seeded
ISLB at 152 shards. The model selects 1,230 shards for Mathlib at 100 GiB; that
extrapolation was not run.

Validation

Focused Rust model tests, the full ix build, and isolated 100-GiB Init, ISLB,
and Batteries calibration runs pass.

Sharding previously guessed a shard count and found out whether it fit
by running out of memory part-way through an FFT. This closes the loop:
every shard's real prover peak is measured on the record its execution
already produced, and over-budget shards are cut in place rather than
failing the run.

The peak model (`AiurSystem::peak_prove_bytes`) evaluates the prover's
three phases over `next_power_of_two` circuit heights, so it responds to
padded rows rather than to work. `suggested_split_parts` inverts it: the
smallest power-of-two part count whose projected per-part peak fits the
budget, over the shard-independent floor the byte gadgets and
preprocessed traces contribute. The estimate is deliberately optimistic
— a part re-executes dependencies shared across the cut — because an
under-split costs one cheap re-execution while an over-split pays the
per-proof floor on every extra part forever.

Prove and check consume it from opposite ends:

- `ix prove --max-ram G` executes a shard once, measures the peak, and
  enters the witness phase only if it fits; otherwise it cuts into the
  suggested part count and re-gates each part. `prove_from_execution`
  is the seam — split out of `prove` and `prove_ixvm`, which duplicated
  everything after execution — and it takes the record BY VALUE so it
  dies before the LDE/commit phases that set the peak.
- `ix check --ram-budget G` runs the same gate on executions alone, wave
  by wave, and reports the leaf partition. `--out-ixes` writes it as a
  corrected manifest: the one `ix verify --ixes` binds the run's proofs
  against, and the one the next run should start from.

Cuts are contiguous equal-block-count runs. Splitting on measured
per-constant cost was implemented and benchmarked against this on a
synthetic fixture, init, and lean, and lost on every axis (parts, depth,
executions): virtual gas tracks time, while the RAM model tracks padded
rows. The rationale stays as a comment on the cutter so it does not get
rebuilt on intuition.

Supporting changes: static shard seeds scale by the `.ixe` block-shape
score instead of a flat byte ratio; `Ix.Aiur.Statistics`'s competing
prover-RAM projection is dropped so one model remains; and the
`aiur-sharded-env` bench row now describes the leaf partition the run
validated, with the measured window still the wave-0 batch call.
@samuelburnham

Copy link
Copy Markdown
Member Author

Superseded by #598

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant