diff --git a/sparsify/fused_encoder.py b/sparsify/fused_encoder.py index 272d47db..854e28ce 100644 --- a/sparsify/fused_encoder.py +++ b/sparsify/fused_encoder.py @@ -3,6 +3,13 @@ import torch import torch.nn.functional as F +# Upper bound on the temporary built inside the weight-gradient accumulation. The +# unchunked form allocates N * k * D elements at once, which for a wide model and +# a large k is the single biggest tensor in the backward pass -- at N=8192, k=32, +# D=1024 it is 1 GiB. Consuming it in row-blocks bounds that without changing the +# result. +BACKWARD_CHUNK_BYTES = 64 * 1024 * 1024 + class EncoderOutput(NamedTuple): top_acts: torch.Tensor @@ -66,15 +73,27 @@ def backward(ctx, grad_values, grad_indices, grad_preacts): # --- Grad w.r.t. weight --- if ctx.needs_input_grad[1]: grad_weight = torch.zeros_like(weight) - # Compute contributions from each top-k element: - # computed as grad_values * input for each top-k location. - contributions = grad_values.unsqueeze(2) * input.unsqueeze(1) - _, _, D = contributions.shape - # Flatten contributions to shape (N*k, D) - contributions = contributions.reshape(-1, D) - - # Accumulate contributions into the correct rows of grad_weight. - grad_weight.index_add_(0, indices.flatten(), contributions.type_as(weight)) + # Each top-k location contributes `grad_values * input` to its row of + # grad_weight. Materialising all of them at once costs N * k * D + # elements, so walk the batch in blocks sized to BACKWARD_CHUNK_BYTES + # and fold each block in as it is built. The arithmetic is unchanged; + # only the accumulation order inside index_add_ differs, and that was + # already unspecified on CUDA. + N, k = grad_values.shape + D = input.shape[-1] + itemsize = torch.promote_types(grad_values.dtype, input.dtype).itemsize + rows = max(1, BACKWARD_CHUNK_BYTES // max(1, k * D * itemsize)) + + for start in range(0, N, rows): + stop = min(start + rows, N) + block = grad_values[start:stop].unsqueeze(2) * input[ + start:stop + ].unsqueeze(1) + grad_weight.index_add_( + 0, + indices[start:stop].flatten(), + block.reshape(-1, D).type_as(weight), + ) # --- Grad w.r.t. bias --- if bias is not None and ctx.needs_input_grad[2]: