From d6d5d9310dc4ee1b3875c3b6f18568cdd468a204 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Mon, 22 Jun 2026 11:36:34 -0500 Subject: [PATCH 01/13] feat: async colocated GRPO with Megatron inference Signed-off-by: Teodor-Dumitru Ene --- ...-30BA3B-2n8g-megatron_async_colocated.yaml | 51 +++++++++++++++ nemo_rl/algorithms/grpo.py | 62 ++++++++++++++++--- ...v3-30BA3B-2n8g-megatron_async_colocated.sh | 38 ++++++++++++ tests/test_suites/nightly.txt | 1 + 4 files changed, 142 insertions(+), 10 deletions(-) create mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.yaml create mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.sh diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.yaml new file mode 100644 index 00000000000..b12830ed31b --- /dev/null +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.yaml @@ -0,0 +1,51 @@ +defaults: ../../grpo_math_1B.yaml +# Async colocated GRPO with Megatron (Nemotron-3-Nano-30B-A3B): training and +# generation share GPUs/workers; the engine stays alive across steps and reads +# weights from the shared model (no refit). Off-policy data via the replay buffer. +grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 8 + max_num_steps: 500 + 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-2n8g-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 + mcore_generation_config: + async_engine: true +logger: + log_dir: logs/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated + wandb_enabled: true + tensorboard_enabled: true + wandb: + project: nemo-rl + name: grpo-nanov3-30BA3B-2n8g-megatron_async_colocated +cluster: + gpus_per_node: 8 + num_nodes: 2 diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index d8d5f26763d..b6c9923a0fd 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -4244,13 +4244,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 unsupported for the desired 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 +4384,21 @@ def async_grpo_train( ) print("⏳ Preparing policy generation for training...", flush=True) - if NEED_REFIT and POLICY_GENERATION_STALE: + if colocated_inference: + # Colocated mode currently does not support weights refit. + print("🔄 Initializing colocated Megatron inference engine...") + try: + policy.offload_before_refit() + policy_generation.prepare_for_generation() + POLICY_GENERATION_STALE = False + print("✅ Colocated Megatron engine initialized") + except Exception as e: + print(f"❌ Colocated Megatron engine initialization failed: {e}") + import traceback + + traceback.print_exc() + return + elif NEED_REFIT and POLICY_GENERATION_STALE: print("🔄 Refitting policy generation with actual model weights...", flush=True) try: refit_policy_generation( @@ -4439,7 +4454,9 @@ def async_grpo_train( processor=processor, ) initial_val_metrics = val_metrics - policy_generation.finish_generation() + # Colocated engine stays alive across steps (preserves KV cache). + if not colocated_inference: + policy_generation.finish_generation() 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 +4842,17 @@ def async_grpo_train( ) train_data.to("cpu") + generation_logger_metrics = None + if colocated_inference: + print("⏸️ Pausing colocated engine + collector for training...") + with timer.time("exposed_generation"): + ray.get(trajectory_collector.prepare_for_refit.remote()) + if policy_generation is not None: + generation_logger_metrics = ( + policy_generation.get_logger_metrics() + ) + policy_generation.finish_generation() + # Training phase (same as sync version) skip_prev_logprobs, skip_reference_logprobs = ( _resolve_logprob_skip_flags(master_config) @@ -4957,8 +4985,17 @@ def async_grpo_train( ) print("🔄 Synchronizing policy weights to trajectory collector…") - generation_logger_metrics = None - if NEED_REFIT: + if colocated_inference: + # Colocated mode currently does not support weights refit. + print("🔄 Resuming colocated engine after training step...") + with timer.time("weight_sync"): + policy.offload_before_refit() + policy_generation.prepare_for_generation() + POLICY_GENERATION_STALE = False + weight_version += 1 + trajectory_collector.set_weight_version.remote(weight_version) + trajectory_collector.resume_after_refit.remote() + elif NEED_REFIT: timer.start("idle/refit_bubble") # Measure pending-generation wait as exposed_generation time @@ -5025,7 +5062,11 @@ 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: + if colocated_inference: + # Colocated engine was already resumed post-training; + # no refit path exists in colocated mode. + pass + elif NEED_REFIT and POLICY_GENERATION_STALE: refit_metrics = refit_policy_generation( policy, policy_generation, @@ -5044,7 +5085,8 @@ def async_grpo_train( logger=logger, processor=processor, ) - policy_generation.finish_generation() + if not colocated_inference: + policy_generation.finish_generation() logger.log_metrics( validation_timings, step + 1, prefix="timing/validation" ) diff --git a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.sh new file mode 100755 index 00000000000..43b19a166e4 --- /dev/null +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.sh @@ -0,0 +1,38 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=8 +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 + +# Only run metrics if the target step is reached +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/reward"]) > 0.0' +fi diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index cfde74908d1..384029a50cb 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -21,6 +21,7 @@ tests/test_suites/llm/grpo-moonlight-16b-automodel-1n8g-ep8.sh tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron.sh tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3.sh tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.sh +tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.sh # Functional 32b run tests/test_suites/llm/grpo-qwen2.5-32b-32n8g-fsdp2tp8-actckpt.v3.sh From b469e9b349b9d614bb339fa9e9a5509d778cc42c Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 29 Jul 2026 00:26:06 -0500 Subject: [PATCH 02/13] Drop unnecessary synchronize Signed-off-by: Teodor-Dumitru Ene Co-authored-by: Anil Thomas --- nemo_rl/models/generation/megatron/megatron_worker.py | 1 - 1 file changed, 1 deletion(-) diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 2bbd1014a84..6c7b1af453f 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -939,7 +939,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.""" From deeac4092b18bfe2aa25ac1193612a4a252d0c44 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 29 Jul 2026 00:29:56 -0500 Subject: [PATCH 03/13] Defer colocated wake past checkpoint saves Signed-off-by: Teodor-Dumitru Ene Co-authored-by: Anil Thomas --- nemo_rl/algorithms/grpo.py | 57 ++++++++++++++++++++++++-------------- 1 file changed, 36 insertions(+), 21 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index b6c9923a0fd..e71bcade1ad 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -4984,17 +4984,32 @@ 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) + ) + should_save_by_timeout = timeout.check_save() + will_save_checkpoint = master_config.checkpointing["enabled"] and ( + should_save_by_step or should_save_by_timeout + ) + print("🔄 Synchronizing policy weights to trajectory collector…") if colocated_inference: # Colocated mode currently does not support weights refit. - print("🔄 Resuming colocated engine after training step...") with timer.time("weight_sync"): policy.offload_before_refit() - policy_generation.prepare_for_generation() POLICY_GENERATION_STALE = False weight_version += 1 trajectory_collector.set_weight_version.remote(weight_version) - trajectory_collector.resume_after_refit.remote() + if will_save_checkpoint: + # Don't wake up engine if we're about to send it to sleep for checkpointing. + print("⏸️ Keeping colocated engine asleep for checkpointing...") + else: + print("🔄 Resuming colocated engine after training step...") + policy_generation.prepare_for_generation() + trajectory_collector.resume_after_refit.remote() elif NEED_REFIT: timer.start("idle/refit_bubble") @@ -5037,7 +5052,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 @@ -5063,9 +5077,10 @@ def async_grpo_train( if should_run_validation: with timer.time("idle/validation"): if colocated_inference: - # Colocated engine was already resumed post-training; - # no refit path exists in colocated mode. - pass + # No refit path exists in colocated mode. + # On save-bound steps, engine stayed asleep after training. + if will_save_checkpoint: + policy_generation.prepare_for_generation() elif NEED_REFIT and POLICY_GENERATION_STALE: refit_metrics = refit_policy_generation( policy, @@ -5085,7 +5100,7 @@ def async_grpo_train( logger=logger, processor=processor, ) - if not colocated_inference: + if not colocated_inference or will_save_checkpoint: policy_generation.finish_generation() logger.log_metrics( validation_timings, step + 1, prefix="timing/validation" @@ -5199,19 +5214,12 @@ 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 + # Early stop saves the final state like a last step; it is only + # known after validation, so it cannot fold into the earlier + # will_save_checkpoint computation. + if will_save_checkpoint or ( + master_config.checkpointing["enabled"] + and early_stop_message is not None ): grpo_save_state.current_step = step + 1 grpo_save_state.total_valid_tokens = total_valid_tokens @@ -5308,6 +5316,13 @@ def async_grpo_train( checkpointer, last_checkpoint_step=step + 1 ) + # On save-bound steps, engine stayed asleep after training. + # (On an early-stop save it is already awake and the loop + # exits right below, so no wake is needed.) + if colocated_inference and will_save_checkpoint: + policy_generation.prepare_for_generation() + 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 From e33e49c76528cfd4c86e097b173e245f1762a159 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 29 Jul 2026 01:46:39 -0500 Subject: [PATCH 04/13] Allow inference_optimized during training Signed-off-by: Teodor-Dumitru Ene Co-authored-by: Anil Thomas --- nemo_rl/models/policy/__init__.py | 2 +- .../policy/workers/megatron_policy_worker.py | 21 +++++++++----- .../generation/test_megatron_generation.py | 29 ++++++++++++++++++- 3 files changed, 42 insertions(+), 10 deletions(-) diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 9506422ad32..5d025ba303b 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -397,7 +397,7 @@ 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. # 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..750adb9e28c 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -496,14 +496,19 @@ 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." - ) + if config["megatron_cfg"].get("transformer_impl") == "inference_optimized": + if ( + config["megatron_cfg"]["tensor_model_parallel_size"] > 1 + and not config["megatron_cfg"]["sequence_parallel"] + ): + config["megatron_cfg"]["sequence_parallel"] = True + print( + "transformer_impl=inference_optimized with TP>1: " + "enabling megatron_cfg.sequence_parallel." + ) + # TODO: Remove the following two lines after Megatron-Bridge#5164 lands. + from megatron.bridge.models.conversion.param_mapping import AutoMapping + AutoMapping.register_module_type("InferenceColumnParallelLinear", "column") runtime_config = validate_and_set_config( config, self.rank, diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index 351f6764462..cd7b70934c1 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -28,6 +28,7 @@ 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 +383,17 @@ 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( + "transformer_impl", ["transformer_engine", "inference_optimized"] +) +def test_megatron_generation_colocated( + cluster, test_input_data, tokenizer, transformer_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 + config["megatron_cfg"]["transformer_impl"] = transformer_impl # construction guard: exactly one of `cluster` / `policy` is required with pytest.raises(AssertionError): @@ -413,6 +420,26 @@ def test_megatron_generation_colocated(cluster, test_input_data, tokenizer): 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 transformer_impl == "inference_optimized": + # Dual-mode: the same shared model must run the trainable TE + # fallback (train step, finite loss) before fast-path generation. + 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"dual-mode 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) From 67d8e1f4ecc2cbc13e2a239e5ad3f8414c74bcdf Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 4 Aug 2026 13:34:20 -0500 Subject: [PATCH 05/13] Add functional test Signed-off-by: Teodor-Dumitru Ene --- .../L1_Functional_Tests_Megatron_4.sh | 1 + ...egatron_generation_colocated_async_grpo.sh | 72 +++++++++++++++++++ 2 files changed, 73 insertions(+) create mode 100644 tests/functional/grpo_megatron_generation_colocated_async_grpo.sh diff --git a/tests/functional/L1_Functional_Tests_Megatron_4.sh b/tests/functional/L1_Functional_Tests_Megatron_4.sh index 2462db4d6e2..889a4cb1dd1 100644 --- a/tests/functional/L1_Functional_Tests_Megatron_4.sh +++ b/tests/functional/L1_Functional_Tests_Megatron_4.sh @@ -53,6 +53,7 @@ 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_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..8847ece11c5 --- /dev/null +++ b/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh @@ -0,0 +1,72 @@ +#!/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=2 \ + 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 save-bound step must defer the engine wake past the checkpoint 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 From b1e36e5e815c6ac9d24bf6cd815d04d3dc77d7b2 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 18 Aug 2026 05:22:44 -0500 Subject: [PATCH 06/13] Handle nightly tests Signed-off-by: Teodor-Dumitru Ene --- ...30BA3B-4n4g-megatron_async_colocated.yaml} | 17 ++++---- ...3-30BA3B-4n4g-megatron_sync_colocated.yaml | 41 +++++++++++++++++++ ...3-30BA3B-4n4g-megatron_async_colocated.sh} | 13 +++--- ...ov3-30BA3B-4n4g-megatron_sync_colocated.sh | 37 +++++++++++++++++ tests/test_suites/nightly.txt | 1 - tests/test_suites/nightly_gb200.txt | 2 + 6 files changed, 95 insertions(+), 16 deletions(-) rename examples/configs/recipes/llm/{grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.yaml => grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml} (72%) create mode 100644 examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.yaml rename tests/test_suites/llm/{grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.sh => grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh} (72%) create mode 100755 tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_sync_colocated.sh diff --git a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.yaml b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml similarity index 72% rename from examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.yaml rename to examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml index b12830ed31b..4ea17c36073 100644 --- a/examples/configs/recipes/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.yaml +++ b/examples/configs/recipes/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.yaml @@ -1,11 +1,11 @@ defaults: ../../grpo_math_1B.yaml # Async colocated GRPO with Megatron (Nemotron-3-Nano-30B-A3B): training and -# generation share GPUs/workers; the engine stays alive across steps and reads -# weights from the shared model (no refit). Off-policy data via the replay buffer. +# 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 - max_num_steps: 500 async_grpo: enabled: true max_trajectory_age_steps: 4 # weight versions a rollout may span @@ -15,7 +15,7 @@ loss_fn: use_importance_sampling_correction: true # required for off-policy replay data checkpointing: enabled: false - checkpoint_dir: results/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated + 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 @@ -40,12 +40,13 @@ policy: mcore_generation_config: async_engine: true logger: - log_dir: logs/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated + log_dir: logs/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated wandb_enabled: true tensorboard_enabled: true wandb: project: nemo-rl - name: grpo-nanov3-30BA3B-2n8g-megatron_async_colocated + name: grpo-nanov3-30BA3B-4n4g-megatron_async_colocated cluster: - gpus_per_node: 8 - num_nodes: 2 + 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/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.sh b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh similarity index 72% rename from tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.sh rename to tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh index 43b19a166e4..2bbe91651a4 100755 --- a/tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.sh +++ b/tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated.sh @@ -3,8 +3,9 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source $SCRIPT_DIR/common.env # ===== BEGIN CONFIG ===== -NUM_NODES=2 -GPUS_PER_NODE=8 +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 @@ -31,8 +32,6 @@ uv run examples/run_grpo.py \ # Convert tensorboard logs to json uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS -# Only run metrics if the target step is reached -if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then - uv run tests/check_metrics.py $JSON_METRICS \ - 'max(data["train/reward"]) > 0.0' -fi +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.txt b/tests/test_suites/nightly.txt index 384029a50cb..cfde74908d1 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -21,7 +21,6 @@ tests/test_suites/llm/grpo-moonlight-16b-automodel-1n8g-ep8.sh tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron.sh tests/test_suites/llm/grpo-qwen3-1.7b-1n8g-megatron-eagle3.sh tests/test_suites/llm/grpo-llama3.2-1b-instruct-1n8g-megatron_generation.sh -tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron_async_colocated.sh # Functional 32b run tests/test_suites/llm/grpo-qwen2.5-32b-32n8g-fsdp2tp8-actckpt.v3.sh 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 From 48d05e52c58964502e9a4a029cffde41a770be9d Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 18 Aug 2026 06:05:24 -0500 Subject: [PATCH 07/13] Remove the NEED_REFIT carve-out Signed-off-by: Teodor-Dumitru Ene --- nemo_rl/algorithms/grpo.py | 71 +++++++++++++-------------------- nemo_rl/algorithms/grpo_sync.py | 15 ++----- 2 files changed, 30 insertions(+), 56 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index e71bcade1ad..f01ed4b09ed 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 @@ -4384,21 +4376,7 @@ def async_grpo_train( ) print("⏳ Preparing policy generation for training...", flush=True) - if colocated_inference: - # Colocated mode currently does not support weights refit. - print("🔄 Initializing colocated Megatron inference engine...") - try: - policy.offload_before_refit() - policy_generation.prepare_for_generation() - POLICY_GENERATION_STALE = False - print("✅ Colocated Megatron engine initialized") - except Exception as e: - print(f"❌ Colocated Megatron engine initialization failed: {e}") - import traceback - - traceback.print_exc() - return - elif NEED_REFIT and POLICY_GENERATION_STALE: + if POLICY_GENERATION_STALE: print("🔄 Refitting policy generation with actual model weights...", flush=True) try: refit_policy_generation( @@ -4996,21 +4974,24 @@ def async_grpo_train( ) print("🔄 Synchronizing policy weights to trajectory collector…") - if colocated_inference: - # Colocated mode currently does not support weights refit. - with timer.time("weight_sync"): + if colocated_inference and will_save_checkpoint: + # 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 trajectory_collector.set_weight_version.remote(weight_version) - if will_save_checkpoint: - # Don't wake up engine if we're about to send it to sleep for checkpointing. - print("⏸️ Keeping colocated engine asleep for checkpointing...") - else: - print("🔄 Resuming colocated engine after training step...") - policy_generation.prepare_for_generation() - trajectory_collector.resume_after_refit.remote() - elif NEED_REFIT: + else: timer.start("idle/refit_bubble") # Measure pending-generation wait as exposed_generation time @@ -5020,7 +5001,11 @@ 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 ( + policy_generation is not None + and generation_logger_metrics is None + ): generation_logger_metrics = ( policy_generation.get_logger_metrics() ) @@ -5076,12 +5061,7 @@ 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 colocated_inference: - # No refit path exists in colocated mode. - # On save-bound steps, engine stayed asleep after training. - if will_save_checkpoint: - policy_generation.prepare_for_generation() - elif NEED_REFIT and POLICY_GENERATION_STALE: + if POLICY_GENERATION_STALE: refit_metrics = refit_policy_generation( policy, policy_generation, @@ -5089,6 +5069,9 @@ def async_grpo_train( ) POLICY_GENERATION_STALE = False else: + # 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, diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 3d843e36b05..bc8c0820ccb 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 @@ -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: @@ -599,7 +590,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 +1002,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, From 7b3a56a88aa8bbe7650e5a03dfecccaefdf8d1cc Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 18 Aug 2026 05:23:05 -0500 Subject: [PATCH 08/13] Handle rebase conflicts Signed-off-by: Teodor-Dumitru Ene --- docs/guides/async-grpo.md | 2 +- ...-30BA3B-4n4g-megatron_async_colocated.yaml | 2 - nemo_rl/algorithms/grpo.py | 41 ++++++--- nemo_rl/models/generation/interfaces.py | 9 ++ .../megatron/megatron_generation.py | 30 ++++++- .../generation/megatron/megatron_worker.py | 16 +++- .../models/generation/vllm/vllm_generation.py | 8 ++ nemo_rl/models/policy/__init__.py | 4 +- .../policy/workers/megatron_policy_worker.py | 13 --- .../generation/test_megatron_generation.py | 83 ++++++++++++++++--- 10 files changed, 163 insertions(+), 45 deletions(-) diff --git a/docs/guides/async-grpo.md b/docs/guides/async-grpo.md index 64838fab625..1877561530c 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: 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 index 4ea17c36073..69abf1b560d 100644 --- 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 @@ -37,8 +37,6 @@ policy: enabled: false generation: backend: megatron - mcore_generation_config: - async_engine: true logger: log_dir: logs/grpo-nanov3-30BA3B-4n4g-megatron_async_colocated wandb_enabled: true diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index f01ed4b09ed..5b4ce3e3b28 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1140,6 +1140,22 @@ def _spinup_nemo_gym(base_urls, model_name): ) mcore_cfg["kv_cache_management_mode"] = "recompute" + # inference_optimized layers hard-require sequence parallelism with TP>1 + # (asserted at model build). Enable it here, driver-side, so the logged + # and checkpointed config reflects the effective value. + if "megatron_cfg" in policy_config and policy_config["megatron_cfg"]["enabled"]: + megatron_cfg = policy_config["megatron_cfg"] + if ( + megatron_cfg.get("transformer_impl") == "inference_optimized" + and megatron_cfg["tensor_model_parallel_size"] > 1 + and not megatron_cfg["sequence_parallel"] + ): + megatron_cfg["sequence_parallel"] = True + print( + "Auto-enabling `policy.megatron_cfg.sequence_parallel=True`: " + "transformer_impl=inference_optimized requires it with TP>1." + ) + # Define initialization functions that will be used in all paths init_reference_model = loss_config.reference_policy_kl_penalty > 0 @@ -4239,7 +4255,7 @@ def async_grpo_train( assert (not colocated_inference) or ( isinstance(policy_generation, MegatronGeneration) - ), "Colocated async GRPO is unsupported for the desired generation backend." + ), "Colocated async GRPO is only supported for the Megatron generation backend." # Initialize advantage estimator adv_estimator = _create_advantage_estimator(master_config) @@ -4432,9 +4448,9 @@ def async_grpo_train( processor=processor, ) initial_val_metrics = val_metrics - # Colocated engine stays alive across steps (preserves KV cache). - if not colocated_inference: - 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(for_training=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: @@ -4821,15 +4837,12 @@ def async_grpo_train( train_data.to("cpu") generation_logger_metrics = None - if colocated_inference: + 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()) - if policy_generation is not None: - generation_logger_metrics = ( - policy_generation.get_logger_metrics() - ) - policy_generation.finish_generation() + generation_logger_metrics = policy_generation.get_logger_metrics() + policy_generation.finish_generation(for_training=True) # Training phase (same as sync version) skip_prev_logprobs, skip_reference_logprobs = ( @@ -5083,8 +5096,12 @@ def async_grpo_train( logger=logger, processor=processor, ) - if not colocated_inference or will_save_checkpoint: - policy_generation.finish_generation() + # 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( + for_training=will_save_checkpoint + ) logger.log_metrics( validation_timings, step + 1, prefix="timing/validation" ) diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index f8c7d25ca02..c5cb0bb6e81 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -438,6 +438,15 @@ 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 clear_logger_metrics(self) -> None: """Clear logger metrics for performance reporting. diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index d21c40bb52a..72e7d06b307 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -260,13 +260,39 @@ 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.""" + """Clean up after generation. + + Accepts `for_training` (default True): when False, a colocated engine + keeps serving instead of standing down (see the worker docstring). + """ futures = self._policy.worker_group.run_all_workers_single_data( - "finish_generation" + "finish_generation", **kwargs ) 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 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 6c7b1af453f..d672de77744 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -368,8 +368,20 @@ 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, *, for_training: bool = True) -> None: + """Wind down a generation cycle. + + Args: + for_training: the caller is about to need 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. Non-colocated engines wind + down either way. + """ + if self.is_generation_colocated and not for_training: + return print(f"[Rank {self.rank}] finishing generation", flush=True) log_gpu_memory("finish_generation START") diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index a8cb78bd13c..f0a7a94ea7a 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1001,6 +1001,14 @@ def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool: print(f"Error during policy preparation: {e}") return False + def blocks_training(self) -> bool: + """Whether the engine must stand down before a training step. + + Colocated vLLM shares the training GPUs (async colocated vLLM is + not supported yet, but the capability answers correctly for it). + """ + return bool(self.cfg["colocated"]["enabled"]) + def finish_generation(self, *args: Any, **kwargs: Any) -> bool: """Sleep workers and reset prefix cache.""" try: diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 5d025ba303b..370be5d78b9 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. + # Transformer implementation backing the model. 'inference_optimized' + # trains through the TE parent path and requires sequence_parallel with + # TP>1 (auto-enabled driver-side for GRPO). # 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 750adb9e28c..fc913dbb8ce 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -496,19 +496,6 @@ def __init__( self.tokenizer.pad_token = self.tokenizer.eos_token # Step 3: Setup model configuration - if config["megatron_cfg"].get("transformer_impl") == "inference_optimized": - if ( - config["megatron_cfg"]["tensor_model_parallel_size"] > 1 - and not config["megatron_cfg"]["sequence_parallel"] - ): - config["megatron_cfg"]["sequence_parallel"] = True - print( - "transformer_impl=inference_optimized with TP>1: " - "enabling megatron_cfg.sequence_parallel." - ) - # TODO: Remove the following two lines after Megatron-Bridge#5164 lands. - from megatron.bridge.models.conversion.param_mapping import AutoMapping - AutoMapping.register_module_type("InferenceColumnParallelLinear", "column") runtime_config = validate_and_set_config( config, self.rank, diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index cd7b70934c1..c0613295eea 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -23,6 +23,9 @@ 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 ( @@ -384,16 +387,34 @@ async def test_megatron_policy_generation_async(cluster, test_input_data, tokeni @pytest.mark.mcore @pytest.mark.timeout(900) @pytest.mark.parametrize( - "transformer_impl", ["transformer_engine", "inference_optimized"] + "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, transformer_impl + 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 - config["megatron_cfg"]["transformer_impl"] = transformer_impl + # 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 # construction guard: exactly one of `cluster` / `policy` is required with pytest.raises(AssertionError): @@ -415,14 +436,22 @@ def test_megatron_generation_colocated( 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 transformer_impl == "inference_optimized": - # Dual-mode: the same shared model must run the trainable TE - # fallback (train step, finite loss) before fast-path generation. + 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( { @@ -436,7 +465,7 @@ def test_megatron_generation_colocated( policy.prepare_for_training() loss = policy.train(train_data, SimpleLossFn())["loss"] assert not torch.isnan(loss).any() and not torch.isinf(loss).any(), ( - f"dual-mode train step produced bad loss: {loss}" + f"pre-generation train step produced bad loss: {loss}" ) policy.finish_training() @@ -445,6 +474,36 @@ def test_megatron_generation_colocated( 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"], + } + ) + 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})" + ) + # 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) From 3fd09e24a9c983334c616bc5219ac6ad7bf8204c Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 19 Aug 2026 22:06:31 -0500 Subject: [PATCH 09/13] Address PR team review Signed-off-by: Teodor-Dumitru Ene --- docs/guides/async-grpo.md | 2 +- nemo_rl/algorithms/grpo.py | 64 ++++++++----- nemo_rl/algorithms/grpo_sync.py | 5 +- nemo_rl/models/generation/interfaces.py | 24 +++++ nemo_rl/models/generation/megatron/config.py | 14 ++- .../megatron/megatron_generation.py | 11 ++- .../generation/megatron/megatron_worker.py | 19 +++- nemo_rl/models/megatron/setup.py | 12 +++ .../L1_Functional_Tests_Megatron_4.sh | 1 + ...egatron_generation_colocated_async_grpo.sh | 2 +- ...generation_colocated_reshard_async_grpo.sh | 91 +++++++++++++++++++ .../generation/test_megatron_generation.py | 4 + 12 files changed, 211 insertions(+), 38 deletions(-) create mode 100644 tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh diff --git a/docs/guides/async-grpo.md b/docs/guides/async-grpo.md index 1877561530c..94b40c215f7 100644 --- a/docs/guides/async-grpo.md +++ b/docs/guides/async-grpo.md @@ -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/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 5b4ce3e3b28..68d35c32e0a 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -4981,13 +4981,23 @@ def async_grpo_train( 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…") - if colocated_inference and will_save_checkpoint: + 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 @@ -5003,7 +5013,9 @@ def async_grpo_train( policy.offload_before_refit() POLICY_GENERATION_STALE = False weight_version += 1 - trajectory_collector.set_weight_version.remote(weight_version) + ray.get( + trajectory_collector.set_weight_version.remote(weight_version) + ) else: timer.start("idle/refit_bubble") @@ -5015,10 +5027,7 @@ def async_grpo_train( # Collect generation logger metrics for performance reporting # inflight batch sizes and num pending samples are collected from each worker # (colocated collects them before the engine sleeps for training). - if ( - policy_generation is not None - and generation_logger_metrics is None - ): + if generation_logger_metrics is None: generation_logger_metrics = ( policy_generation.get_logger_metrics() ) @@ -5096,11 +5105,21 @@ def async_grpo_train( logger=logger, processor=processor, ) + # 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( - for_training=will_save_checkpoint + for_training=saving_this_step ) logger.log_metrics( validation_timings, step + 1, prefix="timing/validation" @@ -5116,11 +5135,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) @@ -5214,13 +5228,7 @@ def async_grpo_train( consumed_samples += master_config.grpo.num_prompts_per_step timeout.mark_iteration() - # Early stop saves the final state like a last step; it is only - # known after validation, so it cannot fold into the earlier - # will_save_checkpoint computation. - if will_save_checkpoint or ( - master_config.checkpointing["enabled"] - and early_stop_message is not None - ): + 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: @@ -5316,12 +5324,20 @@ def async_grpo_train( checkpointer, last_checkpoint_step=step + 1 ) - # On save-bound steps, engine stayed asleep after training. - # (On an early-stop save it is already awake and the loop - # exits right below, so no wake is needed.) - if colocated_inference and will_save_checkpoint: + # 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() - trajectory_collector.resume_after_refit.remote() + ray.get(trajectory_collector.resume_after_refit.remote()) # Logging # Log training data (match sync GRPO logging payload for parity). diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index bc8c0820ccb..6ddd72e1e25 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -379,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, @@ -571,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"): diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index c5cb0bb6e81..7aeb4887ae7 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -372,10 +372,23 @@ 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 `for_training` (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. + Colocated backends must honor the flag; engines on dedicated GPUs may ignore it. + """ pass @abstractmethod @@ -447,6 +460,17 @@ def blocks_training(self) -> bool: """ 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..b1b2169027a 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -75,11 +75,19 @@ 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; + # this is a derived config, so complete it rather than assert. + if ( + merged.get("transformer_impl") == "inference_optimized" + and merged["tensor_model_parallel_size"] > 1 + ): + merged["sequence_parallel"] = True + return merged def dedicated_inference_megatron_cfg( @@ -92,8 +100,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 72e7d06b307..b61b8c7b2db 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -259,14 +259,13 @@ def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool: ] return True - def finish_generation(self, *args: Any, **kwargs: Any) -> bool: + def finish_generation(self, *, for_training: bool = True) -> bool: """Clean up after generation. - Accepts `for_training` (default True): when False, a colocated engine - keeps serving instead of standing down (see the worker docstring). + When `for_training` is False, a colocated engine keeps serving instead of standing down. """ futures = self._policy.worker_group.run_all_workers_single_data( - "finish_generation", **kwargs + "finish_generation", for_training=for_training ) ray.get(futures) return True @@ -279,6 +278,10 @@ def blocks_training(self) -> bool: """ 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. diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index d672de77744..aac55135e1d 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -377,8 +377,9 @@ def finish_generation(self, *, for_training: bool = True) -> None: 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. Non-colocated engines wind - down either way. + 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 for_training: return @@ -414,10 +415,24 @@ def finish_generation(self, *, for_training: bool = True) -> 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"] diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index accee9ec08d..d693a2d3f16 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -329,6 +329,18 @@ def validate_and_set_config( weights_path, optimizer_path, ): + # inference_optimized layers hard-require SP with TP>1; fail here with the config key. + # GRPO's setup() auto-enables SP driver-side, so it never trips this. + 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/tests/functional/L1_Functional_Tests_Megatron_4.sh b/tests/functional/L1_Functional_Tests_Megatron_4.sh index 889a4cb1dd1..484300a0ccf 100644 --- a/tests/functional/L1_Functional_Tests_Megatron_4.sh +++ b/tests/functional/L1_Functional_Tests_Megatron_4.sh @@ -54,6 +54,7 @@ if megatron_generation_supported; then 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 index 8847ece11c5..9227bca5f1b 100644 --- a/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh +++ b/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh @@ -38,7 +38,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE grpo.async_grpo.in_flight_weight_updates=true \ loss_fn.use_importance_sampling_correction=true \ grpo.max_num_steps=3 \ - grpo.val_period=2 \ + grpo.val_period=1 \ grpo.max_val_samples=8 \ grpo.val_batch_size=8 \ cluster.gpus_per_node=2 \ 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..df989637f6d --- /dev/null +++ b/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh @@ -0,0 +1,91 @@ +#!/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/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.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. +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/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index c0613295eea..ef7835f94bf 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -415,6 +415,10 @@ def test_megatron_generation_colocated( # 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 below does not reflect any real use-case. + # We must disable CUDA graphs beforehand just for this test. + config["generation"]["mcore_generation_config"]["cuda_graph_impl"] = "none" # construction guard: exactly one of `cluster` / `policy` is required with pytest.raises(AssertionError): From b3b51bbc2c0782328bd1f02b2ebb8026539e926b Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 20 Aug 2026 05:32:48 -0500 Subject: [PATCH 10/13] Fix CI Signed-off-by: Teodor-Dumitru Ene --- tests/unit/algorithms/test_grpo.py | 5 +++++ tests/unit/models/generation/test_megatron_generation.py | 3 +++ tests/unit/models/policy/test_megatron_worker.py | 2 ++ 3 files changed, 10 insertions(+) diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 51c9d4290ed..5845f33c867 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -2958,6 +2958,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( diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index ef7835f94bf..9ce6957e6e2 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -419,6 +419,9 @@ def test_megatron_generation_colocated( # The parity block below does not reflect any real use-case. # We must disable CUDA graphs beforehand just for this test. 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): 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) From 8772e0419d858f1035a426316f8058d32f48bb87 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 20 Aug 2026 08:08:45 -0500 Subject: [PATCH 11/13] Address reviewer comment Signed-off-by: Teodor-Dumitru Ene --- nemo_rl/algorithms/grpo.py | 37 ++----- nemo_rl/models/generation/interfaces.py | 5 +- nemo_rl/models/generation/megatron/config.py | 13 ++- .../megatron/megatron_generation.py | 6 +- .../generation/megatron/megatron_worker.py | 16 +-- .../models/generation/vllm/vllm_generation.py | 8 -- nemo_rl/models/megatron/setup.py | 3 +- nemo_rl/models/policy/__init__.py | 2 +- .../megatron_weight_synchronizer.py | 12 ++- ...egatron_generation_colocated_async_grpo.sh | 5 +- ...generation_colocated_reshard_async_grpo.sh | 1 + tests/unit/algorithms/test_grpo.py | 102 ++++++++++++++++++ .../weight_sync/test_weight_synchronizer.py | 4 +- 13 files changed, 149 insertions(+), 65 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 68d35c32e0a..61d7f3468b3 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1140,22 +1140,6 @@ def _spinup_nemo_gym(base_urls, model_name): ) mcore_cfg["kv_cache_management_mode"] = "recompute" - # inference_optimized layers hard-require sequence parallelism with TP>1 - # (asserted at model build). Enable it here, driver-side, so the logged - # and checkpointed config reflects the effective value. - if "megatron_cfg" in policy_config and policy_config["megatron_cfg"]["enabled"]: - megatron_cfg = policy_config["megatron_cfg"] - if ( - megatron_cfg.get("transformer_impl") == "inference_optimized" - and megatron_cfg["tensor_model_parallel_size"] > 1 - and not megatron_cfg["sequence_parallel"] - ): - megatron_cfg["sequence_parallel"] = True - print( - "Auto-enabling `policy.megatron_cfg.sequence_parallel=True`: " - "transformer_impl=inference_optimized requires it with TP>1." - ) - # Define initialization functions that will be used in all paths init_reference_model = loss_config.reference_policy_kl_penalty > 0 @@ -4450,7 +4434,7 @@ def async_grpo_train( initial_val_metrics = val_metrics # A colocated engine keeps serving between phases (preserves its # KV/prefix cache); the backend makes that call, not the loop. - policy_generation.finish_generation(for_training=False) + 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: @@ -4842,7 +4826,7 @@ def async_grpo_train( 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(for_training=True) + policy_generation.finish_generation(release_gpu=True) # Training phase (same as sync version) skip_prev_logprobs, skip_reference_logprobs = ( @@ -5083,18 +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 POLICY_GENERATION_STALE: - refit_metrics = refit_policy_generation( - policy, - policy_generation, - colocated_inference, - ) - POLICY_GENERATION_STALE = False - else: - # 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() + # 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, @@ -5119,7 +5094,7 @@ def async_grpo_train( # so the engine must stand down; otherwise a colocated # engine keeps serving (backend's call). policy_generation.finish_generation( - for_training=saving_this_step + release_gpu=saving_this_step ) logger.log_metrics( validation_timings, step + 1, prefix="timing/validation" diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 7aeb4887ae7..0867bc121ed 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -382,12 +382,13 @@ def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool: def finish_generation(self, *args: Any, **kwargs: Any) -> bool: """Wind down after a generation phase. - Callers may pass `for_training` (keyword-only, default True): + 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. - Colocated backends must honor the flag; engines on dedicated GPUs may ignore it. + Only the colocated Megatron backend honors the flag today; other backends + ignore it, as do engines on dedicated GPUs. """ pass diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index b1b2169027a..234357f0a37 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -80,13 +80,20 @@ def merged_inference_megatron_cfg(policy_config: PolicyConfig) -> dict[str, Any] **(generation_config.get("mcore_generation_config") or {}), "activation_checkpointing": False, } - # inference_optimized layers hard-require SP with TP>1; - # this is a derived config, so complete it rather than assert. + # 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"] ): - merged["sequence_parallel"] = True + 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 diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index b61b8c7b2db..4e71222a0b3 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -259,13 +259,13 @@ def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool: ] return True - def finish_generation(self, *, for_training: bool = True) -> bool: + def finish_generation(self, *, release_gpu: bool = True) -> bool: """Clean up after generation. - When `for_training` is False, a colocated engine keeps serving instead of standing down. + 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", for_training=for_training + "finish_generation", release_gpu=release_gpu ) ray.get(futures) return True diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index aac55135e1d..ea08a8483bb 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -368,20 +368,20 @@ 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, *, for_training: bool = True) -> None: + def finish_generation(self, *, release_gpu: bool = True) -> None: """Wind down a generation cycle. Args: - for_training: the caller is about to need 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. + 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 for_training: + 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") diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index f0a7a94ea7a..a8cb78bd13c 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1001,14 +1001,6 @@ def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool: print(f"Error during policy preparation: {e}") return False - def blocks_training(self) -> bool: - """Whether the engine must stand down before a training step. - - Colocated vLLM shares the training GPUs (async colocated vLLM is - not supported yet, but the capability answers correctly for it). - """ - return bool(self.cfg["colocated"]["enabled"]) - def finish_generation(self, *args: Any, **kwargs: Any) -> bool: """Sleep workers and reset prefix cache.""" try: diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index d693a2d3f16..d7e3fdb0175 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -330,7 +330,8 @@ def validate_and_set_config( optimizer_path, ): # inference_optimized layers hard-require SP with TP>1; fail here with the config key. - # GRPO's setup() auto-enables SP driver-side, so it never trips this. + # 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 diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 370be5d78b9..b4048602bb3 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -399,7 +399,7 @@ class MegatronConfig(TypedDict): moe_router_group_topk: NotRequired[int | None] # Transformer implementation backing the model. 'inference_optimized' # trains through the TE parent path and requires sequence_parallel with - # TP>1 (auto-enabled driver-side for GRPO). + # TP>1 (enforced at setup). # Options are 'transformer_engine' and 'inference_optimized'. transformer_impl: NotRequired[str] # CUDA-graph implementation. 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/grpo_megatron_generation_colocated_async_grpo.sh b/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh index 9227bca5f1b..f58d0138f5b 100644 --- a/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh +++ b/tests/functional/grpo_megatron_generation_colocated_async_grpo.sh @@ -38,7 +38,7 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE 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.val_period=3 \ grpo.max_val_samples=8 \ grpo.val_batch_size=8 \ cluster.gpus_per_node=2 \ @@ -58,9 +58,10 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS 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"]' + '"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 diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh b/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh index df989637f6d..4d609b0d9dd 100644 --- a/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh +++ b/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh @@ -80,6 +80,7 @@ if ! grep -q "prepare_for_generation: engine already serving, skipping" $RUN_LOG 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 diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 5845f33c867..0106b43cc49 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -91,6 +91,8 @@ 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 return policy_generation @@ -3075,6 +3077,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/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 From ab0c13788b449740baf90267d405295f28915e0a Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 20 Aug 2026 12:23:13 -0500 Subject: [PATCH 12/13] Fix CI Signed-off-by: Teodor-Dumitru Ene --- tests/unit/algorithms/test_grpo.py | 3 +++ tests/unit/models/generation/test_megatron_generation.py | 8 ++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 0106b43cc49..e0bba7874d9 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -93,6 +93,9 @@ def _mock_policy_generation() -> MagicMock: 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 diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index 9ce6957e6e2..3158bbdc2be 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -416,8 +416,8 @@ def test_megatron_generation_colocated( config["megatron_cfg"]["transformer_impl"] = train_impl config["generation"]["mcore_generation_config"]["transformer_impl"] = gen_impl if train_impl == "inference_optimized": - # The parity block below does not reflect any real use-case. - # We must disable CUDA graphs beforehand just for this test. + # 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. @@ -493,6 +493,8 @@ def test_megatron_generation_colocated( "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) @@ -510,6 +512,8 @@ def test_megatron_generation_colocated( 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 From a53fc02313aec91aebd0db0fec1e691b5eb48d99 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Thu, 20 Aug 2026 18:19:23 -0500 Subject: [PATCH 13/13] Fix CI Signed-off-by: Teodor-Dumitru Ene --- .../grpo_megatron_generation_colocated_reshard_async_grpo.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh b/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh index 4d609b0d9dd..62c8e57f862 100644 --- a/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh +++ b/tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh @@ -27,7 +27,7 @@ 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 \ + 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 \