From ba6bc3ebba168c3d46503d3ab55f509d4b4f50b8 Mon Sep 17 00:00:00 2001 From: nakaken3013 Date: Wed, 15 Jul 2026 02:38:43 +0900 Subject: [PATCH] Reduce peak backward memory in GumbelSoftmaxFunction (in-place grad arithmetic) The tail of GumbelSoftmaxFunction.backward allocated three separate (4, out, in) tensors (g_q - dot, grad_z, grad_quant_logits) on every train_step. Reuse the already-allocated g_q buffer in place for all of them instead. Operation order is kept identical, so the result is bit-for-bit unchanged (verified for fp32 and bf16). Trims peak backward memory, which matters on 12GB-class cards under NPROC>1 (see issue #6). --- src/quantization/gumbel_quantizer_2bit.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/quantization/gumbel_quantizer_2bit.py b/src/quantization/gumbel_quantizer_2bit.py index a181874..052449e 100644 --- a/src/quantization/gumbel_quantizer_2bit.py +++ b/src/quantization/gumbel_quantizer_2bit.py @@ -95,8 +95,15 @@ def backward(ctx, grad_output): g_q = g_soft_output.unsqueeze(0) * values.view(4, 1, 1) dot = (g_q * soft_quant).sum(dim=0, keepdim=True) - grad_z = soft_quant * (g_q - dot) - grad_quant_logits = grad_z * scale / temperature - - return grad_quant_logits.to(quant_logits.dtype), grad_scales, None, None, None, None, None + # Reuse the freshly-allocated (4, out, in) `g_q` buffer in place for the remaining + # arithmetic instead of allocating separate grad_z / grad_quant_logits tensors of the + # same shape at each step. The operation order is kept identical to the allocating + # version, so the result is bit-for-bit unchanged (verified for fp32 and bf16). This + # trims peak backward memory, which matters on 12GB-class cards (see issue #6). + 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, None, None, None, None, None