[v1] Support multimodal data training (#10656)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

tests_v1/core/utils/test_batching.py

@@ -370,3 +370,208 @@ def test_dynamic_padding_free_fill_buffer_restarts_until_micro_batch_is_complete
     assert len(batch) == 1
     assert batch[0]["input_ids"].shape == (1, 18)
     assert len(batch_generator._buffer) == 1
+
+
+def _image_fragment(n_pad: int = 4, merge_sq: int = 4):
+    """Hand-crafted image fragment: vision_start + n_pad image_pad + vision_end."""
+    import torch
+
+    pad, vstart, vend = 9, 8, 7
+    return {
+        "input_ids": [vstart] + [pad] * n_pad + [vend],
+        "mm_token_type_ids": [0] + [1] * n_pad + [0],
+        "pixel_values": torch.zeros((n_pad * merge_sq, 16), dtype=torch.float32),
+        "image_grid_thw": torch.tensor([[1, 2, n_pad * 2]], dtype=torch.long),
+    }
+
+
+def _text_sample(n: int, base: int = 100):
+    s = _make_model_input(n, start=base)
+    s["position_ids"] = list(range(1, n + 1))
+    return s
+
+
+def test_inject_appends_zero_loss_dummy_into_collated_text_batch():
+    import torch
+
+    from llamafactory.v1.core.utils.batching import _collate_micro_batch, _inject_dummy_into_collated
+
+    collated = _collate_micro_batch([_text_sample(20), _text_sample(8)], cutoff_len=4096)
+    assert "pixel_values" not in collated
+    bsz, seqlen = collated["input_ids"].shape
+
+    frag = _image_fragment(n_pad=4)
+    fl = len(frag["input_ids"])
+    _inject_dummy_into_collated(collated, frag, marker=1)
+
+    new_len = seqlen + fl
+    # every sequence field grew by the fragment length, batch size unchanged
+    for key in ("input_ids", "attention_mask", "labels", "loss_weights", "position_ids", "mm_token_type_ids"):
+        assert collated[key].shape == (bsz, new_len)
+
+    # dummy lives only in row 0's tail; other rows are padding (attention 0) there
+    assert collated["input_ids"][0, seqlen:].tolist() == frag["input_ids"]
+    assert collated["attention_mask"][0, seqlen:].tolist() == [1] * fl
+    assert collated["attention_mask"][1, seqlen:].tolist() == [0] * fl
+    # zero loss contribution
+    assert collated["labels"][0, seqlen:].tolist() == [IGNORE_INDEX] * fl
+    assert torch.all(collated["loss_weights"][:, seqlen:] == 0.0)
+    assert collated["mm_token_type_ids"][0, seqlen:].tolist() == frag["mm_token_type_ids"]
+    # pixel features carried verbatim
+    assert torch.equal(collated["pixel_values"], frag["pixel_values"])
+    assert torch.equal(collated["image_grid_thw"], frag["image_grid_thw"])
+
+
+def test_inject_video_concatenates_alongside_existing_image():
+    """Injecting a missing modality leaves the other modality's features intact (dim-0 cat)."""
+    import torch
+
+    from llamafactory.v1.core.utils.batching import _collate_micro_batch, _inject_dummy_into_collated
+
+    img = _text_sample(10)
+    img["pixel_values"] = torch.ones((8, 16), dtype=torch.float32)
+    img["image_grid_thw"] = torch.tensor([[1, 2, 4]], dtype=torch.long)
+    img["mm_token_type_ids"] = [0] * 10
+    collated = _collate_micro_batch([img], cutoff_len=4096)
+
+    video_frag = {
+        "input_ids": [8, 6, 6, 7],
+        "mm_token_type_ids": [0, 2, 2, 0],
+        "pixel_values_videos": torch.zeros((8, 16), dtype=torch.float32),
+        "video_grid_thw": torch.tensor([[1, 2, 4]], dtype=torch.long),
+    }
+    _inject_dummy_into_collated(collated, video_frag, marker=2)
+
+    # image features untouched, video features added
+    assert torch.equal(collated["pixel_values"], torch.ones((8, 16)))
+    assert collated["pixel_values_videos"].shape[0] == 8
+    assert collated["video_grid_thw"].shape[0] == 1
+    assert collated["mm_token_type_ids"][0, -4:].tolist() == [0, 2, 2, 0]
+
+
+def test_collate_creates_mm_token_type_ids_for_pure_text_then_inject():
+    """A pure-text micro batch has no mm_token_type_ids; injection must create it."""
+    from llamafactory.v1.core.utils.batching import _collate_micro_batch, _inject_dummy_into_collated
+
+    collated = _collate_micro_batch([_text_sample(12)], cutoff_len=4096)
+    assert "mm_token_type_ids" not in collated
+    seqlen = collated["input_ids"].shape[1]
+
+    frag = _image_fragment(n_pad=3)
+    _inject_dummy_into_collated(collated, frag, marker=1)
+
+    assert "mm_token_type_ids" in collated
+    assert collated["mm_token_type_ids"].shape == collated["input_ids"].shape
+    # original region all zero (text), dummy region carries the markers
+    assert collated["mm_token_type_ids"][0, :seqlen].tolist() == [0] * seqlen
+    assert collated["mm_token_type_ids"][0, seqlen:].tolist() == frag["mm_token_type_ids"]
+
+
+def _audio_fragment(n_tok: int = 2, n_frames: int = 3000):
+    """Hand-crafted audio fragment: audio_bos + n_tok AUDIO + audio_eos, with feature rows."""
+    import torch
+
+    aud, bos, eos = 50, 51, 52
+    return {
+        "input_ids": [bos] + [aud] * n_tok + [eos],
+        "mm_token_type_ids": [0] + [3] * n_tok + [0],
+        "input_features": torch.zeros((1, 128, n_frames), dtype=torch.float32),
+        "feature_attention_mask": torch.ones((1, n_frames), dtype=torch.long),
+    }
+
+
+def test_inject_audio_dummy_into_text_batch():
+    """A pure-text micro batch gets an audio dummy appended so the audio tower fires on every rank."""
+    import torch
+
+    from llamafactory.v1.core.utils.batching import _collate_micro_batch, _inject_dummy_into_collated
+
+    collated = _collate_micro_batch([_text_sample(12)], cutoff_len=4096)
+    assert "input_features" not in collated
+    seqlen = collated["input_ids"].shape[1]
+
+    frag = _audio_fragment(n_tok=2)
+    fl = len(frag["input_ids"])
+    _inject_dummy_into_collated(collated, frag, marker=3)
+
+    # audio feature tensors carried verbatim; placeholder tokens marked 3 in the dummy tail
+    assert torch.equal(collated["input_features"], frag["input_features"])
+    assert torch.equal(collated["feature_attention_mask"], frag["feature_attention_mask"])
+    assert collated["mm_token_type_ids"][0, seqlen:].tolist() == frag["mm_token_type_ids"]
+    # zero loss contribution from the dummy
+    assert collated["labels"][0, seqlen:].tolist() == [IGNORE_INDEX] * fl
+    assert torch.all(collated["loss_weights"][:, seqlen:] == 0.0)
+
+
+def test_audio_truncation_drops_orphaned_item_and_zeros_tokens():
+    """Truncating mid-audio trims the orphaned feature row and zeros its in-window tokens."""
+    import torch
+
+    from llamafactory.v1.core.utils.collation import _align_multimodal_on_truncation
+
+    aud = 50
+    # text(2) + [audio#0: 4 tok] + text(1) + [audio#1: 4 tok] + text(1)
+    input_ids = [1, 2] + [aud] * 4 + [3] + [aud] * 4 + [4]
+    mm = [0, 0] + [3] * 4 + [0] + [3] * 4 + [0]
+    sample = {
+        "input_ids": input_ids,
+        "labels": input_ids.copy(),
+        "loss_weights": [1.0] * len(input_ids),
+        "mm_token_type_ids": mm,
+        "input_features": torch.zeros((2, 128, 10), dtype=torch.float32),
+        "feature_attention_mask": torch.ones((2, 10), dtype=torch.long),
+    }
+    # audio#1 occupies positions 7..10; cut at 9 so its last token (10) is orphaned, audio#0 intact
+    out = _align_multimodal_on_truncation(dict(sample), max_length=9)
+
+    assert out["input_features"].shape[0] == 1  # only the complete audio#0 survives
+    assert out["feature_attention_mask"].shape[0] == 1
+    # audio#0 tokens (positions 2..5) untouched
+    assert all(out["input_ids"][i] == aud and out["mm_token_type_ids"][i] == 3 for i in range(2, 6))
+    # audio#1's in-window tokens (positions 7,8) zeroed + delabeled (positions >= 9 cut by truncation)
+    for i in (7, 8):
+        assert out["input_ids"][i] == 0
+        assert out["mm_token_type_ids"][i] == 0
+        assert out["labels"][i] == IGNORE_INDEX
+        assert out["loss_weights"][i] == 0.0
+
+
+def test_audio_truncation_keeps_all_when_complete():
+    """No trimming when the cut falls after every audio's last token."""
+    import torch
+
+    from llamafactory.v1.core.utils.collation import _align_multimodal_on_truncation
+
+    aud = 50
+    input_ids = [1] + [aud] * 4 + [2]
+    sample = {
+        "input_ids": input_ids,
+        "labels": input_ids.copy(),
+        "loss_weights": [1.0] * len(input_ids),
+        "mm_token_type_ids": [0] + [3] * 4 + [0],
+        "input_features": torch.zeros((1, 128, 10), dtype=torch.float32),
+        "feature_attention_mask": torch.ones((1, 10), dtype=torch.long),
+    }
+    out = _align_multimodal_on_truncation(dict(sample), max_length=6)
+    assert out["input_features"].shape[0] == 1
+    assert out["input_ids"] == input_ids
+
+
+def test_drop_unsupervised_samples():
+    """Samples whose supervised tokens fall entirely beyond cutoff_len are dropped (warn once)."""
+    from types import SimpleNamespace
+
+    def _s(weights):  # a sample's input_ids length matches its loss_weights length
+        return {"input_ids": list(range(len(weights))), "loss_weights": weights}
+
+    gen = SimpleNamespace(cutoff_len=4, _warned_truncation=False)
+    samples = [
+        _s([0.0, 0.0, 1.0, 1.0]),             # fits cutoff (len 4), supervised -> kept
+        _s([0.0, 0.0, 0.0, 0.0, 1.0, 1.0]),   # len 6 > 4, supervision only beyond cutoff -> dropped
+        _s([1.0, 1.0]),                        # short, fully supervised -> kept
+        _s([0.0, 0.0, 1.0, 1.0, 1.0, 1.0]),   # len 6 > 4 but supervision within cutoff -> kept
+    ]
+    kept = BatchGenerator._drop_unsupervised(gen, samples)
+    assert kept == [samples[0], samples[2], samples[3]]
+    assert gen._warned_truncation is True
+

tests_v1/plugins/data_plugins/test_converter.py

@@ -71,6 +71,148 @@ def test_sharegpt_converter():
     assert DataConverterPlugin("sharegpt")(example) == expected_data
 
 
+def test_sharegpt_converter_multimodal():
+    example = {
+        "conversations": [
+            {"from": "human", "value": "What is <image> and what happens in <video>?"},
+            {"from": "gpt", "value": "An image and a video."},
+        ],
+        "images": ["/p/a.jpg"],
+        "videos": ["/p/v.mp4"],
+    }
+    expected_data = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "value": "What is "},
+                    {"type": "image_url", "value": "/p/a.jpg"},
+                    {"type": "text", "value": " and what happens in "},
+                    {"type": "video_url", "value": "/p/v.mp4"},
+                    {"type": "text", "value": "?"},
+                ],
+                "loss_weight": 0.0,
+            },
+            {"role": "assistant", "content": [{"type": "text", "value": "An image and a video."}], "loss_weight": 1.0},
+        ]
+    }
+    assert DataConverterPlugin("sharegpt")(example) == expected_data
+
+
+def test_sharegpt_converter_multiple_images_in_order():
+    # images are a sample-level list consumed by <image> tags in document order across turns
+    example = {
+        "conversations": [
+            {"from": "human", "value": "<image><image>Compare these."},
+            {"from": "gpt", "value": "Done."},
+        ],
+        "images": ["/p/a.jpg", "/p/b.jpg"],
+    }
+    user = DataConverterPlugin("sharegpt")(example)["messages"][0]
+    assert user["content"] == [
+        {"type": "image_url", "value": "/p/a.jpg"},
+        {"type": "image_url", "value": "/p/b.jpg"},
+        {"type": "text", "value": "Compare these."},
+    ]
+
+
+def test_sharegpt_converter_no_media_unchanged():
+    # backward compatibility: a scalar (non-list) image column and no tags is normalized; with no
+    # media columns at all the output is byte-identical to the text-only path.
+    example = {"conversations": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "yo"}]}
+    assert DataConverterPlugin("sharegpt")(example) == {
+        "messages": [
+            {"role": "user", "content": [{"type": "text", "value": "hi"}], "loss_weight": 0.0},
+            {"role": "assistant", "content": [{"type": "text", "value": "yo"}], "loss_weight": 1.0},
+        ]
+    }
+
+
+def test_alpaca_converter_multimodal():
+    example = {"instruction": "Describe <image>", "input": "", "output": "ok", "images": ["/p/a.jpg"]}
+    user = DataConverterPlugin("alpaca")(example)["messages"][0]
+    assert user["content"] == [
+        {"type": "text", "value": "Describe "},
+        {"type": "image_url", "value": "/p/a.jpg"},
+    ]
+
+
+def test_pair_converter_multimodal_shared_media():
+    # chosen and rejected each reference the same sample-level image
+    example = {
+        "chosen": [
+            {"role": "user", "content": "Look at <image>"},
+            {"role": "assistant", "content": "good"},
+        ],
+        "rejected": [
+            {"role": "user", "content": "Look at <image>"},
+            {"role": "assistant", "content": "bad"},
+        ],
+        "images": ["/p/a.jpg"],
+    }
+    out = DataConverterPlugin("pair")(example)
+    for side in ("chosen_messages", "rejected_messages"):
+        assert out[side][0]["content"] == [
+            {"type": "text", "value": "Look at "},
+            {"type": "image_url", "value": "/p/a.jpg"},
+        ]
+
+
+def test_converter_media_count_mismatch():
+    # more tags than media files
+    with pytest.raises(ValueError, match="More <image> tags"):
+        DataConverterPlugin("sharegpt")(
+            {
+                "conversations": [{"from": "human", "value": "<image><image>"}, {"from": "gpt", "value": "x"}],
+                "images": ["/p/a.jpg"],
+            }
+        )
+    # fewer tags than media files
+    with pytest.raises(ValueError, match="Fewer <image> tags"):
+        DataConverterPlugin("sharegpt")(
+            {
+                "conversations": [{"from": "human", "value": "<image>"}, {"from": "gpt", "value": "x"}],
+                "images": ["/p/a.jpg", "/p/b.jpg"],
+            }
+        )
+
+
+def test_converter_audio_column_and_tag():
+    # an <audio> tag consumes the next path from the audios column, lifted into an audio_url block
+    example = {
+        "conversations": [
+            {"from": "human", "value": "hear <audio>What is this?"},
+            {"from": "gpt", "value": "A bell."},
+        ],
+        "audios": ["/p/a.wav"],
+    }
+    user = DataConverterPlugin("sharegpt")(example)["messages"][0]
+    assert user["content"] == [
+        {"type": "text", "value": "hear "},
+        {"type": "audio_url", "value": "/p/a.wav"},
+        {"type": "text", "value": "What is this?"},
+    ]
+
+
+def test_converter_audio_count_mismatch():
+    # more audio tags than files
+    with pytest.raises(ValueError, match="More <audio> tags"):
+        DataConverterPlugin("sharegpt")(
+            {
+                "conversations": [{"from": "human", "value": "<audio><audio>"}, {"from": "gpt", "value": "x"}],
+                "audios": ["/p/a.wav"],
+            }
+        )
+    # fewer audio tags than files
+    with pytest.raises(ValueError, match="Fewer <audio> tags"):
+        DataConverterPlugin("sharegpt")(
+            {
+                "conversations": [{"from": "human", "value": "<audio>"}, {"from": "gpt", "value": "x"}],
+                "audios": ["/p/a.wav", "/p/b.wav"],
+            }
+        )
+
+
 @pytest.mark.parametrize("num_samples", [16])
 def test_pair_converter(num_samples: int):
     data_args = DataArguments(train_dataset="llamafactory/v1-dataset-info/orca-dpo-pairs.yaml")
@@ -117,3 +259,4 @@ def test_pair_converter(num_samples: int):
             ],
         }
         assert data_engine[index] == {"_dataset_name": "tiny_dataset", **expected_data}
+