| # 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 | |