[trainer, ops] fix: make channel loss capture lifecycle-safe

scripts/ci/accelerator/README.md

@@ -32,6 +32,15 @@ from the test's launch plan. See
 [Registering Tests in CI](../../../docs/testing.md#registering-tests-in-ci) for
 the manifest example, registration rules and collection check command.
 
+The `qwen38-channel_loss_callback-unit` selection is present in both
+`suites/gpu.json` and `suites/npu.json`. It collects
+`tests/trainer/test_channel_loss_callback.py` for the detached channel-loss
+callback, packed-source alignment, and local SP/Gloo checks. The existing
+`tests/ops` selection also collects chunk-loss observer/offload tests. Consult
+the individual case plans, skip conditions, and job logs for the hardware
+actually exercised; inclusion in an accelerator suite alone does not prove a
+native device run.
+
 ## Launch plans
 
 An unmarked case defaults to one card of the selected accelerator backend, or

scripts/ci/accelerator/suites/gpu.json

@@ -9,6 +9,15 @@
       "tests/parallel/ulysses/test_ulysses.py"
     ]
   },
+  {
+    "name": "qwen38-channel_loss_callback-unit",
+    "suite": "unit",
+    "args": [
+      "-s",
+      "-x",
+      "tests/trainer/test_channel_loss_callback.py"
+    ]
+  },
   {
     "name": "02-unit",
     "suite": "unit",

scripts/ci/accelerator/suites/npu.json

@@ -19,6 +19,15 @@
       "tests/parallel/ulysses/test_all_gather.py"
     ]
   },
+  {
+    "name": "qwen38-channel_loss_callback-unit",
+    "suite": "unit",
+    "args": [
+      "-s",
+      "-x",
+      "tests/trainer/test_channel_loss_callback.py"
+    ]
+  },
   {
     "name": "npu-unit-03",
     "suite": "unit",

tests/ops/loss/test_chunk_loss_offload.py

@@ -0,0 +1,158 @@
+"""Exercise the registered chunk CE under production offload/checkpoint hooks."""
+
+from contextlib import contextmanager, nullcontext
+
+import pytest
+import torch
+from torch.utils.checkpoint import checkpoint
+
+from scripts.ci.accelerator.launch_plan import LaunchSpec
+from veomni.distributed.offloading import OffloadPolicy, build_activation_offloading_context
+from veomni.ops import VeomniOp
+from veomni.ops.kernels.loss.cross_entropy_loss.standard import chunk_loss as chunk_loss_impl
+from veomni.utils.device import IS_CUDA_AVAILABLE, IS_NPU_AVAILABLE
+
+
+DEVICES = [
+    "cpu",
+    pytest.param("cuda", marks=pytest.mark.skipif(not IS_CUDA_AVAILABLE, reason="requires CUDA")),
+    pytest.param("npu", marks=pytest.mark.skipif(not IS_NPU_AVAILABLE, reason="requires native NPU")),
+]
+
+
+@contextmanager
+def _capture_chunk_tokens(target):
+    capture = {}
+    token = chunk_loss_impl._CAPTURE.set(capture)
+    try:
+        yield
+        assert "error" not in capture, capture.get("error")
+        assert isinstance(capture.get("per_token"), torch.Tensor)
+        target.append(capture["per_token"])
+    finally:
+        chunk_loss_impl._CAPTURE.reset(token)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _deterministic_npu_comparisons():
+    """Keep exact offload comparisons independent of NPU reduction scheduling."""
+    if not IS_NPU_AVAILABLE:
+        yield
+        return
+    enabled = torch.are_deterministic_algorithms_enabled()
+    warn_only = torch.is_deterministic_algorithms_warn_only_enabled()
+    torch.use_deterministic_algorithms(True)
+    try:
+        yield
+    finally:
+        torch.use_deterministic_algorithms(enabled, warn_only=warn_only)
+
+
+def _run(device, dtype, *, offload, recompute, frozen, all_ignored, bias_kind=None, observed=None):
+    torch.manual_seed(127)
+    hidden = torch.randn(2, 23, 64).to(device=device, dtype=dtype).requires_grad_(frozen != "hidden")
+    weight = torch.randn(137, 64).to(device=device, dtype=dtype).requires_grad_(frozen != "weight")
+    labels = torch.randint(0, 137, (2, 23), device=device)
+    labels[:, 1:8] = -100
+    if all_ignored:
+        labels.fill_(-100)
+    bias = bias_weight = None
+    if bias_kind is not None:
+        bias = torch.randn(2, 23, 4 if bias_kind == "factored" else 137, device=device, dtype=dtype)
+        bias.requires_grad_()
+        if bias_kind == "factored":
+            bias_weight = torch.randn(137, 4, device=device, dtype=dtype).requires_grad_()
+    op = VeomniOp("cross_entropy_loss", "standard", "chunk_loss")
+    # Even the last (2 x 3) token chunk is eligible for the >1 KiB offload policy.
+    assert 2 * 3 * weight.size(0) * torch.float32.itemsize > 1024
+    offloaded, restored = [], []
+    contexts = (nullcontext(), nullcontext())
+    if offload:
+        contexts = build_activation_offloading_context(enable_activation=True, enable_gradient_checkpointing=recompute)
+        for context in contexts:
+            if not hasattr(context, "pack_hook"):
+                continue
+            original_pack, original_unpack = context.pack_hook, context.unpack_hook
+
+            def pack(tensor, original=original_pack):
+                packed = original(tensor)
+                if packed[0] == OffloadPolicy.OFFLOAD:
+                    assert packed[2].device.type == "cpu"
+                    offloaded.append((tuple(tensor.shape), tensor.device.type))
+                return packed
+
+            def unpack(packed, original=original_unpack):
+                tensor = original(packed)
+                if packed[0] == OffloadPolicy.OFFLOAD:
+                    assert tensor.device == packed[1]
+                    restored.append(tensor.device.type)
+                return tensor
+
+            context.pack_hook, context.unpack_hook = pack, unpack
+
+    def body(x, head, bias, bias_weight):
+        return op(x.sin(), labels, head, bias, bias_weight, chunk_size=5)
+
+    capture = _capture_chunk_tokens(observed) if observed is not None else nullcontext()
+    with capture, contexts[0]:
+        inputs = (hidden, weight, bias, bias_weight)
+        loss = checkpoint(body, *inputs, use_reentrant=False) if recompute else body(*inputs)
+    with contexts[1]:
+        loss.backward()
+    if offload:
+        assert offloaded and restored
+        assert all(
+            len(shape) < 2 or shape[-1] != 137 or (bias_kind == "full" and shape == tuple(bias.shape))
+            for shape, _ in offloaded
+        ), "Transient vocabulary logits escaped the chunk"
+        if device != "cpu":
+            assert any(source == device for _, source in offloaded)
+            assert device in restored
+    return (
+        loss.detach(),
+        hidden.grad,
+        weight.grad,
+        bias.grad if bias is not None else None,
+        bias_weight.grad if bias_weight is not None else None,
+    )
+
+
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("recompute", [False, True])
+@pytest.mark.parametrize("frozen", [None, "hidden", "weight"])
+@pytest.mark.parametrize("all_ignored", [False, True])
+def test_registered_chunk_loss_offload_preserves_loss_and_gradients(device, dtype, recompute, frozen, all_ignored):
+    expected = _run(device, dtype, offload=False, recompute=recompute, frozen=frozen, all_ignored=all_ignored)
+    actual = _run(device, dtype, offload=True, recompute=recompute, frozen=frozen, all_ignored=all_ignored)
+    for value, reference in zip(actual, expected, strict=True):
+        if reference is None:
+            assert value is None
+        else:
+            assert torch.isfinite(value).all()
+            torch.testing.assert_close(value, reference, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("recompute", [False, True])
+@pytest.mark.parametrize("bias_kind", ["full", "factored"])
+@pytest.mark.launch_plan(
+    provider=lambda context: LaunchSpec(device="cpu" if context.params["device"] == "cpu" else "accelerator")
+)
+def test_chunk_observer_bias_offload_parity(device, dtype, recompute, bias_kind):
+    """Keep loss, all gradients and observations exact with native bias/offload."""
+    kwargs = dict(recompute=recompute, frozen=None, all_ignored=False, bias_kind=bias_kind)
+    expected = _run(device, dtype, offload=False, **kwargs)
+    reference_tokens, actual_tokens = [], []
+    observed = _run(device, dtype, offload=False, observed=reference_tokens, **kwargs)
+    actual = _run(device, dtype, offload=True, observed=actual_tokens, **kwargs)
+    for result in (observed, actual):
+        for value, reference in zip(result, expected, strict=True):
+            if reference is None:
+                assert value is None
+            else:
+                assert torch.isfinite(value).all()
+                torch.testing.assert_close(value, reference, rtol=0, atol=0)
+    assert len(reference_tokens) == len(actual_tokens) == 1
+    torch.testing.assert_close(actual_tokens[0], reference_tokens[0], rtol=0, atol=0)

tests/ops/loss/test_cross_entropy_loss.py

@@ -16,13 +16,20 @@
 
 from __future__ import annotations
 
+import subprocess
+import sys
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager, nullcontext
+from threading import Barrier
+
 import pytest
 import torch
 import torch.nn.functional as F
 from torch import Tensor
 from torch.utils._python_dispatch import TorchDispatchMode
 from transformers.loss.loss_utils import fixed_cross_entropy
 
+from scripts.ci.accelerator.launch_plan import LaunchSpec
 from tests.ops.tol import (
     CE_FUSED_ATOL,
     CE_FUSED_GRAD_ATOL,
@@ -35,9 +42,159 @@
 )
 from tests.ops.utils import make_grad_leaf
 from veomni.ops import resolve_op
+from veomni.ops.kernels.loss.cross_entropy_loss.standard import chunk_loss as chunk_loss_impl
 from veomni.utils.device import IS_CUDA_AVAILABLE
 
 
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+def test_chunk_loss_standalone_does_not_import_trainer():
+    program = """
+import sys
+import torch
+from veomni.ops import resolve_op
+
+callback = 'veomni.trainer.callbacks.channel_loss_callback'
+assert callback not in sys.modules
+hidden = torch.randn(2, 4, requires_grad=True)
+weight = torch.randn(7, 4, requires_grad=True)
+loss = resolve_op('cross_entropy_loss', 'standard', 'chunk_loss').wrapper(
+    hidden, torch.tensor([1, 2]), weight, chunk_size=1
+)
+loss.backward()
+assert callback not in sys.modules
+"""
+    result = subprocess.run([sys.executable, "-c", program], capture_output=True, text=True, timeout=30)
+    assert result.returncode == 0, result.stderr
+
+
+@contextmanager
+def _capture_chunk_tokens(target):
+    capture = {}
+    token = chunk_loss_impl._CAPTURE.set(capture)
+    try:
+        yield
+        assert "error" not in capture, capture.get("error")
+        assert isinstance(capture.get("per_token"), torch.Tensor)
+        target.append(capture["per_token"])
+    finally:
+        chunk_loss_impl._CAPTURE.reset(token)
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+def test_chunk_capture_is_nested_and_thread_local():
+    barrier = Barrier(2)
+
+    def worker(seed):
+        outer, inner = {}, {}
+        outer_token = chunk_loss_impl._CAPTURE.set(outer)
+        try:
+            inner_token = chunk_loss_impl._CAPTURE.set(inner)
+            try:
+                barrier.wait(timeout=10)
+                hidden = torch.arange(12, dtype=torch.float32).reshape(3, 4) + seed
+                weight = torch.arange(20, dtype=torch.float32).reshape(5, 4) / 11
+                labels = torch.tensor([0, 1, 2])
+                resolve_op("cross_entropy_loss", "standard", "chunk_loss").wrapper(
+                    hidden, labels, weight, chunk_size=2
+                )
+                expected = F.cross_entropy(F.linear(hidden, weight), labels, reduction="none")
+                torch.testing.assert_close(inner["per_token"], expected)
+                assert "per_token" not in outer
+            finally:
+                chunk_loss_impl._CAPTURE.reset(inner_token)
+            assert chunk_loss_impl._CAPTURE.get() is outer
+        finally:
+            chunk_loss_impl._CAPTURE.reset(outer_token)
+        assert chunk_loss_impl._CAPTURE.get() is None
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        list(pool.map(worker, [0, 1]))
+    assert chunk_loss_impl._CAPTURE.get() is None
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+def test_chunk_observer_failure_does_not_change_main_loss_or_gradients(monkeypatch):
+    hidden = torch.randn(5, 4, requires_grad=True)
+    weight = torch.randn(7, 4, requires_grad=True)
+    labels = torch.tensor([0, 1, 2, 3, -100])
+    reference_hidden = hidden.detach().clone().requires_grad_()
+    reference_weight = weight.detach().clone().requires_grad_()
+    op = resolve_op("cross_entropy_loss", "standard", "chunk_loss").wrapper
+    reference = op(reference_hidden, labels, reference_weight, chunk_size=2)
+    reference.backward()
+    original = F.cross_entropy
+
+    def fail_observer(*args, **kwargs):
+        if kwargs.get("reduction") == "none":
+            raise RuntimeError("injected observer failure")
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(F, "cross_entropy", fail_observer)
+    capture = {}
+    token = chunk_loss_impl._CAPTURE.set(capture)
+    try:
+        actual = op(hidden, labels, weight, chunk_size=2)
+    finally:
+        chunk_loss_impl._CAPTURE.reset(token)
+    actual.backward()
+    assert "injected observer failure" in capture["error"]
+    assert "per_token" not in capture
+    torch.testing.assert_close(actual, reference, rtol=0, atol=0)
+    torch.testing.assert_close(hidden.grad, reference_hidden.grad, rtol=0, atol=0)
+    torch.testing.assert_close(weight.grad, reference_weight.grad, rtol=0, atol=0)
+    with pytest.raises(AssertionError, match="injected observer failure"), _capture_chunk_tokens([]):
+        op(hidden.detach(), labels, weight.detach(), chunk_size=2)
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+@pytest.mark.parametrize("bias_kind", ["full", "factored"])
+@pytest.mark.parametrize("trainable", ["all", "bias_only", "none"])
+@pytest.mark.parametrize("all_ignored", [False, True])
+def test_chunk_observer_preserves_logits_bias_and_gradients(bias_kind, trainable, all_ignored):
+    """Observation must preserve both bias forms, including bias-only gradients."""
+    from torch.autograd.graph import saved_tensors_hooks
+
+    torch.manual_seed(712)
+    hidden = torch.randn(2, 7, 8)
+    weight = torch.randn(19, 8)
+    bias = torch.randn(2, 7, 4 if bias_kind == "factored" else 19)
+    tensors = [hidden, weight, bias]
+    if bias_kind == "factored":
+        tensors.append(torch.randn(19, 4))
+    labels = torch.randint(0, 19, (2, 7))
+    labels[:, 1:3] = -100
+    if all_ignored:
+        labels.fill_(-100)
+
+    results, captured = [], []
+    for observe in (False, True):
+        args = [
+            tensor.detach().clone().requires_grad_(trainable == "all" or (trainable == "bias_only" and i >= 2))
+            for i, tensor in enumerate(tensors)
+        ]
+        context = _capture_chunk_tokens(captured) if observe else nullcontext()
+        with context, saved_tensors_hooks(lambda tensor: tensor.detach().clone(), lambda tensor: tensor):
+            loss = resolve_op("cross_entropy_loss", "standard", "chunk_loss").wrapper(
+                args[0], labels, args[1], *args[2:], chunk_size=3
+            )
+            if loss.requires_grad:
+                loss.backward()
+        results.append((loss.detach(), *(arg.grad for arg in args)))
+
+    for actual, expected in zip(results[1], results[0], strict=True):
+        if expected is None:
+            assert actual is None
+        else:
+            torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    logits = F.linear(hidden, weight)
+    logits = logits + (F.linear(bias, tensors[3]) if bias_kind == "factored" else bias)
+    expected_tokens = F.cross_entropy(logits.reshape(-1, 19), labels.reshape(-1), reduction="none").reshape_as(labels)
+    assert len(captured) == 1
+    assert not captured[0].requires_grad
+    assert not torch._C._functorch.is_functorch_wrapped_tensor(captured[0])
+    torch.testing.assert_close(captured[0], expected_tokens)
+
+
 class _FullWeightAddCounter(TorchDispatchMode):
     def __init__(self, weight_shape):
         super().__init__()