multimodel_inference.py

import torch
import threading
from transformers import AutoTokenizer, AutoModelForCausalLM

# Model id (plain base or instruct — your choice)
model_id = "Qwen/Qwen3-0.6B"

# Shared tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

# Load models on separate GPUs
model0 = AutoModelForCausalLM.from_pretrained(
    model_id,
    load_in_8bit=True,
    device_map={"": 0},
    trust_remote_code=True
)
model1 = AutoModelForCausalLM.from_pretrained(
    model_id,
    load_in_8bit=True,
    device_map={"": 1},
    trust_remote_code=True
)

# Prompts
prompt0 = "The mean of a probability density function is defined as"
prompt1 = "The variance of a probability density function is defined as"

# Prepare inputs
inputs0 = tokenizer(prompt0, return_tensors="pt").to("cuda:0")
inputs1 = tokenizer(prompt1, return_tensors="pt").to("cuda:1")

# Results dictionary
results = {}

def run_inference(name, model, inputs, device):
    with torch.no_grad():
        # Generate continuation
        outputs = model.generate(
            **inputs,
            max_new_tokens=100,
            temperature=0.7,
            top_p=0.9,
            return_dict_in_generate=True,
            output_hidden_states=True
        )
    # Get hidden states from last step
    hidden = outputs.hidden_states[-1][-1]   # [batch, seq_len, hidden_dim]
    # Decode output sequence
    decoded = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
    # Save
    results[name] = {"hidden": hidden, "decoded": decoded}

# Create threads for parallel execution
t0 = threading.Thread(target=run_inference, args=("model0", model0, inputs0, "cuda:0"))
t1 = threading.Thread(target=run_inference, args=("model1", model1, inputs1, "cuda:1"))

# Run both
t0.start()
t1.start()
t0.join()
t1.join()

# Print results
print("Model0 answer:")
print(results["model0"]["decoded"], "\n")

print("Model1 answer:")
print(results["model1"]["decoded"], "\n")

print("Hidden0 shape:", results["model0"]["hidden"].shape)
print("Hidden1 shape:", results["model1"]["hidden"].shape)