Skip to content

feat(sft): a VL SFT rollout, so a -vl model config can be fine-tuned - #2377

Open
shifangx wants to merge 1 commit into
THUDM:mainfrom
shifangx:feat/sft-rollout-vl
Open

shifangx wants to merge 1 commit into
THUDM:mainfrom
shifangx:feat/sft-rollout-vl

Conversation

@shifangx

Copy link
Copy Markdown

What

slime has VL RL and text SFT, and they do not meet. slime/rollout/sft_rollout.py
loads a processor at :37 and never refers to it again, its tokens come from the
text-only get_loss_mask at :49, and nothing populates Sample.multimodal_inputs.
Point a -vl model provider at it and the run trains a VL model on text-only batches
while looking entirely healthy — which is why none of the four SFT scripts sources a
-vl config.

This adds slime/rollout/sft_rollout_vl.py: the same rollout with the processor
actually used.

Why it is small

Almost nothing had to be invented — everything downstream of the encode step already
existed and keys on whether the field is set, not on the rollout being RL:

  • Dataset builds multimodal_inputs for every row when --multimodal-keys is set
    (utils/data.py:265)
  • ray/rollout.py:414 forwards multimodal_train_inputs for any rollout function
  • megatron_utils/data.py:136 concatenates them, model.py:434 makes them forward kwargs
  • slime_plugins/models/qwen3_5_vl.py:205 consumes them

The one gap was a caller for mask_utils.py:244's
get_loss_mask_with_multimodal_alignment, which takes input_ids as an input rather
than producing them — it was written for exactly this and had no caller anywhere in the
tree.

So the new code is one encode step, kept as a pure function (encode_multimodal_sample)
so it tests without Ray, a buffer or a GPU: render with add_generation_prompt=False so
the assistant turn stays in the text, run the processor over text + images, hand the
resulting input_ids to the alignment, and package everything the processor returned
except input_ids/attention_mask as multimodal_train_inputs.

Rows without images fall back to the text path, so this is a superset of
sft_rollout.py and a mixed dataset trains correctly.

Guards

Two things the module refuses to do quietly. An all-zero loss mask is counted and warned
per row — a single one can be deliberate via step_loss_mask — but a whole batch of them
raises, because otherwise the step runs and reports a plausible loss while learning
nothing. And loss_mask[-0:] is the whole list rather than the empty one, so the
zero-response case is spelled out instead of silently training on the prompt.

Testing

tests/test_sft_rollout_vl.py, 7 passed. Six wiring tests need neither a checkpoint
nor a GPU. The seventh pins the alignment against a real Qwen3.5-VL template, a real fast
tokenizer and an odd 137-token image expansion: the decoded trainable span is
byte-identical to the one the text-only path produces and carries nothing from the
prompt. It skips unless SLIME_TEST_QWEN3_5_VL_CHECKPOINT points at a staged checkpoint.

Also exercised end to end on 8×B200 (Qwen3.5-35B-A3B, GEO3K, TP2·EP8·PP1·CP1,
--debug-train-only): every rollout logged 128/128 samples with images, and
train/loss fell 2.153 → 0.166 over five optimizer steps with grad_norm 67 → 4.9.

Known limitation

The alignment is a left-pad of zeros onto a mask computed from the text-only projection
of the messages, so it is correct exactly while every token the processor adds sits
before the first trainable one — true for an image in a user turn followed by an
assistant turn, and not a general guarantee. The tests pin that shape.

slime has VL RL and text SFT, and they do not meet: sft_rollout.py loads a
processor at :37 and never uses it again, its tokens come from the text-only
get_loss_mask at :49, and nothing populates Sample.multimodal_inputs. Point a
-vl model provider at it and the run trains a VL model on text-only batches
while looking entirely healthy -- which is why none of the four SFT scripts
sources a -vl config.

Almost nothing had to be invented. Dataset already builds multimodal_inputs for
every row when --multimodal-keys is set (utils/data.py:265), ray/rollout.py:414
forwards multimodal_train_inputs to the trainer for any rollout function,
megatron_utils/data.py:136 concatenates them, model.py:434 makes them forward
kwargs, and qwen3_5_vl.py:205 consumes them. The one gap was a caller for
mask_utils.py:244's get_loss_mask_with_multimodal_alignment, which takes
input_ids as an input rather than producing them -- it was written for exactly
this and had no caller anywhere in the tree.

So the new code is one encode step, kept as a pure function
(encode_multimodal_sample) so it tests without Ray, a buffer or a GPU: render
the messages with add_generation_prompt=False so the assistant turn stays in the
text, run the processor over text + images, hand the resulting input_ids to the
alignment, and package everything the processor returned except input_ids and
attention_mask as multimodal_train_inputs.

Rows without images fall back to the text path, so this is a superset of
sft_rollout.py and a mixed dataset trains correctly.

Two things the module refuses to do quietly. An all-zero loss mask is counted
and warned per row -- a single one can be deliberate via step_loss_mask -- but a
whole batch of them is a chat-template/--loss-mask-type mismatch and raises,
because otherwise the step runs and reports a plausible loss while learning
nothing. And loss_mask[-0:] is the whole list rather than the empty one, so the
zero-response case is spelled out instead of silently training on the prompt.

The alignment is a left-pad of zeros onto a mask computed from the text-only
projection of the messages, so it is correct exactly while every token the
processor adds sits before the first trainable one. The tests pin that: with a
real Qwen3.5-VL template, a real fast tokenizer and an odd 137-token image
expansion, the decoded trainable span is byte-identical to the one the text-only
path produces, and carries nothing from the prompt. That test skips unless
SLIME_TEST_QWEN3_5_VL_CHECKPOINT points at a staged checkpoint; the six wiring
tests need neither a checkpoint nor a GPU.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant