feat(rollout): add vLLM/vLLM-Ascend rollout path for DFlash OPD on Ascend NPU - #8
Open
curnane-lab wants to merge 45 commits into
Open
feat(rollout): add vLLM/vLLM-Ascend rollout path for DFlash OPD on Ascend NPU#8curnane-lab wants to merge 45 commits into
curnane-lab wants to merge 45 commits into
Conversation
Mirror the DFLASH OPD contract of the SGLang rollout so the training side (composed DFLASH student + distillation losses) stays engine-agnostic: - map_opd_draft_weight_name / iter_opd_draft_weights: strip FSDP/wrapper prefixes and keep only the draft_model.* subtree, renamed to draft-relative keys the inference engine expects. - build_dflash_extra_fields: build TokenOutput.extra_fields from engine-side reject metadata (anchor indices / offsets / token ids / teacher logprobs), matching the keys consumed by agent_loop. - merge_dflash_metadata: tolerate None parts from TP ranks.
…apter When the composed DFLASH/EAGLE3 OPD student is enabled, engine_workers passes dflash_draft_only=True. The vLLM ServerAdapter now: - filters the weight stream to the draft_model.* subtree (renamed to draft-relative keys via opd_utils.iter_opd_draft_weights); - forwards draft_model_only=True to the worker extension so weights are loaded into the speculative draft model instead of the frozen target; - counts synced tensors, raises on an empty stream, and returns opd/weight_sync/draft_tensor_count like the SGLang path. Rejects the incompatible combination with LoRA adapter sync.
Extend vLLMColocateWorkerExtension.update_weights_from_ipc with a draft_model_only flag: when set, received tensors are routed to model_runner.drafter.model.load_weights() (e.g. DFlashQwen3ForCausalLM) instead of the frozen target model, and target-only post-load transforms (process_weights_after_loading / MoE weight-loader patch) are skipped. The draft-relative key naming lines up with vllm's DFlash loader, which re-adds the 'model.' prefix itself. Raises a clear error when the server was started without a speculative_config.
OPD training needs per-request reject positions and teacher logprobs of rejected draft tokens, which stock vLLM/vllm-ascend does not expose (the SGLang path gets them from the sglang-dflash fork). This module monkey-patches NPUModelRunner._sample: after each verify step it replays acceptance per request (num_accepted = non-placeholder tokens - 1, robust for random rejection sampling where the recovered token may coincide with the draft token), records the rejected draft token id, its 1-based block offset, and the target logprob of the rejected token (TP-aware gather from the raw target logits, before top-k/top-p). Events accumulate in a per-request registry with cumulative response position tracking (response starts with the prefill token at position 0; each verify step appends num_accepted + 1 tokens). pop() materializes the same anchor/offset/token_id/teacher_logprob contract the SGLang rollout produces. Activated via VERL_DFLASH_OPD=1; failures never break serving.
- vLLMColocateWorkerExtension: apply dflash_opd_patch at worker init when VERL_DFLASH_OPD=1; expose get_dflash_opd_metadata / gc_dflash_opd_metadata for collective_rpc from the rollout frontend. - vLLMHttpServer: detect the composed DFlash/EAGLE3 OPD student (hf_config.verl_composed_dflash_student), force load_format=auto when it is dummy (target weights must be real since only the draft is trained), and set VERL_DFLASH_OPD=1 before engine launch so workers inherit it. - vLLMHttpServer.generate: after the final output, fetch the request's reject metadata via collective_rpc and merge it into extra_fields with opd_utils.build_dflash_extra_fields, producing the same dflash_* keys the SGLang rollout path emits (agent_loop / OPD losses unchanged).
- MtpConfig gains draft_model_path: a standalone draft checkpoint for methods that use a separate draft model (dflash/eagle). It is passed as speculative_config["model"]; None keeps target-integrated methods (mtp). - vLLMHttpServer forwards it into speculative_config alongside method and num_speculative_tokens when mtp rollout is enabled. - run_qwen_gsm8k.sh: ROLLOUT_NAME becomes overridable (default sglang). - New example run_qwen_gsm8k_vllm_npu.sh: Draft-OPD on Ascend NPU with the vLLM(vllm-ascend) rollout engine (method=dflash, K=16 by default); whitelisted in .gitignore like the other public entrypoints.
Cover the engine-agnostic pieces that do not need vLLM/NPU hardware: - map_opd_draft_weight_name / iter_opd_draft_weights: wrapper-prefix stripping, nested draft_model marker, target-weight rejection. - DFlashOpdMetadataRegistry.pop: response-position replay across verify steps (prefill token at position 0, each verify appends accepted+1), offset/token/teacher-logprob passthrough, pop-once semantics. - build_dflash_extra_fields: parity with the SGLang contract keys, out-of-range anchor filtering, empty/missing metadata, length checks.
Clone with --recursive now yields the full vLLM(vllm-ascend) engine stack: - vllm @ 7c5deda2 (curnane-lab/vllm, domino_npu_v0.21.0) - vllm-ascend @ f9fbb0a8 (curnane-lab/vllm-ascend, domino_npu_v0.21.0rc1, which carries the official DFlash speculative method plus the longer-draft buffer fixes the OPD rollouts rely on) install.sh now runs git submodule update --init --recursive (so a plain clone also works) and installs vllm (VLLM_TARGET_DEVICE=empty) and vllm-ascend (with its nested custom-operator submodules) before sglang-dflash and verl. README documents the submodule layout and the Ascend NPU notes (CANN/torch-npu/triton-ascend prerequisites and the run_qwen_gsm8k_vllm_npu.sh entrypoint).
dflash_opd_patch recorder: - Batch the per-step device syncs: one (sampled != PLACEHOLDER).sum(dim=1) plus one draft_token_ids.tolist() per verify step instead of a per-request .item() (syncs are expensive on NPU). - Record plain decode steps (num_draft_tokens == 0) as (0,0,-1) events so the response-position replay stays aligned when a request occasionally has a draft-free step. - Remove duplicated slicing/check block left by the previous edit. run_qwen_gsm8k_vllm_npu.sh: - ROLLOUT_ENFORCE_EAGER (default True) for the first NPU bring-up so aclgraph failure modes are out of scope; flip to False after validation. - ROLLOUT_FREE_CACHE_ENGINE knob (default True; vllm-ascend implements worker.sleep/wake_up) with the documented fallback for hybrid colocate. - Comment the recommended PYTORCH_NPU_ALLOC_CONF=max_split_size_mb:32.
- INSTALL_VLLM / INSTALL_VLLM_ASCEND / INSTALL_SGLANG_DFLASH switches so container images that already ship the engines (e.g. editable installs under /vllm-workspace) can skip those steps instead of clobbering them. - Warn against installing verl with extras: verl[vllm] pins vllm<=0.12.0 and verl[sglang] pins torch==2.9.1, both incompatible with an NPU / vllm-ascend v0.21 stack. Plain 'pip install -e ./verl' is unpinned. - requirements-npu.txt: triton-ascend 3.2.0 -> 3.2.1, matching the vllm-ascend v0.21.0rc1 installation guide.
Auto-detects each component with importlib.metadata and skips what is already installed, so one script covers both pre-provisioned NPU containers (vllm/vllm-ascend already present, e.g. /vllm-workspace) and bare environments (engines installed from the in-repo submodules): - vllm / vllm_ascend: skip when present; warn when the version is not the v0.21.x the OPD patches target. git submodule update only runs when an engine actually needs installing, so plain clones without --recursive work in pre-provisioned containers. - triton-ascend: skip when present, else install ==3.2.1. - sglang: skip when present, else install the vendored sglang-dflash; SKIP_SGLANG_DFLASH=1 opts out. - verl: skip when already editable-installed from this repo; always installed without extras (verl[vllm]/[sglang] pin old engines and would break an NPU / vllm-ascend v0.21 stack). - torch/torch_npu presence is checked and warned about but never auto-installed (must match the CANN toolkit). - FORCE_INSTALL_ENGINES=1 forces engine reinstall from submodules.
The vLLM(vllm-ascend) rollout path does not use the sglang engine, and installing the vendored CUDA-ecosystem sglang-dflash can disturb the container's torch/torch_npu stack via dependency resolution. It is now opt-in (INSTALL_SGLANG_DFLASH=1) for users who need the SGLang reference path; the old SKIP_SGLANG_DFLASH switch is removed.
- install.sh is now a thin wrapper around install_npu.sh so there is a single adaptive implementation (previously it still installed the engines and sglang-dflash by default, which could clobber a pre-provisioned container or disturb its torch/torch_npu stack). - README Install section now documents the adaptive behavior, the sglang opt-in (INSTALL_SGLANG_DFLASH=1), FORCE_INSTALL_ENGINES=1, and the no-extras rule for verl. - CONTRIBUTING.md gains an Ascend NPU note: do not install the [vllm]/[sglang] extras (they pin vllm<=0.12.0 and sglang==0.5.8 + torch==2.9.1 and would downgrade the pre-installed engines).
vLLMHttpServer._init_config instantiates the rollout config via hydra instantiate(_convert_="partial"). rollout.mtp is normally an interpolation to the structured model.mtp (with _target_=MtpConfig), but force-add (++) CLI overrides replace that node with a plain dict lacking _target_, so instantiate leaves mtp as a dict and "self.config.mtp.enable" raised AttributeError: 'dict' object has no attribute 'enable'. Two-sided fix: - vllm_async_server: read mtp fields through _config_node_get, which handles dataclass / DictConfig / plain dict uniformly. - run_qwen_gsm8k_vllm_npu.sh: pass ++actor_rollout_ref.rollout.mtp._target_=verl.workers.config.MtpConfig as the first mtp override so instantiate recreates a real MtpConfig.
DFlashQwen3ForCausalLM.load_weights returns None (it delegates to an inner AutoWeightsLoader and finishes with _build_fused_kv_buffers), so the draft-only sync crashed in the f-string after the load had already succeeded. Count the received weight tensors instead.
The training side crashed with 'DFLASH OPD reject metadata batch size mismatch: got 0 entries' when a rollout sample carried no dflash keys. Three layered fixes: 1. SGLang-parity contract: _collect_dflash_opd_fields now returns the full dflash_* key set with empty values (new opd_utils.empty_dflash_extra_fields) when a request produced no verify events, instead of omitting the keys entirely. A single short-response sample (EOS right after prefill) can no longer kill the whole micro batch. 2. Startup probe: vLLMHttpServer runs a collective_rpc probe right after engine init when VERL_DFLASH_OPD=1 and raises immediately if the _sample patch is not applied or no drafter exists on any worker — fail fast at startup instead of failing at the first training step. 3. Robust activation: VERL_DFLASH_OPD=1 is now set whenever the effective speculative_config uses method=dflash (the composed-student flag is no longer the sole trigger), with a startup log line showing composed_student / speculative_config / env state; generate() gates collection on the env flag; the first missing-metadata case is logged.
Engine core processes died at startup once VERL_DFLASH_OPD=1 activated apply_dflash_opd_patches for the first time. The most likely cause is a vllm-ascend version drift: if the container's NPUModelRunner lacks the _sample attribute the patch wraps, the AttributeError propagated through the worker extension __new__ and killed the engine worker during launch_core_engines. - apply_dflash_opd_patches now catches every exception path (import failure, missing _sample attribute, anything else), logs clearly, and returns False instead of raising. - The worker extension wraps the call in try/except as a second layer, so a patch failure can never kill an engine worker; the startup probe (_assert_dflash_opd_pipeline) is responsible for reporting the broken state with a clear error.
…ptance stats Gated by OPD_PATCH_DEBUG=1. Prints per-request num_draft_tokens, non-placeholder valid counts, the first sampled row, the draft token ids and the derived rejected events for the first verify step, to diagnose why recorded events carry no rejections (empty anchors).
…q id and logits magnitude The single-shot dump fired during the engine startup dummy/profile run (all-zero tensors), which is indistinguishable from real traffic without the request id. Log the first 8 steps instead, include req0 and the raw target logits abs-max to spot degenerate forward passes.
…se token sampling - skip warmup/profile dummy requests at the recording entry so their all-zero rows do not pollute the registry - OPD_PATCH_DEBUG now logs the first 8 verify steps plus every 500th, tagged with an allzero flag for the sampled/draft rows - vLLMHttpServer logs the first 32 token ids of the first 3 finished responses per server, the decisive check for degenerate generation
Match the value validated by standalone vllm-ascend dflash serving; 16 remains available via the NUM_SPECULATIVE_TOKENS env override.
The composed DFLASH/EAGLE3 OPD student re-syncs only the trainable draft subtree after wake (dflash_draft_only); the frozen target is loaded once at engine init and never streamed again. _sleep_hybrid hard-coded sleep level 2, which frees weight memory without a host copy, so after the first sleep the target ran on uninitialized memory: target logits -> inf, argmax collapsed to token 0, every draft token was trivially accepted (no rejection anchors), OPD loss was identically 0 and generation degenerated (responses filled to max_tokens, never emitting EOS). Use sleep level 1 for the composed student (and, as before, for LoRA adapters) so resident weights survive sleep/wake. This mirrors the SGLang path, which already sets enable_weights_cpu_backup=True for the composed student. Co-authored-by: curnane-lab <curnane-lab@users.noreply.github.com>
…eight-sleep-8e24 fix(rollout): keep frozen DFlash target weights across HYBRID sleep (vLLM OPD zero-loss root cause)
…cher Matches the default 7-student + 1-teacher split (TRAIN_PROMPT_BSZ=21 stays divisible by 7 agent workers). ASCEND_RT_VISIBLE_DEVICES can still be overridden from the launching environment.
…just PLACEHOLDER count After the target-weight sleep fix, generation is healthy but OPD still records zero rejections (opd/rollout/empty_reject_ratio=1.0, loss=0). The engine patch inferred num_accepted purely from the PLACEHOLDER_TOKEN_ID padding of sampled_token_ids (n_accepted = n_valid - 1). That relies on stock vLLM v1's rejected-tail padding contract, which vllm-ascend's AscendRejectionSampler does not follow: it returns a fully-filled row, so the placeholder count reports full acceptance every step and every rejection anchor is dropped. Add compute_dflash_num_accepted(), which also uses the matching-prefix signal (an accepted speculative token equals the sampled token at that position; the first mismatch is the rejection point) and takes min() of the two signals so it never over-reports acceptance and never regresses the stock/GPU path. Also log per-row accepted counts under OPD_PATCH_DEBUG. Adds CPU tests covering stock-placeholder and ascend-no-placeholder cases. Co-authored-by: curnane-lab <curnane-lab@users.noreply.github.com>
logits_absmax was always inf because compute_logits() masks the padded/OOV vocab columns with -inf (monkey_patch_compute_logits), so it never reflected forward health. Report the finite-only abs max instead, which is nan only when the whole forward is genuinely non-finite. Co-authored-by: curnane-lab <curnane-lab@users.noreply.github.com>
The dflash rollout path is validated end to end (healthy verify events, real rejection anchors, finite logits). Eager mode costs 2-5x decode time on NPU, so flip the default to graph mode; eager remains available via the env override for bring-up/debugging.
…eights) The HYBRID sleep/wake cycle offloads and restores the full weight set between NPU and host RAM on every step (sleep level=1 since f113a95), which is a multi-second per-step cost per engine. With the 4B student FSDP-sharded over the worker cards, rollout and training memory typically coexist, so keep weights resident by default; flip back to True if update_actor OOMs.
…nt engine Graph capture with the verl default (FULL_AND_PIECEWISE, ~35 sizes) fails with AclmdlRICaptureBegin 207005 (driver stream/queue resource exhaustion) when the student engine colocates with the FSDP trainer. Default to PIECEWISE-only with 4 capture sizes, override via VLLM_COMPILATION_CONFIG. The JSON is quote-escaped so it survives as a single hydra string override (raw JSON fails the override grammar).
…e OOM Resident rollout weights (~27G) plus training state (~28G) leave no room for the 8.5G backward allocation on a 61G card when gradient checkpointing is off. Sleep level=1 is the correct colocate path; document that resident weights require gradient checkpointing (or a smaller rollout memory fraction) to fit.
STUDENT_WORLD_SIZE=14 + TEACHER_WORLD_SIZE=2 matches the 16-card layout. TRAIN_PROMPT_BSZ defaults to 28 so the prompt batch stays evenly divisible by the 14 agent loop workers (21 crashed the equal chunk assertion with 14 workers).
…, kept rejects) Rollout verify events show real rejections in OPD_PATCH_DEBUG while training metrics report empty_reject_ratio=1.0, so the break is downstream of the registry. Log per finished request: whether pop() returned metadata, the raw anchor count and how many rejects survive build_dflash_extra_fields.
The vLLMHttpServer logger (logging.getLogger(__file__)) never surfaces through this actor's logging setup, so the collect-path traces would be invisible like the earlier 'DFLASH OPD rollout' lines. Print to stderr instead so ray captures and forwards them.
Rollout-side recording works (rejected_events non-empty) but training metrics still show empty_reject_ratio=1.0 / loss=0. Root cause: the engine records OPD reject events keyed by model_runner.input_batch.req_ids, which vLLM suffixes as '<base>-<tag>' (e.g. uuid4().hex + '-b50b54a4'), while the frontend pops with the base request id it passed to engine.generate(). The mismatch makes pop() return None, so _collect_dflash_opd_fields falls back to empty_dflash_extra_fields and every sample looks like an empty reject. Make the registry lookup suffix-tolerant: match the base id exactly, else fall back to a '<base>-' prefix match (the base id is unique per request/turn). Adds CPU tests for the suffixed-id and exact-vs-suffix cases. Co-authored-by: curnane-lab <curnane-lab@users.noreply.github.com>
The first real OPD micro-batch crashes in FlexAttention: transformers validates the device and raises ValueError (CUDA/CPU/HPU only). Two layered fixes: - dflash_student: the runtime fallback around the flex draft forward caught only RuntimeError; catch ValueError too so an unsupported device retries the micro-batch with sdpa instead of crashing. - launchers: default DFLASH_ATTENTION_IMPL=sdpa for the vllm NPU script (exported so the exec'd base script sees it), and make the base script's flex_attention default env-overridable.
…e_anchors Uncapped rejection counts (~450 mean, >1k on 4096-token responses at the current acceptance rate) make the trainer's OPD forward process context + anchors*draft_block_size tokens, OOMing the 61G card in update_actor. The cap evenly subsamples anchor/boundary pairs per sample, always keeping the last pair so the tail segment reaches the response end; dropped pairs just merge segments already capped at draft_block_size-1. The vllm NPU launcher defaults the cap to 186 (DFLASH_MAX_RESPONSE_ANCHORS, 0 disables).
…ncher Even with the 186-anchor cap, the trainer process holds ~59G during the composed-student OPD forward (uncached 4609-token target activations, draft segments, vocab-sized loss tensors) and OOMs the 61G card in update_actor. The composed student cascades gradient_checkpointing_enable to its target and draft submodules, cutting activation memory at a ~20-30% update_actor compute cost. GRADIENT_CHECKPOINTING=False opts out.
The dflash speculative method is upstreamed in both projects, so track the official repos instead of the curnane-lab forks: - vllm -> vllm-project/vllm @ v0.21.0 (ad7125a431) - vllm-ascend -> vllm-project/vllm-ascend @ releases/v0.21.0rc (80610e44) README updated to match.
Keep this PR scoped to the vLLM rollout path; the installer and contributing-doc changes are dropped from the diff.
Keep the submodule clone step, the no-extras warning, the engine auto-skip behavior and the NPU trainer entrypoint; drop the opt-in flags, CANN/triton details and duplicated notes.
install_npu.sh is the NPU + vLLM installer (skips pre-installed engines, no verl extras); install.sh stays the standard SGLang-path installer. Keeps every file in the diff referenced and the text accurate.
curnane-lab
changed the title
Rewrites the README diff to pure additions: the original install text is untouched; the submodule clone hint and the NPU installer note become a single appended paragraph.
Collaborator
|
it seems good, but can you give a wandb metric show that accept length is improved in training? just like this: |
…length) Mean tokens advanced per verify step (accepted prefix + bonus token), matching sglang's spec_accept_length semantics: response_len / num_verify_steps per request, aggregated as opd/rollout/accept_length/mean. Rises as the draft distills toward the target, giving a first-class wandb signal for acceptance improvement during OPD training.
Author
Thanks for the quick review! Evidence from the current run (see the attached figure): mean accept length per verify step, derived from per-request OPD metadata sampling (accepted prefix + bonus token per verify step, K7). It rises steadily over the run, consistent with Metric now first-class: Longer runs are coming: several more experiments will be launched in the next few days much longer horizons (5k–6k steps, similar to the setup in your example), including the acceptance-length curve above as a native wandb metric. Updated figures will be posted here as soon as those trainings complete. |
… validation The reviewer's reference panels (test_rollout/<dataset>/sglang_spec_accept_length) aggregate sglang-specific output fields that the vLLM path never emits, so the val loop produced no accept-length curve on the vLLM path. Export dflash_num_verify_steps alongside dflash_accept_length and aggregate sum(response tokens)/sum(verify steps) per benchmark as test_rollout/<benchmark>/vllm_spec_accept_length (token-weighted, same semantics as the sglang panel).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a vLLM / vLLM-Ascend rollout engine path for DFlash on-policy distillation, as an alternative to the existing SGLang path (which is untouched). Validated end to end on Ascend NPU. Opened against the
dev_vllmstaging branch as requested in #7.Engine versions (pinned as git submodules)
vllm-project/vllm@ad7125a431e176d4161099480a66f0169609a690(tagv0.21.0)vllm-project/vllm-ascend@80610e4438dba05011b05f89fc45d91e96992671(branchreleases/v0.21.0rc, tagv0.21.0rc1)dflashspeculative method is upstreamed in both; no fork patches are required.Environment
quay.nju.edu.cn/ascend/vllm-ascend:v0.21.0rc1-a3Qwen3-4B, draftQwen3-4B-DFlash-b16(block_size 16),num_speculative_tokens=7v0.21.0rc1,triton-ascend==3.2.1Install & reproduce
Key launcher defaults (all env-overridable): 16 cards (14 student + 2 teacher), graph mode with PIECEWISE capture
[8,16,32,64], HYBRID sleep/wake, sdpa draft attention, OPD anchor cap 186, gradient checkpointing on.Design & compatibility
ROLLOUT_NAME); the vLLM path is additive and shares the OPD trainer (dflash_student, losses, anchor plan).vLLMHttpServerper student worker (14 replicas, TP=1, AsyncLLM). OPD reject metadata: monkey-patchedNPUModelRunner._samplerecords verify events into a per-process registry → on request finish the server pops them viacollective_rpc→dflash_*extra fields → OPD anchor plan in the composed student.dflash_draft_only); the frozen target is loaded once at engine init. HYBRID sleep uses level 1 for the composed student so the frozen target survives sleep/wake.Tests
CPU unit tests (no NPU needed):
Covers: registry position replay, engine-suffixed request-id matching, acceptance detection for both sampler output formats, and the
dflash_*extra-fields contract.Validation (16x NPU, ~550 steps / 55 logged points, first OPD end-to-end run)
Training dashboard (from the run's wandb history):
Readout of the curves:
opd/rollout/reject_token_count/mean466 → ~360 and still declining — the OPD signal is being consumed as designed.actor/distillation/loss1.94 → ~1.2;forward_kl_loss0.85 → ~0.7;rejected_draft_reverse_kl_loss8.2 → 6.5 withrejected_draft_loss_clamp_fraction0.71 → 0.49.grad_normsettles from early spikes (≤0.28) to ~0.065;empty_reject_ratioandopd_skipped_sample_ratiostay at 0.0 for the whole run;opd_effective_token_ratio~0.3; anchors ~120/sample.response_length/clip_ratio~0.19 (EOS restored; was 1.0 in the broken state),response_length/mean~1400 → 1500.perf/throughput40 → 46–48;timing_s/gen~430 → ~380s as acceptance rises.OPD_PATCH_DEBUGsampling).run-*.wandbor linkCloses #7.