values.yaml

targetPlatform: "gke"

volumes:
  # The VM host path for SSDs is assumed at /mnt/stateful_partition/kube-ephemeral-ssd
  ssdMountPath: "/ssd"

workload:
  # This should be the image built and pushed to the registry (EDIT THIS)
  image: "/address/to/image/built/and/uploaded/<image_name>:<image_tag>"

  torchDistributedTarget: "run_clm.py"

  # HuggingFace tokens (EDIT THIS)
  hfToken: "<hf_xxxxxxxxxxxxxxxxxxxxxxxxxxx>"

  # It will be mounted to /nfs on container startup using GCS fuse (EDIT THIS)
  gcsBucketForDataCataPath: <bucket_name>

  gpus: 16 # This should be one of: {<= 8,  multiple of 8}
  arguments:
  # The argument name will be prefixed with '+' (see https://hydra.cc/docs/advanced/override_grammar/basic/)
  - name: "exp_manager.explicit_log_dir"
    value: "/nemo-experiments/results" 
  - name: "exp_manager.exp_dir"
    value: "/nemo-experiments/"

  # Llama 2 tokenizer (not used)
  #- name: "model.data.data_prefix"
  #  value: "[1.0,/ssd/.cache/wikipedia-tokenized-for-llama2]"    
  #- name: "model.tokenizer.model"
  #  value: "/ssd/.cache/llama-2-7b-megatron-checkpoint/tokenizer.model"     

  # If not 'null', launches a Tensorboard server on first node. By design, the job will then not exit on first node.
  # This is primarly intended for debugging purposes, when a shared file-system or external Tensorboard is unavailable.  
  embeddedTensorboardTarget: null

network:
  stack: "tcpxo" # one of {"tcp", "tcpx", "tcpxo"}

  daemonVersion: "us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpgpudmarxd-dev:v1.0.9"
  pluginVersion: "us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/nccl-plugin-gpudirecttcpx-dev:v1.0.3"
    

  ncclSettings:
  - name: NCCL_DEBUG
    value: "VERSION"
  - name: NCCL_ALGO
    value: "Ring,Tree"

  # The following NCCL settings are recommended for TCPxo only (but tunable):
  - name: NCCL_MIN_NCHANNELS
    value: "4"