From a039e1aca9f6d42f101d78af6e61ab19c27ef32e Mon Sep 17 00:00:00 2001 From: Daniel Mandragona Date: Thu, 27 Aug 2026 16:35:50 -0700 Subject: [PATCH] Add layer-by-layer tensor distribution debugging mode in MaxText. This change introduces: 1. Core Tensor Telemetry (debug_tensor_utils.py): Pure JAX/NumPy utility logging tensor distributions (mean, std, min, max, L2 norm, percentiles [1%, 5%, 25%, 50%, 75%, 95%, 99%], NaN count, Inf count, and continuous MoE expert routing mass vectors) across both Forward (FWD) and Backward (BWD) passes. 2. Interceptor Telemetry Infrastructure (debug_tensor_interceptors.py): Modular Linen interceptor (linen_interceptor_fn) and NNX module wrapper (wrap_nnx_module_for_debug) paired with thread-local context management (debug_telemetry_scope). Automatically constructs hierarchical scope paths (e.g. decoder/layers_0/self_attention/query, decoder/layers_0/MoeBlock_0/router_weights) with zero changes to individual model and layer classes. 3. Centralized NNX Hook in ToLinen (layers/nnx_wrappers.py): Integrates with Flax Linen and NNX hybrid architectures. 4. FWD and BWD Pass Autodiff: Uses @jax.custom_vjp with static trace-time short-circuiting to automatically log forward activations and reverse-mode gradient cotangents with zero overhead when disabled. Tested: - Unit tests in //third_party/py/maxtext/tests/unit:debug_tensor_utils_test. - Integration tests in //third_party/py/maxtext/tests/unit:debug_tensor_integration_test. - End-to-end synthetic Mixtral MoE training run (trainers:train with model_name=mixtral-8x7b, num_experts=4, num_experts_per_tok=2, debug_tensor_distribution=True): ``` [DEBUG_TENSOR FWD] step=0 name=decoder/layer_0/moe/expert_inputs shape=(12, 128, 64) dtype=bfloat16 mean=-8.294509e-03 std=9.999614e-01 min=-3.843750e+00 max=4.312500e+00 l2_norm=3.135334e+02 nan_count=0 inf_count=0 [DEBUG_TENSOR FWD] step=0 name=decoder/layer_0/moe/gate_logits shape=(12, 128, 4) dtype=bfloat16 mean=-4.968616e-02 std=9.442701e-01 min=-3.515625e+00 max=3.187500e+00 l2_norm=7.411775e+01 nan_count=0 inf_count=0 [DEBUG_TENSOR FWD] step=0 name=decoder/layer_0/moe/router_weights shape=(12, 128, 2) dtype=bfloat16 mean=5.000076e-01 std=1.925698e-01 min=1.757812e-02 max=9.843750e-01 l2_norm=2.969752e+01 nan_count=0 inf_count=0 [DEBUG_TENSOR FWD] step=0 name=decoder/layer_0/moe/expert_outputs shape=(12, 128, 64) dtype=bfloat16 mean=-1.053208e-03 std=3.509205e-02 min=-1.699219e-01 max=1.699219e-01 l2_norm=1.100753e+01 nan_count=0 inf_count=0 [DEBUG_TENSOR BWD] step=0 name=decoder/layer_0/moe/expert_outputs/grad shape=(12, 128, 64) dtype=bfloat16 mean=-3.270777e-07 std=7.611285e-05 min=-2.393723e-04 max=2.355576e-04 l2_norm=2.386424e-02 nan_count=0 inf_count=0 [DEBUG_TENSOR BWD] step=0 name=decoder/layer_0/moe/router_weights/grad shape=(12, 128, 2) dtype=bfloat16 mean=-3.831090e-07 std=2.770081e-05 min=-1.196861e-04 max=1.049042e-04 l2_norm=1.535482e-03 nan_count=0 inf_count=0 [DEBUG_TENSOR BWD] step=0 name=decoder/layer_0/moe/gate_logits/grad shape=(12, 128, 4) dtype=bfloat16 mean=0.000000e+00 std=6.066302e-06 min=-4.196167e-05 max=4.196167e-05 l2_norm=4.754990e-04 nan_count=0 inf_count=0 [DEBUG_TENSOR BWD] step=0 name=decoder/layer_0/moe/expert_inputs/grad shape=(12, 128, 64) dtype=bfloat16 mean=-1.991233e-09 std=4.040187e-06 min=-2.288818e-05 max=2.342463e-05 l2_norm=1.266739e-03 nan_count=0 inf_count=0 ``` PiperOrigin-RevId: 972220916 --- src/maxtext/configs/base.yml | 4 + src/maxtext/configs/types.py | 9 + src/maxtext/layers/nnx_wrappers.py | 13 + src/maxtext/trainers/pre_train/train.py | 61 ++- .../utils/debug_tensor_interceptors.py | 302 +++++++++++ src/maxtext/utils/debug_tensor_utils.py | 237 ++++++++ tests/unit/debug_tensor_integration_test.py | 505 ++++++++++++++++++ tests/unit/debug_tensor_utils_test.py | 374 +++++++++++++ 8 files changed, 1481 insertions(+), 24 deletions(-) create mode 100644 src/maxtext/utils/debug_tensor_interceptors.py create mode 100644 src/maxtext/utils/debug_tensor_utils.py create mode 100644 tests/unit/debug_tensor_integration_test.py create mode 100644 tests/unit/debug_tensor_utils_test.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 518ef8fc56..402cafdbcb 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1103,6 +1103,10 @@ vertex_tensorboard_region: "" # If set to true, MaxText will perform extra checks using jax.checkify. Note that this will effect performance. max_checkify: false +# Enable tensor distribution debugging across forward and backward passes. +debug_tensor_distribution: false +debug_tensor_distribution_layers: "all" +debug_tensor_distribution_step_interval: 1 # Inference inference_microbenchmark_prefill_lengths: "64,128,256,512,1024" diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 362a97fa3e..d6a4d53cd0 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2325,6 +2325,15 @@ class DevelopmentAndDebugging(BaseModel): False, description="If True, perform extra checks using jax.checkify, affecting performance.", ) + debug_tensor_distribution: bool = Field( + False, + description="Enable tensor distribution debugging in FWD and BWD passes.", + ) + debug_tensor_distribution_layers: str = Field( + "all", + description=("Filter layers/submodules to debug ('all' or comma-separated names)."), + ) + debug_tensor_distribution_step_interval: int = Field(1, description="Step frequency interval for debug tensor logging.") @classmethod def _clean_empty_string_for_list(cls, v: Any) -> Any: diff --git a/src/maxtext/layers/nnx_wrappers.py b/src/maxtext/layers/nnx_wrappers.py index e204502cb2..1bcb1257ba 100644 --- a/src/maxtext/layers/nnx_wrappers.py +++ b/src/maxtext/layers/nnx_wrappers.py @@ -33,6 +33,7 @@ from flax.nnx.rnglib import Rngs import jax from jax import tree_util as jtu +from maxtext.utils import debug_tensor_interceptors import qwix M = tp.TypeVar("M", bound=Module) @@ -517,6 +518,12 @@ def _module_kwargs(): # update linen variables before call module to save initial state self._update_variables(module) _fix_for_qwix_quantization(module) + if debug_tensor_interceptors.is_debug_telemetry_active(): + dbg_cfg, dbg_step = debug_tensor_interceptors.get_active_telemetry_context() + dbg_parent = "/".join(self.path) if self.scope is not None and self.path else self.name or "" + module = debug_tensor_interceptors.wrap_nnx_module_for_debug( + module, parent_path=dbg_parent, step=dbg_step, config=dbg_cfg + ) method_fn = _get_module_method(module, nnx_method) out = method_fn(module, *args, **kwargs) return out @@ -590,6 +597,12 @@ def maybe_unbox(x): module = nnx.merge(graphdef, full_state) _fix_for_qwix_quantization(module) + if debug_tensor_interceptors.is_debug_telemetry_active(): + dbg_cfg, dbg_step = debug_tensor_interceptors.get_active_telemetry_context() + dbg_parent = "/".join(self.path) if self.scope is not None and self.path else self.name or "" + module = debug_tensor_interceptors.wrap_nnx_module_for_debug( + module, parent_path=dbg_parent, step=dbg_step, config=dbg_cfg + ) method_fn = _get_module_method(module, nnx_method) out = method_fn(module, *args, **kwargs) self._update_variables(module) diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 692690db76..84bd7f30ff 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -48,6 +48,8 @@ from maxtext.configs.types import TeCommGemmOverlapPolicy from maxtext.diffusion.block_diffusion import target_alignment as block_diffusion_target_alignment from maxtext.utils.globals import EPS +from maxtext.utils.debug_tensor_interceptors import debug_telemetry_scope, is_debug_telemetry_active, wrap_nnx_module_for_debug +from maxtext.utils.debug_tensor_utils import debug_tensor from maxtext.utils import elastic_utils # Placeholder: internal @@ -107,7 +109,7 @@ def get_first_step(model, state): # ----------------------------------------------------------------------------- -def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_train=True): +def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_train=True, step=0): """loss_fn for both train and eval. Args: @@ -116,7 +118,9 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr data: Batch of data to apply to the model dropout_rng: A key to use to generate rng for dropout (Linen); unused for NNX. params: Model params (Linen); unused for NNX (params are part of the model). + sparsity_state: Batch stats for sparsity is_train: True for train_step and False for eval_step + step: Current training/evaluation step Returns: loss: average loss @@ -204,19 +208,21 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr model_vars["batch_stats"] = sparsity_state else: model_vars = params - logits, intermediate_outputs = model.apply( - model_vars, - data["inputs"], - data["inputs_position"], - decoder_segment_ids=data["inputs_segmentation"], - **encoder_kwargs, - enable_dropout=config.enable_dropout if is_train else False, - rngs={"dropout": rng1, "params": aqt_rng}, # pyrefly: ignore[bad-argument-type] - mutable=mutable_collections, - decoder_target_tokens=data["targets"], - decoder_target_mask=data["targets_segmentation"], - **forced_routing_kwargs, - ) + with debug_telemetry_scope(config, step=step): + logits, intermediate_outputs = model.apply( + model_vars, + data["inputs"], + data["inputs_position"], + decoder_segment_ids=data["inputs_segmentation"], + **encoder_kwargs, + enable_dropout=config.enable_dropout if is_train else False, + rngs={"dropout": rng1, "params": aqt_rng}, # pyrefly: ignore[bad-argument-type] + mutable=mutable_collections, + decoder_target_tokens=data["targets"], + decoder_target_mask=data["targets_segmentation"], + **forced_routing_kwargs, + ) + logits = debug_tensor(logits, "loss/logits", enabled=config) if (config.use_indexer and not config.indexer_sparse_training) and is_train: # In Dense Warm-up stage, we skip main model loss calculation for efficiency. @@ -237,6 +243,7 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr ) one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) + xent = debug_tensor(xent, "loss/cross_entropy_per_token", enabled=config) xent = sharding.maybe_shard_with_logical( xent, @@ -264,16 +271,20 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr total_z_loss = jnp.sum(z_loss) else: # Flax NNX model: forward pass, then pop Intermediates sown during it. - logits = model( - decoder_input_tokens=data["inputs"], - decoder_positions=data["inputs_position"], - decoder_segment_ids=data["inputs_segmentation"], - **encoder_kwargs, - enable_dropout=config.enable_dropout if is_train else False, - decoder_target_tokens=data["targets"], - decoder_target_mask=data["targets_segmentation"], - **forced_routing_kwargs, - ) + with debug_telemetry_scope(config, step=step): + if is_debug_telemetry_active(): + wrap_nnx_module_for_debug(model, parent_path="", step=step, config=config) + logits = model( + decoder_input_tokens=data["inputs"], + decoder_positions=data["inputs_position"], + decoder_segment_ids=data["inputs_segmentation"], + **encoder_kwargs, + enable_dropout=config.enable_dropout if is_train else False, + decoder_target_tokens=data["targets"], + decoder_target_mask=data["targets_segmentation"], + **forced_routing_kwargs, + ) + logits = debug_tensor(logits, "loss/logits", enabled=config) # mtp_losses and mtp_acceptance subclass nnx.Intermediate, and nnx type filters match # subclasses. Pop them before the generic Intermediate pop below, which would otherwise # take them too and leave the MTP loss silently reading as 0. @@ -318,6 +329,7 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr ) one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) + xent = debug_tensor(xent, "loss/cross_entropy_per_token", enabled=config) xent = sharding.maybe_shard_with_logical( xent, @@ -436,6 +448,7 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr "mtp_loss": mtp_loss, "batch_stats": (intermediate_outputs.get("batch_stats", None) if hasattr(intermediate_outputs, "get") else None), } + loss = debug_tensor(loss, "loss/total_loss", enabled=config) return loss, aux diff --git a/src/maxtext/utils/debug_tensor_interceptors.py b/src/maxtext/utils/debug_tensor_interceptors.py new file mode 100644 index 0000000000..3d4419e2db --- /dev/null +++ b/src/maxtext/utils/debug_tensor_interceptors.py @@ -0,0 +1,302 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Flax Linen and NNX interceptors for automatic tensor distribution debugging. + +This module provides interception and wrapping hooks for monitoring tensor +distributions, activations, and routing statistics across Flax Linen and NNX +model layers in MaxText without modifying layer implementations. + +Why This Is Needed: + Debugging numerical anomalies (e.g., NaNs, Infs, vanishing/exploding + activations, and imbalanced MoE router weights) in large-scale distributed + training traditionally requires manually placing debug logging statements or + probes throughout the model definition. This is error-prone, intrusive, and + clutters architecture code. + + This module solves this by providing non-intrusive, automated instrumentation + for both Flax paradigms used in MaxText: + 1. Flax Linen: Intercepts module execution using `nn.intercept_methods` to + capture intermediate forward outputs and auxiliary tensors (such as KV + caches and MoE load-balancing loss). + 2. Flax NNX: Traverses module object graphs via `nnx.iter_graph` and wraps + module instances, hooking `__call__` and MoE routing methods (`get_topk`, + `reshape_and_update_weights`) to capture layer outputs, router weights, + and combine weights. + + Both interceptors integrate with `debug_telemetry_scope`, introducing zero + overhead at trace time when disabled, and conditionally forwarding tensors + to `debug_tensor` when active according to configured step intervals and + layer filters. + +Key Functions and Classes: + - debug_telemetry_scope: Context manager that manages thread-local telemetry + state (config and current step) and enables Linen method interception. + - linen_interceptor_fn: Linen interceptor callback invoked around module calls + to attach hierarchical tags (e.g., `decoder/layers_0/...`) and instrument + outputs. + - wrap_nnx_module_for_debug: Traverses an NNX module hierarchy and subclasses + modules to transparently intercept forward calls and MoE routing methods. + - is_debug_telemetry_active: Queries whether telemetry logging is currently + active in the local thread. + - get_active_telemetry_context: Returns the active configuration and step. +""" + +import contextlib +import dataclasses +import threading +from typing import Any +from flax import linen as nn +from flax import nnx +import jax +import jax.numpy as jnp +from maxtext.utils.debug_tensor_utils import debug_tensor +from maxtext.utils.debug_tensor_utils import should_debug_tensor +import numpy as np + + +@dataclasses.dataclass +class _TelemetryState: + active: bool = False + config: Any = None + step: int | jax.Array = 0 + + +_TELEMETRY_STATE = threading.local() + + +def _get_telemetry_state() -> _TelemetryState: + if not hasattr(_TELEMETRY_STATE, "state"): + _TELEMETRY_STATE.state = _TelemetryState() + return _TELEMETRY_STATE.state + + +def is_debug_telemetry_active() -> bool: + """Returns True if debug telemetry scope is currently active and enabled.""" + state = _get_telemetry_state() + return state.active and getattr(state.config, "debug_tensor_distribution", False) + + +def get_active_telemetry_context() -> tuple[Any, int | jax.Array]: + """Returns (config, step) of currently active debug telemetry scope.""" + state = _get_telemetry_state() + return state.config, state.step + + +def _instrument_output( + out: Any, + primary_tag: str, + aux_tag: str | None, + step: int | jax.Array, +) -> Any: + """Instruments a tensor or tuple of tensors with primary and optional aux debug tags.""" + if isinstance(out, (jax.Array, np.ndarray, jnp.ndarray)): + return debug_tensor(out, primary_tag, step=step, enabled=True) + if isinstance(out, tuple) and len(out) > 0: + items = list(out) + if isinstance(items[0], (jax.Array, np.ndarray, jnp.ndarray)): + items[0] = debug_tensor(items[0], primary_tag, step=step, enabled=True) + if aux_tag and len(items) > 1 and items[1] is not None and isinstance(items[1], (jax.Array, np.ndarray, jnp.ndarray)): + items[1] = debug_tensor(items[1], aux_tag, step=step, enabled=True) + return tuple(items) + return out + + +def linen_interceptor_fn(next_fun, args, kwargs, context: nn.module.InterceptorContext): + """Flax Linen method interceptor that automatically instruments module outputs.""" + out = next_fun(*args, **kwargs) + if not is_debug_telemetry_active(): + return out + + config, step = get_active_telemetry_context() + path_tuple = getattr(context.module, "path", ()) + if not path_tuple: + return out + + path_tag = "/".join(str(p) for p in path_tuple) + if not should_debug_tensor(config, path_tag, step): + return out + + tag_lower = path_tag.lower() + if any(k in tag_lower for k in ("moe", "router", "expert", "mhc")): + aux_tag = f"{path_tag}/load_balancing_loss" + elif any(k in tag_lower for k in ("attn", "attention")): + aux_tag = f"{path_tag}/kv_cache" + else: + aux_tag = f"{path_tag}/aux_output" + + return _instrument_output(out, path_tag, aux_tag, step) + + +_DECODER_LAYER_NAMES = ( + "DecoderLayer", + "NNXDecoderLayer", + "LlamaDecoderLayer", + "MixtralDecoderLayer", + "DeepSeekDenseLayer", + "DeepSeekMoELayer", + "Gemma4DecoderLayer", + "Qwen3NextDecoderLayer", +) + + +def _get_nnx_tags(node: nnx.Module, tag: str) -> tuple[str, str | None]: + """Determines primary and auxiliary output tags for an NNX module.""" + name = type(node).__name__ + if name == "GateLogit": + return f"{tag}/gate_logits", None + if name == "RoutedAndSharedMoE": + return f"{tag}/combined_outputs", f"{tag}/load_balancing_loss" + if ( + name in ("RoutedMoE", "MoeBlock") + or name.endswith(("MoE", "Moe", "MoeBlock", "SparseMoeBlock", "MoELayer")) + or hasattr(node, "get_topk") + ): + return f"{tag}/expert_outputs", f"{tag}/load_balancing_loss" + if name.endswith("DecoderLayer") or name in _DECODER_LAYER_NAMES: + return f"{tag}/layer_output", None + return tag, None + + +def _is_telemetry_active_for_tag(tag: str) -> tuple[bool, int | jax.Array]: + """Checks if debug telemetry is active and enabled for the given tag.""" + if not is_debug_telemetry_active(): + return False, 0 + config, step = get_active_telemetry_context() + return should_debug_tensor(config, tag, step), step + + +def _wrap_call(fn, primary_tag: str, aux_tag: str | None, node_tag: str): + """Wraps an NNX module's __call__ method to instrument its output.""" + + def wrapped_call(self, *args, **kwargs): + out = fn(self, *args, **kwargs) + active, step = _is_telemetry_active_for_tag(node_tag) + if not active: + return out + return _instrument_output(out, primary_tag, aux_tag, step) + + return wrapped_call + + +def _wrap_topk(fn, node_tag: str): + """Wraps an MoE module's get_topk method to instrument router weights.""" + tag = f"{node_tag}/router_weights" + + def wrapped_topk(self, *args, **kwargs): + res = fn(self, *args, **kwargs) + active, step = _is_telemetry_active_for_tag(tag) + if active and isinstance(res, tuple) and len(res) >= 2: + weights = debug_tensor(res[0], tag, step=step, enabled=True) + return (weights, *res[1:]) + return res + + return wrapped_topk + + +def _wrap_reshape_and_update_weights(fn, node_tag: str): + """Wraps an MoE module's reshape_and_update_weights to instrument combine weights.""" + tag = f"{node_tag}/combine_weights" + + def wrapped_reshape(self, *args, **kwargs): + weights = fn(self, *args, **kwargs) + active, step = _is_telemetry_active_for_tag(tag) + if active and isinstance(weights, (jax.Array, np.ndarray, jnp.ndarray)): + return debug_tensor(weights, tag, step=step, enabled=True) + return weights + + return wrapped_reshape + + +def _instrument_nnx_node(node: nnx.Module, node_tag: str) -> None: + """Subclasses an NNX module node in-place to wrap methods with telemetry.""" + primary_tag, aux_tag = _get_nnx_tags(node, node_tag) + methods = { + "__call__": _wrap_call(node.__class__.__call__, primary_tag, aux_tag, node_tag), + "_is_debug_wrapped": True, + } + if hasattr(node, "get_topk"): + methods["get_topk"] = _wrap_topk(node.__class__.get_topk, node_tag) + if hasattr(node, "reshape_and_update_weights"): + methods["reshape_and_update_weights"] = _wrap_reshape_and_update_weights( + node.__class__.reshape_and_update_weights, node_tag + ) + node.__class__ = type(node.__class__.__name__, (node.__class__,), methods) + + +def _format_node_tag(parent_path: str, path: tuple[Any, ...]) -> str: + """Formats hierarchical tag string from parent path and graph path.""" + subpath = "/".join(str(p) for p in path) if path else "" + if parent_path and subpath: + return f"{parent_path}/{subpath}" + return parent_path or subpath + + +def wrap_nnx_module_for_debug( + module: nnx.Module, + parent_path: str = "", + step: int | jax.Array = 0, + config: Any = None, +) -> nnx.Module: + """Centralized NNX wrapping hook to instrument NNX module outputs with hierarchical tags.""" + if config is None: + config, step = get_active_telemetry_context() + + if not getattr(config, "debug_tensor_distribution", False): + return module + + for path, node in nnx.iter_graph(module): + if not isinstance(node, nnx.Module) or getattr(node, "_is_debug_wrapped", False): + continue + + node_tag = _format_node_tag(parent_path, path) + if not node_tag: + continue + + _instrument_nnx_node(node, node_tag) + + return module + + +@contextlib.contextmanager +def debug_telemetry_scope(config: Any = None, step: int | jax.Array = 0): + """Scope context manager activating Flax Linen & NNX unified interceptors. + + When config is None and no active scope exists, or debug_tensor_distribution + is False, + this yields immediately with zero overhead at trace time. + """ + state = _get_telemetry_state() + active_config = config if config is not None else state.config + if active_config is None or not getattr(active_config, "debug_tensor_distribution", False): + old_step = state.step + state.step = step + try: + yield + finally: + state.step = old_step + return + + old_active, old_config, old_step = state.active, state.config, state.step + state.active = True + state.config = active_config + state.step = step + + try: + with nn.intercept_methods(linen_interceptor_fn): + yield + finally: + state.active = old_active + state.config = old_config + state.step = old_step diff --git a/src/maxtext/utils/debug_tensor_utils.py b/src/maxtext/utils/debug_tensor_utils.py new file mode 100644 index 0000000000..95bd0ee2c1 --- /dev/null +++ b/src/maxtext/utils/debug_tensor_utils.py @@ -0,0 +1,237 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tensor distribution debugging utilities for MaxText.""" + +import functools +from typing import Any +import jax +import jax.numpy as jnp +import numpy as np + + +def _compute_stats(x: Any) -> dict[str, Any]: + """Computes mean, std, min, max, l2_norm, percentiles, NaNs, and Infs for a tensor.""" + x_f32 = jnp.asarray(x, dtype=jnp.float32) + x_flat = jnp.ravel(x_f32) + + mean = jnp.mean(x_flat) + std = jnp.std(x_flat) + min_val = jnp.min(x_flat) + max_val = jnp.max(x_flat) + l2_norm = jnp.linalg.norm(x_flat) + + # Non-optional percentiles: 1%, 5%, 25%, 50%, 75%, 95%, 99% + pct_qs = jnp.array([1.0, 5.0, 25.0, 50.0, 75.0, 95.0, 99.0]) + pcts = jnp.percentile(x_flat, pct_qs) + + nan_count = jnp.sum(jnp.isnan(x_flat)) + inf_count = jnp.sum(jnp.isinf(x_flat)) + + return { + "mean": mean, + "std": std, + "min": min_val, + "max": max_val, + "l2_norm": l2_norm, + "p01": pcts[0], + "p05": pcts[1], + "p25": pcts[2], + "p50": pcts[3], + "p75": pcts[4], + "p95": pcts[5], + "p99": pcts[6], + "nans": nan_count, + "infs": inf_count, + "nan_count": nan_count, + "inf_count": inf_count, + } + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3)) +def _debug_tensor_vjp(x: Any, name: str, step: int | jax.Array, enabled: bool) -> Any: + """Custom VJP wrapper to tap tensor distributions in forward and backward passes.""" + return x + + +def _debug_tensor_fwd(x: Any, name: str, step: int | jax.Array, enabled: bool): + """Forward rule for _debug_tensor_vjp computing and printing tensor statistics.""" + if enabled: + stats = _compute_stats(x) + shape_str = getattr(x, "shape", ()) + dtype_str = getattr(x, "dtype", type(x).__name__) + is_moe_routing = name.endswith(("router_weights", "combine_weights", "gate_logits", "expert_weights")) or ( + "moe" in name and name.endswith(("weights", "logits")) + ) + if is_moe_routing and hasattr(x, "ndim") and x.ndim >= 2: + x_f32 = jnp.asarray(x, dtype=jnp.float32) + mean_expert_weights = jnp.mean(x_f32, axis=tuple(range(x.ndim - 1))) + jax.debug.print( + "[DEBUG_TENSOR FWD] step={step} name={name} shape={shape}" + " dtype={dtype} mean={mean:.6e} std={std:.6e} min={min:.6e}" + " max={max:.6e} l2_norm={l2_norm:.6e} p01={p01:.6e} p05={p05:.6e}" + " p25={p25:.6e} p50={p50:.6e} p75={p75:.6e} p95={p95:.6e}" + " p99={p99:.6e} nan_count={nan_count} inf_count={inf_count}" + " expert_weights={expert_weights}", + step=step, + name=name, + shape=shape_str, + dtype=str(dtype_str), + mean=stats["mean"], + std=stats["std"], + min=stats["min"], + max=stats["max"], + l2_norm=stats["l2_norm"], + p01=stats["p01"], + p05=stats["p05"], + p25=stats["p25"], + p50=stats["p50"], + p75=stats["p75"], + p95=stats["p95"], + p99=stats["p99"], + nan_count=stats["nan_count"], + inf_count=stats["inf_count"], + expert_weights=mean_expert_weights, + ) + else: + jax.debug.print( + "[DEBUG_TENSOR FWD] step={step} name={name} shape={shape}" + " dtype={dtype} mean={mean:.6e} std={std:.6e} min={min:.6e}" + " max={max:.6e} l2_norm={l2_norm:.6e} p01={p01:.6e} p05={p05:.6e}" + " p25={p25:.6e} p50={p50:.6e} p75={p75:.6e} p95={p95:.6e}" + " p99={p99:.6e} nan_count={nan_count} inf_count={inf_count}", + step=step, + name=name, + shape=shape_str, + dtype=str(dtype_str), + mean=stats["mean"], + std=stats["std"], + min=stats["min"], + max=stats["max"], + l2_norm=stats["l2_norm"], + p01=stats["p01"], + p05=stats["p05"], + p25=stats["p25"], + p50=stats["p50"], + p75=stats["p75"], + p95=stats["p95"], + p99=stats["p99"], + nan_count=stats["nan_count"], + inf_count=stats["inf_count"], + ) + return x, None + + +def _debug_tensor_bwd(name: str, step: int | jax.Array, enabled: bool, res: Any, g: Any): + """Backward rule for _debug_tensor_vjp computing and printing gradient statistics.""" + if enabled and g is not None: + stats = _compute_stats(g) + shape_str = getattr(g, "shape", ()) + dtype_str = getattr(g, "dtype", type(g).__name__) + jax.debug.print( + "[DEBUG_TENSOR BWD] step={step} name={name}/grad shape={shape}" + " dtype={dtype} mean={mean:.6e} std={std:.6e} min={min:.6e}" + " max={max:.6e} l2_norm={l2_norm:.6e} p01={p01:.6e} p05={p05:.6e}" + " p25={p25:.6e} p50={p50:.6e} p75={p75:.6e} p95={p95:.6e} p99={p99:.6e}" + " nan_count={nan_count} inf_count={inf_count}", + step=step, + name=name, + shape=shape_str, + dtype=str(dtype_str), + mean=stats["mean"], + std=stats["std"], + min=stats["min"], + max=stats["max"], + l2_norm=stats["l2_norm"], + p01=stats["p01"], + p05=stats["p05"], + p25=stats["p25"], + p50=stats["p50"], + p75=stats["p75"], + p95=stats["p95"], + p99=stats["p99"], + nan_count=stats["nan_count"], + inf_count=stats["inf_count"], + ) + return (g,) + + +_debug_tensor_vjp.defvjp(_debug_tensor_fwd, _debug_tensor_bwd) + + +def _get_active_step() -> int | jax.Array: + """Returns current active step from telemetry scope if active, else 0.""" + try: + # pylint: disable=import-outside-toplevel + from maxtext.utils import debug_tensor_interceptors + + _, step = debug_tensor_interceptors.get_active_telemetry_context() + return step + except (ImportError, AttributeError): + return 0 + + +def should_debug_tensor( + config: Any, + name: str, + step: int | jax.Array | None = None, +) -> bool: + """Determines if debug logging should be enabled for a given tensor name and step.""" + if step is None: + step = _get_active_step() + if config is None: + return False + if not getattr(config, "debug_tensor_distribution", False): + return False + layers_filter = getattr(config, "debug_tensor_distribution_layers", "all") + if layers_filter and layers_filter != "all": + allowed_layers = [l.strip() for l in layers_filter.split(",") if l.strip()] + if not any(layer in name for layer in allowed_layers): + return False + interval = getattr(config, "debug_tensor_distribution_step_interval", 1) + if isinstance(step, (int, np.integer)) and interval > 1: + if step % interval != 0: + return False + return True + + +def debug_tensor( + x: Any, + name: str, + step: int | jax.Array | None = None, + enabled: bool | Any = True, +) -> Any: + """Identity function instrumented with FWD and BWD distribution logging. + + When enabled=False, statically returns x directly at trace time with zero + Jaxpr overhead. + When enabled is a Config object, automatically evaluates should_debug_tensor. + """ + if step is None: + step = _get_active_step() + if not isinstance(enabled, (bool, np.bool_)): + enabled = should_debug_tensor(enabled, name, step) + if not enabled: + return x + return _debug_tensor_vjp(x, name, step, True) + + +def debug_tensor_from_config( + x: Any, + name: str, + config: Any, + step: int | jax.Array | None = None, +) -> Any: + """Convenience wrapper around debug_tensor using MaxText Config.""" + return debug_tensor(x, name, step=step, enabled=config) diff --git a/tests/unit/debug_tensor_integration_test.py b/tests/unit/debug_tensor_integration_test.py new file mode 100644 index 0000000000..4abfb64628 --- /dev/null +++ b/tests/unit/debug_tensor_integration_test.py @@ -0,0 +1,505 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration tests for tensor distribution debugging in MaxText layers.""" + +import sys +import unittest +from absl.testing import absltest +from flax import nnx +import jax +import jax.numpy as jnp +from jax.sharding import Mesh +from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_TRAIN +from maxtext.configs import pyconfig +from maxtext.layers.decoders import DecoderLayer, SequentialBlockDecoderLayers +from maxtext.layers.initializers import nd_dense_init +from maxtext.layers.moe import RoutedMoE +from maxtext.layers.nnx_decoders import NNXDecoderLayer +from maxtext.models.llama2 import LlamaDecoderLayer +from maxtext.models.mixtral import MixtralDecoderLayer +from maxtext.utils import debug_tensor_interceptors +from maxtext.utils import debug_tensor_utils +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path +import numpy as np + +_BASE_CONFIG = { + "per_device_batch_size": 1.0, + "run_name": "debug_tensor_integration_test", + "enable_checkpointing": False, + "base_num_decoder_layers": 1, + "attention": "dot_product", + "max_target_length": 16, + "base_emb_dim": 64, + "base_num_query_heads": 2, + "base_num_kv_heads": 2, + "base_mlp_dim": 128, + "max_prefill_predict_length": 4, + "scan_layers": False, +} + + +def _make_config(**overrides): + merged = {**_BASE_CONFIG, **overrides} + return pyconfig.initialize([sys.argv[0], get_test_config_path()], override_model_config=True, **merged) + + +def _make_mesh(cfg): + devices_array = maxtext_utils.create_device_mesh(cfg) + return Mesh(devices_array, cfg.mesh_axes) + + +class DebugTensorIntegrationTest(unittest.TestCase): + + def setUp(self): + super().setUp() + self.rng = jax.random.PRNGKey(0) + + def _make_inputs(self, cfg): + batch = cfg.global_batch_size_to_train_on + seq_len = cfg.max_target_length + emb_dim = cfg.emb_dim + inputs = jax.random.normal(self.rng, (batch, seq_len, emb_dim), dtype=jnp.float32) + segment_ids = jnp.full((batch, seq_len), DECODING_ACTIVE_SEQUENCE_INDICATOR) + positions = jnp.broadcast_to(jnp.arange(seq_len)[None], (batch, seq_len)) + return inputs, segment_ids, positions + + def test_linen_decoder_layer_debug_tensor(self): + cfg_disabled = _make_config(debug_tensor_distribution=False) + cfg_enabled = _make_config(debug_tensor_distribution=True) + mesh = _make_mesh(cfg_disabled) + + inputs, segment_ids, positions = self._make_inputs(cfg_disabled) + + layer_disabled = DecoderLayer( + config=cfg_disabled, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + ) + layer_enabled = DecoderLayer( + config=cfg_enabled, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + ) + + init_rng = jax.random.PRNGKey(42) + variables = layer_disabled.init( + init_rng, + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + # Forward pass + out_disabled, _ = layer_disabled.apply( + variables, + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + out_enabled, _ = layer_enabled.apply( + variables, + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + np.testing.assert_allclose(np.array(out_disabled), np.array(out_enabled), rtol=1e-6, atol=1e-6) + + # Gradient computation + def loss_disabled(x): + out, _ = layer_disabled.apply( + variables, + x, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return jnp.sum(out**2) + + def loss_enabled(x): + out, _ = layer_enabled.apply( + variables, + x, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return jnp.sum(out**2) + + grad_disabled = jax.grad(loss_disabled)(inputs) + grad_enabled = jax.grad(loss_enabled)(inputs) + + np.testing.assert_allclose(np.array(grad_disabled), np.array(grad_enabled), rtol=1e-6, atol=1e-6) + + def test_nnx_decoder_layer_debug_tensor(self): + cfg_disabled = _make_config(debug_tensor_distribution=False) + cfg_enabled = _make_config(debug_tensor_distribution=True) + mesh = _make_mesh(cfg_disabled) + + inputs, segment_ids, positions = self._make_inputs(cfg_disabled) + + rngs = nnx.Rngs(params=42, dropout=1) + layer_disabled = NNXDecoderLayer( + config=cfg_disabled, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=rngs, + ) + # Clone state into layer_enabled for exact weight parity + rngs_en = nnx.Rngs(params=42, dropout=1) + layer_enabled = NNXDecoderLayer( + config=cfg_enabled, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=rngs_en, + ) + + out_disabled, _ = layer_disabled( + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + out_enabled, _ = layer_enabled( + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + np.testing.assert_allclose(np.array(out_disabled), np.array(out_enabled), rtol=1e-6, atol=1e-6) + + # Gradient computation through NNX + def loss_nnx(layer, x): + out, _ = layer( + x, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return jnp.sum(out**2) + + grad_disabled = jax.grad(loss_nnx, argnums=1)(layer_disabled, inputs) + grad_enabled = jax.grad(loss_nnx, argnums=1)(layer_enabled, inputs) + np.testing.assert_allclose(np.array(grad_disabled), np.array(grad_enabled), rtol=1e-6, atol=1e-6) + + def test_moe_debug_tensor(self): + cfg_disabled = _make_config( + debug_tensor_distribution=False, + base_num_decoder_layers=1, + base_emb_dim=64, + base_mlp_dim=128, + moe_mlp_dim=128, + base_moe_mlp_dim=128, + num_experts=4, + num_experts_per_tok=2, + ) + cfg_enabled = _make_config( + debug_tensor_distribution=True, + base_num_decoder_layers=1, + base_emb_dim=64, + base_mlp_dim=128, + moe_mlp_dim=128, + base_moe_mlp_dim=128, + num_experts=4, + num_experts_per_tok=2, + ) + mesh = _make_mesh(cfg_enabled) + rngs_dis = nnx.Rngs(params=42, dropout=1) + rngs_en = nnx.Rngs(params=42, dropout=1) + + moe_disabled = RoutedMoE( + config=cfg_disabled, + num_experts=4, + num_experts_per_tok=2, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + dtype=cfg_disabled.dtype, + rngs=rngs_dis, + ) + moe_enabled = RoutedMoE( + config=cfg_enabled, + num_experts=4, + num_experts_per_tok=2, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + dtype=cfg_enabled.dtype, + rngs=rngs_en, + ) + + batch = cfg_enabled.global_batch_size_to_train_on + seq_len = cfg_enabled.max_target_length + emb_dim = cfg_enabled.emb_dim + inputs = jax.random.normal(self.rng, (batch, seq_len, emb_dim), dtype=jnp.float32) + + out_dis, lb_loss_dis, _ = moe_disabled(inputs) + out_en, lb_loss_en, _ = moe_enabled(inputs) + + self.assertEqual(out_en.shape, inputs.shape) + np.testing.assert_allclose(np.array(out_dis), np.array(out_en), rtol=1e-6, atol=1e-6) + if lb_loss_dis is not None and lb_loss_en is not None: + np.testing.assert_allclose(np.array(lb_loss_dis), np.array(lb_loss_en), rtol=1e-6, atol=1e-6) + else: + self.assertEqual(lb_loss_dis, lb_loss_en) + + # MoE Gradient computation + def moe_loss(layer, x): + out, lb, _ = layer(x) + loss_val = jnp.sum(out**2) + if lb is not None: + loss_val += lb + return loss_val + + grad_dis = jax.grad(moe_loss, argnums=1)(moe_disabled, inputs) + grad_en = jax.grad(moe_loss, argnums=1)(moe_enabled, inputs) + np.testing.assert_allclose(np.array(grad_dis), np.array(grad_en), rtol=1e-6, atol=1e-6) + + def test_llama2_layer_unroll_debug_tensor(self): + cfg_disabled = _make_config(debug_tensor_distribution=False) + cfg_enabled = _make_config(debug_tensor_distribution=True) + mesh = _make_mesh(cfg_disabled) + + inputs, segment_ids, positions = self._make_inputs(cfg_disabled) + + rngs_0 = nnx.Rngs(params=42, dropout=1) + layer0_en = LlamaDecoderLayer( + config=cfg_enabled, + model_mode=MODEL_MODE_TRAIN, + mesh=mesh, + rngs=rngs_0, + ) + layer0_en = debug_tensor_interceptors.wrap_nnx_module_for_debug( + layer0_en, parent_path="decoder/layer_0", config=cfg_enabled + ) + + rngs_1 = nnx.Rngs(params=42, dropout=1) + layer1_en = LlamaDecoderLayer( + config=cfg_enabled, + model_mode=MODEL_MODE_TRAIN, + mesh=mesh, + rngs=rngs_1, + ) + layer1_en = debug_tensor_interceptors.wrap_nnx_module_for_debug( + layer1_en, parent_path="decoder/layer_1", config=cfg_enabled + ) + + # Forward pass and gradient check + out0, _ = layer0_en( + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + self.assertEqual(out0.shape, inputs.shape) + + rngs_dis = nnx.Rngs(params=42, dropout=1) + layer0_dis = LlamaDecoderLayer( + config=cfg_disabled, + model_mode=MODEL_MODE_TRAIN, + mesh=mesh, + rngs=rngs_dis, + ) + out0_dis, _ = layer0_dis( + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + np.testing.assert_allclose(np.array(out0), np.array(out0_dis), rtol=1e-6, atol=1e-6) + + def test_sequential_decoder_layer_unroll(self): + cfg_disabled = _make_config(debug_tensor_distribution=False, base_num_decoder_layers=2, scan_layers=True) + cfg_enabled = _make_config(debug_tensor_distribution=True, base_num_decoder_layers=2, scan_layers=True) + mesh = _make_mesh(cfg_disabled) + + inputs, segment_ids, positions = self._make_inputs(cfg_disabled) + + seq_layers_dis = SequentialBlockDecoderLayers( + decoder_layer=DecoderLayer, + num_decoder_layers=2, + config=cfg_disabled, + mesh=mesh, + quant=None, + model_mode=MODEL_MODE_TRAIN, + ) + seq_layers_en = SequentialBlockDecoderLayers( + decoder_layer=DecoderLayer, + num_decoder_layers=2, + config=cfg_enabled, + mesh=mesh, + quant=None, + model_mode=MODEL_MODE_TRAIN, + ) + + init_rng = jax.random.PRNGKey(42) + variables = seq_layers_dis.init( + init_rng, + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + out_dis = seq_layers_dis.apply( + variables, + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + with debug_tensor_interceptors.debug_telemetry_scope(cfg_enabled): + out_en = seq_layers_en.apply( + variables, + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + np.testing.assert_allclose(np.array(out_dis[0]), np.array(out_en[0]), rtol=1e-6, atol=1e-6) + + def test_mixtral_decoder_layer_debug_tensor(self): + cfg_disabled = _make_config( + debug_tensor_distribution=False, + base_num_decoder_layers=1, + base_emb_dim=64, + base_mlp_dim=128, + moe_mlp_dim=128, + base_moe_mlp_dim=128, + num_experts=4, + num_experts_per_tok=2, + ) + cfg_enabled = _make_config( + debug_tensor_distribution=True, + base_num_decoder_layers=1, + base_emb_dim=64, + base_mlp_dim=128, + moe_mlp_dim=128, + base_moe_mlp_dim=128, + num_experts=4, + num_experts_per_tok=2, + ) + mesh = _make_mesh(cfg_enabled) + inputs, segment_ids, positions = self._make_inputs(cfg_enabled) + + rngs_dis = nnx.Rngs(params=42, dropout=1) + rngs_en = nnx.Rngs(params=42, dropout=1) + + layer_dis = MixtralDecoderLayer( + config=cfg_disabled, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=rngs_dis, + ) + layer_en = MixtralDecoderLayer( + config=cfg_enabled, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=rngs_en, + ) + layer_en = debug_tensor_interceptors.wrap_nnx_module_for_debug( + layer_en, parent_path="decoder/layers_0", config=cfg_enabled + ) + + with debug_tensor_interceptors.debug_telemetry_scope(cfg_enabled, step=0): + out_en, _ = layer_en( + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + out_dis, _ = layer_dis( + inputs, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + + self.assertEqual(out_en.shape, inputs.shape) + np.testing.assert_allclose(np.array(out_en), np.array(out_dis), rtol=1e-6, atol=1e-6) + + def test_full_moe_telemetry_capture_and_visualizer(self): + cfg_enabled = _make_config( + debug_tensor_distribution=True, + base_num_decoder_layers=1, + base_emb_dim=64, + base_mlp_dim=128, + moe_mlp_dim=128, + base_moe_mlp_dim=128, + num_experts=4, + num_experts_per_tok=2, + ) + mesh = _make_mesh(cfg_enabled) + inputs, segment_ids, positions = self._make_inputs(cfg_enabled) + + rngs = nnx.Rngs(params=42, dropout=1) + layer = MixtralDecoderLayer( + config=cfg_enabled, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=rngs, + ) + layer = debug_tensor_interceptors.wrap_nnx_module_for_debug(layer, parent_path="decoder/layers_0", config=cfg_enabled) + + def log_capture_forward_backward(): + with debug_tensor_interceptors.debug_telemetry_scope(cfg_enabled, step=0): + # Instrument dummy loss and embeddings + emb = debug_tensor_utils.debug_tensor(inputs, "decoder/token_embeddings", step=0, enabled=True) + + def loss_fn(x): + out, _ = layer( + x, + segment_ids, + positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + logits = debug_tensor_utils.debug_tensor(out, "loss/logits", step=0, enabled=True) + loss_val = jnp.sum(logits**2) + return debug_tensor_utils.debug_tensor(loss_val, "loss/cross_entropy_per_token", step=0, enabled=True) + + grad_fn = jax.grad(loss_fn) + g = grad_fn(emb) + return g + + # Run the instrumented forward and backward pass + grad_val = log_capture_forward_backward() + self.assertEqual(grad_val.shape, inputs.shape) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/debug_tensor_utils_test.py b/tests/unit/debug_tensor_utils_test.py new file mode 100644 index 0000000000..5628b83660 --- /dev/null +++ b/tests/unit/debug_tensor_utils_test.py @@ -0,0 +1,374 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# pylint: disable=protected-access +"""Unit tests for debug_tensor_utils.""" + +from absl.testing import absltest +from absl.testing import parameterized +from flax import linen as nn +from flax import nnx +import jax +import jax.numpy as jnp +from maxtext.utils import debug_tensor_interceptors +from maxtext.utils import debug_tensor_utils +import numpy as np + + +class DummyConfig: + """Mock config for testing debug_tensor filtering.""" + + def __init__( + self, + debug_tensor_distribution=False, + debug_tensor_distribution_layers="all", + debug_tensor_distribution_step_interval=1, + ): + self.debug_tensor_distribution = debug_tensor_distribution + self.debug_tensor_distribution_layers = debug_tensor_distribution_layers + self.debug_tensor_distribution_step_interval = debug_tensor_distribution_step_interval + + +class DebugTensorUtilsTest(parameterized.TestCase): + + def test_compute_stats_accuracy(self): + np.random.seed(42) + data = np.random.randn(10, 20).astype(np.float32) + x = jnp.array(data) + + stats = debug_tensor_utils._compute_stats(x) + + np.testing.assert_allclose(float(stats["mean"]), float(np.mean(data)), rtol=1e-5) + np.testing.assert_allclose(float(stats["std"]), float(np.std(data)), rtol=1e-5) + np.testing.assert_allclose(float(stats["min"]), float(np.min(data)), rtol=1e-5) + np.testing.assert_allclose(float(stats["max"]), float(np.max(data)), rtol=1e-5) + np.testing.assert_allclose(float(stats["l2_norm"]), float(np.linalg.norm(data)), rtol=1e-5) + + expected_pcts = np.percentile(data, [1.0, 5.0, 25.0, 50.0, 75.0, 95.0, 99.0]) + np.testing.assert_allclose(float(stats["p01"]), float(expected_pcts[0]), rtol=1e-4) + np.testing.assert_allclose(float(stats["p05"]), float(expected_pcts[1]), rtol=1e-4) + np.testing.assert_allclose(float(stats["p25"]), float(expected_pcts[2]), rtol=1e-4) + np.testing.assert_allclose(float(stats["p50"]), float(expected_pcts[3]), rtol=1e-4) + np.testing.assert_allclose(float(stats["p75"]), float(expected_pcts[4]), rtol=1e-4) + np.testing.assert_allclose(float(stats["p95"]), float(expected_pcts[5]), rtol=1e-4) + np.testing.assert_allclose(float(stats["p99"]), float(expected_pcts[6]), rtol=1e-4) + + self.assertEqual(int(stats["nans"]), 0) + self.assertEqual(int(stats["infs"]), 0) + + def test_compute_stats_nan_and_inf_detection(self): + data = np.array([1.0, 2.0, np.nan, 4.0, np.inf, -np.inf, np.nan], dtype=np.float32) + x = jnp.array(data) + + stats = debug_tensor_utils._compute_stats(x) + self.assertEqual(int(stats["nans"]), 2) + self.assertEqual(int(stats["infs"]), 2) + + def test_vjp_gradient_preservation(self): + x = jnp.array([-2.0, -1.0, 0.5, 2.0, 3.0], dtype=jnp.float32) + + def f_instrumented(val): + d = debug_tensor_utils.debug_tensor(val, "layer/activation", step=0, enabled=True) + return jnp.sum(d**3) + + def f_reference(val): + return jnp.sum(val**3) + + grad_instrumented = jax.grad(f_instrumented)(x) + grad_reference = jax.grad(f_reference)(x) + + np.testing.assert_array_equal(grad_instrumented, grad_reference) + np.testing.assert_array_equal(grad_instrumented, 3.0 * (x**2)) + + def test_compute_stats_scalar(self): + scalar = jnp.array(42.0, dtype=jnp.float32) + stats = debug_tensor_utils._compute_stats(scalar) + + self.assertEqual(float(stats["mean"]), 42.0) + self.assertEqual(float(stats["std"]), 0.0) + self.assertEqual(float(stats["min"]), 42.0) + self.assertEqual(float(stats["max"]), 42.0) + self.assertEqual(float(stats["l2_norm"]), 42.0) + self.assertEqual(float(stats["p50"]), 42.0) + self.assertEqual(int(stats["nan_count"]), 0) + self.assertEqual(int(stats["inf_count"]), 0) + self.assertEqual(int(stats["nans"]), 0) + self.assertEqual(int(stats["infs"]), 0) + + def test_zero_overhead_when_disabled(self): + def f_instrumented_bool(x): + y = x * 2.0 + y = debug_tensor_utils.debug_tensor(y, "my_tensor", enabled=False) + return y + 1.0 + + def f_instrumented_none(x): + y = x * 2.0 + y = debug_tensor_utils.debug_tensor(y, "my_tensor", enabled=None) + return y + 1.0 + + def f_instrumented_cfg_disabled(x): + y = x * 2.0 + cfg = DummyConfig(debug_tensor_distribution=False) + y = debug_tensor_utils.debug_tensor(y, "my_tensor", enabled=cfg) + return y + 1.0 + + def f_plain(x): + y = x * 2.0 + return y + 1.0 + + x_dummy = jnp.zeros((4, 8), dtype=jnp.float32) + jaxpr_plain = jax.make_jaxpr(f_plain)(x_dummy) + jaxpr_bool = jax.make_jaxpr(f_instrumented_bool)(x_dummy) + jaxpr_none = jax.make_jaxpr(f_instrumented_none)(x_dummy) + jaxpr_cfg = jax.make_jaxpr(f_instrumented_cfg_disabled)(x_dummy) + + self.assertEqual(len(jaxpr_bool.eqns), len(jaxpr_plain.eqns)) + self.assertEqual(str(jaxpr_bool.jaxpr), str(jaxpr_plain.jaxpr)) + + self.assertEqual(len(jaxpr_none.eqns), len(jaxpr_plain.eqns)) + self.assertEqual(str(jaxpr_none.jaxpr), str(jaxpr_plain.jaxpr)) + + self.assertEqual(len(jaxpr_cfg.eqns), len(jaxpr_plain.eqns)) + self.assertEqual(str(jaxpr_cfg.jaxpr), str(jaxpr_plain.jaxpr)) + + def test_should_debug_tensor_filtering(self): + # None config + self.assertFalse(debug_tensor_utils.should_debug_tensor(None, "test")) + + # Disabled config + cfg_disabled = DummyConfig(debug_tensor_distribution=False) + self.assertFalse(debug_tensor_utils.should_debug_tensor(cfg_disabled, "test")) + + # Enabled all layers + cfg_all = DummyConfig(debug_tensor_distribution=True, debug_tensor_distribution_layers="all") + self.assertTrue(debug_tensor_utils.should_debug_tensor(cfg_all, "attn/query_proj")) + self.assertTrue(debug_tensor_utils.should_debug_tensor(cfg_all, "mlp/out_proj")) + + # Layer filtering + cfg_filtered = DummyConfig( + debug_tensor_distribution=True, + debug_tensor_distribution_layers="attn,mlp", + ) + self.assertTrue(debug_tensor_utils.should_debug_tensor(cfg_filtered, "attn/query_proj")) + self.assertTrue(debug_tensor_utils.should_debug_tensor(cfg_filtered, "decoder/mlp/out_proj")) + self.assertFalse(debug_tensor_utils.should_debug_tensor(cfg_filtered, "norm/rms_norm")) + self.assertFalse(debug_tensor_utils.should_debug_tensor(cfg_filtered, "embed/token_embeddings")) + + # Step interval filtering + cfg_interval = DummyConfig( + debug_tensor_distribution=True, + debug_tensor_distribution_layers="all", + debug_tensor_distribution_step_interval=5, + ) + self.assertTrue(debug_tensor_utils.should_debug_tensor(cfg_interval, "attn/query", step=0)) + self.assertFalse(debug_tensor_utils.should_debug_tensor(cfg_interval, "attn/query", step=1)) + self.assertFalse(debug_tensor_utils.should_debug_tensor(cfg_interval, "attn/query", step=4)) + self.assertTrue(debug_tensor_utils.should_debug_tensor(cfg_interval, "attn/query", step=5)) + self.assertTrue(debug_tensor_utils.should_debug_tensor(cfg_interval, "attn/query", step=10)) + + def test_debug_tensor_with_config(self): + cfg_enabled = DummyConfig(debug_tensor_distribution=True) + cfg_disabled = DummyConfig(debug_tensor_distribution=False) + + x = jnp.ones((2, 4), dtype=jnp.float32) + + # Enabled via config object + out_enabled = debug_tensor_utils.debug_tensor(x, "tensor_a", step=0, enabled=cfg_enabled) + np.testing.assert_array_equal(out_enabled, x) + + # Disabled via config object + out_disabled = debug_tensor_utils.debug_tensor(x, "tensor_b", step=0, enabled=cfg_disabled) + np.testing.assert_array_equal(out_disabled, x) + + # debug_tensor_from_config helper + out_helper = debug_tensor_utils.debug_tensor_from_config(x, "tensor_c", cfg_enabled, step=0) + np.testing.assert_array_equal(out_helper, x) + + def test_debug_telemetry_scope_activation(self): + cfg_disabled = DummyConfig(debug_tensor_distribution=False) + cfg_enabled = DummyConfig(debug_tensor_distribution=True) + + self.assertFalse(debug_tensor_interceptors.is_debug_telemetry_active()) + + # Disabled scope -> should stay inactive + with debug_tensor_interceptors.debug_telemetry_scope(cfg_disabled, step=1): + self.assertFalse(debug_tensor_interceptors.is_debug_telemetry_active()) + + # Enabled scope -> should be active inside, inactive outside + with debug_tensor_interceptors.debug_telemetry_scope(cfg_enabled, step=3): + self.assertTrue(debug_tensor_interceptors.is_debug_telemetry_active()) + active_cfg, active_step = debug_tensor_interceptors.get_active_telemetry_context() + self.assertEqual(active_cfg, cfg_enabled) + self.assertEqual(active_step, 3) + + self.assertFalse(debug_tensor_interceptors.is_debug_telemetry_active()) + + def test_linen_interceptor_with_flax_module(self): + class SimpleSubmodule(nn.Module): + """Mock Linen submodule for testing.""" + + @nn.compact + def __call__(self, x): + return x * 2.0 + + class SimpleModel(nn.Module): + """Mock Linen root model for testing.""" + + @nn.compact + def __call__(self, x): + sub = SimpleSubmodule(name="sub") + return sub(x) + 1.0 + + cfg_enabled = DummyConfig(debug_tensor_distribution=True) + x = jnp.ones((2, 3), dtype=jnp.float32) + model = SimpleModel() + + with debug_tensor_interceptors.debug_telemetry_scope(cfg_enabled, step=0): + out, _ = model.init_with_output(jax.random.PRNGKey(0), x) + + np.testing.assert_allclose(np.array(out), np.ones((2, 3)) * 3.0) + + def test_wrap_nnx_module_for_debug(self): + class SubNNX(nnx.Module): + """Mock NNX leaf module for testing.""" + + def __init__(self, rngs: nnx.Rngs): + self.w = nnx.Param(jax.random.normal(rngs.params(), (4, 4))) + + def __call__(self, x): + return jnp.dot(x, self.w.value) + + class RootNNX(nnx.Module): + """Mock NNX root module for testing.""" + + def __init__(self, rngs: nnx.Rngs): + self.sub = SubNNX(rngs) + + def __call__(self, x): + return self.sub(x) + + rngs = nnx.Rngs(params=0) + root = RootNNX(rngs) + cfg_enabled = DummyConfig(debug_tensor_distribution=True) + + wrapped = debug_tensor_interceptors.wrap_nnx_module_for_debug( + root, parent_path="decoder/layers_0", step=0, config=cfg_enabled + ) + x = jnp.ones((2, 4), dtype=jnp.float32) + with debug_tensor_interceptors.debug_telemetry_scope(cfg_enabled, step=0): + out = wrapped(x) + self.assertEqual(out.shape, (2, 4)) + + def test_wrap_nnx_moe_routing_methods(self): + class DummyMoE(nnx.Module): + """Mock NNX MoE module for testing.""" + + def __init__(self, rngs: nnx.Rngs): + self.w = nnx.Param(jax.random.normal(rngs.params(), (4, 4))) + + def get_topk(self, gate_logits, pre_bias_logits, rngs=None, input_ids=None): + return gate_logits * 2.0, jnp.zeros_like(gate_logits, dtype=jnp.int32) + + def reshape_and_update_weights(self, weights, indices): + return weights * 1.5 + + def __call__(self, x): + weights, indices = self.get_topk(x, x) + comb_weights = self.reshape_and_update_weights(weights, indices) + return jnp.dot(comb_weights, self.w.value), jnp.array(0.05, dtype=jnp.float32) + + rngs = nnx.Rngs(params=0) + moe = DummyMoE(rngs) + cfg_enabled = DummyConfig(debug_tensor_distribution=True) + + wrapped = debug_tensor_interceptors.wrap_nnx_module_for_debug( + moe, parent_path="decoder/layers_0/moe", step=0, config=cfg_enabled + ) + x = jnp.ones((2, 4), dtype=jnp.float32) + with debug_tensor_interceptors.debug_telemetry_scope(cfg_enabled, step=0): + out, lb = wrapped(x) + self.assertEqual(out.shape, (2, 4)) + self.assertAlmostEqual(float(lb), 0.05, places=5) + + def test_dynamic_step_and_scope_switching(self): + class SimpleNNX(nnx.Module): + """Mock NNX simple module for testing.""" + + def __call__(self, x): + return x * 3.0 + + node = SimpleNNX() + cfg_interval = DummyConfig( + debug_tensor_distribution=True, + debug_tensor_distribution_step_interval=5, + ) + wrapped = debug_tensor_interceptors.wrap_nnx_module_for_debug(node, parent_path="layer_0", config=cfg_interval) + x = jnp.ones((2, 2), dtype=jnp.float32) + + # Step 0: should be active (0 % 5 == 0) + with debug_tensor_interceptors.debug_telemetry_scope(cfg_interval, step=0): + out0 = wrapped(x) + np.testing.assert_array_equal(out0, x * 3.0) + + # Step 1: should be active scope, but filtered out by step interval + with debug_tensor_interceptors.debug_telemetry_scope(cfg_interval, step=1): + out1 = wrapped(x) + np.testing.assert_array_equal(out1, x * 3.0) + + # Outside scope: should be inactive + out_out = wrapped(x) + np.testing.assert_array_equal(out_out, x * 3.0) + + def test_linen_interceptor_auxiliary_outputs(self): + class MoeLinen(nn.Module): + """Mock Linen MoE module for testing.""" + + @nn.compact + def __call__(self, x): + return x * 2.0, jnp.array(0.01, dtype=jnp.float32) + + class AttnLinen(nn.Module): + """Mock Linen attention module for testing.""" + + @nn.compact + def __call__(self, x): + return x * 1.5, jnp.zeros((2, 4), dtype=jnp.float32) + + cfg = DummyConfig(debug_tensor_distribution=True) + moe_mod = MoeLinen(name="moe_layer") + attn_mod = AttnLinen(name="self_attention") + + x = jnp.ones((2, 4), dtype=jnp.float32) + with debug_tensor_interceptors.debug_telemetry_scope(cfg, step=0): + moe_out, _ = moe_mod.init_with_output(jax.random.PRNGKey(0), x) + attn_out, _ = attn_mod.init_with_output(jax.random.PRNGKey(1), x) + + np.testing.assert_array_equal(moe_out[0], x * 2.0) + self.assertAlmostEqual(float(moe_out[1]), 0.01, places=5) + np.testing.assert_array_equal(attn_out[0], x * 1.5) + np.testing.assert_array_equal(attn_out[1], jnp.zeros((2, 4))) + + def test_debug_tensor_moe_routing_expert_weights(self): + # Shape: (batch=2, seq=4, num_experts=8) + np.random.seed(42) + weights = np.random.uniform(0.0, 1.0, size=(2, 4, 8)).astype(np.float32) + x = jnp.array(weights) + cfg_enabled = DummyConfig(debug_tensor_distribution=True) + + out = debug_tensor_utils.debug_tensor(x, "decoder/layer_0/moe/router_weights", step=0, enabled=cfg_enabled) + np.testing.assert_array_equal(out, x) + + +if __name__ == "__main__": + absltest.main()