[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)