add option to add the global grad norm for sub-components of the network to summaries

GitOrigin-RevId: 02f7cfd

Author: pdufter

axlearn/common/learner.py

@@ -12,10 +12,11 @@
 import dataclasses
 import enum
 from collections.abc import Mapping, Sequence
-from typing import Any, Callable, Optional, cast
+from typing import Any, Callable, List, Optional, Union, cast
 
 import jax
 import optax
+from absl import logging
 from jax import numpy as jnp
 
 from axlearn.common.base_layer import ParameterSpec
@@ -47,6 +48,7 @@
     Nested,
     Tensor,
     flatten_items,
+    get_recursively,
     match_regex_rules,
     prune_empty,
     register_per_param_settings,
@@ -166,9 +168,10 @@ class Config(BaseLearner.Config):
         #
         # See optimizers.param_ema for more details.
         ema: InstantiableConfig = config_for_function(param_ema)
-        # Whether to add per variable gradient and norm summaries. Enable it will make
-        # the training slower since the summary is computed for every step.
-        enable_per_variable_summaries: bool = False
+        # Whether to add per variable gradient and norm summaries. Enable it will make the training
+        # slower since the summary is computed for every step. When a list of strings is provided,
+        # the global norm for each  specified component (prefix-matched) is computed.
+        enable_per_variable_summaries: Union[bool, List[str]] = False
         # Optional decorator for the ForwardFn. An example use would be to enable gradient
         # accumulation by using `with_minibatch_steps` decorator defined below. E.g.
         # learner.forward_fn_transformation = config.config_for_function(with_minibatch_steps).set(
@@ -269,7 +272,7 @@ def _compute_updated_params(
         state_updates: Nested[Tensor],
     ) -> Nested[Tensor]:
         cfg = self.config
-        if cfg.enable_per_variable_summaries:
+        if cfg.enable_per_variable_summaries is True:
             param_rms = jax.tree.map(
                 lambda p: optax.safe_root_mean_squares(p.value, min_rms=1e-3), opt_params
             )
@@ -280,6 +283,20 @@ def _compute_updated_params(
             )
             for p, g_n in flatten_items(grad_rms):
                 self.add_summary(f"grad_rms/{p}", g_n)
+        elif isinstance(cfg.enable_per_variable_summaries, list):
+            # Expecting a list of strings for which to get grad norm summaries,
+            # e.g., ["comp_1", "comp_2/weight"]
+            for variable_prefix in cfg.enable_per_variable_summaries:
+                try:
+                    variable_grad_norm = get_recursively(gradients, variable_prefix)
+                except KeyError:
+                    # fail gracefully to not interrupt training
+                    logging.error("Could not get grad norm for %s", variable_prefix)
+                    variable_grad_norm = {}
+
+                # Using global norm when subcomponents are used.
+                variable_grad_norm = optax.global_norm(variable_grad_norm)
+                self.add_summary(f"grad_norm/{variable_prefix}", variable_grad_norm)
 
         # Set `parameter_updates` to 0 if the param is not updated by the optimizer.
         parameter_updates = jax.tree.map(

axlearn/common/learner_test.py

@@ -586,7 +586,11 @@ def loss_fn(params):
 
         self.assertNestedAllClose(dict(v=expected_new_v), updated_params, atol=1e-6)
 
-    def test_per_variable_summaries(self):
+    @parameterized.parameters(
+        (True,),
+        (["comp_1", "comp_2/weight"],),
+    )
+    def test_per_variable_summaries(self, enable_per_variable_summaries):
         sgd_cfg = config_for_function(sgd_optimizer).set(
             learning_rate=1.0,
             decouple_weight_decay=True,
@@ -596,31 +600,58 @@ def test_per_variable_summaries(self):
             args=(config_for_function(clip_by_global_norm), sgd_cfg),
         )
         cfg = Learner.default_config().set(
-            name="test", optimizer=optimizer_cfg, enable_per_variable_summaries=True
+            name="test",
+            optimizer=optimizer_cfg,
+            enable_per_variable_summaries=enable_per_variable_summaries,
         )
         learner: Learner = cfg.instantiate(parent=None)
         params = dict(
-            weight=OptParam(
-                value=jnp.asarray([0, 2, 2, -3], dtype=jnp.float32),
-                factorization_spec=None,
-                weight_decay_scale=1.0,
+            comp_1=dict(
+                weight=OptParam(
+                    value=jnp.asarray([0, 2, 2, -3], dtype=jnp.float32),
+                    factorization_spec=None,
+                    weight_decay_scale=1.0,
+                ),
+                moving_mean=OptParam(
+                    value=jnp.array([0, -1, 0, 0], dtype=jnp.float32),
+                    factorization_spec=None,
+                    weight_decay_scale=0.0,
+                ),
             ),
-            moving_mean=OptParam(
-                value=jnp.array([0, -1, 0, 0], dtype=jnp.float32),
-                factorization_spec=None,
-                weight_decay_scale=0.0,
+            comp_2=dict(
+                weight=OptParam(
+                    value=jnp.asarray([0, 2, 2, -3], dtype=jnp.float32),
+                    factorization_spec=None,
+                    weight_decay_scale=1.0,
+                ),
+                bias=OptParam(
+                    value=jnp.array([0, -1, 0, 0], dtype=jnp.float32),
+                    factorization_spec=None,
+                    weight_decay_scale=0.0,
+                ),
             ),
         )
         state = learner.init(model_params=params)
 
         def loss_fn(x):
-            return -jax.nn.log_softmax(x["weight"] + x["moving_mean"])[1]
+            return -jax.nn.log_softmax(
+                x["comp_1"]["weight"]
+                + x["comp_1"]["moving_mean"]
+                + x["comp_2"]["weight"]
+                + x["comp_2"]["bias"]
+            )[1]
 
         loss, grads = jax.value_and_grad(loss_fn)(jax.tree.map(lambda p: p.value, params))
-        np.testing.assert_allclose(loss, 1.412078, atol=1e-6, rtol=1e-6)
-        expected_grad = jnp.asarray([0.089629, -0.756364, 0.662272, 0.004462])
+        np.testing.assert_allclose(loss, 2.142971, atol=1e-6, rtol=1e-6)
+        expected_grad = jnp.asarray([0.01587562, -0.8826942, 0.86677927, 0.00003935])
         self.assertNestedAllClose(
-            dict(weight=expected_grad, moving_mean=expected_grad), grads, atol=1e-6, rtol=1e-6
+            dict(
+                comp_1=dict(weight=expected_grad, moving_mean=expected_grad),
+                comp_2=dict(weight=expected_grad, bias=expected_grad),
+            ),
+            grads,
+            atol=1e-6,
+            rtol=1e-6,
         )
         _, output_collection = F(
             learner,
@@ -632,22 +663,44 @@ def loss_fn(x):
                 Updates(
                     delta_updates=grads,
                     opt_params=params,
-                    inplace_updates=dict(moving_mean=params["moving_mean"].value + 1),
+                    inplace_updates=dict(
+                        dict(comp_1=dict(moving_mean=params["comp_1"]["moving_mean"].value + 1))
+                    ),
                 )
             ],
         )
+        expected_summaries = {
+            "optimizer/learning_rate": 1.0,
+            "optimizer/lr_schedule_step": 1,
+            "optimizer/gradient_norm": jnp.sqrt(jnp.sum(4 * expected_grad**2)),
+            "optimizer/schedule_step": 1,
+            "optimizer/schedule_scale": -1.0,
+        }
+
+        if enable_per_variable_summaries is True:
+            expected_summaries = {
+                **expected_summaries,
+                **{
+                    "param_rms/comp_1/weight": jnp.sqrt((0 + 4 + 4 + 9) / 4),
+                    "param_rms/comp_1/moving_mean": 0.5,
+                    "param_rms/comp_2/weight": jnp.sqrt((0 + 4 + 4 + 9) / 4),
+                    "param_rms/comp_2/bias": 0.5,
+                    "grad_rms/comp_1/weight": jnp.sqrt(jnp.mean(expected_grad**2)),
+                    "grad_rms/comp_1/moving_mean": jnp.sqrt(jnp.mean(expected_grad**2)),
+                    "grad_rms/comp_2/weight": jnp.sqrt(jnp.mean(expected_grad**2)),
+                    "grad_rms/comp_2/bias": jnp.sqrt(jnp.mean(expected_grad**2)),
+                },
+            }
+        else:
+            expected_summaries = {
+                **expected_summaries,
+                **{
+                    "grad_norm/comp_1": jnp.sqrt(jnp.sum(2 * expected_grad**2)),
+                    "grad_norm/comp_2/weight": jnp.sqrt(jnp.sum(expected_grad**2)),
+                },
+            }
         self.assertNestedAllClose(
-            {
-                "optimizer/learning_rate": 1.0,
-                "optimizer/lr_schedule_step": 1,
-                "optimizer/gradient_norm": jnp.sqrt(jnp.sum(2 * expected_grad**2)),
-                "param_rms/weight": jnp.sqrt((0 + 4 + 4 + 9) / 4),
-                "param_rms/moving_mean": 0.5,
-                "grad_rms/weight": jnp.sqrt(jnp.mean(expected_grad**2)),
-                "grad_rms/moving_mean": jnp.sqrt(jnp.mean(expected_grad**2)),
-                "optimizer/schedule_step": 1,
-                "optimizer/schedule_scale": -1.0,
-            },
+            expected_summaries,
             output_collection.summaries,
         )