Serverless GPUs for bursty training

For most of a month our GPU sat idle. Then a client batch would land, we'd want three LoRA variants trained by the weekend, and suddenly one node wasn't enough. This is the shape of research work at a small company, and it fits an always-on rented box about as well as a suit bought for someone else. The thing that finally pushed us off the default wasn't the bill, it was that on a team small enough that nobody's job is to babysit infrastructure, the box became somebody's job.

What the week actually looked like

Two kinds of work were running through the same hardware.

The first is sustained: reading papers, poking at a ComfyUI graph, running the pipeline against a handful of images over and over to see whether a change to the mask dilation did anything. That work is interactive, it wants a machine that's already warm, and it doesn't care about throughput. It runs for hours at low utilisation and it's completely intolerant of a two minute wait to see a result.

The second is bursty: a finetune kicked off because we got 3,000 new occlusion images, or a batch of client renders that needs to be through the pipeline before a review call. Nothing happens for days, then a lot happens at once, and the only thing that matters is wall clock to finish.

Running both on one rented node means the node is sized for the burst and idle during the rest. Running both on a node sized for the interactive work means the burst takes as long as it takes.

The case against keeping one box on

Cost is the argument everyone reaches for and it's the weakest one, because a reserved node is cheap per hour and it's easy to convince yourself the idle time is fine. Two other things did more damage.

The box accumulates state. Somebody installs a build of xformers to test something, somebody else pins a different diffusers version, a checkpoint gets dropped in /home because the volume was full, and three weeks later a training run fails for a reason nobody can reconstruct. A machine you never rebuild is a machine whose configuration nobody knows.

And a single box serialises the burst. When you want to compare a rank 32 adapter against the same data at a different learning rate, those are two runs, and on one node they're two runs back to back. The whole point of a LoRA at rank 32 is that it's small and cheap enough to train several of, and that advantage disappears if the hardware makes them sequential.

The split we landed on

Dedicated Amazon EC2 nodes for sustained training and R&D. Modal for the bursts, both training and inference. S3 in the middle holding source images, generated outputs and logs, so neither side owns the data. (We also ran diffusion inference on ShaktiCloud, an India-hosted GPU cloud, for work that needed to stay closer to home; that's a separate decision from this one and it didn't change the burst story.)

The dedicated node is where you keep a shell open. The serverless functions are where you run things you want four of.

1. The image is the deployment

The part of Modal that changed how I work isn't the GPU, it's that the container image is defined in Python next to the function that runs in it. There's no Dockerfile drifting out of sync with what's actually installed on the box, and no ambiguity about what a run was executed against.

# modal_app.py
"""Modal app for AuraX burst workloads: LoRA training and batch inference.
 
The image below is the whole environment. If a run reproduces here, it reproduces
anywhere, which is the property the long-lived EC2 node had stopped having.
"""
 
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",   # AdamW8bit
        "peft==0.12.0",
        "boto3==1.34.162",
        "pillow==10.4.0",
    )
    # Kohya-ss is pinned to a commit, not a branch. A training script that moves
    # under you is the same failure as a box whose packages nobody tracks.
    .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 are a volume, not a layer. See below for why.
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")]

Pinning Kohya-ss to a commit rather than a branch is the single line in that file I'd defend hardest. Training scripts in this ecosystem move fast and quietly, and a LoRA you can't retrain is a LoRA you can't fix.

2. Weights live in a volume, not in the image

The rule I settled on is that a layer holds anything that changes when the code changes, and a mount holds anything that changes on its own schedule. System packages, dependencies and the training script are code: if they move, the behaviour of a run moves with them, and you want a rebuild to be the only way that can happen. Datasets, generated outputs and logs are the opposite, and baking them in would mean a rebuild every time somebody adds a hundred images.

Model weights sit across that line, which is why they're the awkward case. They behave like data, being large and opaque and rarely edited. They behave like code, because swapping the checkpoint changes every output, so a run is only reproducible if you know which one it used. Bake them and you buy reproducibility by rebuilding gigabytes for an unrelated dependency bump. Mount them and you buy fast rebuilds at the cost of a mutable input, unless you version by path and never overwrite in place, which is what we ended up doing.

The first version of this baked the checkpoint into the image. It worked, and it was wrong. A base checkpoint is large, it changes rarely, and it has nothing to do with the code, so putting it in a layer means every dependency bump rebuilds and re-uploads gigabytes for no reason. Worse, the image becomes the thing that determines which checkpoint a run used, which is exactly the coupling you were trying to avoid.

# modal_app.py (continued)
 
@app.function(
    image=image,
    volumes={"/weights": weights},
    secrets=secrets,
    timeout=60 * 60,
)
def sync_weights(keys: list[str]) -> None:
    """Pull base checkpoints and adapters from S3 into the shared volume.
 
    Run this once when a checkpoint changes. Every training or inference function
    then mounts /weights read-mostly and starts against whatever is there, so the
    checkpoint in use is a property of the volume rather than of the image.
    """
    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"skip {key} (present, {dest.stat().st_size / 1e9:.1f} GB)")
            continue
        dest.parent.mkdir(parents=True, exist_ok=True)
        print(f"pull {key}")
        s3.download_file(bucket, key, str(dest))
 
    # Commit makes the new files visible to other containers mounting this volume.
    weights.commit()
    print("volume committed")

There's a subtlety here that cost me an afternoon. A volume mounted into a running container is a snapshot; a function that started before commit() will not see the new file. If you kick off training in the same session that uploaded the checkpoint, order it explicitly rather than assuming.

3. The burst entrypoint

With the image and the weights sorted, the actual training function is thin. It shells out to the same Kohya-ss invocation we'd run on the dedicated node, with the same rank, alpha, learning rate and optimiser, so a run is comparable across the two environments.

# train_burst.py
"""Fan out LoRA training runs. One container per run, all of them at once."""
 
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=60 * 60 * 6,
    retries=modal.Retries(max_retries=1, backoff_coefficient=1.0),
)
def train_lora(run: dict) -> str:
    """Train one adapter and return the path of the artifact it produced.
 
    `run` carries only the things we vary between experiments. Everything else is
    fixed on purpose, so that a diff between two runs is a diff in the data or in
    one hyperparameter, never in the environment.
    """
    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"
    print(f"finished {name} -> {artifact}")
    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("ready:", artifact)

train_lora.map(runs) is the whole argument in one line. Four runs, four containers, four GPUs, one wall clock. On a single node that's four runs in sequence and a decision about which one you care about enough to run first. I stopped making that decision, which is worth more to a research loop than it sounds, because the runs you skip are disproportionately the ones that would have surprised you.

Cold starts are a real cost, and I'm not giving you a number

Every burst container pays for pulling the image, mounting the volume and loading a large checkpoint into VRAM before it does any work you care about. For a six hour training run that overhead is noise. For a batch of a few dozen renders it is not noise, and there's a crossover point below which you should just run the job on the node that's already warm.

I would like to tell you where that crossover sits for us. I can't, and I'd rather say so than make one up. We never instrumented cold starts properly, we tuned by feel, and the only latency figure I trust from that whole period is a lab measurement of roughly 12 seconds for a full-body try-on on an RTX 4090, single image, nothing to do with production serving. Every number in the code above is a configuration value we set, not something we measured. If you're evaluating this for your own work, the honest advice is to measure it yourself on your checkpoint size, because that's the variable that dominates and it's the one I can't hand you.

There's a shape of this that throughput numbers would hide even if I had them. You change one line in a sampler config and run it to see what happened. The job is correct, the batch completes, nothing is slow in any sense a dashboard would report. But the gap between pressing enter and seeing a pixel is long enough that you tab away, and by the time the render lands you've lost the thread of why you made the change. So you stop asking one question at a time and start batching them, which is a worse way to work, because the value of an interactive loop is that each result changes what you ask next. It reads like a capacity problem and it's a latency problem, and I spent longer than I'd like to admit throwing parallelism at it.

What I can say qualitatively is which knobs mattered. Keeping containers warm between calls during an active session helps a lot and costs you idle time you're paying for, so it's a session-level decision rather than a permanent setting. Loading the checkpoint once per container and processing many images inside that container beats one image per invocation by a wide margin. And keeping the checkpoint on a volume rather than pulling it from S3 per container removes the largest single chunk of that startup, which is the practical reason the volume design above exists.

# batch_infer.py
"""Batch inference: amortise the load over many images inside one container."""
 
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,      # keep the container warm during an active batch
    max_containers=8,          # ceiling, not a target
)
class TryOnWorker:
    @modal.enter()
    def load(self):
        """Runs once per container, before any request is served."""
        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")
        print("pipeline resident")
 
    @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

The @modal.enter() split is the whole trick: model loading happens once per container lifetime, request handling happens per call. Getting that boundary wrong is the most common way a serverless GPU setup ends up looking slow and expensive, and it looks exactly like a cold start problem when it's really a code structure problem.

What you give up

Everything above is the case for the split. The cost is debugging, and it's a real one.

When a training run dies on a box you own, you SSH in and look. The process is still there, the half-written checkpoint is still on disk, nvidia-smi tells you what the memory actually did, and you can open a shell against the same environment and poke at the tensor that came out wrong. When a burst function dies, the container is gone. What you have is the logs you had the foresight to print, and if you didn't print the thing you now need, you get to run the job again to find out.

That changes how you write the code more than how you debug it. You print more and print earlier, on the theory that a log line you don't need costs nothing while one you didn't write costs a re-run. And you care much more about whether a failure reproduces locally, because a failure that only happens in the burst environment is one you're investigating through a keyhole.

So anything we didn't already understand stayed on the dedicated node until we did. Serverless is a good place to run a job whose failure modes you can predict, and a bad place to find out what a new one is.

What I'd do differently

I'd instrument the startup path on day one, even crudely. A timestamp at container entry, one after the checkpoint loads, one at first output, written to the same S3 bucket as everything else. It costs almost nothing and it's the difference between this post having an answer in it and this post admitting it doesn't.

I'd stop treating the two environments as different systems earlier. Once the training invocation was byte-identical between the node and the burst function, a whole category of "it works on the box but not on Modal" disappeared, and that identity should have been the first thing I built rather than something I converged on.

The next piece of this is what to do with four adapters once you've trained them in parallel, which turns out to be a much harder problem than getting the GPUs.