fix: capture inner vLLM TP-worker nsys traces for ALL ranks

[zswerth]

What does this PR do ?

Fixes multi-rank generation-profiling so that every inner vLLM tensor-parallel worker writes a complete nsys trace (previously TP>1 generation profiling produced zero usable inner traces, and crashed at init).

Issues

No tracked issue; standalone fix
Multi-rank generation (inner vLLM tensor-parallel worker) nsys profiling previously produced NO usable traces: at teardown the inner workers are SIGKILLed, so nsys never writes the stream footer and the .qdstrm is an incomplete, unconvertible capture. On the async engine the inner workers never even received cudaProfilerStart. On aarch64 GB200 the default cuda-graph-trace=node HW-CUPTI path also SIGSEGVs at cudaProfilerStart.

Usage

Enable generation + policy nsys profiling as usual — inner TP-worker traces now finalize for all ranks:

# Profile policy (training) + vLLM (generation) workers, steps 2..4 (left-incl/right-excl).
# cuda-sw is the aarch64 GB200-safe trace target (HW `cuda` SIGSEGVs at cudaProfilerStart).
export NRL_NSYS_WORKER_PATTERNS="*policy*,*vllm*"
export NRL_NSYS_PROFILE_STEP_RANGE="2:5"
export NRL_NSYS_EXTRA_OPTIONS='{"t":"cuda-sw,nvtx"}'

uv run examples/run_grpo.py --config examples/configs/grpo_math_1B_megatron.yaml \
    policy.generation.vllm_cfg.tensor_parallel_size=2 \
    grpo.max_num_steps=5

Result: one 2:5.nsys-rep per rank — megatron_policy_worker_, vllm_generation_worker_* (outer), and one vllm_tp_worker_* per inner TP rank (the traces that were missing before).

Repro of the issue and validation of the fix

qwen3-32b, async engine, vLLM tensor_parallel_size=2 (so real inner Ray TP workers exist), single NVLink-local node, profile window 1:3. The two runs differ only by this patch:

┌──────────────────────────────┬────────────────────────────────────────────┬───────────────────────────────┐
│                              │             BEFORE (unpatched)             │        AFTER (this PR)        │
├──────────────────────────────┼────────────────────────────────────────────┼───────────────────────────────┤
│ inner vllm_tp_worker traces  │ 0                                          │ 32                            │
├──────────────────────────────┼────────────────────────────────────────────┼───────────────────────────────┤
│ outer vllm_generation_worker │ —                                          │ 32                            │
├──────────────────────────────┼────────────────────────────────────────────┼───────────────────────────────┤
│ megatron_policy_worker       │ —                                          │ 32                            │
├──────────────────────────────┼────────────────────────────────────────────┼───────────────────────────────┤
│ segfaults at EngineCore init │ 5 (CUPTI × cuda-graph-trace=node, aarch64) │ 0                             │
├──────────────────────────────┼────────────────────────────────────────────┼───────────────────────────────┤
│ profile-window entry         │ never fired                                │ fired, trained through window │
└──────────────────────────────┴────────────────────────────────────────────┴───────────────────────────────┘

Why it failed before:

1. On aarch64 GB200 the default cuda-graph-trace=node HW-CUPTI path SIGSEGVs at cudaProfilerStart → inner EngineCore workers die at init.
2. On the async engine the inherited sync start/stop_gpu_profiling dropped the collective_rpc(...) coroutine un-awaited → inner workers never received cudaProfilerStart/Stop.
3. At teardown the inner TP workers are SIGKILLed → nsys never writes the stream footer → the .qdstrm is an incomplete, unconvertible capture (at most one rank survived).

Backward-compat (non-regression) — all PASS, 0 crashes:

┌────────────────────────────────────┬───────────────┬─────────────────────────────┬──────────────┐
│              Use case              │ inner workers │           traces            │   outcome    │
├────────────────────────────────────┼───────────────┼─────────────────────────────┼──────────────┤
│ sync + TP=1 (VllmGenerationWorker) │ none (TP=1)   │ 8 gen + 8 policy            │ ✅ unchanged │
├────────────────────────────────────┼───────────────┼─────────────────────────────┼──────────────┤
│ async + TP=1                       │ none (TP=1)   │ gen + policy                │ ✅ unchanged │
├────────────────────────────────────┼───────────────┼─────────────────────────────┼──────────────┤
│ async + TP>1                       │ present       │ full inner + outer + policy │ ✅ fixed     │
└────────────────────────────────────┴───────────────┴─────────────────────────────┴──────────────┘

Container Image: nvcr.io#nvidian/nemo-rl:nightly

All three use cases run the same GRPO entrypoint on a Ray-on-Slurm cluster; only the vLLM parallelism/engine and profile knobs differ. Common profiling env (set on the launch shell so it propagates to every node/worker):

export NRL_NSYS_WORKER_PATTERNS="*policy*,*vllm*"   # profile policy (train) + vLLM (gen) workers
export NRL_NSYS_PROFILE_STEP_RANGE="1:3"            # left-incl / right-excl; fires cudaProfilerStart@1, stop@3
export NRL_NSYS_EXTRA_OPTIONS='{"t":"cuda-sw,nvtx"}' # trace-target override (use on platforms where the HW `cuda` target crashes at cudaProfilerStart)

Traces land per worker under each Ray session dir:
.../-logs/ray//session_*/logs/nsight/.nsys-rep

Use case 1 — async + TP>1 (the fix: inner-worker traces)

uv run examples/run_grpo.py \
  --config examples/configs/recipes/llm/performance/grpo-qwen3-32b.yaml \
  policy.generation.vllm_cfg.async_engine=true \
  policy.generation.vllm_cfg.tensor_parallel_size=2 \
  grpo.max_num_steps=5

Expected: three trace classes — megatron_policy_worker_, vllm_async_generation_worker_ (outer), and vllm_tp_worker_1:3_%p (one per inner TP rank = #engines × TP). The vllm_tp_worker_ prefix (vs stock vLLM's worker_process_%p) is the proof the deferred-capture config reached the inner EngineCore subprocess. Before the fix: the inner workers segfault at init / produce zero finalized inner traces.

Use case 2 — sync + TP=1 (non-regression)

uv run examples/run_grpo.py \
  --config examples/configs/recipes/llm/performance/grpo-qwen3-32b.yaml \
  policy.generation.vllm_cfg.async_engine=false \
  policy.generation.vllm_cfg.tensor_parallel_size=1 \
  grpo.max_num_steps=5

Expected: megatron_policy_worker_* + vllm_generation_worker_* only. At TP=1 each GPU is its own engine (no inner Ray workers), so there is no vllm_tp_worker; the plugin is a no-op. Run completes, 0 crashes.

Use case 3 — async + TP=1 (non-regression)

uv run examples/run_grpo.py \
  --config examples/configs/recipes/llm/performance/grpo-qwen3-32b.yaml \
  policy.generation.vllm_cfg.async_engine=true \
  policy.generation.vllm_cfg.tensor_parallel_size=1 \
  grpo.max_num_steps=5

Expected: same as case 2 — vllm_generation_worker_* + megatron_policy_worker_*, no inner workers, 0 crashes.

Four minimal changes:

• distributed/worker_group_utils.py: drop cuda-graph-trace=node from the policy/outer nsight config (fragile CUPTI on aarch64; user can re-add via NRL_NSYS_EXTRA_OPTIONS).

• models/generation/vllm/vllm_worker.py (_patch_vllm_nsight_config): honor NRL_NSYS_EXTRA_OPTIONS so the trace target is overridable (e.g. cuda-sw on aarch64), drop cuda-graph-trace.

• models/generation/vllm/vllm_worker_async.py: override start/stop_gpu_profiling as async and AWAIT collective_rpc, so cudaProfilerStart/Stop actually reach every inner worker on the async engine (the inherited sync methods dropped the coroutine un-awaited); plus a short grace after stop so nsys can finalize.

• models/generation/vllm/vllm_backend.py (VllmInternalWorkerExtension.stop_gpu_profiling): after cudaProfilerStop, finalize each inner worker's live nsys capture NON-DESTRUCTIVELY via nsys stop --session. This writes a complete .nsys-rep without killing the process, so no worker death and no engine fault-tolerance cascade -- every rank finalizes (vs at most one before). Runs on all inner workers via the existing collective_rpc("stop_gpu_profiling") broadcast.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[zswerth]

A pending PR on NemoRL should fix all the added patches provided here for profiling enablement: #4090