[PyTorch 2.13 regression] Dion3 normalization fails during dynamic-shape compilation

[WindChimeRan]

Summary

While integrating Dion3 in vllm-project/speculators#1031, we found that optimizer.step() fails on PyTorch 2.13.0 when Dion3 processes a second distinct parameter shape.

The affected path is nordion2_normalize_selected_stacked. After automatic dynamic-shape generalization, Inductor generates a Triton scatter epilogue that references a temporary defined only inside the preceding tl.range reduction body:

NameError: tmp19 is not defined

A pure-PyTorch reproducer passes on PyTorch 2.12.0 and 2.12.1 but fails on 2.13.0. PyTorch 2.12.1 and 2.13.0 both use Triton 3.7.1, so the evidence points to an Inductor regression rather than a Dion3 bug.

Upstream report: pytorch/pytorch#194490

Impact

This prevents Dion3 training on PyTorch 2.13.0 when multiple parameter-shape groups reach the compiled per-neuron normalization function. We reproduced it with microsoft/dion@58d38ad, FSDP2, and 8×H100 GPUs, although the upstream reproducer does not require Dion3 or distributed training.

Workaround

Disabling automatic dynamic-shape promotion avoids the invalid code generation:

with torch._dynamo.config.patch(
    automatic_dynamic_shapes=False,
    cache_size_limit=64,
):
    optimizer.step()

This is related to #23: static shapes require one compilation per distinct parameter shape, so the workaround raises cache_size_limit to avoid the recompile limit.