Logger fix

PiperOrigin-RevId: 869952160

Author: priyakasimbeg

hessian/model_debugger.py

@@ -288,6 +288,7 @@ def __init__(self,
                use_pmap=True,
                save_every=1,
                metrics_logger=None,
+               pytree_metrics_logger=None,
                skip_flags=None,
                skip_groups=None):
     """Used to inspect a models forward and backward pass.
@@ -305,6 +306,7 @@ def __init__(self,
         save_every == 0
       metrics_logger: utils.MetricsLogger object. If provided then all
         calculations will be saved to disk.
+      pytree_metrics_logger: utils.PytreeMetricLogger object.
       skip_flags: A list of strings of modules to selectively turn off when
         doing the skip analysis on the backward pass.
       skip_groups: A list of registered functions (defined in partition_tree.py)
@@ -317,6 +319,7 @@ def __init__(self,
                        ' path must be specified when building metrics_logger')
     self._save_every = save_every
     self._metrics_logger = metrics_logger
+    self._pytree_metrics_logger = pytree_metrics_logger
     self._use_pmap = use_pmap
     self.forward_pass = None
     self.grad_fn = None
@@ -340,8 +343,11 @@ def __init__(self,
     self._stored_metrics = {}
 
     # In the case of preemption we want to restore prior metrics.
-    if metrics_logger and metrics_logger.latest_pytree_checkpoint_step():
-      self._stored_metrics = metrics_logger.load_latest_pytree()
+    if (
+        pytree_metrics_logger
+        and pytree_metrics_logger.latest_pytree_checkpoint_step()
+    ):
+      self._stored_metrics = pytree_metrics_logger.load_latest_pytree()
 
   def _grab_statistics(self,
                        step,
@@ -375,9 +381,9 @@ def _maybe_save_metrics(self, step):
     save_dict = self._stored_metrics.copy()
     if self._save_every:
       if step % self._save_every == 0:
-        self._metrics_logger.write_pytree(save_dict, step=step)
+        self._pytree_metrics_logger.write_pytree(save_dict, step=step)
     else:
-      self._metrics_logger.write_pytree(save_dict, step=step)
+      self._pytree_metrics_logger.write_pytree(save_dict, step=step)
 
   @property
   def stored_metrics(self):

hessian/model_debugger_callback.py

@@ -99,7 +99,7 @@ def __init__(
     self.dataset = dataset
     checkpoint_dir = os.path.join(train_dir, 'checkpoints')
     pytree_path = os.path.join(checkpoint_dir, 'debugger')
-    logger = utils.MetricLogger(pytree_path=pytree_path)
+    logger = utils.PytreeMetricLogger(pytree_path=pytree_path)
 
     get_act_stats_fn = model_debugger.create_forward_pass_stats_fn(
         model.apply_on_batch,

hessian/test_model_debugger.py

@@ -286,16 +286,20 @@ def test_model_debugger_restore(self):
 
     rep_variables = set_up_cnn()
     pytree_path = os.path.join(self.test_dir, 'metrics')
-    metrics_logger = utils.MetricLogger(
-        pytree_path=pytree_path, events_dir=self.test_dir)
+    pytree_metrics_logger = utils.PytreeMetricLogger(
+        pytree_path=pytree_path)
+    metrics_logger = utils.MetricLogger(self.test_dir)
 
     # Fake grad_fn for testing.
     def grad_fn(params, batch, rng):
       del params, batch, rng
       return rep_variables['params']
 
     debugger = model_debugger.ModelDebugger(
-        use_pmap=False, grad_fn=grad_fn, metrics_logger=metrics_logger)
+        use_pmap=False,
+        grad_fn=grad_fn,
+        metrics_logger=metrics_logger,
+        pytree_metrics_logger=pytree_metrics_logger)
 
     # eval twice to test the concat
     extra_metrics = {'train_loss': 1.0}
@@ -319,8 +323,8 @@ def grad_fn(params, batch, rng):
         'train_loss',
     ]
 
-    metrics_logger.wait_until_pytree_checkpoint_finished()
-    loaded_metrics = metrics_logger.load_latest_pytree(None)
+    pytree_metrics_logger.wait_until_pytree_checkpoint_finished()
+    loaded_metrics = pytree_metrics_logger.load_latest_pytree(None)
 
     self.assertEqual(set(expected_keys), set(metrics.keys()))
     expected_shape = ()
@@ -337,13 +341,15 @@ def grad_fn(params, batch, rng):
 
     # Test restore of prior metrics.
     new_debugger = model_debugger.ModelDebugger(
-        use_pmap=True, metrics_logger=metrics_logger)
+        use_pmap=True,
+        metrics_logger=metrics_logger,
+        pytree_metrics_logger=pytree_metrics_logger)
     _ = new_debugger.full_eval(
         30,
         params=rep_variables['params'],
         grad=rep_variables['params'],
         extra_scalar_metrics=extra_metrics2)
-    metrics_logger.wait_until_pytree_checkpoint_finished()
+    pytree_metrics_logger.wait_until_pytree_checkpoint_finished()
     self.assertEqual(
         new_debugger.stored_metrics['param_norms_sql2']['Conv_0']
         ['kernel'].shape, (3,))

init2winit/model_lib/test_models.py

@@ -1018,7 +1018,7 @@ def test_vit_model_apply_sharding_overrides_failure(self):
         str(error),
         'Param shape (1536,) is not compatible with sharding '
         'NamedSharding(mesh=Mesh(\'devices\': 8, axis_types=(Auto,)), '
-        'spec=PartitionSpec(None, \'devices\'), memory_kind=device)',
+        'spec=P(None, \'devices\'), memory_kind=device)',
     )
 
     good_overrides = {

init2winit/test_utils.py

@@ -17,7 +17,6 @@
 
 """
 
-import os
 import shutil
 import tempfile
 
@@ -84,24 +83,6 @@ def test_run_in_parallel_on_failing_fn(self):
     with self.assertRaisesRegex(ValueError, 'I always fail.'):
       utils.run_in_parallel(_fn_that_always_fails, [dict(arg='hi')], 10)
 
-  def test_append_pytree(self):
-    """Test appending and loading pytrees."""
-    pytrees = [{'a': i} for i in range(10)]
-    pytree_path = os.path.join(self.test_dir, 'pytrees')
-    logger = utils.MetricLogger(pytree_path=pytree_path)
-
-    for i, pytree in enumerate(pytrees):
-      logger.append_pytree(pytree, step=i)
-
-    latest = logger.load_latest_pytree(
-        target=None,
-    )
-    saved_pytrees = latest if latest else []
-
-    self.assertEqual(
-        pytrees, [saved_pytrees[i] for i in range(len(saved_pytrees))]
-    )
-
   def test_array_append(self):
     """Test appending to an array."""
     np.testing.assert_allclose(

init2winit/utils.py

@@ -106,10 +106,8 @@ def wrapper(*args, **kwargs):
 def set_up_loggers(train_dir, xm_work_unit=None):
   """Creates a logger for eval metrics as well as initialization metrics."""
   csv_path = os.path.join(train_dir, 'measurements.csv')
-  pytree_path = os.path.join(train_dir, 'ttl=180d', 'training_metrics')
   metrics_logger = MetricLogger(
       csv_path=csv_path,
-      pytree_path=pytree_path,
       xm_work_unit=xm_work_unit,
       events_dir=train_dir)
 
@@ -171,6 +169,68 @@ def add_log_file(logfile):
   absl_logger.addHandler(handler)
 
 
+class PytreeMetricLogger(object):
+  """Used to log pytree metrics during training.
+  
+  This class is used to log pytree metrics during training. It is similar to
+  MetricLogger, but it is designed to log pytree metrics.
+  """
+
+  def __init__(
+      self,
+      pytree_path,
+  ):
+    self._pytree_path = pytree_path
+
+    orbax_file_options = ocp.checkpoint_manager.FileOptions(
+        path_permission_mode=0o775,
+        cns2_storage_options=ocp.options.Cns2StorageOptions(
+            choose_store_cell=True,
+        ),
+    )
+    self._orbax_checkpoint_manager = ocp.CheckpointManager(
+        self._pytree_path,
+        options=ocp.CheckpointManagerOptions(
+            file_options=orbax_file_options,
+            max_to_keep=1, create=True,
+        ),
+    )
+
+  def write_pytree(self, pytree, step=0):
+    """Record a serializable pytree to disk at the given step.
+
+    Args:
+      pytree: Any serializable pytree
+      step: Integer. The global step.
+    """
+    state = dict(pytree=pytree)
+    checkpoint.save_checkpoint(
+        step,
+        state=state,
+        orbax_checkpoint_manager=self._orbax_checkpoint_manager,
+    )
+
+  def wait_until_pytree_checkpoint_finished(self):
+    self._orbax_checkpoint_manager.wait_until_finished()
+
+  def latest_pytree_checkpoint_step(self):
+    return self._orbax_checkpoint_manager.latest_step()
+
+  def load_latest_pytree(self, target=None):
+    """Load pytree from checkpoint."""
+    if target:
+      target = dict(pytree=target)
+    logging.info('target: %s', target)
+    loaded_target = checkpoint.load_latest_checkpoint(
+        target=target,
+        orbax_checkpoint_manager=self._orbax_checkpoint_manager,
+    )
+    if loaded_target:
+      return loaded_target['pytree']
+    else:
+      return target
+
+
 # TODO(gdahl,gilmer): Use atomic writes to avoid file corruptions due to
 # preemptions.
 class MetricLogger(object):
@@ -183,7 +243,6 @@ class MetricLogger(object):
   def __init__(self,
                csv_path='',
                json_path='',
-               pytree_path='',
                events_dir=None,
                **logger_kwargs):
     """Create a recorder for metrics, as CSV or JSON.
@@ -192,7 +251,6 @@ def __init__(self,
     Args:
       csv_path: A filepath to a CSV file to append to.
       json_path: An optional filepath to a JSON file to append to.
-      pytree_path: Where to save trees of numeric arrays.
       events_dir: Optional. If specified, save tfevents summaries to this
         directory.
       **logger_kwargs: Optional keyword arguments, whose only valid parameter
@@ -202,7 +260,6 @@ def __init__(self,
     self._measurements = {}
     self._csv_path = csv_path
     self._json_path = json_path
-    self._pytree_path = pytree_path
     if logger_kwargs:
       if len(logger_kwargs.keys()) > 1 or 'xm_work_unit' not in logger_kwargs:
         raise ValueError(
@@ -216,33 +273,6 @@ def __init__(self,
     if events_dir:
       self._tb_metric_writer = metric_writers.create_default_writer(events_dir)
 
-    orbax_file_options = ocp.checkpoint_manager.FileOptions(
-        path_permission_mode=0o775,
-        cns2_storage_options=ocp.options.Cns2StorageOptions(
-            choose_store_cell=True,
-        ),
-    )
-    self._orbax_checkpoint_manager = ocp.CheckpointManager(
-        self._pytree_path,
-        options=ocp.CheckpointManagerOptions(
-            max_to_keep=1, create=True, file_options=orbax_file_options
-        ),
-    )
-
-  def load_latest_pytree(self, target=None):
-    """Load pytree from checkpoint."""
-    if target:
-      target = dict(pytree=target)
-    logging.info('target: %s', target)
-    loaded_target = checkpoint.load_latest_checkpoint(
-        target=target,
-        orbax_checkpoint_manager=self._orbax_checkpoint_manager,
-    )
-    if loaded_target:
-      return loaded_target['pytree']
-    else:
-      return target
-
   def append_scalar_metrics(self, metrics):
     """Record a dictionary of scalar metrics at a given step.
 
@@ -289,54 +319,6 @@ def append_scalar_metrics(self, metrics):
       # size 512. We could only flush at the end of training to optimize this.
       self._tb_metric_writer.flush()
 
-  def write_pytree(self, pytree, step=0):
-    """Record a serializable pytree to disk at the given step.
-
-    Args:
-      pytree: Any serializable pytree
-      step: Integer. The global step.
-    """
-    state = dict(pytree=pytree)
-    checkpoint.save_checkpoint(
-        step,
-        state=state,
-        orbax_checkpoint_manager=self._orbax_checkpoint_manager,
-    )
-
-  def wait_until_pytree_checkpoint_finished(self):
-    self._orbax_checkpoint_manager.wait_until_finished()
-
-  def latest_pytree_checkpoint_step(self):
-    return self._orbax_checkpoint_manager.latest_step()
-
-  def append_pytree(self, pytree, step=0):
-    """Append and record a serializable pytree to disk.
-
-    The pytree will be saved to disk as a list of pytree objects. Everytime
-    this function is called, it will load the previous saved state, append the
-    next pytree to the list, then save the appended list.
-
-    Args:
-      pytree: Any serializable pytree.
-      step: Integer. The global step.
-    """
-    # Read the latest (and only) checkpoint if it exists, then append the new
-    # state to it before saving back to disk.
-    try:
-      old_state = self.load_latest_pytree(target=None)
-    except ValueError:
-      old_state = None
-    # Because we pass target=None, checkpointing will return the raw state
-    # dict, where 'pytree' is a dict with keys ['0', '1', ...] instead of a
-    # list.
-    if old_state:
-      state_list = [old_state[i] for i in range(len(old_state))]
-    else:
-      state_list = []
-    state_list.append(pytree)
-
-    self.write_pytree(state_list, step=step)
-
   def append_json_object(self, json_obj):
     """Append a json serializable object to the json file."""