docker-compose multi-gpu example

[PeterTucker]

Can anyone give an example on how to get TabbyML working with two GPUs?

I have two 3060s. This is what I have so far, but I'm not sure if TabbyML is actually using both GPUs.

version: '3.5'

services:
  tabby:
    restart: always
    image: tabbyml/tabby:0.6.0
    command: serve --model TabbyML/StarCoder-3B  --chat-model TabbyML/WizardCoder-3B --device cuda
    #command: serve --model TabbyML/StarCoder-3B --device cuda
    volumes:
      - "./data:/data"
    ports:
      - 8080:8080
    environment:
      - CUDA_VISIBLE_DEVICES=0,1
      - NVIDIA_VISIBLE_DEVICES=all
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

I also get this error:

CUDA error 1 at /root/workspace/crates/llama-cpp-bindings/llama.cpp/ggml-cuda.cu:6693: invalid argument

[wsxiaoys]

Since we don't do tensor parallelism at the moment, the only way to utilize multiple gpu is separate the workload into different process / different cuda devices.

In this particular case, you might:

1. start serve --model TabbyML/StarCoder-3B --device cuda on two container process.
2. Each container process has visibility of single cuda device.

[PeterTucker]

@wsxiaoys Any chance I could get a docker-compose example? Not sure how that would work if both containers are using port 8080.
I'm also not sure how I would use multiple endpoints in VSCode.

[PeterTucker]

Just separating the models into different containers to use each GPU seems to work. I was hoping for a single model to use two GPUs, but that does not seem possible.

For anyone looking for a multi-gpu setup, here is the docker-compose setup:

version: '3.8'

services:
  tabby1:
    restart: always
    image: tabbyml/tabby:0.6.0
    command: serve --model TabbyML/StarCoder-7B --device cuda
    volumes:
      - ./data:/data
    ports:
      - 8080:8080
    environment:
      - CUDA_VISIBLE_DEVICES=0
      - NVIDIA_VISIBLE_DEVICES=0
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

  tabby2:
    restart: always
    image: tabbyml/tabby:0.6.0
    command: serve  --chat-model TabbyML/WizardCoder-3B --device cuda
    volumes:
      - ./data:/data
    ports:
      - 8081:8080
    environment:
      - CUDA_VISIBLE_DEVICES=1
      - NVIDIA_VISIBLE_DEVICES=1
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

[wsxiaoys]

@wsxiaoys Any chance I could get a docker-compose example? Not sure how that would work if both containers are using port 8080.

Putting a reverse proxy upfront is a natural choice in this case (e.g. Caddy), on the other hand, if you're interested in Tabby's built-in distributed worker support...

#  tabby webserver process
tabby serve

# Worker #1, completion model
tabby worker::completion --model TabbyML/StarCoder-1B --device cuda --port 8082 --url webserver-ip:8080 --token <registration_token>

# Worker #2, chat model
tabby worker::chat --model TabbyML/Mistral-7B --device cuda --port 8082 --url webserver-ip:8080 --token <registration_token>

[wsxiaoys]

A blog post added here: https://tabby.tabbyml.com/blog/2024/03/26/tabby-with-replicas-behind-reverse-proxy