From e1586f791980d95aab346aa9783f60567e0ad2ca Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 6 May 2025 11:18:00 -0400 Subject: [PATCH] make sure gc_steps is used for all trainers (#2638) --- src/axolotl/core/trainer_builder.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index 31ee3cccf..99d97800c 100755 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -168,6 +168,9 @@ class TrainerBuilderBase(abc.ABC): ) ) + if self.cfg.gc_steps: + callbacks.append(GCCallback(gc_steps=self.cfg.gc_steps)) + if self.cfg.use_wandb: callbacks.append( SaveAxolotlConfigtoWandBCallback(self.cfg.axolotl_config_path) @@ -249,9 +252,6 @@ class HFCausalTrainerBuilder(TrainerBuilderBase): if self.cfg.loss_watchdog_threshold is not None: callbacks.append(LossWatchDogCallback(self.cfg)) - if self.cfg.gc_steps: - callbacks.append(GCCallback(gc_steps=self.cfg.gc_steps)) - return callbacks def get_post_trainer_create_callbacks(self, trainer):