From d39b586b6127b1d7047b843bb083e8d17ffe2819 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Tue, 3 Nov 2026 05:12:13 +0800 Subject: [PATCH 1/5] accel: recognize an allocator OOM that carries no vendor exception class `torch.npu` exports no `OutOfMemoryError`, so `oom_error(device)` fell back to `torch.cuda.OutOfMemoryError` and `except oom_error(device)` never matched Ascend's refusal, which arrives as a bare `RuntimeError("NPU out of memory. Tried to allocate ...")`. Both memory probes -- the causal-LM auto-batch resolver and the diffusion one -- were therefore not probes at all on Ascend: the first size that did not fit killed the run instead of fixing the ceiling. Observed on a 910B3 with Qwen3-0.6B at 8k tokens: `--micro-batch-size auto --grad-accum 8` died in `_probe_once` with "Tried to allocate 37.09 GiB" rather than settling on a smaller micro-batch. `accel.is_oom(exc, device)` answers the question the probes actually have -- is this the allocator refusing an allocation -- across the vendor class, torch's generic `OutOfMemoryError`, and the message-only form. The probes re-raise anything else, so an aborted collective still fails loudly instead of being mistaken for a memory ceiling. The existing test suite missed this because its fake npu module defines an `OutOfMemoryError` that the real torch_npu does not. Co-Authored-By: Claude Opus 5.5 --- tests/test_accel.py | 21 +++++++++++++++++++++ tests/test_autobatch.py | 38 ++++++++++++++++++++++++++++++++++++++ yeto/accel.py | 22 ++++++++++++++++++++++ yeto/autobatch.py | 4 +++- yeto/diffusion/learner.py | 4 +++- 5 files changed, 87 insertions(+), 2 deletions(-) diff --git a/tests/test_accel.py b/tests/test_accel.py index bc338cdd8..d964e24f1 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 2e11cbbcd..954c35ecd 100644 --- a/tests/test_autobatch.py +++ b/tests/test_autobatch.py @@ -298,3 +298,41 @@ 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 + ) diff --git a/yeto/accel.py b/yeto/accel.py index 2e3461f0f..9eb7dc67f 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 64b8f7d8b..0370e1279 100644 --- a/yeto/autobatch.py +++ b/yeto/autobatch.py @@ -219,7 +219,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/diffusion/learner.py b/yeto/diffusion/learner.py index 1e4b83e88..cbeea421e 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) From 1f3afacf43404332d85cd904464a97dd685d85ff Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Tue, 3 Nov 2026 05:16:09 +0800 Subject: [PATCH 2/5] Autobatch: probe the native loss instead of a stand-in for it `_probe_once` approximated the native loss with `logits[:, :-1].float().sum()` on the stated assumption that "the shifted f32 logits copy dominates". It does not: `sft_loss` hands that copy to `F.cross_entropy`, which keeps its log_softmax output -- a second tensor of the same full (B, T-1, V) shape -- alive for the backward. The probe therefore reported a ceiling about one logits-sized f32 tensor too high. On a 910B3 with Qwen3-0.6B (151936-token vocabulary) at `--seq-len 8192 --grad-accum 8`, the probe admitted micro-batch 4 and the first real training step died asking for another 18.54 GiB. Running the real loss in the probe -- with all-ones weights, the worst case over packed blocks -- makes the probe answer the question it claims to answer. Co-Authored-By: Claude Opus 5.5 --- tests/test_autobatch.py | 26 ++++++++++++++++++++++++++ yeto/autobatch.py | 19 ++++++++++++++++--- 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/tests/test_autobatch.py b/tests/test_autobatch.py index 954c35ecd..da0b93587 100644 --- a/tests/test_autobatch.py +++ b/tests/test_autobatch.py @@ -336,3 +336,29 @@ def probe(*a, **k): 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/yeto/autobatch.py b/yeto/autobatch.py index 0370e1279..b3c004bd1 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). From a7ef6b88ec5a369631c676597a0dc73a4059fe1e Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Tue, 3 Nov 2026 06:37:12 +0800 Subject: [PATCH 3/5] Learner: name the chunked loss kernel in the causal kernel recipe log The recipe line mapped every non-liger backend to torch-cross-entropy, so --kernel-backend chunked runs were archived as native-loss runs. Map each backend explicitly and reject unknown values. Co-Authored-By: Claude Opus 5.5 --- tests/test_chunked_ce.py | 24 ++++++++++++++++++++++++ yeto/chunked_ce.py | 1 + yeto/learner.py | 25 ++++++++++++++++++++----- 3 files changed, 45 insertions(+), 5 deletions(-) diff --git a/tests/test_chunked_ce.py b/tests/test_chunked_ce.py index 614c53dcd..35a0a4fb9 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/yeto/chunked_ce.py b/yeto/chunked_ce.py index 4c3795395..22c6d0eef 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/learner.py b/yeto/learner.py index 5399f8a51..22f910fe7 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,24 @@ def load_model_and_tokenizer(args, device): return model, tokenizer +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 From b8dbd24e9fa0ab21fabf3a9d742889e380769f57 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Tue, 3 Nov 2026 01:45:09 +0800 Subject: [PATCH 4/5] Learner: enter training mode before enabling gradient checkpointing `from_pretrained` returns an eval-mode module and the learner never left it. Transformers' GradientCheckpointingLayer only recomputes when `self.training` is true, so `--gradient-checkpointing on` was accepted, logged, and silently ineffective: every activation stayed alive for backward. Dropout, LoRA's included, was inert for the same reason. Measured on one Ascend 910B3 card with Qwen3-0.6B, LoRA all-linear, one sequence per micro-batch, activation peak per sequence length: tokens | eval mode | train mode 4096 | 21.5 GB | 9.3 GB 8192 | 33.9 GB | 9.6 GB 16384 | 59.0 GB | 10.1 GB 28000 | OOM | 10.9 GB Which is the difference between OOM and training at long context: a 28k-token SFT run on a 64 GB card is impossible in eval mode and comfortable in train mode. The checkpointing setup moves into `enter_training_mode`, where the ordering is stated once and covered by tests. Co-Authored-By: Claude Opus 5.5 --- tests/test_learner_units.py | 82 +++++++++++++++++++++++++++++++++++++ yeto/learner.py | 52 +++++++++++++++-------- 2 files changed, 118 insertions(+), 16 deletions(-) diff --git a/tests/test_learner_units.py b/tests/test_learner_units.py index 6bc6a6475..1a9b1f35f 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,84 @@ 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 diff --git a/yeto/learner.py b/yeto/learner.py index 22f910fe7..9e4247bfe 100644 --- a/yeto/learner.py +++ b/yeto/learner.py @@ -602,6 +602,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} @@ -827,22 +862,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( From 0d8b68c1c1ffafcfb17394cdd7a27f460fb8a43f Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Tue, 3 Nov 2026 10:43:38 +0800 Subject: [PATCH 5/5] Learner: keep first-step warm-up out of the s/step telemetry The first logging window included the one-off first-step cost (streaming dataset build, compilation, collective setup). On a 2.9 GB dataset that reported 25 s/step for a run whose steady state is 4.75 s/step. Log the first step's duration on its own and start the first window after it. Co-Authored-By: Claude Opus 5.5 --- tests/test_learner_units.py | 37 ++++++++++++++++++++ yeto/learner.py | 68 +++++++++++++++++++++++++++---------- 2 files changed, 88 insertions(+), 17 deletions(-) diff --git a/tests/test_learner_units.py b/tests/test_learner_units.py index 1a9b1f35f..c39c42536 100644 --- a/tests/test_learner_units.py +++ b/tests/test_learner_units.py @@ -1254,3 +1254,40 @@ def test_auto_checkpointing_follows_free_memory(monkeypatch, free, total, expect 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/learner.py b/yeto/learner.py index 9e4247bfe..9c988ef80 100644 --- a/yeto/learner.py +++ b/yeto/learner.py @@ -534,6 +534,42 @@ 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. @@ -1313,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, @@ -1321,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) @@ -1411,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)", @@ -1423,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( @@ -1476,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)