From 8e973d8475df78bb767e72b4d0de0ab3d0b28ef5 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Fri, 10 Apr 2026 02:28:40 -0500 Subject: [PATCH 1/3] Fix RL to once again work with --skip-train --- megatron/training/training.py | 1 + 1 file changed, 1 insertion(+) diff --git a/megatron/training/training.py b/megatron/training/training.py index 133e5f3a937..8e909f3ffe7 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3720,6 +3720,7 @@ def should_disable_forward_pre_hook(args): """Block forward pre-hook for certain configurations.""" return ( not args.use_megatron_fsdp + and not (args.skip_train and (not args.perform_rl_step or args.no_load_optim)) and (args.use_distributed_optimizer or args.use_layer_wise_distributed_optimizer) and args.overlap_param_gather ) From 2048831a5ce1e4633a36709dc5d0c96d7181f124 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Sat, 18 Apr 2026 09:38:04 -0500 Subject: [PATCH 2/3] Explain boolean logic --- megatron/training/training.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/megatron/training/training.py b/megatron/training/training.py index 85c29a15343..a5d8f405a00 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1614,10 +1614,11 @@ def setup_model_and_optimizer( timers = get_timers() one_logger = get_one_logger() - # Skip optimizer when not training. In RL inference-only mode (skip_train + perform_rl_step), - # --no-load-optim controls whether the optimizer is skipped (saving memory) or created - # (required for --rl-offload-optimizer-during-inference). - skip_optimizer = args.skip_train and (not args.perform_rl_step or args.no_load_optim) + # Typically, --skip-train is the only thing needed to disable the optimizer. + has_normal_optimizer = not args.skip_train + # Even with --skip-train, RL still creates an optimizer unless --no-load-optim is set. + has_rl_optimizer = args.perform_rl_step and not args.no_load_optim + skip_optimizer = not (has_normal_optimizer or has_rl_optimizer) wrap_with_ddp = not skip_optimizer model = get_model(model_provider_func, model_type, wrap_with_ddp=wrap_with_ddp) unwrapped_model = unwrap_model(model) @@ -3745,9 +3746,14 @@ def _get_iterator(dataloader_type, dataloader): def should_disable_forward_pre_hook(args): """Block forward pre-hook for certain configurations.""" + # Typically, --skip-train is the only thing needed to disable the optimizer. + has_normal_optimizer = not args.skip_train + # Even with --skip-train, RL still creates an optimizer unless --no-load-optim is set. + has_rl_optimizer = args.perform_rl_step and not args.no_load_optim + has_optimizer = has_normal_optimizer or has_rl_optimizer return ( not args.use_megatron_fsdp - and not (args.skip_train and (not args.perform_rl_step or args.no_load_optim)) + and has_optimizer and (args.use_distributed_optimizer or args.use_layer_wise_distributed_optimizer) and args.overlap_param_gather ) From 8d7da5912be47111d5ebadce22826e6e0da8cdbd Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Sat, 18 Apr 2026 15:35:20 -0500 Subject: [PATCH 3/3] Address reviewer comment --- megatron/training/training.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/megatron/training/training.py b/megatron/training/training.py index a5d8f405a00..c1c2b5ac68e 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3750,6 +3750,8 @@ def should_disable_forward_pre_hook(args): has_normal_optimizer = not args.skip_train # Even with --skip-train, RL still creates an optimizer unless --no-load-optim is set. has_rl_optimizer = args.perform_rl_step and not args.no_load_optim + # The forward pre-hooks are part of the distributed optimizer's overlapped param-gather; + # so in order to disable them, we must check that the optimizer actually exists. has_optimizer = has_normal_optimizer or has_rl_optimizer return ( not args.use_megatron_fsdp