# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
+ "#\n",
+ "# See LICENSE for license information.\n",
+ "\n",
+ """"PercentageGreaterThanThreshold Feature support for nvidia-dlframework-inspect"""\n",
+ "\n",
+ "from typing import Dict, Optional\n",
+ "\n",
+ "import torch\n",
+ "\n",
+ "from nvdlfw_inspect.registry import Registry, api_method\n",
+ "from nvdlfw_inspect.logging import MetricLogger\n",
+ "import nvdlfw_inspect.api as debug_api\n",
+ "\n",
+ "from transformer_engine.debug.features.api import TEConfigAPIMapper\n",
+ "from transformer_engine.pytorch.tensor import QuantizedTensor, Quantizer\n",
+ "\n",
+ "\n",
+ "# Class should inherit from TEConfigAPIMapper and be registered to the transformer_engine namespace.\n",
+ "@Registry.register_feature(namespace="transformer_engine")\n",
+ "class PercentageGreaterThanThreshold(TEConfigAPIMapper):\n",
+ "\n",
+ " @api_method\n",
+ " def inspect_tensor(\n",
+ " self,\n",
+ " config: Dict,\n",
+ " layer_name: str,\n",
+ " tensor_name: str,\n",
+ " iteration: int,\n",
+ " tp_group: torch.distributed.ProcessGroup,\n",
+ " tensor: torch.Tensor,\n",
+ " rowwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None,\n",
+ " columnwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None,\n",
+ " quantizer: Optional[Quantizer] = None,\n",
+ " ):\n",
+ " # API call inspect_tensor is used to gather the data about the tensor.\n",
+ " # All API calls are documented in the `Precision debug tools / API / Calls to Nvidia-DL-Framework-Inspect`\n",
+ " # section of the documentation.\n",
+ "\n",
+ " threshold = config["threshold"]\n",
+ "\n",
+ " # Get the reduction group from the debug tool\n",
+ " # one can set it using debug_api.set_tensor_reduction_group(group)\n",
+ " reduction_group = debug_api.get_tensor_reduction_group()\n",
+ "\n",
+ " # Compute percentage on local tensor\n",
+ " count = (tensor > threshold).sum().float()\n",
+ " total = torch.tensor(tensor.numel(), dtype=torch.float32, device=tensor.device)\n",
+ "\n",
+ " # Perform reduction across the group if needed.\n",
+ " # Note that we perform all_reduce twice per every tensor, which is suboptimal.\n",
+ " # For guidance on implementing efficient statistics reduction, see the implementation in the `LogTensorStats` feature.\n",
+ " # In this tutorial we only showcase basic implementation of the feature.\n",
+ " if reduction_group is not None:\n",
+ " torch.distributed.all_reduce(count, group=reduction_group)\n",
+ " torch.distributed.all_reduce(total, group=reduction_group)\n",
+ "\n",
+ " percentage = count / total\n",
+ "\n",
+ " # MetricLogger is a class from nvidia-dlframework-inspect.\n",
+ " # By using it we can also use functionalities provided by nvidia-dlframework-inspect,\n",
+ " # like logging to TensorBoard, etc.\n",
+ " MetricLogger.log_scalar(\n",
+ " f"{layer_name}_{tensor_name}_percentage_greater_than_threshold", percentage, iteration\n",
+ " )\n",
+ "\n",
+ " @api_method\n",
+ " def inspect_tensor_enabled(\n",
+ " self, config: Dict, layer_name: str, tensor_name: str, iteration: int\n",
+ " ):\n",
+ " # This call is used by TE to determine if the unfused debug layer - which is slower - needs to be run.\n",
+ " # It returns a tuple (bool, int), where the int indicates the next iteration when the feature will be enabled\n",
+ " # and bool indicates if the feature should be enabled at the current iteration.\n",
+ "\n",
+ " run_current = iteration % config["freq"] == 0\n",
+ " # run in next multiple of freq\n",
+ " next_iter = iteration + (config["freq"] - iteration % config["freq"])\n",
+ " return run_current, next_iter\n",
+ "stats:\n",
+ " enabled: True\n",
+ " layers:\n",
+ " layer_name_regex_pattern: .*\n",
+ " transformer_engine:\n",
+ " PercentageGreaterThanThreshold:\n",
+ " enabled: True\n",
+ " tensors: [activation]\n",
+ " threshold: 0.1\n",
+ " freq: 5\n",
+ " LogTensorStats:\n",
+ " enabled: True\n",
+ " tensors: [activation]\n",
+ " stats: [min]\n",
+ " freq: 5\n",
+ "