diff --git a/docker/npu_patch/megatron-bridge.patch b/docker/npu_patch/megatron-bridge.patch index 44f206ce..1d313830 100644 --- a/docker/npu_patch/megatron-bridge.patch +++ b/docker/npu_patch/megatron-bridge.patch @@ -5,6 +5,7 @@ index a0421273..7f9203ab 100644 @@ -1097,15 +1097,19 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]): "LinearCrossEntropyModule", "TEColumnParallelLinear", ++ "MindSpeedTEColumnParallelLinear", "TELayerNormColumnParallelLinear", + "MindSpeedTELayerNormColumnParallelLinear", "TEColumnParallelGroupedLinear", @@ -298,3 +299,77 @@ index 60cc091c..1dd5c808 100644 from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_decoder_block_spec, get_gpt_layer_local_spec, +diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py +index 7775c11c..c7b094cf 100644 +--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py ++++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py +@@ -26,6 +26,7 @@ from megatron.core.packed_seq_params import PackedSeqParams + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.transformer import MegatronModule + from megatron.core.transformer.spec_utils import ModuleSpec ++from megatron.core.utils import nvtx_range_pop, nvtx_range_push + from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig as Qwen3VLConfigHF + + from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention +@@ -294,7 +295,7 @@ class Qwen3VLModel(MegatronModule): + # position ids is computed within the model + position_ids = None + +- torch.cuda.nvtx.range_push("Qwen3VLModel.forward.pre_process") ++ nvtx_range_push(msg="Qwen3VLModel.forward.pre_process") + + cp_rank = self.pg_collection.cp.rank() + cp_size = self.pg_collection.cp.size() +@@ -387,7 +388,7 @@ class Qwen3VLModel(MegatronModule): + combined_embeddings = split_data_cp_rank(combined_embeddings, cp_size, 0, cp_rank) + if packed_seq_params is not None: + if attention_mask is None: +- attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device) ++ attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device) + input_ids_thd, _ = preprocess_packed_seqs( + input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection + ) +@@ -443,7 +444,7 @@ class Qwen3VLModel(MegatronModule): + # convert lm_input_ids to THD format so it matches position_ids. + if packed_seq_params is not None: + if attention_mask is None: +- attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device) ++ attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device) + lm_input_ids, _ = preprocess_packed_seqs( + input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection + ) +@@ -499,8 +500,8 @@ class Qwen3VLModel(MegatronModule): + attention_mask = None + self.language_model.rotary_pos_emb.is_thd_format = True + +- torch.cuda.nvtx.range_pop() +- torch.cuda.nvtx.range_push("Qwen3VLModel.forward.language_model") ++ nvtx_range_pop(msg="Qwen3VLModel.forward.pre_process") ++ nvtx_range_push(msg="Qwen3VLModel.forward.language_model") + + output = self.language_model( + input_ids=lm_input_ids, +@@ -516,6 +517,6 @@ class Qwen3VLModel(MegatronModule): + **(extra_block_kwargs or {}), + **kwargs, + ) +- torch.cuda.nvtx.range_pop() ++ nvtx_range_pop(msg="Qwen3VLModel.forward.language_model") + + return output +diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py +index b714d0f7..42fcab74 100644 +--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py ++++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py +@@ -668,6 +668,11 @@ def preprocess_packed_seqs( + """ + batch_size = input_ids.shape[0] + ++ # Ensure boolean dtype for correct advanced indexing (bool → mask select, ++ # int → fancy index which silently corrupts data when values are 0/1). ++ if attention_mask.dtype != torch.bool: ++ attention_mask = attention_mask.bool() ++ + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + if pg_collection is not None: + tp_size = pg_collection.tp.size() diff --git a/docker/npu_patch/mindspeed.patch b/docker/npu_patch/mindspeed.patch index 6093e0ce..8bf86eb9 100644 --- a/docker/npu_patch/mindspeed.patch +++ b/docker/npu_patch/mindspeed.patch @@ -98,7 +98,7 @@ index f14b231d..4615590e 100644 setattr(full_args, k, v) MindSpeedFeaturesManager.apply_features_pre_patches(full_args) diff --git a/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py b/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py -index ac4eabe5..78c5866d 100644 +index ac4eabe5..061377cc 100644 --- a/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -330,6 +330,8 @@ class DotProductAttention(torch.nn.Module): @@ -110,18 +110,52 @@ index ac4eabe5..78c5866d 100644 ) -> torch.Tensor: """ Dot Product Attention Layer. -@@ -659,7 +661,9 @@ class MindSpeedTEDotProductAttention(DotProductAttention): +@@ -653,18 +655,39 @@ class MindSpeedTEDotProductAttention(DotProductAttention): + packed_seq_params: PackedSeqParams = None, + ): + """Forward.""" ++ packed_seq_kwargs = ( ++ {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params} ++ if packed_seq_params is not None ++ else {} ++ ) ++ ++ # Honor the per-call attn_mask_type. Megatron passes AttnMaskType.no_mask for the ++ # (bidirectional) vision tower and AttnMaskType.causal for the causal LM decoder. ++ # The previous code unconditionally built a triu(2048) causal mask and forwarded ++ # config.attention_mask_type, forcing CAUSAL attention even for no_mask -> the ++ # Qwen3-VL vision tower's full attention was computed causally (wrong image embeds). ++ if attn_mask_type == AttnMaskType.no_mask: ++ # Full (bidirectional) attention: no causal mask; per-image segmentation is ++ # handled by cu_seqlens in packed_seq_params (sparse_mode=0 via 'no_mask'). ++ core_attn_out = super().forward( ++ query, ++ key, ++ value, ++ None, ++ attn_mask_type='no_mask', ++ **packed_seq_kwargs, ++ ) ++ return core_attn_out ++ + if ( + attention_mask is None and + self.attn_mask_type == AttnMaskType.causal ) and not getattr(self.config, 'is_llava', False): self.config.sparse_mode = 2 attention_mask = get_attention_mask(self.config) - +- packed_seq_kwargs = ( +- {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params} +- if packed_seq_params is not None +- else {} +- ) + attention_mask = torch.triu( + torch.ones((2048, 2048), -+ device=query.device, dtype=torch.bool), diagonal=1) - packed_seq_kwargs = ( - {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params} - if packed_seq_params is not None - ++ device=query.device, dtype=torch.bool), diagonal=1) + + core_attn_out = super().forward( + query, diff --git a/mindspeed/core/fusions/fused_rope.py b/mindspeed/core/fusions/fused_rope.py index 70f7cb08..15189c2d 100644 --- a/mindspeed/core/fusions/fused_rope.py diff --git a/examples/geo3k_vlm/run_geo3k_vlm_npu.sh b/examples/geo3k_vlm/run_geo3k_vlm_npu.sh new file mode 100644 index 00000000..08c8aa77 --- /dev/null +++ b/examples/geo3k_vlm/run_geo3k_vlm_npu.sh @@ -0,0 +1,256 @@ +#!/bin/bash + +# Single-turn Qwen3-VL GRPO on geo3k for Ascend NPU. +# Self-contained: mirrors run_geo3k_vlm.sh but inlines the NPU handling that the +# multi-turn NPU example gets via execute_train_npu() (Ascend toolkit env, HCCL +# ports, sglang/slime cleanup, ray start without --num-gpus). No .py orchestrator. +# Usage: +# VIME_SCRIPT_MODEL_NAME=Qwen3-VL-8B-Instruct ./run_geo3k_vlm_npu.sh + +# Configuration +TRAIN_BACKEND="megatron" +MODEL_NAME=${VIME_SCRIPT_MODEL_NAME:-"Qwen3-VL-8B-Instruct"} +DATASET_NAME=${VIME_SCRIPT_DATASET_NAME:-"chenhegu/geo3k_imgurl"} +NUM_GPUS=${VIME_SCRIPT_NUM_GPUS:-8} +DATA_ROOT="/root/datasets/$(basename "$DATASET_NAME")" +MODEL_ROOT="/root/models/${MODEL_NAME}" + +# Validate MODEL_NAME (Qwen3-VL only, matches the NPU .py orchestrator) +VALID_MODELS=" + Qwen3-VL-2B-Instruct + Qwen3-VL-4B-Instruct + Qwen3-VL-8B-Instruct + Qwen3-VL-2B-Thinking + Qwen3-VL-4B-Thinking + Qwen3-VL-8B-Thinking +" +if ! echo "$VALID_MODELS" | grep -qw "$MODEL_NAME"; then + echo "Error: MODEL_NAME must be one of: $VALID_MODELS" + exit 1 +fi + +MODEL_NAME_LOWER=$(echo "$MODEL_NAME" | tr '[:upper:]' '[:lower:]') + +# External Ray flag +if [ -z "$VIME_SCRIPT_EXTERNAL_RAY" ] || [ "$VIME_SCRIPT_EXTERNAL_RAY" = "0" ]; then + USE_EXTERNAL_RAY=0 +else + USE_EXTERNAL_RAY=1 +fi + +# --------------------------------------------------------------------------- +# NPU environment (from examples/geo3k_vlm_multi_turn/run_grpo_npu.sh wrapper) +# --------------------------------------------------------------------------- +export PYTORCH_ALLOC_CONF=expandable_segments:True +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:$PYTHONPATH" +export ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15} +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 +export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 +export HYDRA_FULL_ERROR=1 +export MASTER_PORT=${MASTER_PORT:-$(shuf -i 20000-65000 -n 1)} +export VLLM_ASCEND_ENABLE_NZ=0 + +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY ALL_PROXY all_proxy + +sleep 3 +if [ "$USE_EXTERNAL_RAY" = "0" ]; then + ray stop --force + pkill -9 ray +fi +pkill -9 vime +sleep 3 +if [ "$USE_EXTERNAL_RAY" = "0" ]; then + pkill -9 ray +fi +pkill -9 vime +pkill -9 redis + +set -ex + +export PYTHONUNBUFFERED=1 + +# Download model and dataset if missing +mkdir -p /root/models /root/datasets +if [ ! -d "${MODEL_ROOT}" ]; then + hf download Qwen/${MODEL_NAME} --local-dir ${MODEL_ROOT} +fi +if [ ! -d "${DATA_ROOT}" ]; then + hf download --repo-type dataset ${DATASET_NAME} --local-dir ${DATA_ROOT} +fi + +# Common args +CKPT_ARGS=( + --hf-checkpoint ${MODEL_ROOT} + # qwen3 vl model has rotary base 5000000 + --rotary-base 5000000 +) + +# Single-turn rollout (no multi-turn custom-generate-function / custom-config) +ROLLOUT_ARGS=( + --prompt-data ${DATA_ROOT}/train.parquet + --input-key problem + --label-key answer + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout 200 + --rollout-batch-size 64 + --n-samples-per-prompt 8 + --rollout-max-response-len 4096 + --rollout-temperature 0.8 + --global-batch-size 512 +) + +# required for vlm datasets +MULTIMODAL_KEYS='{"image": "images"}' + +# Eval disabled for NPU bring-up (minimize moving parts; eval is an extra full +# rollout pass that re-exposes the same NPU rollout bugs). test.parquet exists, +# so re-enable this block + the ${EVAL_ARGS[@]} line below once training is stable. +# EVAL_ARGS=( +# --eval-interval 20 +# --eval-prompt-data geo3k_imgurl_processed ${DATA_ROOT}/test.parquet +# --n-samples-per-eval-prompt 1 +# --eval-max-response-len 4096 +# ) + +# Note: no reg anchor (kl-coef 0 + entropy-coef 0). The multi-turn NPU run +# diverged ~step128 / NaN-crashed ~step167 this way, but that's well past this +# short bring-up run (num-rollout 200). Add --kl-coef 0.005 for longer runs. +GRPO_ARGS=( + --advantage-estimator grpo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 1 + --vllm-gpu-memory-utilization ${VLLM_GPU_MEMORY_UTILIZATION:-0.9} + --vllm-max-model-len 16384 +) + +# Wandb args (only if WANDB_API_KEY is set) +if [ -n "$WANDB_API_KEY" ]; then + WANDB_ARGS=( + --use-wandb + --wandb-project vime-geo3k-vlm + --wandb-group ${MODEL_NAME_LOWER}-${TRAIN_BACKEND}-vllm-${NUM_GPUS}npu + --wandb-key ${WANDB_API_KEY} + --disable-wandb-random-suffix + ) +else + WANDB_ARGS=() +fi + +# megatron backend (NPU-tuned) +BACKEND_ARGS=( + --train-backend megatron + --load ${MODEL_ROOT} + --ref-load ${MODEL_ROOT} + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + # P8 workaround: bypass broken THD/varlen attention by using bshd. bshd + # asserts use_dynamic_batch_size is False -> fixed micro-batch (no + # --use-dynamic-batch-size/--max-tokens-per-gpu/--balance-data). + --qkv-format bshd + --micro-batch-size 1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --megatron-to-hf-mode bridge +) + +MISC_ARGS=( + --no-gradient-accumulation-fusion + --use-flash-attn + # Checkpoint planning (disk-aware): + # /workspace (=/dev/sda2) is 96% full (~19G) -> cannot hold even one ckpt. + # /root/.cache is a ~35T NFS share with ~1.9T free -> save there instead. + --save /root/.cache/vime-ckpt/geo3k_vlm_single_turn_npu + # Test run: weights only (~16GB/ckpt) instead of full state (~100GB with + # optimizer). Drop --no-save-optim if you need to resume the optimizer. + --no-save-optim + # No auto-prune of old ckpts. At ~16GB each, keep the interval large and + # clean /root/.cache/vime-ckpt manually if it grows. + --save-interval 50 +) + +# get MODEL_ARGS from scripts/models for megatron backend +VIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)" +MODEL_ARGS_FILE=$(echo "$MODEL_NAME" | sed 's/-Instruct//g; s/-Thinking//g; s/Qwen3-VL-/qwen3-/g; s/-2B/-1.7B/g') +# VL models require rotary-base 5000000 +MODEL_ARGS_ROTARY_BASE=5000000 source "${VIME_DIR}/scripts/models/${MODEL_ARGS_FILE}.sh" + +# Start Ray if not using external Ray +if [ "$USE_EXTERNAL_RAY" = "0" ]; then + export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} + export no_proxy="127.0.0.1,${MASTER_ADDR}" + ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +fi + +# Build runtime env +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES\": \"1\", + \"PYTHONPATH\": \"${PYTHONPATH}\", + \"PYTORCH_ALLOC_CONF\": \"expandable_segments:True\", + \"PYTORCH_NPU_ALLOC_CONF\": \"expandable_segments:True\", + \"VLLM_ASCEND_ENABLE_NZ\": \"0\", + \"ASCEND_TOOLKIT_HOME\": \"/usr/local/Ascend/ascend-toolkit/latest/\", + \"ASCEND_OPP_PATH\": \"/usr/local/Ascend/ascend-toolkit/latest/opp/\", + \"ASCEND_AICPU_PATH\": \"/usr/local/Ascend/ascend-toolkit/latest/\", + \"ASCEND_HOME_PATH\": \"/usr/local/Ascend/ascend-toolkit/latest/\", + \"set_env_path\": \"/usr/local/Ascend/nnal/atb/set_env.sh\", + \"HYDRA_FULL_ERROR\": \"1\", + \"HCCL_HOST_SOCKET_PORT_RANGE\": \"60000-60050\", + \"HCCL_NPU_SOCKET_PORT_RANGE\": \"61000-61050\", + \"no_proxy\": \"127.0.0.1,${MASTER_ADDR}\", + \"MASTER_ADDR\": \"${MASTER_ADDR}\" + } +}" + +LOG_FILE="/root/${MODEL_NAME}-single-turn-${DATE}.log" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node ${NUM_GPUS} \ + --rollout-num-gpus ${NUM_GPUS} \ + --multimodal-keys "${MULTIMODAL_KEYS}" \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${BACKEND_ARGS[@]} \ + ${MISC_ARGS[@]} \ + 2>&1 | tee -a "$LOG_FILE" diff --git a/examples/geo3k_vlm_multi_turn/rollout.py b/examples/geo3k_vlm_multi_turn/rollout.py index eccd2819..7af4a18e 100644 --- a/examples/geo3k_vlm_multi_turn/rollout.py +++ b/examples/geo3k_vlm_multi_turn/rollout.py @@ -3,6 +3,7 @@ import importlib import importlib.util import json +import re import sys import uuid from pathlib import Path @@ -74,7 +75,16 @@ def _build_initial_messages(sample: Sample) -> list[dict]: for image in images: content.append({"type": "image", "image": image}) # Backward-compatible fallback for old runs that pass a raw text prompt containing . - text_prompt = str(sample.prompt).replace("", "").lstrip() + # The image is re-attached above as a structured {"type": "image"} part, so render expands it + # exactly once. We must therefore strip BOTH placeholder forms from the text, otherwise render + # passes any literal placeholder through verbatim and we get a double placeholder (the +1 stray + # <|image_pad|> render "bug"): the raw `` marker AND the chat-template-expanded + # `<|vision_start|><|image_pad|>...<|vision_end|>` (present when sample.prompt was already run + # through tokenizer.apply_chat_template). See _validate_multimodal_train_inputs / docs P3/4. + text_prompt = str(sample.prompt).replace("", "") + text_prompt = re.sub( + r"<\|vision_start\|>(?:<\|image_pad\|>)+<\|vision_end\|>", "", text_prompt + ).lstrip() content.append({"type": "text", "text": text_prompt}) return [{"role": "user", "content": content}] @@ -150,13 +160,84 @@ def _validate_multimodal_train_inputs(sample: Sample, tokenizer: Any, processor: grid = grid.reshape(-1, 3) image_processor = getattr(processor, "image_processor", processor) merge_size = int(getattr(image_processor, "merge_size", 1) or 1) + # Megatron's Qwen3-VL vision tower emits exactly prod(grid)//merge**2 embeddings and + # reorganize_inputs() hard-asserts (#image_pad)*merge**2 == prod(grid), so the rendered + # prompt MUST carry exactly this many <|image_pad|> tokens or training crashes. This is a + # permanent fail-fast guard (NOT a stopgap): it loudly catches any rollout/train placeholder + # mismatch — e.g. a regression that double-sends the image placeholder to render (a literal + # <|image_pad|> in the text PLUS a separate image_url => N+1, the old Problem 3/4) — instead + # of silently corrupting the step. See QWEN3VL_P34_EXPLAINER.md. expected = int((grid.prod(dim=1) // (merge_size * merge_size)).sum().item()) if image_tokens != expected: + # Report the contiguous <|image_pad|> run lengths so a mismatch is diagnosable + # (e.g. a double-send splitting the block as [target] + a stray separated placeholder). + runs: list[int] = [] + tid = int(image_token_id) + k = 0 + toks = sample.tokens + while k < len(toks): + if toks[k] == tid: + m = k + while m < len(toks) and toks[m] == tid: + m += 1 + runs.append(m - k) + k = m + else: + k += 1 raise RuntimeError( "Image token count does not match multimodal render features: " - f"image_tokens={image_tokens}, expected_image_tokens={expected}, image_grid_thw={grid.tolist()}" + f"image_tokens={image_tokens}, expected_image_tokens={expected}, image_grid_thw={grid.tolist()}, " + f"image_pad_run_lengths={runs}" + ) + + +_DELTA_SENTINEL = "\x00vime_assistant_sentinel\x00" + + +def _observation_delta_token_ids( + tokenizer: Any, observation_message: dict, last_token_is_im_end: bool +) -> list[int]: + """Tokenize ONLY the new observation turn's template fragment, to append to the token stream. + + Why not re-render the whole conversation: re-tokenizing already-emitted turns triggers a BPE + re-merge at the ``<|im_start|>assistant\\n`` + generated-content seam (the model's first token + is often ``\\n\\n``; ``assistant\\n`` + ``\\n\\n`` re-tokenizes ``\\n\\n\\n`` into a single + token), which shifts every later token by one and silently desyncs train/infer token streams. + Keeping the incrementally-built stream as the single source of truth and only tokenizing the + new delta avoids that entirely (see QWEN3VL_NPU_GRPO_DEBUG.md Problem 6). + + The fragment is the assistant-turn closer + the observation user turn + the next generation + prompt, derived faithfully from the chat template via a sentinel so we never re-tokenize real + assistant content: + <|im_end|>\\n<|im_start|>user\\n{obs}<|im_end|>\\n<|im_start|>assistant\\n + When the model already emitted its closing ``<|im_end|>`` (the normal multi-turn case), we drop + the leading ``<|im_end|>`` so the stream carries exactly one (no double ``<|im_end|>``). The + fragment then begins with ``\\n`` immediately after an atomic special token, so tokenizing it in + isolation and concatenating is bit-identical to the canonical in-context tokenization. + """ + # The delta path only handles text observations: an image in the observation would need feature + # re-rendering (and image-pad expansion) that this incremental path deliberately avoids. + content = observation_message.get("content") + if isinstance(content, list) and any( + (part.get("type") in ("image", "video")) or ("image" in part) or ("image_url" in part) or ("video" in part) + for part in content + if isinstance(part, dict) + ): + raise NotImplementedError( + "Incremental multi-turn rollout only supports text observations; " + "an observation carried multimodal content." ) + probe = tokenizer.apply_chat_template( + [{"role": "assistant", "content": _DELTA_SENTINEL}, observation_message], + add_generation_prompt=True, + tokenize=False, + ) + fragment = probe.split(_DELTA_SENTINEL, 1)[1] + if last_token_is_im_end and fragment.startswith("<|im_end|>"): + fragment = fragment[len("<|im_end|>") :] + return _coerce_flat_int_token_ids(tokenizer.encode(fragment, add_special_tokens=False)) + async def generate(args: Any, sample: Sample, sampling_params) -> Sample: """Custom multi-turn rollout that interacts with a pluggable environment via the vLLM render route.""" @@ -224,35 +305,30 @@ def sampling_params_for_turn() -> dict | None: try: env.reset() - latest_features = None - pending_obs_offset: int | None = None - rendered_body = await render() - prompt_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids")) + # Render ONCE. The image lives only in the first user message and never moves, and + # observations are text, so sample.tokens[:len(prompt_ids)] (which carries the image-pad + # block + the render features' mm_placeholders) is never mutated. We therefore reuse this + # body's features for every turn and feed the incrementally-built sample.tokens directly as + # the prompt, instead of re-rendering the whole conversation each turn (which re-tokenizes + # the assistant seam and desyncs the stream — see QWEN3VL_NPU_GRPO_DEBUG.md Problem 6). + initial_body = await render() + latest_features = initial_body.get("features") + prompt_ids = _coerce_flat_int_token_ids(initial_body.get("token_ids")) if not sample.tokens: sample.tokens = list(prompt_ids) if args.rollout_max_context_len is not None: max_response_budget = max(0, args.rollout_max_context_len - len(sample.tokens)) - for turn_idx in range(args.max_turns): - input_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids")) - latest_features = rendered_body.get("features") - - if pending_obs_offset is not None: - obs_tokens = input_ids[pending_obs_offset:] - remaining = remaining_budget() - if remaining is not None and len(obs_tokens) > remaining: - append_response_window(obs_tokens[: max(remaining, 0)], [0] * max(remaining, 0)) - sample.status = Sample.Status.TRUNCATED - break - append_response_window(obs_tokens, [0] * len(obs_tokens)) - pending_obs_offset = None + eos_token_id = getattr(state.tokenizer, "eos_token_id", None) + for turn_idx in range(args.max_turns): current_sampling_params = sampling_params_for_turn() if current_sampling_params is None: sample.status = Sample.Status.TRUNCATED break - body = dict(rendered_body) + body = dict(initial_body) + body["token_ids"] = list(sample.tokens) body["sampling_params"] = current_sampling_params output = await post(f"{base_url}/inference/v1/generate", body, headers=headers) choice = output["choices"][0] @@ -269,7 +345,6 @@ def sampling_params_for_turn() -> dict | None: train_logprobs = list(new_logprobs) train_loss_mask = [1] * len(train_tokens) stop = current_sampling_params.get("stop") - eos_token_id = getattr(state.tokenizer, "eos_token_id", None) append_stop_eos = ( stop and eos_token_id is not None @@ -292,8 +367,6 @@ def sampling_params_for_turn() -> dict | None: append_response_window(train_tokens, train_loss_mask, train_logprobs) _apply_vllm_routed_experts(args, sample, choice) - messages.append({"role": "assistant", "content": response_text}) - if finish_reason == "length": sample.status = Sample.Status.TRUNCATED break @@ -310,24 +383,21 @@ def sampling_params_for_turn() -> dict | None: sample.status = Sample.Status.TRUNCATED break + # Append ONLY the new observation turn's tokens to the incremental stream (no re-render). next_user_message = env.format_observation(observation) - messages.append(next_user_message) - render_prefix_len = len(input_ids) + len(new_tokens) - pending_obs_offset = render_prefix_len - rendered_body = await render() - rendered_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids")) - is_prefix_stable = rendered_ids[:pending_obs_offset] == sample.tokens[:pending_obs_offset] + last_token_is_im_end = bool(sample.tokens and eos_token_id is not None and sample.tokens[-1] == eos_token_id) + obs_tokens = _observation_delta_token_ids(state.tokenizer, next_user_message, last_token_is_im_end) + remaining = remaining_budget() + if remaining is not None and len(obs_tokens) > remaining: + append_response_window(obs_tokens[: max(remaining, 0)], [0] * max(remaining, 0)) + sample.status = Sample.Status.TRUNCATED + break + append_response_window(obs_tokens, [0] * len(obs_tokens)) sample.metadata["multiturn_render"] = { - "prefix_stable": is_prefix_stable, - "prefix_len": pending_obs_offset, + "turns": turn_idx + 1, "sample_len": len(sample.tokens), - "rendered_len": len(rendered_ids), + "last_obs_tokens": len(obs_tokens), } - if not is_prefix_stable: - raise RuntimeError( - "Full conversation render is not prefix-stable with the generated token stream: " - f"{sample.metadata['multiturn_render']}" - ) multimodal_train_inputs = _multimodal_train_inputs_from_features(latest_features) _validate_multimodal_train_inputs(sample, state.tokenizer, state.processor, multimodal_train_inputs) diff --git a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py index 31cb94ef..0431638c 100644 --- a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py +++ b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py @@ -82,7 +82,7 @@ def execute(): "--use-precision-aware-optimizer " ) - vllm_args = "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.6 " + vllm_args = "--rollout-num-gpus-per-engine 1 " "--vllm-gpu-memory-utilization 0.6 " "--vllm-max-model-len 16384 " megatron_args = ( "--train-backend megatron " @@ -97,9 +97,11 @@ def execute(): "--recompute-granularity full " "--recompute-method uniform " "--recompute-num-layers 1 " - "--use-dynamic-batch-size " - "--max-tokens-per-gpu 16384 " - "--balance-data " + # P8 workaround: bypass the broken THD/varlen attention path by using bshd. + # bshd asserts use_dynamic_batch_size is False, so use a fixed micro-batch + # instead of --use-dynamic-batch-size/--max-tokens-per-gpu/--balance-data. + "--qkv-format bshd " + "--micro-batch-size 1 " "--attention-dropout 0.0 " "--hidden-dropout 0.0 " "--accumulate-allreduce-grads-in-fp32 " diff --git a/vime/rollout/vllm_rollout.py b/vime/rollout/vllm_rollout.py index 3bf9b4cd..57713be7 100644 --- a/vime/rollout/vllm_rollout.py +++ b/vime/rollout/vllm_rollout.py @@ -57,6 +57,85 @@ def _coerce_flat_int_token_ids(ids: Any) -> list[int]: return [int(x)] +def _decode_vllm_routed_experts(value: str) -> np.ndarray: + raw = base64.b64decode(value.encode("ascii"), validate=True) + return np.load(io.BytesIO(raw), allow_pickle=False) + + +def _apply_vllm_routed_experts( + args: Namespace, + sample: Sample, + choice: dict, +) -> None: + """Populate ``sample.rollout_routed_experts`` from vLLM ``/inference/v1/generate`` (R3 only). + + vLLM's contract is a single base64 `.npy` buffer on `choices[].routed_experts` with decoded + shape `(len(tokens) - 1, num_layers, top_k)`. + """ + if not getattr(args, "use_rollout_routing_replay", False): + return + + routed = choice.get("routed_experts") + if sample.status == Sample.Status.ABORTED and sample.response_length == 0: + return + if routed is None: + raise RuntimeError( + "vLLM routing replay: missing choices[0].routed_experts on /inference/v1/generate response. " + "Check vLLM 0.22+ was launched with --enable-return-routed-experts." + ) + if not isinstance(routed, str): + raise RuntimeError( + f"vLLM routing replay: choices[0].routed_experts must be base64 npy str, got {type(routed)}" + ) + + arr = _decode_vllm_routed_experts(routed) + if arr.ndim != 3: + raise RuntimeError(f"vLLM routing replay: routed_experts ndim={arr.ndim}, expected 3, shape={arr.shape}") + + expected_rows = max(0, len(sample.tokens) - 1) + if arr.shape[0] != expected_rows: + raise RuntimeError( + f"vLLM routing replay: routed_experts rows {arr.shape[0]} != expected {expected_rows} (len(tokens)-1)." + ) + + nl = getattr(args, "num_layers", None) + mtk = getattr(args, "moe_router_topk", None) + if nl is not None and mtk is not None and (arr.shape[1] != nl or arr.shape[2] != mtk): + raise RuntimeError(f"vLLM routing replay: routed_experts shape {arr.shape} != (rows,{nl},{mtk}) from args.") + sample.rollout_routed_experts = np.ascontiguousarray(arr.astype(np.int32, copy=True)) + + +def _inference_generate_tokens_and_logprobs(choice: dict[str, Any]) -> tuple[list[int], list[float]]: + """Parse ``token_ids`` and ``logprobs.content`` from a vLLM ``/inference/v1/generate`` choice.""" + tids_raw = choice.get("token_ids") + if isinstance(tids_raw, list) and tids_raw and all(isinstance(x, int) for x in tids_raw): + tids = [int(x) for x in tids_raw] + lps: list[float] = [] + lp = choice.get("logprobs") + if not isinstance(lp, dict): + return tids, [0.0] * len(tids) + + content = lp.get("content") + if isinstance(content, list) and len(content) == len(tids): + for item in content: + if isinstance(item, dict): + lps.append(float(item.get("logprob", 0.0))) + else: + lps.append(0.0) + return tids, lps + if isinstance(content, list) and content: + # Partial content: pad / truncate to token_ids length if possible. + for i in range(len(tids)): + if i < len(content) and isinstance(content[i], dict): + lps.append(float(content[i].get("logprob", 0.0))) + else: + lps.append(0.0) + return tids, lps + + return tids, [0.0] * len(tids) + + return [], [] + def _prepare_prompt_ids(sample: Sample, tokenizer, processor: Any) -> list[int]: raw_multimodal_inputs = sample.multimodal_inputs or {} has_multimodal_inputs = any(value is not None for value in raw_multimodal_inputs.values())