[Fix] Key compiled artifacts by the sizes they were specialised on - #50
[Fix] Key compiled artifacts by the sizes they were specialised on#50cennn wants to merge 1 commit into
Conversation
An artifact was keyed by runtime_shape alone, which is None for every dynamically compiled graph. The dimensions nobody marked dynamic are not dynamic at all -- _mark_static_shapes burns them into the graph as constants -- so one artifact was handed to inputs it was never traced for, and the mismatch surfaced far from its cause, as a stride assertion inside the generated kernel:
AssertionError: expected size 75==7, stride 8160==8160 at dim=2
Here a graph traced for a 1s video clip (7 latent frames) was loaded for a 12s one (75), because both resolved to artifact_shape_None_subgraph_0 under the same model tag.
The entry and the artifact path now carry a digest of the sizes that were specialised. Traces that differ only in those sizes no longer collide: the second one misses, compiles, and stores under its own key. Sizes that were marked dynamic arrive as SymInt and are left out of the digest, so genuinely shape-agnostic artifacts keep being reused as before.
|
Closing: this targets the wrong layer. compile_sizes defaults to an empty list and nothing sets it (athena has no occurrence; the only reference in this repo is a test passing []), so concrete_size_entries is always empty and PiecewiseBackend.call always falls through to compiled_graph_for_general_shape. There is only ever one graph in memory, so a cache key that distinguishes shapes only helps on the narrow reload-after-restart path. Measured: with this change already applied, a VAE window traced at 8 latent frames was still handed a 9-frame input in the same process and died on the same kind of assertion (expected size 9==8 at dim=2). The key did not prevent it. The real constraint is that a magi_compile'd callable admits exactly one combination of specialised sizes, because _mark_static_shapes burns every undeclared dimension into the graph while the backend treats the first trace as general. Missing dynamic_arg_dims should fail loudly and name the dimension, not be papered over by a key. The crash that motivated this was fixed on the athena side by declaring the time dimension dynamic (world-sim-dev/athena PR #955). |
Stacked on #49.
Problem
An artifact is keyed by
runtime_shapealone, which isNonefor every dynamically compiled graph. But the dimensions nobody marked dynamic are not dynamic:_mark_static_shapesburns them into the graph as constants. So a graph traced for one input size is silently handed to another, and the mismatch surfaces far from its cause, as a stride assertion inside the generated kernel:A VAE decode graph traced for a 1s video clip (7 latent frames) was loaded for a 12s one (75 frames): both resolve to
artifact_shape_None_subgraph_0under the same model tag. The graph signature confirms only two dims generalise —submod_0(..., l_z_, s37, s38, ...)— the declareddynamic_arg_dims={"z": [3, 4]}, while dim 2 is a constant.Nothing reports a miss; the load succeeds and the failure looks like a bad meta kernel.
Fix
CacheEntryand the artifact path now carry a digest of the sizes that were specialised. Traces differing only in those sizes stop colliding: the second one misses, compiles, and stores under its own key. Dimensions marked dynamic arrive asSymIntand are excluded from the digest, so genuinely shape-agnostic artifacts are reused exactly as before, and no extra compilation is introduced — dynamo already retraces on those shape changes, this only stops the disk cache from serving the wrong result across processes.Verification
End to end on athena's gaga4 pipeline: the 12s clip that reproduced the assertion above now decodes through, 8 requests at 73s each.
Same class of defect as #42 (parallel topology into the cache directory) and #49 (compile state per topology): a key that omits something the artifact actually depends on.