Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / callbacks / base.py: 70%
583 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Callback base class and built-in callbacks.
17dispatched explicitly in ``on_step_end`` etc. Engineer sees all callbacks and
18order at a glance.
20``checkpoint_callback.py`` (242 lines) + ``trace_callback.py`` (231 lines).
21"""
22import copy
23import gc
24import json
25import logging
26import math
27import os
28import re
29import threading
30import time
31from typing import TYPE_CHECKING, Optional
33import torch
34from torch.utils.tensorboard import SummaryWriter
36from hyper_parallel import get_platform
37from hyper_parallel.core.distributed_checkpoint import load as dcp_load, save as dcp_save
38from hyper_parallel.core.distributed_checkpoint.offline_transform import (
39 save_state_dict_as_huggingface_format,
40)
41from hyper_parallel.core.fully_shard.api import get_model_state_dict
43platform = get_platform()
45if TYPE_CHECKING:
46 from hyper_parallel.trainer.base import BaseTrainer, TrainerState
48logger = logging.getLogger(__name__)
51class Callback:
52 """Base class for all trainer callbacks.
54 Each callback holds a reference to the trainer for accessing model,
55 optimizer, state, and config. Subclass and override the hooks you need.
57 Args:
58 trainer: The BaseTrainer instance.
59 """
61 def __init__(self, trainer: "BaseTrainer") -> None:
62 self.trainer = trainer
64 # ------------------------------------------------------------------
65 # Lifecycle hooks
66 # ------------------------------------------------------------------
68 def on_init_end(self, state: "TrainerState", **kwargs) -> None:
69 """Called once at the end of ``BaseTrainer.__init__`` / subclass init.
71 At this point every ``_build_*`` has run — model is parallelised,
72 optimizer/scheduler/dataloader are built, callbacks are constructed.
73 Use this for one-shot setup that must see the FINAL trainer state
74 (e.g. logging the parameter count, opening a TensorBoard writer
75 keyed by run_id, validating user config against the built model).
76 """
78 def on_train_begin(self, state: "TrainerState", **kwargs) -> None:
79 """Called at the start of ``train()`` (before any optimizer.step).
81 ``CheckpointCallback`` runs resume here, so when this hook fires
82 ``state.global_step`` may already be > 0 if a checkpoint was loaded.
83 """
85 def on_train_end(self, state: "TrainerState", **kwargs) -> None:
86 """Called at the end of training (before ``destroy_process_group``).
88 Final checkpoints, profiler stops, W&B finish, etc. happen here.
89 """
91 def on_epoch_begin(self, state: "TrainerState", **kwargs) -> None:
92 """Called at the start of each epoch."""
94 def on_epoch_end(self, state: "TrainerState", **kwargs) -> None:
95 """Called at the end of each epoch."""
97 def on_step_begin(self, state: "TrainerState", **kwargs) -> None:
98 """Called at the start of each training step (before fwd of mb 0)."""
100 def on_step_end(self, state: "TrainerState", *, loss: float = None,
101 grad_norm: float = None, **kwargs) -> None:
102 """Called at the end of each training step (after optimizer.step)."""
104 def on_substep_end(self, state: "TrainerState", **kwargs) -> None:
105 """Called after each micro-batch fwd+bwd (gradient accumulation)."""
107 def on_pre_optimizer_step(self, state: "TrainerState", *,
108 grad_norm: float = None, **kwargs) -> None:
109 """Called after grad clip, before ``optimizer.step``.
111 ``grad_norm`` here is the post-clip scalar produced by hyper's
112 DTensor-aware clipper — use it to detect NaN/Inf or to log the
113 effective clip ratio.
114 """
116 def on_log(self, state: "TrainerState", *, metrics: dict, **kwargs) -> None:
117 """Called when ``LoggingCallback`` emits a structured metrics record.
119 Reuse this hook in TensorBoard / W&B / external metric sinks so
120 every logging backend sees the SAME record. Avoids three callbacks
121 each computing throughput / lr independently.
123 Args:
124 metrics: Dict containing at minimum ``step``, ``loss``,
125 ``grad_norm``, ``lr``, ``step_time``; throughput fields
126 (``tokens_per_sec``, ``tflops``, ``mfu``) are present iff
127 ``logging.report_throughput`` is on.
128 """
130 def on_save(self, state: "TrainerState", *, checkpoint_dir: str,
131 **kwargs) -> None:
132 """Called immediately after ``CheckpointCallback`` finishes a save.
134 Use to upload to remote storage, register the ckpt with an
135 experiment tracker, or trigger downstream eval jobs. ``checkpoint_dir``
136 is the on-disk path containing model shards + optimizer/scheduler/RNG/
137 dataloader/extra_state.
138 """
140 def on_load(self, state: "TrainerState", *, checkpoint_dir: str,
141 **kwargs) -> None:
142 """Called immediately after ``CheckpointCallback`` finishes a resume.
144 Use to verify the resumed step matches expectations, log the
145 restore event, or seed downstream callbacks with the resumed state.
146 """
148 def on_evaluate(self, state: "TrainerState", *, metrics: dict = None,
149 **kwargs) -> None:
150 """Called when an evaluation pass completes.
152 Currently triggered as a stub from ``EvalCallback``; once a real
153 eval loop lands the callback will pass back the eval ``metrics``
154 dict for sinks (TensorBoard / W&B) to log.
155 """
158class LoggingCallback(Callback):
159 """Log training metrics: loss, grad_norm, lr, throughput.
161 """
163 def __init__(self, trainer: "BaseTrainer") -> None:
164 super().__init__(trainer)
165 train_cfg = getattr(trainer.args, 'train', None)
166 log_cfg = getattr(train_cfg, 'logging', None)
167 if log_cfg is None:
168 log_cfg = getattr(trainer.args, 'logging', None)
169 self.log_steps = getattr(log_cfg, 'log_steps', 10) if log_cfg else 10
170 self.report_global_loss = (
171 getattr(log_cfg, 'report_global_loss', False) if log_cfg else False
172 )
173 self.report_throughput = (
174 getattr(log_cfg, 'report_throughput', True) if log_cfg else True
175 )
176 self.model_flops_per_token = (
177 getattr(log_cfg, 'model_flops_per_token', None) if log_cfg else None
178 )
179 self.peak_tflops = (
180 getattr(log_cfg, 'peak_tflops', None) if log_cfg else None
181 )
182 # Estimate per-step tokens as upper bound (batch × seq_len). Real
183 # token count is available per step via ``last_global_tokens`` that
184 # ``BaseTrainer.train_step`` stashes onto the trainer.
185 gbs = getattr(trainer.args.train, 'global_batch_size', 1)
186 seq_len = getattr(trainer.args.data, 'max_seq_len', 1)
187 self._tokens_per_step_est = int(gbs) * int(seq_len)
188 self._step_start_time = 0.0
190 def on_step_begin(self, state: "TrainerState", **kwargs) -> None:
191 self._step_start_time = time.time()
193 def on_step_end(self, state: "TrainerState", *, loss: float = None,
194 grad_norm: float = None, **kwargs) -> None:
195 if state.global_step % self.log_steps != 0:
196 return
198 elapsed = max(time.time() - self._step_start_time, 1e-9)
199 lr = 0.0
200 if self.trainer.lr_scheduler is not None:
201 lr = self.trainer.lr_scheduler.get_last_lr()[0]
203 metrics = {
204 "step": state.global_step,
205 # 8-decimal precision keeps fp32 sub-bf16 differences visible
206 # in the log for sanity comparisons across runs.
207 "loss": f"{loss:.8f}" if loss is not None else "N/A",
208 "grad_norm": (
209 f"{grad_norm:.8f}" if grad_norm is not None else "N/A"
210 ),
211 "lr": f"{lr:.2e}",
212 "step_time": f"{elapsed:.2f}s",
213 }
215 tokens_per_sec = None
216 if self.report_throughput:
217 # Prefer real per-step token count stashed by train_step; fall back
218 # to the estimate until the first step sets it (declared None).
219 tokens = getattr(self.trainer, '_last_global_tokens', None)
220 if tokens is None:
221 tokens = self._tokens_per_step_est
222 tokens_per_sec = tokens / elapsed
223 metrics["tokens_per_sec"] = f"{tokens_per_sec:,.0f}"
225 if self.model_flops_per_token and self.peak_tflops:
226 # Observed TFLOPS = tokens/sec × flops/token / 1e12.
227 # MFU = observed / (peak × world_size).
228 world = max(platform.get_world_size(), 1)
229 observed_tflops = (
230 tokens_per_sec * self.model_flops_per_token / 1e12
231 )
232 mfu = observed_tflops / (self.peak_tflops * world)
233 metrics["tflops"] = f"{observed_tflops:.1f}"
234 metrics["mfu"] = f"{mfu * 100:.1f}%"
236 # Include aux_loss from MoEMonitorCallback when available.
237 moe_cb = getattr(self.trainer, 'moe_monitor_callback', None)
238 aux_loss = getattr(moe_cb, 'last_mean_aux_loss', None) if moe_cb is not None else None
239 if aux_loss is not None:
240 metrics["aux_loss"] = f"{aux_loss:.6f}"
242 message = " | ".join(f"{k}={v}" for k, v in metrics.items())
243 monitor_cb = getattr(self.trainer, 'training_state_monitor_callback', None)
244 monitor_active = monitor_cb is not None and getattr(monitor_cb, 'active', False)
245 if monitor_active:
246 logger.info(message)
247 else:
248 logger.info_rank0(message)
250 record = {
251 "step": state.global_step,
252 "loss": loss,
253 "grad_norm": grad_norm,
254 "lr": lr,
255 "step_time": elapsed,
256 "tokens_per_sec": tokens_per_sec,
257 "aux_loss": aux_loss,
258 }
259 state.log_history.append(record)
261 # Fan-out to other log-event listeners (TB / W&B / sinks).
262 dispatch = getattr(self.trainer, "dispatch_log_event", None)
263 if dispatch is not None:
264 dispatch(record)
267class CheckpointCallback(Callback):
268 """Save distributed checkpoints and handle resume.
270 Uses hyper's own DCP ``save`` / ``load`` APIs.
271 """
273 def __init__(self, trainer: "BaseTrainer") -> None:
274 super().__init__(trainer)
275 train_cfg = getattr(trainer.args, 'train', None)
276 ckpt_cfg = getattr(train_cfg, 'checkpoint', None)
277 if ckpt_cfg is None:
278 ckpt_cfg = getattr(trainer.args, 'checkpoint', None)
279 self.save_steps = getattr(ckpt_cfg, 'save_steps', 0) if ckpt_cfg else 0
280 self.output_dir = (
281 getattr(ckpt_cfg, 'output_dir', 'outputs') if ckpt_cfg else 'outputs'
282 )
283 self.load_path = (
284 getattr(ckpt_cfg, 'load_path', None) if ckpt_cfg else None
285 )
286 self.save_async = (
287 getattr(ckpt_cfg, 'save_async', False) if ckpt_cfg else False
288 )
289 self._last_saved_step = -1
290 self._save_thread = None # async save worker
292 def on_train_begin(self, state: "TrainerState", **kwargs) -> None:
293 """Resume from checkpoint: model + optimizer + lr_scheduler + step + RNG.
295 RFC DoD: "Save → resume → 续训 loss 一致（含 dataloader + RNG 恢复）"
296 """
297 if not self.load_path:
298 return
299 try:
300 # pylint: disable=C0415
301 # Non-model artifacts (optimizer/scheduler/RNG) are plain dicts —
302 # use torch.save/load, matching the save side.
304 if not os.path.isdir(self.load_path):
305 logger.warning("Checkpoint path not found: %s", self.load_path)
306 return
308 # 1. Restore model via hyper DCP
309 model_sd = self.trainer.model.state_dict()
310 dcp_load(model_sd, checkpoint_id=self.load_path, use_collectives=False)
311 self.trainer.model.load_state_dict(model_sd)
312 logger.info("Model restored from %s", self.load_path)
314 # 2. Restore extra state (step, epoch)
315 extra_path = os.path.join(self.load_path, "extra_state.json")
316 if os.path.isfile(extra_path):
317 with open(extra_path, encoding="utf-8") as f:
318 extra = json.load(f)
319 state.global_step = extra.get("global_step", 0)
320 state.epoch = extra.get("epoch", 0)
321 logger.info("Resumed at step=%d, epoch=%d",
322 state.global_step, state.epoch)
324 # 3. Restore optimizer
325 optim_path = os.path.join(self.load_path, f"optimizer_rank{platform.get_rank()}.pt")
326 if os.path.isfile(optim_path) and self.trainer.optimizer:
327 optim_sd = torch.load(optim_path, map_location="cpu", weights_only=True)
328 self.trainer.optimizer.load_state_dict(optim_sd)
329 logger.info("Optimizer restored")
331 # 4. Restore LR scheduler
332 sched_path = os.path.join(self.load_path, "scheduler.pt")
333 if os.path.isfile(sched_path) and self.trainer.lr_scheduler:
334 sched_sd = torch.load(sched_path, map_location="cpu", weights_only=True)
335 self.trainer.lr_scheduler.load_state_dict(sched_sd)
336 logger.info("LR scheduler restored")
338 # 5. Restore RNG state
339 rng_path = os.path.join(self.load_path, f"rng_rank{platform.get_rank()}.pt")
340 if os.path.isfile(rng_path):
341 rng_state = torch.load(rng_path, map_location="cpu", weights_only=True)
342 platform.set_rng_state(rng_state)
343 logger.info("RNG state restored")
345 # 6. Restore dataloader position (StatefulDataLoader)
346 dl_path = os.path.join(self.load_path, f"dataloader_rank{platform.get_rank()}.pt")
347 if os.path.isfile(dl_path) and hasattr(self.trainer, 'train_dataloader'):
348 dl_state = torch.load(dl_path, map_location="cpu", weights_only=False)
349 self.trainer.train_dataloader.load_state_dict(dl_state)
350 logger.info("Dataloader state restored")
352 # Fan-out the load event so other callbacks (TensorBoard /
353 # W&B / external trackers) can record the resume.
354 dispatch = getattr(self.trainer, "dispatch_load_event", None)
355 if dispatch is not None:
356 dispatch(self.load_path)
358 except (OSError, RuntimeError, ValueError) as exc:
359 logger.warning("Failed to load checkpoint from %s: %s", self.load_path, exc)
361 def on_step_end(self, state: "TrainerState", *, loss: float = None,
362 grad_norm: float = None, **kwargs) -> None:
363 if self.save_steps <= 0:
364 return
365 if state.global_step % self.save_steps != 0:
366 return
367 if state.global_step == self._last_saved_step:
368 return
369 self._dispatch_save(state)
371 def on_train_end(self, state: "TrainerState", **kwargs) -> None:
372 """Save final checkpoint (synchronously, to guarantee completion)."""
373 # Wait for any outstanding async save first so the two don't race on
374 # the same directory / state-dict iterator.
375 self._join_pending()
376 if self.save_steps > 0 and state.global_step != self._last_saved_step:
377 # Final save always sync — the process is about to exit.
378 self._save(state)
380 # --- async plumbing -------------------------------------------------
381 def _dispatch_save(self, state: "TrainerState") -> None:
382 """Route to sync or async save based on ``save_async`` flag."""
383 if not self.save_async:
384 self._save(state)
385 return
386 # Wait for previous save to finish before starting a new one; saving
387 # twice concurrently would double RAM and race the filesystem.
388 self._join_pending()
389 # pylint: disable=C0415
390 # Snapshot state fields so the worker doesn't see later mutations.
391 snap_step = state.global_step
392 snap_epoch = state.epoch
393 state_snapshot = copy.copy(state)
394 state_snapshot.global_step = snap_step
395 state_snapshot.epoch = snap_epoch
396 self._save_thread = threading.Thread(
397 target=self._save,
398 args=(state_snapshot,),
399 name=f"ckpt-save-step{snap_step}",
400 daemon=True,
401 )
402 self._save_thread.start()
403 logger.info_rank0(
404 "Checkpoint save for step %d dispatched async (thread=%s)",
405 snap_step, self._save_thread.name,
406 )
408 def _join_pending(self) -> None:
409 """Block until any running async save finishes."""
410 t = self._save_thread
411 if t is not None and t.is_alive():
412 logger.info_rank0(
413 "Waiting for prior async ckpt save (%s)...", t.name,
414 )
415 t.join()
416 self._save_thread = None
418 def _save(self, state: "TrainerState") -> None:
419 """Save complete training state: model + optimizer + scheduler + step + RNG.
421 RFC DoD: "Save → resume → 续训 loss 一致（含 dataloader + RNG 恢复）"
422 """
423 # Optimizer/scheduler/RNG state dicts are plain Python dicts, not
424 # nn.Module — platform.save_checkpoint expects Module (safetensors).
425 # Use torch.save/load for these non-model artifacts.
426 save_dir = os.path.join(self.output_dir, f"step_{state.global_step}")
427 os.makedirs(save_dir, exist_ok=True)
428 rank = platform.get_rank()
430 try:
431 # 1. Model — via hyper DCP (each rank saves its own shards)
432 model_sd = self.trainer.model.state_dict()
433 dcp_save(model_sd, checkpoint_id=save_dir, use_collectives=False)
435 # 2. Optimizer — per-rank
436 if self.trainer.optimizer:
437 optim_path = os.path.join(save_dir, f"optimizer_rank{rank}.pt")
438 torch.save(self.trainer.optimizer.state_dict(), optim_path)
440 # 3. LR scheduler
441 if self.trainer.lr_scheduler and rank == 0:
442 sched_path = os.path.join(save_dir, "scheduler.pt")
443 torch.save(self.trainer.lr_scheduler.state_dict(), sched_path)
445 # 4. Extra state: global_step, epoch
446 if rank == 0:
447 extra = {
448 "global_step": state.global_step,
449 "epoch": state.epoch,
450 }
451 extra_path = os.path.join(save_dir, "extra_state.json")
452 with open(extra_path, "w", encoding="utf-8") as f:
453 json.dump(extra, f)
455 # 5. RNG state — per-rank via platform API
456 rng_state = platform.get_rng_state()
457 rng_path = os.path.join(save_dir, f"rng_rank{rank}.pt")
458 torch.save(rng_state, rng_path)
460 # 6. Dataloader position — per-rank (StatefulDataLoader)
461 if hasattr(self.trainer, 'train_dataloader') and hasattr(
462 self.trainer.train_dataloader, 'state_dict'
463 ):
464 dl_path = os.path.join(save_dir, f"dataloader_rank{rank}.pt")
465 torch.save(self.trainer.train_dataloader.state_dict(), dl_path)
467 self._last_saved_step = state.global_step
468 logger.info_rank0("Checkpoint saved to %s", save_dir)
470 # Fan-out the save event so other callbacks (W&B artifact
471 # upload, remote-storage sync, downstream eval triggers) can
472 # observe the new checkpoint without coupling to ckpt internals.
473 dispatch = getattr(self.trainer, "dispatch_save_event", None)
474 if dispatch is not None:
475 dispatch(save_dir)
477 except (OSError, RuntimeError, ValueError) as exc:
478 logger.warning("Failed to save checkpoint: %s", exc)
480 # HF format export is handled by SafetensorsExportCallback (separate concern).
483class SafetensorsExportCallback(Callback):
484 """Export model weights in HuggingFace safetensor format.
486 Separated from CheckpointCallback per RFC Section 5.2.
487 Uses ``get_model_state_dict`` with ``full_state_dict=True`` to gather
488 all FSDP shards into a full state dict before saving.
490 """
492 def __init__(self, trainer: "BaseTrainer") -> None:
493 super().__init__(trainer)
494 train_cfg = getattr(trainer.args, 'train', None)
495 ckpt_cfg = getattr(train_cfg, 'checkpoint', None)
496 if ckpt_cfg is None:
497 ckpt_cfg = getattr(trainer.args, 'checkpoint', None)
498 self.enabled = getattr(ckpt_cfg, 'save_hf_weights', False) if ckpt_cfg else False
499 self.save_steps = getattr(ckpt_cfg, 'save_steps', 0) if ckpt_cfg else 0
500 self.output_dir = getattr(ckpt_cfg, 'output_dir', 'outputs') if ckpt_cfg else 'outputs'
501 self._last_saved_step = -1
503 def on_step_end(self, state: "TrainerState", *, loss: Optional[float] = None,
504 grad_norm: Optional[float] = None, **kwargs) -> None:
505 if not self.enabled or self.save_steps <= 0:
506 return
507 if state.global_step % self.save_steps != 0:
508 return
509 if state.global_step == self._last_saved_step:
510 return
511 self._export(state)
513 def on_train_end(self, state: "TrainerState", **kwargs) -> None:
514 if self.enabled and self.save_steps > 0 and state.global_step != self._last_saved_step:
515 self._export(state)
517 def _export(self, state: "TrainerState") -> None:
518 """Gather full state dict from FSDP shards and save in HF format.
520 Routes through ``spec.state_dict_adapter().save_hf_state_dict`` when
521 the model's ``ModelSpec`` provides one, so per-model HF tensor
522 renaming and per-expert packing live in the model package, not in
523 this generic callback. Falls back to the legacy
524 ``save_state_dict_as_huggingface_format`` path when the spec has no
525 adapter (keeps ad-hoc / template models working).
526 """
527 # pylint: disable=C0415
529 rank = platform.get_rank()
530 save_dir = os.path.join(self.output_dir, f"step_{state.global_step}", "hf_ckpt")
532 try:
533 # ``StateDictOptions`` is a torch-backend type; hyper does not yet
534 # provide a wrapper, so the trainer reaches into torch directly.
535 # pylint: disable=C0415
536 from torch.distributed.checkpoint.state_dict import StateDictOptions
537 # full_state_dict=True gathers all FSDP shards; cpu_offload avoids OOM
538 options = StateDictOptions(full_state_dict=True, cpu_offload=True)
539 full_sd = get_model_state_dict(self.trainer.model, options=options)
541 if rank == 0:
542 os.makedirs(save_dir, exist_ok=True)
544 # Prefer the model-specific save adapter (closes the load/save
545 # loop via the ModelSpec contract). When absent, fall back to
546 # the generic offline-transform path.
547 spec = getattr(self.trainer, "spec", None)
548 adapter_cls = getattr(spec, "state_dict_adapter", None) if spec else None
549 save_fn = (
550 getattr(adapter_cls(), "save_hf_state_dict", None)
551 if adapter_cls is not None else None
552 )
553 if save_fn is not None:
554 hf_sd = save_fn(full_sd, self.trainer.model.config)
555 from safetensors.torch import save_file # pylint: disable=C0415
556 save_file(hf_sd, os.path.join(save_dir, "model.safetensors"))
557 logger.info(
558 "HF checkpoint saved via %s.save_hf_state_dict to %s",
559 adapter_cls.__name__, save_dir,
560 )
561 else:
562 save_state_dict_as_huggingface_format(full_sd, save_dir)
563 logger.info(
564 "HF checkpoint saved (no adapter on spec) to %s", save_dir,
565 )
567 self._last_saved_step = state.global_step
569 except (OSError, RuntimeError, ValueError) as exc:
570 logger.warning_rank0("Failed to save HF checkpoint: %s", exc)
573class EvalCallback(Callback):
574 """Evaluation callback stub.
576 Full evaluation is not yet implemented. This stub logs a warning whenever
577 an evaluation trigger is received so the absence of eval is visible in
578 training logs rather than silently skipped.
579 """
581 def on_step_end(self, state: "TrainerState", *, loss: Optional[float] = None,
582 grad_norm: Optional[float] = None, **kwargs) -> None:
583 eval_cfg = getattr(self.trainer.args, 'eval', None)
584 eval_steps = getattr(eval_cfg, 'eval_steps', 0) if eval_cfg else 0
585 if eval_steps > 0 and state.global_step % eval_steps == 0:
586 if platform.get_rank() == 0:
587 logger.warning(
588 "EvalCallback: evaluation not implemented (step=%d)", state.global_step
589 )
592class ProfilerCallback(Callback):
593 """Training profiler callback — STUB (not verified).
595 Hook reserved for ``torch.profiler.profile`` integration. Not yet
596 verified against the trainer; if you enable ``args.profiler.enabled``
597 we emit a one-time warning so the absence of profiling traces is
598 visible. To implement: wire ``torch.profiler.profile`` start/step/stop
599 in ``on_train_begin`` / ``on_step_end`` / ``on_train_end``.
600 """
602 def __init__(self, trainer: "BaseTrainer") -> None:
603 super().__init__(trainer)
604 prof_cfg = getattr(trainer.args, 'profiler', None)
605 if getattr(prof_cfg, 'enabled', False) and platform.get_rank() == 0:
606 logger.warning(
607 "ProfilerCallback: enabled=True but the implementation is "
608 "a stub — torch profiler is NOT started. Implement before "
609 "relying on traces."
610 )
613class WandbCallback(Callback):
614 """Weights & Biases logging callback — STUB (not verified).
616 Hook reserved for W&B integration. Not yet verified; if you enable
617 ``args.wandb.enabled`` we emit a one-time warning so missing W&B logs
618 are visible. To implement: wire ``wandb.init`` / ``wandb.log`` /
619 ``wandb.finish`` in ``on_train_begin`` / ``on_step_end`` /
620 ``on_train_end`` and verify against a real W&B run.
621 """
623 def __init__(self, trainer: "BaseTrainer") -> None:
624 super().__init__(trainer)
625 wandb_cfg = getattr(trainer.args, 'wandb', None)
626 if getattr(wandb_cfg, 'enabled', False) and platform.get_rank() == 0:
627 logger.warning(
628 "WandbCallback: enabled=True but the implementation is a "
629 "stub — nothing is sent to W&B. Implement before relying on "
630 "W&B dashboards."
631 )
634class ProgressCallback(Callback):
635 """tqdm progress bar callback (rank 0 only).
637 Displays a progress bar over training steps with live loss and grad_norm
638 metrics. Requires ``tqdm``; degrades gracefully if not installed.
639 """
641 def __init__(self, trainer: "BaseTrainer") -> None:
642 super().__init__(trainer)
643 self._pbar = None
645 def on_train_begin(self, state: "TrainerState", **kwargs) -> None:
646 if platform.get_rank() != 0:
647 return
648 try:
649 # pylint: disable=C0415
650 from tqdm import tqdm # pylint: disable=C0415 # optional dep
651 self._pbar = tqdm(
652 total=state.max_steps,
653 initial=state.global_step,
654 desc="Training",
655 unit="step",
656 dynamic_ncols=True,
657 )
658 except ImportError:
659 logger.warning("ProgressCallback: 'tqdm' not installed — progress bar disabled")
661 def on_step_end(self, state: "TrainerState", *, loss: Optional[float] = None,
662 grad_norm: Optional[float] = None, **kwargs) -> None:
663 if self._pbar is None:
664 return
665 postfix = {}
666 if loss is not None:
667 postfix["loss"] = f"{loss:.4f}"
668 if grad_norm is not None:
669 postfix["gnorm"] = f"{grad_norm:.4f}"
670 self._pbar.set_postfix(postfix)
671 self._pbar.update(1)
673 def on_train_end(self, state: "TrainerState", **kwargs) -> None:
674 if self._pbar is not None:
675 self._pbar.close()
676 self._pbar = None
679class MoEMonitorCallback(Callback):
680 """Mixture-of-Experts load-balancing monitor.
682 Delegates to :class:`~hyper_parallel.core.moe_utils.MoEMonitorCallback`
683 for expert bias updates and aux_loss aggregation. Exposes
684 ``last_mean_aux_loss`` so that :class:`LoggingCallback` can include it
685 in the main training loss log line.
687 Config: ``cfg.train.moe_monitor.*`` (see :class:`MoEMonitorConfig`).
688 """
690 def __init__(self, trainer: "BaseTrainer") -> None:
691 """Initialize MoEMonitorCallback from trainer config."""
692 super().__init__(trainer)
693 moe_cfg = getattr(trainer.args, 'moe_monitor', None)
694 self.enabled = getattr(moe_cfg, 'enabled', False) if moe_cfg else False
695 self._impl = None
697 if self.enabled:
698 from hyper_parallel.core.moe_utils import ( # pylint: disable=C0415
699 MoEMonitorCallback as _CoreMoEMonitorCallback,
700 )
701 from hyper_parallel.core.fully_shard.hsdp_utils import ( # pylint: disable=C0415
702 GroupInfo,
703 )
704 lr = getattr(moe_cfg, 'lr', 1e-3)
705 num_recomputations = getattr(moe_cfg, 'num_recomputations', 1)
707 # Resolve DP/TP/CP groups from trainer's device mesh.
708 dp_group = getattr(self.trainer, '_dp_group_info', None)
709 tp_group = None
710 cp_group = None
711 mesh = getattr(self.trainer, 'mesh', None)
712 if mesh is not None:
713 for name, attr_name in [("tp", "tp_group"), ("cp", "cp_group")]:
714 try:
715 raw_group = mesh.get_group(name)
716 group_info = GroupInfo(
717 group_name=name, group=raw_group,
718 rank_size=raw_group.size(),
719 )
720 if attr_name == "tp_group":
721 tp_group = group_info
722 else:
723 cp_group = group_info
724 except (KeyError, ValueError, AttributeError):
725 pass
727 self._impl = _CoreMoEMonitorCallback(
728 model=self.trainer.model,
729 lr=lr,
730 dp_group=dp_group,
731 tp_group=tp_group,
732 cp_group=cp_group,
733 num_recomputations=num_recomputations,
734 )
736 @property
737 def last_mean_aux_loss(self) -> Optional[float]:
738 """Mean aux_loss across MoE layers from the last ``on_step_end``."""
739 if self._impl is not None:
740 return self._impl.last_mean_aux_loss
741 return None
743 def on_train_begin(self, state: "TrainerState", **kwargs) -> None:
744 """Log one-time confirmation when MoE monitoring is enabled."""
745 if self.enabled and platform.get_rank() == 0:
746 logger.info("MoEMonitorCallback: MoE expert-load monitoring enabled")
748 def on_step_end(self, state: "TrainerState", *, loss: float = None,
749 grad_norm: float = None, **kwargs) -> None:
750 """Delegate expert bias update to core MoEMonitorCallback."""
751 if self._impl is not None:
752 self._impl.on_step_end()
754 def on_substep_end(self, state: "TrainerState", **kwargs) -> None:
755 """No-op; expert bias updates happen in on_step_end."""
758class TrainingStateMonitorCallback(Callback):
759 """Save training loss / gradient scalar monitors to TensorBoard and log.
761 Config: ``cfg.train.monitor.*``. This mirrors the first-stage MindFormers
762 training-state monitor surface for loss and gradient norms while reusing
763 HyperParallel's callback lifecycle.
764 """
766 _SUPPORTED_FORMATS = frozenset({"tensorboard", "log"})
768 def __init__(self, trainer: "BaseTrainer") -> None:
769 super().__init__(trainer)
770 train_cfg = getattr(trainer.args, "train", None)
771 self.cfg = getattr(train_cfg, "monitor", None)
772 if self.cfg is None:
773 self.cfg = getattr(trainer.args, "monitor", None)
775 self.enabled = bool(getattr(self.cfg, "monitor_on", False)) if self.cfg else False
776 self.dump_path = getattr(self.cfg, "dump_path", "./dump") if self.cfg else "./dump"
777 self.step_interval = int(getattr(self.cfg, "step_interval", 1) if self.cfg else 1)
778 if self.step_interval < 1:
779 raise ValueError("train.monitor.step_interval must be >= 1.")
781 self.local_loss_format = self._parse_formats("local_loss_format")
782 self.device_local_loss_format = self._parse_formats("device_local_loss_format")
783 self.local_norm_format = self._parse_formats("local_norm_format")
784 self.device_local_norm_format = self._parse_formats("device_local_norm_format")
786 raw_patterns = getattr(self.cfg, "target", None) if self.cfg else None
787 if isinstance(raw_patterns, str):
788 raw_patterns = [raw_patterns]
789 self._target_patterns = (
790 [re.compile(pattern) for pattern in raw_patterns]
791 if raw_patterns else None
792 )
793 self._invert = bool(getattr(self.cfg, "invert", False)) if self.cfg else False
794 self._rank_writer = None
795 self._global_writer = None
796 self._rank = None
797 self._step_weighted_loss_sum = 0.0
798 self._step_token_count = 0
799 self.last_local_loss = None
800 self.last_local_grad_norm = None
801 self.last_monitor_step = -1
802 self._state = None
803 self._pending_log_metrics = {}
804 self._grad_hook_handles = []
805 self._hook_grad_stats = {}
808 def _parse_formats(self, field_name: str) -> tuple[str, ...]:
809 """Parse configured monitor output formats for a metric field."""
810 value = getattr(self.cfg, field_name, None) if self.cfg else None
811 if value is None:
812 return ()
813 if isinstance(value, str):
814 formats = (value,)
815 else:
816 formats = tuple(value)
817 unknown = sorted(set(formats) - self._SUPPORTED_FORMATS)
818 if unknown:
819 raise ValueError(
820 f"train.monitor.{field_name} only supports "
821 f"{sorted(self._SUPPORTED_FORMATS)}, got {unknown}."
822 )
823 return formats
825 @staticmethod
826 def _to_scalar(value) -> float:
827 """Convert tensor-like values to a Python float."""
828 if value is None:
829 return 0.0
830 if hasattr(value, "to_local"):
831 value = value.to_local()
832 if hasattr(value, "detach"):
833 value = value.detach()
834 if hasattr(value, "float"):
835 value = value.float()
836 if hasattr(value, "item"):
837 return float(value.item())
838 return float(value)
840 @staticmethod
841 def _sanitize_tag(value: str) -> str:
842 value = value.replace(".", "/").replace(" ", "_")
843 return re.sub(r"[^A-Za-z0-9_/\-]", "_", value)
845 def _should_record_param(self, name: str) -> bool:
846 patterns = self._target_patterns
847 if not patterns:
848 matched = True
849 else:
850 matched = any(pattern.search(name) for pattern in patterns)
851 return not matched if self._invert else matched
853 def _uses_tensorboard(self) -> bool:
854 return any(
855 "tensorboard" in formats
856 for formats in (
857 self.local_loss_format,
858 self.device_local_loss_format,
859 self.local_norm_format,
860 self.device_local_norm_format,
861 )
862 )
864 def _write(self, formats: tuple[str, ...], tag: str, value: float,
865 step: int, *, global_metric: bool = False) -> None:
866 if not formats:
867 return
868 if "tensorboard" in formats:
869 writer = self._global_writer if global_metric else self._rank_writer
870 if writer is not None:
871 writer.add_scalar(tag, value, step)
872 if "log" in formats:
873 self._pending_log_metrics[tag] = value
875 def _flush_step_log(self, step: int) -> None:
876 """Print one compact rank-local monitor line for console output."""
877 if not self._pending_log_metrics:
878 return
879 field_map = (
880 ("loss", "loss/local_loss"),
881 ("accum_loss", "loss/device_accum_local_loss"),
882 ("grad_norm", "grad/device_local_norm"),
883 )
884 raw_tags = {tag for _, tag in field_map}
885 parts = []
886 for display_name, tag in field_map:
887 if tag in self._pending_log_metrics:
888 value = self._pending_log_metrics[tag]
889 parts.append(f"{display_name}={float(value):.8f}")
890 for tag in sorted(self._pending_log_metrics):
891 if tag in raw_tags:
892 continue
893 value = self._pending_log_metrics[tag]
894 parts.append(f"{tag}={float(value):.8f}")
895 timestamp = time.strftime("%H:%M:%S")
896 print(
897 f"[{timestamp}][rank{self._rank}][INFO] local_state: "
898 f"step={step} | " + " | ".join(parts),
899 flush=True,
900 )
901 self._pending_log_metrics = {}
903 @staticmethod
904 def _compute_grad_stats(grad) -> dict:
905 """Compute norm-only scalar stats for a hook-captured local gradient."""
906 local_grad = grad.to_local() if hasattr(grad, "to_local") else grad
907 local_float = local_grad.detach().float()
908 sum_sq = local_float.pow(2).sum()
909 return {
910 "sum_sq": float(sum_sq.item()),
911 }
913 def _accumulate_hook_grad_stats(self, name: str, grad) -> None:
914 """Accumulate local gradient stats captured before FSDP/HSDP communication."""
915 stats = self._compute_grad_stats(grad)
916 cached = self._hook_grad_stats.setdefault(
917 name,
918 {
919 "sum_sq": 0.0,
920 },
921 )
922 cached["sum_sq"] += stats["sum_sq"]
923 if (
924 "log" in self.local_norm_format
925 and self._state is not None
926 and self._state.global_step % self.step_interval == 0
927 ):
928 timestamp = time.strftime("%H:%M:%S")
929 print(
930 f"[{timestamp}][rank{self._rank}][INFO] parameter_grad: "
931 f"step={self._state.global_step} | "
932 f"name={name} | grad_norm={math.sqrt(stats['sum_sq']):.8f}",
933 flush=True,
934 )
936 def _clear_hook_grad_stats(self) -> None:
937 """Drop hook-collected local gradient stats for the current step."""
938 self._hook_grad_stats = {}
940 def _remove_grad_hooks(self) -> None:
941 """Remove parameter backward hooks registered by the monitor."""
942 for handle in self._grad_hook_handles:
943 if hasattr(handle, "remove"):
944 handle.remove()
945 self._grad_hook_handles = []
947 def _register_grad_hooks(self) -> None:
948 """Register parameter hooks that capture pre-communication local grads."""
949 if not (self.enabled and any((
950 self.local_norm_format,
951 self.device_local_norm_format,))):
952 return
953 self._remove_grad_hooks()
954 self._clear_hook_grad_stats()
955 for name, param in self.trainer.model.named_parameters():
956 if not self._should_record_param(name) or not getattr(param, "requires_grad", False):
957 continue
958 register_hook = getattr(param, "register_hook", None)
959 if not callable(register_hook):
960 logger.warning(
961 "TrainingStateMonitor: parameter %s does not support register_hook, skip local grad monitor.",
962 name,
963 )
964 continue
966 def hook_fn(grad, param_name=name):
967 self._accumulate_hook_grad_stats(param_name, grad)
968 return grad
970 self._grad_hook_handles.append(register_hook(hook_fn))
972 def on_train_begin(self, state: "TrainerState", **kwargs) -> None:
973 if not self.enabled:
974 return
975 self._state = state
976 self._rank = platform.get_rank()
977 if self._uses_tensorboard():
978 tb_root = os.path.join(self.dump_path, "tensorboard")
979 self._rank_writer = SummaryWriter(
980 os.path.join(tb_root, f"rank_{self._rank}")
981 )
982 if self._rank == 0:
983 self._global_writer = SummaryWriter(os.path.join(tb_root, "global"))
984 logger.info_rank0(
985 "TrainingStateMonitor enabled: dump_path=%s step_interval=%d",
986 self.dump_path, self.step_interval,
987 )
988 self._register_grad_hooks()
990 def on_train_end(self, state: "TrainerState", **kwargs) -> None:
991 self._remove_grad_hooks()
992 self._clear_hook_grad_stats()
993 for writer in (self._rank_writer, self._global_writer):
994 if writer is not None:
995 writer.close()
996 self._rank_writer = None
997 self._global_writer = None
998 self._state = None
1000 def on_substep_end(self, state: "TrainerState", **kwargs) -> None:
1001 if not self.enabled or state.global_step % self.step_interval != 0:
1002 return
1003 substep_info = getattr(state, "substep_info", {})
1004 raw_loss = kwargs.get("raw_loss", substep_info.get("raw_loss"))
1005 if raw_loss is None:
1006 return
1007 micro_tokens = int(kwargs.get("micro_tokens", substep_info.get("micro_tokens", 1)))
1008 loss_value = self._to_scalar(raw_loss)
1009 self._step_weighted_loss_sum += loss_value * micro_tokens
1010 self._step_token_count += micro_tokens
1011 self.last_local_loss = self._step_weighted_loss_sum / max(self._step_token_count, 1)
1012 self.last_monitor_step = state.global_step
1014 tag = "loss/local_loss"
1015 self._write(
1016 self.local_loss_format,
1017 tag,
1018 loss_value,
1019 state.global_step,
1020 global_metric=False,
1021 )
1023 def _record_hook_grad_stats(self, state: "TrainerState") -> None:
1024 """Record local gradient stats captured by backward hooks."""
1025 device_sum_sq = 0.0
1026 for name, _ in self.trainer.model.named_parameters():
1027 if not self._should_record_param(name):
1028 continue
1029 stats = self._hook_grad_stats.get(name)
1030 if stats is None:
1031 continue
1032 sum_sq = float(stats["sum_sq"])
1033 norm = math.sqrt(sum_sq)
1035 device_sum_sq += sum_sq
1037 clean_name = self._sanitize_tag(name)
1038 self._write(
1039 tuple(fmt for fmt in self.local_norm_format if fmt != "log"),
1040 f"grad/local_norm/{clean_name}",
1041 norm,
1042 state.global_step,
1043 )
1045 device_norm = math.sqrt(device_sum_sq)
1046 self.last_local_grad_norm = device_norm
1047 self.last_monitor_step = state.global_step
1048 self._write(
1049 self.device_local_norm_format,
1050 "grad/device_local_norm",
1051 device_norm,
1052 state.global_step,
1053 )
1055 def on_pre_optimizer_step(self, state: "TrainerState", **kwargs) -> None:
1056 if not self.enabled:
1057 return
1058 if state.global_step % self.step_interval != 0:
1059 self._clear_hook_grad_stats()
1060 return
1061 try:
1062 self._record_hook_grad_stats(state)
1063 finally:
1064 self._clear_hook_grad_stats()
1066 def on_step_end(self, state: "TrainerState", *, loss: Optional[float] = None,
1067 grad_norm: Optional[float] = None, **kwargs) -> None:
1068 if not self.enabled:
1069 return
1070 if state.global_step % self.step_interval == 0:
1071 if self._step_token_count > 0:
1072 device_loss = self._step_weighted_loss_sum / self._step_token_count
1073 self._write(
1074 self.device_local_loss_format,
1075 "loss/device_accum_local_loss",
1076 device_loss,
1077 state.global_step,
1078 )
1079 if self._rank == 0 and loss is not None:
1080 self._write(
1081 ("tensorboard",) if self._uses_tensorboard() else (),
1082 "loss/global_loss",
1083 float(loss),
1084 state.global_step,
1085 global_metric=True,
1086 )
1087 if self._rank == 0 and grad_norm is not None:
1088 self._write(
1089 ("tensorboard",) if self._uses_tensorboard() else (),
1090 "grad/global_grad_norm",
1091 float(grad_norm),
1092 state.global_step,
1093 global_metric=True,
1094 )
1095 self._flush_step_log(state.global_step)
1097 self._step_weighted_loss_sum = 0.0
1098 self._step_token_count = 0
1101class GradientHealthCallback(Callback):
1102 """Detect NaN / Inf grad_norm and raise / warn.
1104 Hooks ``on_pre_optimizer_step`` — which fires after ``clip_grad_norm_``
1105 and before ``optimizer.step()``. ``grad_norm`` at that point is a plain
1106 scalar produced by hyper's DTensor-aware clipper. If it's not finite, the
1107 optimizer.step() would silently corrupt weights with NaN; we want to
1108 surface it immediately.
1110 Config: ``cfg.train.debug.check_nan_inf``.
1111 """
1113 def __init__(self, trainer: "BaseTrainer") -> None:
1114 super().__init__(trainer)
1115 train_cfg = getattr(trainer.args, "train", None)
1116 debug_cfg = getattr(train_cfg, "debug", None) if train_cfg is not None else None
1117 if debug_cfg is None:
1118 debug_cfg = getattr(trainer.args, 'debug', None)
1119 self.enabled = (
1120 getattr(debug_cfg, 'check_nan_inf', False) if debug_cfg else False
1121 )
1123 def on_pre_optimizer_step(self, state: "TrainerState", *,
1124 grad_norm: Optional[float] = None,
1125 **kwargs) -> None:
1126 if not self.enabled or grad_norm is None:
1127 return
1128 if math.isnan(grad_norm) or math.isinf(grad_norm):
1129 # Always log on every rank — divergence may be rank-local.
1130 logger.error(
1131 "GradientHealthCallback: grad_norm=%s at step %d "
1132 "(NaN/Inf). Optimizer.step would corrupt weights.",
1133 grad_norm, state.global_step,
1134 )
1135 # Raise on rank 0 only; other ranks will be torn down by NCCL.
1136 if platform.get_rank() == 0:
1137 raise RuntimeError(
1138 f"Non-finite grad_norm={grad_norm} at "
1139 f"step {state.global_step}. "
1140 "Disable cfg.train.debug.check_nan_inf to skip this guard."
1141 )
1144class GCCallback(Callback):
1145 """Explicit garbage-collection scheduler.
1147 Python's cyclic GC can stall large training jobs when it decides to run;
1148 forcing a collection every N steps — outside the compute hot path —
1149 keeps pauses predictable.).
1151 Config: ``cfg.train.debug.gc_steps`` (``0`` disables).
1152 """
1154 def __init__(self, trainer: "BaseTrainer") -> None:
1155 super().__init__(trainer)
1156 debug_cfg = getattr(trainer.args, 'debug', None)
1157 self.gc_steps = (
1158 getattr(debug_cfg, 'gc_steps', 0) if debug_cfg else 0
1159 )
1160 if self.gc_steps > 0:
1161 # Disable the automatic generational collector; we'll drive it.
1162 gc.disable()
1163 logger.info("GCCallback: Python gc.collect every %d steps "
1164 "(auto GC disabled)", self.gc_steps)
1166 def on_step_end(self, state: "TrainerState", *,
1167 loss: Optional[float] = None,
1168 grad_norm: Optional[float] = None, **kwargs) -> None:
1169 if self.gc_steps <= 0:
1170 return
1171 if state.global_step % self.gc_steps != 0:
1172 return
1173 gc.collect()
1176class TensorBoardCallback(Callback):
1177 """TensorBoard scalar writer — STUB (not verified).
1179 Hook reserved for ``torch.utils.tensorboard.SummaryWriter`` integration.
1180 Not yet verified; if you enable ``args.tensorboard.enabled`` we emit
1181 a one-time warning so missing TB scalars are visible. To implement:
1182 open SummaryWriter in ``on_train_begin``, write scalars in ``on_log``,
1183 close in ``on_train_end``.
1184 """
1186 def __init__(self, trainer: "BaseTrainer") -> None:
1187 super().__init__(trainer)
1188 tb_cfg = getattr(trainer.args, 'tensorboard', None)
1189 if getattr(tb_cfg, 'enabled', False) and platform.get_rank() == 0:
1190 logger.warning(
1191 "TensorBoardCallback: enabled=True but the implementation "
1192 "is a stub — nothing is written to TensorBoard. Implement "
1193 "before relying on TB scalars."
1194 )
1197class MemoryMonitorCallback(Callback):
1198 """Peak / current device memory monitor — STUB (not verified).
1200 Hook reserved for ``platform.get_device_handle().memory_allocated`` /
1201 ``max_memory_allocated`` polling. Not yet verified; if you enable
1202 ``args.memory_monitor.enabled`` we emit a one-time warning so missing
1203 memory logs are visible. To implement: poll the device handle in
1204 ``on_step_end`` gated by ``log_steps`` and log
1205 ``cur=...GB peak=...GB``.
1206 """
1208 def __init__(self, trainer: "BaseTrainer") -> None:
1209 super().__init__(trainer)
1210 cfg = getattr(trainer.args, 'memory_monitor', None)
1211 if getattr(cfg, 'enabled', False) and platform.get_rank() == 0:
1212 logger.warning(
1213 "MemoryMonitorCallback: enabled=True but the implementation "
1214 "is a stub — no memory stats are emitted. Implement before "
1215 "relying on these logs."
1216 )