[BUG] clamped_swiglu silently reuses a stale clamp_value under default torch.compile/Inductor configuration #6918

Describe the bug

The original Megatron-LM compiled functions

• megatron.core.fusions.fused_bias_swiglu.clamped_swiglu
• megatron.core.fusions.fused_bias_swiglu.clamped_swiglu_back

can silently reuse a stale Python-float clamp_value under the default torch.compile/Inductor configuration.

After the compiled function observes clamp_value=0.5, a later invocation with clamp_value=1.0 can execute Inductor code that still behaves as if the clamp value were 0.5.

This produces silently incorrect results in:

1. clamped_swiglu forward;
2. its AOTAutograd-generated backward;
3. the explicitly compiled Megatron clamped_swiglu_back.

No exception or warning is raised.

The reproducer:

• directly imports the original functions from Megatron-LM;
• uses the undecorated functions retained by torch.compile as eager references;
• does not copy or reimplement the SwiGLU/clamp expression;
• does not call torch._dynamo.reset();
• does not modify or patch specialize_float;
• runs under PyTorch's default specialize_float=False configuration;
• reproduces in a fresh Python process on both CPU and CUDA.

Tested Megatron-LM revision:

efb6ed573603656a08923e5c3e1ccd29ac108591

Clamped SwiGLU was introduced through:

#5940 Complete clamped SwiGLU across expert paths

The underlying PyTorch issue is tracked at:

pytorch/pytorch#194976

@NVIDIA/mcore-oncall

Steps/Code to reproduce bug

Check out the tested Megatron-LM revision:

git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
git checkout efb6ed573603656a08923e5c3e1ccd29ac108591

Save the following as repro_clamped_swiglu.py in the repository root:

import torch

from megatron.core.fusions.fused_bias_swiglu import (
    clamped_swiglu,
    clamped_swiglu_back,
)


# torch.compile retains the original undecorated Megatron functions here.
# They are used as eager references, so this reproducer does not copy or
# reimplement the SwiGLU/clamp computation.
eager_clamped_swiglu = clamped_swiglu._torchdynamo_orig_callable
eager_clamped_swiglu_back = clamped_swiglu_back._torchdynamo_orig_callable

device = "cuda" if torch.cuda.is_available() else "cpu"

print("torch:", torch.__version__)
print("torch git:", torch.version.git_version)
print("device:", device)
print("specialize_float:", torch._dynamo.config.specialize_float)

# This bug reproduces with PyTorch's default configuration.
assert torch._dynamo.config.specialize_float is False

cases = [
    ((1, 2), False, 1.0),
    ((1, 4), True, 0.5),
    ((1, 4), True, 1.0),
]

for shape, requires_grad, clamp_value in cases:
    x = torch.full(
        shape,
        0.75,
        device=device,
        dtype=torch.float32,
        requires_grad=requires_grad,
    )
    x_ref = x.detach().clone().requires_grad_(requires_grad)

    # Original Megatron torch.compile entry point.
    actual = clamped_swiglu(x, clamp_value)

    # Original undecorated Megatron function.
    expected = eager_clamped_swiglu(x_ref, clamp_value)

    forward_diff = (actual - expected).abs().max().item()
    autograd_diff = None

    if requires_grad:
        actual.sum().backward()
        expected.sum().backward()
        autograd_diff = (x.grad - x_ref.grad).abs().max().item()

    # Independently exercise Megatron's explicitly implemented and
    # torch.compile-decorated backward entry point.
    grad_output = torch.ones_like(expected)

    actual_explicit_back = clamped_swiglu_back(
        grad_output,
        x.detach(),
        clamp_value,
    )
    expected_explicit_back = eager_clamped_swiglu_back(
        grad_output,
        x_ref.detach(),
        clamp_value,
    )

    explicit_back_diff = (
        actual_explicit_back - expected_explicit_back
    ).abs().max().item()

    print(
        f"shape={shape}, "
        f"requires_grad={requires_grad}, "
        f"clamp_value={clamp_value}, "
        f"forward_diff={forward_diff}, "
        f"autograd_diff={autograd_diff}, "
        f"explicit_back_diff={explicit_back_diff}"
    )

Run the reproducer in a fresh Python process:

PYTHONPATH=. python repro_clamped_swiglu.py

The invocation order is relevant:

1. shape (1, 2), no autograd, clamp_value=1.0;
2. shape (1, 4), autograd enabled, clamp_value=0.5;
3. shape (1, 4), autograd enabled, clamp_value=1.0.

No Dynamo cache reset or configuration patch is needed.

CUDA result

The exact script above was run in a fresh Python process using CUDA tensors:

torch: 2.8.0a0+34c6371d24.nv25.08
torch git: Unknown
device: cuda
specialize_float: False

shape=(1, 2), requires_grad=False, clamp_value=1.0,
forward_diff=0.0,
autograd_diff=None,
explicit_back_diff=0.0

shape=(1, 4), requires_grad=True, clamp_value=0.5,
forward_diff=0.0,
autograd_diff=0.0,
explicit_back_diff=0.0

shape=(1, 4), requires_grad=True, clamp_value=1.0,
forward_diff=0.22642318904399872,
autograd_diff=0.26196935772895813,
explicit_back_diff=0.26196935772895813

Latest PyTorch nightly CPU result

The same exact script was also run in another fresh Python process using CPU Inductor and the latest tested official PyTorch nightly:

torch: 2.15.0.dev20260826+cpu
torch git: 68a5d90dd35d45ca457fda21dbae885987fb1e05
device: cpu
specialize_float: False

shape=(1, 2), requires_grad=False, clamp_value=1.0,
forward_diff=0.0,
autograd_diff=None,
explicit_back_diff=0.0

shape=(1, 4), requires_grad=True, clamp_value=0.5,
forward_diff=0.0,
autograd_diff=0.0,
explicit_back_diff=0.0

shape=(1, 4), requires_grad=True, clamp_value=1.0,
forward_diff=0.22642318904399872,
autograd_diff=0.26196935772895813,
explicit_back_diff=0.26196935772895813

The CPU and CUDA differences are identical. This indicates that the root cause is not specific to CUDA, Triton, Transformer Engine, or a GPU kernel.

Expected behavior

Every invocation of clamped_swiglu and clamped_swiglu_back must use the supplied runtime clamp_value.

For the reproducer above, all of the following should be zero:

forward_diff
autograd_diff
explicit_back_diff

When specialize_float=False, PyTorch should either:

1. keep clamp_value genuinely dynamic throughout the compiled graph; or
2. specialize it with an appropriate runtime guard and recompile when the value changes.

It must not embed and reuse an unguarded trace-time value.

Different tensor shapes, autograd states, or invocation order must not cause clamped_swiglu to reuse another invocation's clamp value.

Additional context

Ablations

On the same CUDA environment, enabling Python-float specialization produces:

device: cuda
specialize_float: True

shape=(1, 4), requires_grad=True, clamp_value=1.0,
forward_diff=0.0,
autograd_diff=0.0,
explicit_back_diff=0.0

Using aot_eager with specialize_float=False also produces zero differences:

device: cuda
backend: aot_eager
specialize_float: False

shape=(1, 4), requires_grad=True, clamp_value=1.0,
forward_diff=0.0,
autograd_diff=0.0,
explicit_back_diff=0.0

Verified temporary workaround

A concrete workaround is to enable Dynamo's Python-float specialization during process initialization:

import torch

# Work around pytorch/pytorch#194976.
# Guard Python float values and compile a separate graph when clamp_value
# changes instead of attempting to keep it automatically dynamic.
torch._dynamo.config.specialize_float = True

This should be configured before the first compiled Megatron model invocation:

import torch

torch._dynamo.config.specialize_float = True

from megatron.core.fusions.fused_bias_swiglu import (
    clamped_swiglu,
    clamped_swiglu_back,
)

# Initialize and execute the model after enabling the workaround.

The setting does not strictly need to precede the import, because Dynamo normally traces on the first invocation. However, setting it during process initialization is the safest approach and ensures it remains active for every initial trace and recompile.

If the process has already invoked compiled Megatron functions, restart it and set specialize_float=True before model execution. A fresh process avoids dependence on previously compiled graphs.

Workaround trade-off

With specialize_float=True, Dynamo guards the exact Python float value and normally compiles a separate graph for every distinct clamp_value.

The expected cost is:

• an additional compilation when a previously unseen float value is encountered;
• additional compiled-graph cache entries;
• potentially higher startup or warm-up time if many distinct float values are used.

There should be no inherent per-call steady-state runtime regression after the corresponding graph has been compiled. Since activation_func_clamp_value is normally a stable model configuration value with few distinct values, the compilation and cache overhead should generally be limited.

This is a correctness workaround, not the intended long-term behavior. The desired upstream fix should keep supported clamp uses genuinely dynamic and safely specialize a Python float whenever any remaining use cannot be tensorified.

Suspected upstream cause

The suspected root cause is in PyTorch's device-independent Dynamo/AOTAutograd Python-scalar tensorification path.

clamp_value becomes an automatically dynamic backed SymFloat. Some uses of that symbol can be tensorified, while the clamp uses are not handled consistently. The symbol may then be treated as fully tensorified even though an unsupported use still embeds its trace-time value.

This can allow Inductor code generated using clamp_value=0.5 to be reused when the runtime argument is 1.0.

The upstream analysis and proposed PyTorch fix are tracked at:

pytorch/pytorch#194976

Why track this in Megatron-LM

Although the root cause is in PyTorch, the affected compiled entry points are part of Megatron-LM main and are used by the clamped SwiGLU MoE paths.

The bug can be triggered whenever the shared compiled callable observes different clamp values in the same process, for example across model/config instances or tests.

Because the failure is silent, a downstream correctness safeguard and regression test would be useful until the minimum supported PyTorch version contains the upstream fix.

Suggested Megatron-LM follow-up:

1. Add a regression test that invokes the original clamped_swiglu and clamped_swiglu_back using the sequence above.

2. Until the upstream fix is available in supported PyTorch releases, document or apply:

   torch._dynamo.config.specialize_float = True

   as a correctness workaround for compiled clamped SwiGLU.

3. Consider changing clamp_value to a 0-D tensor on the same device, after validating forward, backward, dtype and performance behavior.

4. As a conservative fallback, disable Inductor compilation for the clamped variants until the upstream fix is available.

5. Do not rely on creating separate closure or callable objects for isolation. Dynamo automatic-dynamic state can be associated with the shared Python code object.

Environments

CUDA reproduction:

Megatron-LM:
  branch: main
  commit: efb6ed573603656a08923e5c3e1ccd29ac108591

PyTorch:
  version: 2.8.0a0+34c6371d24.nv25.08
  git version: Unknown
  CUDA used to build PyTorch: 13.0

Triton:
  version: 3.3.1

GPU reported by nvidia-smi:
  NVIDIA H200

NVIDIA driver:
  570.148.08

Latest nightly CPU reproduction:

Megatron-LM:
  branch: main
  commit: efb6ed573603656a08923e5c3e1ccd29ac108591

PyTorch:
  version: 2.15.0.dev20260826+cpu
  git version: 68a5d90dd35d45ca457fda21dbae885987fb1e05

Python:
  3.12.3

OS:
  Ubuntu 24.04.1 LTS, x86_64