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

**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`](https://github.com/NVIDIA/Megatron-LM/commit/efb6ed573603656a08923e5c3e1ccd29ac108591)

Clamped SwiGLU was introduced through:

[#5940 Complete clamped SwiGLU across expert paths](https://github.com/NVIDIA/Megatron-LM/pull/5940)

The underlying PyTorch issue is tracked at:

[pytorch/pytorch#194976](https://github.com/pytorch/pytorch/issues/194976)

@NVIDIA/mcore-oncall

**Steps/Code to reproduce bug**

Check out the tested Megatron-LM revision:

```bash
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:

```python
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:

```bash
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:

```text
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:

```text
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:

```text
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

| Backend/configuration | Forward | AOTAutograd backward | Explicit `clamped_swiglu_back` |
|---|---:|---:|---:|
| Inductor, default `specialize_float=False` | Incorrect | Incorrect | Incorrect |
| Inductor, `specialize_float=True` | Correct | Correct | Correct |
| `aot_eager`, `specialize_float=False` | Correct | Correct | Correct |

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

```text
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:

```text
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:

```python
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:

```python
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](https://github.com/pytorch/pytorch/issues/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:

```python
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:

```text
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:

```text
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
```