1. Summary

Currently, in the TransformersModel class of smolagents, the model is being allocated to a device using the .to(device) method. This approach restricts the usage to only one CUDA card. I would like to request the addition of an option to use device_map="auto" to enable the utilization of multiple GPUs.

2. Problem Description

1. Limited GPU Utilization: With the current implementation, smolagents can only make use of a single CUDA device. This is a significant limitation, especially when dealing with large models that could benefit from parallel processing across multiple GPUs. For example, running a large language model for text generation tasks would be much faster if multiple GPUs could be used.
2. Lack of Scalability: As the size of models and the complexity of tasks grow, the ability to scale by using multiple GPUs becomes crucial. The current .to(device) method does not provide this scalability.

There is the error I meet. Although I have many cudas, it only work on cuda:0.

OutOfMemoryError: CUDA out of memory. Tried to allocate 384.00 MiB. GPU 0 has a total capacity of 23.65 GiB of which 14.06 MiB is free. Including non-PyTorch memory, this process has 23.63 GiB memory in use. Of the allocated memory 23.12 GiB is allocated by PyTorch, and 60.91 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

3. Proposed Solution

1. Add device_map="auto" Option: In the TransformersModel class, when initializing the model, add an option to use device_map="auto" instead of just .to(device). This would allow the model to automatically distribute across available GPUs, optimizing performance. For example, the code could be updated as follows:

try:
    self.model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
except Exception as e:
    # Handle the exception as before
    pass

2. Configuration for Users: Provide a way for users to easily configure this option. This could be through a configuration file or an additional parameter when initializing the relevant classes in smolagents.

4. Benefits

1. Performance Improvement: Using multiple GPUs will significantly improve the inference speed of the models in smolagents. This will lead to faster response times, especially for computationally intensive tasks.
2. Scalability: It will make smolagents more suitable for large - scale projects and research, where the ability to scale the computing resources is essential.

Thank you for considering this feature request. I believe it will greatly enhance the capabilities of smolagents.