Add more log to measure time

Author: Charles Li

src/maxtext/trainers/pre_train/train.py

@@ -23,6 +23,7 @@
 import datetime
 import functools
 import os
+import warnings
 
 from absl import app
 
@@ -525,6 +526,7 @@ def train_loop(config, recorder, state=None):
   params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
 
   with jax.set_mesh(mesh), mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    start_time = datetime.datetime.now()
     p_train_step, p_eval_step = train_utils.jit_train_and_eval_step(
         config,
         model,
@@ -536,6 +538,9 @@ def train_loop(config, recorder, state=None):
         eval_data_iterator,
         params_shardings,
     )
+    seconds_elapsed = (datetime.datetime.now() - start_time).total_seconds()
+    print(f"train jit {seconds_elapsed} secs")
+    warnings.warn(f"DEUBG: train jit cost {seconds_elapsed} secs")
     shaped_batch = maxtext_utils.get_shaped_batch(config)
     if config.shard_optimizer_over_data:
       state = sharding.maybe_shard_with_name(state, state_mesh_shardings, config.shard_mode)

src/maxtext/trainers/pre_train/train_compile.py

@@ -21,9 +21,11 @@
 as you would on the target hardware.
 """
 
+import datetime
 import functools
 import os
 import pickle
+import warnings
 from typing import Sequence
 
 from absl import app
@@ -163,18 +165,32 @@ def jit_and_compile(
   # Use both jax.set_mesh (new API) and `with mesh:` (old API) so that drjax,
   # which reads from pxla.thread_resources.env.physical_mesh, can find the mesh.
   with jax.set_mesh(mesh), mesh, logical_axis_rules:
+    start_time = datetime.datetime.now()
     jitted = jax.jit(
         func,
         in_shardings=in_shardings,
         out_shardings=out_shardings,
         static_argnums=static_argnums,
         donate_argnums=donate_argnums,
     )
+    seconds_elapsed = (datetime.datetime.now() - start_time).total_seconds()
+    print(f"train_compile jit {seconds_elapsed} secs")
+    warnings.warn(f"DEUBG: train_compile jit cost {seconds_elapsed} secs")
     maxtext_utils.maybe_dump_jaxpr(config, jitted, func_input_args)
+
+    start_time = datetime.datetime.now()
     lowered = jitted.lower(*func_input_args, **func_input_kwargs)
+    seconds_elapsed = (datetime.datetime.now() - start_time).total_seconds()
+    print(f"train_compile lower {seconds_elapsed} secs")
+    warnings.warn(f"DEUBG: train_compile lower cost {seconds_elapsed} secs")
   # Import libtpu flags as compiler options. Defaults to empty dict if string is empty.
   compiler_options = max_utils.parse_libtpu_flags_to_dict(config.compile_xla_flags)
+
+  start_time = datetime.datetime.now()
   compiled = lowered.compile(compiler_options=compiler_options)
+  seconds_elapsed = (datetime.datetime.now() - start_time).total_seconds()
+  print(f"train_compile compile {seconds_elapsed} secs")
+  warnings.warn(f"DEUBG: train_compile compile cost {seconds_elapsed} secs")
   return compiled
 
 

src/maxtext/utils/train_utils.py

@@ -15,7 +15,9 @@
 # pylint: disable=bare-except, consider-using-generator
 """Utils that are only interesting for training in MaxText."""
 
+import datetime
 import os
+import warnings
 from functools import partial
 
 import jax
@@ -176,9 +178,13 @@ def jit_train_and_eval_step(
     train_step_partial = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings)
     train_step = diloco.build_diloco_train_step(config, train_step_partial, mesh=mesh)
   data_sharding = sharding.get_input_data_sharding(config, mesh)
+  start_time = datetime.datetime.now()
   p_train_step = jit_train_step(
       config, model, state, state_mesh_shardings, data_sharding, train_step, params_shardings, mesh=mesh
   )
+  seconds_elapsed = (datetime.datetime.now() - start_time).total_seconds()
+  print(f"train_util train_step jit {seconds_elapsed} secs")
+  warnings.warn(f"DEUBG: train_util train_step jit cost {seconds_elapsed} secs")
   p_eval_step = None
   if eval_data_iterator:
     p_eval_step = jit_eval_step(config, model, state_mesh_shardings, data_sharding, eval_step)