GPU acceleration
If you have code or use libraries that are GPU accelerated, you can attach the
first available GPU to your function by passing the gpu="any" argument to the
@stub.function decorator:
import modal
stub = modal.Stub()
@stub.function(gpu="any")
def my_function():
# code here will be executed on a machine with an available GPU
...Specifying GPU type
When gpu="any" is specified, your function runs in a container with access to
a supported GPU. Currently this gives you Nvidia
Tesla T4 or
A10G instances. If
you need more control, you can pick a specific GPU type by changing this
argument.
@stub.function(gpu="T4")
def my_t4_function():
...
@stub.function(gpu="A10G")
def my_a10g_function():
...
@stub.function(gpu="A100")
def my_a100_function():
...
@stub.function(gpu="L4")
def my_l4_function():
...For information on all valid values for the gpu parameter see
the modal.gpu reference page.
Specifying GPU count
You may also specify the number of GPUs to attach to your function by using the
object form of the gpu parameter for your desired
GPU:
@stub.function(gpu=modal.gpu.A10G(count=2))
def my_a10g_function():
...Currently A100 and T4 instances support up to 8 GPUs, and A10G support up to 4 GPUs. Note that requesting more than 2 GPUs per container will usually result in larger wait times.
A100 GPUs
Modal’s fastest GPUs are the A100s, which are NVIDIA’s flagship data center chip. They have beefier hardware and more GPU memory.
To request an A100 with 40 GB of GPU memory, replace the gpu="any" argument
with gpu="A100":
@stub.function(gpu="A100")
def my_a100_function():
...Modal also support 80GB A100s:
@stub.function(gpu=modal.gpu.A100(memory=80))
def my_a100_80GB_function():
...Examples
Take a look at some of our examples that use GPUs: