How to use `device_map="auto"` option

[bongjunj]

Hi,

I usually run a model with device_map="auto" option to utilize all the cards I have.

output = model.generate(
    inputs,
    max_length=100,
    num_return_sequences=1,
    device_map="auto",
    attention_mask=torch.ones_like(inputs),
    pad_token_id=tokenizer.eos_token_id,
)

But it seems like using llm = Syncode(...); llm.infer(...) API, there is no direct way to turn the option on.
Is there any way to do so in the current setting?

Thank you

[Sindre01]

^Any updates on this?

Tried adding device_map="auto" to load_model() function, but still only utilizes one gpu:

´´´
def load_model(model_name, device, quantize, auto_device_map=True)

    if model_name == 'test':
        model = AutoModelForCausalLM.from_pretrained('bigcode/tiny_starcoder_py').to(device)
    elif model_name == 'test-instruct':
        model = AutoModelForCausalLM.from_pretrained("rahuldshetty/tiny-starcoder-instruct")
    else:
        if (quantize):
            if auto_device_map:
                model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", cache_dir=HF_CACHE, token=HF_ACCESS_TOKEN, trust_remote_code=True).eval().to(device) 

....
´´´

[shubhamugare]

It should be fixed now!

[Sindre01]

Possible to quickly release this feature to pypi package?