From 2cddfd5e02f99ce164d422e86d5f29d5e4e6f094 Mon Sep 17 00:00:00 2001 From: David Hall Date: Mon, 12 Feb 2024 23:05:53 -0800 Subject: [PATCH] sigh --- src/levanter/trainer.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/levanter/trainer.py b/src/levanter/trainer.py index 8a118f549..9c8f2d2f2 100644 --- a/src/levanter/trainer.py +++ b/src/levanter/trainer.py @@ -358,7 +358,6 @@ def init_state(partial_state, model_init, training_key, is_trainable): return eqx.combine(partial_state, fresh_state) state = init_state(state, model_init, training_key, self.is_trainable_param) - print(state.step.sharding) return state @@ -385,8 +384,6 @@ def training_steps( with capture_time() as loading_time: example = next(iter_data) - levanter.tracker.log_metrics({"throughput/loading_time": loading_time()}, step=state.step) - info = self.train_step(state, example) state = info.state @@ -394,7 +391,9 @@ def training_steps( with capture_time() as hook_time: self.run_hooks(info) - levanter.tracker.log_metrics({"throughput/hook_time": hook_time()}, step=state.step) + levanter.tracker.log_metrics({"throughput/hook_time": hook_time()}, step=info.step) + + levanter.tracker.log_metrics({"throughput/loading_time": loading_time()}, step=info.step) yield info