A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 07:45 UTC
[ROCm][CI] Support async weight transfer example with platform-aware determinism (#35710)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 3, 2026 at 19:44 UTC
f7da9cdffca2d7f11882249550cfa20605a0ca04
2 files changed
+91
-33
.buildkite/test-amd.yaml
+9
-3
@@ -1339,6 +1339,7 @@ steps:
1339
- tests/v1/entrypoints/openai/test_multi_api_servers.py
1340
- tests/v1/shutdown
1341
- tests/v1/worker/test_worker_memory_snapshot.py
1342
+ - examples/offline_inference/new_weight_syncing/
1343
commands:
1344
# Work around HIP bug tracked here: https://github.com/ROCm/hip/issues/3876
1345
# TODO: Remove when the bug is fixed in a future ROCm release
@@ -1970,8 +1971,10 @@ steps:
1971
1972
- label: Distributed Tests (4 GPUs) # 35min
1973
timeout_in_minutes: 50
1973
- mirror_hardwares: [amdexperimental]
1974
+ mirror_hardwares: [amdexperimental, amdproduction]
1975
agent_pool: mi355_4
1976
+ optional: true
1977
+ # grade: Blocking
1978
working_dir: "/vllm-workspace/tests"
1979
num_gpus: 4
1980
source_file_dependencies:
@@ -2025,7 +2028,8 @@ steps:
2028
- popd
2029
# NEW rlhf examples
2030
- pushd ../examples/offline_inference/new_weight_syncing
2028
- - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
2031
+ - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_nccl.py
2032
+ - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_ipc.py
2033
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
2034
- popd
2035
@@ -2989,8 +2993,10 @@ steps:
2993
2994
- label: Distributed Tests (2 GPUs) # 68min
2995
timeout_in_minutes: 90
2992
- mirror_hardwares: [amdexperimental]
2996
+ mirror_hardwares: [amdexperimental, amdproduction]
2997
agent_pool: mi355_2
2998
+ optional: true
2999
+ # grade: Blocking
3000
working_dir: "/vllm-workspace/tests"
3001
num_gpus: 2
3002
source_file_dependencies:
examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py
+82
-30
@@ -47,12 +47,14 @@ from vllm.distributed.weight_transfer.nccl_engine import (
47
NCCLWeightTransferInitInfo,
48
NCCLWeightTransferUpdateInfo,
49
)
50
+from vllm.platforms import current_platform
51
from vllm.utils.network_utils import get_ip, get_open_port
52
from vllm.v1.executor import Executor
53
54
MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base"
55
MODEL_NAME_V2 = "Qwen/Qwen3-1.7B"
56
PAUSE_TOKEN_THRESHOLD = 10
57
+ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "FLASH_ATTN"
58
59
60
class MyLLM(vllm.AsyncLLMEngine):
@@ -116,10 +118,16 @@ class TrainModel:
118
from vllm.model_executor.layers.batch_invariant import (
119
init_batch_invariance,
120
)
121
+ from vllm.platforms import current_platform
122
from vllm.v1.attention.backends.registry import AttentionBackendEnum
123
124
# need to init all env vars for batch invariance which affect nccl ops
122
- init_batch_invariance(AttentionBackendEnum.FLASH_ATTN)
125
+ attn_backend = (
126
+ AttentionBackendEnum.TRITON_ATTN
127
+ if current_platform.is_rocm()
128
+ else AttentionBackendEnum.FLASH_ATTN
129
+ )
130
+ init_batch_invariance(attn_backend)
131
132
self.model = AutoModelForCausalLM.from_pretrained(
133
model_name, dtype=torch.bfloat16
@@ -175,39 +183,56 @@ class TrainModel:
183
return new_token_ids
184
185
178
-ray.init(
179
- runtime_env={
180
- "env_vars": {
181
- # enable batch invariance for deterministic outputs
182
- "VLLM_BATCH_INVARIANT": "1",
183
- # prevent ray from setting CUDA_VISIBLE_DEVICES
184
- "RAY_EXPERIMENTAL_NOSET_CUDA_ENV_VAR": "1",
185
- }
186
- }
187
-)
186
+# Build platform-specific env vars for Ray
187
+ray_env_vars = {
188
+ # Prevent Ray from setting CUDA_VISIBLE_DEVICES
189
+ "RAY_EXPERIMENTAL_NOSET_CUDA_ENV_VAR": "1",
190
+}
191
+
192
+if current_platform.is_rocm():
193
+ # For ROCm, BATCH_INVARIANT vllm is not supported
194
+ ray_env_vars["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
195
+else:
196
+ # Enable batch invariance for deterministic outputs on NVIDIA
197
+ ray_env_vars["VLLM_BATCH_INVARIANT"] = "1"
198
+
199
+ray.init(runtime_env={"env_vars": ray_env_vars})
200
201
# Launch the training model actor. Ray's resource scheduler will allocate
202
# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
203
train_model = TrainModel.remote(MODEL_NAME_V2)
204
193
-# Launch the vLLM inference engine. The `enforce_eager` flag reduces
194
-# start-up latency.
195
-# With data_parallel_backend="ray", vLLM's CoreEngineActorManager creates
196
-# its own placement groups internally for each DP rank, so we must NOT
197
-# create an outer placement group (it would reserve GPUs and hide them
198
-# from the internal DP resource check).
199
-llm = ray.remote(
200
- num_cpus=0,
201
- num_gpus=0,
202
-)(MyLLM).remote(
205
+rocm_determinism_kwargs = {}
206
+if current_platform.is_rocm():
207
+ # ROCm: To minimize non-determinism, we set fixed seed, no prefix caching, and
208
+ # sequential request processing (max_num_seqs=1).
209
+ rocm_determinism_kwargs = {
210
+ "seed": 0,
211
+ "enable_prefix_caching": False,
212
+ "max_num_seqs": 1,
213
+ }
214
+
215
+# Build platform-specific LLM kwargs
216
+llm_kwargs = dict(
217
model=MODEL_NAME_V1,
218
enforce_eager=True,
219
max_model_len=8192,
220
distributed_executor_backend="ray",
207
- attention_backend="FLASH_ATTN",
221
+ attention_backend=ATTN_BACKEND,
222
gpu_memory_utilization=0.75,
223
weight_transfer_config=WeightTransferConfig(backend="nccl"),
224
)
225
+llm_kwargs.update(rocm_determinism_kwargs)
226
+
227
+# Launch the vLLM inference engine.
228
+# With data_parallel_backend="ray", vLLM's CoreEngineActorManager creates
229
+# its own placement groups internally for each DP rank, so we must NOT
230
+# create an outer placement group (it would reserve GPUs and hide them
231
+# from the internal DP resource check).
232
+llm = ray.remote(
233
+ num_cpus=0,
234
+ num_gpus=0,
235
+)(MyLLM).remote(**llm_kwargs)
236
237
PROMPTS = [
238
"The president of the United States is",
@@ -304,25 +329,42 @@ for i, (output, pause_idx) in enumerate(results):
329
print(f" New weights ({n_after} tokens): {after_text!r}")
330
331
# ── Phase 2: validate with a fresh V2 vLLM instance ────────────────
332
+# This validation relies on batch-invariant (deterministic) generation to
333
+# compare outputs from the weight-synced engine against a fresh V2 instance.
334
+# On NVIDIA, batch invariance is fully supported, so we require 100% exact
335
+# token match. On ROCm, batch invariance is not yet fully implemented
336
+# (see https://github.com/vllm-project/vllm/issues/27433 and
337
+# https://github.com/vllm-project/vllm/issues/33123), so residual
338
+# non-determinism (e.g. GEMM accumulation order, missing kernel overrides)
339
+# can cause single-token divergences that don't indicate a weight-sync
340
+# failure. We relax the pass rate to 90% on ROCm to accommodate this; a
341
+# real regression (broken weight transfer) would cause ~0% pass rate, not 90%+.
342
+MIN_PASS_RATE = 1.0 if not current_platform.is_rocm() else 0.9
343
+
344
print(f"\n{'=' * 50}")
345
print("VALIDATION: comparing weight-synced vLLM with fresh V2 instance")
346
+if current_platform.is_rocm():
347
+ print(f" (ROCm mode: requiring >= {MIN_PASS_RATE:.0%} exact match rate)")
348
print(f"{'=' * 50}")
349
350
ray.get(llm.shutdown.remote())
351
ray.kill(llm)
352
ray.kill(train_model)
353
315
-llm_v2 = ray.remote(
316
- num_cpus=0,
317
- num_gpus=0,
318
-)(MyLLM).remote(
354
+llm_v2_kwargs = dict(
355
model=MODEL_NAME_V2,
356
enforce_eager=True,
357
max_model_len=8192,
358
gpu_memory_utilization=0.75,
359
distributed_executor_backend="ray",
324
- attention_backend="FLASH_ATTN",
360
+ attention_backend=ATTN_BACKEND,
361
)
362
+llm_v2_kwargs.update(rocm_determinism_kwargs)
363
+
364
+llm_v2 = ray.remote(
365
+ num_cpus=0,
366
+ num_gpus=0,
367
+)(MyLLM).remote(**llm_v2_kwargs)
368
369
val_futures = [
370
llm_v2.do_generate.remote(
@@ -335,16 +377,17 @@ val_futures = [
377
]
378
val_results = ray.get(val_futures)
379
338
-all_pass = True
380
+num_pass = 0
381
+num_total = len(results)
382
for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_results)):
383
expected = list(output.outputs[0].token_ids)[pause_idx:]
384
actual = list(val_output.outputs[0].token_ids)
385
match = actual == expected
386
387
if match:
388
+ num_pass += 1
389
print(f" [PASS] {PROMPTS[i]!r}")
390
else:
347
- all_pass = False
391
print(f" [FAIL] {PROMPTS[i]!r}")
392
print(f" weight-synced vLLM: {tokenizer.decode(expected)!r}")
393
print(f" V2 vLLM: {tokenizer.decode(actual)!r}")
@@ -359,5 +402,14 @@ for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_resu
402
403
ray.get(llm_v2.shutdown.remote())
404
ray.kill(llm_v2)
362
-assert all_pass, "Some prompts failed validation, see above for details"
405
+
406
+pass_rate = num_pass / num_total
407
+print(f"\n Result: {num_pass}/{num_total} prompts passed ({pass_rate:.0%})")
408
+print(f" Required: >= {MIN_PASS_RATE:.0%}")
409
+
410
+assert pass_rate >= MIN_PASS_RATE, (
411
+ f"Validation pass rate {pass_rate:.0%} ({num_pass}/{num_total}) "
412
+ f"is below the required {MIN_PASS_RATE:.0%} threshold. "
413
+ f"See failures above for details."
414
+)
415
print("=" * 50)