version: '3'
services:
  gpt-neox:
    command: nvidia-smi -q --loop=10
    image: leogao2/gpt-neox:main
    shm_size: 1g
    ulimits:
      memlock:
        soft: -1
        hard: -1
    runtime: nvidia
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              capabilities: [gpu]
    logging:
      options:
        max-size: "100m"
        max-file: "3"
    volumes:
      - ${NEOX_DATA_PATH}:/home/mchorse/data
      - ${NEOX_CHECKPOINT_PATH}:/home/mchorse/chk
      - .:/home/mchorse/gpt-neox