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: