valohai.yaml

- step:
    name: train-torchrun
    image: sofiiavalohai/valohai-llms:v1.5
    environment: aws-eu-west-1-p3-8xlarge
    command:
      - pip install valohai-utils
      - torchrun --nproc-per-node=4 train-torchrun.py {parameters}
    parameters:
      - name: model-ckpt
        default: "facebook/bart-large-cnn"
      - name: output-dir
        default: "bart-large-cnn"
      - name: batch-size
        default: 1
      - name: num-epochs
        default: 1
      - name: warmup-steps
        default: 500
      - name: evaluation-steps
        default: 500
    inputs:
      - name: dataset
        default:
        - https://valohai-demo-library-data.s3.eu-west-1.amazonaws.com/distributed-llms/train.json
        - https://valohai-demo-library-data.s3.eu-west-1.amazonaws.com/distributed-llms/val.json

- step:
    name: train-accelerator
    image: sofiiavalohai/valohai-llms:v1.5
    environment: aws-eu-west-1-p3-8xlarge
    command:
      - pip install accelerate evaluate absl-py rouge-score
      - python -m accelerate.commands.launch --num_processes=4 train-accelerator.py {parameters}
    parameters:
      - name: model-ckpt
        default: "facebook/bart-large-cnn"
      - name: output-dir
        default: "bart-large-cnn"
      - name: batch-size
        default: 1
      - name: num-epochs
        default: 1
      - name: warmup-steps
        default: 500
      - name: evaluation-steps
        default: 500
    inputs:
      - name: dataset
        default:
        - https://valohai-demo-library-data.s3.eu-west-1.amazonaws.com/distributed-llms/train.json
        - https://valohai-demo-library-data.s3.eu-west-1.amazonaws.com/distributed-llms/val.json

- step:
    name: train-task
    image: sofiiavalohai/valohai-llms:v1.5
    environment: sharkdemo-aws-eu-central-1-g4dn-xlarge
    command:
      - pip install accelerate evaluate absl-py rouge-score
      - python train-task.py {parameters}
    environment-variables:
      - name: VH_DOCKER_NETWORK
        default: host
    parameters:
      - name: model-ckpt
        default: "facebook/bart-large-cnn"
      - name: output-dir
        default: "bart-large-cnn"
      - name: batch-size
        default: 6
      - name: num-epochs
        default: 1
      - name: warmup-steps
        default: 500
      - name: evaluation-steps
        default: 500
    inputs:
      - name: dataset
        default:
        - https://valohai-demo-library-data.s3.eu-west-1.amazonaws.com/distributed-llms/train.json
        - https://valohai-demo-library-data.s3.eu-west-1.amazonaws.com/distributed-llms/val.json

- task:
    step: train-task
    name: task-distributed
    type: distributed
    parameters: []
    execution-count: 4