SFT learner fixes: training mode for gradient checkpointing, Ascend OOM detection, real-loss autobatch probe

tests/test_accel.py

@@ -359,3 +359,24 @@ def test_npu_without_torch_npu_is_a_clear_error(monkeypatch):
 def test_cpu_family_binds_nothing(fake_cuda):
     assert accel.bind_family_device("cpu", local_rank=5) == torch.device("cpu")
     assert fake_cuda.bound == []
+
+
+def test_is_oom_matches_the_vendor_exception(npu):
+    assert accel.is_oom(FakeOutOfMemory("npu"), device("npu", 0))
+    assert accel.is_oom(torch.cuda.OutOfMemoryError("cuda"), device("cuda", 0))
+
+
+def test_is_oom_matches_a_message_only_allocator_refusal(monkeypatch):
+    # torch_npu exports no OutOfMemoryError: its allocator raises a bare
+    # RuntimeError, which an `except oom_error(device)` would let escape.
+    monkeypatch.setattr(accel, "backend", lambda _device: SimpleNamespace())
+    refusal = RuntimeError(
+        "NPU out of memory. Tried to allocate 37.09 GiB (NPU 0; 60.96 GiB total capacity)"
+    )
+    assert accel.is_oom(refusal, device("npu", 0))
+
+
+def test_is_oom_does_not_swallow_an_unrelated_runtime_error(monkeypatch):
+    monkeypatch.setattr(accel, "backend", lambda _device: SimpleNamespace())
+    assert not accel.is_oom(RuntimeError("HCCL communicator was aborted"), device("npu", 0))
+    assert not accel.is_oom(ValueError("not an allocator problem"), device("npu", 0))

tests/test_autobatch.py

@@ -298,3 +298,67 @@ def test_probe_preserves_cpu_and_cuda_rng():
 
     assert torch.equal(torch.get_rng_state(), cpu_rng)
     assert torch.equal(torch.cuda.get_rng_state(device), cuda_rng)
+
+
+def test_probe_backs_off_on_a_message_only_npu_oom(monkeypatch):
+    """torch_npu raises a bare RuntimeError; the probe must treat it as OOM."""
+    from yeto import accel
+
+    sizes = []
+
+    def probe(model, params, opt, seq_len, vocab, device, mb, *, loss_forward=None):
+        sizes.append(mb)
+        if mb > 2:
+            raise RuntimeError(
+                "NPU out of memory. Tried to allocate 37.09 GiB (NPU 0; 60.96 GiB total capacity)"
+            )
+
+    monkeypatch.setattr(autobatch, "_probe_once", probe)
+    monkeypatch.setattr(accel, "backend", lambda _device: SimpleNamespace(empty_cache=lambda: None))
+    micro_batch = autobatch.resolve_micro_batch_size(
+        _args(grad_accum=8), None, {}, None, Tok(), SimpleNamespace(type="npu"), 1
+    )
+
+    assert sizes == [1, 2, 4]
+    assert micro_batch == 2
+    assert autobatch.exact_grad_accum(8, micro_batch) == 4
+
+
+def test_probe_reraises_a_non_oom_runtime_error(monkeypatch):
+    from yeto import accel
+
+    def probe(*a, **k):
+        raise RuntimeError("HCCL communicator was aborted")
+
+    monkeypatch.setattr(autobatch, "_probe_once", probe)
+    monkeypatch.setattr(accel, "backend", lambda _device: SimpleNamespace(empty_cache=lambda: None))
+    with pytest.raises(RuntimeError, match="HCCL communicator was aborted"):
+        autobatch.resolve_micro_batch_size(
+            _args(), None, {}, None, Tok(), SimpleNamespace(type="npu"), 1
+        )
+
+
+def test_native_probe_runs_the_real_loss_not_a_surrogate():
+    """The probe must measure sft_loss itself: F.cross_entropy keeps a
+    second full-size f32 tensor alive that a logits-copy stand-in misses."""
+    model = TinyCausalLM()
+    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, foreach=False, fused=False)
+    params = dict(model.named_parameters())
+    seen = {}
+    real = autobatch.sft_loss
+
+    def spy(logits, labels, loss_function="cross_entropy", weights=None):
+        seen["loss_function"] = loss_function
+        seen["logits_shape"] = tuple(logits.shape)
+        seen["weights_all_ones"] = weights is not None and bool((weights == 1).all())
+        return real(logits, labels, loss_function, weights)
+
+    autobatch.sft_loss = spy
+    try:
+        autobatch._probe_once(model, params, opt, 6, 30, torch.device("cpu"), 2)
+    finally:
+        autobatch.sft_loss = real
+
+    assert seen["loss_function"] == "cross_entropy"
+    assert seen["logits_shape"] == (2, 6, 32)
+    assert seen["weights_all_ones"]

tests/test_chunked_ce.py

@@ -193,3 +193,27 @@ def test_validate_kernel_request_chunked_any_device():
     with pytest.raises(ValueError):
         validate_kernel_request("chunked", "cross_entropy", torch.device("cpu"), torch.float32,
                                 shard="fsdp")
+
+
+@pytest.mark.parametrize(
+    "backend, expected",
+    [
+        ("native", "torch-cross-entropy"),
+        ("liger", "liger-fused-linear-cross-entropy"),
+        ("chunked", chunked_ce.CHUNKED_LOSS_IMPLEMENTATION),
+    ],
+)
+def test_kernel_recipe_names_the_loss_that_runs(backend, expected):
+    from yeto.learner import loss_implementation_for
+
+    assert loss_implementation_for(backend) == expected
+
+
+def test_every_kernel_backend_has_a_recipe_loss_name():
+    from yeto.causal_kernels import KERNEL_BACKENDS
+    from yeto.learner import loss_implementation_for
+
+    names = {loss_implementation_for(b) for b in KERNEL_BACKENDS}
+    assert len(names) == len(KERNEL_BACKENDS)
+    with pytest.raises(ValueError):
+        loss_implementation_for("bogus")

tests/test_learner_units.py

@@ -5,6 +5,7 @@
 import pytest
 import torch
 
+import yeto.learner as learner_module
 from yeto.learner import (
     _loss_metric_dtype,
     allreduce_trainable_grads,
@@ -1172,3 +1173,121 @@ def test_throughput_is_measured_per_window_not_cumulatively():
     assert first["train/raw_tokens_total"] == counters.raw_tokens // 2
     window_tokens = second["train/tokens_per_sec"] * second["train/sec_per_step"] * 10
     assert window_tokens == pytest.approx(counters.raw_tokens / 2, rel=1e-6)
+
+
+def _checkpointing_args(mode, tuning="lora"):
+    return SimpleNamespace(gradient_checkpointing=mode, tuning=tuning)
+
+
+class _RecordingModel(torch.nn.Module):
+    """Mimics the Transformers/peft surface enter_training_mode touches."""
+
+    def __init__(self):
+        super().__init__()
+        self.checkpointing_kwargs = None
+        self.input_grads_enabled = False
+        self.eval()
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        # Transformers only recomputes when the module is in training mode, so
+        # a checkpointing setup done in eval mode is silently ineffective.
+        assert self.training, "checkpointing enabled on an eval-mode model"
+        self.checkpointing_kwargs = gradient_checkpointing_kwargs
+
+    def enable_input_require_grads(self):
+        self.input_grads_enabled = True
+
+
+def test_entering_training_mode_precedes_checkpointing_setup():
+    model = _RecordingModel()
+    assert not model.training
+
+    enabled = learner_module.enter_training_mode(
+        model, _checkpointing_args("on"), torch.device("cpu")
+    )
+
+    assert enabled
+    assert model.training
+    assert model.checkpointing_kwargs == {"use_reentrant": False}
+    assert model.input_grads_enabled
+
+
+def test_training_mode_is_entered_even_without_checkpointing():
+    model = _RecordingModel()
+
+    enabled = learner_module.enter_training_mode(
+        model, _checkpointing_args("off"), torch.device("cpu")
+    )
+
+    assert not enabled
+    assert model.training
+    assert model.checkpointing_kwargs is None
+    assert not model.input_grads_enabled
+
+
+def test_full_tuning_skips_the_input_grad_hook():
+    model = _RecordingModel()
+
+    assert learner_module.enter_training_mode(
+        model, _checkpointing_args("on", tuning="full"), torch.device("cpu")
+    )
+    assert model.checkpointing_kwargs == {"use_reentrant": False}
+    assert not model.input_grads_enabled
+
+
+@pytest.mark.parametrize(
+    "free,total,expected",
+    [(10, 100, True), (90, 100, False), (None, None, False)],
+)
+def test_auto_checkpointing_follows_free_memory(monkeypatch, free, total, expected):
+    monkeypatch.setattr(
+        learner_module.accel,
+        "mem_get_info",
+        lambda _device: None if free is None else (free, total),
+    )
+    model = _RecordingModel()
+
+    assert (
+        learner_module.enter_training_mode(
+            model, _checkpointing_args("auto"), torch.device("cpu")
+        )
+        is expected
+    )
+    assert model.training
+
+
+class _FakeClock:
+    def __init__(self):
+        self.t = 0.0
+
+    def __call__(self):
+        return self.t
+
+
+def test_step_window_timer_excludes_first_step_warmup():
+    from yeto.learner import StepWindowTimer
+
+    clock = _FakeClock()
+    timer = StepWindowTimer(0, 0, clock=clock)
+    clock.t = 200.0  # first step: dataset build + compile
+    assert timer.first_step_done(1, 100) == 200.0
+    for step in range(2, 11):
+        clock.t += 5.0
+        assert timer.first_step_done(step, step * 100) is None
+    n, dt, toks = timer.window(10, 1000)
+    assert (n, dt, toks) == (9, 45.0, 900)
+    assert dt / n == 5.0  # steady state, not (200 + 45) / 10
+    clock.t += 50.0
+    assert timer.window(20, 2000) == (10, 50.0, 1000)
+
+
+def test_step_window_timer_resume_counts_from_resumed_step():
+    from yeto.learner import StepWindowTimer
+
+    clock = _FakeClock()
+    timer = StepWindowTimer(500, 7000, clock=clock)
+    assert timer.first_step_done(500, 7000) is None  # no step taken yet
+    clock.t = 30.0
+    assert timer.first_step_done(501, 7100) == 30.0
+    clock.t = 39.0
+    assert timer.window(510, 8000) == (9, 9.0, 900)

yeto/accel.py

@@ -208,6 +208,28 @@ def oom_error(device: torch.device) -> type[BaseException]:
     return getattr(backend(device), "OutOfMemoryError", torch.cuda.OutOfMemoryError)
 
 
+# torch_npu's caching allocator raises a bare RuntimeError ("NPU out of
+# memory. Tried to allocate ...") and `torch.npu` exports no
+# `OutOfMemoryError`, so an `except oom_error(device)` written for CUDA never
+# matches on Ascend and an allocator OOM escapes as a hard failure. Probes
+# must ask `is_oom` instead of naming an exception class.
+_OOM_MESSAGE = "out of memory"
+
+
+def is_oom(exc: BaseException, device: torch.device) -> bool:
+    """Whether ``exc`` is ``device``'s allocator refusing an allocation.
+
+    Covers the vendor exception class where one exists, torch's generic
+    ``OutOfMemoryError``, and the message-only form that torch_npu raises.
+    """
+    if isinstance(exc, oom_error(device)):
+        return True
+    generic = getattr(torch, "OutOfMemoryError", None)
+    if generic is not None and isinstance(exc, generic):
+        return True
+    return isinstance(exc, RuntimeError) and _OOM_MESSAGE in str(exc).lower()
+
+
 def loss_metric_dtype(device: torch.device) -> torch.dtype:
     """Highest-precision loss telemetry dtype supported by its collective.
 

yeto/autobatch.py

@@ -26,6 +26,7 @@
 import torch.distributed as dist
 
 from . import accel
+from .losses import sft_loss
 
 log = logging.getLogger("learner")
 
@@ -145,9 +146,21 @@ def _probe_once(
             ids = torch.randint(0, vocab, (micro_batch, seq_len), device=device)
             if loss_forward is None:
                 out = model(input_ids=ids)
-                # Mirror native sft_loss's memory peak: the shifted f32
-                # logits copy dominates the loss computation.
-                loss = out.logits[:, :-1].float().sum()
+                # Run the native loss itself rather than a surrogate for it.
+                # A hand-rolled stand-in underestimates the peak: the f32
+                # logits copy is not the only full-size tensor, because
+                # F.cross_entropy keeps its log_softmax output alive for the
+                # backward. On a 151k-token vocabulary that missing tensor is
+                # another ~20 GB at micro-batch 4 x 8k tokens, enough for the
+                # probe to admit a size whose first real step then OOMs.
+                # All-ones weights make every position a target, which is the
+                # worst case over real packed blocks.
+                loss, _ = sft_loss(
+                    out.logits,
+                    ids,
+                    "cross_entropy",
+                    torch.ones_like(ids, dtype=out.logits.dtype),
+                )
             else:
                 # Optimized backends can have a materially different peak
                 # (notably fused linear CE, which never materializes logits).
@@ -219,7 +232,9 @@ def resolve_micro_batch_size(
                 size,
                 loss_forward=loss_forward,
             )
-        except accel.oom_error(device):
+        except Exception as exc:  # noqa: BLE001 -- re-raised unless it is an OOM
+            if not accel.is_oom(exc, device):
+                raise
             ok = False
         accel.empty_cache(device)
         if world > 1:

yeto/chunked_ce.py

@@ -24,6 +24,7 @@
 from torch.utils.checkpoint import checkpoint
 
 DEFAULT_LOSS_CHUNK_TOKENS = 4096
+CHUNKED_LOSS_IMPLEMENTATION = "chunked-linear-cross-entropy"
 
 # Config attributes that change how HF turns hidden states into logits.
 # The chunked lane applies only a plain linear head, so it refuses these.

yeto/diffusion/learner.py

@@ -3020,7 +3020,9 @@ def _probe_max_micro_batch(
         ok = True
         try:
             _probe_diffusion_once(pipe, params, opt, probe_rows, args, device, size, adapter)
-        except accel.oom_error(device):
+        except Exception as exc:  # noqa: BLE001 -- re-raised unless it is an OOM
+            if not accel.is_oom(exc, device):
+                raise
             ok = False
         opt.zero_grad(set_to_none=True)
         accel.empty_cache(device)