awsome-distributed-training/3.test_cases/pytorch/neuronx-distributed/llama3/kubernetes/llama3_train.yaml-template at cc7480152e06defb9ae9d51a5cf185a07d1376b5 · flamingofugang/awsome-distributed-training · 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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
apiVersion: v1
kind: Service
metadata:
  name: etcd
spec:
  ports:
    - name: etcd-client-port
      port: 2379
      protocol: TCP
      targetPort: 2379
  selector:
    app: etcd

---
apiVersion: apps/v1
kind: Deployment
metadata:
  labels:
    app: etcd
  name: etcd
spec:
  replicas: 1
  selector:
    matchLabels:
      app: etcd
  template:
    metadata:
      labels:
        app: etcd
    spec:
      containers:
        - name: etcd
          command: ["/usr/local/bin/etcd"]
          args:
            - "--data-dir"
            - "/var/lib/etcd"
            - "--enable-v2"
            - "--listen-client-urls"
            - "http://0.0.0.0:2379"
            - "--advertise-client-urls"
            - "http://0.0.0.0:2379"
            - "--initial-cluster-state"
            - "new"
          image: quay.io/coreos/etcd:v3.5.19
          ports:
            - containerPort: 2379
              name: client
              protocol: TCP
            - containerPort: 2380
              name: server
              protocol: TCP
      restartPolicy: Always
---
apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: trn1-llama3
spec:
  elasticPolicy:
    rdzvBackend: etcd
    rdzvHost: etcd
    rdzvPort: 2379
    minReplicas: 1
    maxReplicas: 64
    maxRestarts: 100
    metrics:
      - type: Resource
        resource:
          name: cpuyeah
          target:
            type: Utilization
            averageUtilization: 90
  pytorchReplicaSpecs:
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        metadata:
          labels:
            app: trn1-llama3
        spec:
          volumes:
            - name: shmem
              hostPath:
                path: /dev/shm
            - name: persistent-storage
              persistentVolumeClaim:
                claimName: ${FSX_CLAIM}
            - name: local
              hostPath:
                path: /dev
            - name: hyperpod
              hostPath:
                path: /var/log/aws/clusters
          nodeSelector:
           node.kubernetes.io/instance-type: ${INSTANCE_TYPE}
          containers:
            - name: pytorch
              image: ${IMAGE_URI}
              imagePullPolicy: Always
              resources:
                requests:
                  aws.amazon.com/neuron: ${NEURON_PER_NODE}
                  vpc.amazonaws.com/efa: ${EFA_PER_NODE}
                limits:
                  aws.amazon.com/neuron: ${NEURON_PER_NODE}
                  vpc.amazonaws.com/efa: ${EFA_PER_NODE}
              env:
              - name: LOGLEVEL
                value: "DEBUG"
              - name: FI_PROVIDER
                value: efa
              - name: FI_EFA_USE_DEVICE_RDMA
                value: "1"
              - name: FI_EFA_FORK_SAFE
                value: "1"
              - name: FI_LOG_LEVEL
                value: "1"
              - name: FI_EFA_ENABLE_SHM_TRANSFER
                value: "1"
              - name: NEURON_RT_NUM_CORES
                value: "32"
              - name: NUM_NEURONCORES
                value: "32"
              - name: TPU_NUM_DEVICES
                value: "32"
              - name: TPU_CHIPS_PER_HOST_BOUNDS
                value: "32"
              - name: TORCH_NCCL_DEBUG_INFO_TEMP_FILE
                value: "/local/nccl_trace_rank_"
              - name: PYTORCH_CUDA_ALLOC_CONF
                value: "expandable_segments:True"
              - name: MALLOC_ARENA_MAX
                value: "64"
              - name: NCCL_SOCKET_IFNAME
                value: "^lo"
              - name: NEURON_RT_ASYNC_EXEC_MAX_INFLIGHT_REQUESTS
                value: "3"
              - name: NEURON_FUSE_SOFTMAX
                value: "1"
              - name: NEURON_CC_FLAGS
                value: "--model-type transformer --distribution-strategy=llm-training --cache_dir=${NEURON_CACHE_DIR}"
              command:
                - torchrun
                - --nproc_per_node=32
                - --nnodes=$NUM_NODES
                - train.py
                - --model_path=${MODEL_PATH}
                - --data_dir=${TOKENIZED_DATA_PATH}/${DATASET_NAME}_llama3_tokenized_8k
                - --tensor_parallel_size=32
                - --batch_size=${BATCH_SIZE}
                - --steps_this_run=${STEPS_THIS_RUN}
                - --max_steps=${MAX_STEPS}
                - --warmup_steps=100
                - --lr=1.5e-4
                - --grad_accum_usteps=16
                - --seq_len=8192
                - --sequence_parallel_enabled
                - --selective_checkpoint_enabled
                - --logging_interval=10
                - --qkv_linear
                - --kv_replicator=4
                - --use_flash_attention=1
                - --use_zero_1
                - --use_mix_precision
                - --checkpoint_freq=${CHECKPOINT_FREQ}
                - --num_kept_checkpoint=${NUM_KEPT_CHECKPOINTS}
                - --checkpoint_dir=${CHECKPOINT_DIR}
              volumeMounts:
                - name: shmem
                  mountPath: /dev/shm
                - name: persistent-storage
                  mountPath: /fsx
                - name: hyperpod
                  mountPath: /var/log/aws/clusters