smpv2.yaml-template

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:latest
          ports:
            - containerPort: 2379
              name: client
              protocol: TCP
            - containerPort: 2380
              name: server
              protocol: TCP
      restartPolicy: Always
---
apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: ${JOB_NAME}
spec:
  elasticPolicy:
    rdzvBackend: etcd
    rdzvHost: etcd
    rdzvPort: 2379
    minReplicas: 1
    maxReplicas: 64
    maxRestarts: 100
    metrics:
      - type: Resource
        resource:
          name: cpu
          target:
            type: Utilization
            averageUtilization: 90
  pytorchReplicaSpecs:
    Worker:
      replicas: $NUM_NODES
      restartPolicy: OnFailure
      template:
        metadata:
          labels:
            app: ${JOB_NAME}
        spec:
          volumes:
            - name: shmem
              hostPath: 
                path: /dev/shm
            - name: local
              hostPath:
                path: /mnt/k8s-disks/0
            - 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:
                  nvidia.com/gpu: $GPU_PER_NODE
                  vpc.amazonaws.com/efa: $EFA_PER_NODE
                limits:
                  nvidia.com/gpu: ${GPU_PER_NODE}
                  vpc.amazonaws.com/efa: ${EFA_PER_NODE}
              env:
              # for P5 FI_* should be commented out
              - name: LOGLEVEL
                value: "DEBUG"
              #- name: FI_PROVIDER
              #  value: $FI_PROVIDER
              #- 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: TORCH_DISTRIBUTED_DEBUG
                value: "DETAIL"
              - name: TORCH_NCCL_ENABLE_MONITORING
                value: "1"
              - name: TORCH_NCCL_TRACE_BUFFER_SIZE
                value: "20000"
              - name: TORCH_NCCL_DUMP_ON_TIMEOUT
                value: "1"
              - name: TORCH_NCCL_DEBUG_INFO_TEMP_FILE
                value: "/local/nccl_trace_rank_"
              - name: PYTORCH_CUDA_ALLOC_CONF
                value: "expandable_segments:True"
              - name: NCCL_DEBUG
                value: "INFO"
              - name: NCCL_SOCKET_IFNAME
                value: "^lo"
              - name: TORCH_NCCL_ASYNC_ERROR_HANDLING
                value: "1"
              #- name: TORCH_DIST_INIT_BARRIER
              #  value: "1"
              #- name: NCCL_IGNORE_DISABLED_P2P
              #  value: "1"
              #- name: NCCL_NVLS_ENABLE
              #  value: "0"
              command: 
                - torchrun
                - --nproc_per_node=$GPU_PER_NODE
                - --nnodes=$NUM_NODES
                - /workspace/train_external.py
                - --train_batch_size=${TRAIN_BATCH_SIZE}
                - --max_steps 100
                - --hidden_width=${HIDDEN_WIDTH}
                - --num_layers=${NUM_LAYERS}
                - --num_heads=${NUM_HEADS}
                - --llama_intermediate_size=${LLAMA_INTERMEDIATE_SIZE}
                - --shard_degree=${SHARD_DEGREE}
                - --model_type=llama_v2
                - --profile_nsys=1
                - --use_smp_implementation=1
                - --max_context_width=4096
                - --tensor_parallel_degree=1
                - --use_synthetic_data=${USE_SYNTHETIC_DATA}
                - --training_dir=${TRAINING_DIR}
                - --test_dir=${TEST_DIR}
                - --dataset_type=hf
                - --checkpoint_dir=${CHECKPOINT_DIR}
                - --checkpoint_freq=100
                - --intermediate_size=11008
                - --fp8=${USE_FP8}
              volumeMounts:
                - name: shmem
                  mountPath: /dev/shm
                - name: local
                  mountPath: /local
                - name: hyperpod
                  mountPath: /var/log/aws/clusters