tensorboard-plugin-profile 2.18.0__tar.gz → 2.19.0__tar.gz

tensorboard_plugin_profile-2.19.0/PKG-INFO

@@ -0,0 +1,110 @@
+Metadata-Version: 2.1
+Name: tensorboard_plugin_profile
+Version: 2.19.0
+Summary: Profile Tensorboard Plugin
+Home-page: https://github.com/tensorflow/profiler
+Author: Google Inc.
+Author-email: packages@tensorflow.org
+License: Apache 2.0
+Keywords: jax pytorch xla tensorflow tensorboard xprof profile plugin
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 2
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: Software Development :: Libraries
+Requires-Python: >= 2.7, != 3.0.*, != 3.1.*
+Description-Content-Type: text/markdown
+
+# Tensorboard Profiler Plugin
+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.
+
+The profiler plugin offers a number of tools to analyse and visualize the
+performance of your model across multiple devices. Some of the tools include:
+
+*   **Overview**: A high-level overview of the performance of your model. This
+    is an aggregated overview for your host and all devices. It includes:
+    *   Performance summary and breakdown of step times.
+    *   A graph of individual step times.
+    *   A table of the top 10 most expensive operations.
+*   **Trace Viewer**: Displays a timeline of the execution of your model that shows:
+    *   The duration of each op.
+    *   Which part of the system (host or device) executed an op.
+    *   The communication between devices.
+*   **Memory Profile Viewer**: Monitors the memory usage of your model.
+*   **Graph Viewer**: A visualization of the graph structure of HLOs of your model.
+
+## Demo
+First time user? Come and check out this [Colab Demo](https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras).
+
+## Prerequisites
+* TensorFlow >= 2.18.0
+* TensorBoard >= 2.18.0
+* tensorboard-plugin-profile >= 2.18.0
+
+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).
+Some charts and tables may be missing if you run TensorBoard entirely offline on
+your local machine, behind a corporate firewall, or in a datacenter.
+
+To profile on a **single GPU** system, the following NVIDIA software must be installed on your system:
+
+1. NVIDIA GPU drivers and CUDA Toolkit:
+    * CUDA 12.5 requires 525.60.13 and higher.
+2. Ensure that CUPTI 10.1 exists on the path.
+
+   ```shell
+   $ /sbin/ldconfig -N -v $(sed 's/:/ /g' <<< $LD_LIBRARY_PATH) | grep libcupti
+   ```
+
+   If you don't see `libcupti.so.12.5` on the path, prepend its installation directory to the $LD_LIBRARY_PATH environmental variable:
+
+   ```shell
+   $ export LD_LIBRARY_PATH=/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
+   ```
+   Run the ldconfig command above again to verify that the CUPTI 12.5 library is found.
+
+   If this doesn't work, try:
+   ```shell
+   $ sudo apt-get install libcupti-dev
+   ```
+
+To profile a system with **multiple GPUs**, see this [guide](docs/profile_multi_gpu.md) for details.
+
+To profile multi-worker GPU configurations, profile individual workers independently.
+
+To profile cloud TPUs, you must have access to Google Cloud TPUs.
+
+## Quick Start
+The profiler plugin follows the TensorFlow versioning scheme. As a result, the
+`tensorboard-plugin-profile` PyPI package can be behind the `tbp-nightly` PyPI
+package. In order to get the latest version of the profiler plugin, you can
+install the nightly package..
+
+To install the nightly version of profiler:
+
+```
+$ pip uninstall tensorboard-plugin-profile
+$ pip install tbp-nightly
+```
+
+Run TensorBoard:
+
+```
+$ tensorboard --logdir=profiler/demo
+```
+If you are behind a corporate firewall, you may need to include the `--bind_all`
+tensorboard flag.
+
+Go to `localhost:6006/#profile` of your browser, you should now see the demo overview page show up.
+![Overview Page](docs/images/overview_page.png)
+Congratulations! You're now ready to capture a profile.
+
+## Next Steps
+* JAX Profiling Guide: https://jax.readthedocs.io/en/latest/profiling.html
+* TensorFlow Profiling Guide:  https://tensorflow.org/guide/profiler
+* Cloud TPU Profiling Guide: https://cloud.google.com/tpu/docs/cloud-tpu-tools
+* Colab Tutorial: https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras
+* MiniGPT Example: https://docs.jaxstack.ai/en/latest/JAX_for_LLM_pretraining.html

tensorboard_plugin_profile-2.19.0/README.md

@@ -0,0 +1,89 @@
+# Tensorboard Profiler Plugin
+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.
+
+The profiler plugin offers a number of tools to analyse and visualize the
+performance of your model across multiple devices. Some of the tools include:
+
+*   **Overview**: A high-level overview of the performance of your model. This
+    is an aggregated overview for your host and all devices. It includes:
+    *   Performance summary and breakdown of step times.
+    *   A graph of individual step times.
+    *   A table of the top 10 most expensive operations.
+*   **Trace Viewer**: Displays a timeline of the execution of your model that shows:
+    *   The duration of each op.
+    *   Which part of the system (host or device) executed an op.
+    *   The communication between devices.
+*   **Memory Profile Viewer**: Monitors the memory usage of your model.
+*   **Graph Viewer**: A visualization of the graph structure of HLOs of your model.
+
+## Demo
+First time user? Come and check out this [Colab Demo](https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras).
+
+## Prerequisites
+* TensorFlow >= 2.18.0
+* TensorBoard >= 2.18.0
+* tensorboard-plugin-profile >= 2.18.0
+
+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).
+Some charts and tables may be missing if you run TensorBoard entirely offline on
+your local machine, behind a corporate firewall, or in a datacenter.
+
+To profile on a **single GPU** system, the following NVIDIA software must be installed on your system:
+
+1. NVIDIA GPU drivers and CUDA Toolkit:
+    * CUDA 12.5 requires 525.60.13 and higher.
+2. Ensure that CUPTI 10.1 exists on the path.
+
+   ```shell
+   $ /sbin/ldconfig -N -v $(sed 's/:/ /g' <<< $LD_LIBRARY_PATH) | grep libcupti
+   ```
+
+   If you don't see `libcupti.so.12.5` on the path, prepend its installation directory to the $LD_LIBRARY_PATH environmental variable:
+
+   ```shell
+   $ export LD_LIBRARY_PATH=/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
+   ```
+   Run the ldconfig command above again to verify that the CUPTI 12.5 library is found.
+
+   If this doesn't work, try:
+   ```shell
+   $ sudo apt-get install libcupti-dev
+   ```
+
+To profile a system with **multiple GPUs**, see this [guide](docs/profile_multi_gpu.md) for details.
+
+To profile multi-worker GPU configurations, profile individual workers independently.
+
+To profile cloud TPUs, you must have access to Google Cloud TPUs.
+
+## Quick Start
+The profiler plugin follows the TensorFlow versioning scheme. As a result, the
+`tensorboard-plugin-profile` PyPI package can be behind the `tbp-nightly` PyPI
+package. In order to get the latest version of the profiler plugin, you can
+install the nightly package..
+
+To install the nightly version of profiler:
+
+```
+$ pip uninstall tensorboard-plugin-profile
+$ pip install tbp-nightly
+```
+
+Run TensorBoard:
+
+```
+$ tensorboard --logdir=profiler/demo
+```
+If you are behind a corporate firewall, you may need to include the `--bind_all`
+tensorboard flag.
+
+Go to `localhost:6006/#profile` of your browser, you should now see the demo overview page show up.
+![Overview Page](docs/images/overview_page.png)
+Congratulations! You're now ready to capture a profile.
+
+## Next Steps
+* JAX Profiling Guide: https://jax.readthedocs.io/en/latest/profiling.html
+* TensorFlow Profiling Guide:  https://tensorflow.org/guide/profiler
+* Cloud TPU Profiling Guide: https://cloud.google.com/tpu/docs/cloud-tpu-tools
+* Colab Tutorial: https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras
+* MiniGPT Example: https://docs.jaxstack.ai/en/latest/JAX_for_LLM_pretraining.html

{tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/setup.py

@@ -30,11 +30,12 @@ REQUIRED_PACKAGES = [
     'setuptools >= 41.0.0',
     'six >= 1.10.0',
     'werkzeug >= 0.11.15',
+    'etils[epath] >= 1.0.0'
 ]
 
 
 def get_readme():
-  with open('README.rst') as f:
+  with open('README.md') as f:
     return f.read()
 
 
@@ -43,6 +44,7 @@ setuptools.setup(
     version=VERSION,
     description='Profile Tensorboard Plugin',
     long_description=get_readme(),
+    long_description_content_type='text/markdown',
     author='Google Inc.',
     author_email='packages@tensorflow.org',
     url='https://github.com/tensorflow/profiler',
@@ -53,6 +55,7 @@ setuptools.setup(
     ),
     package_data={
         'tensorboard_plugin_profile': ['static/**'],
+        '': ['_pywrap_profiler_plugin.so'],
     },
     entry_points={
         'tensorboard_plugins': [
@@ -78,5 +81,5 @@ setuptools.setup(
         'Topic :: Software Development :: Libraries',
     ],
     license='Apache 2.0',
-    keywords='tensorflow tensorboard xprof profile plugin',
+    keywords='jax pytorch xla tensorflow tensorboard xprof profile plugin',
 )

tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/convert/hlo_stats_proto_to_gviz.py

@@ -0,0 +1,162 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+# ==============================================================================
+"""For conversion of HloStats protos to GViz DataTables.
+
+Usage:
+    gviz_data_tables = generate_hlo_stats_table(hlo_stats_db)
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from typing import Any
+
+import gviz_api
+
+from tensorboard_plugin_profile.protobuf import hlo_stats_pb2
+
+
+def get_hlo_op_name_from_expression(hlo_op_expression: str) -> str:
+  """Returns the hlo op name from the hlo op expression.
+
+  Args:
+    hlo_op_expression: A string of hlo op expression.
+
+  Returns:
+    Returns the hlo op name extracted from the hlo op expression.
+  """
+  # The parse logic based on the assumption that the hlo op text is in format of
+  # '%op_name = <long name>'
+  parts = hlo_op_expression.split(" = ")
+  hlo_op_name = parts[0]
+  if hlo_op_name and hlo_op_name[0] == "%":
+    hlo_op_name = hlo_op_name[1:]
+  return hlo_op_name
+
+
+TableColumnDescription = tuple[str, str, str]
+TableRow = list[Any]
+TableProperties = dict[str, Any]
+
+
+def get_hlo_stats_table_args(
+    hlo_stats_db: hlo_stats_pb2.HloStatsDatabase,
+) -> tuple[list[TableColumnDescription], list[TableRow], TableProperties]:
+  """Creates hlo op stats table from a hlo stats proto.
+
+  Args:
+    hlo_stats_db: A HloStatsDatabase proto.
+
+  Returns:
+    Returns table description(column defubutuibs), data(rows data) and custom
+    properties for preparing the HloStats gviz table.
+  """
+
+  table_description = [
+      ("rank", "number", "Rank"),
+      ("program_id", "string", "Program id"),
+      ("category", "string", "HLO op category"),
+      ("hlo_op_name", "string", "HLO op name"),
+      ("hlo_op_expression", "string", "HLO op text"),
+      ("tf_op_name", "string", "Framework op name"),
+      ("occurrences", "number", "#Occurrences"),
+      ("total_time", "number", "Total time (us)"),
+      ("avg_time", "number", "Avg. time (us)"),
+      ("total_self_time", "number", "Total self time (us)"),
+      ("avg_self_time", "number", "Avg. self time (us)"),
+      ("total_self_time_percent", "number", "Total self time (%)"),
+      (
+          "cumulative_total_self_time_percent",
+          "number",
+          "Cumulative total self time (%)",
+      ),
+      ("dma_stall_percent", "number", "%time stalled by DMA"),
+      ("model_flop_rate", "number", "Model GFLOP/s"),
+      ("normalized_flop_rate", "number", "Normalized GFLOP/s"),
+      ("measured_memory_bw", "number", "Measured memory BW (GiB/s)"),
+      ("hbm_bw", "number", "HBM BW (GiB/s)"),
+      ("cmem_read_bw", "number", "CMEM Read BW (GiB/s)"),
+      ("cmem_write_bw", "number", "CMEM Write BW (GiB/s)"),
+      ("operational_intensity", "number", "Operational intensity (FLOPS/Byte)"),
+      ("bound_by", "string", "Bound by"),
+      ("hlo_rematerialization", "string", "Rematerialization"),
+      ("outside_compilation", "string", "Outside Compilation"),
+      ("autotuned", "string", "Autotuned"),
+  ]
+
+  data = []
+  for record in hlo_stats_db.hlo_stats_record:
+    row = [
+        record.rank,
+        str(record.program_id),
+        record.hlo_category,
+        get_hlo_op_name_from_expression(record.hlo_expression),
+        record.hlo_expression,
+        record.tf_op_name,
+        record.occurrences,
+        record.total_time_in_us,
+        record.avg_time_in_us,
+        record.total_self_time_in_us,
+        record.avg_self_time_in_us,
+        record.total_self_time_as_fraction,
+        record.cumulative_total_self_time_as_fraction,
+        record.dma_stall_fraction,
+        record.model_flop_rate,
+        record.measured_flop_rate,
+        record.measured_memory_bw,
+        record.hbm_bw,
+        record.cmem_read_bw,
+        record.cmem_write_bw,
+        record.operational_intensity,
+        record.bound_by,
+        "Yes" if record.rematerialization else "No",
+        "Yes" if record.outside_compilation else "No",
+        "Yes" if record.autotuned else "No",
+    ]
+    data.append(row)
+
+  return (table_description, data, {})
+
+
+def generate_hlo_stats_table(
+    hlo_stats_db: hlo_stats_pb2.HloStatsDatabase,
+) -> gviz_api.DataTable:
+  """Converts a HloStatsDb proto to a GViz DataTable.
+
+  Args:
+    hlo_stats_db: A HloStatsDatabase proto.
+
+  Returns:
+    Returns a GViz DataTable for HloStats.
+  """
+  (table_description, data, custom_properties) = get_hlo_stats_table_args(
+      hlo_stats_db
+  )
+  return gviz_api.DataTable(table_description, data, custom_properties)
+
+
+def to_json(raw_data: bytes) -> str:
+  """Converts a serialized HloStatsDb bytes to json.
+
+  Args:
+    raw_data: Bytes of serialized HloStatsDb.
+
+  Returns:
+    Returns a json string of HloStats Table.
+  """
+  hlo_stats_db = hlo_stats_pb2.HloStatsDatabase()
+  hlo_stats_db.ParseFromString(raw_data)
+  return generate_hlo_stats_table(hlo_stats_db).ToJSon()