diff --git a/tests/test_accel.py b/tests/test_accel.py index bc338cdd..d964e24f 100644 --- a/tests/test_accel.py +++ b/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)) diff --git a/tests/test_autobatch.py b/tests/test_autobatch.py index 2e11cbbc..da0b9358 100644 --- a/tests/test_autobatch.py +++ b/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"] diff --git a/tests/test_chunked_ce.py b/tests/test_chunked_ce.py index 614c53dc..35a0a4fb 100644 --- a/tests/test_chunked_ce.py +++ b/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") diff --git a/tests/test_learner_units.py b/tests/test_learner_units.py index 6bc6a647..c39c4253 100644 --- a/tests/test_learner_units.py +++ b/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) diff --git a/yeto/accel.py b/yeto/accel.py index 2e3461f0..9eb7dc67 100644 --- a/yeto/accel.py +++ b/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. diff --git a/yeto/autobatch.py b/yeto/autobatch.py index 64b8f7d8..b3c004bd 100644 --- a/yeto/autobatch.py +++ b/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: diff --git a/yeto/chunked_ce.py b/yeto/chunked_ce.py index 4c379539..22c6d0ee 100644 --- a/yeto/chunked_ce.py +++ b/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. diff --git a/yeto/diffusion/learner.py b/yeto/diffusion/learner.py index 1e4b83e8..cbeea421 100644 --- a/yeto/diffusion/learner.py +++ b/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) diff --git a/yeto/learner.py b/yeto/learner.py index 5399f8a5..9c988ef8 100644 --- a/yeto/learner.py +++ b/yeto/learner.py @@ -48,6 +48,7 @@ from .finalization import finalize_torch_island from .fragments import build_layout from .chunked_ce import ( + CHUNKED_LOSS_IMPLEMENTATION, DEFAULT_LOSS_CHUNK_TOKENS, chunked_sft_forward, require_chunked_ce_support, @@ -485,11 +486,7 @@ def load_model_and_tokenizer(args, device): f"failed to bind the isolated fused-linear-CE loss to {model_id!r}" ) from exc resolved_attention = resolved_attention_backend(model, attention_backend) - loss_implementation = ( - FUSED_LINEAR_CE_IMPLEMENTATION - if kernel_backend == "liger" - else NATIVE_LOSS_IMPLEMENTATION - ) + loss_implementation = loss_implementation_for(kernel_backend) log.info( "causal kernel recipe: attention requested=%s resolved=%s; " "layers=%s; loss=%s%s", @@ -537,6 +534,60 @@ def load_model_and_tokenizer(args, device): return model, tokenizer +class StepWindowTimer: + """Wall-clock windows for the s/step and tokens/s telemetry. + + The first local step of a process carries one-off costs (streaming + dataset build, graph compilation, collective setup) that can be minutes + on a large dataset. It is reported once on its own and the first window + starts after it, so s/step reflects steady-state throughput only. + """ + + def __init__(self, steps, tokens, clock=time.monotonic): + self._clock = clock + self._start_steps = steps + self._first_pending = True + self._reset(steps, tokens) + + def _reset(self, steps, tokens): + self._t = self._clock() + self._steps = steps + self._tokens = tokens + + def first_step_done(self, steps, tokens): + """Seconds the first step took (once), else None; restarts the window.""" + if not self._first_pending or steps <= self._start_steps: + return None + self._first_pending = False + elapsed = self._clock() - self._t + self._reset(steps, tokens) + return elapsed + + def window(self, steps, tokens): + """(steps, seconds, tokens) since the last window; starts a new one.""" + out = (steps - self._steps, self._clock() - self._t, tokens - self._tokens) + self._reset(steps, tokens) + return out + + +def loss_implementation_for(kernel_backend: str) -> str: + """Name of the loss kernel a ``--kernel-backend`` value actually runs. + + This is what the ``causal kernel recipe`` log line records, so every + backend must map to its own implementation; an unknown value is an error + rather than a silent fallback to the native name. + """ + implementations = { + "native": NATIVE_LOSS_IMPLEMENTATION, + "liger": FUSED_LINEAR_CE_IMPLEMENTATION, + "chunked": CHUNKED_LOSS_IMPLEMENTATION, + } + try: + return implementations[kernel_backend] + except KeyError: + raise ValueError(f"unknown kernel backend {kernel_backend!r}") from None + + # Attention projection names across the architectures we run (Llama/Qwen/ # Gemma-style q/k/v/o, fused qkv variants, DeepSeek's low-rank q_a/q_b and # kv_a/kv_b split). peft treats a string target as a regex fullmatch on the @@ -587,6 +638,41 @@ def resolve_lora_targets(choice: str, config) -> str: return "all-linear" +def enter_training_mode(model, args, device) -> bool: + """Put ``model`` in training mode and apply ``--gradient-checkpointing``. + + The order matters. ``from_pretrained`` returns an eval-mode module, and + Transformers' ``GradientCheckpointingLayer`` skips recomputation whenever + ``self.training`` is False, so enabling checkpointing on an eval-mode + model is a silent no-op: the flag is accepted, the log line is printed, + and every activation is still kept for backward (Qwen3-0.6B, one 16k-token + sequence: 59 GB of activations in eval mode against 10 GB in train mode, + which is the difference between OOM and training at long context). + Dropout -- LoRA's included -- is inert in eval mode for the same reason. + + Returns whether checkpointing ended up enabled. + """ + model.train() + + grad_ckpt = args.gradient_checkpointing == "on" + if args.gradient_checkpointing == "auto": + # The base is fully on-device here (load ends with model.to(device)), + # so free memory directly reflects what activations must fit into. + memory = accel.mem_get_info(device) + grad_ckpt = memory is not None and memory[0] < memory[1] / 2 + if grad_ckpt: + # Non-reentrant checkpointing composes with FSDP and peft; the input + # grad hook keeps the graph alive from the embeddings down to the + # first adapter when the base is frozen. + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + if args.tuning == "lora": + model.enable_input_require_grads() + log.info("gradient checkpointing enabled") + return grad_ckpt + + def trainable_params(model) -> dict[str, torch.Tensor]: return {n: p for n, p in model.named_parameters() if p.requires_grad} @@ -812,22 +898,7 @@ def main(argv=None) -> None: torch.manual_seed(training_seed) accel.manual_seed_all(device, training_seed) - grad_ckpt = args.gradient_checkpointing == "on" - if args.gradient_checkpointing == "auto": - # The base is fully on-device here (load ends with model.to(device)), - # so free memory directly reflects what activations must fit into. - memory = accel.mem_get_info(device) - grad_ckpt = memory is not None and memory[0] < memory[1] / 2 - if grad_ckpt: - # Non-reentrant checkpointing composes with FSDP and peft; the input - # grad hook keeps the graph alive from the embeddings down to the - # first adapter when the base is frozen. - model.gradient_checkpointing_enable( - gradient_checkpointing_kwargs={"use_reentrant": False} - ) - if args.tuning == "lora": - model.enable_input_require_grads() - log.info("gradient checkpointing enabled") + grad_ckpt = enter_training_mode(model, args, device) params = trainable_params(model) layout = build_layout( @@ -1278,7 +1349,7 @@ def train_metrics(loss_sum, targets, window, seconds, tokens_since): "train/lr": sched.get_last_lr()[0], "train/raw_tokens_total": tokens_total, "train/target_tokens_total": target_tokens_total, - "train/sec_per_step": seconds / window, + "train/sec_per_step": seconds / window if window else 0.0, "train/tokens_per_sec": tokens_since / seconds if seconds > 0 else 0.0, "train/epoch": epoch, "local_step": steps_total, @@ -1286,8 +1357,7 @@ def train_metrics(loss_sum, targets, window, seconds, tokens_since): } epoch = 0 - t_last = time.monotonic() - tokens_at_last_log = 0 + timer = StepWindowTimer(steps_total, tokens_total) while not sync_state.shutdown and steps_total < args.max_local_steps: steps_at_epoch_start = steps_total sampler = getattr(loader, "sampler", None) @@ -1376,11 +1446,17 @@ def train_metrics(loss_sum, targets, window, seconds, tokens_since): tokens_total += global_raw target_tokens_total += global_targets + warmup_s = timer.first_step_done(steps_total, tokens_total) + if warmup_s is not None and rank == 0: + log.info( + "first local step took %.2f s (data loading, compilation " + "and collective setup; excluded from s/step)", + warmup_s, + ) if steps_total % TELEMETRY_EVERY == 0: loss_sum = _global_loss_sum(step_loss_local, world) if rank == 0: - dt = time.monotonic() - t_last - t_last = time.monotonic() + n_steps, dt, d_tokens = timer.window(steps_total, tokens_total) log.info( "local_step=%d global_step=%d loss/token=%.4f " "target_tokens=%d (%.2f s/step)", @@ -1388,18 +1464,17 @@ def train_metrics(loss_sum, targets, window, seconds, tokens_since): sync_state.global_step, loss_sum / global_targets, target_tokens_total, - dt / TELEMETRY_EVERY, + dt / max(n_steps, 1), ) wandb_run.log( train_metrics( loss_sum, global_targets, - TELEMETRY_EVERY, + n_steps, dt, - tokens_total - tokens_at_last_log, + d_tokens, ) ) - tokens_at_last_log = tokens_total # --- fragment sync at the step boundary (never blocks) --- if sync_diloco_boundary( @@ -1441,17 +1516,11 @@ def train_metrics(loss_sum, targets, window, seconds, tokens_since): and steps_total > 0 and steps_total % TELEMETRY_EVERY != 0 ): - window = steps_total % TELEMETRY_EVERY loss_sum = _global_loss_sum(step_loss_local, world) if rank == 0: + window, dt, d_tokens = timer.window(steps_total, tokens_total) wandb_run.log( - train_metrics( - loss_sum, - global_targets, - window, - time.monotonic() - t_last, - tokens_total - tokens_at_last_log, - ) + train_metrics(loss_sum, global_targets, window, dt, d_tokens) ) learner_budget_steps = getattr(args, "learner_budget_steps", None)