Skip to content

Aggregated-delta fast-path for O(1) mid-run rollout joins (deferred) #62

Description

@nanjiangwill

Problem

A fresh rollout engine that joins mid-run (elastic scale-up) has to bring its host-local
checkpoint from base up to the live version. Today it folds the whole published d1..vN delta
chain from base in one pass (_apply_range in the sglang receiver) — correct, but O(N) in the
number of versions. Late in a long run (large N) that fold gets slow.

Idea: an aggregated delta

Keep aggregate/ = base ⊕ vM — the entire chain folded into a single xor delta
(base_version 0) — maintained in the background. A joiner then:

  1. seeds base,
  2. applies the one aggregate to reach vM (one apply, independent of N),
  3. folds only the short d(M+1)..target tail.

So the join cost becomes ~constant instead of O(N).

Prototype (built, then removed)

  • Producer src/stitch/aggregate.py — transport-agnostic advance(root, target): resume from
    the current aggregate or rebuild from the chain, telescoping xor fold (d1⊕…⊕dM = base⊕vM),
    in-place write. [aggregate] extra (numpy/zstandard).
  • Task cookbook/common/aggregator.py — an opt-in CPU Modal task following the pointer, paced
    by AGGREGATE_THRESHOLD (fold every K versions → a joiner folds a ≤K tail). Default off.
  • Consumer (sglang fork receiver) — a fresh host applies the aggregate then folds the tail. The
    aggregate is one in-place dir on an eventually-consistent transport, so every shard carries a
    __version__ stamp and a reader that sees a disagreement (a torn mid-update read whose per-shard
    checksums still self-verify) skips it and folds from base. The aggregate is a cache, never
    required for correctness.

Reference commits (removed from their branches; git show <sha> while still in reflog):
stitch 688ea45 (aggregate: producer + Aggregator for O(1) mid-run joins); sglang fork
modal-projects/sglang 44c81a6e (aggregate fast-forward folded into the /pull_weights
receiver). The design above is enough to rebuild if the objects are gone.

Why deferred

Measured on glm45-air-fp8 (113 GB base, real fp8 deltas): one incremental fold costs ~70s
(decompress the aggregate's touched tensors + XOR + recompress — a near-full-model recompress),
and a rebuild of d1..d8 was ~322s. Keeping the aggregate current therefore means a continuous
~70s-per-training-step background recompress
(plus a transport rewrite + a torn-read window each
update), for a whole extra producer component + producer/consumer format contract + torn-read
handling.

Against that: the fold already makes a mid-run join correct (fold the chain from base in one
pass — O(N) decompress, not O(N) mmap read-modify-writes), and joins are rare (elastic
scale-up, not steady state). Spending a permanent expensive background workload + extra code to
turn a rare, already-tolerable join into a fast one isn't worth it yet.

Revisit when

  • late-run joins (large N — hundreds of deltas) make the fold-from-base a real bottleneck, or
  • joins become frequent.

Then reintroduce, tuning the update cadence so the producer's O(full) recompress is amortized (and
reconsider whether the trainer can publish only changed tensors, which would make one fold cheap
and might make per-step aggregation viable).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions