@terrykong @ashors1 Created a draft PR (duplicate of #521) to see if CI passes on this instead.

adding @ashors1 @yfw to review/approve

I tried a run on the clevr dataset and noticed the token_mult_prob_error was a bit high. This is a measure of the difference in logprobs between vllm and dtensor. For non-vlm, we generally expect this number to be < 1.05 (for qwen 2.5-1b, we see around 1.02). A higher number usually indicates some issue with the refit so vllm and dtensor aren't exactly matching. Is there anything specific about the vlm setup that could cause this?

I'm not entirely sure what causes the token_mult_prob_error to be relatively higher than the LLM only case. I did not fiddle with the either vllm or the dtensor model forward setup so I'm not sure what causes this.

For the functional test cases, I had to choose higher thresholds for the token_mult_prob_error errors. There are also a few spikes like you have shown that I cannot quite explain (the loss at these iterations does not spike though).

I'm not entirely sure what causes the token_mult_prob_error to be relatively higher than the LLM only case. I did not fiddle with the either vllm or the dtensor model forward setup so I'm not sure what causes this.

For the functional test cases, I had to choose higher thresholds for the token_mult_prob_error errors. There are also a few spikes like you have shown that I cannot quite explain (the loss at these iterations does not spike though).

One thought is that we do the preprocessing of the image before calling the model in the dtensor path whereas vllm does preprocessing of the image internally (if I understand what is happening correctly). We may need to make sure whatever preprocessing vllm is doing matches exactly what we're doing in the dtensor path.

This will take me a while to analyse since I don't know exactly how the vllm engine processes the images internally.

For the policy, the typical multimodal pipeline is to use the processor to encode the chat template dict into a sequence of text tokens, multimodal tokens indexed by the key pixel_values and metadata image_grid_thw to compute mRoPE embeddings (keys would be different for videos or audio). The pixel_values item takes a PIL Image as input --> processes it into patches.

For vllm, the message log is simply reformatted into the format specified in this tutorial. The same sequence of PIL Images is provided to the vLLM frontend.
I assume the token ids and multimodal tokens must be the same, but I will have to double check. In most cases the logprobs mult error is still close to 1, so it could also be numerical differences between the vllm frontend and the dtensor policy.

From today's meeting, the remaining blockers on this PR:

• understanding the logprob error
• address review on API VLM processor keys

re: Remaining blockers:

1. understanding the logprob error: This is something I want to chalk up to how vllm loads multimodal image embeddings in the image processor. For LLM-only, I noted that vllm takes the same list of token_ids (int value list) that the policy consumes (i.e. going through the same text embedding layer, etc.). However, for multimodal images, vllm processes the images internally. There could also be differences in how sampling is done differently. I found the following excerpt from vllm docs https://docs.vllm.ai/en/v0.9.1/usage/v1_guide.html#feature-model

Logprobs in V1 are now returned immediately once computed from the model’s raw output (i.e. before applying any logits post-processing such as temperature scaling or penalty adjustments). As a result, the returned logprobs do not reflect the final adjusted probabilities used during sampling.
Support for logprobs with post-sampling adjustments is in progress and will be added in future updates.

I prefer handling this issue in a separate PR (and merging an initial support first) for three four reasons:

• this discrepancy is isolated to multimodal models only, so a "fix" can be shipped independently
• multiple VLMs converge on three different datasets despite the apparent discrepancy. It is equivalent to training GRPO with a slightly off-policy model, but it does not seem to be very unstable or destructive to the learning process
• other PRs break multimodal support regularly (every 2-3 days) and I have to rollback / fix those changes in my PR to make my scripts work. Merging this PR or at least the test cases will prevent other PRs from breaking multimodal support
• the PR has gotten very big as it is, and adding more fixes will add additional overhead to the review process

2. PR has now migrated (again) to feat: GRPO + SFT Dtensor support for multimodal training  #712, and is tested on 4 families of multimodal models and 3 datasets. This rollsback the passing around of the vlm_kwargs list throughout the training process and instead proposes a PackedGenericDataItem to handle non-sequence data items (most of them would be multimodal tensors). The single implementation seems to work for multiple multimodal models without any additional modifications to the config.