diff --git a/docs/guides/async-grpo.md b/docs/guides/async-grpo.md
index 64838fab625..94b40c215f7 100644
--- a/docs/guides/async-grpo.md
+++ b/docs/guides/async-grpo.md
@@ -24,7 +24,7 @@ loss_fn:
   use_importance_sampling_correction: true
 ```
 
-3. **Disable colocated inference** (required for async mode):
+3. **Disable colocated inference** (required for async mode with the vLLM backend; the Megatron backend supports colocated async — see `examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml`):
 ```yaml
 policy:
   generation:
@@ -189,7 +189,7 @@ If no `replay_buffer.pt` file is found in the latest checkpoint directory, train
 
 3. **Resource Allocation**: Ensure sufficient GPU memory for both the training and generation clusters
 
-4. **In-Flight Weight Updates**: Enable `in_flight_weight_updates: true` when using `async_engine: true` for updating the weights of vLLM engine during generation. This prevents stalling training pipeline until longest generation finishes and provides significant performance benefits.
+4. **In-Flight Weight Updates**: Enable `in_flight_weight_updates: true` to update engine weights during generation; with vLLM this requires `async_engine: true`, while the Megatron backend is always async-engine. This prevents stalling the training pipeline until the longest generation finishes and provides significant performance benefits.
 
 5. **Recompute KV Cache After Weight Updates**: A user can choose whether to invalidate and recompute KV caches after weight updates by setting the `recompute_kv_cache_after_weight_updates` configuration. This is applicable to async GRPO and independent of in-flight updates.
 
diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml
new file mode 100644
index 00000000000..69abf1b560d
--- /dev/null
+++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml
@@ -0,0 +1,50 @@
+defaults: ../../grpo_math_1B.yaml
+# Async colocated GRPO with Megatron (Nemotron-3-Nano-30B-A3B): training and
+# generation share GPUs/workers; the engine sleeps across training steps and
+# serves the shared weights (resharding into a dedicated inference layout when
+# one is configured). Off-policy data via the replay buffer.
+grpo:
+  num_prompts_per_step: 2
+  num_generations_per_prompt: 8
+  async_grpo:
+    enabled: true
+    max_trajectory_age_steps: 4  # weight versions a rollout may span
+    in_flight_weight_updates: true
+    # For AREAL-style KV invalidation, add recompute_kv_cache_after_weight_updates: true.
+loss_fn:
+  use_importance_sampling_correction: true  # required for off-policy replay data
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated
+  save_period: 100
+policy:
+  model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16
+  tokenizer:
+    name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+  train_global_batch_size: 16
+  train_micro_batch_size: 1
+  logprob_batch_size: 1
+  max_total_sequence_length: 2048
+  megatron_cfg:
+    enabled: true
+    bias_activation_fusion: false
+    tensor_model_parallel_size: 2
+    expert_model_parallel_size: 8
+    sequence_parallel: true
+  dtensor_cfg:
+    enabled: false
+  sequence_packing:
+    enabled: false
+  generation:
+    backend: megatron
+logger:
+  log_dir: logs/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl
+    name: grpo-nanov3-30BA3B-4n4g-megatron_async_colocated
+cluster:
+  gpus_per_node: 4
+  num_nodes: 4
+  segment_size: 2
diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.yaml
new file mode 100644
index 00000000000..78f43ccb3c5
--- /dev/null
+++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.yaml
@@ -0,0 +1,41 @@
+defaults: ../../grpo_math_1B.yaml
+grpo:
+  num_prompts_per_step: 2
+  num_generations_per_prompt: 8
+loss_fn:
+  use_importance_sampling_correction: true  # matches the async sibling so the pair isolates the loop
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated
+  save_period: 100
+policy:
+  model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-Base-BF16
+  tokenizer:
+    name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
+  train_global_batch_size: 16
+  train_micro_batch_size: 1
+  logprob_batch_size: 1
+  max_total_sequence_length: 2048
+  megatron_cfg:
+    enabled: true
+    bias_activation_fusion: false
+    tensor_model_parallel_size: 2
+    expert_model_parallel_size: 8
+    sequence_parallel: true
+  dtensor_cfg:
+    enabled: false
+  sequence_packing:
+    enabled: false
+  generation:
+    backend: megatron
+logger:
+  log_dir: logs/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated
+  wandb_enabled: true
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl
+    name: grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated
+cluster:
+  gpus_per_node: 4
+  num_nodes: 4
+  segment_size: 2
diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index d8d5f26763d..61d7f3468b3 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/nemo_rl/algorithms/grpo.py
@@ -2782,10 +2782,6 @@ def grpo_train(
 
     kv_scales_cache = None  # Cache reused for computed kv scales
 
-    NEED_REFIT = not (
-        isinstance(policy_generation, MegatronGeneration)
-        and master_config.policy["generation"]["colocated"]["enabled"]
-    )
     assert policy_generation is not None
 
     # Check if we need to sync KV cache scales
@@ -2827,7 +2823,7 @@ def grpo_train(
         print("\n🔍 Running initial validation...", flush=True)
         memory_tracker.snapshot_start_of_stage("Initial validation", dir())
 
-        if NEED_REFIT and POLICY_GENERATION_STALE:
+        if POLICY_GENERATION_STALE:
             refit_policy_generation(
                 policy,
                 policy_generation,
@@ -2949,7 +2945,7 @@ def grpo_train(
                     flush=True,
                 )
                 with timer.time("prepare_for_generation/total"):
-                    if NEED_REFIT and POLICY_GENERATION_STALE:
+                    if POLICY_GENERATION_STALE:
                         # Compute KV scales if needed for FP8 quantization
                         if sync_kv_scales and kv_scales_cache is None:
                             print("▶ Computing KV cache scales...", flush=True)
@@ -3448,7 +3444,7 @@ def grpo_train(
                 # Run validation if it's a validation step or last step with val_at_end
                 if should_run_validation:
                     memory_tracker.snapshot_start_of_stage("Validation", dir())
-                    if NEED_REFIT and POLICY_GENERATION_STALE:
+                    if POLICY_GENERATION_STALE:
                         refit_metrics = refit_policy_generation(
                             policy,
                             policy_generation,
@@ -4225,10 +4221,6 @@ def async_grpo_train(
         fit_last_save_time=True,
     )
     timeout.start_iterations()
-    NEED_REFIT = not (
-        isinstance(policy_generation, MegatronGeneration)
-        and master_config.policy["generation"]["colocated"]["enabled"]
-    )
     assert policy_generation is not None
 
     # Training state
@@ -4244,13 +4236,14 @@ def async_grpo_train(
     colocated_inference = master_config.policy["generation"]["colocated"]["enabled"]
     stop_at_validation_threshold = master_config.grpo.stop_at_validation_threshold
     stop_at_validation_metric = master_config.grpo.stop_at_validation_metric
+
+    assert (not colocated_inference) or (
+        isinstance(policy_generation, MegatronGeneration)
+    ), "Colocated async GRPO is only supported for the Megatron generation backend."
+
     # Initialize advantage estimator
     adv_estimator = _create_advantage_estimator(master_config)
 
-    assert not colocated_inference, (
-        "Colocated inference is not supported for async GRPO. Please use non-colocated inference."
-    )
-
     # Calculate minimum buffer size from training requirements
     # In per-prompt buffer mode, one buffer entry is 1 prompt * num_generations_per_prompt
     num_prompts_per_step = master_config.grpo.num_prompts_per_step
@@ -4383,7 +4376,7 @@ def async_grpo_train(
     )
 
     print("⏳ Preparing policy generation for training...", flush=True)
-    if NEED_REFIT and POLICY_GENERATION_STALE:
+    if POLICY_GENERATION_STALE:
         print("🔄 Refitting policy generation with actual model weights...", flush=True)
         try:
             refit_policy_generation(
@@ -4439,7 +4432,9 @@ def async_grpo_train(
                 processor=processor,
             )
             initial_val_metrics = val_metrics
-            policy_generation.finish_generation()
+            # A colocated engine keeps serving between phases (preserves its
+            # KV/prefix cache); the backend makes that call, not the loop.
+            policy_generation.finish_generation(release_gpu=False)
             logger.log_metrics(val_metrics, step, prefix="validation")
             logger.log_metrics(validation_timings, step, prefix="timing/validation")
             if master_config.grpo.debug_payload_metrics:
@@ -4825,6 +4820,14 @@ def async_grpo_train(
                     )
                     train_data.to("cpu")
 
+                generation_logger_metrics = None
+                if policy_generation.blocks_training():
+                    print("⏸️ Pausing colocated engine + collector for training...")
+                    with timer.time("exposed_generation"):
+                        ray.get(trajectory_collector.prepare_for_refit.remote())
+                    generation_logger_metrics = policy_generation.get_logger_metrics()
+                    policy_generation.finish_generation(release_gpu=True)
+
                 # Training phase (same as sync version)
                 skip_prev_logprobs, skip_reference_logprobs = (
                     _resolve_logprob_skip_flags(master_config)
@@ -4956,9 +4959,48 @@ def async_grpo_train(
                         timer=timer,
                     )
 
+                is_last_step = step + 1 == master_config.grpo.max_num_steps
+                should_save_by_step = (
+                    is_last_step
+                    or (step + 1) % master_config.checkpointing["save_period"] == 0
+                    or (ft_save_period is not None and (step + 1) % ft_save_period == 0)
+                )
+                # Checked pre-validation so the wake-deferral below can see it.
+                # A crossing during refit/validation is caught by the lookahead in check_save.
+                should_save_by_timeout = timeout.check_save()
+                will_save_checkpoint = master_config.checkpointing["enabled"] and (
+                    should_save_by_step or should_save_by_timeout
+                )
+                # An early stop (known only after validation) also saves.
+                saving_this_step = will_save_checkpoint
+                # Save-bound colocated steps leave the engine asleep through save with no transfer.
+                defer_wake_for_save = (
+                    policy_generation.blocks_training()
+                    and will_save_checkpoint
+                    and policy_generation.wake_carries_weight_updates()
+                )
+
                 print("🔄 Synchronizing policy weights to trajectory collector…")
-                generation_logger_metrics = None
-                if NEED_REFIT:
+                if defer_wake_for_save:
+                    # Wake-deferral (checkpoint scheduling, which the backend
+                    # cannot see): the engine is about to be saved, so leave it
+                    # asleep; just drop training-only buffers and version-stamp
+                    # the weights. The post-save block wakes it and resumes
+                    # collection.
+                    print("⏸️ Keeping colocated engine asleep for checkpointing...")
+                    # Seed the category with 0.0 (no refit wake happens on
+                    # save-bound steps) so efficiency summaries, which skip
+                    # missing keys, stay comparable across modes.
+                    with timer.time("idle/refit_bubble"):
+                        pass
+                    with timer.time("offload_before_refit"):
+                        policy.offload_before_refit()
+                    POLICY_GENERATION_STALE = False
+                    weight_version += 1
+                    ray.get(
+                        trajectory_collector.set_weight_version.remote(weight_version)
+                    )
+                else:
                     timer.start("idle/refit_bubble")
 
                     # Measure pending-generation wait as exposed_generation time
@@ -4968,7 +5010,8 @@ def async_grpo_train(
 
                     # Collect generation logger metrics for performance reporting
                     # inflight batch sizes and num pending samples are collected from each worker
-                    if policy_generation is not None:
+                    # (colocated collects them before the engine sleeps for training).
+                    if generation_logger_metrics is None:
                         generation_logger_metrics = (
                             policy_generation.get_logger_metrics()
                         )
@@ -5000,7 +5043,6 @@ def async_grpo_train(
 
                 # Validation
                 val_metrics, validation_timings = None, None
-                is_last_step = step + 1 == master_config.grpo.max_num_steps
                 should_run_validation = (
                     val_period > 0
                     and (step + 1) >= val_start_at
@@ -5025,15 +5067,9 @@ def async_grpo_train(
                 # Run validation if it's a validation step or last step with val_at_end
                 if should_run_validation:
                     with timer.time("idle/validation"):
-                        if NEED_REFIT and POLICY_GENERATION_STALE:
-                            refit_metrics = refit_policy_generation(
-                                policy,
-                                policy_generation,
-                                colocated_inference,
-                            )
-                            POLICY_GENERATION_STALE = False
-                        else:
-                            policy_generation.prepare_for_generation()
+                        # No-op on an already-running engine;
+                        # wakes the colocated engine when it stayed asleep for a save-bound step.
+                        policy_generation.prepare_for_generation()
                         val_metrics, validation_timings = validate(
                             policy_generation,
                             val_dataloader,
@@ -5044,7 +5080,22 @@ def async_grpo_train(
                             logger=logger,
                             processor=processor,
                         )
-                        policy_generation.finish_generation()
+                        # An early stop triggers a save; must note before engine wake/resume.
+                        early_stop_message = _validation_early_stop_message(
+                            val_metrics,
+                            stop_at_validation_threshold,
+                            stop_at_validation_metric,
+                        )
+                        saving_this_step = will_save_checkpoint or (
+                            master_config.checkpointing["enabled"]
+                            and early_stop_message is not None
+                        )
+                        # Save-bound steps need the GPUs for checkpointing,
+                        # so the engine must stand down; otherwise a colocated
+                        # engine keeps serving (backend's call).
+                        policy_generation.finish_generation(
+                            release_gpu=saving_this_step
+                        )
                         logger.log_metrics(
                             validation_timings, step + 1, prefix="timing/validation"
                         )
@@ -5059,11 +5110,6 @@ def async_grpo_train(
                                     step + 1,
                                     prefix="validation",
                                 )
-                        early_stop_message = _validation_early_stop_message(
-                            val_metrics,
-                            stop_at_validation_threshold,
-                            stop_at_validation_metric,
-                        )
                         if early_stop_message is not None:
                             # Exit at the end of this step, after checkpointing.
                             print(early_stop_message, flush=True)
@@ -5157,20 +5203,7 @@ def async_grpo_train(
                 consumed_samples += master_config.grpo.num_prompts_per_step
                 timeout.mark_iteration()
 
-                # +1 because step is 0-indexed
-                should_save_by_step = (
-                    is_last_step
-                    # Early stop saves the final state like a last step.
-                    or early_stop_message is not None
-                    or (step + 1) % master_config.checkpointing["save_period"] == 0
-                    or (ft_save_period is not None and (step + 1) % ft_save_period == 0)
-                )
-                # Check if timeout-based checkpointing is enabled in config.
-                should_save_by_timeout = timeout.check_save()
-
-                if master_config.checkpointing["enabled"] and (
-                    should_save_by_step or should_save_by_timeout
-                ):
+                if saving_this_step:
                     grpo_save_state.current_step = step + 1
                     grpo_save_state.total_valid_tokens = total_valid_tokens
                     if val_metrics is not None:
@@ -5266,6 +5299,21 @@ def async_grpo_train(
                             checkpointer, last_checkpoint_step=step + 1
                         )
 
+                    # On save-bound steps, engine stayed asleep after training;
+                    # wake it unless the loop exits right below (last step, timeout, early stop),
+                    # where a wake would only feed the teardown.
+                    # The intervening logging runs with the collector paused either way.
+                    if defer_wake_for_save and not (
+                        is_last_step
+                        or should_save_by_timeout
+                        or early_stop_message is not None
+                    ):
+                        # The save onloaded model+optimizer;
+                        # generation windows must start from the offloaded state.
+                        policy.offload_after_refit()
+                        policy_generation.prepare_for_generation()
+                        ray.get(trajectory_collector.resume_after_refit.remote())
+
             # Logging
             # Log training data (match sync GRPO logging payload for parity).
             # NeMo Gym responses can be very large and expensive to log; when
diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py
index 3d843e36b05..6ddd72e1e25 100644
--- a/nemo_rl/algorithms/grpo_sync.py
+++ b/nemo_rl/algorithms/grpo_sync.py
@@ -80,7 +80,6 @@
 from nemo_rl.environments.interfaces import EnvironmentInterface
 from nemo_rl.experience.sync_rollout_actor import SyncRolloutActor
 from nemo_rl.models.generation.interfaces import GenerationInterface
-from nemo_rl.models.generation.megatron import MegatronGeneration
 from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface
 from nemo_rl.utils.checkpoint import CheckpointManager
 from nemo_rl.utils.logger import Logger, print_message_log_samples
@@ -380,7 +379,7 @@ def _compute_seq_logprob_error_metrics(
 
 def grpo_train_sync(
     policy: ColocatablePolicyInterface,
-    policy_generation: Optional[GenerationInterface],
+    policy_generation: GenerationInterface,
     wrapped_dataloader,
     val_dataloader: Optional[StatefulDataLoader],
     tokenizer,
@@ -415,14 +414,6 @@ def grpo_train_sync(
 
     kv_scales_cache = None  # Cache reused for computed kv scales
 
-    NEED_REFIT = not (
-        isinstance(policy_generation, MegatronGeneration)
-        and master_config.policy["generation"]["colocated"]["enabled"]
-    )
-    # If policy_generation is None, use the policy as the generation interface (megatron framework backend)
-    if policy_generation is None:
-        policy_generation = policy  # type: ignore
-        NEED_REFIT = False
     POLICY_GENERATION_STALE = True
     assert policy_generation is not None
 
@@ -511,7 +502,7 @@ def grpo_train_sync(
         print("\n🔍 Running initial validation...", flush=True)
         memory_tracker.snapshot_start_of_stage("Initial validation", dir())
 
-        if NEED_REFIT and POLICY_GENERATION_STALE:
+        if POLICY_GENERATION_STALE:
             refit_policy_generation(policy, policy_generation, colocated_inference)
             POLICY_GENERATION_STALE = False
         else:
@@ -580,8 +571,7 @@ def grpo_train_sync(
                 )
 
             maybe_gpu_profile_step(policy, total_steps + 1)
-            if policy != policy_generation:
-                maybe_gpu_profile_step(policy_generation, total_steps + 1)
+            maybe_gpu_profile_step(policy_generation, total_steps + 1)
             val_metrics, validation_timings = None, None
 
             with timer.time("total_step_time"):
@@ -599,7 +589,7 @@ def grpo_train_sync(
                     flush=True,
                 )
                 with timer.time("prepare_for_generation/total"):
-                    if NEED_REFIT and POLICY_GENERATION_STALE:
+                    if POLICY_GENERATION_STALE:
                         if sync_kv_scales and kv_scales_cache is None:
                             # KV-scale calibration uses message_log of the
                             # current step's PROMPTS (pre-generation), which
@@ -1011,7 +1001,7 @@ def grpo_train_sync(
                     and (total_steps + 1) % val_period == 0
                 ) or (val_at_end and is_last_step):
                     memory_tracker.snapshot_start_of_stage("Validation", dir())
-                    if NEED_REFIT and POLICY_GENERATION_STALE:
+                    if POLICY_GENERATION_STALE:
                         refit_policy_generation(
                             policy,
                             policy_generation,
diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py
index f8c7d25ca02..0867bc121ed 100644
--- a/nemo_rl/models/generation/interfaces.py
+++ b/nemo_rl/models/generation/interfaces.py
@@ -372,10 +372,24 @@ def generate(
 
     @abstractmethod
     def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool:
+        """Ready the engine for a generation phase (start or wake it).
+
+        Idempotent wake: calling this on an already-running engine must be safe and cheap.
+        """
         pass
 
     @abstractmethod
     def finish_generation(self, *args: Any, **kwargs: Any) -> bool:
+        """Wind down after a generation phase.
+
+        Callers may pass `release_gpu` (keyword-only, default True):
+        True means the caller needs the GPUs for itself (a training step or a checkpoint save),
+        so even a colocated engine must fully stand down;
+        False means the phase is merely over, and a colocated engine must keep serving
+        usable with no intervening prepare_for_generation.
+        Only the colocated Megatron backend honors the flag today; other backends
+        ignore it, as do engines on dedicated GPUs.
+        """
         pass
 
     @abstractmethod
@@ -438,6 +452,26 @@ def attach_fleet_health(self, monitor: Any, selector: Any) -> None:
     def invalidate_kv_cache(self) -> bool:
         return False
 
+    def blocks_training(self) -> bool:
+        """Whether this engine must stand down before a training step.
+
+        True when generation shares GPUs with training (colocated): the
+        training loop then pauses collection and winds the engine down
+        before training. Engines on dedicated GPUs never block training.
+        """
+        return False
+
+    def wake_carries_weight_updates(self) -> bool:
+        """Whether prepare_for_generation alone serves the latest weights.
+
+        True when waking the engine suffices for it to serve weights updated while it slept
+        (colocated Megatron: the wake reshards, or the engine shares the training tensors outright).
+        The async loop may then defer a wake past a checkpoint save and advance
+        the collector's weight version with no explicit transfer.
+        Backends whose wake does not reload weights must return False so the loop refits instead.
+        """
+        return False
+
     def clear_logger_metrics(self) -> None:
         """Clear logger metrics for performance reporting.
 
diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py
index 4d2511523b1..234357f0a37 100644
--- a/nemo_rl/models/generation/megatron/config.py
+++ b/nemo_rl/models/generation/megatron/config.py
@@ -75,11 +75,26 @@ class MCoreGenerationConfig(GenerationConfig):
 def merged_inference_megatron_cfg(policy_config: PolicyConfig) -> dict[str, Any]:
     """The `megatron_cfg` a dedicated inference model runs with."""
     generation_config = cast(MCoreGenerationConfig, policy_config["generation"])
-    return {
+    merged: dict[str, Any] = {
         **cast(dict[str, Any], policy_config["megatron_cfg"]),
         **(generation_config.get("mcore_generation_config") or {}),
         "activation_checkpointing": False,
     }
+    # inference_optimized layers hard-require SP with TP>1. Raise with the
+    # config key: the colocated build bypasses validate_and_set_config, so this
+    # merge is the only spot the inference cfg gets a named error instead of a
+    # raw MCore assert at model build.
+    if (
+        merged.get("transformer_impl") == "inference_optimized"
+        and merged["tensor_model_parallel_size"] > 1
+        and not merged["sequence_parallel"]
+    ):
+        raise ValueError(
+            "transformer_impl=inference_optimized requires sequence parallelism "
+            "with TP>1 on the generation model: set "
+            "policy.generation.mcore_generation_config.sequence_parallel=true."
+        )
+    return merged
 
 
 def dedicated_inference_megatron_cfg(
@@ -92,8 +107,8 @@ def dedicated_inference_megatron_cfg(
     builds a second model and reshards into it on every wake. Inference never
     uses CP, so CP is pinned to 1 (CP>1 training therefore always differs).
 
-    Returns None when the resolved config matches training (dual-mode: generate
-    directly on the shared training model).
+    Returns None when the resolved config matches training (reshardless:
+    generate directly on the shared training model).
     """
     inference_mcfg = merged_inference_megatron_cfg(policy_config)
     inference_mcfg["context_parallel_size"] = 1
diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py
index d21c40bb52a..4e71222a0b3 100644
--- a/nemo_rl/models/generation/megatron/megatron_generation.py
+++ b/nemo_rl/models/generation/megatron/megatron_generation.py
@@ -259,14 +259,43 @@ def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool:
             ]
         return True
 
-    def finish_generation(self, *args: Any, **kwargs: Any) -> bool:
-        """Clean up after generation."""
+    def finish_generation(self, *, release_gpu: bool = True) -> bool:
+        """Clean up after generation.
+
+        When `release_gpu` is False, a colocated engine keeps serving instead of standing down.
+        """
         futures = self._policy.worker_group.run_all_workers_single_data(
-            "finish_generation"
+            "finish_generation", release_gpu=release_gpu
         )
         ray.get(futures)
         return True
 
+    def blocks_training(self) -> bool:
+        """Whether the engine must stand down before a training step.
+
+        Colocated generation shares the training GPUs, so the training
+        loop must wind the engine down before it can train.
+        """
+        return bool(self.cfg["colocated"]["enabled"])
+
+    def wake_carries_weight_updates(self) -> bool:
+        """The colocated wake reshards (or shares tensors); see the ABC."""
+        return bool(self.cfg["colocated"]["enabled"])
+
+    def invalidate_kv_cache(self) -> bool:
+        """Report whether weight updates invalidate the KV cache.
+
+        Under "recompute" mode the engine drops and rebuilds its KV cache
+        across the suspend/resume that brackets every weight update, so
+        invalidation is genuinely handled; report it truthfully instead of
+        inheriting the interface's `False` (which makes the trajectory
+        collector warn every step).
+        """
+        return (
+            self.cfg["mcore_generation_config"].get("kv_cache_management_mode")
+            == "recompute"
+        )
+
     def preinit_nvshmem_collective(self) -> list[ray.ObjectRef]:
         """Pre-initialize NVShmem collectively after CUDA graph capture.
 
diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py
index 2bbd1014a84..ea08a8483bb 100644
--- a/nemo_rl/models/generation/megatron/megatron_worker.py
+++ b/nemo_rl/models/generation/megatron/megatron_worker.py
@@ -368,8 +368,21 @@ def _run_async_coordinator_start(self):
             print(f"[Rank {torch.distributed.get_rank()}] HTTP Server not started")
             self.base_url = None
 
-    def finish_generation(self) -> None:
-        """Wind down a generation cycle."""
+    def finish_generation(self, *, release_gpu: bool = True) -> None:
+        """Wind down a generation cycle.
+
+        Args:
+            release_gpu: the caller needs the GPUs for itself (a training
+                step or a checkpoint save), so even a colocated engine must
+                fully stand down. Pass False between generation phases (e.g.
+                after validation) to let a colocated engine keep serving;
+                tearing it down there would discard KV/prefix caches and CUDA
+                graphs for no reason.
+                For non-colocated engines this body reduces to a rotary-cache clear
+                (their engine pause happens in suspend_for_refit).
+        """
+        if self.is_generation_colocated and not release_gpu:
+            return
         print(f"[Rank {self.rank}] finishing generation", flush=True)
         log_gpu_memory("finish_generation START")
 
@@ -402,10 +415,24 @@ def finish_generation(self) -> None:
     def prepare_for_generation(self, tags=None, **kwargs) -> None:
         """Enter inference mode and start (or wake) the inference engine.
 
+        Idempotent wake: a plain call (no tags) on an already-serving engine returns immediately.
+        Refit-protocol calls (tags) are never skipped.
+
         Called in both colocated and non-colocated setups.
         Even in non-colocated mode, Megatron's engine has to be intentionally paused before a refit
         (and its weights are not detachable), so we have to switch modes around every refit.
         """
+        if (
+            tags is None
+            and self._inference_engine_initialized
+            and not self._inference_engine_asleep
+        ):
+            print(
+                f"[Rank {self.rank}] prepare_for_generation: engine already "
+                "serving, skipping",
+                flush=True,
+            )
+            return
         log_gpu_memory("prepare_for_generation START")
         mcore_generation_config = self.cfg["generation"]["mcore_generation_config"]
 
@@ -939,7 +966,6 @@ def suspend_for_refit(self) -> None:
         if not self._inference_engine_initialized:
             return
         self._sleep()
-        torch.cuda.synchronize()
 
     def resume_after_refit(self) -> None:
         """Resume+unpause the inference engine after a weight refit."""
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
index accee9ec08d..d7e3fdb0175 100644
--- a/nemo_rl/models/megatron/setup.py
+++ b/nemo_rl/models/megatron/setup.py
@@ -329,6 +329,19 @@ def validate_and_set_config(
     weights_path,
     optimizer_path,
 ):
+    # inference_optimized layers hard-require SP with TP>1; fail here with the config key.
+    # This guards the training cfg; the inference cfg is guarded in
+    # merged_inference_megatron_cfg (the colocated build bypasses this function).
+    if (
+        config["megatron_cfg"].get("transformer_impl") == "inference_optimized"
+        and config["megatron_cfg"]["tensor_model_parallel_size"] > 1
+        and not config["megatron_cfg"]["sequence_parallel"]
+    ):
+        raise ValueError(
+            "transformer_impl=inference_optimized requires sequence parallelism "
+            "with TP>1: set policy.megatron_cfg.sequence_parallel=true."
+        )
+
     # Handle generation configuration
     is_generation_colocated = None
     sampling_params = None
diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py
index 9506422ad32..b4048602bb3 100644
--- a/nemo_rl/models/policy/__init__.py
+++ b/nemo_rl/models/policy/__init__.py
@@ -397,7 +397,9 @@ class MegatronConfig(TypedDict):
     # (used when transformer_impl='inference_optimized')
     moe_router_num_groups: NotRequired[int | None]
     moe_router_group_topk: NotRequired[int | None]
-    # Transformer implementation backing the model. Only valid on generation workers.
+    # Transformer implementation backing the model. 'inference_optimized'
+    # trains through the TE parent path and requires sequence_parallel with
+    # TP>1 (enforced at setup).
     # Options are 'transformer_engine' and 'inference_optimized'.
     transformer_impl: NotRequired[str]
     # CUDA-graph implementation.
diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index e1b96eeced6..fc913dbb8ce 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -496,14 +496,6 @@ def __init__(
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
         # Step 3: Setup model configuration
-        # Training workers cannot use inference_optimized transformer spec.
-        if init_optimizer:
-            assert (
-                config["megatron_cfg"].get("transformer_impl") != "inference_optimized"
-            ), (
-                "transformer_impl=inference_optimized must not be set on training workers. "
-                "Use policy.generation.mcore_generation_config.transformer_impl=inference_optimized instead."
-            )
         runtime_config = validate_and_set_config(
             config,
             self.rank,
diff --git a/nemo_rl/weight_sync/megatron_weight_synchronizer.py b/nemo_rl/weight_sync/megatron_weight_synchronizer.py
index 4688ba26fb0..dab017fce56 100644
--- a/nemo_rl/weight_sync/megatron_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/megatron_weight_synchronizer.py
@@ -25,7 +25,7 @@ class MegatronWeightSynchronizer(WeightSynchronizer):
     """Weight synchronization for the Megatron generation backend, both colocation modes.
 
     Colocated is the degenerate path: generation either aliases the training
-    weights outright (dual-mode) or re-partitions them into the worker's
+    weights outright (reshardless) or re-partitions them into the worker's
     dedicated inference model inside ``prepare_for_generation`` (when the
     configured inference layout/impl differs) — a genuine parallelism-changing
     transfer, but one the worker performs internally on wake. Sync therefore
@@ -98,11 +98,13 @@ def sync_weights(
         kv_scales: Optional[dict[str, float]] = None,
     ) -> Optional[dict[str, float]]:
         if self._colocated:
-            # The wake below carries any configured reshard; the loop already
-            # slept the engine before training (or it has not started yet),
-            # so no suspend is needed.
+            # The wake below carries any configured reshard; the loop already slept the engine
+            # before training, so no suspend is needed.
+            # Tagging the call bypasses the worker's engine-awake early-return, so the reshard
+            # copy riding this wake cannot be skipped. Any tag except "weights" works: the worker
+            # treats "weights" as the wake-suppressing mid-refit call.
             self._policy.offload_before_refit()
-            self._generation.prepare_for_generation()
+            self._generation.prepare_for_generation(tags=["colocated_refit"])
             self._stale = False
             return {}
 
diff --git a/tests/functional/L1_Functional_Tests_Megatron_4.sh b/tests/functional/L1_Functional_Tests_Megatron_4.sh
index 2462db4d6e2..484300a0ccf 100644
--- a/tests/functional/L1_Functional_Tests_Megatron_4.sh
+++ b/tests/functional/L1_Functional_Tests_Megatron_4.sh
@@ -53,6 +53,8 @@ if megatron_generation_supported; then
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_topology.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_non_colocated.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_reshard.sh
+    run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_async_grpo.sh
+    run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_async_gym.sh
     run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_topp_topk.sh
     # Disabled: token_mult_prob_error ~2.0 > 1.1 under top_p/top_k after the
diff --git a/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh b/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh
new file mode 100644
index 00000000000..f58d0138f5b
--- /dev/null
+++ b/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh
@@ -0,0 +1,73 @@
+#!/bin/bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetches metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
+
+set -eou pipefail
+
+EXP_NAME=$(basename $0 .sh)
+EXP_DIR=$SCRIPT_DIR/$EXP_NAME
+LOG_DIR=$EXP_DIR/logs
+JSON_METRICS=$EXP_DIR/metrics.json
+RUN_LOG=$EXP_DIR/run.log
+CKPT_DIR=$EXP_DIR/ckpts
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
+
+rm -rf $EXP_DIR $LOG_DIR
+mkdir -p $EXP_DIR $LOG_DIR $CKPT_DIR
+
+# clean up checkpoint directory on exit
+trap "rm -rf $CKPT_DIR" EXIT
+
+# async colocated
+cd $PROJECT_ROOT
+uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
+    $PROJECT_ROOT/examples/run_grpo.py \
+    --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \
+    policy.model_name=Qwen/Qwen2.5-0.5B \
+    grpo.num_prompts_per_step=2 \
+    grpo.num_generations_per_prompt=4 \
+    policy.train_global_batch_size=4 \
+    policy.logprob_batch_size=4 \
+    policy.train_micro_batch_size=1 \
+    policy.generation.backend=megatron \
+    grpo.async_grpo.enabled=true \
+    grpo.async_grpo.max_trajectory_age_steps=1 \
+    grpo.async_grpo.in_flight_weight_updates=true \
+    loss_fn.use_importance_sampling_correction=true \
+    grpo.max_num_steps=3 \
+    grpo.val_period=3 \
+    grpo.max_val_samples=8 \
+    grpo.val_batch_size=8 \
+    cluster.gpus_per_node=2 \
+    logger.tensorboard_enabled=true \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=false \
+    logger.monitor_gpus=true \
+    checkpointing.enabled=true \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    checkpointing.save_period=2 \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Smoke-level threshold (matches grpo_megatron_generation_async_gym.sh); tighten after CI runs.
+uv run tests/check_metrics.py $JSON_METRICS \
+    'median(data["train/gen_kl_error"]) < 1.3' \
+    '"3" in data["train/loss"]' \
+    '"3" in data["validation/accuracy"]'
+
+# The save-bound step must defer the engine wake past the checkpoint save.
+# `val_period=3` gives us a step 2 that does not wake/sleep cycle the engine before save.
+if ! grep -q "Keeping colocated engine asleep for checkpointing" $RUN_LOG; then
+    echo "FAIL: deferred-wake log line not found (colocated checkpoint path not exercised)"
+    exit 1
+fi
+
+if [[ ! -f $CKPT_DIR/step_2/replay_buffer.pt ]]; then
+    echo "FAIL: replay_buffer.pt not found in step_2 checkpoint"
+    exit 1
+fi
diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh b/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh
new file mode 100644
index 00000000000..62c8e57f862
--- /dev/null
+++ b/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh
@@ -0,0 +1,92 @@
+#!/bin/bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetches metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
+
+set -eou pipefail
+
+EXP_NAME=$(basename $0 .sh)
+EXP_DIR=$SCRIPT_DIR/$EXP_NAME
+LOG_DIR=$EXP_DIR/logs
+JSON_METRICS=$EXP_DIR/metrics.json
+RUN_LOG=$EXP_DIR/run.log
+CKPT_DIR=$EXP_DIR/ckpts
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
+
+rm -rf $EXP_DIR $LOG_DIR
+mkdir -p $EXP_DIR $LOG_DIR $CKPT_DIR
+
+# clean up checkpoint directory on exit
+trap "rm -rf $CKPT_DIR" EXIT
+
+# async colocated, reshard mode: TE TP2 training; inference_optimized TP1
+# generation on a dedicated model, resharded into on every wake.
+cd $PROJECT_ROOT
+uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
+    $PROJECT_ROOT/examples/run_grpo.py \
+    --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \
+    policy.model_name=Qwen/Qwen3-0.6B \
+    grpo.num_prompts_per_step=2 \
+    grpo.num_generations_per_prompt=4 \
+    policy.train_global_batch_size=4 \
+    policy.logprob_batch_size=4 \
+    policy.train_micro_batch_size=1 \
+    policy.megatron_cfg.tensor_model_parallel_size=2 \
+    policy.generation.backend=megatron \
+    ++policy.generation.mcore_generation_config.transformer_impl=inference_optimized \
+    ++policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \
+    policy.generation.mcore_generation_config.refit_backend=nccl \
+    grpo.async_grpo.enabled=true \
+    grpo.async_grpo.max_trajectory_age_steps=1 \
+    grpo.async_grpo.in_flight_weight_updates=true \
+    loss_fn.use_importance_sampling_correction=true \
+    grpo.max_num_steps=3 \
+    grpo.val_period=1 \
+    grpo.max_val_samples=8 \
+    grpo.val_batch_size=8 \
+    cluster.gpus_per_node=2 \
+    logger.tensorboard_enabled=true \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=false \
+    logger.monitor_gpus=true \
+    checkpointing.enabled=true \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    checkpointing.save_period=2 \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Smoke-level threshold (matches grpo_megatron_generation_async_gym.sh); tighten after CI runs.
+uv run tests/check_metrics.py $JSON_METRICS \
+    'median(data["train/gen_kl_error"]) < 1.3' \
+    '"3" in data["train/loss"]' \
+    '"2" in data["validation/accuracy"]'
+
+# The dedicated inference model must actually be built — guard against this
+# leg silently degenerating to the matched-impl (reshardless) path.
+if ! grep -q "\[colocated-reshard\] building dedicated inference model" $RUN_LOG; then
+    echo "FAIL: dedicated-model build log line not found (reshard path not exercised)"
+    exit 1
+fi
+
+# The non-save validation (step 1) wakes an already-serving engine; the
+# worker must skip it (guards against redundant per-validation resharding).
+if ! grep -q "prepare_for_generation: engine already serving, skipping" $RUN_LOG; then
+    echo "FAIL: idempotent-wake skip log line not found (redundant reshard on validation?)"
+    exit 1
+fi
+
+# The save-bound step must defer the engine wake past the checkpoint save.
+# With `val_period=1`, the validation always intervenes before the save.
+if ! grep -q "Keeping colocated engine asleep for checkpointing" $RUN_LOG; then
+    echo "FAIL: deferred-wake log line not found (colocated checkpoint path not exercised)"
+    exit 1
+fi
+
+if [[ ! -f $CKPT_DIR/step_2/replay_buffer.pt ]]; then
+    echo "FAIL: replay_buffer.pt not found in step_2 checkpoint"
+    exit 1
+fi
diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh
new file mode 100755
index 00000000000..2bbe91651a4
--- /dev/null
+++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=4
+GPUS_PER_NODE=4
+SEGMENT_SIZE=2
+STEPS_PER_RUN=10
+MAX_STEPS=10
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=60
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Run the experiment
+cd $PROJECT_ROOT
+uv run examples/run_grpo.py \
+    --config $CONFIG_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=False \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+uv run tests/check_metrics.py $JSON_METRICS \
+    'max(data["train/reward"]) > 0.0' \
+    'median(data["train/gen_kl_error"]) < 1.3'
diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.sh
new file mode 100755
index 00000000000..2bbe91651a4
--- /dev/null
+++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=4
+GPUS_PER_NODE=4
+SEGMENT_SIZE=2
+STEPS_PER_RUN=10
+MAX_STEPS=10
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=60
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Run the experiment
+cd $PROJECT_ROOT
+uv run examples/run_grpo.py \
+    --config $CONFIG_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=False \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+uv run tests/check_metrics.py $JSON_METRICS \
+    'max(data["train/reward"]) > 0.0' \
+    'median(data["train/gen_kl_error"]) < 1.3'
diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt
index aa2dccea837..0cc19199652 100644
--- a/tests/test_suites/nightly_gb200.txt
+++ b/tests/test_suites/nightly_gb200.txt
@@ -94,3 +94,5 @@ tests/test_suites/llm/distillation-qwen3-32b-to-1.7b-base-1n4g-megatron-tp1pp2cp
 # Nano3 hybrid MoE/Mamba ModelOpt layer-spec smoke. Keeps
 # policy.disable_modelopt_layer_spec=false to cover modelopt_mamba_stack_spec.
 tests/test_suites/llm/distillation-nano3-30ba3b-4n4g-megatron-qa-nvfp4-modelopt-spec.sh
+tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh
+tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.sh
diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py
index 51c9d4290ed..e0bba7874d9 100644
--- a/tests/unit/algorithms/test_grpo.py
+++ b/tests/unit/algorithms/test_grpo.py
@@ -91,6 +91,11 @@ def _mock_policy_generation() -> MagicMock:
     policy_generation = MagicMock(spec=MegatronGeneration)
     policy_generation.requires_kv_scale_sync = False
     policy_generation.get_logger_metrics.return_value = {}
+    policy_generation.blocks_training.return_value = False
+    policy_generation.wake_carries_weight_updates.return_value = False
+    policy_generation.weight_synchronizer = MagicMock()
+    policy_generation.weight_synchronizer.is_stale = True
+    policy_generation.weight_synchronizer.sync_weights.return_value = {}
     return policy_generation
 
 
@@ -2958,6 +2963,11 @@ def test_grpo_train_shutdown_on_epoch_completion(mock_grpo_components, tmp_path)
             "nemo_rl.algorithms.grpo.compute_and_apply_seq_logprob_error_masking",
             return_value=_mock_seq_logprob_error_result(),
         ),
+        # Refit runs unconditionally when generation is stale.
+        patch(
+            "nemo_rl.algorithms.grpo.refit_policy_generation",
+            return_value={},
+        ),
         patch("nemo_rl.algorithms.grpo.torch.save"),
     ):
         grpo_mod.grpo_train(
@@ -3070,6 +3080,106 @@ def test_grpo_ft_save_period_triggers_periodic_saves(
     assert saved_steps == [2, 4, 5]
 
 
+def test_async_grpo_colocated_save_defers_wake_until_after_checkpoint(
+    mock_grpo_components, tmp_path
+):
+    """Colocated save steps keep the engine asleep through the checkpoint.
+
+    With a backend that blocks training and whose wake carries the weight
+    updates (colocated Megatron), a save-bound step must version-stamp the
+    weights with the engine still asleep, save, and only then wake the engine
+    and resume collection. The final step skips the wake (the loop exits).
+    """
+    mock_batch = next(iter(mock_grpo_components["train_dataloader"]))
+    mock_rollout_metrics = {"mean_gen_tokens_per_sample": 2.0}
+    policy = mock_grpo_components["policy"]
+    checkpointer = mock_grpo_components["checkpointer"]
+
+    master_config = mock_grpo_components["master_config"]
+    master_config.grpo.max_num_steps = 3
+    master_config.grpo.max_num_epochs = 1
+    master_config.grpo.val_period = 0
+    master_config.grpo.val_at_start = False
+    master_config.grpo.val_at_end = False
+    master_config.grpo.use_dynamic_sampling = False
+    master_config.checkpointing["enabled"] = True
+    # Step 2 saves via save_period; step 3 saves as the last step.
+    master_config.checkpointing["save_period"] = 2
+    master_config.checkpointing["metric_name"] = None
+    master_config.policy["generation"]["colocated"]["enabled"] = True
+
+    events = []
+    policy_generation = _mock_policy_generation()
+    policy_generation.blocks_training.return_value = True
+    policy_generation.wake_carries_weight_updates.return_value = True
+    policy_generation.finish_generation.side_effect = lambda *a, **k: events.append(
+        ("finish_generation", k.get("release_gpu", True))
+    )
+    policy_generation.prepare_for_generation.side_effect = (
+        lambda *a, **k: events.append("wake_engine")
+    )
+    policy.offload_before_refit.side_effect = lambda *a, **k: events.append(
+        "offload_before_refit"
+    )
+    policy.offload_after_refit.side_effect = lambda *a, **k: events.append(
+        "offload_after_refit"
+    )
+
+    def record_save(step, *args, **kwargs):
+        events.append(("save", step))
+        return "/tmp/checkpoint"
+
+    checkpointer.init_tmp_checkpoint.side_effect = record_save
+    checkpointer.checkpoint_dir = tmp_path
+
+    with (
+        mock_async_grpo_infrastructure(
+            mock_batch, mock_rollout_metrics, collector_events=events
+        ),
+        _patched_logprob_phase(policy),
+        patch("nemo_rl.algorithms.grpo.torch.save"),
+    ):
+        async_grpo_train(
+            policy,
+            policy_generation,
+            mock_grpo_components["train_dataloader"],
+            mock_grpo_components["val_dataloader"],
+            mock_grpo_components["tokenizer"],
+            mock_grpo_components["loss_fn"],
+            mock_grpo_components["task_to_env"],
+            mock_grpo_components["val_task_to_env"],
+            mock_grpo_components["logger"],
+            checkpointer,
+            _initial_grpo_save_state(),
+            master_config,
+        )
+
+    assert events == [
+        # Startup: the initial refit is patched out; the collector still gets
+        # the version stamp before collection starts.
+        "set_weight_version",
+        "start_collection",
+        # Step 1 (no save): stand down for training, then refit-arm stamp + resume.
+        ("finish_generation", True),
+        "set_weight_version",
+        "resume_after_refit",
+        # Step 2 (save-bound): version-stamp with the engine asleep, save,
+        # then wake and resume.
+        ("finish_generation", True),
+        "offload_before_refit",
+        "set_weight_version",
+        ("save", 2),
+        "offload_after_refit",
+        "wake_engine",
+        "resume_after_refit",
+        # Step 3 (last step saves): same deferral, but no wake — the loop exits.
+        ("finish_generation", True),
+        "offload_before_refit",
+        "set_weight_version",
+        ("save", 3),
+    ]
+
+
 @pytest.mark.parametrize("train_func", [grpo_train, async_grpo_train])
 def test_grpo_train_skips_reference_policy_logprobs(mock_grpo_components, train_func):
     """Regression test for issue #1968 (Bug 1) and PRs #2174 / #2178.
diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py
index 351f6764462..3158bbdc2be 100644
--- a/tests/unit/models/generation/test_megatron_generation.py
+++ b/tests/unit/models/generation/test_megatron_generation.py
@@ -23,11 +23,15 @@
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.virtual_cluster import RayVirtualCluster
 from nemo_rl.models.generation.megatron import MegatronGeneration
+from nemo_rl.models.generation.megatron.config import (
+    dedicated_inference_megatron_cfg,
+)
 from nemo_rl.models.policy import PolicyConfig
 from nemo_rl.models.policy.lm_policy import Policy
 from nemo_rl.weight_sync.megatron_weight_synchronizer import (
     MegatronWeightSynchronizer,
 )
+from tests.unit.test_utils import SimpleLossFn
 
 model_name = "Qwen/Qwen3-0.6B"
 
@@ -382,11 +386,42 @@ async def test_megatron_policy_generation_async(cluster, test_input_data, tokeni
 
 @pytest.mark.mcore
 @pytest.mark.timeout(900)
-def test_megatron_generation_colocated(cluster, test_input_data, tokenizer):
+@pytest.mark.parametrize(
+    "train_impl, gen_impl",
+    [
+        ("transformer_engine", "transformer_engine"),
+        ("transformer_engine", "inference_optimized"),
+        ("inference_optimized", "inference_optimized"),
+    ],
+)
+def test_megatron_generation_colocated(
+    cluster, test_input_data, tokenizer, train_impl, gen_impl
+):
     """Colocated Megatron generation: wrap an existing training policy without owning it."""
     config = deepcopy(basic_megatron_test_config)
     config["generation"]["colocated"]["enabled"] = True
-    config["generation"]["mcore_generation_config"]["expose_http_server"] = True
+    # Eager engine startup (expose_http_server) flips MLM's process-wide
+    # InferenceMode on at construction; the legs that train before any
+    # generate/suspend cycle must construct engine-less.
+    expose_http_server = (
+        train_impl == "transformer_engine" and gen_impl == "transformer_engine"
+    )
+    config["generation"]["mcore_generation_config"]["expose_http_server"] = (
+        expose_http_server
+    )
+    # Matched impls => reshardless colocated (shared model; the
+    # inference_optimized pair trains through the TE parent path); differing
+    # impls => the worker builds a dedicated resharded inference model on
+    # the shared GPUs.
+    config["megatron_cfg"]["transformer_impl"] = train_impl
+    config["generation"]["mcore_generation_config"]["transformer_impl"] = gen_impl
+    if train_impl == "inference_optimized":
+        # The parity block's sleep/wake cycle would tear down and recapture
+        # CUDA graphs mid-test; keep them off here.
+        config["generation"]["mcore_generation_config"]["cuda_graph_impl"] = "none"
+        # The parity block's 2-sample batch shards to 1 sample per DP rank
+        # (DP=2 on the 2-GPU cluster); the logprob microbatch must divide it.
+        config["logprob_batch_size"] = 1
 
     # construction guard: exactly one of `cluster` / `policy` is required
     with pytest.raises(AssertionError):
@@ -408,16 +443,78 @@ def test_megatron_generation_colocated(cluster, test_input_data, tokenizer):
         assert "max_tokens" not in config["megatron_cfg"]
         assert config["megatron_cfg"] == megatron_cfg_before
 
-        # setup() hands dp_openai_server_base_urls to NeMo Gym right after
-        # construction, so the colocated constructor must have collected them.
-        assert mg.dp_openai_server_base_urls, "no OpenAI server URLs collected"
-        assert all(url.startswith("http") for url in mg.dp_openai_server_base_urls)
+        # The selector: matched impls => reshardless (no dedicated model).
+        assert (dedicated_inference_megatron_cfg(config) is None) == (
+            train_impl == gen_impl
+        )
+
+        if expose_http_server:
+            # setup() hands dp_openai_server_base_urls to NeMo Gym right after
+            # construction, so the colocated constructor must have collected them.
+            assert mg.dp_openai_server_base_urls, "no OpenAI server URLs collected"
+            assert all(url.startswith("http") for url in mg.dp_openai_server_base_urls)
+
+        if gen_impl == "inference_optimized":
+            # Both inference_optimized legs take a train step (finite loss)
+            # before the engine ever starts; the reshard leg then generates
+            # on the dedicated model built at first wake, the matched-impl
+            # leg directly on the shared training model.
+            torch.manual_seed(42)
+            train_data = BatchedDataDict(
+                {
+                    "input_ids": torch.randint(0, 32000, (4, 64)),
+                    "input_lengths": torch.full((4,), 64, dtype=torch.int32),
+                    "attention_mask": torch.ones(4, 64),
+                    "labels": torch.randint(0, 32000, (4, 64)),
+                    "sample_mask": torch.ones(4),
+                }
+            )
+            policy.prepare_for_training()
+            loss = policy.train(train_data, SimpleLossFn())["loss"]
+            assert not torch.isnan(loss).any() and not torch.isinf(loss).any(), (
+                f"pre-generation train step produced bad loss: {loss}"
+            )
+            policy.finish_training()
 
         # re-entering generation mode must be a no-op on the running engine
         mg.prepare_for_generation()
         outputs = mg.generate(test_input_data, greedy=True)
         _assert_valid_generation_output(outputs, test_input_data)
 
+        if train_impl == "inference_optimized":
+            # 3490-review follow-up: bound token mult-prob error on the
+            # matched-impl leg — generation and recomputed policy logprobs
+            # run the same inference kernels on the same shared weights.
+            # Greedy must be off: processed logprobs are ~0 under top_k=1.
+            sampled = mg.generate(test_input_data, greedy=False)
+            fprop_data = BatchedDataDict(
+                {
+                    "input_ids": sampled["output_ids"],
+                    "input_lengths": sampled["unpadded_sequence_lengths"],
+                }
+            )
+            # Production ordering: the engine stands down before any training-path forward.
+            mg.finish_generation(release_gpu=True)
+            policy.prepare_for_lp_inference()
+            lp_logprobs = policy.get_logprobs(fprop_data)["logprobs"]
+            gen_mask = torch.zeros_like(sampled["logprobs"], dtype=torch.bool)
+            for i, (start, end) in enumerate(
+                zip(
+                    test_input_data["input_lengths"],
+                    sampled["unpadded_sequence_lengths"],
+                )
+            ):
+                gen_mask[i, start:end] = True
+            abs_diff = (sampled["logprobs"] - lp_logprobs).abs().masked_select(gen_mask)
+            avg_prob_mult_error = torch.exp(abs_diff).mean()
+            assert avg_prob_mult_error <= 1.05, (
+                f"matched-impl inference_optimized: generation logprobs "
+                f"diverge from policy logprobs (avg prob mult error "
+                f"{avg_prob_mult_error:.4f})"
+            )
+            # Wake the engine again for the post-shutdown generation check.
+            mg.prepare_for_generation()
+
         # ownership guard: shutdown is a no-op, so the wrapped policy keeps generating
         assert mg.shutdown() is True
         after_shutdown = mg.generate(test_input_data, greedy=True)
diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py
index a79d01961aa..b8eb001a3c9 100644
--- a/tests/unit/models/policy/test_megatron_worker.py
+++ b/tests/unit/models/policy/test_megatron_worker.py
@@ -681,6 +681,8 @@ def test_prepare_for_generation_disables_param_gather_hook_before_wake(
         lambda *, param_sync=False: events.append(("disable_hook", param_sync))
     )
     worker._inference_engine_initialized = True
+    # Asleep, so the idempotent-wake guard falls through to the full wake path.
+    worker._inference_engine_asleep = True
     worker._wake = lambda: events.append("wake_engine")
 
     monkeypatch.setattr(megatron_worker, "log_gpu_memory", lambda *_: None)
diff --git a/tests/unit/weight_sync/test_weight_synchronizer.py b/tests/unit/weight_sync/test_weight_synchronizer.py
index 3335db7a3d4..ea1fa5fef03 100644
--- a/tests/unit/weight_sync/test_weight_synchronizer.py
+++ b/tests/unit/weight_sync/test_weight_synchronizer.py
@@ -555,7 +555,9 @@ def test_colocated_sync_is_offload_and_wake(self):
         assert sync.is_stale
         assert sync.sync_weights() == {}
         policy.offload_before_refit.assert_called_once()
-        gen.prepare_for_generation.assert_called_once_with()
+        # The refit-protocol tag makes the wake bypass the worker's
+        # engine-awake early-return (the reshard copy rides this wake).
+        gen.prepare_for_generation.assert_called_once_with(tags=["colocated_refit"])
         gen.suspend_for_refit.assert_not_called()
         policy.swap_weights_via_reshard.assert_not_called()
         assert not sync.is_stale