[Bug Report] "unsloth/llama-3.2-3b-instruct" with 2x 3060 throws RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cuda:1)

[jasonrichdarmawan]

If you are submitting a bug report, please fill in the following details and use the tag [bug].

Describe the bug
My requirement is to load unsloth/llama-3.2-3b-instruct (drop-in replacement to meta-llama/llama-3.2-3b-instruct because my request was rejected) from a local folder because I put the model weights on a NFS

Code example

# %%

model_name = "Llama-3.2-3B-Instruct"
print(f"Loading {model_name} model")
model_path = os.path.join(
  "/media/tao/disk4T/jason/transformers",
  f"unsloth/{model_name}",
)
hf_model = AutoModelForCausalLM.from_pretrained(
  model_path,
)
tokenizer = AutoTokenizer.from_pretrained(
  model_path,
)

# %%

model = HookedTransformer.from_pretrained(
  model_name=f"meta-llama/{model_name}",
  hf_model=hf_model,
  tokenizer=tokenizer,
  device="cuda",
  n_devices=2,
)

model.generate("hello world", max_new_tokens=200, do_sample=True, temperature=0.3)

System Info
Describe the characteristic of your environment:

• Describe how transformer_lens was installed: pip
• What OS are you using? Linux
• Python version: 3.11

Additional context
Add any other context about the problem here.

Checklist

• I have checked that there is no similar issue in the repo (required)

[jasonrichdarmawan]

This is a minimal-impact issue

However, I hope that TransformerLens can better support GPU-limited researchers. My definition of success is: simpler ideas are easier to build on and are more likely to serve as useful baselines for future researchers. I know—it’s a bit ironic coming from me, as I’m also a GPU-poor researcher

PS: My development node has 2× 3060 GPUs, which I can use at no cost. I can rent a single 3090, but it would put a hole in my wallet since I’m a student. I only use a more powerful GPU when absolutely necessary. For context, renting a 3090 for 12 hours costs $2.64, which is equivalent to about two meals

PS2: I’m new to the TransformerLens library. I plan to try fixing this issue in August

PS3: Why does TransformerLens hard-code the list of supported models? I believe the nnsight library doesn’t do this; if I’m not mistaken, nnsight delegates the device_map argument to PyTorch. I haven’t used nnsight or TransformerLens in a while, because I was concerned that using TransformerLens might increase the learning curve of my research and make it less accessible to others. So far, I’ve been relying on raw PyTorch hooks. However, someone at ARENA recommended using TransformerLens, so I think it’s becoming inevitable that I make the switch

[jlarson4]

Hello @jasonrichdarmawan! Sorry it has taken so long for someone to get back to you on this, I recently took over maintaining this project and am reviewing the issue back log.

We had the same concern about the narrow model support for HookedTransformer, and have been working since around the time of this issue posting to create a new system that is more flexible. TransformerLens 3.0 is now available, and the TransformerBridge system now supports by architecture, rather than per model. This has allowed us to extend our reach to more models. TransformerBridge currently supports about 9,000 models, with more on the way. The specific unsloth/lllama-3.2 model you reference has not been explicitly verified yet, but I think your model should be runnable on the new system. Here is the migration guide, if you are still interested in trying it out. I also have a PR open for multi-device support on TransformerBridge, but have not yet had time to test it on a multi-gpu device. It should in theory solve your problem when included in the next release.

Thanks!