disable step time logging

AI-Hypercomputer · Aug 20, 2024 · f9ccbc7 · f9ccbc7
1 parent 8cbea66
commit f9ccbc7
Showing 1 changed file with 10 additions and 10 deletions.
diff --git a/MaxText/train.py b/MaxText/train.py
@@ -148,13 +148,13 @@ def write_metrics_to_tensorboard(writer, metrics, step, config):
 
     full_log = step % config.log_period == 0
 
-    max_logging.log(
-        f"completed step: {step}, seconds: {metrics['scalar']['perf/step_time_seconds']:.3f}, "
-        f"TFLOP/s/device: {metrics['scalar']['perf/per_device_tflops_per_sec']:.3f}, "
-        f"Tokens/s/device: {metrics['scalar']['perf/per_device_tokens_per_sec']:.3f}, "
-        f"total_weights: {metrics['scalar']['learning/total_weights']}, "
-        f"loss: {metrics['scalar']['learning/loss']:.3f}"
-    )
+    # max_logging.log(
+    #     f"completed step: {step}, seconds: {metrics['scalar']['perf/step_time_seconds']:.3f}, "
+    #     f"TFLOP/s/device: {metrics['scalar']['perf/per_device_tflops_per_sec']:.3f}, "
+    #     f"Tokens/s/device: {metrics['scalar']['perf/per_device_tokens_per_sec']:.3f}, "
+    #     f"total_weights: {metrics['scalar']['learning/total_weights']}, "
+    #     f"loss: {metrics['scalar']['learning/loss']:.3f}"
+    # )
 
     if full_log and jax.process_index() == 0:
       max_logging.log(f"To see full metrics 'tensorboard --logdir={config.tensorboard_dir}'")
@@ -642,9 +642,9 @@ def map_fn(key_path, value):
     new_time = datetime.datetime.now()
     record_scalar_metrics(metrics, new_time - last_step_completion, per_device_tflops, learning_rate_schedule(step), per_device_tokens)
     step_time_delta = new_time - last_step_completion
-    max_logging.log(f"completed step: {step}, seconds: {step_time_delta.total_seconds()}, "
-          f"TFLOP/s/device: {per_device_tflops / step_time_delta.total_seconds()}, "
-          f"loss: {metrics['scalar']['learning/loss']:.3f}")
+    # max_logging.log(f"completed step: {step}, seconds: {step_time_delta.total_seconds()}, "
+    #       f"TFLOP/s/device: {per_device_tflops / step_time_delta.total_seconds()}, "
+    #       f"loss: {metrics['scalar']['learning/loss']:.3f}")
     last_step_completion = new_time
 
     if checkpoint_manager is not None: