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: