perf: overlap non-colocated MInf + polict init

nemo_rl/algorithms/grpo.py

@@ -1268,12 +1268,39 @@ def init_megatron_generation(
             cluster=None if colocated_inference else inference_cluster,
             policy=policy if colocated_inference else None,
             processor=processor,
-            weights_path=weights_path,
             skip_weight_load=not colocated_inference,
             reserved_http_server_port=reserved_http_server_port,
         )
         return mg, time.perf_counter() - t0
 
+    def init_megatron_weight_synchronizer(
+        policy: ColocatablePolicyInterface,
+        policy_generation: MegatronGeneration,
+    ) -> None:
+        """Initialize Megatron weight synchronizer.
+
+        For non-colocated inference, also performs the initial weight sync.
+        """
+        t0 = time.perf_counter()
+        weight_synchronizer = create_weight_synchronizer(
+            policy=policy,
+            generation=policy_generation,
+            generation_backend="megatron",
+            colocated=colocated_inference,
+            train_cluster=train_cluster,
+            inference_cluster=None if colocated_inference else inference_cluster,
+        )
+        policy_generation.weight_synchronizer = weight_synchronizer
+        weight_synchronizer.init_communicator()
+        setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0
+        if not colocated_inference:
+            # The skip-load inference engine gets its final weight buffers here.
+            # Its first prepare_for_generation also starts the HTTP server only after the refit,
+            # so CUDA graphs can capture those persistent buffers.
+            t0 = time.perf_counter()
+            weight_synchronizer.sync_weights()
+            setup_timing_metrics.weight_sync_time_s = time.perf_counter() - t0
+
     def initialize_generation_with_policy(
         init_generation_fn,
         colocated_inference: bool,
@@ -1349,41 +1376,55 @@ def initialize_generation_with_policy(
             setup_timing_metrics.generation_init_reserve_time_s = reserve_time
             print(f"  ✓ Reserved Megatron server URL: {reserved_url}", flush=True)
 
-            def init_megatron_stack():
-                """Init policy then generation; rank 0 holds the reserved port."""
-                p, policy_t = init_policy(
-                    reserved_http_server_port=reserved_http_server_port
-                    if colocated_inference
-                    else None
-                )
-                pg, gen_t = init_megatron_generation(
-                    p,
-                    reserved_http_server_port=None
-                    if colocated_inference
-                    else reserved_http_server_port,
-                )
-                return p, policy_t, pg, gen_t
-
             def init_nemo_gym():
                 """Spin up NeMo Gym servers against the reserved URL."""
                 return _spinup_nemo_gym([reserved_url], generation_config["model_name"])
 
-            init_tasks = {
-                "megatron": init_megatron_stack,
-                "nemo_gym": init_nemo_gym,
-            }
-            print(f"  ⚡ Init tasks: {', '.join(init_tasks.keys())}", flush=True)
+            # Exactly one task adopts the reserved port: the policy when colocated
+            # (generation wraps it), else the dedicated generation policy.
+            policy_port, generation_port = (
+                (reserved_http_server_port, None)
+                if colocated_inference
+                else (None, reserved_http_server_port)
+            )
+
+            def init_megatron_generation_task(policy_future):
+                """Colocated generation waits; non-colocated inits in parallel."""
+                if colocated_inference:
+                    p, _ = policy_future.result()
+                    return init_megatron_generation(p)
+                return init_megatron_generation(
+                    reserved_http_server_port=generation_port
+                )
+
+            print("  ⚡ Init tasks: policy, megatron_generation, nemo_gym", flush=True)
+            init_tasks_t0 = time.perf_counter()
             try:
-                with ThreadPoolExecutor(max_workers=len(init_tasks)) as executor:
-                    submitted = {k: executor.submit(fn) for k, fn in init_tasks.items()}
-                    results = {k: f.result() for k, f in submitted.items()}
+                with ThreadPoolExecutor(max_workers=3) as executor:
+                    policy_future = executor.submit(
+                        init_policy, reserved_http_server_port=policy_port
+                    )
+                    generation_future = executor.submit(
+                        init_megatron_generation_task, policy_future
+                    )
+                    nemo_gym_future = executor.submit(init_nemo_gym)
+                    policy, policy_time = policy_future.result()
+                    policy_generation, megatron_gen_time = generation_future.result()
+                    if not colocated_inference:
+                        setup_timing_metrics.parallel_wall_time_s = (
+                            time.perf_counter() - init_tasks_t0
+                        )
+                        setup_timing_metrics.parallel_init_enabled = 1.0
+                        # NeMo Gym probes the pre-published endpoint before its future completes.
+                        # A skip-load Megatron endpoint starts only during this initial refit,
+                        # so it must happen while Gym is waiting rather than after it resolves.
+                        init_megatron_weight_synchronizer(policy, policy_generation)
+                    nemo_gym_actor, nemo_gym_time = nemo_gym_future.result()
             finally:
                 ray.kill(port_holder)
 
-            policy, policy_time, policy_generation, megatron_gen_time = results[
-                "megatron"
-            ]
-            nemo_gym_actor, nemo_gym_time = results["nemo_gym"]
+            if colocated_inference:
+                setup_timing_metrics.parallel_init_enabled = 0.0
             setup_timing_metrics.policy_init_time_s = policy_time
             setup_timing_metrics.generation_init_time_s = (
                 reserve_time + megatron_gen_time
@@ -1392,13 +1433,20 @@ def init_nemo_gym():
             setup_timing_metrics.nemo_gym_init_time_s = nemo_gym_time
 
         else:
-            # Initialize training first so checkpoint conversion completes before inference starts.
-            policy, policy_time = init_policy()
-            setup_timing_metrics.policy_init_time_s = policy_time
+            if not colocated_inference:
+                policy_generation, policy = initialize_generation_with_policy(
+                    init_megatron_generation,
+                    colocated_inference,
+                    setup_timing_metrics,
+                )
+            else:
+                # Colocated generation wraps the training policy.
+                policy, policy_time = init_policy()
+                setup_timing_metrics.policy_init_time_s = policy_time
 
-            # Colocated wraps the training policy; non-colocated builds a dedicated inference policy.
-            policy_generation, megatron_gen_time = init_megatron_generation(policy)
-            setup_timing_metrics.generation_init_time_s = megatron_gen_time
+                policy_generation, megatron_gen_time = init_megatron_generation(policy)
+                setup_timing_metrics.generation_init_time_s = megatron_gen_time
+                setup_timing_metrics.parallel_init_enabled = 0.0
 
         print(
             f"  ✓ Using {backend} backend for generation with {policy_config['model_name']}",
@@ -1643,22 +1691,8 @@ def init_dynamo():
         )
 
     if backend == "megatron":
-        t0 = time.perf_counter()
-        policy_generation.weight_synchronizer = create_weight_synchronizer(
-            policy=policy,
-            generation=policy_generation,
-            generation_backend=backend,
-            colocated=colocated_inference,
-            train_cluster=train_cluster,
-            inference_cluster=None if colocated_inference else inference_cluster,
-        )
-        policy_generation.weight_synchronizer.init_communicator()
-        setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0
-        if not colocated_inference:
-            # Load the model weights now.
-            t0 = time.perf_counter()
-            policy_generation.weight_synchronizer.sync_weights()
-            setup_timing_metrics.weight_sync_time_s = time.perf_counter() - t0
+        if policy_generation.weight_synchronizer is None:
+            init_megatron_weight_synchronizer(policy, policy_generation)
         if enable_nemo_gym:
             served_urls = policy_generation.dp_openai_server_base_urls
             if served_urls != [reserved_url]:

nemo_rl/models/generation/megatron/config.py

@@ -70,10 +70,8 @@ class MCoreGenerationSpecificArgs(TypedDict):
     # parity checks should select raw_logprobs explicitly.
     logprobs_mode: Literal["processed_logprobs", "raw_logprobs"]
 
-    # FP8/MXFP8 for the dedicated (non-colocated) inference model; merged into its
-    # `megatron_cfg` by `merged_inference_megatron_cfg`. When enabled, the first
-    # refit re-quantizes the weight buffers, so engine construction (CUDA-graph
-    # capture) is deferred until after that refit (#3731).
+    # FP8/MXFP8 for the dedicated (non-colocated) inference model;
+    # merged into its `megatron_cfg` by `merged_inference_megatron_cfg`.
     fp8_cfg: NotRequired[Fp8Config]
 
 

nemo_rl/models/generation/megatron/megatron_generation.py

@@ -147,7 +147,6 @@ def __init__(
         policy: Optional["Policy"] = None,
         name_prefix: str = "megatron_generation",
         processor: Optional[AutoProcessor] = None,
-        weights_path: Optional[str] = None,
         skip_weight_load: bool = False,
         reserved_http_server_port: Optional[int] = None,
     ):
@@ -162,7 +161,6 @@ def __init__(
             policy: Existing training Policy to reuse for generation.
             name_prefix: Prefix for naming the worker group (non-colocated only).
             processor: Optional processor for VLMs (non-colocated only).
-            weights_path: Optional path to model weights (non-colocated only).
             skip_weight_load: Do not load the weights from the checkpoint; refit will do it.
             reserved_http_server_port: Driver-reserved OpenAI server port for non-colocated.
         """
@@ -214,16 +212,14 @@ def __init__(
             processor=processor,
             init_optimizer=False,
             init_reference_model=False,
-            weights_path=weights_path,
             skip_weight_load=skip_weight_load,
             reserved_http_server_port=reserved_http_server_port,
         )
 
-        # MXFP8 inference re-quantizes weights at the first refit, so CUDA graphs must
-        # capture the post-refit buffers (#3731). Everything else starts the engine +
-        # HTTP server now: the NeMo-Gym overlap blocks on the server URL in setup (#3569).
-        gen_fp8_cfg = self.cfg["mcore_generation_config"].get("fp8_cfg")
-        if not (skip_weight_load and gen_fp8_cfg and gen_fp8_cfg["enabled"]):
+        # Skip-load models do not have their final refit weight buffers yet.
+        # Defer engine initialization so CUDA graphs capture the persistent buffers.
+        # The engine + HTTP server then first come up at the initial refit.
+        if not skip_weight_load:
             self.prepare_for_generation()
 
     def init_collective(

nemo_rl/models/megatron/community_import.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 
 import os
+import shutil
+import threading
+import uuid
 from collections.abc import Iterator
 from contextlib import contextmanager
 from typing import Any, Callable, Optional
@@ -62,6 +65,76 @@ def to_torch_dtype(dtype: str | torch.dtype) -> torch.dtype:
     raise ValueError(f"Unknown dtype: {dtype}")
 
 
+def megatron_conversion_is_complete(pretrained_path: str) -> bool:
+    """Whether a completed HF->Megatron conversion exists at `pretrained_path`."""
+    return os.path.exists(
+        os.path.join(pretrained_path, "iter_0000000", "run_config.yaml")
+    )
+
+
+def publish_megatron_conversion(
+    staging_path: str, pretrained_path: str, *, overwrite: bool = False
+) -> None:
+    """Atomically publish a staged HF->Megatron conversion.
+
+    Args:
+        staging_path: Directory the conversion was saved into.
+        pretrained_path: Final conversion-cache path.
+        overwrite: Replace any existing complete conversion.
+    """
+    displaced: Optional[str] = None
+    for attempt in range(2):
+        if not overwrite and megatron_conversion_is_complete(pretrained_path):
+            # A concurrent producer sharing the cache won; use its artifact.
+            print(
+                f"Completed conversion already published at {pretrained_path}; "
+                "discarding the staged copy.",
+                flush=True,
+            )
+            threading.Thread(
+                target=shutil.rmtree,
+                args=(staging_path,),
+                kwargs={"ignore_errors": True},
+                daemon=True,
+            ).start()
+            break
+        try:
+            os.rename(staging_path, pretrained_path)
+            break
+        except OSError:
+            if attempt == 1:
+                if not overwrite and megatron_conversion_is_complete(pretrained_path):
+                    # A peer published into the slot we freed; use its artifact.
+                    threading.Thread(
+                        target=shutil.rmtree,
+                        args=(staging_path,),
+                        kwargs={"ignore_errors": True},
+                        daemon=True,
+                    ).start()
+                    break
+                raise
+        if not overwrite and megatron_conversion_is_complete(pretrained_path):
+            # A complete artifact raced in between the probe and the rename;
+            # concede to it on the next pass.
+            continue
+        # The final path is occupied (stale partial conversion, or a complete one under overwrite);
+        # displace it, retry the rename, then delete the displaced copy.
+        displaced = f"{pretrained_path}.displaced-{uuid.uuid4().hex[:8]}"
+        try:
+            os.rename(pretrained_path, displaced)
+        except FileNotFoundError:
+            # A concurrent publisher displaced it first; retry from the top.
+            displaced = None
+    if displaced is not None:
+        # Doomed uniquely-named copy nothing reads; delete without blocking startup.
+        threading.Thread(
+            target=shutil.rmtree,
+            args=(displaced,),
+            kwargs={"ignore_errors": True},
+            daemon=True,
+        ).start()
+
+
 @contextmanager
 def _prefer_nvrx_for_dist_ckpt_save():
     """Prefer NVRx async strategy for torch_dist save in HF->Megatron import.
@@ -106,6 +179,8 @@ def import_model_from_hf_name(
     model_post_wrap_hook: Optional[Callable] = None,
     transformer_layer_spec: Optional[ModuleSpec | Callable] = None,
     mamba_stack_spec: Optional[ModuleSpec | Callable] = None,
+    *,
+    overwrite: bool = False,
     **config_overrides: Any,
 ):
     """Import a Hugging Face model into Megatron checkpoint format and save the Megatron checkpoint to the output path.
@@ -123,6 +198,7 @@ def import_model_from_hf_name(
         mamba_stack_spec: Optional Megatron ``ModuleSpec`` (or callable
             returning one) overriding the default Mamba stack spec selected by
             Mamba model providers.
+        overwrite: Publish over an existing complete conversion.
         **config_overrides: Extra keyword arguments forwarded to
             ``AutoBridge.from_hf_pretrained``.
     """
@@ -225,8 +301,35 @@ def import_model_from_hf_name(
     config.num_layers_in_last_pipeline_stage = orig_num_layers_in_last_pipeline_stage
     config.pipeline_dtype = orig_pipeline_dtype
 
+    # Stage the save next to the final path, then atomically rename into place
+    # so concurrent readers of the shared cache never see a partial checkpoint.
+    output_path = os.path.normpath(output_path)
+    output_parent = os.path.dirname(output_path)
+    if output_parent:
+        os.makedirs(output_parent, exist_ok=True)
+    dist_active = (
+        torch.distributed.is_available() and torch.distributed.is_initialized()
+    )
+    staging_token = uuid.uuid4().hex[:8]
+    if dist_active:
+        token_box = [staging_token]
+        torch.distributed.broadcast_object_list(token_box, src=0)
+        staging_token = token_box[0]
+    staging_path = os.path.join(
+        output_parent, f".{os.path.basename(output_path)}.staging-{staging_token}"
+    )
+
     with _prefer_nvrx_for_dist_ckpt_save():
-        bridge.save_megatron_model(megatron_model, output_path)
+        bridge.save_megatron_model(megatron_model, staging_path)
+
+    # Every rank must finish writing before rank 0 publishes the staging dir,
+    # and no rank may read output_path before the rename lands.
+    if dist_active:
+        torch.distributed.barrier()
+    if not dist_active or torch.distributed.get_rank() == 0:
+        publish_megatron_conversion(staging_path, output_path, overwrite=overwrite)
+    if dist_active:
+        torch.distributed.barrier()
 
     # resetting mcore state
     import megatron.core.rerun_state_machine