Repository navigation
fix(lowering): keep dynamic dims through complex lowering - #4896
Conversation
complex_decomposition_adapter hands the graph to upstream's decompose_complex_in_graph, which retraces it with make_fx in real mode. There an op that builds a tensor of symbolic size, such as arange(seq) for rotary positions, needs a concrete value, so the retrace guards the symbol to its traced value in the exported program's ShapeEnv. Every engine built afterwards is static: Llama 4 and DeepSeek-V2 compiled with a dynamic batch and sequence refuse any other shape at runtime. Graphs with symbolic input dims now take the legacy complex_graph_detection pass, which keeps them symbolic. The check runs before the retrace, since the guard it adds outlives the retrace. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
github-actions
Bot
added
component: tests
shoumikhin
left a comment
shoumikhin
left a comment
There was a problem hiding this comment.
The fix works for the rotary case, and the new test fails without it and passes with it. But it makes some dynamic models that are correct today return wrong numbers, see inline. Tested on A100.
| ) | ||
|
|
||
| flat_args = _fake_flat_args(gm) | ||
| if any(free_symbols(arg) for arg in flat_args): |
There was a problem hiding this comment.
This sends every complex graph with a dynamic dim to the legacy rewriter, and the legacy rewriter gets some ops wrong. So some dynamic models that are correct today now compile without error and return wrong numbers. With a dynamic batch, main vs this PR:
torch.prod(z, dim=1): correct before, wrong afterz + (1+2j)andz * (2-1j): correct before, wrong aftertorch.complex(z.real * 2, z.imag): correct before, wrong output shape after
These ops are also wrong on the legacy path with static shapes, so the bugs are in the legacy rewriter. For example, its prod.dim_int rewrite multiplies the real parts and the imaginary parts separately. Refit with verify_output=True also fails now for a model with a dynamic complex input.
Could the fallback be limited to graphs the retrace really specializes? Or could those legacy ops be fixed first? A dynamic-shape test with complex prod would catch this kind of regression.
There was a problem hiding this comment.
Agreed—the blanket dynamic-shape fallback regressed operations that upstream already handles correctly. Fixed in bc118cf. Dynamic graphs now use upstream's ComplexTensor dispatch with make_fx(tracing_mode="symbolic"), preserving arange(seq) and other symbolic factories without routing supported ops through the legacy rewriter. The small wrapper retains upstream's mutation check; static graphs keep the upstream entry point.
Added regression coverage for all four examples (prod, complex scalar add/mul, and torch.complex), with and without a symbolic arange in the same graph, at two batch sizes. The tests also check that complex inputs and output metadata remain intact for refit's reference module. All eight cases fail against the original PR fallback and pass with this change; all 20 adapter tests pass locally on PyTorch 2.15 using the workspace adapter and decomposition table.
Also extended the complex-buffer refit test to dynamic batch sizes, asserted that verify_output=True reports success, and checked other batch sizes after refitting. That GPU test is added for CI; CUDA is unavailable locally, so I have not claimed an end-to-end refit pass. All commit hooks passed.
|
LGTM now with the new commit. As @shoumikhin had pointed the change earlier regressed complex prod, scalar add/mul and torch.complex. But then yeah we should have a followup PR to correct those. Also could also file an upstream pytorch change maybe to support this symbolic tracing mode. |
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Lan Luo <lanl@nvidia.com>
Summary
PyTorch's
decompose_complex_in_graphretraces withmake_fxin real mode. Tensor factories such asarange(seq)can then specialize the exported program's symbolic dimensions, leaving complex rotary models unable to accept other input shapes.Dynamic graphs now use the upstream complex dispatch and mutation check through a small wrapper that selects
tracing_mode="symbolic". Static graphs continue to use the upstream entry point. Unsupported-op and older-PyTorch fallbacks remain as before.This replaces the original blanket fallback to the legacy rewriter, which regressed complex
prod, complex scalar arithmetic, complex construction, and refit's reference input layout.Validation
arange, preserve complex input/output metadata, and check two batch sizes. They fail against the original blanket fallback and pass with symbolic tracing.arange(seq)test passes at a different sequence length.verify_output=True, and checked alternate batch sizes after refitting. GPU execution remains for CI because CUDA is unavailable locally.