47 lines
1.0 KiB
Python
47 lines
1.0 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from swift.megatron.trainers import BaseMegatronTrainer
|
|
|
|
|
|
class MegatronCallback:
|
|
|
|
def __init__(self, trainer: 'BaseMegatronTrainer'):
|
|
self.trainer = trainer
|
|
self.args = trainer.args
|
|
self.state = trainer.state
|
|
|
|
def on_train_begin(self):
|
|
pass
|
|
|
|
def on_train_end(self):
|
|
pass
|
|
|
|
def on_step_begin(self):
|
|
pass
|
|
|
|
def on_step_end(self):
|
|
pass
|
|
|
|
def on_log(self, logs):
|
|
pass
|
|
|
|
def on_eval_begin(self):
|
|
pass
|
|
|
|
def on_eval_end(self):
|
|
pass
|
|
|
|
def on_eval_step(self):
|
|
pass
|
|
|
|
def on_save(self, output_dir):
|
|
"""Called after save_checkpoint() returns.
|
|
|
|
Note: When async_save is enabled, the checkpoint may not be fully
|
|
written to disk yet. Use only for non-I/O-dependent logic, or
|
|
ensure async_save is disabled if you need to read the checkpoint.
|
|
"""
|
|
pass
|