chore: import upstream snapshot with attribution

2026-07-13 13:34:58 +08:00
commit a203934033
1368 changed files with 175001 additions and 0 deletions
@@ -0,0 +1,43 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import gc
+
+from .base import MegatronCallback
+
+
+class DefaultFlowCallback(MegatronCallback):
+
+    def on_train_begin(self):
+        args = self.args
+        if args.manual_gc:
+            gc.disable()
+            gc.collect()
+
+    def on_step_end(self):
+        args = self.args
+        state = self.state
+
+        state.consumed_train_samples += args.global_batch_size
+
+        if state.iteration == 1 or state.iteration % args.logging_steps == 0:
+            state.should_log = True
+        if args.eval_steps and state.iteration % args.eval_steps == 0 and args.eval_iters > 0:
+            state.should_eval = True
+        if args.save_steps and state.iteration % args.save_steps == 0:
+            state.should_save = True
+
+        if state.iteration >= args.train_iters:
+            if args.eval_iters > 0:
+                state.should_eval = True
+            state.should_save = True
+        if args.manual_gc and args.manual_gc_steps != 0 and state.iteration % args.manual_gc_steps == 0:
+            gc.collect()
+
+    def on_eval_begin(self):
+        args = self.args
+        if args.manual_gc and args.manual_gc_eval:
+            gc.collect()
+
+    def on_eval_end(self):
+        args = self.args
+        if args.manual_gc and args.manual_gc_eval:
+            gc.collect(generation=0)