SFT: device-agnostic chunked linear+CE loss and --gradient-checkpointing forwarding

docs/A100_KERNELS.md

@@ -8,7 +8,9 @@ The supported controls are:
 | control | values | default | meaning |
 |---|---|---|---|
 | `--attention-backend` | `auto`, `sdpa`, `flash-attn-2` | `auto` | Hugging Face attention implementation |
-| `--kernel-backend` | `native`, `liger` | `native` | SFT loss implementation; model layers remain native; fused requires LoRA/DDP |
+| `--kernel-backend` | `native`, `liger`, `chunked` | `native` | SFT loss implementation; model layers remain native; fused requires LoRA/DDP; `chunked` is device-agnostic (see [ASCEND.md](ASCEND.md#long-context-sft-loss)) |
+| `--gradient-checkpointing` | `auto`, `on`, `off` | `auto` | learner activation recomputation; forwarded only when not `auto` |
+| `--loss-chunk-tokens` | integer | `4096` | logits rows per chunk for `--kernel-backend chunked` |
 
 An explicit request is fail-closed. A missing or differently-versioned
 dependency, an unsupported model implementation, FlashAttention with FP32,

docs/ASCEND.md

@@ -108,10 +108,50 @@ actually been executed on a 910x.
 | diffusion `--shard fsdp` | ✅ (FSDP2/hccl) | `maybe_wrap_for_distributed` |
 | diffusion adapter reload and image/video sampling | ✅ | `yeto.diffusion.sample` |
 | causal-LM `--base-quantization nf4` | ❌ refused | `validate_base_quantization` |
+| `--kernel-backend chunked` (long-context SFT loss) | ✅ | `yeto/chunked_ce.py` |
 | `--kernel-backend liger` | ❌ refused | Triton; `validate_kernel_request` |
 | `--attention-backend flash-attn-2` | ❌ refused | `attention_load_kwargs` |
 | causal-LM `--shard fsdp` | ❌ refused | no validated sharding evidence |
 
+## Long-context SFT loss
+
+With the native loss the model returns full `[seq, vocab]` logits, and the
+cross-entropy adds an FP32 copy plus softmax intermediates. For Qwen3 (vocab
+~151k) at a 28k-token sequence that is tens of GB per sample, which exceeds a
+64 GB 910B card even with LoRA, gradient checkpointing, and micro-batch 1.
+
+`--kernel-backend chunked` avoids this on every device family (NPU, CUDA,
+MPS, CPU) using plain PyTorch ops. The model runs to its final hidden states
+(`logits_to_keep` is an empty index, so the head emits nothing), then the loss
+computes `hidden_chunk @ lm_head.weight^T` and CE over `--loss-chunk-tokens`
+target tokens at a time. Each chunk runs under non-reentrant activation
+checkpointing, so backward recomputes one chunk's logits at a time; only
+positions with non-zero loss weight are fed to the head. Peak loss memory is
+roughly `chunk * vocab * 12` bytes (bf16 logits, FP32 logits, FP32 grad):
+about 7 GB at 4096 tokens and 3.7 GB at 2048 for a 151k vocabulary.
+
+The loss is identical to the native `cross_entropy` SFT loss: same label
+shift, assistant mask and fractional token weights, FP32 CE, local token-sum
+reduction, and target-token count. Tied input/output embeddings are
+supported (gradients reach the shared weight). Requirements: the built-in
+`cross_entropy` loss, `--shard ddp`, no quantized base, a plain `nn.Linear`
+head (models that soft-cap or scale logits, or LoRA targeting `lm_head`, are
+refused), and a Transformers model whose forward accepts `logits_to_keep`.
+
+Suggested 910B flags for Qwen3-0.6B at 28k context:
+
+```
+--kernel-backend chunked --loss-chunk-tokens 4096 \
+--tuning lora --shard ddp --micro-batch-size 1 --gradient-checkpointing on \
+--attention-backend sdpa --seq-len 28000
+```
+
+Lower `--loss-chunk-tokens` to 2048 if the card is still tight; the cost is
+more, smaller head matmuls. `--gradient-checkpointing` is accepted by both the launcher and
+`yeto.learner`. Its `auto` default enables recomputation only when the base
+occupies more than half of device memory, which a 0.6B model on a 64 GB card
+never does, so pass `on` explicitly for long contexts.
+
 The causal-LM NF4, Liger, and FlashAttention paths still reject every non-CUDA
 device. Diffusion quantization is a separate Diffusers load path: it is enabled
 only when the installed bitsandbytes explicitly reports NPU support. The

tests/test_chunked_ce.py

@@ -0,0 +1,195 @@
+import weakref
+
+import pytest
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+
+transformers = pytest.importorskip("transformers")
+
+from yeto import chunked_ce  # noqa: E402
+from yeto.causal_kernels import validate_kernel_request  # noqa: E402
+from yeto.losses import sft_loss  # noqa: E402
+
+VOCAB = 512
+HIDDEN = 64
+
+
+def _tiny_qwen3(dtype=torch.float32, tie=True, seed=0):
+    torch.manual_seed(seed)
+    config = transformers.Qwen3Config(
+        vocab_size=VOCAB,
+        hidden_size=HIDDEN,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        max_position_embeddings=256,
+        tie_word_embeddings=tie,
+    )
+    config._attn_implementation = "sdpa"
+    model = transformers.Qwen3ForCausalLM(config).to(dtype)
+    return model
+
+
+def _batch(seq=97, batch=2, seed=1):
+    g = torch.Generator().manual_seed(seed)
+    ids = torch.randint(0, VOCAB, (batch, seq), generator=g)
+    weights = (torch.rand(batch, seq, generator=g) > 0.4).float()
+    return ids, weights
+
+
+def _grads(model):
+    return {n: p.grad.detach().clone() for n, p in model.named_parameters() if p.grad is not None}
+
+
+def _native(model, ids, weights):
+    model.zero_grad(set_to_none=True)
+    loss, n = sft_loss(model(input_ids=ids).logits, ids, "cross_entropy", weights)
+    loss.backward()
+    return loss.detach(), n, _grads(model)
+
+
+def _chunked(model, ids, weights, chunk):
+    model.zero_grad(set_to_none=True)
+    loss, n = chunked_ce.chunked_sft_forward(model, ids, weights, chunk_tokens=chunk)
+    loss.backward()
+    return loss.detach(), n, _grads(model)
+
+
+def _assert_match(a, b, rtol, atol):
+    assert int(a[1]) == int(b[1])
+    torch.testing.assert_close(a[0].float(), b[0].float(), rtol=rtol, atol=atol)
+    assert a[2].keys() == b[2].keys()
+    for name in a[2]:
+        torch.testing.assert_close(
+            a[2][name].float(), b[2][name].float(), rtol=rtol, atol=atol, msg=name
+        )
+
+
+@pytest.mark.parametrize("tie", [True, False])
+@pytest.mark.parametrize("chunk", [1, 7, 64, 4096])
+def test_fp32_loss_and_grads_match_native(tie, chunk):
+    model = _tiny_qwen3(tie=tie)
+    ids, weights = _batch()
+    _assert_match(_native(model, ids, weights), _chunked(model, ids, weights, chunk), 1e-5, 1e-5)
+
+
+def test_fractional_weights_and_no_weights_match():
+    model = _tiny_qwen3()
+    ids, weights = _batch()
+    weights = weights * torch.rand_like(weights) * 2
+    _assert_match(_native(model, ids, weights), _chunked(model, ids, weights, 13), 1e-5, 1e-5)
+    _assert_match(_native(model, ids, None), _chunked(model, ids, None, 13), 1e-5, 1e-5)
+
+
+def test_ignore_index_labels_match():
+    # Input IDs never contain -100, so exercise ignore_index on the helper.
+    model = _tiny_qwen3()
+    hidden = torch.randn(50, 64, dtype=torch.float32, requires_grad=True)
+    weight = model.lm_head.weight
+    targets = torch.randint(0, VOCAB, (50,))
+    targets[::3] = -100
+    w = torch.rand(50)
+    ref = (torch.nn.functional.cross_entropy(
+        hidden @ weight.T, targets, ignore_index=-100, reduction="none") * w).sum()
+    got = chunked_ce.chunked_linear_cross_entropy(hidden, weight, targets, w, chunk_tokens=8)
+    torch.testing.assert_close(got, ref)
+
+
+def test_all_masked_batch_backward():
+    model = _tiny_qwen3()
+    ids, _ = _batch()
+    weights = torch.zeros_like(ids, dtype=torch.float32)
+    loss, n = chunked_ce.chunked_sft_forward(model, ids, weights, chunk_tokens=8)
+    loss.backward()
+    assert float(loss) == 0.0 and int(n) == 0
+
+
+def test_bf16_matches_native():
+    model = _tiny_qwen3(dtype=torch.bfloat16)
+    ids, weights = _batch()
+    _assert_match(_native(model, ids, weights), _chunked(model, ids, weights, 32), 2e-2, 2e-2)
+
+
+def test_lora_with_gradient_checkpointing_matches():
+    peft = pytest.importorskip("peft")
+    model = _tiny_qwen3()
+    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    model = peft.get_peft_model(
+        model, peft.LoraConfig(r=4, lora_alpha=8, target_modules="all-linear", init_lora_weights=False)
+    )
+    model.train()
+    ids, weights = _batch()
+    chunked_ce.require_chunked_ce_support(model)
+    a = _native(model, ids, weights)
+    b = _chunked(model, ids, weights, 16)
+    assert a[2] and all("lora" in k for k in a[2])
+    _assert_match(a, b, 1e-5, 1e-5)
+
+
+def test_lora_on_lm_head_is_refused():
+    peft = pytest.importorskip("peft")
+    model = peft.get_peft_model(
+        _tiny_qwen3(tie=False), peft.LoraConfig(r=4, target_modules=["lm_head", "q_proj"])
+    )
+    with pytest.raises(ValueError, match="plain"):
+        chunked_ce.require_chunked_ce_support(model)
+
+
+class _VocabTensorTracker(TorchDispatchMode):
+    """Record the row count of every vocab-wide tensor any op produces and the
+    peak of live vocab-wide bytes, across forward and backward."""
+
+    def __init__(self):
+        super().__init__()
+        self.max_rows = 0
+        self.live = {}
+        self.peak_bytes = 0
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        out = func(*args, **(kwargs or {}))
+        for t in torch.utils._pytree.tree_leaves(out):
+            if (
+                isinstance(t, torch.Tensor)
+                and t.dim()
+                and t.shape[-1] == VOCAB
+                and t.numel() != VOCAB * HIDDEN  # lm_head weight views
+            ):
+                self.max_rows = max(self.max_rows, t.numel() // VOCAB)
+                key = id(t)
+                self.live[key] = t.numel() * t.element_size()
+                weakref.finalize(t, self.live.pop, key, None)
+                self.peak_bytes = max(self.peak_bytes, sum(self.live.values()))
+        return out
+
+
+def _track(model, fn):
+    model.zero_grad(set_to_none=True)
+    tracker = _VocabTensorTracker()
+    with tracker:
+        loss, _ = fn()
+        loss.backward()
+    return tracker
+
+
+def test_full_sequence_logits_never_materialized():
+    model = _tiny_qwen3(tie=False)
+    ids, _ = _batch(seq=257, batch=2)
+    weights = torch.ones_like(ids, dtype=torch.float32)
+    chunk = 16
+    native = _track(model, lambda: sft_loss(model(input_ids=ids).logits, ids, "cross_entropy", weights))
+    chunked = _track(model, lambda: chunked_ce.chunked_sft_forward(model, ids, weights, chunk_tokens=chunk))
+    assert native.max_rows >= 2 * 256
+    assert chunked.max_rows <= chunk
+    assert chunked.peak_bytes * 10 < native.peak_bytes
+
+
+def test_validate_kernel_request_chunked_any_device():
+    validate_kernel_request("chunked", "cross_entropy", torch.device("cpu"), torch.float32,
+                            tuning="lora", shard="ddp")
+    with pytest.raises(ValueError):
+        validate_kernel_request("chunked", "custom:x.py", torch.device("cpu"), torch.float32)
+    with pytest.raises(ValueError):
+        validate_kernel_request("chunked", "cross_entropy", torch.device("cpu"), torch.float32,
+                                shard="fsdp")

tests/test_island_backend.py

@@ -88,6 +88,39 @@ def test_torch_backend_installs_only_explicit_kernel_dependencies():
     assert "--no-build-isolation" in task.setup
 
 
+def test_chunked_backend_forwards_chunk_size_without_extra_dependencies():
+    task = make_learner_task(
+        _args(island_backend="torch", kernel_backend="chunked", shard="ddp",
+              loss_chunk_tokens=2048),
+        _SPEC,
+        0,
+        1,
+        "1.2.3.4:29400",
+    )
+    assert "--kernel-backend chunked --loss-chunk-tokens 2048" in task.run
+    assert "liger-kernel" not in task.setup
+
+
+def test_gradient_checkpointing_forwarded_only_when_explicit():
+    default = make_learner_task(_args(island_backend="torch"), _SPEC, 0, 1, "1.2.3.4:29400")
+    assert "--gradient-checkpointing" not in default.run
+    for mode in ("on", "off"):
+        task = make_learner_task(
+            _args(island_backend="torch", gradient_checkpointing=mode),
+            _SPEC, 0, 1, "1.2.3.4:29400",
+        )
+        assert f"--gradient-checkpointing {mode}" in task.run
+
+
+def test_chunked_launcher_rejects_fsdp_and_custom_loss():
+    with pytest.raises(ValueError, match="only --shard ddp"):
+        make_learner_task(_args(kernel_backend="chunked", shard="fsdp"),
+                          _SPEC, 0, 1, "1.2.3.4:29400")
+    with pytest.raises(ValueError, match="only the built-in cross_entropy"):
+        make_learner_task(_args(kernel_backend="chunked", loss_function="pickle:loss.pkl"),
+                          _SPEC, 0, 1, "1.2.3.4:29400")
+
+
 def test_liger_launcher_rejects_non_builtin_loss():
     with pytest.raises(ValueError, match="only the built-in cross_entropy"):
         make_learner_task(

yeto/causal_kernels.py

@@ -27,7 +27,7 @@
 )
 
 ATTENTION_BACKENDS = ("auto", "sdpa", "flash-attn-2")
-KERNEL_BACKENDS = ("native", "liger")
+KERNEL_BACKENDS = ("native", "liger", "chunked")
 
 NATIVE_LAYER_BACKEND = "transformers-native"
 NATIVE_LOSS_IMPLEMENTATION = "torch-cross-entropy"
@@ -840,6 +840,25 @@ def validate_kernel_request(
         raise ValueError(f"unknown kernel backend {kernel_backend!r}; choose from {KERNEL_BACKENDS}")
     if kernel_backend == "native":
         return
+    if kernel_backend == "chunked":
+        # Plain PyTorch ops: valid on every device family.
+        if loss_function != "cross_entropy":
+            raise ValueError(
+                "--kernel-backend chunked supports only the built-in "
+                "cross_entropy loss; use --kernel-backend native for custom, "
+                "pickled, or RL losses"
+            )
+        if base_quantization != "none":
+            raise ValueError(
+                "--kernel-backend chunked does not support a quantized base; "
+                "use --base-quantization none or --kernel-backend native"
+            )
+        if shard is not None and shard != "ddp":
+            raise ValueError(
+                "--kernel-backend chunked supports only --shard ddp; FSDP "
+                "shards lm_head outside the wrapped forward"
+            )
+        return
     if device.type != "cuda":
         raise RuntimeError("--kernel-backend liger fused-linear-CE requires CUDA")
     _require_a100(device)