Add --distribute_latents: shard the latent dimension across ranks

[alepot55]

Summary

Adds --distribute_latents, a third way to spread SAE training across GPUs: shard the latent dimension of every coder instead of assigning whole hookpoints to ranks. Rank r owns latents [r*M/W, (r+1)*M/W) of every coder.

I opened this after profiling where the memory in torchrun training actually goes, so every number below is measured rather than modelled. Happy to split, rename, or restructure it if the shape is wrong for you.

Stacked on #141, which it needs for the chunked weight-gradient. That commit is the first of the two here; the second is the feature.

Why not --distribute_modules

--distribute_modules requires the hookpoint count to divide evenly by the world size, and it leaves each rank materialising the world-sized activation for every hookpoint, including the ones it does not own. Sharding latents instead keeps every hookpoint on every rank, so any GPU count works, and the [N, num_latents] pre-activation shrinks with the world size rather than staying fixed.

Design

Two collectives per coder per forward, none in the backward:

1. Each rank offers its local top-k as candidates and all ranks agree on the same global winners. The global top-k of size k can draw at most k entries from any single rank, so local candidates are sufficient.
2. Each rank decodes only the winners it owns; the partial reconstructions are summed. The backward of that sum is the identity, because every rank computes the same loss from the same summed output.

b_dec and W_skip are replicated rather than sharded, so each rank contributes 1/W of them and only their gradients are reduced. The sharded rows see every token on every rank, so their gradients are already complete.

Two details worth flagging for review:

• Survivors are kept as a ragged (token, latent, value) list, not as the [N, k] candidate block with losers zeroed. With the block layout both the decode and the encoder weight-gradient stay O(N*k*d) however few entries survive, so per-rank work does not fall as the world grows, which defeats the point. Compacting makes them O(nnz*d) with nnz = N*k/world_size. Measured, at 8 GPUs: 944 -> 430 ms per step at batch 8, and 6988 -> 2752 ms at batch 64. Padding to the busiest token instead would have recovered only about 2.5x at k=32 and 8 ranks, against world_size here.
• Shards are drawn from the global initialization and sliced. Every rank runs the same seed, so drawing directly at the sharded shape hands all of them identical rows and the dictionary silently collapses to num_latents / world_size distinct latents. It also makes a sharded run reproduce an unsharded run with the same seed, row for row. test_shard_initialization_matches_global_draw guards this.

Checkpoints are gathered on rank 0 and written unsharded, so they load with the ordinary SparseCoder.load_from_disk and a sharded run is interchangeable with a single-GPU one.

Measurements

SmolLM2-135M, 30 hookpoints, batch 8 per rank, ctx 512, A100-40GB, 8-bit Adam. Peak allocated per rank and median step time. Every column processes the same number of tokens per step.

GPUs	DDP	--distribute_modules	--distribute_latents
2	9.04 GiB / 0.33 s	3.86 GiB / 0.38 s	3.65 GiB / 0.33 s
4	9.04 GiB	not runnable	2.24 GiB / 0.36 s
8	9.04 GiB / 0.30 s	not runnable	1.69 GiB / 0.43 s

"Not runnable" is AssertionError: Number of modules must be divisible by world size, since 30 is not divisible by 4 or 8.

The trade-off, stated plainly

Against --distribute_modules this is better on every axis I measured: lighter, faster, and it runs at GPU counts where the other refuses to start.

Against DDP it buys memory with throughput. At 8 GPUs and batch 8 that is 5.4x less memory for 1.4x the step time. The slowdown grows with the batch, because the collectives scale with the gathered batch while the per-rank shard does not: at batch 64 the same run is 1.9x lighter and 2.1x slower. A phase breakdown at 8 GPUs and batch 64 puts the top-k all-gather at 9.2% of the step and the reconstruction all-reduce at 6.8%, with the rest being work DDP does too. The README says all of this and suggests preferring a smaller per-rank batch with more gradient accumulation at high rank counts.

Step times on this setup vary about 25% run to run, so I would only trust the large ratios above, not differences of a few percent.

Correctness

New tests/test_latent_parallel.py, the first distributed tests in the suite. They run gloo on CPU, so they need no GPU and no launcher; the equivalence being checked is algebraic. Each case builds a full coder and a shard holding the same weights, then asserts the shard matches the unsharded coder on the reconstruction, every loss term, the gradients of the rows it owns, the reduced gradients of the replicated parameters, and the gathered checkpoint. Covered: k in {1, 4, 32}, multi-TopK, skip connections, AuxK, and AuxK with skip connections.

End to end on GPU, FVU at the first step matches DDP across world sizes (0.6385 to 0.6408 against DDP's 0.6391 to 0.6407).

Existing suite on an A100: main gives 1 failed, 10 passed (the failure is #142), this branch gives the same single failure and 18 passed.

Not supported

groupmax partitions the latent dimension itself, so it cannot be split this way and is rejected with an explicit error. --distribute_latents cannot be combined with --distribute_modules, or with the ce and kl losses, matching the existing restriction on --distribute_modules.

[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.}