Serverless GPUs for bursty training

For long stretches, our GPU instance sat idle. Then client data arrived, we needed three LoRA variants trained concurrently by the weekend, and a single reserved node created an immediate bottleneck.

The cost was only part of the problem. On small engineering teams without dedicated MLOps staff, maintaining a long-lived GPU server turns into a configuration liability.

A server you never rebuild is a machine whose configuration state is unknown. That debt appears when an experimental run fails for reasons that cannot be reproduced.

Chapter 0: Defining burst workloads

A burst workload has two defining characteristics: it is separated by long idle windows, and its primary efficiency metric is total wall-clock time to completion.

Interactive research (debugging ComfyUI nodes, testing diffusion mask dilations on test images) requires an always-warm instance with zero cold-start overhead. Burst jobs (training four LoRA adapters in parallel across parameter sweeps) require elastic fan-out.

Sizing a static instance for burst peaks leaves hardware idle most of the month; sizing it for interactive work forces training experiments into sequential queues.

Step 1: Decouple interactive R&D from elastic batch jobs

Averaging interactive and batch requirements yields an infrastructure setup that fails both constraints:

The hybrid architecture uses dedicated EC2 instances for interactive experimentation, Modal for parallelized training sweeps, and Amazon S3 as the centralized artifact and dataset store.

The dedicated instance maintains an active development environment. Serverless workers handle parallelized hyperparameter sweeps.

A single static node serializes research: testing four rank-32 adapters at different learning rates requires running them back-to-back. Serverless fan-out executes all four runs concurrently.

Step 2: Declare infrastructure environments in code

Define the container in Python next to the function it runs, and every training run executes against the same pinned specification:

# modal_app.py
import modal
 
app = modal.App("aurax-burst")
 
CUDA_TAG = "12.4.1-devel-ubuntu22.04"
 
image = (
    modal.Image.from_registry(f"nvidia/cuda:{CUDA_TAG}", add_python="3.11")
    .apt_install("git", "libgl1", "libglib2.0-0")
    .pip_install(
        "torch==2.4.0",
        "diffusers==0.30.0",
        "transformers==4.44.0",
        "accelerate==0.33.0",
        "safetensors==0.4.4",
        "bitsandbytes==0.43.3",
        "peft==0.12.0",
        "boto3==1.34.162",
        "pillow==10.4.0",
    )
    .run_commands(
        "git clone https://github.com/kohya-ss/sd-scripts /opt/sd-scripts",
        "cd /opt/sd-scripts && git checkout 5b6c2e2 && pip install -r requirements.txt",
    )
    .env({"HF_HOME": "/weights/hf", "PYTHONUNBUFFERED": "1"})
)
 
weights = modal.Volume.from_name("aurax-weights", create_if_missing=True)
outputs = modal.Volume.from_name("aurax-outputs", create_if_missing=True)
 
secrets = [modal.Secret.from_name("aws-s3"), modal.Secret.from_name("hf-token")]

Pin the third-party training repo to an exact commit hash. Otherwise an upstream change lands in the middle of a sweep and the run you want to reproduce no longer exists.

Step 3: Decouple model checkpoints from container layers

Separating code dependencies from large binary weights prevents unnecessary image rebuilds:

Baking 20GB+ foundation model weights into container image layers forces multi-gigabyte uploads on minor dependency changes. Storing weights in network volumes enables sub-second container startups with cached filesystem layers.

# modal_app.py (continued)
 
@app.function(
    image=image,
    volumes={"/weights": weights},
    secrets=secrets,
    timeout=3600,
)
def sync_weights(keys: list[str]) -> None:
    """Synchronizes base checkpoints and adapters from S3 to shared volume."""
    import os
    from pathlib import Path
    import boto3
 
    s3 = boto3.client("s3")
    bucket = os.environ["AURAX_WEIGHTS_BUCKET"]
 
    for key in keys:
        dest = Path("/weights") / key
        if dest.exists():
            print(f"Skipping {key} (already present)")
            continue
        dest.parent.mkdir(parents=True, exist_ok=True)
        print(f"Downloading {key} from S3...")
        s3.download_file(bucket, key, str(dest))
 
    weights.commit()
    print("Volume synchronized and committed.")

Step 4: Fan out parameter sweeps across independent GPU containers

With container environments and weight storage decoupled, training sweeps execute in parallel across independent L40S/A100 instances:

# train_burst.py
import modal
from modal_app import app, image, outputs, secrets, weights
 
@app.function(
    image=image,
    gpu="L40S",
    volumes={"/weights": weights, "/out": outputs},
    secrets=secrets,
    timeout=21600,
    retries=modal.Retries(max_retries=1, backoff_coefficient=1.0),
)
def train_lora(run: dict) -> str:
    """Executes single adapter training run inside an isolated container."""
    import subprocess
 
    name = run["name"]
    cmd = [
        "accelerate", "launch", "--num_cpu_threads_per_process", "8",
        "/opt/sd-scripts/sdxl_train_network.py",
        "--pretrained_model_name_or_path", "/weights/flux/base.safetensors",
        "--train_data_dir", f"/weights/datasets/{run['dataset']}",
        "--output_dir", "/out/lora",
        "--output_name", name,
        "--resolution", "1024,1024",
        "--network_module", "networks.lora",
        "--network_dim", "32",
        "--network_alpha", "16",
        "--learning_rate", str(run.get("lr", 1e-4)),
        "--lr_scheduler", "cosine",
        "--train_batch_size", "4",
        "--gradient_accumulation_steps", "4",
        "--gradient_checkpointing",
        "--optimizer_type", "AdamW8bit",
        "--mixed_precision", "bf16",
        "--save_every_n_epochs", "1",
    ]
    subprocess.run(cmd, check=True)
    outputs.commit()
 
    artifact = f"/out/lora/{name}.safetensors"
    return artifact
 
@app.local_entrypoint()
def main():
    runs = [
        {"name": "drape_lr1e4", "dataset": "drape_5k", "lr": 1e-4},
        {"name": "drape_lr5e5", "dataset": "drape_5k", "lr": 5e-5},
        {"name": "occ_lr1e4", "dataset": "occlusion_3k", "lr": 1e-4},
        {"name": "occ_lr5e5", "dataset": "occlusion_3k", "lr": 5e-5},
    ]
    for artifact in train_lora.map(runs):
        print("Completed run artifact:", artifact)

train_lora.map(runs) allocates four containers simultaneously. Total wall-clock time equals the duration of a single run rather than the sum of four sequential jobs.

Step 5: Amortize model loading overhead in batch inference

For batch inference, container cold-start overhead must be amortized across multiple requests:

# batch_infer.py
import modal
from modal_app import app, image, outputs, secrets, weights
 
@app.cls(
    image=image,
    gpu="L40S",
    volumes={"/weights": weights, "/out": outputs},
    secrets=secrets,
    scaledown_window=300,      # Retain warm container for 5 minutes
    max_containers=8,
)
class TryOnWorker:
    @modal.enter()
    def load(self):
        """Loads weights once during container initialization."""
        import torch
        from diffusers import FluxFillPipeline
 
        self.pipe = FluxFillPipeline.from_pretrained(
            "/weights/flux/fill", torch_dtype=torch.bfloat16
        ).to("cuda")
        self.pipe.load_lora_weights("/weights/lora/aurax_merged.safetensors")
 
    @modal.method()
    def render(self, job: dict) -> str:
        out = self.pipe(
            prompt=job["prompt"],
            image=job["image"],
            mask_image=job["mask"],
            num_inference_steps=30,
            height=1024,
            width=1024,
        ).images[0]
        path = f"/out/renders/{job['id']}.png"
        out.save(path)
        return path

@modal.enter() loads weights into GPU VRAM once per container lifecycle, so every later render call pays only for the forward pass.

Operational considerations and failure modes

• Cold-start latency boundaries: If individual inference jobs execute in under 2 seconds, a 30-second cold-start penalty dominates total latency unless warm pools (scaledown_window) are maintained.
• Ephemeral container debugging: Terminated serverless containers cannot be inspected via SSH. Structure scripts to log memory statistics, loss values, and error traces to S3 or persistent network volumes on failure.

When this is the wrong choice

• The GPU is busy most of the month. This split trades a lower average price for cold starts, an image build, and a control plane. A node running at steady utilization has no idle hours to reclaim, and reserved capacity is cheaper per GPU-hour than per-second billing.
• The work is interactive. Debugging ComfyUI nodes or checking mask dilations on test images means a live shell and a fast edit-run loop. Paying 30 to 60 seconds of container init per attempt makes that loop unusable. That workload belongs on the dedicated instance, which is why the architecture keeps one.
• There is only one run. Fan-out pays because four adapters finish in the wall-clock time of one. A single training job gets no speedup from train_lora.map, and you have added an image definition, a weights sync, and a deploy step to a script that would have run on the box you already had.
• You cannot get diagnostics out of a dead container. If the failure you are chasing needs a live process to inspect, a terminated serverless container gives you nothing back. Reproduce it on the EC2 node first, then move the sweep.