From 8d1fd879902cca7d21c5bc807cd3ca3302ec8038 Mon Sep 17 00:00:00 2001 From: saivishal1999 Date: Tue, 28 Jul 2026 16:18:44 -0700 Subject: [PATCH 1/9] Add Nemotron Nano pytorch profiler TOML and chakra tracing docs - Add conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml: Nemotron 3 Nano 3B (2 nodes, 8 GPUs) with PyTorch kineto profiling via [extra_cmd_args] Hydra overrides at steps 45-50. Validated end-to-end on Lyris GB200: profiler fires, rank-0.json.gz and memory_profile.pickle produced. - Add doc/workloads/chakra_tracing.rst: Documents how to enable Kineto and Chakra tracing in MegatronBridge, how to disable nsys (incompatible with Kineto), and full parameter reference. - Modify megatron_bridge.py: Add _copy_profiler_traces() to was_run_successful(). NemoRun places traces at experiments////torch_profile/ internally; this method copies them to /torch_profile/ after job completion so MegatronBridge and MegatronRun produce traces at a consistent location. --- ...bridge_nemotron_nano_pytorch_profiler.toml | 53 +++++++ doc/workloads/chakra_tracing.rst | 136 ++++++++++++++++++ doc/workloads/index.rst | 1 + .../megatron_bridge/megatron_bridge.py | 30 ++++ 4 files changed, 220 insertions(+) create mode 100644 conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml create mode 100644 doc/workloads/chakra_tracing.rst diff --git a/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml b/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml new file mode 100644 index 000000000..0708fbe1d --- /dev/null +++ b/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +name = "megatron_bridge_nemotron_nano_pytorch_profiler" +description = "Megatron-Bridge Nemotron 3 Nano 3B with PyTorch kineto profiling enabled" +test_template_name = "MegatronBridge" + +extra_container_mounts = [] + +[[git_repos]] +url = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git" +commit = "8a34622655951e44e1847796d9b2a0ce647999ee" +mount_as = "/opt/Megatron-Bridge" +init_submodules = true + +[cmd_args] +gpu_type = "gb200" +container_image = "nvcr.io/nvidia/nemo:26.06.00" +model_family_name = "nemotronh" +model_recipe_name = "nemotron_3_nano" +num_gpus = 8 +domain = "llm" +task = "pretrain" +compute_dtype = "fp8_mx" +hf_token = "" +enable_vboost = true +max_steps = 60 +wandb_experiment_name = "megatron_bridge_nemotron_nano_pytorch_profiler" +tp = 1 +pp = 1 +ep = 8 + +[extra_cmd_args] +"profiling.use_pytorch_profiler" = "true" +"profiling.profile_step_start" = "45" +"profiling.profile_step_end" = "50" +"profiling.profile_ranks" = "[0]" +"profiling.record_shapes" = "true" +"profiling.record_memory_history" = "true" +"profiling.memory_snapshot_path" = "memory_profile.pickle" diff --git a/doc/workloads/chakra_tracing.rst b/doc/workloads/chakra_tracing.rst new file mode 100644 index 000000000..d48bcf7f0 --- /dev/null +++ b/doc/workloads/chakra_tracing.rst @@ -0,0 +1,136 @@ +Chakra / Kineto Tracing with MegatronBridge +============================================ + +MegatronBridge supports two profiling modes that **cannot run simultaneously**: + +- **nsys (Nsight Systems)** — GPU timeline profiling via ``enable_nsys = true`` +- **Kineto / Chakra** — PyTorch kineto trace + optional Chakra trace via ``[extra_cmd_args]`` + +This page explains how to enable Kineto/Chakra tracing and disable nsys. + + +Background +---------- + +Kineto trace (``rank-0.json.gz``) is a Chrome-format trace viewable in tools like +`Perfetto `_ or ``chrome://tracing``. + +Chakra trace is an additional execution graph format generated *on top of* the Kineto +trace. To produce a Chakra trace you must first enable the Kineto trace, then set +``pytorch_profiler_collect_chakra = true``. + + +Disabling nsys +-------------- + +The existing MegatronBridge TOMLs in ``conf/staging/`` have ``enable_nsys = true``. +To switch to Kineto/Chakra profiling, either: + +- **Omit** ``enable_nsys`` from the TOML (it defaults to ``false``) +- **Or** explicitly set ``enable_nsys = false`` + +Both are equivalent. Do not set ``enable_nsys = true`` alongside Kineto flags — the +two profilers conflict and Megatron-Bridge will raise an assertion error. + + +Enabling Kineto Trace +--------------------- + +Add the following to ``[extra_cmd_args]`` in your Test TOML. These are Hydra-style +overrides passed directly to the Megatron-Bridge training script: + +.. code-block:: toml + + [extra_cmd_args] + "profiling.use_pytorch_profiler" = "true" + "profiling.profile_step_start" = "45" + "profiling.profile_step_end" = "50" + "profiling.profile_ranks" = "[0]" + "profiling.record_shapes" = "true" + "profiling.record_memory_history" = "true" + "profiling.memory_snapshot_path" = "memory_profile.pickle" + +This produces ``rank-0.json.gz`` (kineto trace) and ``memory_profile.pickle`` +(memory snapshot) under ``/torch_profile/``. + +.. note:: + + CloudAI automatically copies the trace from NemoRun's nested experiment directory + to ``/torch_profile/`` after the job completes, making the location + consistent with MegatronRun. + + +Enabling Chakra Trace +--------------------- + +Add ``pytorch_profiler_collect_chakra = true`` alongside the Kineto flags: + +.. code-block:: toml + + [extra_cmd_args] + "profiling.use_pytorch_profiler" = "true" + "profiling.profile_step_start" = "45" + "profiling.profile_step_end" = "50" + "profiling.profile_ranks" = "[0]" + "profiling.record_shapes" = "true" + "profiling.record_memory_history" = "true" + "profiling.memory_snapshot_path" = "memory_profile.pickle" + "profiling.pytorch_profiler_collect_chakra" = "true" + +Setting ``pytorch_profiler_collect_chakra = false`` disables the Chakra trace while +keeping the Kineto trace (``rank-0.json.gz``) and memory snapshot. + + +Reference TOML +-------------- + +A ready-to-use TOML is provided at: + +.. code-block:: text + + conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml + +This runs Nemotron 3 Nano 3B (2 nodes, 8 GPUs, 60 steps) with Kineto profiling at +steps 45–50 and is suitable for quickly validating the profiling setup. + + +Parameter Reference +------------------- + +.. list-table:: + :header-rows: 1 + :widths: 40 15 45 + + * - Parameter + - Type + - Description + * - ``profiling.use_pytorch_profiler`` + - bool + - Enable PyTorch kineto profiler. Required for all Kineto/Chakra tracing. + * - ``profiling.profile_step_start`` + - int + - Training step at which profiling begins. + * - ``profiling.profile_step_end`` + - int + - Training step at which profiling ends and trace is exported. + * - ``profiling.profile_ranks`` + - list[int] + - GPU ranks to profile (e.g. ``[0]`` for rank 0 only). + * - ``profiling.record_shapes`` + - bool + - Record tensor shapes in the kineto trace. + * - ``profiling.record_memory_history`` + - bool + - Record memory allocation history (produces ``memory_profile.pickle``). + * - ``profiling.memory_snapshot_path`` + - str + - Filename for the memory snapshot pickle file. + * - ``profiling.pytorch_profiler_collect_chakra`` + - bool + - Collect Chakra execution graph trace (requires Kineto enabled). + * - ``profiling.pytorch_profiler_collect_callstack`` + - bool + - Include Python call stack in the kineto trace (increases trace size). + * - ``enable_nsys`` + - bool + - Enable Nsight Systems profiling. **Cannot be used with Kineto.** Default: ``false``. diff --git a/doc/workloads/index.rst b/doc/workloads/index.rst index 916d596cc..64707cea7 100644 --- a/doc/workloads/index.rst +++ b/doc/workloads/index.rst @@ -14,6 +14,7 @@ Available Workloads ":doc:`ai_dynamo`", "✅", "✅", "❌", "❌" ":doc:`bash_cmd`", "✅", "❌", "❌", "❌" ":doc:`chakra_replay`", "✅", "❌", "❌", "❌" + ":doc:`chakra_tracing`", "✅", "❌", "❌", "❌" ":doc:`ddlb`", "✅", "❌", "❌", "❌" ":doc:`deepep`", "✅", "❌", "❌", "❌" ":doc:`dynamo_mocker`", "❌", "❌", "❌", "✅" diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index 7dfd65532..9bbdf6589 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -14,9 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import glob import logging import os import re +import shutil from typing import List, Optional, Union, cast from pydantic import Field, ValidationInfo, field_validator @@ -575,8 +577,36 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: if not step_times_s: return JobStatusResult(is_successful=False, error_message="\n".join(log_data.splitlines()[-40:])) + self._copy_profiler_traces(tr) + return JobStatusResult(is_successful=True) + def _copy_profiler_traces(self, tr: TestRun) -> None: + """ + Copy pytorch kineto traces from NemoRun's nested experiment directory to the CloudAI output root. + + MegatronBridge runs via NemoRun which creates experiments////torch_profile/ internally. + This copies the traces to /torch_profile/ so both MegatronRun and MegatronBridge produce + traces at a consistent location. + """ + experiment_name = tr.test.cmd_args.wandb_experiment_name + if not experiment_name: + return + + pattern = os.path.join( + str(tr.output_path), "experiments", experiment_name, "*", experiment_name, "torch_profile" + ) + matches = [p for p in glob.glob(pattern) if os.path.isdir(p)] + if not matches: + return + + # Take the most recently modified in case of retries + src = max(matches, key=os.path.getmtime) + dst = tr.output_path / "torch_profile" + if not dst.exists(): + shutil.copytree(src, dst) + logging.info(f"Copied pytorch profiler traces from {src} to {dst}") + def extract_mbridge_metrics(logs: str) -> tuple[list[float], list[float]]: step_times_s: list[float] = [] From ef98ae4fa6890d434f8f6996414da4d073c8716b Mon Sep 17 00:00:00 2001 From: saivishal1999 Date: Tue, 28 Jul 2026 16:41:47 -0700 Subject: [PATCH 2/9] Also copy chakra/ traces to output root; update chakra tracing docs - _copy_profiler_traces now copies both torch_profile/ and chakra/ from NemoRun's nested experiment dir to /. chakra/ is only copied if present (produced when pytorch_profiler_collect_chakra=true). - Update chakra_tracing.rst to clarify that kineto trace is always produced and pytorch_profiler_collect_chakra adds a separate Chakra ExecutionTrace. --- doc/workloads/chakra_tracing.rst | 39 +++++-------------- .../megatron_bridge/megatron_bridge.py | 37 +++++++++++------- 2 files changed, 31 insertions(+), 45 deletions(-) diff --git a/doc/workloads/chakra_tracing.rst b/doc/workloads/chakra_tracing.rst index d48bcf7f0..f28d7e54c 100644 --- a/doc/workloads/chakra_tracing.rst +++ b/doc/workloads/chakra_tracing.rst @@ -1,25 +1,6 @@ Chakra / Kineto Tracing with MegatronBridge ============================================ -MegatronBridge supports two profiling modes that **cannot run simultaneously**: - -- **nsys (Nsight Systems)** — GPU timeline profiling via ``enable_nsys = true`` -- **Kineto / Chakra** — PyTorch kineto trace + optional Chakra trace via ``[extra_cmd_args]`` - -This page explains how to enable Kineto/Chakra tracing and disable nsys. - - -Background ----------- - -Kineto trace (``rank-0.json.gz``) is a Chrome-format trace viewable in tools like -`Perfetto `_ or ``chrome://tracing``. - -Chakra trace is an additional execution graph format generated *on top of* the Kineto -trace. To produce a Chakra trace you must first enable the Kineto trace, then set -``pytorch_profiler_collect_chakra = true``. - - Disabling nsys -------------- @@ -50,20 +31,17 @@ overrides passed directly to the Megatron-Bridge training script: "profiling.record_memory_history" = "true" "profiling.memory_snapshot_path" = "memory_profile.pickle" -This produces ``rank-0.json.gz`` (kineto trace) and ``memory_profile.pickle`` -(memory snapshot) under ``/torch_profile/``. - -.. note:: +This produces the following under ``/``: - CloudAI automatically copies the trace from NemoRun's nested experiment directory - to ``/torch_profile/`` after the job completes, making the location - consistent with MegatronRun. +- ``torch_profile/rank-0.json.gz`` — Chrome-format kineto trace (always produced) +- ``memory_profile.pickle`` — memory allocation snapshot (in nested experiment dir) Enabling Chakra Trace --------------------- -Add ``pytorch_profiler_collect_chakra = true`` alongside the Kineto flags: +The Chakra trace is a PyTorch ``ExecutionTraceObserver`` recording — a separate file +from the kineto trace. Add ``pytorch_profiler_collect_chakra = true``: .. code-block:: toml @@ -77,8 +55,9 @@ Add ``pytorch_profiler_collect_chakra = true`` alongside the Kineto flags: "profiling.memory_snapshot_path" = "memory_profile.pickle" "profiling.pytorch_profiler_collect_chakra" = "true" -Setting ``pytorch_profiler_collect_chakra = false`` disables the Chakra trace while -keeping the Kineto trace (``rank-0.json.gz``) and memory snapshot. +This produces an additional ``chakra/rank-0.json.gz`` under ``/``. +The kineto trace (``torch_profile/rank-0.json.gz``) is always produced regardless of +this flag. Reference TOML @@ -127,7 +106,7 @@ Parameter Reference - Filename for the memory snapshot pickle file. * - ``profiling.pytorch_profiler_collect_chakra`` - bool - - Collect Chakra execution graph trace (requires Kineto enabled). + - Collect Chakra ExecutionTraceObserver trace (``chakra/rank-N.json.gz``). The kineto trace is always produced; this adds an additional file. * - ``profiling.pytorch_profiler_collect_callstack`` - bool - Include Python call stack in the kineto trace (increases trace size). diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index 9bbdf6589..7b6426a08 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -583,29 +583,36 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: def _copy_profiler_traces(self, tr: TestRun) -> None: """ - Copy pytorch kineto traces from NemoRun's nested experiment directory to the CloudAI output root. + Copy pytorch kineto and Chakra traces from NemoRun's nested experiment directory to the CloudAI output root. - MegatronBridge runs via NemoRun which creates experiments////torch_profile/ internally. - This copies the traces to /torch_profile/ so both MegatronRun and MegatronBridge produce - traces at a consistent location. + MegatronBridge runs via NemoRun which places profiler output at + experiments////{torch_profile,chakra}/ internally. + This copies them to /{torch_profile,chakra}/ so both MegatronRun and MegatronBridge + produce traces at a consistent location. + + torch_profile/ is always present when use_pytorch_profiler=true. + chakra/ is only present when pytorch_profiler_collect_chakra=true. """ experiment_name = tr.test.cmd_args.wandb_experiment_name if not experiment_name: return - pattern = os.path.join( - str(tr.output_path), "experiments", experiment_name, "*", experiment_name, "torch_profile" - ) - matches = [p for p in glob.glob(pattern) if os.path.isdir(p)] - if not matches: + base_pattern = os.path.join(str(tr.output_path), "experiments", experiment_name, "*", experiment_name) + task_dirs = [p for p in glob.glob(base_pattern) if os.path.isdir(p)] + if not task_dirs: return - # Take the most recently modified in case of retries - src = max(matches, key=os.path.getmtime) - dst = tr.output_path / "torch_profile" - if not dst.exists(): - shutil.copytree(src, dst) - logging.info(f"Copied pytorch profiler traces from {src} to {dst}") + # Take the most recently modified task dir in case of retries + task_dir = max(task_dirs, key=os.path.getmtime) + + for trace_dir in ("torch_profile", "chakra"): + src = os.path.join(task_dir, trace_dir) + if not os.path.isdir(src): + continue + dst = tr.output_path / trace_dir + if not dst.exists(): + shutil.copytree(src, dst) + logging.info(f"Copied {trace_dir} traces from {src} to {dst}") def extract_mbridge_metrics(logs: str) -> tuple[list[float], list[float]]: From 2d796976e64cf0afbdc15e8da3415137be3dd2e2 Mon Sep 17 00:00:00 2001 From: saivishal1999 Date: Tue, 28 Jul 2026 16:47:31 -0700 Subject: [PATCH 3/9] Catch OSError/shutil.Error in _copy_profiler_traces to not fail job Copy failures (e.g. disk full, permissions) should not mark a successful training job as failed. --- src/cloudai/workloads/megatron_bridge/megatron_bridge.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index 7b6426a08..0ee1c0d8d 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -577,7 +577,10 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: if not step_times_s: return JobStatusResult(is_successful=False, error_message="\n".join(log_data.splitlines()[-40:])) - self._copy_profiler_traces(tr) + try: + self._copy_profiler_traces(tr) + except (OSError, shutil.Error) as e: + logging.warning(f"Failed to copy profiler traces: {e}") return JobStatusResult(is_successful=True) From 53b87a8b0552b17f057c4d3ce4196cd2e7a771ec Mon Sep 17 00:00:00 2001 From: saivishal1999 Date: Tue, 28 Jul 2026 17:20:40 -0700 Subject: [PATCH 4/9] Rename pytorch_profile/ to tensorboard/ instead of copying to root Per reviewer feedback: copying traces to the output root risks overwriting on retries and goes against NemoRun's directory structure. Instead, rename pytorch_profile/ to tensorboard/ within the experiment directory to match MegatronRun's naming convention (pytorch_profile/ contains TensorBoard events, not profiling traces). --- doc/workloads/chakra_tracing.rst | 8 ++++- .../megatron_bridge/megatron_bridge.py | 35 +++++++------------ 2 files changed, 20 insertions(+), 23 deletions(-) diff --git a/doc/workloads/chakra_tracing.rst b/doc/workloads/chakra_tracing.rst index f28d7e54c..5668b29cd 100644 --- a/doc/workloads/chakra_tracing.rst +++ b/doc/workloads/chakra_tracing.rst @@ -34,7 +34,13 @@ overrides passed directly to the Megatron-Bridge training script: This produces the following under ``/``: - ``torch_profile/rank-0.json.gz`` — Chrome-format kineto trace (always produced) -- ``memory_profile.pickle`` — memory allocation snapshot (in nested experiment dir) +- ``memory_profile.pickle`` — memory allocation snapshot + +.. note:: + + When using the ``-pyp`` flag (``pytorch_profiler = true``), MegatronBridge also creates + a ``pytorch_profile/`` directory with TensorBoard events. CloudAI renames this to + ``tensorboard/`` after job completion to match MegatronRun's naming convention. Enabling Chakra Trace diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index 0ee1c0d8d..cf5220d7f 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -578,23 +578,19 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: return JobStatusResult(is_successful=False, error_message="\n".join(log_data.splitlines()[-40:])) try: - self._copy_profiler_traces(tr) + self._rename_pytorch_profile(tr) except (OSError, shutil.Error) as e: - logging.warning(f"Failed to copy profiler traces: {e}") + logging.warning(f"Failed to rename pytorch_profile directory: {e}") return JobStatusResult(is_successful=True) - def _copy_profiler_traces(self, tr: TestRun) -> None: + def _rename_pytorch_profile(self, tr: TestRun) -> None: """ - Copy pytorch kineto and Chakra traces from NemoRun's nested experiment directory to the CloudAI output root. + Rename pytorch_profile/ to tensorboard/ inside the NemoRun experiment directory. - MegatronBridge runs via NemoRun which places profiler output at - experiments////{torch_profile,chakra}/ internally. - This copies them to /{torch_profile,chakra}/ so both MegatronRun and MegatronBridge - produce traces at a consistent location. - - torch_profile/ is always present when use_pytorch_profiler=true. - chakra/ is only present when pytorch_profiler_collect_chakra=true. + PyTorchProfilerPlugin writes TensorBoard events to pytorch_profile/, which is a + misleading name — the data is tensorboard events, not profiling traces. This renames + it to tensorboard/ to match the MegatronRun output convention. """ experiment_name = tr.test.cmd_args.wandb_experiment_name if not experiment_name: @@ -605,17 +601,12 @@ def _copy_profiler_traces(self, tr: TestRun) -> None: if not task_dirs: return - # Take the most recently modified task dir in case of retries - task_dir = max(task_dirs, key=os.path.getmtime) - - for trace_dir in ("torch_profile", "chakra"): - src = os.path.join(task_dir, trace_dir) - if not os.path.isdir(src): - continue - dst = tr.output_path / trace_dir - if not dst.exists(): - shutil.copytree(src, dst) - logging.info(f"Copied {trace_dir} traces from {src} to {dst}") + for task_dir in task_dirs: + src = os.path.join(task_dir, "pytorch_profile") + dst = os.path.join(task_dir, "tensorboard") + if os.path.isdir(src) and not os.path.exists(dst): + os.rename(src, dst) + logging.info(f"Renamed {src} to {dst}") def extract_mbridge_metrics(logs: str) -> tuple[list[float], list[float]]: From 14d495d71156daa2da2c874d935270f8eb37d56e Mon Sep 17 00:00:00 2001 From: saivishal1999 Date: Tue, 28 Jul 2026 21:28:22 -0700 Subject: [PATCH 5/9] Add chakra flag to nemotron nano profiler TOML; remove pytorch_profile rename --- ...bridge_nemotron_nano_pytorch_profiler.toml | 1 + .../megatron_bridge/megatron_bridge.py | 31 ------------------- 2 files changed, 1 insertion(+), 31 deletions(-) diff --git a/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml b/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml index 0708fbe1d..b934c8eb5 100644 --- a/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml +++ b/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml @@ -51,3 +51,4 @@ ep = 8 "profiling.record_shapes" = "true" "profiling.record_memory_history" = "true" "profiling.memory_snapshot_path" = "memory_profile.pickle" +"profiling.pytorch_profiler_collect_chakra" = "true" diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index cf5220d7f..7dfd65532 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -14,11 +14,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import glob import logging import os import re -import shutil from typing import List, Optional, Union, cast from pydantic import Field, ValidationInfo, field_validator @@ -577,37 +575,8 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: if not step_times_s: return JobStatusResult(is_successful=False, error_message="\n".join(log_data.splitlines()[-40:])) - try: - self._rename_pytorch_profile(tr) - except (OSError, shutil.Error) as e: - logging.warning(f"Failed to rename pytorch_profile directory: {e}") - return JobStatusResult(is_successful=True) - def _rename_pytorch_profile(self, tr: TestRun) -> None: - """ - Rename pytorch_profile/ to tensorboard/ inside the NemoRun experiment directory. - - PyTorchProfilerPlugin writes TensorBoard events to pytorch_profile/, which is a - misleading name — the data is tensorboard events, not profiling traces. This renames - it to tensorboard/ to match the MegatronRun output convention. - """ - experiment_name = tr.test.cmd_args.wandb_experiment_name - if not experiment_name: - return - - base_pattern = os.path.join(str(tr.output_path), "experiments", experiment_name, "*", experiment_name) - task_dirs = [p for p in glob.glob(base_pattern) if os.path.isdir(p)] - if not task_dirs: - return - - for task_dir in task_dirs: - src = os.path.join(task_dir, "pytorch_profile") - dst = os.path.join(task_dir, "tensorboard") - if os.path.isdir(src) and not os.path.exists(dst): - os.rename(src, dst) - logging.info(f"Renamed {src} to {dst}") - def extract_mbridge_metrics(logs: str) -> tuple[list[float], list[float]]: step_times_s: list[float] = [] From e3426b934ee7bec014eab59d10adc82570886458 Mon Sep 17 00:00:00 2001 From: saivishal1999 Date: Wed, 29 Jul 2026 01:18:16 -0700 Subject: [PATCH 6/9] Rename tb_logs/ to tensorboard/ after job completion; remove stale note from docs --- doc/workloads/chakra_tracing.rst | 7 ------- .../megatron_bridge/megatron_bridge.py | 21 +++++++++++++++++++ 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/doc/workloads/chakra_tracing.rst b/doc/workloads/chakra_tracing.rst index 5668b29cd..9ee22b108 100644 --- a/doc/workloads/chakra_tracing.rst +++ b/doc/workloads/chakra_tracing.rst @@ -36,13 +36,6 @@ This produces the following under ``/``: - ``torch_profile/rank-0.json.gz`` — Chrome-format kineto trace (always produced) - ``memory_profile.pickle`` — memory allocation snapshot -.. note:: - - When using the ``-pyp`` flag (``pytorch_profiler = true``), MegatronBridge also creates - a ``pytorch_profile/`` directory with TensorBoard events. CloudAI renames this to - ``tensorboard/`` after job completion to match MegatronRun's naming convention. - - Enabling Chakra Trace --------------------- diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index 7dfd65532..44cc288f3 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -14,9 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import glob import logging import os import re +import shutil from typing import List, Optional, Union, cast from pydantic import Field, ValidationInfo, field_validator @@ -575,8 +577,27 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: if not step_times_s: return JobStatusResult(is_successful=False, error_message="\n".join(log_data.splitlines()[-40:])) + try: + self._rename_tb_logs(tr) + except (OSError, shutil.Error) as e: + logging.warning(f"Failed to rename tb_logs directory: {e}") + return JobStatusResult(is_successful=True) + def _rename_tb_logs(self, tr: TestRun) -> None: + experiment_name = tr.test.cmd_args.wandb_experiment_name + if not experiment_name: + return + + base_pattern = os.path.join(str(tr.output_path), "experiments", experiment_name, "*", experiment_name) + task_dirs = [p for p in glob.glob(base_pattern) if os.path.isdir(p)] + for task_dir in task_dirs: + src = os.path.join(task_dir, "tb_logs") + dst = os.path.join(task_dir, "tensorboard") + if os.path.isdir(src) and not os.path.exists(dst): + os.rename(src, dst) + logging.info(f"Renamed {src} to {dst}") + def extract_mbridge_metrics(logs: str) -> tuple[list[float], list[float]]: step_times_s: list[float] = [] From 9764e98774a8dd7afb5bbda63c73c4ea0b359ee2 Mon Sep 17 00:00:00 2001 From: saivishal1999 Date: Wed, 29 Jul 2026 01:56:00 -0700 Subject: [PATCH 7/9] Dynamically find and rename tensorboard dir instead of hardcoding tb_logs --- .../megatron_bridge/megatron_bridge.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index 44cc288f3..cda17ca2f 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -578,13 +578,14 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: return JobStatusResult(is_successful=False, error_message="\n".join(log_data.splitlines()[-40:])) try: - self._rename_tb_logs(tr) + self._rename_tensorboard_dir(tr) except (OSError, shutil.Error) as e: - logging.warning(f"Failed to rename tb_logs directory: {e}") + logging.warning(f"Failed to rename tensorboard directory: {e}") return JobStatusResult(is_successful=True) - def _rename_tb_logs(self, tr: TestRun) -> None: + def _rename_tensorboard_dir(self, tr: TestRun) -> None: + """Rename whichever directory contains TensorBoard events to tensorboard/.""" experiment_name = tr.test.cmd_args.wandb_experiment_name if not experiment_name: return @@ -592,11 +593,16 @@ def _rename_tb_logs(self, tr: TestRun) -> None: base_pattern = os.path.join(str(tr.output_path), "experiments", experiment_name, "*", experiment_name) task_dirs = [p for p in glob.glob(base_pattern) if os.path.isdir(p)] for task_dir in task_dirs: - src = os.path.join(task_dir, "tb_logs") dst = os.path.join(task_dir, "tensorboard") - if os.path.isdir(src) and not os.path.exists(dst): - os.rename(src, dst) - logging.info(f"Renamed {src} to {dst}") + if os.path.exists(dst): + continue + for entry in os.scandir(task_dir): + if not entry.is_dir() or entry.name == "tensorboard": + continue + if glob.glob(os.path.join(entry.path, "events.out.tfevents.*")): + os.rename(entry.path, dst) + logging.info(f"Renamed {entry.path} to {dst}") + break def extract_mbridge_metrics(logs: str) -> tuple[list[float], list[float]]: From 6bfec86bfcbf6c91638a7d0be64206f3ee25ea9b Mon Sep 17 00:00:00 2001 From: saivishal1999 Date: Wed, 29 Jul 2026 02:30:31 -0700 Subject: [PATCH 8/9] Fix chakra_tracing docs: correct output paths and add tensorboard/ to listing --- doc/workloads/chakra_tracing.rst | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/doc/workloads/chakra_tracing.rst b/doc/workloads/chakra_tracing.rst index 9ee22b108..d26394169 100644 --- a/doc/workloads/chakra_tracing.rst +++ b/doc/workloads/chakra_tracing.rst @@ -31,9 +31,10 @@ overrides passed directly to the Megatron-Bridge training script: "profiling.record_memory_history" = "true" "profiling.memory_snapshot_path" = "memory_profile.pickle" -This produces the following under ``/``: +This produces the following under the NemoRun experiment directory (``/experiments////``): - ``torch_profile/rank-0.json.gz`` — Chrome-format kineto trace (always produced) +- ``tensorboard/`` — TensorBoard events (renamed from whatever directory NeMo writes them to) - ``memory_profile.pickle`` — memory allocation snapshot Enabling Chakra Trace @@ -54,7 +55,7 @@ from the kineto trace. Add ``pytorch_profiler_collect_chakra = true``: "profiling.memory_snapshot_path" = "memory_profile.pickle" "profiling.pytorch_profiler_collect_chakra" = "true" -This produces an additional ``chakra/rank-0.json.gz`` under ``/``. +This produces an additional ``chakra/rank-0.json.gz`` in the same experiment directory. The kineto trace (``torch_profile/rank-0.json.gz``) is always produced regardless of this flag. @@ -68,7 +69,7 @@ A ready-to-use TOML is provided at: conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml -This runs Nemotron 3 Nano 3B (2 nodes, 8 GPUs, 60 steps) with Kineto profiling at +This runs Nemotron 3 Nano 3B (8 GPUs, 60 steps) with Kineto and Chakra profiling at steps 45–50 and is suitable for quickly validating the profiling setup. From 13ad2c1155c8b1a178ca42ab7d1cb7d4b54b6969 Mon Sep 17 00:00:00 2001 From: saivishal1999 Date: Wed, 29 Jul 2026 17:00:05 -0700 Subject: [PATCH 9/9] Remove rename logic; move chakra/kineto docs into MegatronBridge docs; drop expensive profiling flags from example TOML --- ...bridge_nemotron_nano_pytorch_profiler.toml | 3 - doc/workloads/chakra_tracing.rst | 115 ------------------ doc/workloads/index.rst | 1 - doc/workloads/megatron_bridge.rst | 111 +++++++++++++++++ .../megatron_bridge/megatron_bridge.py | 27 ---- 5 files changed, 111 insertions(+), 146 deletions(-) delete mode 100644 doc/workloads/chakra_tracing.rst diff --git a/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml b/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml index b934c8eb5..276e8b35d 100644 --- a/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml +++ b/conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml @@ -48,7 +48,4 @@ ep = 8 "profiling.profile_step_start" = "45" "profiling.profile_step_end" = "50" "profiling.profile_ranks" = "[0]" -"profiling.record_shapes" = "true" -"profiling.record_memory_history" = "true" -"profiling.memory_snapshot_path" = "memory_profile.pickle" "profiling.pytorch_profiler_collect_chakra" = "true" diff --git a/doc/workloads/chakra_tracing.rst b/doc/workloads/chakra_tracing.rst deleted file mode 100644 index d26394169..000000000 --- a/doc/workloads/chakra_tracing.rst +++ /dev/null @@ -1,115 +0,0 @@ -Chakra / Kineto Tracing with MegatronBridge -============================================ - -Disabling nsys --------------- - -The existing MegatronBridge TOMLs in ``conf/staging/`` have ``enable_nsys = true``. -To switch to Kineto/Chakra profiling, either: - -- **Omit** ``enable_nsys`` from the TOML (it defaults to ``false``) -- **Or** explicitly set ``enable_nsys = false`` - -Both are equivalent. Do not set ``enable_nsys = true`` alongside Kineto flags — the -two profilers conflict and Megatron-Bridge will raise an assertion error. - - -Enabling Kineto Trace ---------------------- - -Add the following to ``[extra_cmd_args]`` in your Test TOML. These are Hydra-style -overrides passed directly to the Megatron-Bridge training script: - -.. code-block:: toml - - [extra_cmd_args] - "profiling.use_pytorch_profiler" = "true" - "profiling.profile_step_start" = "45" - "profiling.profile_step_end" = "50" - "profiling.profile_ranks" = "[0]" - "profiling.record_shapes" = "true" - "profiling.record_memory_history" = "true" - "profiling.memory_snapshot_path" = "memory_profile.pickle" - -This produces the following under the NemoRun experiment directory (``/experiments////``): - -- ``torch_profile/rank-0.json.gz`` — Chrome-format kineto trace (always produced) -- ``tensorboard/`` — TensorBoard events (renamed from whatever directory NeMo writes them to) -- ``memory_profile.pickle`` — memory allocation snapshot - -Enabling Chakra Trace ---------------------- - -The Chakra trace is a PyTorch ``ExecutionTraceObserver`` recording — a separate file -from the kineto trace. Add ``pytorch_profiler_collect_chakra = true``: - -.. code-block:: toml - - [extra_cmd_args] - "profiling.use_pytorch_profiler" = "true" - "profiling.profile_step_start" = "45" - "profiling.profile_step_end" = "50" - "profiling.profile_ranks" = "[0]" - "profiling.record_shapes" = "true" - "profiling.record_memory_history" = "true" - "profiling.memory_snapshot_path" = "memory_profile.pickle" - "profiling.pytorch_profiler_collect_chakra" = "true" - -This produces an additional ``chakra/rank-0.json.gz`` in the same experiment directory. -The kineto trace (``torch_profile/rank-0.json.gz``) is always produced regardless of -this flag. - - -Reference TOML --------------- - -A ready-to-use TOML is provided at: - -.. code-block:: text - - conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml - -This runs Nemotron 3 Nano 3B (8 GPUs, 60 steps) with Kineto and Chakra profiling at -steps 45–50 and is suitable for quickly validating the profiling setup. - - -Parameter Reference -------------------- - -.. list-table:: - :header-rows: 1 - :widths: 40 15 45 - - * - Parameter - - Type - - Description - * - ``profiling.use_pytorch_profiler`` - - bool - - Enable PyTorch kineto profiler. Required for all Kineto/Chakra tracing. - * - ``profiling.profile_step_start`` - - int - - Training step at which profiling begins. - * - ``profiling.profile_step_end`` - - int - - Training step at which profiling ends and trace is exported. - * - ``profiling.profile_ranks`` - - list[int] - - GPU ranks to profile (e.g. ``[0]`` for rank 0 only). - * - ``profiling.record_shapes`` - - bool - - Record tensor shapes in the kineto trace. - * - ``profiling.record_memory_history`` - - bool - - Record memory allocation history (produces ``memory_profile.pickle``). - * - ``profiling.memory_snapshot_path`` - - str - - Filename for the memory snapshot pickle file. - * - ``profiling.pytorch_profiler_collect_chakra`` - - bool - - Collect Chakra ExecutionTraceObserver trace (``chakra/rank-N.json.gz``). The kineto trace is always produced; this adds an additional file. - * - ``profiling.pytorch_profiler_collect_callstack`` - - bool - - Include Python call stack in the kineto trace (increases trace size). - * - ``enable_nsys`` - - bool - - Enable Nsight Systems profiling. **Cannot be used with Kineto.** Default: ``false``. diff --git a/doc/workloads/index.rst b/doc/workloads/index.rst index 64707cea7..916d596cc 100644 --- a/doc/workloads/index.rst +++ b/doc/workloads/index.rst @@ -14,7 +14,6 @@ Available Workloads ":doc:`ai_dynamo`", "✅", "✅", "❌", "❌" ":doc:`bash_cmd`", "✅", "❌", "❌", "❌" ":doc:`chakra_replay`", "✅", "❌", "❌", "❌" - ":doc:`chakra_tracing`", "✅", "❌", "❌", "❌" ":doc:`ddlb`", "✅", "❌", "❌", "❌" ":doc:`deepep`", "✅", "❌", "❌", "❌" ":doc:`dynamo_mocker`", "❌", "❌", "❌", "✅" diff --git a/doc/workloads/megatron_bridge.rst b/doc/workloads/megatron_bridge.rst index 5a746c223..1f184e212 100644 --- a/doc/workloads/megatron_bridge.rst +++ b/doc/workloads/megatron_bridge.rst @@ -84,6 +84,117 @@ Test-in-Scenario example: domain = "llm" compute_dtype = "fp8_mx" +Chakra / Kineto Tracing +----------------------- + +Disabling nsys +~~~~~~~~~~~~~~ + +The existing MegatronBridge TOMLs in ``conf/staging/`` have ``enable_nsys = true``. +To switch to Kineto/Chakra profiling, either: + +- **Omit** ``enable_nsys`` from the TOML (it defaults to ``false``) +- **Or** explicitly set ``enable_nsys = false`` + +Both are equivalent. Do not set ``enable_nsys = true`` alongside Kineto flags — the +two profilers conflict and Megatron-Bridge will raise an assertion error. + + +Enabling Kineto Trace +~~~~~~~~~~~~~~~~~~~~~ + +Add the following to ``[extra_cmd_args]`` in your Test TOML. These are Hydra-style +overrides passed directly to the Megatron-Bridge training script: + +.. code-block:: toml + + [extra_cmd_args] + "profiling.use_pytorch_profiler" = "true" + "profiling.profile_step_start" = "45" + "profiling.profile_step_end" = "50" + "profiling.profile_ranks" = "[0]" + +This produces the following under the NemoRun experiment directory (``/experiments////``): + +- ``torch_profile/rank-0.json.gz`` — Chrome-format kineto trace (always produced) +- ``tb_logs/`` — TensorBoard events + + +Enabling Chakra Trace +~~~~~~~~~~~~~~~~~~~~~ + +The Chakra trace is a PyTorch ``ExecutionTraceObserver`` recording — a separate file +from the kineto trace. Add ``pytorch_profiler_collect_chakra = true``: + +.. code-block:: toml + + [extra_cmd_args] + "profiling.use_pytorch_profiler" = "true" + "profiling.profile_step_start" = "45" + "profiling.profile_step_end" = "50" + "profiling.profile_ranks" = "[0]" + "profiling.pytorch_profiler_collect_chakra" = "true" + +This produces an additional ``chakra/rank-0.json.gz`` in the same experiment directory. +The kineto trace (``torch_profile/rank-0.json.gz``) is always produced regardless of +this flag. + + +Reference TOML +~~~~~~~~~~~~~~ + +A ready-to-use TOML is provided at: + +.. code-block:: text + + conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml + +This runs Nemotron 3 Nano 3B (8 GPUs, 60 steps) with Kineto and Chakra profiling at +steps 45–50 and is suitable for quickly validating the profiling setup. + + +Profiling Parameter Reference +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 40 15 45 + + * - Parameter + - Type + - Description + * - ``profiling.use_pytorch_profiler`` + - bool + - Enable PyTorch kineto profiler. Required for all Kineto/Chakra tracing. + * - ``profiling.profile_step_start`` + - int + - Training step at which profiling begins. + * - ``profiling.profile_step_end`` + - int + - Training step at which profiling ends and trace is exported. + * - ``profiling.profile_ranks`` + - list[int] + - GPU ranks to profile (e.g. ``[0]`` for rank 0 only). + * - ``profiling.record_shapes`` + - bool + - Record tensor shapes in the kineto trace. Expensive — off by default. + * - ``profiling.record_memory_history`` + - bool + - Record memory allocation history (produces ``memory_profile.pickle``). Expensive — off by default. + * - ``profiling.memory_snapshot_path`` + - str + - Filename for the memory snapshot pickle file. + * - ``profiling.pytorch_profiler_collect_chakra`` + - bool + - Collect Chakra ExecutionTraceObserver trace (``chakra/rank-N.json.gz``). The kineto trace is always produced; this adds an additional file. + * - ``profiling.pytorch_profiler_collect_callstack`` + - bool + - Include Python call stack in the kineto trace (increases trace size). + * - ``enable_nsys`` + - bool + - Enable Nsight Systems profiling. **Cannot be used with Kineto.** Default: ``false``. + + API Documentation ----------------- diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index cda17ca2f..7dfd65532 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -14,11 +14,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -import glob import logging import os import re -import shutil from typing import List, Optional, Union, cast from pydantic import Field, ValidationInfo, field_validator @@ -577,33 +575,8 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: if not step_times_s: return JobStatusResult(is_successful=False, error_message="\n".join(log_data.splitlines()[-40:])) - try: - self._rename_tensorboard_dir(tr) - except (OSError, shutil.Error) as e: - logging.warning(f"Failed to rename tensorboard directory: {e}") - return JobStatusResult(is_successful=True) - def _rename_tensorboard_dir(self, tr: TestRun) -> None: - """Rename whichever directory contains TensorBoard events to tensorboard/.""" - experiment_name = tr.test.cmd_args.wandb_experiment_name - if not experiment_name: - return - - base_pattern = os.path.join(str(tr.output_path), "experiments", experiment_name, "*", experiment_name) - task_dirs = [p for p in glob.glob(base_pattern) if os.path.isdir(p)] - for task_dir in task_dirs: - dst = os.path.join(task_dir, "tensorboard") - if os.path.exists(dst): - continue - for entry in os.scandir(task_dir): - if not entry.is_dir() or entry.name == "tensorboard": - continue - if glob.glob(os.path.join(entry.path, "events.out.tfevents.*")): - os.rename(entry.path, dst) - logging.info(f"Renamed {entry.path} to {dst}") - break - def extract_mbridge_metrics(logs: str) -> tuple[list[float], list[float]]: step_times_s: list[float] = []