Summary

When the weight does not require grad — a frozen base in LoRA or frozen-layer finetuning — autograd discards grad_weight, so computing it in mx_mm.backward (the wgrad GEMM plus its two dim1 MXFP8 casts) is pure overhead.

This guards the weight-gradient computation on ctx.needs_input_grad[1] and returns grad_weight=None for a frozen weight. It's the same idiom already used by the int8_mixed_precision and bitnet training linears.

# input_t @ grad_output = grad_weight
if not ctx.needs_input_grad[1]:
    grad_weight = None
elif wgrad_with_hp:
    ...

Eager vs compile

This is an eager-mode win. Under torch.compile, AOTAutograd already specializes the backward on requires_grad and prunes the unused grad_weight subgraph via DCE, so the compiled frozen path is already optimal; this change just brings eager to parity. It does not change compiled behavior, and ctx.needs_input_grad is a Python bool (a trace-time constant), so the branch does not introduce a graph break.

Measured on an RTX 5060 Ti (sm120), frozen vs trainable fwd+bwd for a frozen linear, ratio = trainable / frozen (>1 means the wgrad work is skipped). Absolute numbers vary run-to-run with GPU clocks; the within-run ratio is the signal:

So on main eager does not skip the wgrad (≈1.0x) while compile already does (≈1.2–1.5x); this PR closes the eager gap.

Test

test_mxfp8_linear_frozen_weight_skips_wgrad (EMULATED + AUTO): a frozen weight gets grad_weight=None, a trainable one does not, grad_input is bit-identical between the two, and the frozen path performs strictly fewer dim1 casts (verified by counting _to_mxfp8_dim1_kernel_wrapper calls). This fails on main (the frozen path still runs the wgrad casts) and passes here.

test_mxfp8_linear_frozen_weight_compile (EMULATED + AUTO): compiles the frozen module with fullgraph=True and checks the result matches eager and weight.grad is None, locking in compile-safety.

Verified locally on sm120: the full test_mxfp8_linear.py suite passes; ruff check and ruff format --check clean.

Context

Frozen-base / parallel-adapter (LoRA) training on a quantized base is the motivating use case; see #4376 for the analogous float8 discussion and #4022 for the MXFP8 training roadmap. Touches the same mx_mm.backward as #4470 (a non-contiguous grad_output fix) but is independent of it.

🤖 Generated with Claude Code