Skip to content

Commit af224d6

Browse files
adibarraadibarra
authored
Add context parallel config options (#2110)
1 parent da54b42 commit af224d6

31 files changed

Lines changed: 368 additions & 59 deletions

‎.github/workflows/benchmark-tmpl.yml‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,15 @@ on:
3232
tp:
3333
required: true
3434
type: string
35+
dcp-size:
36+
required: false
37+
type: string
38+
default: '1'
39+
pcp-size:
40+
required: false
41+
type: string
42+
default: '1'
43+
3544
ep:
3645
required: true
3746
type: string
@@ -105,6 +114,8 @@ env:
105114
FRAMEWORK: ${{ inputs.framework }}
106115
PRECISION: ${{ inputs.precision }}
107116
TP: ${{ inputs.tp }}
117+
DCP_SIZE: ${{ inputs.dcp-size }}
118+
PCP_SIZE: ${{ inputs.pcp-size }}
108119
EP_SIZE: ${{ inputs.ep }}
109120
DP_ATTENTION: ${{ inputs.dp-attn }}
110121
CONC: ${{ inputs.conc }}
@@ -134,7 +145,7 @@ jobs:
134145
timeout-minutes: 500
135146
name: >-
136147
${{ inputs.model-prefix }} ${{ inputs.precision }} ${{ inputs.runner }} ${{ inputs.framework == 'sglang' && 'sgl' || inputs.framework == 'dynamo-sglang' && 'dyn-sgl' || inputs.framework == 'sglang-disagg' && 'sgl-disagg' || inputs.framework }}
137-
TP${{ inputs.tp }}${{ inputs.ep != '' && inputs.ep != '1' && format('/EP{0}', inputs.ep) || '' }}${{ inputs.dp-attn && '/DPA' || '' }}
148+
TP${{ inputs.tp }}/DCP${{ inputs.dcp-size }}/PCP${{ inputs.pcp-size }}${{ inputs.ep != '' && inputs.ep != '1' && format('/EP{0}', inputs.ep) || '' }}${{ inputs.dp-attn && '/DPA' || '' }}
138149
${{ inputs.spec-decoding != 'none' && inputs.spec-decoding || '' }}
139150
${{ inputs.kv-offloading != '' && inputs.kv-offloading != 'none' && format('{0} KV offload', inputs.kv-offloading) || '' }}
140151
${{ inputs.kv-offload-backend != '' && inputs.kv-offload-backend != 'none' && inputs.kv-offload-backend != 'default' && inputs.kv-offload-backend || '' }}
@@ -175,12 +186,15 @@ jobs:
175186
env:
176187
RUNNER_NAME: ${{ runner.name }}
177188
RUNNER_TYPE: ${{ inputs.runner }}
178-
# Hash uniquely on {EXP_NAME}_{PRECISION}_{FRAMEWORK}_tp{}-ep{}-dpa{}_disagg-{}_spec-{}_conc{}_{runner}
179-
RESULT_FILENAME: ${{ env.EXP_NAME }}_${{ env.PRECISION }}_${{ env.FRAMEWORK }}_tp${{ env.TP }}-ep${{ env.EP_SIZE }}-dpa${{ env.DP_ATTENTION }}_disagg-${{ env.DISAGG }}_spec-${{ env.SPEC_DECODING }}_conc${{ env.CONC }}_${{ runner.name }}
189+
# Hash uniquely on {EXP_NAME}_{PRECISION}_{FRAMEWORK}_tp{}-dcp{}-pcp{}-ep{}-dpa{}_disagg-{}_spec-{}_conc{}_{runner}
190+
RESULT_FILENAME: ${{ env.EXP_NAME }}_${{ env.PRECISION }}_${{ env.FRAMEWORK }}_tp${{ env.TP }}-dcp${{ env.DCP_SIZE }}-pcp${{ env.PCP_SIZE }}-ep${{ env.EP_SIZE }}-dpa${{ env.DP_ATTENTION }}_disagg-${{ env.DISAGG }}_spec-${{ env.SPEC_DECODING }}_conc${{ env.CONC }}_${{ runner.name }}
180191
# Suppress per-job eval markdown from being appended to the step summary.
181192
# We'll publish a single combined eval table in the collection job instead.
182193
GITHUB_STEP_SUMMARY: ''
183194
run: |
195+
export GPU_COUNT=$((TP * PCP_SIZE))
196+
echo "GPU_COUNT=${GPU_COUNT}" >> "$GITHUB_ENV"
197+
184198
# Export RESULT_FILENAME early so it's available for artifact uploads even if cancelled
185199
echo "RESULT_FILENAME=${RESULT_FILENAME}" >> $GITHUB_ENV
186200

‎.github/workflows/e2e-tests.yml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,8 @@ jobs:
185185
framework: ${{ matrix.config.framework }}
186186
precision: ${{ matrix.config.precision }}
187187
tp: ${{ matrix.config.tp }}
188+
dcp-size: ${{ matrix.config.dcp-size }}
189+
pcp-size: ${{ matrix.config.pcp-size }}
188190
ep: ${{ matrix.config.ep }}
189191
dp-attn: ${{ matrix.config.dp-attn }}
190192
conc: ${{ matrix.config.conc }}
@@ -267,6 +269,8 @@ jobs:
267269
framework: ${{ matrix.config.framework }}
268270
precision: ${{ matrix.config.precision }}
269271
tp: ${{ matrix.config.tp }}
272+
dcp-size: ${{ matrix.config.dcp-size }}
273+
pcp-size: ${{ matrix.config.pcp-size }}
270274
ep: ${{ matrix.config.ep }}
271275
dp-attn: ${{ matrix.config.dp-attn }}
272276
conc: ${{ matrix.config.conc }}
@@ -297,6 +301,8 @@ jobs:
297301
framework: ${{ matrix.config.framework }}
298302
precision: ${{ matrix.config.precision }}
299303
tp: ${{ matrix.config.tp }}
304+
dcp-size: ${{ matrix.config.dcp-size }}
305+
pcp-size: ${{ matrix.config.pcp-size }}
300306
ep: ${{ matrix.config.ep }}
301307
dp-attn: ${{ matrix.config.dp-attn }}
302308
conc: ${{ matrix.config.conc }}

‎.github/workflows/profile.yml‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,9 @@ jobs:
9898
fail-fast: false
9999
matrix:
100100
config: ${{ fromJson(needs.get-jobs.outputs.filtered-matrix) }}
101+
name: >-
102+
${{ matrix.config.model-prefix }} ${{ matrix.config.precision }} ${{ matrix.config.runner }} ${{ matrix.config.framework }}
103+
TP${{ matrix.config.tp }}/DCP${{ matrix.config.dcp-size }}/PCP${{ matrix.config.pcp-size }} c${{ matrix.config.conc }}
101104
runs-on: ${{ matrix.config.runner }}
102105
env:
103106
EXP_NAME: ${{ matrix.config.exp-name }}
@@ -110,6 +113,8 @@ jobs:
110113
FRAMEWORK: ${{ matrix.config.framework }}
111114
PRECISION: ${{ matrix.config.precision }}
112115
TP: ${{ matrix.config.tp }}
116+
DCP_SIZE: ${{ matrix.config.dcp-size }}
117+
PCP_SIZE: ${{ matrix.config.pcp-size }}
113118
EP_SIZE: ${{ matrix.config.ep }}
114119
DP_ATTENTION: ${{ matrix.config['dp-attn'] }}
115120
CONC: ${{ matrix.config.conc }}
@@ -159,8 +164,10 @@ jobs:
159164
shell: bash
160165
run: |
161166
set -euo pipefail
167+
export GPU_COUNT=$((TP * PCP_SIZE))
168+
echo "GPU_COUNT=${GPU_COUNT}" >> "$GITHUB_ENV"
162169
ep_val="${EP_SIZE:-1}"
163-
res_name="${EXP_NAME}_${PRECISION}_${FRAMEWORK}_tp${TP}_ep${ep_val}_dpa_${DP_ATTENTION}_conc${CONC}_${RUNNER_NAME}"
170+
res_name="${EXP_NAME}_${PRECISION}_${FRAMEWORK}_tp${TP}-dcp${DCP_SIZE}-pcp${PCP_SIZE}_ep${ep_val}_dpa_${DP_ATTENTION}_conc${CONC}_${RUNNER_NAME}"
164171
export RESULT_FILENAME="${res_name}"
165172
echo "RESULT_FILENAME=${res_name}" >> "$GITHUB_ENV"
166173
@@ -252,21 +259,21 @@ jobs:
252259
run: |
253260
set -euo pipefail
254261
255-
dest_dir="storage/profiles/${GITHUB_SHA}/${{ matrix.config.runner }}/${{ matrix.config.framework }}/${{ matrix.config['exp-name'] }}_${{ matrix.config.precision }}_tp${{ matrix.config.tp }}_ep${{ matrix.config.ep || 1 }}_conc${{ matrix.config.conc }}"
262+
dest_dir="storage/profiles/${GITHUB_SHA}/${{ matrix.config.runner }}/${{ matrix.config.framework }}/${{ matrix.config['exp-name'] }}_${{ matrix.config.precision }}_tp${{ matrix.config.tp }}_dcp${{ matrix.config.dcp-size }}_pcp${{ matrix.config.pcp-size }}_ep${{ matrix.config.ep || 1 }}_conc${{ matrix.config.conc }}"
256263
mkdir -p "$dest_dir"
257264
cp "$TRACE_LOCAL" "$dest_dir/trace.json.gz"
258265
259266
pushd storage >/dev/null
260267
git config user.name "github-actions"
261268
git config user.email "github-actions@github.com"
262269
git add -A
263-
git commit -m "Add profile: ${GITHUB_SHA} ${{ matrix.config['exp-name'] }} tp${{ matrix.config.tp }} ep${{ matrix.config.ep || 1 }} conc${{ matrix.config.conc }}" || echo "Nothing to commit"
270+
git commit -m "Add profile: ${GITHUB_SHA} ${{ matrix.config['exp-name'] }} tp${{ matrix.config.tp }} dcp${{ matrix.config.dcp-size }} pcp${{ matrix.config.pcp-size }} ep${{ matrix.config.ep || 1 }} conc${{ matrix.config.conc }}" || echo "Nothing to commit"
264271
git push
265272
STORAGE_SHA="$(git rev-parse HEAD)"
266273
popd >/dev/null
267274
268-
export RAW_URL="https://raw.githubusercontent.com/SemiAnalysisAI/InferenceX-trace-storage/${STORAGE_SHA}/profiles/${GITHUB_SHA}/${{ matrix.config.runner }}/${{ matrix.config.framework }}/${{ matrix.config['exp-name'] }}_${{ matrix.config.precision }}_tp${{ matrix.config.tp }}_ep${{ matrix.config.ep || 1 }}_conc${{ matrix.config.conc }}/trace.json.gz"
269-
export TITLE="${{ matrix.config['exp-name'] }}_${{ matrix.config.precision }}_tp${{ matrix.config.tp }}_ep${{ matrix.config.ep || 1 }}_conc${{ matrix.config.conc }}"
275+
export RAW_URL="https://raw.githubusercontent.com/SemiAnalysisAI/InferenceX-trace-storage/${STORAGE_SHA}/profiles/${GITHUB_SHA}/${{ matrix.config.runner }}/${{ matrix.config.framework }}/${{ matrix.config['exp-name'] }}_${{ matrix.config.precision }}_tp${{ matrix.config.tp }}_dcp${{ matrix.config.dcp-size }}_pcp${{ matrix.config.pcp-size }}_ep${{ matrix.config.ep || 1 }}_conc${{ matrix.config.conc }}/trace.json.gz"
276+
export TITLE="${{ matrix.config['exp-name'] }}_${{ matrix.config.precision }}_tp${{ matrix.config.tp }}_dcp${{ matrix.config.dcp-size }}_pcp${{ matrix.config.pcp-size }}_ep${{ matrix.config.ep || 1 }}_conc${{ matrix.config.conc }}"
270277
271278
enc_src="$(python3 -c 'import os,urllib.parse; print(urllib.parse.quote(os.environ["RAW_URL"], safe=""))')"
272279
enc_title="$(python3 -c 'import os,urllib.parse; print(urllib.parse.quote(os.environ["TITLE"], safe=""))')"

‎.github/workflows/run-sweep.yml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,8 @@ jobs:
331331
framework: ${{ matrix.config.framework }}
332332
precision: ${{ matrix.config.precision }}
333333
tp: ${{ matrix.config.tp }}
334+
dcp-size: ${{ matrix.config.dcp-size }}
335+
pcp-size: ${{ matrix.config.pcp-size }}
334336
ep: ${{ matrix.config.ep }}
335337
dp-attn: ${{ matrix.config.dp-attn }}
336338
conc: ${{ matrix.config.conc }}
@@ -434,6 +436,8 @@ jobs:
434436
framework: ${{ matrix.config.framework }}
435437
precision: ${{ matrix.config.precision }}
436438
tp: ${{ matrix.config.tp }}
439+
dcp-size: ${{ matrix.config.dcp-size }}
440+
pcp-size: ${{ matrix.config.pcp-size }}
437441
ep: ${{ matrix.config.ep }}
438442
dp-attn: ${{ matrix.config.dp-attn }}
439443
conc: ${{ matrix.config.conc }}
@@ -487,6 +491,8 @@ jobs:
487491
framework: ${{ matrix.config.framework }}
488492
precision: ${{ matrix.config.precision }}
489493
tp: ${{ matrix.config.tp }}
494+
dcp-size: ${{ matrix.config.dcp-size }}
495+
pcp-size: ${{ matrix.config.pcp-size }}
490496
ep: ${{ matrix.config.ep }}
491497
dp-attn: ${{ matrix.config.dp-attn }}
492498
conc: ${{ matrix.config.conc }}
@@ -582,6 +588,8 @@ jobs:
582588
framework: ${{ matrix.config.framework }}
583589
precision: ${{ matrix.config.precision }}
584590
tp: ${{ matrix.config.tp }}
591+
dcp-size: ${{ matrix.config.dcp-size }}
592+
pcp-size: ${{ matrix.config.pcp-size }}
585593
ep: ${{ matrix.config.ep }}
586594
dp-attn: ${{ matrix.config.dp-attn }}
587595
conc: ${{ matrix.config.conc }}

‎benchmarks/benchmark_lib.sh‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1074,6 +1074,8 @@ append_lm_eval_summary() {
10741074
"precision": "${prec:-unknown}",
10751075
"spec_decoding": "${SPEC_DECODING}",
10761076
"tp": ${TP:-1},
1077+
"dcp_size": ${DCP_SIZE:-1},
1078+
"pcp_size": ${PCP_SIZE:-1},
10771079
"conc": ${metadata_conc},
10781080
${batch_metadata} "ep": ${EP_SIZE:-1},
10791081
"dp_attention": ${dp_json},

‎benchmarks/single_node/agentic/dsv4_fp4_b200_vllm.sh‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,23 @@ source "$(dirname "$0")/../../benchmark_lib.sh"
2727

2828
check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE DP_ATTENTION
2929

30+
DCP_SIZE="${DCP_SIZE:-1}"
31+
PCP_SIZE="${PCP_SIZE:-1}"
32+
VLLM_CP_ARGS=()
33+
if [ "$DCP_SIZE" -gt 1 ]; then
34+
VLLM_CP_ARGS+=(--decode-context-parallel-size "$DCP_SIZE")
35+
fi
36+
if [ "$PCP_SIZE" -gt 1 ]; then
37+
VLLM_CP_ARGS+=(--prefill-context-parallel-size "$PCP_SIZE")
38+
fi
39+
40+
GPU_COUNT="${GPU_COUNT:-$((TP * PCP_SIZE))}"
41+
if [[ ! "$GPU_COUNT" =~ ^[1-9][0-9]*$ ]]; then
42+
echo "Error: GPU_COUNT must be a positive integer, got '$GPU_COUNT'" >&2
43+
exit 1
44+
fi
45+
export GPU_COUNT
46+
3047
if [[ -n "${SLURM_JOB_ID:-}" ]]; then
3148
echo "JOB $SLURM_JOB_ID running on ${SLURMD_NODENAME:-unknown}"
3249
fi
@@ -95,7 +112,7 @@ OFFLOAD_ARGS=()
95112
if require_agentic_kv_offload_backend mooncake; then
96113
# Embedded mode contributes one segment per GPU rank to a shared
97114
# distributed store, so pre-divide the aggregate host-memory budget.
98-
PER_RANK_GB=$((TOTAL_CPU_DRAM_GB / TP))
115+
PER_RANK_GB=$((TOTAL_CPU_DRAM_GB / GPU_COUNT))
99116

100117
MOONCAKE_VERSION=0.3.11.post1
101118
agentic_pip_install --quiet --no-cache-dir --no-deps \
@@ -183,6 +200,7 @@ VLLM_CMD=(
183200
--kv-cache-dtype fp8
184201
--block-size 256
185202
"${PARALLEL_ARGS[@]}"
203+
"${VLLM_CP_ARGS[@]}"
186204
"${EP_ARGS[@]}"
187205
--compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}'
188206
--attention_config.use_fp4_indexer_cache=True

‎benchmarks/single_node/agentic/dsv4_fp4_b300_vllm.sh‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,23 @@ source "$(dirname "$0")/../../benchmark_lib.sh"
2626

2727
check_env_vars MODEL TP CONC KV_OFFLOADING TOTAL_CPU_DRAM_GB RESULT_DIR DURATION EP_SIZE DP_ATTENTION
2828

29+
DCP_SIZE="${DCP_SIZE:-1}"
30+
PCP_SIZE="${PCP_SIZE:-1}"
31+
VLLM_CP_ARGS=()
32+
if [ "$DCP_SIZE" -gt 1 ]; then
33+
VLLM_CP_ARGS+=(--decode-context-parallel-size "$DCP_SIZE")
34+
fi
35+
if [ "$PCP_SIZE" -gt 1 ]; then
36+
VLLM_CP_ARGS+=(--prefill-context-parallel-size "$PCP_SIZE")
37+
fi
38+
39+
GPU_COUNT="${GPU_COUNT:-$((TP * PCP_SIZE))}"
40+
if [[ ! "$GPU_COUNT" =~ ^[1-9][0-9]*$ ]]; then
41+
echo "Error: GPU_COUNT must be a positive integer, got '$GPU_COUNT'" >&2
42+
exit 1
43+
fi
44+
export GPU_COUNT
45+
2946
if declare -p SLURM_JOB_ID >/dev/null 2>&1 && [ -n "$SLURM_JOB_ID" ]; then
3047
SLURM_NODE=unknown
3148
if declare -p SLURMD_NODENAME >/dev/null 2>&1 && [ -n "$SLURMD_NODENAME" ]; then
@@ -99,7 +116,7 @@ if require_agentic_kv_offload_backend mooncake; then
99116
# Mooncake embedded mode contributes one global segment per GPU rank to
100117
# a shared distributed store. Pre-divide the aggregate host budget
101118
# across those rank-contributed segments.
102-
PER_RANK_GB=$((TOTAL_CPU_DRAM_GB / TP))
119+
PER_RANK_GB=$((TOTAL_CPU_DRAM_GB / GPU_COUNT))
103120

104121
MOONCAKE_VERSION=0.3.11.post1
105122
agentic_pip_install --quiet --no-cache-dir --no-deps \
@@ -186,6 +203,7 @@ vllm serve "$MODEL_PATH" --served-model-name "$MODEL" \
186203
--kv-cache-dtype fp8 \
187204
--block-size 256 \
188205
"${PARALLEL_ARGS[@]}" \
206+
"${VLLM_CP_ARGS[@]}" \
189207
"${EP_ARGS[@]}" \
190208
--compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}' \
191209
--attention_config.use_fp4_indexer_cache=True \

‎configs/CONFIGS.md‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@ entry-name:
1818
osl: int
1919
search-space:
2020
- { tp: int, conc-start: int, conc-end: int }
21-
# Optionally, specify 'ep' (expert-parallelism) and 'dp-attn' (data parallel attention)
22-
- { tp: int, ep: int, dp-attn: bool, conc-start: int, conc-end: int }
21+
# Optionally, specify expert/data-parallel attention and context-parallel sizes
22+
- { tp: int, ep: int, dp-attn: bool, dcp-size: int, pcp-size: int, conc-start: int, conc-end: int }
2323
- ...
2424
- ...
2525
agentic-coding: # optional
@@ -86,6 +86,9 @@ The below list describes what each field is:
8686
- Note: the step factor between `conc-start` and `conc-end` is 2, so if `conc-start` is 4 and `conc-end` is 128, all concurrencies `4, 8, 16, 32, ..., 128` will be run.
8787
- (Optional) `ep`: An integer representing the expert parallelism level that the configuration will be served at. Default is 1 (no expert parallelism) when not specified.
8888
- (Optional) `dp-attn`: A boolean representing whether or not to activate data parallel attention for the configuration. Default is false when not specified.
89+
- (Optional) `dcp-size`: Decode context-parallel size. Default is 1. It must be a positive divisor of `tp`; DCP reuses the TP GPUs.
90+
- (Optional) `pcp-size`: Prefill context-parallel size. Default is 1. A single-node job allocates `tp * pcp-size` GPUs.
91+
- `dcp-size` and `pcp-size` are single-node fields. They are not accepted inside multinode `prefill` or `decode` worker blocks.
8992
- `agentic-coding`: Agentic trace replay benchmarks using real conversation traces. Each entry must have:
9093
- `trace-source`: Identifier for the trace dataset to use.
9194
- `search-space`: Same structure as `fixed-seq-len` search-space entries.
@@ -96,7 +99,7 @@ input.
9699

97100
Notes:
98101
- No extra fields besides the ones listed may be specified, or else the benchmarks will fail to run.
99-
- Setting the fields above, particularly `ep` and `dp-attn`, only guarantee that the respective values will be passed as environment variables to the benchmark scripts! Actually using those environment variables is an implementation detail at the level of the benchmark Bash script.
102+
- Setting the fields above only guarantees that their values are passed as environment variables to benchmark scripts (`ep` as `EP_SIZE`, `dp-attn` as `DP_ATTENTION`, `dcp-size` as `DCP_SIZE`, and `pcp-size` as `PCP_SIZE`). Actually using those variables is an implementation detail of the benchmark Bash script.
100103

101104
## Runners
102105

‎runners/launch_b200-cw.sh‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,11 @@ else
3030
CONTAINER_MOUNT_DIR=/workspace
3131
fi
3232

33+
export GPU_COUNT="${GPU_COUNT:-${TP:?TP must be set}}"
34+
3335
set -x
3436

35-
JOB_ID=$(salloc --partition=$PARTITION --gres=gpu:b200:$TP --time=180 --no-shell --job-name="$RUNNER_NAME" 2>&1 | tee /dev/stderr | grep -oP 'Granted job allocation \K[0-9]+')
37+
JOB_ID=$(salloc --partition=$PARTITION --gres=gpu:b200:$GPU_COUNT --time=180 --no-shell --job-name="$RUNNER_NAME" 2>&1 | tee /dev/stderr | grep -oP 'Granted job allocation \K[0-9]+')
3638

3739
if [ -z "$JOB_ID" ]; then
3840
echo "ERROR: salloc failed to allocate a job"

‎runners/launch_b200-dgxc.sh‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -419,8 +419,10 @@ else
419419
# and gpu-15 names no longer exist. gpu-2 currently has 10 fully-idle GPU
420420
# nodes (all of gpu-2-[0-9]); gpu-1 has 2 drained (gpu-1-4, gpu-1-8). We
421421
# land on gpu-2 to avoid drained nodes and skip the per-node excludes.
422+
export GPU_COUNT="${GPU_COUNT:-${TP:?TP must be set}}"
423+
422424
SALLOC_TIME_LIMIT="${SALLOC_TIME_LIMIT:-480}"
423-
salloc --partition=$SLURM_PARTITION --account=$SLURM_ACCOUNT --gres=gpu:$TP --exclusive --mem=0 --time="$SALLOC_TIME_LIMIT" --no-shell --job-name="$RUNNER_NAME"
425+
salloc --partition=$SLURM_PARTITION --account=$SLURM_ACCOUNT --gres=gpu:$GPU_COUNT --exclusive --mem=0 --time="$SALLOC_TIME_LIMIT" --no-shell --job-name="$RUNNER_NAME"
424426
JOB_ID=$(squeue --name="$RUNNER_NAME" -u "$USER" -h -o %A | head -n1)
425427

426428
# DSv4 is also staged on the compute nodes' local RAID. Loading the 806 GB

0 commit comments

Comments
 (0)