Add Nemotron Nano pytorch profiler TOML and chakra tracing docs

conf/experimental/megatron_bridge/test/gb200/megatron_bridge_nemotron_nano_pytorch_profiler.toml

@@ -0,0 +1,51 @@
+# 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.pytorch_profiler_collect_chakra" = "true"

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 (``<cloudai_output>/experiments/<experiment_name>/<run_id>/<experiment_name>/``):
+
+- ``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
 -----------------