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():
    ...

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, A10G and T4 instances all support up to 4 GPUs, but we are working on supporting more options in the future. There will be increased startup times for functions requesting more than 2 GPUs.

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 although this is a beta feature and may have latency issues:

@stub.function(gpu=modal.gpu.A100(memory=80))
def my_a100_80GB_function():
    ...

Multi-GPU tasks

Modal also supports multi-GPU tasks. This means each container will have access to multiple GPUs. We support up to 4 or 6 GPUs, depending on the type.

@stub.function(gpu=modal.gpu.A100(count=4))
def my_a100_80GB_function():
    ...

Examples

Take a look at some of our examples that use GPUs: