This commit is contained in:
Xi Yan 2024-10-28 16:10:36 -07:00
parent 859f01c604
commit 2990303fd6

View file

@ -9,16 +9,12 @@ services:
devices: devices:
- nvidia.com/gpu=all - nvidia.com/gpu=all
environment: environment:
- CUDA_VISIBLE_DEVICES=0 - CUDA_VISIBLE_DEVICES=0,1,2,3,4
- HF_HOME=/data
- HF_DATASETS_CACHE=/data
- HF_MODULES_CACHE=/data
- HF_HUB_CACHE=/data
- NUM_SHARD=4 - NUM_SHARD=4
- MAX_BATCH_PREFILL_TOKENS=32768 - MAX_BATCH_PREFILL_TOKENS=32768
- MAX_INPUT_TOKENS=8000 - MAX_INPUT_TOKENS=8000
- MAX_TOTAL_TOKENS=8192 - MAX_TOTAL_TOKENS=8192
command: ["--dtype", "bfloat16", "--usage-stats", "on", "--sharded", "false", "--model-id", "meta-llama/Llama-3.1-8B-Instruct", "--port", "5009", "--cuda-memory-fraction", "0.3"] command: []
deploy: deploy:
resources: resources:
reservations: reservations:
@ -26,14 +22,13 @@ services:
- driver: nvidia - driver: nvidia
# that's the closest analogue to --gpus; provide # that's the closest analogue to --gpus; provide
# an integer amount of devices or 'all' # an integer amount of devices or 'all'
count: 1 count: all
# Devices are reserved using a list of capabilities, making # Devices are reserved using a list of capabilities, making
# capabilities the only required field. A device MUST # capabilities the only required field. A device MUST
# satisfy all the requested capabilities for a successful # satisfy all the requested capabilities for a successful
# reservation. # reservation.
capabilities: [gpu] capabilities: [gpu]
runtime: nvidia runtime: nvidia
shm_size: '1gb'
llamastack: llamastack:
depends_on: depends_on:
text-generation-inference: text-generation-inference: