Bound the weight-gradient temporary in FusedEncoder.backward

[alepot55]

What

FusedEncoder.backward builds grad_values[:, :, None] * input[:, None, :], an [N, k, D] tensor holding every top-k contribution at once, and then folds it into grad_weight with a single index_add_. This consumes it in row-blocks sized to a fixed byte budget instead.

Why

It is the largest single allocation in SAE training. Some sizes, all measured on an A100:

shape	temporary
batch 8, ctx 512, d_model 576, k 32 (repo defaults)	288 MiB
N=8192, D=1024, k=32 (what tests/test_encode.py already runs)	1 GiB

It is linear in k, so the --k 192 configuration the README recommends for Llama 3 8B pays six times the first row.

Effect

Peak allocated per rank, SmolLM2-135M, 30 hookpoints, 2x A100-40GB, batch 8:

mode	before	after
DDP	9.19 GiB	9.04 GiB
--distribute_modules	4.24 GiB	3.86 GiB

Median step time is unchanged within run-to-run variance, which I measured at about 25% on this setup, so I would not read a throughput difference either way from these runs.

Correctness

The arithmetic is unchanged. Only the accumulation order inside index_add_ differs, and that was already unspecified on CUDA. tests/test_encode.py compares the encoder gradient against a naive autograd reference at N=8192, D=1024, k=32 and passes.

Full suite on an A100, this branch against main: same results, with the one pre-existing test_auxk_loss failure on both (that one is #142).

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}