Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / callbacks / base.py: 69%

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. 

16 

17dispatched explicitly in ``on_step_end`` etc. Engineer sees all callbacks and 

18order at a glance. 

19 

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 threading 

29import time 

30from typing import TYPE_CHECKING, Optional 

31 

32import torch 

33 

34from hyper_parallel import get_platform 

35from hyper_parallel.core.distributed_checkpoint import load as dcp_load, save as dcp_save 

36from hyper_parallel.core.distributed_checkpoint.offline_transform import ( 

37 save_state_dict_as_huggingface_format, 

38) 

39from hyper_parallel.core.fully_shard.api import get_model_state_dict 

40 

41platform = get_platform() 

42 

43if TYPE_CHECKING: 

44 from hyper_parallel.trainer.base import BaseTrainer, TrainerState 

45 

46logger = logging.getLogger(__name__) 

47 

48 

49class Callback: 

50 """Base class for all trainer callbacks. 

51 

52 Each callback holds a reference to the trainer for accessing model, 

53 optimizer, state, and config. Subclass and override the hooks you need. 

54 

55 Args: 

56 trainer: The BaseTrainer instance. 

57 """ 

58 

59 def __init__(self, trainer: "BaseTrainer") -> None: 

60 self.trainer = trainer 

61 

62 # ------------------------------------------------------------------ 

63 # Lifecycle hooks 

64 # ------------------------------------------------------------------ 

65 

66 def on_init_end(self, state: "TrainerState", **kwargs) -> None: 

67 """Called once at the end of ``BaseTrainer.__init__`` / subclass init. 

68 

69 At this point every ``_build_*`` has run — model is parallelised, 

70 optimizer/scheduler/dataloader are built, callbacks are constructed. 

71 Use this for one-shot setup that must see the FINAL trainer state 

72 (e.g. logging the parameter count, opening a TensorBoard writer 

73 keyed by run_id, validating user config against the built model). 

74 """ 

75 

76 def on_train_begin(self, state: "TrainerState", **kwargs) -> None: 

77 """Called at the start of ``train()`` (before any optimizer.step). 

78 

79 ``CheckpointCallback`` runs resume here, so when this hook fires 

80 ``state.global_step`` may already be > 0 if a checkpoint was loaded. 

81 """ 

82 

83 def on_train_end(self, state: "TrainerState", **kwargs) -> None: 

84 """Called at the end of training (before ``destroy_process_group``). 

85 

86 Final checkpoints, profiler stops, W&B finish, etc. happen here. 

87 """ 

88 

89 def on_epoch_begin(self, state: "TrainerState", **kwargs) -> None: 

90 """Called at the start of each epoch.""" 

91 

92 def on_epoch_end(self, state: "TrainerState", **kwargs) -> None: 

93 """Called at the end of each epoch.""" 

94 

95 def on_step_begin(self, state: "TrainerState", **kwargs) -> None: 

96 """Called at the start of each training step (before fwd of mb 0).""" 

97 

98 def on_step_end(self, state: "TrainerState", *, loss: float = None, 

99 grad_norm: float = None, **kwargs) -> None: 

100 """Called at the end of each training step (after optimizer.step).""" 

101 

102 def on_substep_end(self, state: "TrainerState", **kwargs) -> None: 

103 """Called after each micro-batch fwd+bwd (gradient accumulation).""" 

104 

105 def on_pre_optimizer_step(self, state: "TrainerState", *, 

106 grad_norm: float = None, **kwargs) -> None: 

107 """Called after grad clip, before ``optimizer.step``. 

108 

109 ``grad_norm`` here is the post-clip scalar produced by hyper's 

110 DTensor-aware clipper — use it to detect NaN/Inf or to log the 

111 effective clip ratio. 

112 """ 

113 

114 def on_log(self, state: "TrainerState", *, metrics: dict, **kwargs) -> None: 

115 """Called when ``LoggingCallback`` emits a structured metrics record. 

116 

117 Reuse this hook in TensorBoard / W&B / external metric sinks so 

118 every logging backend sees the SAME record. Avoids three callbacks 

119 each computing throughput / lr independently. 

120 

121 Args: 

122 metrics: Dict containing at minimum ``step``, ``loss``, 

123 ``grad_norm``, ``lr``, ``step_time``; throughput fields 

124 (``tokens_per_sec``, ``tflops``, ``mfu``) are present iff 

125 ``logging.report_throughput`` is on. 

126 """ 

127 

128 def on_save(self, state: "TrainerState", *, checkpoint_dir: str, 

129 **kwargs) -> None: 

130 """Called immediately after ``CheckpointCallback`` finishes a save. 

131 

132 Use to upload to remote storage, register the ckpt with an 

133 experiment tracker, or trigger downstream eval jobs. ``checkpoint_dir`` 

134 is the on-disk path containing model shards + optimizer/scheduler/RNG/ 

135 dataloader/extra_state. 

136 """ 

137 

138 def on_load(self, state: "TrainerState", *, checkpoint_dir: str, 

139 **kwargs) -> None: 

140 """Called immediately after ``CheckpointCallback`` finishes a resume. 

141 

142 Use to verify the resumed step matches expectations, log the 

143 restore event, or seed downstream callbacks with the resumed state. 

144 """ 

145 

146 def on_evaluate(self, state: "TrainerState", *, metrics: dict = None, 

147 **kwargs) -> None: 

148 """Called when an evaluation pass completes. 

149 

150 Currently triggered as a stub from ``EvalCallback``; once a real 

151 eval loop lands the callback will pass back the eval ``metrics`` 

152 dict for sinks (TensorBoard / W&B) to log. 

153 """ 

154 

155 

156class LoggingCallback(Callback): 

157 """Log training metrics: loss, grad_norm, lr, throughput. 

158 

159 """ 

160 

161 def __init__(self, trainer: "BaseTrainer") -> None: 

162 super().__init__(trainer) 

163 log_cfg = getattr(trainer.args, 'logging', None) 

164 self.log_steps = getattr(log_cfg, 'log_steps', 10) if log_cfg else 10 

165 self.report_global_loss = ( 

166 getattr(log_cfg, 'report_global_loss', False) if log_cfg else False 

167 ) 

168 self.report_throughput = ( 

169 getattr(log_cfg, 'report_throughput', True) if log_cfg else True 

170 ) 

171 self.model_flops_per_token = ( 

172 getattr(log_cfg, 'model_flops_per_token', None) if log_cfg else None 

173 ) 

174 self.peak_tflops = ( 

175 getattr(log_cfg, 'peak_tflops', None) if log_cfg else None 

176 ) 

177 # Estimate per-step tokens as upper bound (batch × seq_len). Real 

178 # token count is available per step via ``last_global_tokens`` that 

179 # ``BaseTrainer.train_step`` stashes onto the trainer. 

180 gbs = getattr(trainer.args.train, 'global_batch_size', 1) 

181 seq_len = getattr(trainer.args.data, 'max_seq_len', 1) 

182 self._tokens_per_step_est = int(gbs) * int(seq_len) 

183 self._step_start_time = 0.0 

184 

185 def on_step_begin(self, state: "TrainerState", **kwargs) -> None: 

186 self._step_start_time = time.time() 

187 

188 def on_step_end(self, state: "TrainerState", *, loss: float = None, 

189 grad_norm: float = None, **kwargs) -> None: 

190 if state.global_step % self.log_steps != 0: 

191 return 

192 

193 elapsed = max(time.time() - self._step_start_time, 1e-9) 

194 lr = 0.0 

195 if self.trainer.lr_scheduler is not None: 

196 lr = self.trainer.lr_scheduler.get_last_lr()[0] 

197 

198 metrics = { 

199 "step": state.global_step, 

200 # 8-decimal precision keeps fp32 sub-bf16 differences visible 

201 # in the log for sanity comparisons across runs. 

202 "loss": f"{loss:.8f}" if loss is not None else "N/A", 

203 "grad_norm": ( 

204 f"{grad_norm:.8f}" if grad_norm is not None else "N/A" 

205 ), 

206 "lr": f"{lr:.2e}", 

207 "step_time": f"{elapsed:.2f}s", 

208 } 

209 

210 tokens_per_sec = None 

211 if self.report_throughput: 

212 # Prefer real per-step token count stashed by train_step; fall back 

213 # to the estimate until the first step sets it (declared None). 

214 tokens = getattr(self.trainer, '_last_global_tokens', None) 

215 if tokens is None: 

216 tokens = self._tokens_per_step_est 

217 tokens_per_sec = tokens / elapsed 

218 metrics["tokens_per_sec"] = f"{tokens_per_sec:,.0f}" 

219 

220 if self.model_flops_per_token and self.peak_tflops: 

221 # Observed TFLOPS = tokens/sec × flops/token / 1e12. 

222 # MFU = observed / (peak × world_size). 

223 world = max(platform.get_world_size(), 1) 

224 observed_tflops = ( 

225 tokens_per_sec * self.model_flops_per_token / 1e12 

226 ) 

227 mfu = observed_tflops / (self.peak_tflops * world) 

228 metrics["tflops"] = f"{observed_tflops:.1f}" 

229 metrics["mfu"] = f"{mfu * 100:.1f}%" 

230 

231 # Include aux_loss from MoEMonitorCallback when available. 

232 moe_cb = getattr(self.trainer, 'moe_monitor_callback', None) 

233 aux_loss = getattr(moe_cb, 'last_mean_aux_loss', None) if moe_cb is not None else None 

234 if aux_loss is not None: 

235 metrics["aux_loss"] = f"{aux_loss:.6f}" 

236 

237 logger.info_rank0(" | ".join(f"{k}={v}" for k, v in metrics.items())) 

238 

239 record = { 

240 "step": state.global_step, 

241 "loss": loss, 

242 "grad_norm": grad_norm, 

243 "lr": lr, 

244 "step_time": elapsed, 

245 "tokens_per_sec": tokens_per_sec, 

246 "aux_loss": aux_loss, 

247 } 

248 state.log_history.append(record) 

249 

250 # Fan-out to other log-event listeners (TB / W&B / sinks). 

251 dispatch = getattr(self.trainer, "dispatch_log_event", None) 

252 if dispatch is not None: 

253 dispatch(record) 

254 

255 

256class CheckpointCallback(Callback): 

257 """Save distributed checkpoints and handle resume. 

258 

259 Uses hyper's own DCP ``save`` / ``load`` APIs. 

260 """ 

261 

262 def __init__(self, trainer: "BaseTrainer") -> None: 

263 super().__init__(trainer) 

264 ckpt_cfg = getattr(trainer.args, 'checkpoint', None) 

265 self.save_steps = getattr(ckpt_cfg, 'save_steps', 0) if ckpt_cfg else 0 

266 self.output_dir = ( 

267 getattr(ckpt_cfg, 'output_dir', 'outputs') if ckpt_cfg else 'outputs' 

268 ) 

269 self.load_path = ( 

270 getattr(ckpt_cfg, 'load_path', None) if ckpt_cfg else None 

271 ) 

272 self.save_async = ( 

273 getattr(ckpt_cfg, 'save_async', False) if ckpt_cfg else False 

274 ) 

275 self._last_saved_step = -1 

276 self._save_thread = None # async save worker 

277 

278 def on_train_begin(self, state: "TrainerState", **kwargs) -> None: 

279 """Resume from checkpoint: model + optimizer + lr_scheduler + step + RNG. 

280 

281 RFC DoD: "Save → resume → 续训 loss 一致（含 dataloader + RNG 恢复）" 

282 """ 

283 if not self.load_path: 

284 return 

285 try: 

286 # pylint: disable=C0415 

287 # Non-model artifacts (optimizer/scheduler/RNG) are plain dicts — 

288 # use torch.save/load, matching the save side. 

289 

290 if not os.path.isdir(self.load_path): 

291 logger.warning("Checkpoint path not found: %s", self.load_path) 

292 return 

293 

294 # 1. Restore model via hyper DCP 

295 model_sd = self.trainer.model.state_dict() 

296 dcp_load(model_sd, checkpoint_id=self.load_path, use_collectives=False) 

297 self.trainer.model.load_state_dict(model_sd) 

298 logger.info("Model restored from %s", self.load_path) 

299 

300 # 2. Restore extra state (step, epoch) 

301 extra_path = os.path.join(self.load_path, "extra_state.json") 

302 if os.path.isfile(extra_path): 

303 with open(extra_path, encoding="utf-8") as f: 

304 extra = json.load(f) 

305 state.global_step = extra.get("global_step", 0) 

306 state.epoch = extra.get("epoch", 0) 

307 logger.info("Resumed at step=%d, epoch=%d", 

308 state.global_step, state.epoch) 

309 

310 # 3. Restore optimizer 

311 optim_path = os.path.join(self.load_path, f"optimizer_rank{platform.get_rank()}.pt") 

312 if os.path.isfile(optim_path) and self.trainer.optimizer: 

313 optim_sd = torch.load(optim_path, map_location="cpu", weights_only=True) 

314 self.trainer.optimizer.load_state_dict(optim_sd) 

315 logger.info("Optimizer restored") 

316 

317 # 4. Restore LR scheduler 

318 sched_path = os.path.join(self.load_path, "scheduler.pt") 

319 if os.path.isfile(sched_path) and self.trainer.lr_scheduler: 

320 sched_sd = torch.load(sched_path, map_location="cpu", weights_only=True) 

321 self.trainer.lr_scheduler.load_state_dict(sched_sd) 

322 logger.info("LR scheduler restored") 

323 

324 # 5. Restore RNG state 

325 rng_path = os.path.join(self.load_path, f"rng_rank{platform.get_rank()}.pt") 

326 if os.path.isfile(rng_path): 

327 rng_state = torch.load(rng_path, map_location="cpu", weights_only=True) 

328 platform.set_rng_state(rng_state) 

329 logger.info("RNG state restored") 

330 

331 # 6. Restore dataloader position (StatefulDataLoader) 

332 dl_path = os.path.join(self.load_path, f"dataloader_rank{platform.get_rank()}.pt") 

333 if os.path.isfile(dl_path) and hasattr(self.trainer, 'train_dataloader'): 

334 dl_state = torch.load(dl_path, map_location="cpu", weights_only=False) 

335 self.trainer.train_dataloader.load_state_dict(dl_state) 

336 logger.info("Dataloader state restored") 

337 

338 # Fan-out the load event so other callbacks (TensorBoard / 

339 # W&B / external trackers) can record the resume. 

340 dispatch = getattr(self.trainer, "dispatch_load_event", None) 

341 if dispatch is not None: 

342 dispatch(self.load_path) 

343 

344 except (OSError, RuntimeError, ValueError) as exc: 

345 logger.warning("Failed to load checkpoint from %s: %s", self.load_path, exc) 

346 

347 def on_step_end(self, state: "TrainerState", *, loss: float = None, 

348 grad_norm: float = None, **kwargs) -> None: 

349 if self.save_steps <= 0: 

350 return 

351 if state.global_step % self.save_steps != 0: 

352 return 

353 if state.global_step == self._last_saved_step: 

354 return 

355 self._dispatch_save(state) 

356 

357 def on_train_end(self, state: "TrainerState", **kwargs) -> None: 

358 """Save final checkpoint (synchronously, to guarantee completion).""" 

359 # Wait for any outstanding async save first so the two don't race on 

360 # the same directory / state-dict iterator. 

361 self._join_pending() 

362 if self.save_steps > 0 and state.global_step != self._last_saved_step: 

363 # Final save always sync — the process is about to exit. 

364 self._save(state) 

365 

366 # --- async plumbing ------------------------------------------------- 

367 def _dispatch_save(self, state: "TrainerState") -> None: 

368 """Route to sync or async save based on ``save_async`` flag.""" 

369 if not self.save_async: 

370 self._save(state) 

371 return 

372 # Wait for previous save to finish before starting a new one; saving 

373 # twice concurrently would double RAM and race the filesystem. 

374 self._join_pending() 

375 # pylint: disable=C0415 

376 # Snapshot state fields so the worker doesn't see later mutations. 

377 snap_step = state.global_step 

378 snap_epoch = state.epoch 

379 state_snapshot = copy.copy(state) 

380 state_snapshot.global_step = snap_step 

381 state_snapshot.epoch = snap_epoch 

382 self._save_thread = threading.Thread( 

383 target=self._save, 

384 args=(state_snapshot,), 

385 name=f"ckpt-save-step{snap_step}", 

386 daemon=True, 

387 ) 

388 self._save_thread.start() 

389 logger.info_rank0( 

390 "Checkpoint save for step %d dispatched async (thread=%s)", 

391 snap_step, self._save_thread.name, 

392 ) 

393 

394 def _join_pending(self) -> None: 

395 """Block until any running async save finishes.""" 

396 t = self._save_thread 

397 if t is not None and t.is_alive(): 

398 logger.info_rank0( 

399 "Waiting for prior async ckpt save (%s)...", t.name, 

400 ) 

401 t.join() 

402 self._save_thread = None 

403 

404 def _save(self, state: "TrainerState") -> None: 

405 """Save complete training state: model + optimizer + scheduler + step + RNG. 

406 

407 RFC DoD: "Save → resume → 续训 loss 一致（含 dataloader + RNG 恢复）" 

408 """ 

409 # Optimizer/scheduler/RNG state dicts are plain Python dicts, not 

410 # nn.Module — platform.save_checkpoint expects Module (safetensors). 

411 # Use torch.save/load for these non-model artifacts. 

412 save_dir = os.path.join(self.output_dir, f"step_{state.global_step}") 

413 os.makedirs(save_dir, exist_ok=True) 

414 rank = platform.get_rank() 

415 

416 try: 

417 # 1. Model — via hyper DCP (each rank saves its own shards) 

418 model_sd = self.trainer.model.state_dict() 

419 dcp_save(model_sd, checkpoint_id=save_dir, use_collectives=False) 

420 

421 # 2. Optimizer — per-rank 

422 if self.trainer.optimizer: 

423 optim_path = os.path.join(save_dir, f"optimizer_rank{rank}.pt") 

424 torch.save(self.trainer.optimizer.state_dict(), optim_path) 

425 

426 # 3. LR scheduler 

427 if self.trainer.lr_scheduler and rank == 0: 

428 sched_path = os.path.join(save_dir, "scheduler.pt") 

429 torch.save(self.trainer.lr_scheduler.state_dict(), sched_path) 

430 

431 # 4. Extra state: global_step, epoch 

432 if rank == 0: 

433 extra = { 

434 "global_step": state.global_step, 

435 "epoch": state.epoch, 

436 } 

437 extra_path = os.path.join(save_dir, "extra_state.json") 

438 with open(extra_path, "w", encoding="utf-8") as f: 

439 json.dump(extra, f) 

440 

441 # 5. RNG state — per-rank via platform API 

442 rng_state = platform.get_rng_state() 

443 rng_path = os.path.join(save_dir, f"rng_rank{rank}.pt") 

444 torch.save(rng_state, rng_path) 

445 

446 # 6. Dataloader position — per-rank (StatefulDataLoader) 

447 if hasattr(self.trainer, 'train_dataloader') and hasattr( 

448 self.trainer.train_dataloader, 'state_dict' 

449 ): 

450 dl_path = os.path.join(save_dir, f"dataloader_rank{rank}.pt") 

451 torch.save(self.trainer.train_dataloader.state_dict(), dl_path) 

452 

453 self._last_saved_step = state.global_step 

454 logger.info_rank0("Checkpoint saved to %s", save_dir) 

455 

456 # Fan-out the save event so other callbacks (W&B artifact 

457 # upload, remote-storage sync, downstream eval triggers) can 

458 # observe the new checkpoint without coupling to ckpt internals. 

459 dispatch = getattr(self.trainer, "dispatch_save_event", None) 

460 if dispatch is not None: 

461 dispatch(save_dir) 

462 

463 except (OSError, RuntimeError, ValueError) as exc: 

464 logger.warning("Failed to save checkpoint: %s", exc) 

465 

466 # HF format export is handled by SafetensorsExportCallback (separate concern). 

467 

468 

469class SafetensorsExportCallback(Callback): 

470 """Export model weights in HuggingFace safetensor format. 

471 

472 Separated from CheckpointCallback per RFC Section 5.2. 

473 Uses ``get_model_state_dict`` with ``full_state_dict=True`` to gather 

474 all FSDP shards into a full state dict before saving. 

475 

476 """ 

477 

478 def __init__(self, trainer: "BaseTrainer") -> None: 

479 super().__init__(trainer) 

480 ckpt_cfg = getattr(trainer.args, 'checkpoint', None) 

481 self.enabled = getattr(ckpt_cfg, 'save_hf_weights', False) if ckpt_cfg else False 

482 self.save_steps = getattr(ckpt_cfg, 'save_steps', 0) if ckpt_cfg else 0 

483 self.output_dir = getattr(ckpt_cfg, 'output_dir', 'outputs') if ckpt_cfg else 'outputs' 

484 self._last_saved_step = -1 

485 

486 def on_step_end(self, state: "TrainerState", *, loss: Optional[float] = None, 

487 grad_norm: Optional[float] = None, **kwargs) -> None: 

488 if not self.enabled or self.save_steps <= 0: 

489 return 

490 if state.global_step % self.save_steps != 0: 

491 return 

492 if state.global_step == self._last_saved_step: 

493 return 

494 self._export(state) 

495 

496 def on_train_end(self, state: "TrainerState", **kwargs) -> None: 

497 if self.enabled and self.save_steps > 0 and state.global_step != self._last_saved_step: 

498 self._export(state) 

499 

500 def _export(self, state: "TrainerState") -> None: 

501 """Gather full state dict from FSDP shards and save in HF format. 

502 

503 Routes through ``spec.state_dict_adapter().save_hf_state_dict`` when 

504 the model's ``ModelSpec`` provides one, so per-model HF tensor 

505 renaming and per-expert packing live in the model package, not in 

506 this generic callback. Falls back to the legacy 

507 ``save_state_dict_as_huggingface_format`` path when the spec has no 

508 adapter (keeps ad-hoc / template models working). 

509 """ 

510 # pylint: disable=C0415 

511 

512 rank = platform.get_rank() 

513 save_dir = os.path.join(self.output_dir, f"step_{state.global_step}", "hf_ckpt") 

514 

515 try: 

516 # ``StateDictOptions`` is a torch-backend type; hyper does not yet 

517 # provide a wrapper, so the trainer reaches into torch directly. 

518 # pylint: disable=C0415 

519 from torch.distributed.checkpoint.state_dict import StateDictOptions 

520 # full_state_dict=True gathers all FSDP shards; cpu_offload avoids OOM 

521 options = StateDictOptions(full_state_dict=True, cpu_offload=True) 

522 full_sd = get_model_state_dict(self.trainer.model, options=options) 

523 

524 if rank == 0: 

525 os.makedirs(save_dir, exist_ok=True) 

526 

527 # Prefer the model-specific save adapter (closes the load/save 

528 # loop via the ModelSpec contract). When absent, fall back to 

529 # the generic offline-transform path. 

530 spec = getattr(self.trainer, "spec", None) 

531 adapter_cls = getattr(spec, "state_dict_adapter", None) if spec else None 

532 save_fn = ( 

533 getattr(adapter_cls(), "save_hf_state_dict", None) 

534 if adapter_cls is not None else None 

535 ) 

536 if save_fn is not None: 

537 hf_sd = save_fn(full_sd, self.trainer.model.config) 

538 from safetensors.torch import save_file # pylint: disable=C0415 

539 save_file(hf_sd, os.path.join(save_dir, "model.safetensors")) 

540 logger.info( 

541 "HF checkpoint saved via %s.save_hf_state_dict to %s", 

542 adapter_cls.__name__, save_dir, 

543 ) 

544 else: 

545 save_state_dict_as_huggingface_format(full_sd, save_dir) 

546 logger.info( 

547 "HF checkpoint saved (no adapter on spec) to %s", save_dir, 

548 ) 

549 

550 self._last_saved_step = state.global_step 

551 

552 except (OSError, RuntimeError, ValueError) as exc: 

553 logger.warning_rank0("Failed to save HF checkpoint: %s", exc) 

554 

555 

556class EvalCallback(Callback): 

557 """Evaluation callback stub. 

558 

559 Full evaluation is not yet implemented. This stub logs a warning whenever 

560 an evaluation trigger is received so the absence of eval is visible in 

561 training logs rather than silently skipped. 

562 """ 

563 

564 def on_step_end(self, state: "TrainerState", *, loss: Optional[float] = None, 

565 grad_norm: Optional[float] = None, **kwargs) -> None: 

566 eval_cfg = getattr(self.trainer.args, 'eval', None) 

567 eval_steps = getattr(eval_cfg, 'eval_steps', 0) if eval_cfg else 0 

568 if eval_steps > 0 and state.global_step % eval_steps == 0: 

569 if platform.get_rank() == 0: 

570 logger.warning( 

571 "EvalCallback: evaluation not implemented (step=%d)", state.global_step 

572 ) 

573 

574 

575class ProfilerCallback(Callback): 

576 """Training profiler callback — STUB (not verified). 

577 

578 Hook reserved for ``torch.profiler.profile`` integration. Not yet 

579 verified against the trainer; if you enable ``args.profiler.enabled`` 

580 we emit a one-time warning so the absence of profiling traces is 

581 visible. To implement: wire ``torch.profiler.profile`` start/step/stop 

582 in ``on_train_begin`` / ``on_step_end`` / ``on_train_end``. 

583 """ 

584 

585 def __init__(self, trainer: "BaseTrainer") -> None: 

586 super().__init__(trainer) 

587 prof_cfg = getattr(trainer.args, 'profiler', None) 

588 if getattr(prof_cfg, 'enabled', False) and platform.get_rank() == 0: 

589 logger.warning( 

590 "ProfilerCallback: enabled=True but the implementation is " 

591 "a stub — torch profiler is NOT started. Implement before " 

592 "relying on traces." 

593 ) 

594 

595 

596class WandbCallback(Callback): 

597 """Weights & Biases logging callback — STUB (not verified). 

598 

599 Hook reserved for W&B integration. Not yet verified; if you enable 

600 ``args.wandb.enabled`` we emit a one-time warning so missing W&B logs 

601 are visible. To implement: wire ``wandb.init`` / ``wandb.log`` / 

602 ``wandb.finish`` in ``on_train_begin`` / ``on_step_end`` / 

603 ``on_train_end`` and verify against a real W&B run. 

604 """ 

605 

606 def __init__(self, trainer: "BaseTrainer") -> None: 

607 super().__init__(trainer) 

608 wandb_cfg = getattr(trainer.args, 'wandb', None) 

609 if getattr(wandb_cfg, 'enabled', False) and platform.get_rank() == 0: 

610 logger.warning( 

611 "WandbCallback: enabled=True but the implementation is a " 

612 "stub — nothing is sent to W&B. Implement before relying on " 

613 "W&B dashboards." 

614 ) 

615 

616 

617class ProgressCallback(Callback): 

618 """tqdm progress bar callback (rank 0 only). 

619 

620 Displays a progress bar over training steps with live loss and grad_norm 

621 metrics. Requires ``tqdm``; degrades gracefully if not installed. 

622 """ 

623 

624 def __init__(self, trainer: "BaseTrainer") -> None: 

625 super().__init__(trainer) 

626 self._pbar = None 

627 

628 def on_train_begin(self, state: "TrainerState", **kwargs) -> None: 

629 if platform.get_rank() != 0: 

630 return 

631 try: 

632 # pylint: disable=C0415 

633 from tqdm import tqdm # pylint: disable=C0415 # optional dep 

634 self._pbar = tqdm( 

635 total=state.max_steps, 

636 initial=state.global_step, 

637 desc="Training", 

638 unit="step", 

639 dynamic_ncols=True, 

640 ) 

641 except ImportError: 

642 logger.warning("ProgressCallback: 'tqdm' not installed — progress bar disabled") 

643 

644 def on_step_end(self, state: "TrainerState", *, loss: Optional[float] = None, 

645 grad_norm: Optional[float] = None, **kwargs) -> None: 

646 if self._pbar is None: 

647 return 

648 postfix = {} 

649 if loss is not None: 

650 postfix["loss"] = f"{loss:.4f}" 

651 if grad_norm is not None: 

652 postfix["gnorm"] = f"{grad_norm:.4f}" 

653 self._pbar.set_postfix(postfix) 

654 self._pbar.update(1) 

655 

656 def on_train_end(self, state: "TrainerState", **kwargs) -> None: 

657 if self._pbar is not None: 

658 self._pbar.close() 

659 self._pbar = None 

660 

661 

662class MoEMonitorCallback(Callback): 

663 """Mixture-of-Experts load-balancing monitor. 

664 

665 Delegates to :class:`~hyper_parallel.core.moe_utils.MoEMonitorCallback` 

666 for expert bias updates and aux_loss aggregation. Exposes 

667 ``last_mean_aux_loss`` so that :class:`LoggingCallback` can include it 

668 in the main training loss log line. 

669 

670 Config: ``cfg.train.moe_monitor.*`` (see :class:`MoEMonitorConfig`). 

671 """ 

672 

673 def __init__(self, trainer: "BaseTrainer") -> None: 

674 """Initialize MoEMonitorCallback from trainer config.""" 

675 super().__init__(trainer) 

676 moe_cfg = getattr(trainer.args, 'moe_monitor', None) 

677 self.enabled = getattr(moe_cfg, 'enabled', False) if moe_cfg else False 

678 self._impl = None 

679 

680 if self.enabled: 

681 from hyper_parallel.core.moe_utils import ( # pylint: disable=C0415 

682 MoEMonitorCallback as _CoreMoEMonitorCallback, 

683 ) 

684 from hyper_parallel.core.fully_shard.hsdp_utils import ( # pylint: disable=C0415 

685 GroupInfo, 

686 ) 

687 lr = getattr(moe_cfg, 'lr', 1e-3) 

688 num_recomputations = getattr(moe_cfg, 'num_recomputations', 1) 

689 

690 # Resolve DP/TP/CP groups from trainer's device mesh. 

691 dp_group = getattr(self.trainer, '_dp_group_info', None) 

692 tp_group = None 

693 cp_group = None 

694 mesh = getattr(self.trainer, 'mesh', None) 

695 if mesh is not None: 

696 for name, attr_name in [("tp", "tp_group"), ("cp", "cp_group")]: 

697 try: 

698 raw_group = mesh.get_group(name) 

699 group_info = GroupInfo( 

700 group_name=name, group=raw_group, 

701 rank_size=raw_group.size(), 

702 ) 

703 if attr_name == "tp_group": 

704 tp_group = group_info 

705 else: 

706 cp_group = group_info 

707 except (KeyError, ValueError, AttributeError): 

708 pass 

709 

710 self._impl = _CoreMoEMonitorCallback( 

711 model=self.trainer.model, 

712 lr=lr, 

713 dp_group=dp_group, 

714 tp_group=tp_group, 

715 cp_group=cp_group, 

716 num_recomputations=num_recomputations, 

717 ) 

718 

719 @property 

720 def last_mean_aux_loss(self) -> Optional[float]: 

721 """Mean aux_loss across MoE layers from the last ``on_step_end``.""" 

722 if self._impl is not None: 

723 return self._impl.last_mean_aux_loss 

724 return None 

725 

726 def on_train_begin(self, state: "TrainerState", **kwargs) -> None: 

727 """Log one-time confirmation when MoE monitoring is enabled.""" 

728 if self.enabled and platform.get_rank() == 0: 

729 logger.info("MoEMonitorCallback: MoE expert-load monitoring enabled") 

730 

731 def on_step_end(self, state: "TrainerState", *, loss: float = None, 

732 grad_norm: float = None, **kwargs) -> None: 

733 """Delegate expert bias update to core MoEMonitorCallback.""" 

734 if self._impl is not None: 

735 self._impl.on_step_end() 

736 

737 def on_substep_end(self, state: "TrainerState", **kwargs) -> None: 

738 """No-op; expert bias updates happen in on_step_end.""" 

739 

740 

741class GradientHealthCallback(Callback): 

742 """Detect NaN / Inf grad_norm and raise / warn. 

743 

744 Hooks ``on_pre_optimizer_step`` — which fires after ``clip_grad_norm_`` 

745 and before ``optimizer.step()``. ``grad_norm`` at that point is a plain 

746 scalar produced by hyper's DTensor-aware clipper. If it's not finite, the 

747 optimizer.step() would silently corrupt weights with NaN; we want to 

748 surface it immediately. 

749 

750 Config: ``cfg.train.debug.check_nan_inf``. 

751 """ 

752 

753 def __init__(self, trainer: "BaseTrainer") -> None: 

754 super().__init__(trainer) 

755 debug_cfg = getattr(trainer.args, 'debug', None) 

756 self.enabled = ( 

757 getattr(debug_cfg, 'check_nan_inf', False) if debug_cfg else False 

758 ) 

759 

760 def on_pre_optimizer_step(self, state: "TrainerState", *, 

761 grad_norm: Optional[float] = None, 

762 **kwargs) -> None: 

763 if not self.enabled or grad_norm is None: 

764 return 

765 if math.isnan(grad_norm) or math.isinf(grad_norm): 

766 # Always log on every rank — divergence may be rank-local. 

767 logger.error( 

768 "GradientHealthCallback: grad_norm=%s at step %d " 

769 "(NaN/Inf). Optimizer.step would corrupt weights.", 

770 grad_norm, state.global_step, 

771 ) 

772 # Raise on rank 0 only; other ranks will be torn down by NCCL. 

773 if platform.get_rank() == 0: 

774 raise RuntimeError( 

775 f"Non-finite grad_norm={grad_norm} at " 

776 f"step {state.global_step}. " 

777 "Disable cfg.train.debug.check_nan_inf to skip this guard." 

778 ) 

779 

780 

781class GCCallback(Callback): 

782 """Explicit garbage-collection scheduler. 

783 

784 Python's cyclic GC can stall large training jobs when it decides to run; 

785 forcing a collection every N steps — outside the compute hot path — 

786 keeps pauses predictable.). 

787 

788 Config: ``cfg.train.debug.gc_steps`` (``0`` disables). 

789 """ 

790 

791 def __init__(self, trainer: "BaseTrainer") -> None: 

792 super().__init__(trainer) 

793 debug_cfg = getattr(trainer.args, 'debug', None) 

794 self.gc_steps = ( 

795 getattr(debug_cfg, 'gc_steps', 0) if debug_cfg else 0 

796 ) 

797 if self.gc_steps > 0: 

798 # Disable the automatic generational collector; we'll drive it. 

799 gc.disable() 

800 logger.info("GCCallback: Python gc.collect every %d steps " 

801 "(auto GC disabled)", self.gc_steps) 

802 

803 def on_step_end(self, state: "TrainerState", *, 

804 loss: Optional[float] = None, 

805 grad_norm: Optional[float] = None, **kwargs) -> None: 

806 if self.gc_steps <= 0: 

807 return 

808 if state.global_step % self.gc_steps != 0: 

809 return 

810 gc.collect() 

811 

812 

813class TensorBoardCallback(Callback): 

814 """TensorBoard scalar writer — STUB (not verified). 

815 

816 Hook reserved for ``torch.utils.tensorboard.SummaryWriter`` integration. 

817 Not yet verified; if you enable ``args.tensorboard.enabled`` we emit 

818 a one-time warning so missing TB scalars are visible. To implement: 

819 open SummaryWriter in ``on_train_begin``, write scalars in ``on_log``, 

820 close in ``on_train_end``. 

821 """ 

822 

823 def __init__(self, trainer: "BaseTrainer") -> None: 

824 super().__init__(trainer) 

825 tb_cfg = getattr(trainer.args, 'tensorboard', None)