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

1# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates 

2# Copyright 2026 Huawei Technologies Co., Ltd 

3# 

4# Licensed under the Apache License, Version 2.0 (the "License"); 

5# you may not use this file except in compliance with the License. 

6# You may obtain a copy of the License at 

7# 

8# http://www.apache.org/licenses/LICENSE-2.0 

9# 

10# Unless required by applicable law or agreed to in writing, software 

11# distributed under the License is distributed on an "AS IS" BASIS, 

12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

13# See the License for the specific language governing permissions and 

14# limitations under the License. 

15# ============================================================================ 

16"""Trainer callback lifecycle contracts.""" 

17 

18 

19from typing import TYPE_CHECKING, Any, Dict, List 

20 

21from hyper_parallel.trainer.state import TrainerState 

22 

23 

24if TYPE_CHECKING: 

25 from hyper_parallel.trainer.base import BaseTrainer 

26 

27 

28class Callback: 

29 """Base class defining the trainer callback lifecycle hooks.""" 

30 

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

32 """Bind the callback to its owning trainer.""" 

33 self.trainer = trainer 

34 self.mesh = trainer.mesh 

35 

36 def on_step_begin(self, state: TrainerState, micro_batches: List[Dict[str, Any]] = None, **kwargs: Any) -> None: 

37 """Hook invoked at the start of each training step.""" 

38 

39 def on_step_end( 

40 self, state: TrainerState, loss: float, loss_dict: Dict[str, float], grad_norm: float, **kwargs: Any 

41 ) -> None: 

42 """Hook invoked at the end of each training step.""" 

43 

44 def on_micro_step_begin(self, state: TrainerState, micro_batch: Dict[str, Any], **kwargs: Any) -> None: 

45 """Hook invoked at the start of each gradient-accumulation micro step.""" 

46 

47 def on_micro_step_end(self, state: TrainerState, loss_dict: Dict[str, Any], **kwargs: Any) -> None: 

48 """Hook invoked after a successful gradient-accumulation micro step. 

49 

50 Args: 

51 state: Current trainer state. 

52 loss_dict: Named losses and metrics produced by this micro step. 

53 **kwargs: Additional callback context. 

54 """ 

55 

56 def on_epoch_begin(self, state: TrainerState, **kwargs: Any) -> None: 

57 """Hook invoked at the start of each epoch.""" 

58 

59 def on_epoch_end(self, state: TrainerState, **kwargs: Any) -> None: 

60 """Hook invoked at the end of each epoch.""" 

61 

62 def on_train_begin(self, state: TrainerState, **kwargs: Any) -> None: 

63 """Hook invoked at the start of training.""" 

64 

65 def on_train_end(self, state: TrainerState, **kwargs: Any) -> None: 

66 """Hook invoked at the end of training."""