final_run.yaml

apiVersion: v1
kind: Pod
metadata:
  generateName: resnet-custom-backprop-
spec:
  #nodeName: gpu8-vm13 #9
  containers: 
  - args:
    - |
      torchrun --nproc_per_node=2 transformer_main.py --algo 1f1b-1 --msbp ;
      torchrun --nproc_per_node=4 transformer_main.py --algo 1f1b-1 --msbp ;
      torchrun --nproc_per_node=8 transformer_main.py --algo 1f1b-1 --msbp ;
    command: [/bin/bash, -c, --]
    image: bigballoon8/custom-backprop
    name: custom-backprop
    env:
    - name: OMP_NUM_THREADS
      value: "4"
    #- name: NCCL_DEBUG
    #  value: "INFO"
    resources:
      limits:
        cpu: 32
        memory: 64Gi
        nvidia.com/gpu: 8
    volumeMounts:
    - mountPath: /dev/shm
      name: devshm
    workingDir: /workspace/experiment3(pipeline_parallelism)/
  nodeSelector:
    nvidia.com/gpu.product: NVIDIA-A100-SXM4-40GB
  restartPolicy: Never
  volumes:
  - emptyDir:
      medium: Memory
    name: devshm