ptl_ddp_multi_gpu.yml

name: ptl-mct-multi-gpu
image: mosaicml/composer:latest
compute:
 gpus: 8 # Number of GPUs to use
 cluster: r14z3p1 # Insert cluster name
 gpu_type: h100_80gb # Type of GPU to use. We use a100_80gb in our experiments
integrations:
 - integration_type: "git_repo"
   git_repo: SamRaymond/composer_test
   git_branch: main
command: |-
  cd composer_test/
  mkdir -p /tmp/dataset/
  mkdir -p /tmp/dataset/train/
  mkdir -p /tmp/dataset/test/
  chmod -R 777 /tmp/dataset/test/ /tmp/dataset/train/
  pip install -Uq torchvision transformers lightning
      torchrun \
    --nproc_per_node=8 \
    --nnodes=1 \
    ptl_ddp_multi_gpu.py || (echo "Command failed - killing python" && pkill python && exit 1)
  # devices = 8 
  # Epoch 9: 100%|██████████| 196/196 [00:02<00:00, 87.67it/s, v_num=0]