DistributedTraining/docker-compose.yml at main · abeenoch/DistributedTraining · 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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
version: "3.9"

services:
  parameter_server:
    build:
      context: .
      dockerfile: docker/Dockerfile
    container_name: parameter_server
    environment:
      SERVICE_ROLE: parameter_server
      PS_HOST: 0.0.0.0
      PS_PORT: 50051
      NUM_WORKERS: 2
      AGGREGATION_MODE: sync
      PS_LEARNING_RATE: 0.01
      MODEL_INPUT_DIM: 2
      MODEL_OUTPUT_DIM: 2
    ports:
      - "50051:50051"
    volumes:
      - ./checkpoints:/app/checkpoints
      - ./logs:/app/logs
    networks:
      - training_net

  worker_0:
    build:
      context: .
      dockerfile: docker/Dockerfile
    container_name: worker_0
    depends_on:
      - parameter_server
    environment:
      SERVICE_ROLE: worker
      WORKER_ID: worker_0
      WORKER_RANK: 0
      WORLD_SIZE: 2
      PARAMETER_SERVER_ADDRESS: parameter_server:50051
      BATCH_SIZE: 8
      NUM_EPOCHS: 1
      LEARNING_RATE: 0.01
      MODEL_INPUT_DIM: 2
      MODEL_OUTPUT_DIM: 2
      SYNTHETIC_DATASET_SIZE: 64
      COMPRESSION_ENABLED: "true"
      COMPRESSION_TYPE: quantization
      COMPRESSION_RATIO: 0.1
    volumes:
      - ./checkpoints:/app/checkpoints
      - ./logs:/app/logs
    networks:
      - training_net

  worker_1:
    build:
      context: .
      dockerfile: docker/Dockerfile
    container_name: worker_1
    depends_on:
      - parameter_server
    environment:
      SERVICE_ROLE: worker
      WORKER_ID: worker_1
      WORKER_RANK: 1
      WORLD_SIZE: 2
      PARAMETER_SERVER_ADDRESS: parameter_server:50051
      BATCH_SIZE: 8
      NUM_EPOCHS: 1
      LEARNING_RATE: 0.01
      MODEL_INPUT_DIM: 2
      MODEL_OUTPUT_DIM: 2
      SYNTHETIC_DATASET_SIZE: 64
      COMPRESSION_ENABLED: "true"
      COMPRESSION_TYPE: quantization
      COMPRESSION_RATIO: 0.1
    volumes:
      - ./checkpoints:/app/checkpoints
      - ./logs:/app/logs
    networks:
      - training_net

networks:
  training_net:
    driver: bridge