Multi-GPU Support

[RiverGao]

Hi! I am trying circuit-tracer in this notebook with Qwen3-8B, and find that the device allocation is "cuda" (i.e. GPU0 by default) for all tensors, which causes OOM in my case. Is there a supported way to use multiple GPUs instead of only GPU0? Thank you!

[hannamw]

Hi! There's no multi-GPU support yet. Basically, you would need to load the underlying TransformerLens model across GPUs. Then, you'd ideally make sure that your transcoders are either on the same GPU as the corresponding model layer, or at least load their weights to the same GPU (if they use lazy loading).

There's some multi-GPU support in TransformerLens at the moment, but I found it a bit slow when I tried to add multi-GPU circuit-tracing a few months ago. Since TL's developers are currently making big changes to the library in order to make it more flexible / adaptable to new models, I am planning to wait until that is finished, and see what the state of multi-GPU support is then, before trying to add it here.