Tonic commited on
Commit
34b58cd
·
verified ·
1 Parent(s): 5743cc8

fixes on init end bug

Browse files
Files changed (1) hide show
  1. monitoring.py +14 -0
monitoring.py CHANGED
@@ -261,6 +261,10 @@ class SmolLM3Monitor:
261
  def __init__(self, monitor):
262
  self.monitor = monitor
263
 
 
 
 
 
264
  def on_log(self, args, state, control, logs=None, **kwargs):
265
  """Called when logs are created"""
266
  if logs:
@@ -277,6 +281,16 @@ class SmolLM3Monitor:
277
  """Called when evaluation is performed"""
278
  if metrics:
279
  self.monitor.log_evaluation_results(metrics, state.global_step)
 
 
 
 
 
 
 
 
 
 
280
 
281
  return TrackioCallback(self)
282
 
 
261
  def __init__(self, monitor):
262
  self.monitor = monitor
263
 
264
+ def on_init_end(self, args, state, control, **kwargs):
265
+ """Called when training initialization is complete"""
266
+ logger.info("Training initialization completed")
267
+
268
  def on_log(self, args, state, control, logs=None, **kwargs):
269
  """Called when logs are created"""
270
  if logs:
 
281
  """Called when evaluation is performed"""
282
  if metrics:
283
  self.monitor.log_evaluation_results(metrics, state.global_step)
284
+
285
+ def on_train_begin(self, args, state, control, **kwargs):
286
+ """Called when training begins"""
287
+ logger.info("Training started")
288
+
289
+ def on_train_end(self, args, state, control, **kwargs):
290
+ """Called when training ends"""
291
+ logger.info("Training completed")
292
+ if self.monitor:
293
+ self.monitor.close()
294
 
295
  return TrackioCallback(self)
296