diff --git a/docs/A100_KERNELS.md b/docs/A100_KERNELS.md index 86b80710..0461d4b7 100644 --- a/docs/A100_KERNELS.md +++ b/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, diff --git a/docs/ASCEND.md b/docs/ASCEND.md index bf6aaece..1b0a23f6 100644 --- a/docs/ASCEND.md +++ b/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 diff --git a/tests/test_chunked_ce.py b/tests/test_chunked_ce.py new file mode 100644 index 00000000..614c53dc --- /dev/null +++ b/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") diff --git a/tests/test_island_backend.py b/tests/test_island_backend.py index a746777b..265d8a01 100644 --- a/tests/test_island_backend.py +++ b/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( diff --git a/yeto/causal_kernels.py b/yeto/causal_kernels.py index 277cd518..c4481dd6 100644 --- a/yeto/causal_kernels.py +++ b/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) diff --git a/yeto/chunked_ce.py b/yeto/chunked_ce.py new file mode 100644 index 00000000..4c379539 --- /dev/null +++ b/yeto/chunked_ce.py @@ -0,0 +1,212 @@ +"""Device-agnostic chunked linear + cross-entropy SFT loss. + +The native loss path materializes the full ``[batch, seq, vocab]`` logits +tensor (plus its FP32 copy and the softmax intermediates of the backward +pass). For long sequences and large vocabularies that dominates peak memory. + +This lane runs the model only up to its final hidden states and then computes +``hidden_chunk @ lm_head.weight^T -> cross_entropy`` over bounded token +chunks. Each chunk runs under non-reentrant activation checkpointing, so the +autograd graph retains only the chunk's hidden slice and its scalar loss; the +backward pass recomputes one chunk's logits at a time. The arithmetic mirrors +:func:`yeto.losses.sft_loss` exactly (same label shift, ignore index, token +weights, FP32 CE, local token-SUM loss and positive-weight token count). + +Only plain PyTorch ops are used, so it runs on CUDA, Ascend NPU, MPS and CPU. +""" + +from __future__ import annotations + +import inspect + +import torch +import torch.nn.functional as F +from torch.utils.checkpoint import checkpoint + +DEFAULT_LOSS_CHUNK_TOKENS = 4096 + +# Config attributes that change how HF turns hidden states into logits. +# The chunked lane applies only a plain linear head, so it refuses these. +_UNSUPPORTED_HEAD_CONFIG = ( + "final_logit_softcapping", + "logit_scale", + "logits_scaling", + "output_multiplier_scale", +) + + +def _unwrap(model): + """Return the underlying HF causal-LM module (through DDP/FSDP/PEFT).""" + seen = set() + while id(model) not in seen: + seen.add(id(model)) + if hasattr(model, "module") and isinstance(model.module, torch.nn.Module): + model = model.module + continue + if hasattr(model, "get_base_model"): + model = model.get_base_model() + continue + break + return model + + +def _decoder(causal_lm): + get_decoder = getattr(causal_lm, "get_decoder", None) + decoder = get_decoder() if callable(get_decoder) else None + if decoder is None: + decoder = getattr(causal_lm, "model", None) + if not isinstance(decoder, torch.nn.Module): + raise RuntimeError( + "--kernel-backend chunked requires a causal LM exposing its decoder " + "via get_decoder() or .model" + ) + return decoder + + +def require_chunked_ce_support(model) -> dict: + """Validate that the model's head is a plain linear layer we can chunk. + + Returns a small description for logging. Raises with a clear message on + unsupported heads (soft-capped/scaled logits, adapted lm_head, models + without ``logits_to_keep``). + """ + causal_lm = _unwrap(model) + config = getattr(causal_lm, "config", None) + for name in _UNSUPPORTED_HEAD_CONFIG: + value = getattr(config, name, None) + if value not in (None, 1, 1.0): + raise ValueError( + f"--kernel-backend chunked does not support models with " + f"config.{name}={value!r}; use --kernel-backend native" + ) + head = causal_lm.get_output_embeddings() + if type(head) is not torch.nn.Linear: + raise ValueError( + "--kernel-backend chunked requires lm_head to be a plain " + f"torch.nn.Linear, got {type(head).__name__}; an adapted or " + "quantized lm_head (e.g. LoRA targeting lm_head) is unsupported, " + "use --kernel-backend native" + ) + params = inspect.signature(causal_lm.forward).parameters + if "logits_to_keep" not in params: + raise ValueError( + f"--kernel-backend chunked requires a forward() accepting " + f"logits_to_keep; {type(causal_lm).__name__} does not" + ) + _decoder(causal_lm) + input_emb = causal_lm.get_input_embeddings() + tied = input_emb is not None and getattr(input_emb, "weight", None) is head.weight + return { + "loss_backend": "chunked", + "model_class": type(causal_lm).__name__, + "tied_embeddings": bool(tied), + } + + +def _chunk_loss(hidden, weight, bias, targets, token_weights): + logits = F.linear(hidden, weight, bias).float() + per_token = F.cross_entropy(logits, targets, ignore_index=-100, reduction="none") + return (per_token * token_weights).sum() + + +def chunked_linear_cross_entropy( + hidden: torch.Tensor, + weight: torch.Tensor, + targets: torch.Tensor, + token_weights: torch.Tensor, + *, + bias: torch.Tensor | None = None, + chunk_tokens: int = DEFAULT_LOSS_CHUNK_TOKENS, +) -> torch.Tensor: + """Sum of ``token_weights * CE(hidden @ weight^T + bias, targets)``. + + hidden: (N, H); targets: (N,) with -100 ignored; token_weights: (N,) FP32. + At most ``chunk_tokens`` rows of logits exist at any time, in forward and + in backward. + """ + if chunk_tokens < 1: + raise ValueError(f"chunk_tokens must be positive, got {chunk_tokens}") + total = hidden.new_zeros((), dtype=torch.float32) + needs_grad = torch.is_grad_enabled() and ( + hidden.requires_grad + or weight.requires_grad + or (bias is not None and bias.requires_grad) + ) + for start in range(0, hidden.shape[0], chunk_tokens): + stop = start + chunk_tokens + args = (hidden[start:stop], weight, bias, targets[start:stop], token_weights[start:stop]) + if needs_grad: + part = checkpoint(_chunk_loss, *args, use_reentrant=False) + else: + part = _chunk_loss(*args) + total = total + part + return total + + +def chunked_sft_forward( + model, + input_ids: torch.Tensor, + weights: torch.Tensor | None = None, + *, + chunk_tokens: int = DEFAULT_LOSS_CHUNK_TOKENS, +): + """Run the model to its final hidden states and return the chunked SFT loss. + + Semantics match ``sft_loss(model(input_ids).logits, input_ids, + "cross_entropy", weights)``: returns (local token-SUM loss, number of + shifted positions with positive weight). + + The full model (including any DDP/FSDP/PEFT wrapper) is called so wrapper + hooks still run; ``logits_to_keep`` is an empty index so the HF head + produces no logits, and a forward hook on the decoder captures its final + (normalized) hidden states. + """ + causal_lm = _unwrap(model) + decoder = _decoder(causal_lm) + head = causal_lm.get_output_embeddings() + captured = [] + + def _capture(_module, _inputs, output): + hidden = getattr(output, "last_hidden_state", None) + if hidden is None: + hidden = output[0] if isinstance(output, (tuple, list)) else output + captured.append(hidden) + + handle = decoder.register_forward_hook(_capture) + try: + model( + input_ids=input_ids, + use_cache=False, + logits_to_keep=torch.empty(0, dtype=torch.long, device=input_ids.device), + ) + finally: + handle.remove() + if len(captured) != 1: + raise RuntimeError( + f"expected one decoder forward, observed {len(captured)}" + ) + hidden = captured[0] + + labels = input_ids + shift_labels = labels[:, 1:] + mask = shift_labels != -100 + w = mask.to(torch.float32) + if weights is not None: + w = w * weights[:, 1:].to(torch.float32) + n_tokens = (w > 0).sum() + + flat_w = w.reshape(-1) + keep = torch.nonzero(flat_w != 0, as_tuple=True)[0] + shift_hidden = hidden[:, :-1].reshape(-1, hidden.shape[-1]) + loss = chunked_linear_cross_entropy( + shift_hidden.index_select(0, keep), + head.weight, + shift_labels.reshape(-1).index_select(0, keep), + flat_w.index_select(0, keep), + bias=head.bias, + chunk_tokens=chunk_tokens, + ) + if keep.numel() == 0: + # Keep the graph connected so backward() works on an all-masked batch. + loss = loss + hidden.sum() * 0.0 + return loss, n_tokens diff --git a/yeto/cli.py b/yeto/cli.py index 510fade2..8abe6420 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -265,12 +265,28 @@ def loss_spec(value: str) -> str: ) tune.add_argument( "--kernel-backend", - choices=["native", "liger"], + choices=["native", "liger", "chunked"], default="native", - help="causal SFT loss kernel: native (default) or the pinned, " - "binary-mask-only instance-scoped Liger fused-linear-CE lane; " - "model layers remain native; the fused lane currently requires " - "--tuning lora --shard ddp", + help="causal SFT loss kernel: native (default), the pinned, " + "binary-mask-only instance-scoped Liger fused-linear-CE lane " + "(CUDA, --tuning lora --shard ddp), or chunked: a device-agnostic " + "linear+CE that never materializes full-sequence logits " + "(any device including Ascend, --shard ddp); model layers remain native", + ) + tune.add_argument( + "--gradient-checkpointing", + choices=["auto", "on", "off"], + default="auto", + help="learner activation recomputation; 'auto' (default) enables it " + "when the loaded base occupies more than half of VRAM; long-context " + "runs on small models usually need 'on'", + ) + tune.add_argument( + "--loss-chunk-tokens", + type=int, + default=4096, + help="tokens per logits chunk for --kernel-backend chunked; peak " + "loss memory is about chunk * vocab * 12 bytes", ) def int_or_auto(value: str): diff --git a/yeto/launch/learner_task.py b/yeto/launch/learner_task.py index 64fdcc71..255db28b 100644 --- a/yeto/launch/learner_task.py +++ b/yeto/launch/learner_task.py @@ -113,6 +113,12 @@ def make_learner_task(args, spec: ClusterSpec, learner_id: int, num_learners: in raise ValueError( "--base-quantization nf4 requires --kernel-backend native" ) + if kernel_backend == "chunked" and loss_function != "cross_entropy": + raise ValueError( + "--kernel-backend chunked supports only the built-in cross_entropy loss" + ) + if kernel_backend == "chunked" and args.shard != "ddp": + raise ValueError("--kernel-backend chunked supports only --shard ddp") if kernel_backend == "liger" and loss_function != "cross_entropy": raise ValueError( "--kernel-backend liger fused-linear-CE supports only the built-in " @@ -205,6 +211,13 @@ def make_learner_task(args, spec: ClusterSpec, learner_id: int, num_learners: in f" --attention-backend {attention_backend}" f" --kernel-backend {kernel_backend}" ) + gradient_checkpointing = getattr(args, "gradient_checkpointing", "auto") + if gradient_checkpointing != "auto": + learner_flags += f" --gradient-checkpointing {gradient_checkpointing}" + if kernel_backend == "chunked": + learner_flags += ( + f" --loss-chunk-tokens {int(getattr(args, 'loss_chunk_tokens', 4096))}" + ) else: if getattr(args, "island_backend", "torch") != "torch": raise ValueError("diffusion model-kind uses the torch island backend, not megatron") diff --git a/yeto/learner.py b/yeto/learner.py index b4e99029..5399f8a5 100644 --- a/yeto/learner.py +++ b/yeto/learner.py @@ -47,6 +47,11 @@ from .diloco_sync import DiLoCoSyncState, sync_diloco_boundary from .finalization import finalize_torch_island from .fragments import build_layout +from .chunked_ce import ( + DEFAULT_LOSS_CHUNK_TOKENS, + chunked_sft_forward, + require_chunked_ce_support, +) from .losses import load_custom_loss, load_pickled_loss, sft_loss from .wandb_logger import TELEMETRY_EVERY from .models import MODEL_ALIASES as MODEL_ALIASES @@ -121,10 +126,18 @@ def parse_args(argv=None): "--kernel-backend", choices=KERNEL_BACKENDS, default="native", - help="causal SFT loss kernel: native (default) or the pinned, " - "binary-mask-only instance-scoped Liger fused-linear-CE lane; " - "model layers remain native; the fused lane currently requires " - "--tuning lora --shard ddp", + help="causal SFT loss kernel: native (default), the pinned, " + "binary-mask-only instance-scoped Liger fused-linear-CE lane " + "(CUDA, --tuning lora --shard ddp), or chunked: a device-agnostic " + "linear+CE that never materializes full-sequence logits " + "(any device, --shard ddp); model layers remain native", + ) + p.add_argument( + "--loss-chunk-tokens", + type=int, + default=DEFAULT_LOSS_CHUNK_TOKENS, + help="tokens per logits chunk for --kernel-backend chunked; peak " + "loss memory is about chunk * vocab * 12 bytes", ) p.add_argument( "--train-on", @@ -517,6 +530,10 @@ def load_model_and_tokenizer(args, device): model.to(device) if kernel_backend == "liger": validate_lora_production_envelope(model) + if kernel_backend == "chunked": + log.info("chunked linear-CE loss: %s chunk_tokens=%d", + require_chunked_ce_support(model), + getattr(args, "loss_chunk_tokens", DEFAULT_LOSS_CHUNK_TOKENS)) return model, tokenizer @@ -963,6 +980,12 @@ def probe_loss_forward(probe_model, input_ids): weights = torch.ones_like(input_ids, dtype=torch.float32) loss, _ = liger_sft_forward(probe_model, input_ids, weights) return loss + elif getattr(args, "kernel_backend", "native") == "chunked": + def probe_loss_forward(probe_model, input_ids): + loss, _ = chunked_sft_forward( + probe_model, input_ids, chunk_tokens=args.loss_chunk_tokens + ) + return loss args.micro_batch_size = resolve_micro_batch_size( args, @@ -1225,7 +1248,7 @@ def finalize_sync(): # c_tokens uses tokens_total, which advances by the exact island-global # raw-token count accepted into each optimizer step. kernel_backend = getattr(args, "kernel_backend", "native") - if kernel_backend == "liger": + if kernel_backend in ("liger", "chunked"): compute_loss = None elif args.loss_function.startswith("pickle:"): compute_loss = load_pickled_loss( @@ -1305,6 +1328,11 @@ def train_metrics(loss_sum, targets, window, seconds, tokens_since): loss, batch_target_tokens = liger_sft_forward( model, input_ids, weights ) + elif kernel_backend == "chunked": + loss, batch_target_tokens = chunked_sft_forward( + model, input_ids, weights, + chunk_tokens=args.loss_chunk_tokens, + ) else: out = model(input_ids=input_ids) loss, batch_target_tokens = compute_loss(