Reduce peak backward memory in GumbelSoftmaxFunction (in-place grad arithmetic)

[nakaken3013-code]

Follow-up to #6 — the in-place backward micro-optimization I offered there.

What

In GumbelSoftmaxFunction.backward (src/quantization/gumbel_quantizer_2bit.py), the tail of the gradient computation allocated three separate (4, out, in) tensors — g_q - dot, grad_z, and grad_quant_logits — on every train_step. This reuses the already-allocated g_q buffer in place for all of them.

# before
grad_z = soft_quant * (g_q - dot)
grad_quant_logits = grad_z * scale / temperature
return grad_quant_logits.to(quant_logits.dtype), grad_scales, ...

# after
g_q.sub_(dot)            # g_q - dot
g_q.mul_(soft_quant)     # soft_quant * (g_q - dot)  == grad_z
g_q.mul_(scale)
g_q.div_(temperature)    # grad_z * scale / temperature == grad_quant_logits
return g_q.to(quant_logits.dtype), grad_scales, ...

Why

On 2×12GB consumer GPUs (the setup in #6), NPROC=2 OOMs on the first mlp.down_proj backward by a small margin. out × in here is the full weight shape (e.g. down_proj, intermediate_size=17408), so three transient (4, out, in) buffers per step are non-trivial; dropping them lowers peak backward memory.

Correctness

The operation order is kept identical to the original, so the result is bit-for-bit unchanged. I verified on random inputs for both fp32 and bfloat16 — torch.equal(original, in_place) → True, max abs diff 0.0 — including the * scale / / temperature split kept as two separate ops (folding them into one *(scale/temperature) does perturb bf16 rounding, so I did not do that). The change touches only the transient gradient arithmetic; it does not alter the forward pass, the RNG replay in backward, or grad_scales.

g_q is a freshly allocated, contiguous, non-leaf tensor local to backward (g_soft_output.unsqueeze(0) * values.view(4, 1, 1)), so mutating it in place is safe — no saved tensor or autograd-tracked buffer is touched.

Deliberately left out

I did not include the other change from #6 (generating torch.randn in logits_dtype instead of Q.dtype in GumbelQuantizer2Bit.__init__). As you pointed out, that shifts the init RNG draw and changes results, so it's better measured on your side than bundled into a behavior-preserving PR.