{ "cells": [ { "cell_type": "markdown", "id": "85a94734", "metadata": {}, "source": "# Adding custom feature to precision debug tools\n \nTE comes with several built-in features, such as `LogFp8TensorStats`, which can log statistics for each tensor involved in matrix multiplication (GEMM) operations.\nIn this tutorial, we'll demonstrate how to extend TE by adding a custom feature. This custom feature will log the percentage of elements in a tensor whose absolute values exceed a configurable threshold `t`, as specified in the config file.\n\nCustom features can be used for example for:\n\n1. Logging custom statistics.\n2. Dumping intermediate tensors.\n3. Experiments with modifying intermediate tensors.\n\nHow to add custom feature:\n\n1. Add Python with feature class definition which inherits from `transformer_engine.debug.features.api.TEConfigAPIMapper`.\n2. Wrap the class with `@Registry.register_feature(namespace=\"transformer_engine\")`.\n3. Implement some of API calls to nvidia-dl-framework-inspect described [here](3_api_te_calls.rst).\n\nLet's define a new file at `.../custom_feature_dir/percentage_greater_than_threshold.py` containing the following code:\n" }, { "cell_type": "code", "execution_count": 1, "id": "b4e7562d", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
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",
"