Skip to content

fix(vlm): forward image payloads in async rollouts - #4113

Merged
terrykong merged 6 commits into
NVIDIA-NeMo:mainfrom
jinglinglingling:linglinj/fix-async-vlm-image-payload
Sep 16, 2026
Merged

terrykong merged 6 commits into
NVIDIA-NeMo:mainfrom
jinglinglingling:linglinj/fix-async-vlm-image-payload

Conversation

@jinglinglingling

@jinglinglingling jinglinglingling commented Sep 12, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Preserves native multimodal side-channel fields when SingleController dispatches async VLM rollouts, so vLLM and the training forward receive the same image input.

The async per-sample path previously rebuilt generation input from tokenized message data but dropped fields such as vllm_content and vllm_images. This made vLLM generation image-blind while Megatron recomputed log probabilities with the image, producing large token_mult_prob_error values.

For multi-turn rollouts, later turns retain native media while clearing stale vllm_content, allowing the updated tokenized prefix to be used.

This PR does not modify recipes or generation lengths. The existing defaults remain Qwen max_new_tokens=1024 and Nemotron max_new_tokens=4096.

Validation results:

Reproduction commands

Qwen native-length 100-step validation:

uv run --locked examples/run_grpo_single_controller.py \
  --config examples/configs/recipes/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n8g-megatron-single-controller-async.v1.yaml \
  grpo.max_num_steps=100 \
  logger.wandb_enabled=true \
  logger.tensorboard_enabled=true \
  logger.monitor_gpus=false \
  logger.wandb.project=rohit-async_grpo_v2 \
  logger.wandb.name=linglinj-qwen-image-fix-native1024

Nemotron native-length 100-step validation (3 nodes / 8 GPUs per node):

uv run --locked examples/run_grpo_single_controller.py \
  --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-3n8g-megatron-tp8ep8-single-controller-async.v1.yaml \
  grpo.max_num_steps=100 \
  policy.train_global_batch_size=128 \
  logger.wandb_enabled=true \
  logger.tensorboard_enabled=true \
  logger.monitor_gpus=false \
  logger.wandb.project=rohit-async_grpo_v2 \
  logger.wandb.name=linglinj-nemotron-native4096-gbs128

Signed-off-by: Linglin Jing <50938792+jinglinglingling@users.noreply.github.com>
@jinglinglingling
jinglinglingling requested review from a team as code owners September 12, 2026 02:09
@copy-pr-bot

copy-pr-bot Bot commented Sep 12, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@jinglinglingling jinglinglingling changed the title fix(vlm): preserve native media in async rollouts fix(vlm): forward image payloads in async rollouts Sep 12, 2026
@jinglinglingling jinglinglingling added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Sep 12, 2026
@NVIDIA-NeMo NVIDIA-NeMo deleted a comment from copy-pr-bot Bot Sep 14, 2026
@jinglinglingling

Copy link
Copy Markdown
Contributor Author

/ok to test 2bb044e

Comment thread tests/unit/experience/test_rollout_manager.py
Comment thread nemo_rl/experience/rollout_manager.py
@rohitrango rohitrango added CI:L2 Run doctests, unit tests, functional tests, and convergence tests and removed CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) labels Sep 14, 2026
@rohitrango

Copy link
Copy Markdown
Contributor

/ok to test 167ad7b

Signed-off-by: rohitrango <rohit.rango@gmail.com>
@rohitrango

Copy link
Copy Markdown
Contributor

/ok to test 188c841

@jinglinglingling

Copy link
Copy Markdown
Contributor Author

Addressed the remaining review comment with a two-turn _run_single_rollout regression test in ec12bc64. The PR is labeled CI:L2; @rohitrango, could you approve the latest SHA with /ok to test ec12bc64 so the L2 suite can run again?

@rohitrango

Copy link
Copy Markdown
Contributor

could you sign your latest commit? once it passes DCO I will re-run

jinglinglingling and others added 2 commits September 14, 2026 19:50
Signed-off-by: Linglin Jing <50938792+jinglinglingling@users.noreply.github.com>
@jinglinglingling
jinglinglingling force-pushed the linglinj/fix-async-vlm-image-payload branch from ba19807 to 2037327 Compare September 15, 2026 02:51
@rohitrango

Copy link
Copy Markdown
Contributor

/ok to test 2037327

@rohitrango

Copy link
Copy Markdown
Contributor

/ok to test 6d89cac

@terrykong
terrykong enabled auto-merge (squash) September 16, 2026 03:44
@terrykong
terrykong merged commit 26f3b19 into NVIDIA-NeMo:main Sep 16, 2026
103 checks passed
seonjinn added a commit to seonjinn/RL that referenced this pull request Sep 16, 2026
NVIDIA-NeMo#3803 replaced NATIVE_MULTIMODAL_KEYS with VLLM_PROMPT_KEYS in
nemo_rl/data/multimodal_utils.py; NVIDIA-NeMo#4113 landed 37 minutes later still
importing the old name into nemo_rl/experience/rollout_manager.py. Both
merged clean on their own, so main carries an unresolvable import that
reds the pyrefly lint job on every open PR.

rollouts.py:1320 already builds the same per-sample multimodal dict from
VLLM_PROMPT_KEYS; this restores parity with it.

Signed-off-by: seonjinn <sna@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:L2 Run doctests, unit tests, functional tests, and convergence tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants