Why the new rl.checkpoint.async_save tag says true on every save

NVIDIA-NeMo/RL #4607. All permalinks pin the PR head 2f6b955.

Before
There was no finalize span and no tag.
This PR
Adds the rl.checkpoint.finalize span and tags it with rl.checkpoint.async_save, so you can tell a slow async write from a quick rename after a sync save.
Status
Open, not fixed in this PR. The tag is true for sync saves too. Raised in the review.

Two pages on this PR, one question each. This one: the async_save tag bug.

The tag is set from whether the caller passed a wait function: async_save=wait_fn is not None. But every training loop passes one on every save, sync or async (e.g. grpo.py:4125; all 12 callers do the same). With async save off, that function does nothing, but it is still not None. The PR's test for the sync case (test_checkpoint_spans.py:196) passes no wait function at all, which no real caller does.

The same tag for two different saves

Two Megatron runs, step 12 saves a checkpoint checkpoint.async_save: false checkpoint.async_save: true begin_finalization(tmp_step_12, wait_fn=policy.finalize_async_save) begin_finalization(tmp_step_12, wait_fn=policy.finalize_async_save) wait_fn() → returns at once bytes were already written during the save wait_fn() → waits for the writers bytes were still being written rename tmp_step_12 → step_12 rename tmp_step_12 → step_12 rl.checkpoint.async_save = true ✗ rl.checkpoint.async_save = true ✓ checkpoint.py:489 — async_save=wait_fn is not None wrong right only by luck

The tag answers "was a function passed", which is always yes. So a short finalize span still cannot be read as "sync save" or "async save that had already finished", which is the question the tag was added for (span-groups.md:130).

The fix — proposed, not in the PR

Where the fix lands: RL only. Nothing below is implemented. Remove the tag. The real setting lives in two backend-specific places (Megatron's checkpoint.async_save, Automodel's is_async), and passing it through all 12 save sites only to label a span costs more than the tag is worth. The same change also deletes the constant at instrumentation.py:101, the docs row, and the sync-case test.

@@ utils/checkpoint.py L489 · begin_finalization._finalize
             try:
-                with _finalize_telemetry(carrier, step, async_save=wait_fn is not None):
+                with _finalize_telemetry(carrier, step):
                     if wait_fn is not None:
                         wait_fn()
@@ utils/checkpoint.py L57-L59 and L77 · _finalize_telemetry
 def _finalize_telemetry(
-    carrier: Mapping[str, str], step: Optional[int], *, async_save: bool
+    carrier: Mapping[str, str], step: Optional[int]
 ) -> Iterator[None]:
 ...
-        stack.enter_context(checkpoint_finalize_span(step, async_save=async_save))
+        stack.enter_context(checkpoint_finalize_span(step))
@@ telemetry/instrumentation.py L970-L974 and L1008 · checkpoint_finalize_span
 def checkpoint_finalize_span(
     iteration: Optional[int],
     *,
-    async_save: bool,
     tracer: Optional[Tracer] = None,
 ...
-    attrs: dict[str, Any] = {RL_CHECKPOINT_ASYNC_ATTR: async_save, **attributes}
+    attrs: dict[str, Any] = dict(attributes)
     if iteration is not None:

🔴 review comment — checkpoint.py:489

The finalize span itself stays. Its duration is still the time from staging to "resumable", and a reader who needs the save mode can read it from the run config.

So what. Nothing breaks in training. But every dashboard that splits checkpoint time by rl.checkpoint.async_save puts all sync saves in the async bucket, so the one question the tag exists for gets the wrong answer, quietly.