perf: back-pointer DP for AutoScheme bit allocation to cut path-copy RAM

[SuperMarioYL]

What

choose_bits_per_layer_with_path (the AutoScheme knapsack DP that picks the
per-layer bit allocation) kept the full chosen path inside every DP state and
rebuilt it on every transition with:

new_path = cur_path + [(layer_names, scheme)]

That cur_path + [...] copies the entire path (length up to L) once per
surviving Pareto state per option, so the search is O(L² · S) in both time
and live path memory for L layers and S surviving states. For models with
many layers and incommensurate sizes, the per-state path lists are a large part
of the AutoScheme RAM footprint.

This replaces the per-state path list with an immutable back-pointer cons
cell (layer_names, scheme, prev_tail). Appending a layer is now an O(1)
tuple allocation that structurally shares the unchanged prefix with all
sibling states, and the full path is reconstructed once at the end by walking
the chain in O(L). Pareto pruning and the max_states beam are untouched.

Why

#1913 (AutoScheme RAM usage is too large, still open) reports a multi-GB RSS
blow-up on this code path; #1916 brought peak RAM from ~90 GB to ~50 GB and the
maintainer noted that is "not enough". The path-copy is one of the remaining
sources of that footprint and was left unchanged by #1916, so this is a
complementary, low-risk reduction rather than a competing approach.

This change is intentionally scoped to the back-pointer rewrite only — it does
not change the max_states default or any accuracy-affecting behaviour.

Correctness

The returned (loss, path) is identical to the previous implementation:

• Differential test vs. the old implementation on 20,000 randomized instances
  (1–9 layers, 1–4 options each, random bit/loss costs, and every max_states
  setting incl. None/1/2/3/5) — 0 mismatches.
• New test/test_cpu/schemes/test_choose_bits_per_layer.py (the function had no
  coverage): DP optimum matches an exhaustive brute force; grouped layer-name
  tuples survive reconstruction; infeasible budgets return (None, None); the
  beam cap still yields a feasible, full-length path.

Testing

pytest test/test_cpu/schemes/test_choose_bits_per_layer.py -q
# 4 passed

black/isort (line-length 120) and ruff (E4,E7,E9,F,NPY,FURB) are clean
on both files.

[wenhuach21]

Thanks for the PR! Do you have an estimate of the RAM savings for an 8B model?

[wenhuach21]

I have tested the state number before, it's not too large as the state has been pruned

[SuperMarioYL]

Thanks @wenhuach21! You're right that once the state set is pruned the count stays modest — the saving here isn't in the number of states (that's what #1916's beam bounds), it's in the per-state path bookkeeping, which is independent of the model weights.

The old code stored a full list of (layer_names, scheme) for every surviving state and rebuilt it with cur_path + [...] on each transition — O(L·S) live path memory for L layers and S states, plus an O(L) copy per transition. The back-pointer makes each state a single cons cell that structurally shares its prefix with siblings, so per-state path cost is O(1) and the per-transition copy disappears.

To put a number on it I measured the DP's path memory in isolation (synthetic layers at 8B-scale decision counts — Llama-3-8B is ~224 per-module points — tracemalloc peak; both impls return bit-identical loss/path):

regime	~states S	OLD peak	NEW peak	ratio	saved
beam=1024 (pruned)	~1024	12.8 MB	5.3 MB	2.4×	7.5 MB
beam=4096 (pruned)	~4096	16.5 MB	10.2 MB	1.6×	6.4 MB
no beam, fine-grained costs	~4.7k	67.5 MB	24.4 MB	2.8×	43 MB
no beam, L=300	~7.3k	139.5 MB	47.0 MB	3.0×	93 MB

So:

• With a tight beam (your case): single-digit MB, ~1.6–2.4× — modest in absolute terms, exactly as you'd expect from a well-pruned state set.
• In the default/unbounded regime [Bug]: AutoScheme ram usage is too large #1913 hits (S growing into the thousands→millions), it's a ~3× reduction of the path component that scales linearly with L·S — that's where the GB-scale numbers come from, and it stacks on top of Refine AutoScheme logic for gguf, reduce memory consumption and improve speed #1916's beam rather than competing with it.

The other half of the win is removing the cur_path + [...] allocation itself (a fresh O(L) list per surviving state per layer → O(L²·S) churn), which shows up as allocator/GC pressure and peak transient RSS even when steady-state S is small.

These are synthetic DP-isolation numbers (no GPU/real model needed, since the path cost is weight-independent). Happy to run an end-to-end AutoScheme RSS measurement if you can point me at a repro config from #1913.

[wenhuach21]

@n1ck-guo @xin3he I've forgotten some of the DP details. Could you please review this when you have time?

[wenhuach21]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[SuperMarioYL]

Rebased onto the current main to clear the merge conflict. It was pure base drift in auto_round/auto_scheme/delta_loss.py: main had switched the DP path representation from a list to a tuple (cur_path + ((layer_names, scheme),)), which still copies the whole path on every transition. I re-applied the back-pointer cons-cell optimization on top of that current baseline, so the O(1)-append / structural-sharing behavior is preserved and the diff is unchanged in intent.

I re-verified the rebased implementation returns bit-for-bit identical (loss, path) to the current main version across 240k randomized instances and every max_states setting (including the beam caps), so this is a behavior-preserving change. DCO sign-off is intact. Now MERGEABLE again — happy to address anything from the DP review.

[wenhuach21]

We’ll merge this PR later since we’re under code freeze