@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 596ed1f02e

Fork 0

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, 21:30 UTC

[RL] Validation for pause_mode='keep' (#34992)

Signed-off-by: ahao-anyscale <ahao@anyscale.com>

Aaron Hao committed Feb 23, 2026 at 13:30 UTC 596ed1f02ef2a068e459a019d7d3f54593cc7e9e

2 files changed +182 -106

.buildkite/test_areas/distributed.yaml

+1 -1

@@ -104,7 +104,6 @@ steps:

104 # NEW rlhf examples

105 - cd new_weight_syncing

106 - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py

107 - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py

107

108 - label: Distributed Tests (8 GPUs)(H100)

109 timeout_in_minutes: 10

@@ -146,6 +145,7 @@ steps:

145 num_devices: 2

146 commands:

147 - pytest -v -s tests/distributed/test_context_parallel.py

148 + - cd examples/offline_inference/new_weight_syncing && VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py

149 - VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput

150 - pytest -v -s tests/v1/distributed/test_dbo.py

151

examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py

+181 -105

@@ -26,14 +26,12 @@ workloads. Residual GPU activity interferes with vLLM memory profiling and

26 causes unexpected behavior.

27 """

29 -import os

29 +import asyncio

30 import uuid

31 from dataclasses import asdict

33 import ray

34 import torch

35 -from ray.util.placement_group import placement_group

36 -from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

35 from transformers import AutoModelForCausalLM, AutoTokenizer

37 import vllm

@@ -51,14 +49,15 @@ from vllm.distributed.weight_transfer.nccl_engine import (

49 from vllm.utils.network_utils import get_ip, get_open_port

50 from vllm.v1.executor import Executor

54 -MODEL_NAME = "facebook/opt-125m"

52 +MODEL_NAME_V1 = "Qwen/Qwen3-1.7B-Base"

53 +MODEL_NAME_V2 = "Qwen/Qwen3-1.7B"

54 +PAUSE_TOKEN_THRESHOLD = 10

57 class MyLLM(vllm.AsyncLLMEngine):

58 """Configure the vLLM worker for Ray placement group execution."""

60 def __init__(self, **kwargs):

61 - os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1"

61 engine_args = vllm.AsyncEngineArgs(**kwargs)

62 vllm_config = engine_args.create_engine_config()

63 executor_class = Executor.get_class(vllm_config)

@@ -68,26 +67,44 @@ class MyLLM(vllm.AsyncLLMEngine):

67 log_requests=engine_args.enable_log_requests,

68 log_stats=not engine_args.disable_log_stats,

69 )

70 + self._generation_paused = False

71 + self._request_pause_flag = False

72 - async def generate_with_retry(

73 + async def do_generate(

74 self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams

74 - ) -> vllm.RequestOutput:

75 - finish_reason = "abort"

76 - while finish_reason == "abort":

77 - async for request_output in self.generate(

78 - {"prompt_token_ids": prompt_token_ids},

79 - sampling_params,

80 - request_id=str(uuid.uuid4()),

75 + ) -> tuple[vllm.RequestOutput, int]:

76 + """Generate a single request, setting the request pause flag once the

77 + token count reaches the threshold.

78 +

79 + Returns (output, pause_token_index). pause_token_index is the number

80 + of tokens generated before the weight change, or -1 if no pause.

81 + """

82 + pause_token_index = -1

83 + prev_token_count = 0

84 + async for request_output in self.generate(

85 + {"prompt_token_ids": prompt_token_ids},

86 + sampling_params,

87 + request_id=str(uuid.uuid4()),

88 + ):

89 + output = request_output

90 + cur_token_count = len(output.outputs[0].token_ids)

91 + if (

92 + cur_token_count >= PAUSE_TOKEN_THRESHOLD

93 + and not self._request_pause_flag

94 ):

82 - output = request_output

83 - finish_reason = output.outputs[0].finish_reason

84 - if finish_reason == "abort":

85 - print(

86 - f"ABORT, prompt_token_ids: {prompt_token_ids}, "

87 - f"generated token_ids: {list(output.outputs[0].token_ids)}"

88 - )

89 - prompt_token_ids = prompt_token_ids + list(output.outputs[0].token_ids)

90 - return output

95 + self._request_pause_flag = True

96 + if self._generation_paused and pause_token_index == -1:

97 + pause_token_index = prev_token_count

98 + prev_token_count = cur_token_count

99 + return output, pause_token_index

100 +

101 + async def pause_after_n_tokens(self):

102 + """Wait for any request to set the pause flag, then pause."""

103 + while not self._request_pause_flag:

104 + await asyncio.sleep(0)

105 + await super().pause_generation(mode="keep")

106 + await asyncio.sleep(0.2)

107 + self._generation_paused = True

108

109

110 @ray.remote(num_gpus=1)

@@ -95,6 +112,14 @@ class TrainModel:

112 """Ray actor that wraps the training model on a dedicated GPU."""

113

114 def __init__(self, model_name: str):

115 + from vllm.model_executor.layers.batch_invariant import (

116 + init_batch_invariance,

117 + )

118 + from vllm.v1.attention.backends.registry import AttentionBackendEnum

119 +

120 + # need to init all env vars for batch invariance which affect nccl ops

121 + init_batch_invariance(AttentionBackendEnum.FLASH_ATTN)

122 +

123 self.model = AutoModelForCausalLM.from_pretrained(

124 model_name, dtype=torch.bfloat16

125 ).to("cuda:0")

@@ -133,70 +158,80 @@ class TrainModel:

158 packed=packed,

159 )

160

136 -

137 -# Initialize Ray and set the visible devices. The vLLM engine will

138 -# be placed on GPUs 1 and 2.

139 -ray.init()

161 + @torch.inference_mode()

162 + def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]:

163 + """Greedy-decode max_new_tokens from the given context."""

164 + input_ids = torch.tensor([token_ids], device="cuda:0")

165 + output = self.model.generate(

166 + input_ids,

167 + max_new_tokens=max_new_tokens,

168 + do_sample=False,

169 + )

170 + new_token_ids = output[0, len(token_ids) :].tolist()

171 + return new_token_ids

172 +

173 +

174 +ray.init(

175 + runtime_env={

176 + "env_vars": {

177 + # enable batch invariance for deterministic outputs

178 + "VLLM_BATCH_INVARIANT": "1",

179 + # prevent ray from setting CUDA_VISIBLE_DEVICES

180 + "RAY_EXPERIMENTAL_NOSET_CUDA_ENV_VAR": "1",

181 + }

182 + }

183 +)

184

185 # Launch the training model actor. Ray's resource scheduler will allocate

186 # 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.

143 -train_model = TrainModel.remote(MODEL_NAME)

144 -

145 -# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.

146 -# Learn more about Ray placement groups:

147 -# https://docs.ray.io/en/latest/placement-groups.html

148 -

149 -pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)

150 -ray.get(pg_inference.ready())

151 -scheduling_inference = PlacementGroupSchedulingStrategy(

152 - placement_group=pg_inference,

153 - placement_group_capture_child_tasks=True,

154 - placement_group_bundle_index=0,

155 -)

187 +train_model = TrainModel.remote(MODEL_NAME_V2)

188

189 # Launch the vLLM inference engine. The `enforce_eager` flag reduces

190 # start-up latency.

159 -# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights)

160 -# are now native to vLLM workers.

191 +# With data_parallel_backend="ray", vLLM's CoreEngineActorManager creates

192 +# its own placement groups internally for each DP rank, so we must NOT

193 +# create an outer placement group (it would reserve GPUs and hide them

194 +# from the internal DP resource check).

195 llm = ray.remote(

196 num_cpus=0,

197 num_gpus=0,

164 - scheduling_strategy=scheduling_inference,

198 )(MyLLM).remote(

166 - model=MODEL_NAME,

199 + model=MODEL_NAME_V1,

200 enforce_eager=True,

168 - tensor_parallel_size=2,

201 + max_model_len=8192,

202 distributed_executor_backend="ray",

170 - load_format="dummy",

203 + attention_backend="FLASH_ATTN",

204 + gpu_memory_utilization=0.75,

205 weight_transfer_config=WeightTransferConfig(backend="nccl"),

206 )

207

174 -# Generate text from the prompts.

175 -prompts = [

176 - "My name is",

208 +PROMPTS = [

209 "The president of the United States is",

210 "The capital of France is",

179 - "The future of AI is",

211 + "The largest ocean on Earth is",

212 + "The speed of light in a vacuum is",

213 + "The chemical formula for water is",

214 + "The tallest mountain in the world is",

215 + "The first person to walk on the moon was",

216 + "The Great Wall of China was built to",

217 + "Photosynthesis is the process by which",

218 + "The theory of general relativity was proposed by",

219 + "The boiling point of water at sea level is",

220 + "The largest planet in our solar system is",

221 + "DNA stands for deoxyribonucleic acid and it",

222 ]

223

182 -# Tokenize prompts to token IDs

183 -tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

184 -prompt_token_ids_list = [

185 - tokenizer.encode(prompt, add_special_tokens=False) for prompt in prompts

224 +tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_V1)

225 +batch_prompt_token_ids = [

226 + tokenizer.encode(prompt, add_special_tokens=False) for prompt in PROMPTS

227 ]

228

188 -sampling_params = [

189 - SamplingParams(temperature=0, max_tokens=2),

190 - SamplingParams(temperature=0, max_tokens=32),

191 - SamplingParams(temperature=0, max_tokens=32),

192 - SamplingParams(temperature=0, max_tokens=32),

193 -]

229

230 # Set up the communication channel between the training process and the

231 # inference engine.

232 master_address, master_port = ray.get(train_model.get_master_address_and_port.remote())

233

199 -world_size = 3 # 1 trainer + 2 inference workers (tensor_parallel_size=2)

234 +world_size = 2 # 1 trainer + 1 inference worker

235 inference_handle = llm.init_weight_transfer_engine.remote(

236 WeightTransferInitRequest(

237 init_info=asdict(

@@ -215,22 +250,28 @@ train_handle = train_model.init_weight_transfer_group.remote(world_size)

250 ray.get([train_handle, inference_handle])

251

252

218 -generation_futures = [

219 - llm.generate_with_retry.remote(prompt_token_ids, params)

220 - for prompt_token_ids, params in zip(prompt_token_ids_list, sampling_params)

221 -]

253 +N_NEW_TOKENS = 100

254

223 -finished, pending = ray.wait(generation_futures, num_returns=1)

255 +# Collect weight metadata once

256 +names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())

257

225 -# Pause generation in preparation for weight sync

226 -ray.get(llm.pause_generation.remote(wait_for_inflight_requests=False))

258 +# ── Phase 1: concurrent requests with weight sync ───────────────────

259 +print(f"\n{'=' * 50}")

260 +print(f"Prompts ({len(PROMPTS)}):")

261 +for p in PROMPTS:

262 + print(f" - {p!r}")

263 +print(f"{'=' * 50}")

264

228 -# Synchronize the updated weights to the inference engine using batched API.

229 -# Collect all weight metadata from the training actor

230 -names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())

265 +sampling_params = SamplingParams(

266 + temperature=0, max_tokens=PAUSE_TOKEN_THRESHOLD + N_NEW_TOKENS

267 +)

268 +

269 +gen_futures = [

270 + llm.do_generate.remote(ptids, sampling_params) for ptids in batch_prompt_token_ids

271 +]

272 +

273 +ray.get(llm.pause_after_n_tokens.remote())

274

232 -# Issue update_weights call with NCCL-specific update info

233 -# packed=True enables efficient batched tensor broadcasting

275 inference_handle = llm.update_weights.remote(

276 WeightTransferUpdateRequest(

277 update_info=asdict(

@@ -243,41 +284,76 @@ inference_handle = llm.update_weights.remote(

284 )

285 )

286 )

246 -

247 -# Broadcast all weights from trainer using the weight transfer API

287 train_handle = train_model.broadcast_weights.remote(packed=True)

288 ray.get([train_handle, inference_handle])

289

251 -# Resume generation since weight sync is complete

290 ray.get(llm.resume_generation.remote())

291 +results = ray.get(gen_futures)

292 +

293 +for i, (output, pause_idx) in enumerate(results):

294 + all_token_ids = list(output.outputs[0].token_ids)

295 + before_text = tokenizer.decode(all_token_ids[:pause_idx])

296 + after_text = tokenizer.decode(all_token_ids[pause_idx:])

297 + print(f"\n Request {i} ({PROMPTS[i]!r}):")

298 + print(f" Old weights ({pause_idx} tokens): {before_text!r}")

299 + n_after = len(all_token_ids) - pause_idx

300 + print(f" New weights ({n_after} tokens): {after_text!r}")

301 +

302 +# ── Phase 2: validate with a fresh V2 vLLM instance ────────────────

303 +print(f"\n{'=' * 50}")

304 +print("VALIDATION: comparing weight-synced vLLM with fresh V2 instance")

305 +print(f"{'=' * 50}")

306 +

307 +ray.get(llm.shutdown.remote())

308 +ray.kill(llm)

309 +ray.kill(train_model)

310 +

311 +llm_v2 = ray.remote(

312 + num_cpus=0,

313 + num_gpus=0,

314 +)(MyLLM).remote(

315 + model=MODEL_NAME_V2,

316 + enforce_eager=True,

317 + max_model_len=8192,

318 + gpu_memory_utilization=0.75,

319 + distributed_executor_backend="ray",

320 + attention_backend="FLASH_ATTN",

321 +)

322 +

323 +val_futures = [

324 + llm_v2.do_generate.remote(

325 + list(output.prompt_token_ids) + list(output.outputs[0].token_ids)[:pause_idx],

326 + SamplingParams(

327 + temperature=0, max_tokens=len(output.outputs[0].token_ids) - pause_idx

328 + ),

329 + )

330 + for output, pause_idx in results

331 +]

332 +val_results = ray.get(val_futures)

333 +

334 +all_pass = True

335 +for i, ((output, pause_idx), (val_output, _)) in enumerate(zip(results, val_results)):

336 + expected = list(output.outputs[0].token_ids)[pause_idx:]

337 + actual = list(val_output.outputs[0].token_ids)

338 + match = actual == expected

339 +

340 + if match:

341 + print(f" [PASS] {PROMPTS[i]!r}")

342 + else:

343 + all_pass = False

344 + print(f" [FAIL] {PROMPTS[i]!r}")

345 + print(f" weight-synced vLLM: {tokenizer.decode(expected)!r}")

346 + print(f" V2 vLLM: {tokenizer.decode(actual)!r}")

347 + for j, (e, a) in enumerate(zip(expected, actual)):

348 + if e != a:

349 + print(

350 + f" first divergence at output token {j}: "

351 + f"expected {e} ({tokenizer.decode([e])!r}) vs "

352 + f"actual {a} ({tokenizer.decode([a])!r})"

353 + )

354 + break

355

254 -# Get outputs separately - finished completed before pause, pending were paused/resumed

255 -finished_outputs = ray.get(finished)

256 -pending_outputs = ray.get(pending)

257 -

258 -# Requests that finished before the pause: all generation used original weights

259 -print("-" * 50)

260 -print("Requests that completed BEFORE weight change:")

261 -print("-" * 50)

262 -for output in finished_outputs:

263 - prompt_text = tokenizer.decode(output.prompt_token_ids)

264 - print(f"Prompt: {prompt_text!r}")

265 - print(f"Generated (with original weights): {output.outputs[0].text!r}")

266 - print("-" * 50)

267 -

268 -# Requests that were paused mid-generation: some text before, some after weight change

269 -print("Requests that were PAUSED and RESUMED after weight change:")

270 -print("-" * 50)

271 -for output in pending_outputs:

272 - # Decode the full prompt token IDs (original + generated before pause)

273 - full_prompt_text = tokenizer.decode(output.prompt_token_ids)

274 - # Find the original prompt by checking which one this output started with

275 - original_prompt = next(p for p in prompts if full_prompt_text.startswith(p))

276 - # output.prompt_token_ids contains original prompt + tokens generated before pause

277 - # output.outputs[0].text is what was generated after resuming with new weights