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

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 train_cfg = getattr(trainer.args, 'train', None) 

164 log_cfg = getattr(train_cfg, 'logging', None) 

165 if log_cfg is None: 

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

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

168 self.report_global_loss = ( 

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

170 ) 

171 self.report_throughput = ( 

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

173 ) 

174 self.model_flops_per_token = ( 

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

176 ) 

177 self.peak_tflops = ( 

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

179 ) 

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

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

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

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

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

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

186 self._step_start_time = 0.0 

187 

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

189 self._step_start_time = time.time() 

190 

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

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

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

194 return 

195 

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

197 lr = 0.0 

198 if self.trainer.lr_scheduler is not None: 

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

200 

201 metrics = { 

202 "step": state.global_step, 

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

204 # in the log for sanity comparisons across runs. 

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

206 "grad_norm": ( 

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

208 ), 

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

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

211 } 

212 

213 tokens_per_sec = None 

214 if self.report_throughput: 

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

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

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

218 if tokens is None: 

219 tokens = self._tokens_per_step_est 

220 tokens_per_sec = tokens / elapsed 

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

222 

223 if self.model_flops_per_token and self.peak_tflops: 

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

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

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

227 observed_tflops = ( 

228 tokens_per_sec * self.model_flops_per_token / 1e12 

229 ) 

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

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

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

233 

234 # Include aux_loss from MoEMonitorCallback when available. 

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

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

237 if aux_loss is not None: 

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

239 

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

241 

242 record = { 

243 "step": state.global_step, 

244 "loss": loss, 

245 "grad_norm": grad_norm, 

246 "lr": lr, 

247 "step_time": elapsed, 

248 "tokens_per_sec": tokens_per_sec, 

249 "aux_loss": aux_loss, 

250 } 

251 state.log_history.append(record) 

252 

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

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

255 if dispatch is not None: 

256 dispatch(record) 

257 

258 

259class CheckpointCallback(Callback): 

260 """Save distributed checkpoints and handle resume. 

261 

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

263 """ 

264 

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

266 super().__init__(trainer) 

267 train_cfg = getattr(trainer.args, 'train', None) 

268 ckpt_cfg = getattr(train_cfg, 'checkpoint', None) 

269 if ckpt_cfg is None: 

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

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

272 self.output_dir = ( 

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

274 ) 

275 self.load_path = ( 

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

277 ) 

278 self.save_async = ( 

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

280 ) 

281 self._last_saved_step = -1 

282 self._save_thread = None # async save worker 

283 

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

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

286 

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

288 """ 

289 if not self.load_path: 

290 return 

291 try: 

292 # pylint: disable=C0415 

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

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

295 

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

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

298 return 

299 

300 # 1. Restore model via hyper DCP 

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

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

303 self.trainer.model.load_state_dict(model_sd) 

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

305 

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

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

308 if os.path.isfile(extra_path): 

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

310 extra = json.load(f) 

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

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

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

314 state.global_step, state.epoch) 

315 

316 # 3. Restore optimizer 

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

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

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

320 self.trainer.optimizer.load_state_dict(optim_sd) 

321 logger.info("Optimizer restored") 

322 

323 # 4. Restore LR scheduler 

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

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

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

327 self.trainer.lr_scheduler.load_state_dict(sched_sd) 

328 logger.info("LR scheduler restored") 

329 

330 # 5. Restore RNG state 

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

332 if os.path.isfile(rng_path): 

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

334 platform.set_rng_state(rng_state) 

335 logger.info("RNG state restored") 

336 

337 # 6. Restore dataloader position (StatefulDataLoader) 

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

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

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

341 self.trainer.train_dataloader.load_state_dict(dl_state) 

342 logger.info("Dataloader state restored") 

343 

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

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

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

347 if dispatch is not None: 

348 dispatch(self.load_path) 

349 

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

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

352 

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

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

355 if self.save_steps <= 0: 

356 return 

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

358 return 

359 if state.global_step == self._last_saved_step: 

360 return 

361 self._dispatch_save(state) 

362 

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

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

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

366 # the same directory / state-dict iterator. 

367 self._join_pending() 

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

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

370 self._save(state) 

371 

372 # --- async plumbing ------------------------------------------------- 

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

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

375 if not self.save_async: 

376 self._save(state) 

377 return 

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

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

380 self._join_pending() 

381 # pylint: disable=C0415 

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

383 snap_step = state.global_step 

384 snap_epoch = state.epoch 

385 state_snapshot = copy.copy(state) 

386 state_snapshot.global_step = snap_step 

387 state_snapshot.epoch = snap_epoch 

388 self._save_thread = threading.Thread( 

389 target=self._save, 

390 args=(state_snapshot,), 

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

392 daemon=True, 

393 ) 

394 self._save_thread.start() 

395 logger.info_rank0( 

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

397 snap_step, self._save_thread.name, 

398 ) 

399 

400 def _join_pending(self) -> None: 

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

402 t = self._save_thread 

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

404 logger.info_rank0( 

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

406 ) 

407 t.join() 

408 self._save_thread = None 

409 

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

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

412 

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

414 """ 

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

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

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

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

419 os.makedirs(save_dir, exist_ok=True) 

420 rank = platform.get_rank() 

421 

422 try: 

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

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

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

426 

427 # 2. Optimizer — per-rank 

428 if self.trainer.optimizer: 

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

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

431 

432 # 3. LR scheduler 

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

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

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

436 

437 # 4. Extra state: global_step, epoch 

438 if rank == 0: 

439 extra = { 

440 "global_step": state.global_step, 

441 "epoch": state.epoch, 

442 } 

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

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

445 json.dump(extra, f) 

446 

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

448 rng_state = platform.get_rng_state() 

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

450 torch.save(rng_state, rng_path) 

451 

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

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

454 self.trainer.train_dataloader, 'state_dict' 

455 ): 

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

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

458 

459 self._last_saved_step = state.global_step 

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

461 

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

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

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

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

466 if dispatch is not None: 

467 dispatch(save_dir) 

468 

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

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

471 

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

473 

474 

475class SafetensorsExportCallback(Callback): 

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

477 

478 Separated from CheckpointCallback per RFC Section 5.2. 

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

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

481 

482 """ 

483 

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

485 super().__init__(trainer) 

486 train_cfg = getattr(trainer.args, 'train', None) 

487 ckpt_cfg = getattr(train_cfg, 'checkpoint', None) 

488 if ckpt_cfg is None: 

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

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

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

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

493 self._last_saved_step = -1 

494 

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

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

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

498 return 

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

500 return 

501 if state.global_step == self._last_saved_step: 

502 return 

503 self._export(state) 

504 

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

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

507 self._export(state) 

508 

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

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

511 

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

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

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

515 this generic callback. Falls back to the legacy 

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

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

518 """ 

519 # pylint: disable=C0415 

520 

521 rank = platform.get_rank() 

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

523 

524 try: 

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

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

527 # pylint: disable=C0415 

528 from torch.distributed.checkpoint.state_dict import StateDictOptions 

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

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

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

532 

533 if rank == 0: 

534 os.makedirs(save_dir, exist_ok=True) 

535 

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

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

538 # the generic offline-transform path. 

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

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

541 save_fn = ( 

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

543 if adapter_cls is not None else None 

544 ) 

545 if save_fn is not None: 

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

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

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

549 logger.info( 

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

551 adapter_cls.__name__, save_dir, 

552 ) 

553 else: 

554 save_state_dict_as_huggingface_format(full_sd, save_dir) 

555 logger.info( 

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

557 ) 

558 

559 self._last_saved_step = state.global_step 

560 

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

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

563 

564 

565class EvalCallback(Callback): 

566 """Evaluation callback stub. 

567 

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

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

570 training logs rather than silently skipped. 

571 """ 

572 

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

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

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

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

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

578 if platform.get_rank() == 0: 

579 logger.warning( 

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

581 ) 

582 

583 

584class ProfilerCallback(Callback): 

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

586 

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

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

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

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

591 in ``on_train_begin`` / ``on_step_end`` / ``on_train_end``. 

592 """ 

593 

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

595 super().__init__(trainer) 

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

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

598 logger.warning( 

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

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

601 "relying on traces." 

602 ) 

603 

604 

605class WandbCallback(Callback): 

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

607 

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

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

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

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

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

613 """ 

614 

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

616 super().__init__(trainer) 

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

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

619 logger.warning( 

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

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

622 "W&B dashboards." 

623 ) 

624 

625 

626class ProgressCallback(Callback): 

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

628 

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

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

631 """ 

632 

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

634 super().__init__(trainer) 

635 self._pbar = None 

636 

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

638 if platform.get_rank() != 0: 

639 return 

640 try: 

641 # pylint: disable=C0415 

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

643 self._pbar = tqdm( 

644 total=state.max_steps, 

645 initial=state.global_step, 

646 desc="Training", 

647 unit="step", 

648 dynamic_ncols=True, 

649 ) 

650 except ImportError: 

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

652 

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

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

655 if self._pbar is None: 

656 return 

657 postfix = {} 

658 if loss is not None: 

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

660 if grad_norm is not None: 

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

662 self._pbar.set_postfix(postfix) 

663 self._pbar.update(1) 

664 

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

666 if self._pbar is not None: 

667 self._pbar.close() 

668 self._pbar = None 

669 

670 

671class MoEMonitorCallback(Callback): 

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

673 

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

675 for expert bias updates and aux_loss aggregation. Exposes 

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

677 in the main training loss log line. 

678 

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

680 """ 

681 

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

683 """Initialize MoEMonitorCallback from trainer config.""" 

684 super().__init__(trainer) 

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

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

687 self._impl = None 

688 

689 if self.enabled: 

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

691 MoEMonitorCallback as _CoreMoEMonitorCallback, 

692 ) 

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

694 GroupInfo, 

695 ) 

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

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

698 

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

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

701 tp_group = None 

702 cp_group = None 

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

704 if mesh is not None: 

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

706 try: 

707 raw_group = mesh.get_group(name) 

708 group_info = GroupInfo( 

709 group_name=name, group=raw_group, 

710 rank_size=raw_group.size(), 

711 ) 

712 if attr_name == "tp_group": 

713 tp_group = group_info 

714 else: 

715 cp_group = group_info 

716 except (KeyError, ValueError, AttributeError): 

717 pass 

718 

719 self._impl = _CoreMoEMonitorCallback( 

720 model=self.trainer.model, 

721 lr=lr, 

722 dp_group=dp_group, 

723 tp_group=tp_group, 

724 cp_group=cp_group, 

725 num_recomputations=num_recomputations, 

726 ) 

727 

728 @property 

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

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

731 if self._impl is not None: 

732 return self._impl.last_mean_aux_loss 

733 return None 

734 

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

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

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

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

739 

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

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

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

743 if self._impl is not None: 

744 self._impl.on_step_end() 

745 

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

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

748 

749 

750class GradientHealthCallback(Callback): 

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

752 

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

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

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

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

757 surface it immediately. 

758 

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

760 """ 

761 

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

763 super().__init__(trainer) 

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

765 self.enabled = ( 

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

767 ) 

768 

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

770 grad_norm: Optional[float] = None, 

771 **kwargs) -> None: 

772 if not self.enabled or grad_norm is None: 

773 return 

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

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

776 logger.error( 

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

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

779 grad_norm, state.global_step, 

780 ) 

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

782 if platform.get_rank() == 0: 

783 raise RuntimeError( 

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

785 f"step {state.global_step}. " 

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

787 ) 

788 

789 

790class GCCallback(Callback): 

791 """Explicit garbage-collection scheduler. 

792 

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

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

795 keeps pauses predictable.). 

796 

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

798 """ 

799 

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

801 super().__init__(trainer) 

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

803 self.gc_steps = ( 

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

805 ) 

806 if self.gc_steps > 0: 

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

808 gc.disable() 

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

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

811 

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

813 loss: Optional[float] = None, 

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

815 if self.gc_steps <= 0: 

816 return 

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

818 return 

819 gc.collect()