Use strings for GPU config (#1066)

Co-authored-by: Charles Frye <charles@modal.com>

Author: erikbern

06_gpu_and_ml/dreambooth/diffusers_lora_finetune.py

@@ -265,9 +265,7 @@ class TrainConfig(SharedConfig):
 
 @app.function(
     image=image,
-    gpu=modal.gpu.A100(  # fine-tuning is VRAM-heavy and requires a high-VRAM GPU
-        count=1, size="80GB"
-    ),
+    gpu="A100-80GB",  # fine-tuning is VRAM-heavy and requires a high-VRAM GPU
     volumes={MODEL_DIR: volume},  # stores fine-tuned model
     timeout=1800,  # 30 minutes
     secrets=[huggingface_secret]

06_gpu_and_ml/embeddings/text_embeddings_inference.py

@@ -14,7 +14,7 @@
 
 import modal
 
-GPU_CONFIG = modal.gpu.A10G()
+GPU_CONFIG = "A10G"
 MODEL_ID = "BAAI/bge-base-en-v1.5"
 BATCH_SIZE = 32
 DOCKER_IMAGE = (

06_gpu_and_ml/embeddings/wikipedia/main.py

@@ -7,7 +7,7 @@
 # We first set out configuration variables for our script.
 ## Embedding Containers Configuration
 GPU_CONCURRENCY = 100
-GPU_CONFIG = modal.gpu.A10G()
+GPU_CONFIG = "A10G"
 MODEL_ID = "BAAI/bge-small-en-v1.5"
 MODEL_SLUG = MODEL_ID.split("/")[-1]
 BATCH_SIZE = 512

06_gpu_and_ml/llm-serving/chat_with_pdf_vision.py

@@ -158,7 +158,7 @@ def download_model():
 
 @app.cls(
     image=model_image,
-    gpu=modal.gpu.A100(size="80GB"),
+    gpu="A100-80GB",
     container_idle_timeout=10 * MINUTES,  # spin down when inactive
     volumes={"/vol/pdfs/": pdf_volume, CACHE_DIR: cache_volume},
 )

06_gpu_and_ml/llm-serving/trtllm_llama.py

@@ -154,7 +154,7 @@ def download_model():
 # about two quadrillion per second on an H100 SXM.
 
 N_GPUS = 1  # Heads up: this example has not yet been tested with multiple GPUs
-GPU_CONFIG = modal.gpu.H100(count=N_GPUS)
+GPU_CONFIG = f"H100:{N_GPUS}"
 
 DTYPE = "float16"  # format we download in, regular fp16
 QFORMAT = "fp8"  # format we quantize the weights to

06_gpu_and_ml/llm-serving/vllm_inference.py

@@ -89,7 +89,7 @@
 
 @app.function(
     image=vllm_image,
-    gpu=modal.gpu.H100(count=N_GPU),
+    gpu=f"H100:{N_GPU}",
     container_idle_timeout=5 * MINUTES,
     timeout=24 * HOURS,
     allow_concurrent_inputs=1000,

06_gpu_and_ml/llm-structured/jsonformer_generate.py

@@ -58,7 +58,7 @@ def download_model():
 # The generate function takes two arguments `prompt` and `json_schema`, where
 # `prompt` is used to describe the domain of your data (for example, "plants")
 # and the schema contains the JSON schema you want to populate.
-@app.function(gpu=modal.gpu.A10G(), image=image)
+@app.function(gpu="A10G", image=image)
 def generate(prompt: str, json_schema: dict[str, Any]) -> dict[str, Any]:
     from jsonformer import Jsonformer
     from transformers import AutoModelForCausalLM, AutoTokenizer

06_gpu_and_ml/llm-structured/outlines_generate.py

@@ -106,7 +106,7 @@ def import_model(model_name):
 # We specify that we want to use the Mistral-7B model, and then ask for a character, and we'll receive structured data with the right schema.
 
 
-@app.function(image=outlines_image, gpu=modal.gpu.A100(size="40GB"))
+@app.function(image=outlines_image, gpu="A100-40GB")
 def generate(
     prompt: str = "Amiri, a 53 year old warrior woman with a sword and leather armor.",
 ):

06_gpu_and_ml/stable_diffusion/image_to_image.py

@@ -89,7 +89,7 @@ def download_models():
     snapshot_download("stabilityai/sdxl-turbo", ignore_patterns=ignore)
 
 
-@app.cls(gpu=modal.gpu.A10G(), container_idle_timeout=240)
+@app.cls(gpu="A10G", container_idle_timeout=240)
 class Model:
     @modal.enter()
     def enter(self):

06_gpu_and_ml/text-to-video/mochi.py

@@ -133,7 +133,7 @@ def download_model(revision="83359d26a7e2bbe200ecbfda8ebff850fd03b545"):
         OUTPUTS_PATH: outputs,  # videos will be saved to a distributed volume
         MODEL_PATH: model,
     },
-    gpu=modal.gpu.H100(count=1),
+    gpu="H100",
     timeout=1 * HOURS,
 )
 class Mochi: