[v1] Support multimodal Ulysses CP and memory-efficient chunk loss for SFT (#10762)

tests_v1/plugins/model_plugins/test_chunk_loss.py

@@ -0,0 +1,154 @@
+# Copyright 2026 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from copy import deepcopy
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+from transformers.modeling_outputs import CausalLMOutput
+
+from llamafactory.v1.plugins.model_plugins.chunk_loss import LossPlugin, _ChunkedLinearCrossEntropy
+from llamafactory.v1.trainers.sft_trainer import SFTTrainer
+from llamafactory.v1.utils.constants import IGNORE_INDEX
+from llamafactory.v1.utils.env import find_available_port
+from llamafactory.v1.utils.pytest import dist_env
+
+
+class _TinyCausalLM(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(31, 16)
+        self.lm_head = nn.Linear(16, 31, bias=False)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def forward(self, input_ids, **_):
+        return CausalLMOutput(logits=self.lm_head(self.embed_tokens(input_ids)))
+
+
+def _make_model():
+    return _TinyCausalLM()
+
+
+def _assert_gradients_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
+    error = torch.linalg.vector_norm(actual.float() - expected.float())
+    reference = torch.linalg.vector_norm(expected.float())
+    assert error <= 2 * torch.finfo(expected.dtype).eps * reference
+
+
+def _weighted_cross_entropy(logits, labels, loss_weights):
+    losses = F.cross_entropy(logits.flatten(0, -2).float(), labels.flatten(), reduction="none")
+    return (losses * loss_weights.flatten()).sum()
+
+
+@pytest.mark.parametrize("frozen_head", [False, True])
+def test_chunk_loss_matches_eager_loss_and_gradients(frozen_head):
+    torch.manual_seed(0)
+    eager_head = nn.Linear(4, 7).to(torch.bfloat16)
+    eager_head.requires_grad_(not frozen_head)
+    chunk_head = deepcopy(eager_head)
+    eager_hidden = torch.randn(2, 5, 4, dtype=torch.bfloat16, requires_grad=True)
+    chunk_hidden = eager_hidden.detach().clone().requires_grad_()
+    labels = torch.tensor([[0, 1, IGNORE_INDEX, 3, 4], [5, 6, 0, 1, 2]])
+    loss_weights = torch.tensor([[0.0, 0.25, 1.0, 0.75, 1.0], [1.0, 0.5, 0.0, 0.25, 1.0]])
+
+    eager_loss = _weighted_cross_entropy(eager_head(eager_hidden), labels, loss_weights)
+    chunk_loss = _ChunkedLinearCrossEntropy.apply(
+        chunk_hidden, chunk_head.weight, chunk_head.bias, labels, loss_weights, 3
+    )
+    scale = 0.07 / (loss_weights.sum() + 1e-6)
+    (eager_loss * scale).backward()
+    (chunk_loss * scale).backward()
+
+    torch.testing.assert_close(chunk_loss, eager_loss)
+    _assert_gradients_close(chunk_hidden.grad, eager_hidden.grad)
+    for actual, expected in zip(chunk_head.parameters(), eager_head.parameters()):
+        if frozen_head:
+            assert actual.grad is None
+        else:
+            _assert_gradients_close(actual.grad, expected.grad)
+
+
+@pytest.mark.parametrize("zero_supervision", [False, True])
+def test_chunk_sft_loss_matches_reference(zero_supervision):
+    model = _make_model()
+    input_ids = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
+    labels = input_ids.clone()
+    labels[0, 2] = IGNORE_INDEX
+    batch = {
+        "input_ids": input_ids,
+        "attention_mask": torch.ones_like(input_ids),
+        "position_ids": torch.arange(5).expand(2, -1),
+        "labels": labels,
+        "loss_weights": torch.tensor([[0.0, 0.25, 0.0, 1.0, 0.5], [0.0, 0.0, 0.75, 0.0, 1.0]]),
+    }
+    if zero_supervision:
+        batch["labels"].fill_(IGNORE_INDEX)
+        batch["loss_weights"].zero_()
+    original_batch = {key: value.clone() for key, value in batch.items()}
+
+    trainer = object.__new__(SFTTrainer)
+    trainer.model = model
+    trainer.device = torch.device("cpu")
+    trainer.cp_size = 1
+    trainer._uses_mrope = False
+    trainer._chunk_loss = None
+    eager_loss = trainer.compute_loss(batch)
+
+    chunk_model = deepcopy(model)
+    trainer.model = chunk_model
+    trainer._chunk_loss = LossPlugin("chunk_loss")(chunk_model, chunk_size=3)
+    chunk_loss = trainer.compute_loss(batch)
+
+    torch.testing.assert_close(chunk_loss, eager_loss)
+    for key in batch:
+        torch.testing.assert_close(batch[key], original_batch[key])
+    torch.testing.assert_close(chunk_model(input_ids=input_ids).logits, model(input_ids=input_ids).logits)
+
+
+@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason="Requires the CPU Gloo backend.")
+def test_chunk_loss_preserves_ddp_output_backward_hooks():
+    torch.manual_seed(7)
+    eager_model = _make_model()
+    chunk_model = deepcopy(eager_model)
+    loss_fn = LossPlugin("chunk_loss")(chunk_model, chunk_size=2)
+    input_ids = torch.tensor([[1, 2, 3, 4, 5]])
+    model_inputs = {"input_ids": input_ids, "use_cache": False}
+    labels = torch.tensor([[2, 3, 4, 5, IGNORE_INDEX]])
+    loss_weights = torch.tensor([[0.0, 0.25, 1.0, 0.5, 0.0]])
+    outer_scale = 0.3 / loss_weights.sum()
+    outer_outputs = []
+
+    def capture_outer_output(_model, _args, output):
+        outer_outputs.append(output.logits)
+
+    with dist_env(master_port=find_available_port()):
+        dist.init_process_group("gloo")
+        wrapped_model = DDP(chunk_model, find_unused_parameters=True)
+        wrapped_model.register_forward_hook(capture_outer_output)
+        eager_loss = _weighted_cross_entropy(eager_model(**model_inputs).logits, labels, loss_weights)
+        chunk_loss = loss_fn(wrapped_model, model_inputs, labels, loss_weights)
+
+        assert loss_fn._active_state is None
+        assert chunk_loss is outer_outputs.pop()
+        torch.testing.assert_close(chunk_loss, eager_loss)
+        (eager_loss * outer_scale).backward()
+        (chunk_loss * outer_scale).backward()
+        for expected, actual in zip(eager_model.parameters(), chunk_model.parameters(), strict=True):
+            torch.testing.assert_close(actual.grad, expected.grad)

tests_v1/plugins/model_plugins/test_ulysses_cp.py

@@ -12,19 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 import torch.multiprocessing as mp
+from torch import nn
 
+import llamafactory.v1.plugins.model_plugins.parallelization.hook as hook_module
 from llamafactory.v1.accelerator.interface import DistributedInterface
 from llamafactory.v1.config.model_args import ModelArguments
 from llamafactory.v1.config.training_args import TrainingArguments
 from llamafactory.v1.core.model_engine import ModelEngine
 from llamafactory.v1.plugins.model_plugins.parallelization import ulysses
+from llamafactory.v1.plugins.model_plugins.parallelization.batch import prepare_sequence_parallel_batch
 from llamafactory.v1.plugins.model_plugins.parallelization.sequence_parallel import (
     SequenceParallelModelPlugin,
     sequence_parallel_loss,
 )
+from llamafactory.v1.utils.constants import IGNORE_INDEX
 from llamafactory.v1.utils.env import find_available_port
 from llamafactory.v1.utils.pytest import dist_env
 
@@ -99,3 +105,102 @@ def test_sequence_parallel_loss(cp_size, dp_size, batch_size):
     mp.spawn(
         _test_sequence_parallel_loss, args=(world_size, master_port, cp_size, dp_size, batch_size), nprocs=world_size
     )
+
+
+def test_non_causal_multimodal_encoder_attention_bypasses_ulysses():
+    captured_is_causal = None
+
+    def fake_native_attention(query, _key, _value, _attention_mask, **kwargs):
+        nonlocal captured_is_causal
+        captured_is_causal = kwargs["is_causal"]
+        return query + 1
+
+    query = torch.zeros(1, 4, 2, 8)
+    output = ulysses.new_flash_attn_forward(query, query, query, None, is_causal=False, attn_fn=fake_native_attention)
+
+    torch.testing.assert_close(output, query + 1)
+    assert captured_is_causal is False
+
+
+def _device_mesh(rank=0, size=2):
+    return {"cp": SimpleNamespace(size=lambda: size, get_local_rank=lambda: rank)}
+
+
+class _RecordingLanguageModel(nn.Module):
+    def forward(self, **kwargs):
+        return kwargs
+
+
+def test_multimodal_sequence_parallel_hook(monkeypatch):
+    # One shard gets non-contiguous visual rows while the next shard is empty.
+    cp_rank = [1]
+    device_mesh = {"cp": SimpleNamespace(size=lambda: 3, get_local_rank=lambda: cp_rank[0])}
+    distributed = SimpleNamespace(get_device_mesh=lambda _dim: device_mesh)
+    monkeypatch.setattr(hook_module, "DistributedInterface", lambda: distributed)
+
+    model = nn.Module()
+    model.model = core = nn.Module()
+    boundary = _RecordingLanguageModel()
+    core.visual = nn.Identity()
+    core.language_model = boundary
+    hook_module.install_sequence_parallel_hook(SimpleNamespace(get_base_model=lambda: model))
+
+    fused_inputs = torch.arange(24, dtype=torch.float32).view(2, 6, 2)
+    attention_mask = torch.tensor([[1, 1, 1, 1, 1, 0], [1, 1, 1, 1, 0, 0]])
+    position_ids = torch.arange(36).view(3, 2, 6)
+    visual_mask = torch.tensor([[True, False, False, True, False, False], [True, True, True, False, False, False]])
+    visual_embeds = torch.arange(10, dtype=torch.float32).view(5, 2).requires_grad_()
+    model_inputs = {
+        "input_ids": None,
+        "inputs_embeds": fused_inputs,
+        "attention_mask": attention_mask,
+        "position_ids": position_ids,
+        "visual_pos_masks": visual_mask,
+        "deepstack_visual_embeds": [visual_embeds],
+    }
+
+    outputs = boundary(**model_inputs)
+    torch.testing.assert_close(outputs["inputs_embeds"], fused_inputs[:, 2:4])
+    torch.testing.assert_close(outputs["attention_mask"], attention_mask[:, 2:4])
+    torch.testing.assert_close(outputs["position_ids"], position_ids[..., 2:4])
+    torch.testing.assert_close(outputs["visual_pos_masks"], visual_mask[:, 2:4])
+    torch.testing.assert_close(outputs["deepstack_visual_embeds"][0], visual_embeds[[1, 4]])
+    assert outputs["use_cache"] is False
+
+    outputs["deepstack_visual_embeds"][0].sum().backward()
+    expected_grad = torch.zeros_like(visual_embeds)
+    expected_grad[[1, 4]] = 1
+    torch.testing.assert_close(visual_embeds.grad, expected_grad)
+
+    cp_rank[0] = 2
+    visual_embeds.grad = None
+    empty_visual_embeds = boundary(**model_inputs)["deepstack_visual_embeds"][0]
+    assert empty_visual_embeds.shape == (0, 2)
+    empty_visual_embeds.sum().backward()
+    torch.testing.assert_close(visual_embeds.grad, torch.zeros_like(visual_embeds))
+
+
+def test_prepare_multimodal_sequence_parallel_batch_preserves_encoder_inputs_and_shifts_targets():
+    pixel_values = torch.arange(12, dtype=torch.float32).view(3, 4)
+    batch = {
+        "input_ids": torch.tensor([[1, 2, 3]]),
+        "attention_mask": torch.ones(1, 3, dtype=torch.long),
+        "position_ids": torch.tensor([[0, 1, 2]]),
+        "mm_token_type_ids": torch.tensor([[0, 1, 1]]),
+        "labels": torch.tensor([[1, 2, 3]]),
+        "loss_weights": torch.tensor([[9.0, 0.5, 2.0]]),
+        "pixel_values": pixel_values,
+    }
+    for rank in range(2):
+        prepared = prepare_sequence_parallel_batch(
+            batch, device=torch.device("cpu"), device_mesh=_device_mesh(rank), uses_mrope=True
+        )
+        assert prepared.model_inputs["input_ids"].tolist() == [[1, 2, 3, 0]]
+        assert prepared.model_inputs["mm_token_type_ids"].tolist() == [[0, 1, 1, 0]]
+        assert "labels" not in prepared.model_inputs and "loss_weights" not in prepared.model_inputs
+        assert "position_ids" not in prepared.model_inputs
+        assert prepared.model_inputs["attention_mask"].tolist() == [[1, 1, 1, 0]]
+        torch.testing.assert_close(prepared.model_inputs["pixel_values"], pixel_values)
+        assert prepared.local_shift_labels.tolist() == ([[2, 3]] if rank == 0 else [[IGNORE_INDEX, IGNORE_INDEX]])
+        assert prepared.local_shift_loss_weights.tolist() == ([[0.5, 2.0]] if rank == 0 else [[0.0, 0.0]])
+        assert prepared.global_loss_weight_sum.item() == 2.5