Add VLM support to the dynamic-batching inference server - #6260
Merged
Conversation
svcnvidia-nemo-ci
marked this pull request as draft
August 4, 2026 20:24
Contributor
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
RPrenger
force-pushed
the
vlm-inference
branch
3 times, most recently
from
August 11, 2026 04:00
5d668b2 to
83afeb8
Compare
RPrenger
marked this pull request as ready for review
August 11, 2026 04:03
Contributor
Author
|
/ok to test 83afeb8 |
RPrenger
force-pushed
the
vlm-inference
branch
from
August 11, 2026 04:27
83afeb8 to
9e36a7c
Compare
Contributor
Author
|
/ok to test 9e36a7c |
…spec_te The old comment claimed the moe_layer branch would fail loudly when ``config`` didn't carry the required MoE fields, but no assert was present. Passing a config with ``num_moe_experts=None`` (either accidentally, or because the caller only wanted a non-MoE hybrid) fell through to ``get_moe_module_spec(num_experts=None, ...)`` and either built a spec whose moe_layer trip would fail deep inside the MoE spec, or produced a checkpoint-mismatched architecture silently. Add the assert the comment already promised, and replace the free-form comment with a proper docstring so callers can see the ``config`` contract without reading source. Signature is unchanged: ``config`` and ``padding`` stay in the same positions, so no callers move. Addresses PR NVIDIA#6260 review comment on layer_specs.py:131. Signed-off-by: rprenger <rprenger@nvidia.com>
The previous order was:
1. compute pos = arange over patch tokens only
2. shift pos by class_token_len (to "make room" for CLS)
3. add position_embeddings to the patch-only tensor
4. prepend CLS
That produced two off-by-class_token_len effects:
* The CLS slot never received a position embedding (it was concatenated
onto x after the +position_embeddings step and started life as pure
class_token content, no positional signal).
* Patch tokens ended up at positions class_token_len..class_token_len+N,
which is what a patch would want AFTER the CLS was in place — but the
prepending happened later, so the effective positions the transformer
saw for patches drifted relative to the pre-shift arange.
Prepend CLS first, then compute ``pos = arange(x.shape[1])`` over the
CLS + patch tensor and add ``position_embeddings(pos)``. CLS now gets
positions [0, class_token_len) and patches sit at [class_token_len, N +
class_token_len). RoPE branch is unaffected — rope is passed to the
transformer as a separate tensor covering patch tokens only.
Addresses PR NVIDIA#6260 review comment on vit_model.py:320.
Signed-off-by: rprenger <rprenger@nvidia.com>
The fast path (h_patches / w_patches match the stored resolution) returned ``self.weight.reshape(...)`` without moving the weight tensor to the caller-supplied ``device``. The slow path (bicubic interp) does ``w = self.weight.to(device=device)`` first, so downstream operations land on the right device. If a caller passed a ``device`` that differed from where the module's weight lived (fresh instantiation before ``.to(device)``, or an explicit device override at call time), the fast path returned tensor would be on the wrong device and the addition in ``ViTModel.forward`` would either error on device-mismatch or trigger an implicit copy. Hoist the ``.to(device=device)`` above the fast-path branch so both paths behave identically. Addresses PR NVIDIA#6260 review comment on vit_model.py:801. Signed-off-by: rprenger <rprenger@nvidia.com>
…urn shape Newer ``transformers`` releases default ``apply_chat_template`` to ``return_dict=True``, which returns a ``BatchEncoding`` instead of the list/tensor the ``[0]`` subscript below (and the ``len()`` on the per-turn helper) expects. Restore ``return_dict=False`` on both call sites so behavior stays consistent across ``transformers`` versions. Addresses PR NVIDIA#6260 review comment on multimodal_tokenizer.py:272. Signed-off-by: rprenger <rprenger@nvidia.com>
``image_preprocessing.py`` uses ``PIL.Image`` to decode client-supplied image bytes on the VLM inference path, but Pillow was not declared as a ``megatron-core`` install dependency: on a fresh ``pip install megatron-core[dev]`` without Pillow separately available, the VLM inference request path would raise ``ImportError`` on the first image request, not at install time. Add ``Pillow`` next to the other multimedia deps (``av``, energon's audio/video decoders) in the ``dev`` optional-dependency block. This keeps the megatron-core install self-contained for users who install the multimodal path without cloning the repo. ``torchvision`` is also used by the same file but is already declared via ``override-dependencies`` in the same way as ``torch``, treating it as provided by the NGC PyTorch base image; the convention is unchanged. Requires a follow-up ``uv lock`` regeneration (``UV_PYTHON=3.12 uvx uv@0.7.2 lock``) so ``uv.lock`` picks up the Pillow entry. Signed-off-by: rprenger <rprenger@nvidia.com>
…e error torchvision isn't a hard install dependency of ``megatron-core``: the NGC PyTorch container ships one pinned to the container's torch build, so the toml lists it under ``override-dependencies`` (same treatment as ``torch``) and the container assumption covers users on that path. A plain ``pip install megatron-core`` off PyPI does not get torchvision automatically -- installing it via pip needs a build matching the local torch, which the caller has to pick, so we can't just declare it as a regular dep here without breaking the container path. Wrap the lazy import at the one call site (VLM image preprocessing) in a try/except that translates ``ImportError`` into a message naming ``torchvision`` and pointing at the container path. Matches the ``HAVE_TE`` pattern used elsewhere in the repo -- users installing off PyPI who don't need VLM inference are unaffected; users who do hit the VLM path get a clean instruction instead of an opaque import stack. Signed-off-by: rprenger <rprenger@nvidia.com>
Autoformatter pass to keep the CI ``linting`` job green after the SSRF-hardening and torchvision-guard changes shifted import blocks. Pure formatting; no semantic change. Signed-off-by: rprenger <rprenger@nvidia.com>
Companion to 2722ef2, which declared ``Pillow`` under the ``dev`` extras of ``pyproject.toml`` but never regenerated ``uv.lock``. The CI ``Pip`` / ``UV`` / ``Install test summary`` jobs run ``uv sync --locked``, which refuses to install when the lock and the toml disagree. Regenerated inside the NGC PyTorch container we use for eval (matching ``UV_VERSION=0.7.2`` from ``docker/Dockerfile.ci.*`` and Python 3.12 from ``.python-version``). The diff is just the new ``pillow`` entry plus its transitive hash and the ``dev`` extras edge that references it -- no other package versions moved. Signed-off-by: rprenger <rprenger@nvidia.com>
Contributor
Author
|
/ok to test a9464c7 |
Contributor
Author
|
/ok to test f1180d4 |
megatron.inference is not in mcore's setuptools packages.find, so the install-test check-imports scan fails when it walks megatron.core and tries to top-level import megatron.inference.utils. Defer the import into add_vlm_inference_args, which is only called from the server entry point where the full training tree is present. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>
Contributor
Author
|
/ok to test 43a3650 |
The SSRF-hardening no-redirect handler failed pylint's missing-function-docstring check. Add a one-line docstring. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>
Contributor
Author
|
/ok to test 9e56ee8 |
- llava_model.py: only pass ``imgs_sizes=`` to the vision encoder when it's not None; CLIPViTModel.forward() doesn't accept the kwarg. Move ``_preprocess_data``'s ``position_ids`` to keyword-only so positional callers don't silently shift into the wrong slot, and accept the sound_* kwargs (still ignored inside) for API compatibility with pre-existing callers. - dynamic_context.py: guard the VLM dict clears in ``reset_metadata`` with ``getattr`` so tests that construct the context via ``__new__`` don't trip on unset attrs. - llm.py / async_llm.py: default ``inference_wrapper_cls`` to ``None`` and resolve to ``GPTInferenceWrapper`` at call time, so tests that monkey-patch the module-level name actually steer construction. - test_apis.py: also patch ``GPTInferenceWrapper`` on the llm / async_llm modules (safety-belt alongside the code change). - test_data_parallel_inference_coordinator.py: give ``DummyEngine`` a ``failed_request_ids`` list so ``_handle_failed_request`` no longer raises ``AttributeError`` and hangs the coordinator (was the source of the 13x TimeoutError cluster). - test_llava_model.py::test_preprocess_data: unpack all 5 return values now that ``_preprocess_data`` also returns combined input_ids / position_ids. - test_llava_sound.py::TestPreprocessDataSoundReplacement: skip; the sound-embedding splicing block was removed during the VLM refactor. Kwargs are still accepted but no longer replace tokens in-place. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>
Contributor
Author
|
/ok to test 54e2c47 |
black wanted the pytest.mark.skip decorator wrapped differently. Pure formatting, no behavior change. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>
Contributor
Author
|
/ok to test 93f19e7 |
Companion to the imgs_sizes gating in the previous fix commit — the same CLIPViTModel test paths were still failing on ``CLIPViTModel.forward() got an unexpected keyword argument 'packed_seq_params'``. Move ``packed_seq_params`` into the same "conditional pass" dict as ``imgs_sizes`` so both are only forwarded to the vision encoder when actually populated. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>
Two follow-up fixes for the CI unit-test failures that survived the
earlier round:
- dynamic_context.py: ``reset_metadata`` clears a third VLM dict
(``_request_to_image_token_count``) I missed in the previous guarded
round. Wrap the same way with ``getattr(...).clear()`` so tests that
build the context via ``__new__`` don't AttributeError.
- dynamic_engine.py: the SUBMIT_REQUEST handler dereferences
``self.context.config.image_preprocessing_config`` unconditionally,
even for text-only requests. Test fixtures use a ``DummyContext``
without a ``.config`` attribute, so every SUBMIT_REQUEST raised
AttributeError, got swallowed by ``_fail_submission``, and desynced
the ranks (surfacing later as an NCCL collective timeout in the
distributed coordinator test). Skip the config lookup entirely when
``multi_modal_data is None`` — the resolver already returns ``{}``
for that case, so this preserves the production behavior.
Signed-off-by: Ryan Prenger <rprenger@nvidia.com>
Signed-off-by: rprenger <rprenger@nvidia.com>
Contributor
Author
|
/ok to test 7d2cb12 |
Neither GPTModel.forward() nor HybridModel.forward() accepts ``mtp_source_loss_mask`` — it was never consumed anywhere in-tree. The line just TypeError'd every LLaVA-with-labels forward call, which manifested as the ``test_llava_model.py::test_forward*`` and ``test_cuda_graphs.py::test_llava_cudagraph_is_last_layer_logic`` failures. Drop the kwarg; keep the ``loss_mask`` line since that kwarg IS accepted. If a downstream fork ever adds MTP support to HybridModel and expects this signal, re-thread it there — the plumbing is trivial and the correct wiring depends on where MTP consumes it. Signed-off-by: Ryan Prenger <rprenger@nvidia.com> Signed-off-by: rprenger <rprenger@nvidia.com>
Contributor
Author
|
/ok to test 11abf32 |
Contributor
Author
|
/ok to test cffcdbc |
santhnm2
approved these changes
Aug 18, 2026
jaredcasper
approved these changes
Aug 19, 2026
dimapihtar
approved these changes
Aug 20, 2026
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/32399934814 |
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
12 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.
Adds an end-to-end VLM inference path on top of the DynamicInferenceEngine:
engine plumbing accepts per-request image inputs (imgs, imgs_sizes, num_tiles, num_img_embeddings_per_tile); the engine expands![]()
placeholders into pad tokens, runs the vision encoder on the first PP stage, and attaches per-request image embeddings and an image-token mask to the DynamicInferenceContext so the decoder forward can splice them back in.
wire schema between InferenceClient, the coordinator, and the engine drain grows an optional 5th slot for raw image bytes; text-only callers keep the 4-slot payload. Prefix-cache routing is skipped for image-bearing requests so text-identical prompts with different images do not falsely share kv-cache prefixes.
image preprocessing (dynamic-resolution and static tiling), pixel-stat encoder registry, chat_template plumbing on /v1/completions and /v1/chat/completions, and a VLM-aware run_dynamic_text_generation_server entrypoint that auto-detects VLM checkpoints and builds the LLaVA- wrapped model with VLMInferenceWrapper.
LLaVAModel gets a forward_lm_only entry point for the dynamic path, plus a few attributes the wrapper reads. Upstream audio/video params (sound_model, sound_projection, sound_token_index, temporal_patch_dim, separate_video_embedder, temporal_ckpt_compat) are preserved as no-op stubs so the constructor signature stays source-compatible.
Text-only inference is unaffected: none of the new engine or context work fires unless a caller passes imgs/imgs_sizes/num_tiles.
What does this PR do?
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.