20 lines
633 B
Bash
20 lines
633 B
Bash
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
available="$(python -c 'import torch; print(torch.cuda.device_count())')"
|
|
if [[ "$available" -lt 1 ]]; then
|
|
echo "No CUDA devices are visible. Launch the container with --gpus all." >&2
|
|
exit 1
|
|
fi
|
|
|
|
read -ra counts <<< "${H3_GPU_COUNTS:-1 2 4 6 8}"
|
|
read -ra modes <<< "${H3_DISTRIBUTED_MODES:-ulysses tensor}"
|
|
for count in "${counts[@]}"; do
|
|
if [[ "$count" -gt "$available" ]]; then
|
|
echo "Skipping ${count} GPUs; only ${available} are visible." >&2
|
|
continue
|
|
fi
|
|
for mode in "${modes[@]}"; do
|
|
tools/run_distributed_t2va.sh "$mode" "$count" "${H3_ATTENTION:-sdpa}"
|
|
done
|
|
done
|