File size: 1,749 Bytes
b0734ca e5c09aa b0734ca e5c09aa b0734ca e5c09aa b0734ca e5c09aa b0734ca e5c09aa b0734ca e5c09aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | # ROMA inference on 4x Quadro RTX 6000 (Turing, x86_64).
# Usage (run from the repo root on the workstation):
# docker compose -f docker/docker-rtx6000/docker-compose.yml build
# docker compose -f docker/docker-rtx6000/docker-compose.yml run --rm --service-ports roma bash
services:
roma:
build:
dockerfile: ./docker/docker-rtx6000/Dockerfile
context: ../..
args:
BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
TRANSFORMERS_REF: roma_patch
PIP_INDEX: https://pypi.org/simple
TORCH_INDEX: https://download.pytorch.org/whl/cu124
image: roma-rtx6000:latest
container_name: roma-rtx6000
volumes:
# ~16-22GB checkpoint, HF cache and demo media live on the host, not in the image.
- ../../whole_model:/app/whole_model
- ../../hf_cache:/root/.cache/huggingface
- ../../demo_media:/app/demo_media
- ../../output:/app/output
ports:
- "7860:7860"
- "8000:8000"
environment:
- GRADIO_SERVER_NAME=0.0.0.0
- HF_HUB_ENABLE_HF_TRANSFER=1
# fp16 sharded across all 4 GPUs (no quantization). See README-RTX6000.md for fallbacks.
- ROMA_DTYPE=float16
- ROMA_ATTN=sdpa
- ROMA_LOAD_8BIT=0
# NOTE: intentionally NOT setting CUDA_VISIBLE_DEVICES, so device_map="auto" shards across all 4 GPUs.
# - HF_TOKEN=${HF_TOKEN} # uncomment if the checkpoint pull is rate-limited/gated
ipc: host
shm_size: "16gb"
tty: true
stdin_open: true
command: bash
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: all # all 4 Quadro RTX 6000 cards
capabilities: [gpu]
restart: unless-stopped
|