[PyTorch] Enable fused activation recompute for ScaledTanhSReLU

[wanyingw]

Description

Enable fused activation recompute for ScaledTanhSReLU

• forward, dgrad, dprobs and FC1 wgrad are required to be BITWISE identical -- recompute changes only what is stashed;
• FC2 wgrad cannot be, and is bounded against the unclamped op rather than a constant, since ScaledSReLU has shipped with recompute enabled all along and defines what this feature's regeneration error already is;
• the two runs are proven to differ, by tracing use_dsrelu_reuse into the cuDNN wrapper, so the comparison cannot pass by having saved FC2's input on both sides;
• memory actually drops, measured between forward and backward. Every other check would pass for a "recompute" that regenerates and also keeps the saved tensor.

Both recipes run: the regeneration is two code paths, since NVFP4 has cuDNN return BF16 for TE to requantize. Options that change what wgrad consumes each get a case -- bias, single_grouped_weight, and delay_wgrad_compute, the last because the regenerated tensor has to survive past where the non-delayed path clears it.

Measured on GB300, deviation of the regenerated FC2 wgrad as a fraction of that tensor's RMS:

		max-abs	RMS	mean/rms
MXFP8	SReLU	14%	1.6%	+0.00
	TanhSReLU	19%	1.5%	-0.00
NVFP4	SReLU	62%	8.8%	-0.00
	TanhSReLU	33%	4.1%	-0.00

The max column is a tail statistic and the RMS column is the one that matters. The BF16 saved FC1 output perturbs FC2's input by ~0.8%, and the quantizer turns that into occasional full-ULP bin flips, so the error scales with bin width rather than with BF16 precision -- which is why NVFP4, whose ULP is 4x larger, lands 5.9x higher. The mean is zero to two decimals everywhere, so this is noise rather than an offset that would accumulate into the weights. For scale, TE's own dtype_tols accepts rtol=0.125 for e4m3 on these tensors.

The clamped op comes out ahead of the unclamped one under NVFP4 (4.1% against 8.8%): bounding the activation at s**2 compresses exactly the dynamic range 4-bit quantization handles worst.

Validated end to end as well, driving Megatron-Core -> TE -> cuDNN on GB300 under MXFP8, since MCore is where the flag originates and no TE test exercises that wiring. Four configurations -- unclamped and clamped, each with recompute off and on -- all fused, with use_dsrelu_reuse matching what MCore asked for and the clamp reaching the kernel. Against MCore's own unfused reference all four sit at 1.7-1.8% rms, and against a closed-form fp32 evaluation of the activation at 5.7%. Recompute against its own no-recompute twin was bitwise equal on the forward, both input gradients and FC1 wgrad, with FC2 wgrad moving 0.92% clamped against 0.86% unclamped -- the clamped path is indistinguishable from the one already shipping.

The weight_requires_grad arm of the gate is covered by a strict xfail. It is unreachable today because a frozen-weight module whose input still needs a gradient raises in mark_grouped_tensor, before the gate is consulted -- a pre-existing defect unrelated to recompute or to the clamp. The marker comes off when that is fixed.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[timmoon10]

This is AI slop. The entire point of comments is to make it easier to understand code, and it's unreasonable for a human to read all this.

Including raw experiment logs and unprocessed thoughts is especially bad practice since they'll become stale and irrelevant very quickly. If this is important enough for a human to read, please summarize it yourself and rewrite as a GitHub comment.

Suggested change

	# Recompute-vs-save deviation in FC2's wgrad, as a fraction of that tensor's own RMS,
	# measured on GB300 (hidden 256, 4 groups):
	#
	# max-abs RMS of error
	# MXFP8 SReLU 14% 1.6%
	# TanhSReLU 19% 1.5%
	# NVFP4_RHT SReLU 62% 8.8%
	# TanhSReLU 33% 4.1%
	#
	# The max column looks alarming and the RMS column is the one that matters. The forward
	# applies the activation to its FP32 accumulator while the backward applies it to the
	# saved BF16 FC1 output, so the two disagree by that round trip however well the two
	# copies of the activation match; under a 4-bit recipe the requantized result then lands
	# in a different bin often enough that the tail reaches 60% while the typical
	# disagreement stays below 10%. So the binding check is on RMS and max is only a ceiling.
	#
	# The clamped op comes out *ahead* of the unclamped one under NVFP4 (4.1% against 8.8%),
	# which is the expected direction rather than luck: bounding the activation at s**2
	# compresses exactly the dynamic range 4-bit quantization handles worst.
	#
	# Plain ScaledSReLU has shipped with recompute enabled all along, carrying these numbers,
	# so it is the accepted baseline rather than a target. That is why the binding check is
	# the comparison against it -- it stays meaningful across recipes, shapes and hardware,
	# unlike a constant that would have to be re-derived for each. The absolute ceilings exist
	# only to catch a regression that degrades both at once.

[timmoon10]

We add a bunch of recompute tests for SReLU-like activations, but there's nothing special about SReLU. No reason we couldn't use a similar kernel for GeLU or sigmoid or any other unary activation. In fact, even the GLU activations support the same activation_recompute_in_mlp option.

Rather than making hyperspecific tests just for TanhSReLU and activation recompute, it would be better to add general tests for activation recompute that are independent of the activation. Even if we only run it for SReLU and TanhSReLU for now, it would be more logical and easier to extend if we expand support for activation recompute in the future.

[timmoon10]

Claude's unfinished thoughts don't belong in production code.

[timmoon10]

I guess checking the error RMS is a nice approach. Grouped MLP is large and numerical errors tend to accumulate, so we need excessively loose tols for the maximum error. We should consider generalizing so that other tests can check error RMS.

[timmoon10]

Shouldn't we expect the bound for the error RMS should to be dependent on the quantization recipe?