tensorboard-plugin-profile 2.18.0__tar.gz → 2.19.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- tensorboard_plugin_profile-2.19.0/PKG-INFO +110 -0
- tensorboard_plugin_profile-2.19.0/README.md +89 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/setup.py +5 -2
- tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/convert/hlo_stats_proto_to_gviz.py +162 -0
- tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/convert/inference_stats_proto_to_gviz.py +327 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/input_pipeline_proto_to_gviz.py +238 -10
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/overview_page_proto_to_gviz.py +44 -3
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/raw_to_tool_data.py +24 -4
- tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/convert/roofline_model_proto_to_gviz.py +392 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/trace_events_json.py +0 -1
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/integration_tests/tpu/tensorflow/tpu_tf2_keras_test.py +3 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/profile_plugin.py +60 -26
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/profile_plugin_loader.py +8 -3
- tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/protobuf/hlo_stats_pb2.py +27 -0
- tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/protobuf/inference_stats_pb2.py +75 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/overview_page_pb2.py +7 -3
- tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/protobuf/roofline_model_pb2.py +30 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/tf_stats_pb2.py +2 -2
- tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/protobuf/tpu_input_pipeline_pb2.py +32 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/bundle.js +3816 -682
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/trace_viewer_index.html +46 -1
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/trace_viewer_index.js +258 -26
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/version.py +1 -1
- tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile.egg-info/PKG-INFO +110 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile.egg-info/SOURCES.txt +8 -1
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile.egg-info/requires.txt +1 -0
- tensorboard_plugin_profile-2.18.0/PKG-INFO +0 -21
- tensorboard_plugin_profile-2.18.0/README.rst +0 -1
- tensorboard_plugin_profile-2.18.0/tensorboard_plugin_profile.egg-info/PKG-INFO +0 -21
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/setup.cfg +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/__init__.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/dcn_collective_stats_proto_to_gviz.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/diagnostics.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/kernel_stats_proto_to_gviz.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/tf_data_stats_proto_to_gviz.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/tf_stats_proto_to_gviz.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/integration_tests/tf_mnist.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/integration_tests/tf_profiler_session.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/dcn_slack_analysis_pb2.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/diagnostics_pb2.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/input_pipeline_pb2.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/kernel_stats_pb2.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/power_metrics_pb2.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/tf_data_stats_pb2.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/trace_events_pb2.py +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/index.html +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/index.js +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/materialicons.woff2 +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/styles.css +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/zone.js +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile.egg-info/dependency_links.txt +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile.egg-info/entry_points.txt +0 -0
- {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile.egg-info/top_level.txt +0 -0
|
@@ -0,0 +1,110 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: tensorboard_plugin_profile
|
|
3
|
+
Version: 2.19.0
|
|
4
|
+
Summary: Profile Tensorboard Plugin
|
|
5
|
+
Home-page: https://github.com/tensorflow/profiler
|
|
6
|
+
Author: Google Inc.
|
|
7
|
+
Author-email: packages@tensorflow.org
|
|
8
|
+
License: Apache 2.0
|
|
9
|
+
Keywords: jax pytorch xla tensorflow tensorboard xprof profile plugin
|
|
10
|
+
Classifier: Intended Audience :: Developers
|
|
11
|
+
Classifier: Intended Audience :: Education
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
14
|
+
Classifier: Programming Language :: Python :: 2
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Topic :: Scientific/Engineering :: Mathematics
|
|
17
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
18
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
19
|
+
Requires-Python: >= 2.7, != 3.0.*, != 3.1.*
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# Tensorboard Profiler Plugin
|
|
23
|
+
The profiler includes a suite of tools for [JAX](https://jax.readthedocs.io/), [TensorFlow](https://www.tensorflow.org/), and [PyTorch/XLA](https://github.com/pytorch/xla). These tools help you understand, debug and optimize programs to run on CPUs, GPUs and TPUs.
|
|
24
|
+
|
|
25
|
+
The profiler plugin offers a number of tools to analyse and visualize the
|
|
26
|
+
performance of your model across multiple devices. Some of the tools include:
|
|
27
|
+
|
|
28
|
+
* **Overview**: A high-level overview of the performance of your model. This
|
|
29
|
+
is an aggregated overview for your host and all devices. It includes:
|
|
30
|
+
* Performance summary and breakdown of step times.
|
|
31
|
+
* A graph of individual step times.
|
|
32
|
+
* A table of the top 10 most expensive operations.
|
|
33
|
+
* **Trace Viewer**: Displays a timeline of the execution of your model that shows:
|
|
34
|
+
* The duration of each op.
|
|
35
|
+
* Which part of the system (host or device) executed an op.
|
|
36
|
+
* The communication between devices.
|
|
37
|
+
* **Memory Profile Viewer**: Monitors the memory usage of your model.
|
|
38
|
+
* **Graph Viewer**: A visualization of the graph structure of HLOs of your model.
|
|
39
|
+
|
|
40
|
+
## Demo
|
|
41
|
+
First time user? Come and check out this [Colab Demo](https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras).
|
|
42
|
+
|
|
43
|
+
## Prerequisites
|
|
44
|
+
* TensorFlow >= 2.18.0
|
|
45
|
+
* TensorBoard >= 2.18.0
|
|
46
|
+
* tensorboard-plugin-profile >= 2.18.0
|
|
47
|
+
|
|
48
|
+
Note: The Tensorboard Profiler Plugin requires access to the Internet to load the [Google Chart library](https://developers.google.com/chart/interactive/docs/basic_load_libs#basic-library-loading).
|
|
49
|
+
Some charts and tables may be missing if you run TensorBoard entirely offline on
|
|
50
|
+
your local machine, behind a corporate firewall, or in a datacenter.
|
|
51
|
+
|
|
52
|
+
To profile on a **single GPU** system, the following NVIDIA software must be installed on your system:
|
|
53
|
+
|
|
54
|
+
1. NVIDIA GPU drivers and CUDA Toolkit:
|
|
55
|
+
* CUDA 12.5 requires 525.60.13 and higher.
|
|
56
|
+
2. Ensure that CUPTI 10.1 exists on the path.
|
|
57
|
+
|
|
58
|
+
```shell
|
|
59
|
+
$ /sbin/ldconfig -N -v $(sed 's/:/ /g' <<< $LD_LIBRARY_PATH) | grep libcupti
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
If you don't see `libcupti.so.12.5` on the path, prepend its installation directory to the $LD_LIBRARY_PATH environmental variable:
|
|
63
|
+
|
|
64
|
+
```shell
|
|
65
|
+
$ export LD_LIBRARY_PATH=/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
|
|
66
|
+
```
|
|
67
|
+
Run the ldconfig command above again to verify that the CUPTI 12.5 library is found.
|
|
68
|
+
|
|
69
|
+
If this doesn't work, try:
|
|
70
|
+
```shell
|
|
71
|
+
$ sudo apt-get install libcupti-dev
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
To profile a system with **multiple GPUs**, see this [guide](docs/profile_multi_gpu.md) for details.
|
|
75
|
+
|
|
76
|
+
To profile multi-worker GPU configurations, profile individual workers independently.
|
|
77
|
+
|
|
78
|
+
To profile cloud TPUs, you must have access to Google Cloud TPUs.
|
|
79
|
+
|
|
80
|
+
## Quick Start
|
|
81
|
+
The profiler plugin follows the TensorFlow versioning scheme. As a result, the
|
|
82
|
+
`tensorboard-plugin-profile` PyPI package can be behind the `tbp-nightly` PyPI
|
|
83
|
+
package. In order to get the latest version of the profiler plugin, you can
|
|
84
|
+
install the nightly package..
|
|
85
|
+
|
|
86
|
+
To install the nightly version of profiler:
|
|
87
|
+
|
|
88
|
+
```
|
|
89
|
+
$ pip uninstall tensorboard-plugin-profile
|
|
90
|
+
$ pip install tbp-nightly
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
Run TensorBoard:
|
|
94
|
+
|
|
95
|
+
```
|
|
96
|
+
$ tensorboard --logdir=profiler/demo
|
|
97
|
+
```
|
|
98
|
+
If you are behind a corporate firewall, you may need to include the `--bind_all`
|
|
99
|
+
tensorboard flag.
|
|
100
|
+
|
|
101
|
+
Go to `localhost:6006/#profile` of your browser, you should now see the demo overview page show up.
|
|
102
|
+

|
|
103
|
+
Congratulations! You're now ready to capture a profile.
|
|
104
|
+
|
|
105
|
+
## Next Steps
|
|
106
|
+
* JAX Profiling Guide: https://jax.readthedocs.io/en/latest/profiling.html
|
|
107
|
+
* TensorFlow Profiling Guide: https://tensorflow.org/guide/profiler
|
|
108
|
+
* Cloud TPU Profiling Guide: https://cloud.google.com/tpu/docs/cloud-tpu-tools
|
|
109
|
+
* Colab Tutorial: https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras
|
|
110
|
+
* MiniGPT Example: https://docs.jaxstack.ai/en/latest/JAX_for_LLM_pretraining.html
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
# Tensorboard Profiler Plugin
|
|
2
|
+
The profiler includes a suite of tools for [JAX](https://jax.readthedocs.io/), [TensorFlow](https://www.tensorflow.org/), and [PyTorch/XLA](https://github.com/pytorch/xla). These tools help you understand, debug and optimize programs to run on CPUs, GPUs and TPUs.
|
|
3
|
+
|
|
4
|
+
The profiler plugin offers a number of tools to analyse and visualize the
|
|
5
|
+
performance of your model across multiple devices. Some of the tools include:
|
|
6
|
+
|
|
7
|
+
* **Overview**: A high-level overview of the performance of your model. This
|
|
8
|
+
is an aggregated overview for your host and all devices. It includes:
|
|
9
|
+
* Performance summary and breakdown of step times.
|
|
10
|
+
* A graph of individual step times.
|
|
11
|
+
* A table of the top 10 most expensive operations.
|
|
12
|
+
* **Trace Viewer**: Displays a timeline of the execution of your model that shows:
|
|
13
|
+
* The duration of each op.
|
|
14
|
+
* Which part of the system (host or device) executed an op.
|
|
15
|
+
* The communication between devices.
|
|
16
|
+
* **Memory Profile Viewer**: Monitors the memory usage of your model.
|
|
17
|
+
* **Graph Viewer**: A visualization of the graph structure of HLOs of your model.
|
|
18
|
+
|
|
19
|
+
## Demo
|
|
20
|
+
First time user? Come and check out this [Colab Demo](https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras).
|
|
21
|
+
|
|
22
|
+
## Prerequisites
|
|
23
|
+
* TensorFlow >= 2.18.0
|
|
24
|
+
* TensorBoard >= 2.18.0
|
|
25
|
+
* tensorboard-plugin-profile >= 2.18.0
|
|
26
|
+
|
|
27
|
+
Note: The Tensorboard Profiler Plugin requires access to the Internet to load the [Google Chart library](https://developers.google.com/chart/interactive/docs/basic_load_libs#basic-library-loading).
|
|
28
|
+
Some charts and tables may be missing if you run TensorBoard entirely offline on
|
|
29
|
+
your local machine, behind a corporate firewall, or in a datacenter.
|
|
30
|
+
|
|
31
|
+
To profile on a **single GPU** system, the following NVIDIA software must be installed on your system:
|
|
32
|
+
|
|
33
|
+
1. NVIDIA GPU drivers and CUDA Toolkit:
|
|
34
|
+
* CUDA 12.5 requires 525.60.13 and higher.
|
|
35
|
+
2. Ensure that CUPTI 10.1 exists on the path.
|
|
36
|
+
|
|
37
|
+
```shell
|
|
38
|
+
$ /sbin/ldconfig -N -v $(sed 's/:/ /g' <<< $LD_LIBRARY_PATH) | grep libcupti
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
If you don't see `libcupti.so.12.5` on the path, prepend its installation directory to the $LD_LIBRARY_PATH environmental variable:
|
|
42
|
+
|
|
43
|
+
```shell
|
|
44
|
+
$ export LD_LIBRARY_PATH=/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
|
|
45
|
+
```
|
|
46
|
+
Run the ldconfig command above again to verify that the CUPTI 12.5 library is found.
|
|
47
|
+
|
|
48
|
+
If this doesn't work, try:
|
|
49
|
+
```shell
|
|
50
|
+
$ sudo apt-get install libcupti-dev
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
To profile a system with **multiple GPUs**, see this [guide](docs/profile_multi_gpu.md) for details.
|
|
54
|
+
|
|
55
|
+
To profile multi-worker GPU configurations, profile individual workers independently.
|
|
56
|
+
|
|
57
|
+
To profile cloud TPUs, you must have access to Google Cloud TPUs.
|
|
58
|
+
|
|
59
|
+
## Quick Start
|
|
60
|
+
The profiler plugin follows the TensorFlow versioning scheme. As a result, the
|
|
61
|
+
`tensorboard-plugin-profile` PyPI package can be behind the `tbp-nightly` PyPI
|
|
62
|
+
package. In order to get the latest version of the profiler plugin, you can
|
|
63
|
+
install the nightly package..
|
|
64
|
+
|
|
65
|
+
To install the nightly version of profiler:
|
|
66
|
+
|
|
67
|
+
```
|
|
68
|
+
$ pip uninstall tensorboard-plugin-profile
|
|
69
|
+
$ pip install tbp-nightly
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
Run TensorBoard:
|
|
73
|
+
|
|
74
|
+
```
|
|
75
|
+
$ tensorboard --logdir=profiler/demo
|
|
76
|
+
```
|
|
77
|
+
If you are behind a corporate firewall, you may need to include the `--bind_all`
|
|
78
|
+
tensorboard flag.
|
|
79
|
+
|
|
80
|
+
Go to `localhost:6006/#profile` of your browser, you should now see the demo overview page show up.
|
|
81
|
+

|
|
82
|
+
Congratulations! You're now ready to capture a profile.
|
|
83
|
+
|
|
84
|
+
## Next Steps
|
|
85
|
+
* JAX Profiling Guide: https://jax.readthedocs.io/en/latest/profiling.html
|
|
86
|
+
* TensorFlow Profiling Guide: https://tensorflow.org/guide/profiler
|
|
87
|
+
* Cloud TPU Profiling Guide: https://cloud.google.com/tpu/docs/cloud-tpu-tools
|
|
88
|
+
* Colab Tutorial: https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras
|
|
89
|
+
* MiniGPT Example: https://docs.jaxstack.ai/en/latest/JAX_for_LLM_pretraining.html
|
|
@@ -30,11 +30,12 @@ REQUIRED_PACKAGES = [
|
|
|
30
30
|
'setuptools >= 41.0.0',
|
|
31
31
|
'six >= 1.10.0',
|
|
32
32
|
'werkzeug >= 0.11.15',
|
|
33
|
+
'etils[epath] >= 1.0.0'
|
|
33
34
|
]
|
|
34
35
|
|
|
35
36
|
|
|
36
37
|
def get_readme():
|
|
37
|
-
with open('README.
|
|
38
|
+
with open('README.md') as f:
|
|
38
39
|
return f.read()
|
|
39
40
|
|
|
40
41
|
|
|
@@ -43,6 +44,7 @@ setuptools.setup(
|
|
|
43
44
|
version=VERSION,
|
|
44
45
|
description='Profile Tensorboard Plugin',
|
|
45
46
|
long_description=get_readme(),
|
|
47
|
+
long_description_content_type='text/markdown',
|
|
46
48
|
author='Google Inc.',
|
|
47
49
|
author_email='packages@tensorflow.org',
|
|
48
50
|
url='https://github.com/tensorflow/profiler',
|
|
@@ -53,6 +55,7 @@ setuptools.setup(
|
|
|
53
55
|
),
|
|
54
56
|
package_data={
|
|
55
57
|
'tensorboard_plugin_profile': ['static/**'],
|
|
58
|
+
'': ['_pywrap_profiler_plugin.so'],
|
|
56
59
|
},
|
|
57
60
|
entry_points={
|
|
58
61
|
'tensorboard_plugins': [
|
|
@@ -78,5 +81,5 @@ setuptools.setup(
|
|
|
78
81
|
'Topic :: Software Development :: Libraries',
|
|
79
82
|
],
|
|
80
83
|
license='Apache 2.0',
|
|
81
|
-
keywords='tensorflow tensorboard xprof profile plugin',
|
|
84
|
+
keywords='jax pytorch xla tensorflow tensorboard xprof profile plugin',
|
|
82
85
|
)
|
tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/convert/hlo_stats_proto_to_gviz.py
ADDED
|
@@ -0,0 +1,162 @@
|
|
|
1
|
+
# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
# ==============================================================================
|
|
15
|
+
"""For conversion of HloStats protos to GViz DataTables.
|
|
16
|
+
|
|
17
|
+
Usage:
|
|
18
|
+
gviz_data_tables = generate_hlo_stats_table(hlo_stats_db)
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
from __future__ import absolute_import
|
|
22
|
+
from __future__ import division
|
|
23
|
+
from __future__ import print_function
|
|
24
|
+
|
|
25
|
+
from typing import Any
|
|
26
|
+
|
|
27
|
+
import gviz_api
|
|
28
|
+
|
|
29
|
+
from tensorboard_plugin_profile.protobuf import hlo_stats_pb2
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def get_hlo_op_name_from_expression(hlo_op_expression: str) -> str:
|
|
33
|
+
"""Returns the hlo op name from the hlo op expression.
|
|
34
|
+
|
|
35
|
+
Args:
|
|
36
|
+
hlo_op_expression: A string of hlo op expression.
|
|
37
|
+
|
|
38
|
+
Returns:
|
|
39
|
+
Returns the hlo op name extracted from the hlo op expression.
|
|
40
|
+
"""
|
|
41
|
+
# The parse logic based on the assumption that the hlo op text is in format of
|
|
42
|
+
# '%op_name = <long name>'
|
|
43
|
+
parts = hlo_op_expression.split(" = ")
|
|
44
|
+
hlo_op_name = parts[0]
|
|
45
|
+
if hlo_op_name and hlo_op_name[0] == "%":
|
|
46
|
+
hlo_op_name = hlo_op_name[1:]
|
|
47
|
+
return hlo_op_name
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
TableColumnDescription = tuple[str, str, str]
|
|
51
|
+
TableRow = list[Any]
|
|
52
|
+
TableProperties = dict[str, Any]
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def get_hlo_stats_table_args(
|
|
56
|
+
hlo_stats_db: hlo_stats_pb2.HloStatsDatabase,
|
|
57
|
+
) -> tuple[list[TableColumnDescription], list[TableRow], TableProperties]:
|
|
58
|
+
"""Creates hlo op stats table from a hlo stats proto.
|
|
59
|
+
|
|
60
|
+
Args:
|
|
61
|
+
hlo_stats_db: A HloStatsDatabase proto.
|
|
62
|
+
|
|
63
|
+
Returns:
|
|
64
|
+
Returns table description(column defubutuibs), data(rows data) and custom
|
|
65
|
+
properties for preparing the HloStats gviz table.
|
|
66
|
+
"""
|
|
67
|
+
|
|
68
|
+
table_description = [
|
|
69
|
+
("rank", "number", "Rank"),
|
|
70
|
+
("program_id", "string", "Program id"),
|
|
71
|
+
("category", "string", "HLO op category"),
|
|
72
|
+
("hlo_op_name", "string", "HLO op name"),
|
|
73
|
+
("hlo_op_expression", "string", "HLO op text"),
|
|
74
|
+
("tf_op_name", "string", "Framework op name"),
|
|
75
|
+
("occurrences", "number", "#Occurrences"),
|
|
76
|
+
("total_time", "number", "Total time (us)"),
|
|
77
|
+
("avg_time", "number", "Avg. time (us)"),
|
|
78
|
+
("total_self_time", "number", "Total self time (us)"),
|
|
79
|
+
("avg_self_time", "number", "Avg. self time (us)"),
|
|
80
|
+
("total_self_time_percent", "number", "Total self time (%)"),
|
|
81
|
+
(
|
|
82
|
+
"cumulative_total_self_time_percent",
|
|
83
|
+
"number",
|
|
84
|
+
"Cumulative total self time (%)",
|
|
85
|
+
),
|
|
86
|
+
("dma_stall_percent", "number", "%time stalled by DMA"),
|
|
87
|
+
("model_flop_rate", "number", "Model GFLOP/s"),
|
|
88
|
+
("normalized_flop_rate", "number", "Normalized GFLOP/s"),
|
|
89
|
+
("measured_memory_bw", "number", "Measured memory BW (GiB/s)"),
|
|
90
|
+
("hbm_bw", "number", "HBM BW (GiB/s)"),
|
|
91
|
+
("cmem_read_bw", "number", "CMEM Read BW (GiB/s)"),
|
|
92
|
+
("cmem_write_bw", "number", "CMEM Write BW (GiB/s)"),
|
|
93
|
+
("operational_intensity", "number", "Operational intensity (FLOPS/Byte)"),
|
|
94
|
+
("bound_by", "string", "Bound by"),
|
|
95
|
+
("hlo_rematerialization", "string", "Rematerialization"),
|
|
96
|
+
("outside_compilation", "string", "Outside Compilation"),
|
|
97
|
+
("autotuned", "string", "Autotuned"),
|
|
98
|
+
]
|
|
99
|
+
|
|
100
|
+
data = []
|
|
101
|
+
for record in hlo_stats_db.hlo_stats_record:
|
|
102
|
+
row = [
|
|
103
|
+
record.rank,
|
|
104
|
+
str(record.program_id),
|
|
105
|
+
record.hlo_category,
|
|
106
|
+
get_hlo_op_name_from_expression(record.hlo_expression),
|
|
107
|
+
record.hlo_expression,
|
|
108
|
+
record.tf_op_name,
|
|
109
|
+
record.occurrences,
|
|
110
|
+
record.total_time_in_us,
|
|
111
|
+
record.avg_time_in_us,
|
|
112
|
+
record.total_self_time_in_us,
|
|
113
|
+
record.avg_self_time_in_us,
|
|
114
|
+
record.total_self_time_as_fraction,
|
|
115
|
+
record.cumulative_total_self_time_as_fraction,
|
|
116
|
+
record.dma_stall_fraction,
|
|
117
|
+
record.model_flop_rate,
|
|
118
|
+
record.measured_flop_rate,
|
|
119
|
+
record.measured_memory_bw,
|
|
120
|
+
record.hbm_bw,
|
|
121
|
+
record.cmem_read_bw,
|
|
122
|
+
record.cmem_write_bw,
|
|
123
|
+
record.operational_intensity,
|
|
124
|
+
record.bound_by,
|
|
125
|
+
"Yes" if record.rematerialization else "No",
|
|
126
|
+
"Yes" if record.outside_compilation else "No",
|
|
127
|
+
"Yes" if record.autotuned else "No",
|
|
128
|
+
]
|
|
129
|
+
data.append(row)
|
|
130
|
+
|
|
131
|
+
return (table_description, data, {})
|
|
132
|
+
|
|
133
|
+
|
|
134
|
+
def generate_hlo_stats_table(
|
|
135
|
+
hlo_stats_db: hlo_stats_pb2.HloStatsDatabase,
|
|
136
|
+
) -> gviz_api.DataTable:
|
|
137
|
+
"""Converts a HloStatsDb proto to a GViz DataTable.
|
|
138
|
+
|
|
139
|
+
Args:
|
|
140
|
+
hlo_stats_db: A HloStatsDatabase proto.
|
|
141
|
+
|
|
142
|
+
Returns:
|
|
143
|
+
Returns a GViz DataTable for HloStats.
|
|
144
|
+
"""
|
|
145
|
+
(table_description, data, custom_properties) = get_hlo_stats_table_args(
|
|
146
|
+
hlo_stats_db
|
|
147
|
+
)
|
|
148
|
+
return gviz_api.DataTable(table_description, data, custom_properties)
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
def to_json(raw_data: bytes) -> str:
|
|
152
|
+
"""Converts a serialized HloStatsDb bytes to json.
|
|
153
|
+
|
|
154
|
+
Args:
|
|
155
|
+
raw_data: Bytes of serialized HloStatsDb.
|
|
156
|
+
|
|
157
|
+
Returns:
|
|
158
|
+
Returns a json string of HloStats Table.
|
|
159
|
+
"""
|
|
160
|
+
hlo_stats_db = hlo_stats_pb2.HloStatsDatabase()
|
|
161
|
+
hlo_stats_db.ParseFromString(raw_data)
|
|
162
|
+
return generate_hlo_stats_table(hlo_stats_db).ToJSon()
|