Function nsys_profile

src/fused_mm_sampling/modal_lib/modal_nsys.py:29–110

(
    name: str,
    n_hidden_states: str,
    case: str,
    n_procs: int,
    gpu_name: str,
    postfix: str,
    n_runs_benchmark: str,
)

@app.function(gpu=cfg.gpu_spec, image=nsys_image, volumes=make_volumes(), timeout=cfg.timeout)
def nsys_profile(
    name: str,
    n_hidden_states: str,
    case: str,
    n_procs: int,
    gpu_name: str,
    postfix: str,
    n_runs_benchmark: str,
):
    set_volume_caches()
    report_dir = f"{volume_path}/nsys-profiles/{gpu_name}/tp{n_procs}/case-{case}/bsz{n_hidden_states}{postfix}"
    os.makedirs(report_dir, exist_ok=True)
    report_path = f"{report_dir}/{name}"

    speed_test_args = [
        "/opt/fmms/speed_test.py",
        "--name",
        name,
        "--n_hidden_states",
        n_hidden_states,
        "--n_runs_warmup",
        "25",
        "--n_runs_benchmark",
        n_runs_benchmark,
        "--bench_fn=own",
        f"--case={case}",
        f"--n_procs={n_procs}",
        "--nsys_profile=true",
    ]

    if n_procs > 1:
        # Per-rank nsys: torchrun launches the wrapper, each rank gets its
        # own nsys instance and .nsys-rep file. This avoids the known nsys
        # limitation where cudaProfilerApi capture misses child-process GPU
        # activity when wrapping torchrun from the outside.
        cmd = [
            "torchrun",
            f"--nproc_per_node={n_procs}",
            "/opt/fmms/nsys_wrapper.py",
            "--report-dir",
            report_dir,
            "--name",
            name,
            "--",
            *speed_test_args,
        ]
        env = None
    else:
        cmd = [
            "nsys",
            "profile",
            "-o",
            report_path,
            "--force-overwrite=true",
            "--capture-range=cudaProfilerApi",
            "--cuda-memory-usage=true",
            "python",
            *speed_test_args,