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..276e8b35d --- /dev/null +++ b/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" 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 -----------------