Merge pull request #456 from zswerth/gpu_metrics_support

akoumpa · web-flow · commit 4abd481402ad · 2025-01-28T23:21:11.000-08:00
Add GPU metrics support for nsight profiling
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:

Original file line number	Diff line number	Diff line change
`@@ -290,6 +290,9 @@ def _make_nsys_command_prefix(self, results_dir: str) -> str:`
`290`	`290`	`f"--capture-range-end=stop "`
`291`	`291`	`f"--cuda-graph-trace=node "`
`292`	`292`	`)`
	`293`	`+ if nsys_cfg.get("gpu_metrics", False):`
	`294`	`+ slurm_local_rank = "\${SLURM_LOCALID}"`
	`295`	`+ nsys_prefix += f"--gpu-metrics-device={slurm_local_rank} "`
`293`	`296`	`return nsys_prefix`
`294`	`297`
`295`	`298`	`def _make_container_mounts_string(self) -> str:`