Multi-node Inference

Certain AI models require a specialized type of deployment which is multi-node inference.

Lightning AI automatically sets all necessary environment variables to properly configure multi-node inference and ensures seamless collaboration between nodes.

Here are the environment variables:

Name Description
MASTER_ADDR The address of the main node.
MASTER_PORT The port of the main node.
WORLD_SIZE The total number of available devices.
NODE_RANK The rank of the current node.
NUM_NODES The total number of available nodes.

Example: SGLang multi-node inference

This example demonstrates how the container’s bash start script passes environment variables to sglang to enable multi-node inference for a given model.

Dockerfile

# Dockerfile
FROM lmsysorg/sglang:latest

RUN pip install litmodels

WORKDIR /app
COPY . .

ENTRYPOINT ["/app/start.sh"]

Bash start script

# /app/start.sh

#!/bin/bash

LIGHTNING_CLOUD_URL=https://lightning.ai lightning model download "${MODEL_NAME}" --download_dir=${MODEL_DIR}

echo "python3 -m sglang.launch_server --dist-init-addr ${MASTER_ADDR}:5000 --nnodes ${NUM_NODES} --node-rank ${NODE_RANK} --trust-remote-code --model-path ${MODEL_DIR} $@"
exec python3 -m sglang.launch_server --dist-init-addr ${MASTER_ADDR}:5000 --nnodes ${NUM_NODES} --node-rank ${NODE_RANK} --trust-remote-code --model-path ${MODEL_DIR} "$@"

Was this page helpful?

What went wrong?

↑↓ Navigate↵ Open⌘↵ New tabesc Close