Repository navigation
Pc perf followups - #58
Merged
Merged
Conversation
`eval_top_down_probs` cached its captured graphs on `id()` of the buffers they bake in, the pattern `tensorcircuit.py` had already moved away from. A reallocated `param_flows` (or denominator flows) coming back under a recycled id replayed the graph into the old, freed memory and never wrote the new buffer; and `node_mars` reached the key as a fresh `.view(-1)` on every call, so its id identified nothing. The key is now `cuda_graph_key` (address, shape, strides, dtype), moved to `pyjuice.utils.util` and shared with `TensorCircuit`'s forward/backward graph caches; it also covers `element_flows`, which the capture bakes in but the old key left out. The regression test swaps the memory under the same `param_flows` object (`Tensor.set_`): with the old key the new buffer stays zero.
…er other options Two defects in `TensorCircuit`'s captured inner-layer passes. The step that recorded a backward graph came out wrong. Its three warm-up runs execute the pass for real and accumulate into the parameter flows, and each re-initialized `node_flows` with 0.0 -- an all-ones flow in log space -- so that step's parameter flows were ~4x, and hundreds of times too large in log space. Recording now snapshots everything the pass writes to the host, warms up, captures and restores, as `eval_top_down_probs` already did. A recorded graph is replayed whenever its signature matches, and `apply_cudagraph` is on by default, but the signatures left out options that change the captured work: an MPE forward after a recorded LL forward replayed LL, and a backward with `compute_param_flows = False` (what the queries use) after a recorded training backward replayed a graph that writes parameter flows. The signatures now carry the propagation algorithm and its `alpha`, the bf16/fp32 forcing, the flow temperature (its old term, `(abs(t) - 1.0) < 1e-6`, was true for every |t| < 1), `negate_pflows`, `compute_param_flows` and the denominator flows. A backward given layer callbacks runs eagerly, since a callback runs inside the captured region.
…A graphs are reused `_init_buffer` reallocated `node_mars` / `element_mars` / the flow buffers whenever the batch size changed, and a captured graph is keyed by the addresses it baked in. In a loop whose batch changes every step -- CoDD's PC batch does, on nearly every micro-batch -- the buffers landed at new addresses, so with graphs on every step captured a new graph and none was ever freed: 182 graphs after 80 steps and 21.7 ms per step, against 7.0 ms eager. Each buffer is now re-cut from one backing storage per name, grown only when too small, and comes back zeroed on a shape change as before. Growing a storage moves its views, so the recorded graphs are dropped then (each holds a memory pool). With graphs on, the same loop records one forward and one backward graph per batch size and takes 5.9 ms per step (median of 120; batch 10-26, 17 sizes). A forward at a smaller batch no longer moves the buffers, so the sampler test that needs them to move now uses a larger one.
With the inner layers graphed, rebuilding the dense top-k index is the whole remaining overhead of a CoDD PC step (0.6-1.2 ms of it): the host waits on the build's syncs while ~105 small launches run. Its per-reference half -- the position in the category's list, the piece and table row it falls in, the ratio slot, the gradient offset, `exp` of the evidence and three scatters -- is now one Triton kernel (`_dense_index_scatter_kernel`) instead of about two dozen torch ops over every reference. The tables are bit-identical to before, checked entry by entry on tied and untied (one block per variable) models, single and length-sharded layouts, batch 1-32, lists cut into pieces of 32 and of 3, and float32 and bfloat16 evidence; the kernel uses libdevice's `exp` for that, `tl.exp` being the fast approximation. Build time on the CoDD shape: 0.52/0.71/0.83 -> 0.48/0.64/0.73 ms at batch 1/12/32, 85-105 -> 67-83 launches.
The element-flow backward's heuristic ties its edge tile to the batch tile: TILE_SIZE_K 128 up to a
16-wide batch tile, 64 above. Past batch 16 the 128 family is often much faster, but not everywhere,
so it is now a tuned candidate instead of a new default -- `_ele_tile_family` evaluated for a 16-wide
batch tile -- adopted per shape only when it measures 10% faster, and cached process-wide like the
other candidates. Shapes where it does not win keep exactly today's config.
Element-flow kernel ms at batch 17/32/64/256/512, previous -> now (tuned on both):
CoDD HMM (1024 latents) 0.79/0.78/0.78/0.79(b128)/1.04(b256) -> 0.47/0.46/0.48/0.60/1.04
HMM, 256 latents, untied 0.24/0.31/0.31/0.33/0.38 -> 0.17/0.16/0.17/0.19/0.38
HCLT-256 0.066/0.074/0.081/0.115/0.180 -> 0.050/0.050/0.061/0.110/0.179
PD 28x28, HCLT-64 unchanged (the heuristic is kept; forcing the 128 family there costs
up to 10%, which is why it is not the default)
The candidate regroups the parent reduction, so where it is taken the flows move at rounding level.
The new test forces every config the tuner may pick, on HMM, HCLT and PD at batch 17/33/100, and
requires each to be no less accurate than the heuristic's against the exact fp32 sparse kernels (all
of them measure the same: 2.4e-5 / 3.5e-4 relative at batch 17 / 100). The per-tile parent tables are
now built per edge tile (`_build_ele_parent_tables`); the CUDA forks still use the heuristic's.
The first step at every new batch size cost the CoDD PC 1.1-2.7 s with tuning on (0.1 s with it off), nearly all of it the softevi dense-vs-scattered choice running its scattered candidate -- 10-30x slower than the dense one -- ten times. CoDD's PC batch changes every micro-batch, so that was ~27 s over the 14 sizes of a short run, per process. `best_of` (under every tuned launch) now measures in two stages: each candidate runs once untimed and once timed, any that comes out `PRUNE` (2x) slower than the fastest is dropped -- it cannot clear the 10% `MARGIN` -- and the rest get the remaining warm-up and the usual timed runs. A surviving candidate runs exactly as many times as before, so close races are measured as before; across repeated runs the cached choices differ no more than the old code differs from itself (2-7 of 62, near-ties between tile configs that agree to rounding). First step per new batch size: 1.06-2.46 s -> 0.44-0.81 s; the 6-size sequence measured, 10.7 -> 3.7 s.
The dense top-k path serves any circuit whose soft-evidence leaves form one run of nodes per variable, and every built-in structure (HMM tied or untied, HCLT, PD, RAT-SPN) does; but it had only been tested on HMMs. The new test forces it, in both index layouts, on the other three and compares a whole step (LL, parameter flows, evidence gradient) with the scattered path. It fails on all three if the variables' parameter blocks are mixed up (every variable mapped to the first block).
The low-rank forward leaves `log Z` and its batch size on the layer (`_lr_bw_state`) for the backward, but only `_build_cuda_plan` set it, and the plans are cached per batch size: a cache HIT left whatever the previous forward had stored. Forwards at batch 32, 16, 32 and then a backward used batch 16's state -- parameter flows off by up to 737 (against values of at most 1.1) and wrong node flows, silently. The plan now carries its state and a hit restores it; a gated forward that builds no CUDA plan (the torch or Triton fallback) clears it, so the backward raises its "the CUDA forward did not run" error instead of reading another call's state; and the backward checks the state's batch size. The new test runs forwards at 32, 16, 32 and requires the backward's parameter and node flows to equal those of a backward that directly follows its own forward (it fails without the fix). A sequence without batch changes computes exactly what it did.
…he device `_gate_nstride` (gated layers with several gates per node block) cached its table under a key that read `nids[0]` and `nids[-1]`, to tell a recycled address apart from the rows it cached -- two device-to-host synchronizations on every call, from the backward's per-call launch setup as well as the plan builds, and the reason those passes could never be captured in a CUDA graph. It is now keyed by the rows' address, shape, strides and dtype, and the entry keeps `nids` itself alive, so that memory cannot be reused by another tensor while the key exists: a hit is the same compiled rows, which are never written. The values are unchanged (a node-axis gated forward and backward measured bit-identical to before, with autotuning off, across repeated runs).
…ges or the circuit moves A recorded graph is replayed whenever its key matches, and nothing in the key says which nodes the inner layers compute. A graph recorded under partial evaluation was therefore replayed by the next plain forward after `disable_partial_evaluation()`, computing only part of the circuit: log-likelihoods off by up to 38 nats (with `apply_cudagraph` on by default, the caller did nothing wrong). The other way round, a graph of the whole circuit was replayed under partial evaluation. `enable_partial_evaluation` / `disable_partial_evaluation` and `.to()` (which moves the layers' tables that a graph baked in) now drop the recorded graphs. The new test covers both directions and fails without the change.
…rded with; opt-in compile flag Recorded CUDA graphs of the inner layers were keyed on a hand-written list of options, so anything left off it -- a new option, a kwarg such as `mode`, state the layers keep between passes -- could replay the wrong computation silently. The mechanism (`TensorCircuit._run_inner_pass`) now holds by construction: - Key: the buffers by address, and the very dict the sum layers are called with. The input layers' per-call evidence (`Distribution.call_kwargs`, `InputLayer.call_kwargs`, declared per distribution) is left out of the key so fresh evidence does not stop a replay; the inner layers still receive it as before, and a test checks that no inner-layer source reads one of those names. - State handed from forward to backward through Python attributes (BlockScale `_bs_bw_state`, LowRank `_lr_bw_state`) is declared (`Layer.cuda_graph_state`), saved after a capture and restored after every replay, and a backward graph is keyed on the forward state it was recorded with. - A graph keeps alive everything the inner layers owned when it was captured, so a layer evicting a cached plan cannot leave it writing into freed memory. - Recording runs the call's own pass once for real, then captures (executes nothing): no warm-up runs to undo and no host round trip of the flows. 0.6-1.4 s -> 43-53 ms per recording on an HCLT with 256 latents, 17-30 -> 10-16 ms on the CoDD PC. It no longer goes through `torch.cuda.graph`, which empties the allocator's cache on every capture. Graphs hold no device memory (measured 0 MB). - At most `max_cuda_graphs` (128) are kept, least recently replayed dropped first; a buffer layout recorded more than `max_cuda_graphs_per_layout` (8) times stops recording and names the option that keeps changing. Opt-in by default-off flag: `juice.compile(..., cuda_graphs = True)` (OFF by default) records a pass the second time it is seen with the same buffers and options, so a one-off shape never pays a capture; an explicit `record_cudagraph` still wins (True records on first sight). Wherever a graph cannot be used the pass runs eagerly instead of failing: a capture that fails (a host sync on the path) warns once, keeps the call's result and is not retried; inside a caller's own capture and on the CPU no graph is used; and an external parameterization must declare `cuda_graph_safe = True` (BlockScale and LowRank do). Captures use "thread_local" mode: in "global" mode another thread's CUDA call -- a DataLoader pinning host memory -- failed both that thread and the capture. Eager behaviour is unchanged: with graphs off, 20 workloads (HMM at several batch sizes, HCLT, PD, RAT-SPN, MPE / GeneralLL / tempering / partial evaluation / sparse mode, soft evidence, BlockScale, LowRank, sampling) give the same outputs as the code before this series -- bit-identical where the kernels are deterministic, within the old code's own run-to-run spread where they are not (atomics, autotuning on) -- apart from the LowRank fix two commits back. With every circuit compiled with `cuda_graphs = True` the suite passes except three tests that check the default itself, a one-time warning an earlier call has already given, or allocation counts. Measured on the CoDD PC (real LLaDA-8B evidence, GSM8K, gbs 4), PC forward + backward per micro-batch: 6.60 ms (main) -> 4.62 ms with graphs.
…ct layers In Triton 3.7 `triton.cdiv` and `triton.next_power_of_2` are constexpr functions, and a call from host code goes through a generic wrapper that unwraps every argument: ~0.8 us a call against ~0.1 us for the arithmetic. `SumLayer` and `ProdLayer` call them while setting up every launch -- 788 calls per eager forward + backward of the CoDD PC, 3265 on an HCLT with 256 latents, 99% of all host-side calls -- about 13% of an eager step's host time. Their methods now use `host_cdiv` / `host_next_power_of_2` (`pyjuice.utils.util`), the same formulas, so the same values; kernels keep Triton's. Only the two layer classes' methods were rewritten, by an AST pass that skips anything decorated as a kernel: distribution static methods are device code spliced into generated kernels, so they keep Triton's. (`prod_layer` already imports a tensor-valued `next_power_of_2` from `compilation`, hence the `host_` names.) Eager forward + backward per step, before -> after (interleaved processes, every after-run faster than every before-run): HMM with 1024 latents 4.13 -> 3.68 ms (batch 16), 4.35 -> 4.05 (128); HCLT-256 18.8 -> 16.7 (16), 21.4 -> 19.4 (128); CoDD PC 6.50 -> 5.96 (12 blocks), 6.90 -> 6.39 (24).
Since 47c943d the circuit's buffers (`node_mars`, `element_mars`, the flows, ...) are cut from one backing storage per name that only ever grew, so that batch sizes that vary keep one address and their CUDA graphs. One large batch -- an evaluation pass, say -- therefore held its buffers for the rest of the process: on an HMM with 1024 latents, 3.1 GB still allocated back at batch 16 after one batch-4096 step, against 1.0 GB on main, where every shape got its own allocation. A storage more than `_BUFFER_SHRINK_RATIO` (4) times what the shape needs is now reallocated at the needed size, unless graphs of the inner layers are recorded (their replays need the addresses to stay). The same measurement now gives 1.0 GB, as on main. Batch sizes that vary within the ratio keep one storage as before -- CoDD's micro-batches span 10-26 blocks -- so the measured throughput is unchanged. The new test checks all three cases (a much smaller batch releases the storage, a batch within the ratio keeps it, recorded graphs keep it) and fails if either half of the rule is removed.
…d the kernel launcher At small batch an eager step is bound by the host: a sum layer's backward spent ~45 us of Python for two kernel launches whose C launch is ~7 us. Four changes, none of which caches a dispatch decision: - The autotune candidate lists (sum forward, element- and parameter-flow backward, product layer) are built only when the choice is not cached yet; the cache is keyed on the heuristic default as before. - `FastJITFunction3x`: one pass over the keyword arguments binds them and sets the launch meta-parameters aside (was a filtered copy plus a second pass), the signature is one comprehension, and the active driver is fetched once per launch. The bound arguments and the signature are identical to before on every one of 1114 launches checked (HMM, HCLT, PD, RAT-SPN steps). - The partial-evaluation selections are declared None at class level and the shared-memory fallback flags are set in `SumLayer.__init__`: `provided()` / `getattr(..., default)` used to find them missing through `nn.Module.__getattr__`'s AttributeError, 257 times a step at ~0.4 us each. - `FastParamList.__getitem__` reads `_parameters` directly; anything not there takes the old route, so an invalid index fails exactly as before. Host time per call, before -> after: product layer ~15 -> ~12 us, sum forward 12 -> 9.5, sum backward 52.5 -> 43.5. Eager step: HMM with 1024 latents 3.70 -> 2.98 ms (batch 16), 3.80 -> 3.32 (128); HCLT-256 16.6 -> 14.2 (16); CoDD PC 6.07 -> 5.46 (12 blocks). Outputs: 15 of 20 eager workloads bit-identical to before, the other 5 within the code's own run-to-run spread (atomics).
The test drew its DiscreteLogistic parameters from whatever RNG state earlier tests left, and a rare draw puts `log(cdf_high - cdf_low)` where float32 cancellation alone exceeds the 1e-4 tolerance: unseeded, 1-3 failures in 150 runs, on main's code and this branch's alike. Seeded as its neighbour `test_discrete_logistic_nodes_behavior` is (239829), whose worst error is 1.4e-6; 0 failures in 150 runs from scrambled prior RNG states.
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
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.
PC performance follow-ups: safe opt-in CUDA graphs, eager host-overhead cuts, and five silent-correctness fixes
Summary
This branch follows up the leftover performance items from #57. It makes three kinds of change:
main, are fixed.With graphs off (the default), outputs match
mainexcept where a bug was fixed. CoDD 2.0 training is about 3% faster per step, and the PC part of the step is about 20% faster.Bug fixes (all silent on
main)5fd39c5e). The three warm-up runs before capture accumulated into the parameter flows and resetnode_flows. The recording step came out with about 4× its parameter flows, and far more with log-space flows.5fd39c5e):compute_param_flows=False(as queries use) replayed a graph that writes parameter flows.40457fc6). Afterdisable_partial_evaluation(), a plainpc(x)replayed the graph recorded under partial evaluation. Log-likelihoods were off by up to 38 nats.461eeb17, eager mode, no graphs involved). Forwards at batch 32, then 16, then 32, followed by a backward, used batch 16's state. Parameter flows were off by up to 737 (correct values are at most 1.1), and node flows were wrong.id()(8d603206). A recycled id could replay into a freed buffer.CUDA graphs: safe by construction, opt-in
The key covers everything the recorded layers see. It is built from the very dict of options the inner layers are called with, plus each buffer's address, shape, strides and dtype. No option can reach a recorded pass without being in its key.
Per-call evidence is left out of the key. Evidence that only input layers read is now declared per distribution (
Distribution.call_kwargs). A source-scan test checks that no inner layer reads those names.State handed from forward to backward is declared. BlockScale's and LowRank's forward state (
Layer.cuda_graph_state) is saved after recording and restored after every replay. A backward graph only replays against the forward state it was recorded with.A graph keeps alive everything the layers owned at capture time, so evicting a cached plan can't leave a graph writing into freed memory.
Recording is cheap. It runs the call's own pass once, then captures; there's no snapshot and no allocator-cache flush.
Graphs hold no GPU memory.
The number of graphs is bounded. At most 128 are kept (LRU). A pass shape recorded more than 8 times stops recording and warns, naming the option that keeps changing.
Opt-in:
juice.compile(..., cuda_graphs=True), off by default. With it on, a pass is recorded the second time it's seen; an explicitrecord_cudagraph=Truerecords on the first. Wherever a graph can't be used, the pass runs eagerly:cuda_graph_safe = True(BlockScale and LowRank do).Capture uses
thread_localmode. In the default global mode, a DataLoader pinning host memory on another thread broke both itself and the capture.Buffers keep stable addresses across batch sizes, so varying batches reuse their graphs. A buffer more than 4× larger than needed is released when no graphs are recorded (
1751341f), so one large evaluation batch no longer holds memory for the rest of the process.Eager performance
triton.cdiv/next_power_of_2(7d67325c). In Triton 3.7 these cost about 0.8 µs per host call, and the layers made 788–3,265 calls per step.179dc80c):AttributeErrorround-trips per step.abce8bf6; its tables are bit-identical);304e1c6d);af237451), which cuts first-step tuning from 1.1–2.7 s to 0.44–0.81 s;beed338f).Eager step, wall time, against
main:mainResults against
mainCoDD 2.0 training, full epoch, 1,868 steps:
mainValidation log-likelihood is equal across arms. The PC is now about 2% of the step (it was 2.4%).
Standalone training, HCLT-256 on ImageNet32 with mini-batch EM:
mainLog-likelihoods match.
Verification
mainand the branch, 2–4 runs each, with autotuning on and off:main's own run-to-run noise where they aren't, apart from the LowRank fix.cuda_graphs=True, the suite passes except three tests that check the default itself, a one-time warning, or allocation counts.test_discrete_logistic_nodesfailed 1–3 times in 150 runs onmaintoo (36b330ca).Notes for reviewers
record_cudagraphnow defaults toNone, meaning "use the circuit'scuda_graphssetting". ExplicitTrue/Falsebehave exactly as before.cuda_graphs=True, LLaDA's kernels run about 1% slower in the same process, which cancels the PC's 30% gain. The cause isn't understood; graphs are opt-in, so the default isn't affected.main). The cause is a near-tie between two numerically equivalent kernel choices.