[CI][ROCm] Prefetch safetensors weights in AMD CI (#54898)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Andreas Karatzas committed
Sep 2, 2026 at 03:28 UTC
c6005000d326fe2726aed62553f5837d42c5e04c
5 files changed
+15
-5
.buildkite/lm-eval-harness/configs/Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml
+1
-1
@@ -1,7 +1,7 @@
1
# For hf script, without -t option (tensor parallel size).
2
# bash .buildkite/lm-eval-harness/run-lm-eval-mmlupro-vllm-baseline.sh -m meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8 -l 250 -t 8 -f 5
3
model_name: "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
4
-rocm_safetensors_load_strategy: lazy
4
+rocm_safetensors_load_strategy: prefetch
5
required_gpu_arch:
6
- gfx942
7
- gfx950
tests/weight_loading/models-amd.txt
+2
-1
@@ -1,6 +1,7 @@
1
+# quantization, model, revision[, minimum capability[, safetensors load strategy]]
2
fp8, amd/Meta-Llama-3.1-8B-Instruct-FP8-KV, main
3
None, amd/Llama-3.2-1B-Instruct-FP8-KV, main
3
-fp8, amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV, main
4
+fp8, amd/Mixtral-8x7B-Instruct-v0.1-FP8-KV, main, 80, prefetch
5
gptq, robertgshaw2/zephyr-7b-beta-channelwise-gptq, main
6
gptq, TheBloke/Llama-2-7B-GPTQ, main
7
gptq, TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ, gptq-8bit--1g-actorder_True
tests/weight_loading/models-large-amd.txt
+4
-3
@@ -1,3 +1,4 @@
1
-fp8, amd/Meta-Llama-3.1-70B-Instruct-FP8-KV, main
2
-None, microsoft/phi-4, main
3
-fp8, amd/Mixtral-8x22B-Instruct-v0.1-FP8-KV, main
1
+# quantization, model, revision[, minimum capability[, safetensors load strategy]]
2
+fp8, amd/Meta-Llama-3.1-70B-Instruct-FP8-KV, main, 80, prefetch
3
+None, microsoft/phi-4, main, 80, prefetch
4
+fp8, amd/Mixtral-8x22B-Instruct-v0.1-FP8-KV, main, 80, prefetch
tests/weight_loading/run_model_weight_loading_test.sh
+4
@@ -28,6 +28,7 @@ do
28
29
echo "=== RUNNING MODEL: $MODEL_CONFIG ==="
30
31
+ unset WEIGHT_LOADING_TEST_SAFETENSORS_LOAD_STRATEGY
32
export QUANTIZATION=${array[0]}
33
export MODEL_NAME=${array[1]}
34
export REVISION=${array[2]}
@@ -35,6 +36,9 @@ do
36
if [ ${#array[@]} -gt 3 ]; then
37
export MIN_CAPABILITY=${array[3]}
38
fi
39
+ if [ ${#array[@]} -gt 4 ]; then
40
+ export WEIGHT_LOADING_TEST_SAFETENSORS_LOAD_STRATEGY=${array[4]}
41
+ fi
42
pytest -s weight_loading/test_weight_loading.py || LOCAL_SUCCESS=$?
43
44
if [[ $LOCAL_SUCCESS == 0 ]]; then
tests/weight_loading/test_weight_loading.py
+4
@@ -16,6 +16,9 @@ MODEL_NAME = os.environ.get(
16
REVISION = os.environ.get("REVISION", "main")
17
QUANTIZATION = os.environ.get("QUANTIZATION", "gptq_marlin")
18
MIN_CAPABILITY = os.environ.get("MIN_CAPABILITY", "80")
19
+SAFETENSORS_LOAD_STRATEGY = os.environ.get(
20
+ "WEIGHT_LOADING_TEST_SAFETENSORS_LOAD_STRATEGY"
21
+)
22
23
24
@pytest.mark.skipif(
@@ -42,6 +45,7 @@ def test_weight_loading(vllm_runner):
45
quantization=None if QUANTIZATION == "None" else QUANTIZATION,
46
max_model_len=MAX_MODEL_LEN,
47
tensor_parallel_size=2,
48
+ safetensors_load_strategy=SAFETENSORS_LOAD_STRATEGY,
49
# Generating 20 tokens needs a tiny KV cache, while sizing and
50
# allocating a full-size one dominates this test on large devices.
51
kv_cache_memory_bytes=2 * GiB_bytes,