Add GPU metrics support for nsight profiling

NVIDIA · Jan 23, 2025 · 1122df8 · 1122df8
1 parent 855b40e
commit 1122df8
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/launcher_scripts/nemo_launcher/core/stages.py b/launcher_scripts/nemo_launcher/core/stages.py
@@ -290,6 +290,9 @@ def _make_nsys_command_prefix(self, results_dir: str) -> str:
                 f"--capture-range-end=stop "
                 f"--cuda-graph-trace=node "
             )
+            if nsys_cfg.get("gpu_metrics", False):
+                slurm_local_rank = "\${SLURM_LOCALID}"
+                nsys_prefix += (f"--gpu-metrics-device={slurm_local_rank} ")
         return nsys_prefix
 
     def _make_container_mounts_string(self) -> str: