Using more than one GPU
Run a worker across several cards, where every rank executes your handler and only the first one answers.
Introduction
A worker across several cards runs one serving process per GPU. That one fact is the whole difference from a single-GPU handler, and it has nothing to do with the arithmetic.
There are no weights here, so a failure is the world coming up, not a model.
The app
"""One serving process per GPU, and a handler that every one of them runs."""
from __future__ import annotations
from runware_serverless import endpoint, serve
DEFAULT_SIZE = 1 << 22
@serve
class MultiGPU:
requirements = ["torch==2.14.1"]
def load(self) -> None:
# Nothing here joins the world. @serve runs this inside the platform's
# initialization, which sets this rank's device and brings the process
# group up before load, then waits on a barrier after it. A second
# init_process_group raises rather than being ignored.
import torch
import torch.distributed as dist
self.world = dist.get_world_size() if dist.is_initialized() else 1
self.rank = dist.get_rank() if dist.is_initialized() else 0
self.device = (
torch.device("cuda", torch.cuda.current_device())
if torch.cuda.is_available()
else torch.device("cpu")
)
@endpoint
def world_view(self) -> dict[str, object]:
import torch
import torch.distributed as dist
mine = {
"rank": self.rank,
"device": str(self.device),
"card": torch.cuda.get_device_name(self.device)
if torch.cuda.is_available()
else "cpu",
}
if self.world == 1:
return {"worldSize": 1, "ranks": [mine]}
gathered: list[object] = [None] * self.world
dist.all_gather_object(gathered, mine)
return {"worldSize": self.world, "ranks": gathered}
@endpoint
def reduce(self, size: int = DEFAULT_SIZE) -> dict[str, object]:
import torch
import torch.distributed as dist
size = DEFAULT_SIZE if size is None else size
start, end = _slice(size, self.rank, self.world)
# On this rank's own card, and only this rank's slice of the range. The
# whole point of the world is that no single card holds all of it.
mine = torch.arange(start, end, dtype=torch.float64, device=self.device)
partial = mine.sum()
if self.world == 1:
return {"total": partial.item(), "worldSize": 1, "size": size}
# Every rank reaches this, because every rank ran the same request. A rank
# that returned early would hang the ones that arrived here.
total = partial.clone()
dist.all_reduce(total)
return {
"total": total.item(),
"worldSize": self.world,
"size": size,
"partials": _partials(partial, self.world),
}
def _slice(size: int, rank: int, world: int) -> tuple[int, int]:
"""This rank's half-open share of ``range(size)``, the remainder spread early."""
share, remainder = divmod(size, world)
start = rank * share + min(rank, remainder)
return start, start + share + (1 if rank < remainder else 0)
def _partials(partial: object, world: int) -> list[object]:
import torch.distributed as dist
gathered: list[object] = [None] * world
dist.all_gather_object(gathered, partial.item()) # pyright: ignore[reportAttributeAccessIssue]
return gatheredWhat each rank runs
Only rank 0 is reachable from outside the sandbox, so it copies each invocation to its siblings and holds all of them at a barrier before any of them reaches your code.
That is what makes a collective inside your handler safe: by the time your first line runs, every rank is there with the same request.
It is also the constraint. A rank that returns early hangs every rank that arrived at the collective it skipped, so never put a shortcut in a handler, not for an empty share and not for an input one rank has nothing to do with.
The split here gives the remainder to the early ranks, and a rank whose share is empty still reaches the all-reduce.
The process group
load runs inside the platform's initialization, which sets your rank's device and calls init_process_group("nccl") before it, then waits on a barrier after.
A second init_process_group raises instead of being ignored, so your load only reads the world it was handed. For the same reason torch.cuda.current_device() is already your rank's card, so allocate with it and you keep off your siblings' memory.
Which rank answers
The siblings answer rank 0 and it discards them, and nothing breaks when the ranks disagree. If your response should carry every rank's work, gather it.
That is what partials is for. A worker that ignored the world would answer with one entry holding everything.
Deploying it
runware serverless deploy model.py \
--id multi-gpu \
--gpu-type rtx-pro-6000 \
--gpus-per-worker 2You set the count with --gpus-per-worker, which takes 1, 2, 4 or 8. Nothing in the file names one, so the same file serves a world of one.
A worker of N GPUs needs one node with N of them free, and not every type has one. A build knows nothing about where its worker will land, so this fails at rollout, not at build: the build stays ready and the app goes to failed.
runware serverless apps events <appId> says so in one line. Reach for it before reading your own code, because a worker that was never scheduled ran none of it.
Calling it
Call world-view first. It takes no payload and says whether the world came up, and on which cards.
import { createClient } from '@runware/sdk'
const client = await createClient({ apiKey: process.env.RUNWARE_API_KEY })
const task = await client.invoke({
appId: 'multi-gpu',
endpointPath: 'world-view',
payload: {}
})import asyncio
import os
from runware import Runware
async def main():
client = Runware(api_key=os.environ["RUNWARE_API_KEY"])
try:
task = await client.invoke(
"multi-gpu",
"world-view",
{},
)
finally:
await client.close()
asyncio.run(main())curl https://api.serverless.runware.ai/v1/apps/multi-gpu/invoke-async/world-view \
-H "Authorization: Bearer $RUNWARE_API_KEY" \
-H "Content-Type: application/json" \
-d '{
"taskId": "7c9e6679-7425-40de-944b-e07fc1f90ae7",
"payload": {}
}'echo '{}' | runware serverless apps invoke multi-gpu world-view -f -{
"taskId": "7c9e6679-7425-40de-944b-e07fc1f90ae7",
"payload": {}
}The endpoint path is the method name with underscores turned into hyphens, which is why world_view answers at world-view.
Then reduce, which sums a range with each rank holding its own slice of it.
import { createClient } from '@runware/sdk'
const client = await createClient({ apiKey: process.env.RUNWARE_API_KEY })
const task = await client.invoke({
appId: 'multi-gpu',
endpointPath: 'reduce',
payload: {
size: 4194304
}
})import asyncio
import os
from runware import Runware
async def main():
client = Runware(api_key=os.environ["RUNWARE_API_KEY"])
try:
task = await client.invoke(
"multi-gpu",
"reduce",
{
"size": 4194304
},
)
finally:
await client.close()
asyncio.run(main())curl https://api.serverless.runware.ai/v1/apps/multi-gpu/invoke-async/reduce \
-H "Authorization: Bearer $RUNWARE_API_KEY" \
-H "Content-Type: application/json" \
-d '{
"taskId": "7c9e6679-7425-40de-944b-e07fc1f90ae7",
"payload": {
"size": 4194304
}
}'echo '{"size":4194304}' | runware serverless apps invoke multi-gpu reduce -f -{
"taskId": "7c9e6679-7425-40de-944b-e07fc1f90ae7",
"payload": {
"size": 4194304
}
}Measured on two RTX PRO 6000:
{
"total": 8796090925056,
"partials": [2199022206976, 6597068718080],
"worldSize": 2,
"size": 4194304
}Two different partials, each summed on its own card, and a total that matches the same sum computed anywhere else.
The first call is slow because a worker carrying torch pulls a large image before it runs anything, which took about eight minutes from the first invoke to the result. runware serverless apps workers <appId> reports pulling while that happens, so you can see a long first call instead of guessing at it. The next call answers in 5 seconds.