Add nsys command argument to profile cuda graph workload (#138)

* Add nsys command argument to profile cuda graph workload; fix nsys profile export path Signed-off-by: Guyue Huang <[email protected]> * Revert a change to nsys profile path because the bug is fixed in tot Signed-off-by: Guyue Huang <[email protected]> * Fix test_slurm nsys command Signed-off-by: Guyue Huang <[email protected]> --------- Signed-off-by: Guyue Huang <[email protected]> Signed-off-by: Guyue Huang <[email protected]> Co-authored-by: Guyue Huang <[email protected]>
NVIDIA · Jan 23, 2025 · 5ed6128 · 5ed6128
1 parent 11e0d2f
commit 5ed6128
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/src/nemo_run/core/execution/base.py b/src/nemo_run/core/execution/base.py
@@ -50,6 +50,7 @@ def get_nsys_prefix(self, profile_dir: str) -> Optional[list[str]]:
                 "true",
                 "--capture-range=cudaProfilerApi",
                 "--capture-range-end=stop",
+                "--cuda-graph-trace=node",
             ]
             return args
 

diff --git a/test/core/execution/test_slurm.py b/test/core/execution/test_slurm.py
@@ -508,6 +508,7 @@ def test_dummy_batch_request_nsys(
             "true",
             "--capture-range=cudaProfilerApi",
             "--capture-range-end=stop",
+            "--cuda-graph-trace=node",
         ]
 
     def test_dummy_batch_request_warn(