23 lines
651 B
Bash
23 lines
651 B
Bash
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
world_size="${1:-${H3_WORLD_SIZE:-}}"
|
|
if [[ -z "$world_size" ]]; then
|
|
world_size="$(python -c 'import torch; print(torch.cuda.device_count())')"
|
|
fi
|
|
if [[ "$world_size" -lt 1 ]]; then
|
|
echo "No CUDA devices are visible. Launch the container with --gpus all." >&2
|
|
exit 1
|
|
fi
|
|
backend="${2:-sdpa}"
|
|
sequence="${3:-20000}"
|
|
output="${H3_BENCHMARK_OUTPUT:-/output/h3-baselines/ulysses-${world_size}gpu-${backend}-${sequence}t.json}"
|
|
|
|
torchrun \
|
|
--standalone \
|
|
--nnodes=1 \
|
|
--nproc-per-node="$world_size" \
|
|
tools/benchmark_ulysses.py \
|
|
--sequence "$sequence" \
|
|
--backend "$backend" \
|
|
--output "$output"
|