ForgeLM/docker-compose.yaml at main · HodeTech/ForgeLM · GitHub

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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
# ==============================================================================
# ForgeLM Docker Compose
#
# Place your config file(s) under ./configs/ on the host — the `train`
# service mounts that whole directory read-only at /workspace/configs
# (see the `volumes:` block below), so --config must point inside it.
#
# Usage:
#   Training:
#     docker compose run --rm train --config /workspace/configs/my_config.yaml
#
#   Dry-run:
#     docker compose run --rm train --config /workspace/configs/my_config.yaml --dry-run
#
#   Benchmark only:
#     docker compose run --rm train --config /workspace/configs/my_config.yaml \
#       --benchmark-only /workspace/output/final_model
#
#   Wizard (generate config):
#     docker compose run --rm train --wizard
#
#   TensorBoard:
#     docker compose up tensorboard
#     Open http://localhost:6006
# ==============================================================================

services:
  train:
    build:
      context: .
      args:
        INSTALL_EVAL: "true"
    volumes:
      # Mount your config, data, and output directories
      - ./configs:/workspace/configs:ro
      - ./data:/workspace/data:ro
      - ./output:/workspace/output
      # Optional: mount HF cache to avoid re-downloading models
      - hf_cache:/root/.cache/huggingface
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    # Override entrypoint args via: docker compose run train --config ...
    entrypoint: ["forgelm"]

  tensorboard:
    # Pinned (not :latest) for reproducibility, matching the project's other
    # supply-chain pins (publish.yml's pinned-by-SHA action, pyproject.toml's
    # bounded dependency ranges). 2.15.0 tracks the `tensorboard>=2.15.0,<3.0.0`
    # floor declared in pyproject.toml's [project].dependencies — bump both
    # together when the floor moves.
    image: tensorflow/tensorflow:2.15.0
    ports:
      - "6006:6006"
    volumes:
      - ./output:/workspace/output:ro
    command: ["tensorboard", "--logdir=/workspace/output", "--host=0.0.0.0"]

volumes:
  hf_cache: