Use Case: Multi-GPU support for 35B SAE training & minor fix for legacy n_devices bug

[summer0517]

Description:

Following the v3.0 release notes regarding the removal of n_devices and the recent comment that multi-device support is now possible via device_map, I want to share our use case to highlight the critical need for multi-GPU support.

1. Our Use Case: 35B Model SAE Training via sae_lens

We run Sparse Autoencoder (SAE) training on large models (e.g., 35B). Even when offloading activations to the CPU (act_store_device="cpu"), the sheer size of the model weights and the forward pass instantly cause OOM on a single GPU.

Previously, we bypassed this by sharding layers across multiple GPUs, passing n_devices indirectly through sae_lens:

model_from_pretrained_kwargs={"trust_remote_code": True, "n_devices": 8}

I understand device_map is now the standard under TransformerBridge, and I plan to test it. Before I do, I have a quick question: I previously tried "device_map": "auto" in older versions, but unlike n_devices, it failed to prevent OOM errors. Does the new implementation include specific optimizations for multi-GPU training to address this?

Thank you for your incredible work on this library!

---

2. Side Note: Quick Fix for Legacy "invalid device ordinal" Bug

If n_devices remains as a fallback, there is a mathematical flaw in transformers_lens/utilities/devices causing crashes when n_layers is not perfectly divisible by n_devices.

The Bug:

offset = index // (cfg.n_layers // cfg.n_devices)

(Example: 62 layers, 8 devices. For layer 60, offset evaluates to 60 // 7 = 8, which is out of bounds.)

The Fix:
Using multiplication first guarantees the resulting offset is strictly < n_devices:

device_offset = (index * cfg.n_devices) // cfg.n_layers

(Example: (60 * 8) // 62 = 7, safely within bounds.)

[jlarson4]

Hi @summer0517!

Thank you for the detailed write up on this, it helps to have true use-case examples for features like this. I am going to put together a PR resolving the issue with n_devices on 3.x for anyone still using the HookedTransformer endpoint.

For your question about device_map="auto" on TransformerBridge:

There aren't new training-specific multi-GPU optimizations beyond what HF/accelerate gives us. The bridge's implementation constrains the dispatch to GPU-only and validates user-supplied maps. n_devices=8 still works, including via model_from_pretrained_kwargs. Under TransformerBridge it now resolves to device_map="balanced" with max_memory={0: "auto", 1: "auto", ..., 7: "auto"} — see resolve_device_map. So your existing kwarg should give you the same balanced GPU-only dispatch you had before.

What I believe is the reason for your past issues with device_map: HF's device_map="auto" is allowed to spill to CPU (and disk) when it doesn't think the model fits on available GPUs. n_devices=N always stayed GPU-only. TransformerLens has handling for this scenario, if you pass device_map="auto" and accelerate decides to spill to CPU, the bridge catches that at load time and raises a clear error rather than letting it OOM later.

I'd recommend trying these options for your 35B SAE setup, in this order:

• "n_devices": 8 – TransformerBridge treats this as balanced + per-GPU max_memory. Likely the closest match to what you had.
• "device_map": "balanced" – explicit equivalent.
• "device_map": "balanced_low_0" – puts fewer model weights on GPU 0, which leaves headroom there for activation buffers and the SAE optimizer state.
• "device_map": {...} + max_memory={0: "70GiB", ...} – if balanced still OOMs because activation memory is asymmetric across layers (e.g. some MoE 35B variants), explicit per-device caps let you reserve room for activations.

Avoid device_map="auto", the bridge will reject it at load time if accelerate spills to CPU.

Let us know how the balanced path works for the 35B run. If you hit any issues, happy to help investigate. This feature is largely untested and having a validated use case is invaluable to TransformerLens. Please let me know if any of these solutions work for you.