From 8ae06df300686faa98d7a48bf14666a6bdc311d6 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Mon, 1 Dec 2025 12:37:23 -0800 Subject: [PATCH 01/18] Turn on training cudagraphs for RL --- .../abstract_model_inference_wrapper.py | 2 +- .../text_generation_controller.py | 2 +- megatron/core/transformer/cuda_graphs.py | 2 +- megatron/rl/rl_utils.py | 8 -------- 4 files changed, 3 insertions(+), 11 deletions(-) diff --git a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py index 95d476a9f83..4c3d02b1fc5 100644 --- a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py +++ b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py @@ -348,7 +348,7 @@ def forward_pass_with_pipeline_parallel_large_input_batch( # NOTE: Only returns the logits on the last pipeline stage return logits - @torch.inference_mode() + @torch.no_grad() def run_one_forward_step( self, inference_input: Dict[str, Any], recv_buffer_seq_len: Optional[int] = None ) -> torch.Tensor: diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index 6ee050ad62e..6f2396db3be 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -550,7 +550,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor) active_request_count = context.total_request_count - context.paused_request_count - with torch.inference_mode(): + with torch.no_grad(): logits = self.inference_wrapped_model.run_one_forward_step( {"tokens": input_ids, "position_ids": position_ids, "attention_mask": None} ) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 7b81eb723ed..1595be3adc8 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1319,7 +1319,7 @@ def __call__(self, megatron_module, args, kwargs): # Now replay the graph out = runner.replay_graph_capture(self.is_first_microbatch, args, kwargs) - elif self.training: + elif self.training and torch.is_grad_enabled(): # Training mode runner = self.get_cudagraph_runner(megatron_module, args, kwargs) # check if a layer is frozen during training. diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index e40b0b3220d..b16795fa9c0 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -2275,10 +2275,6 @@ def megatron_rl_inference_mode( with nvtx_range("offload-optimizer-before-inference"): optimizer.offload_to_cpu() - # TODO: Remove this if statement once a change to `toggle_cuda_graphs` makes it safe to. - if cuda_graph_impl != "none": - toggle_cuda_graphs(lang_module, cuda_graph_impl, reset_cuda_graphs=reset_cuda_graphs) - inference_interface = get_inference_interface(args, loop, model) with nvtx_range("onload-kv-cache-before-inference"): @@ -2324,10 +2320,6 @@ def megatron_rl_inference_mode( elif remove_kv_cache_during_training: inference_interface._inference_engine.context.memory_buffer = None - # TODO: Remove this if statement once a change to `toggle_cuda_graphs` makes it safe to. - if cuda_graph_impl != "none": - toggle_cuda_graphs(lang_module, 'none', reset_cuda_graphs=reset_cuda_graphs) - if offload_optimizer_during_inference: with nvtx_range("onload-optimizer-after-inference"): optimizer.restore_from_cpu() From 97266e32b028cef53b1d154ee584c9f5f323e33a Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 2 Dec 2025 08:24:25 -0800 Subject: [PATCH 02/18] Add data pointer comparison to tensors in ArgMetadata --- megatron/core/transformer/cuda_graphs.py | 1 + 1 file changed, 1 insertion(+) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 1595be3adc8..390b826e5da 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -99,6 +99,7 @@ def __init__(self, arg): self.shape = arg.shape self.dtype = arg.dtype self.device = arg.device + self.value = arg.data_ptr() else: self.value = arg From e59966f83c83fea5b4222e4a191fcd951a3ec457 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Fri, 5 Dec 2025 09:09:36 -0800 Subject: [PATCH 03/18] persist cudagraphs --- megatron/core/inference/engines/dynamic_engine.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index be3ba80d5c3..d5800b9e53e 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -551,12 +551,6 @@ def suspend(self): ): self.context.deallocate_all_tensors() - # Delete cuda graphs when not using unified memory at all (level 0). For - # levels 1 and 2, the context's tensors maintain static memory addresses, - # so the cuda graphs are re-used. - if self.unified_memory_level == 0: - delete_cuda_graphs() - # Maintain references to requests before reset. waiting_request_ids = list(self.waiting_request_ids) active_request_ids = set(self.requests.keys()) - set(waiting_request_ids) From 836bda6adf775ca0e28a28e50ed8839a2d0fed84 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Mon, 8 Dec 2025 08:25:39 -0800 Subject: [PATCH 04/18] =?UTF-8?q?Gate=20training=20cudagraph=20persistency?= =?UTF-8?q?=20via=20-=E2=80=93rl-persist-cuda-graphs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- megatron/core/inference/engines/dynamic_engine.py | 6 ++++++ megatron/training/arguments.py | 4 ++++ 2 files changed, 10 insertions(+) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index d5800b9e53e..180fa628a1f 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -551,6 +551,12 @@ def suspend(self): ): self.context.deallocate_all_tensors() + # Delete cuda graphs when not using unified memory at all (level 0) and + # `rl-persist-cuda-graphs` is not passed. For UVM levels 1 and 2, the context's tensors + # maintain static memory addresses, so the cuda graphs are re-used. + if self.unified_memory_level == 0 and not args.rl_persist_cuda_graphs: + delete_cuda_graphs() + # Maintain references to requests before reset. waiting_request_ids = list(self.waiting_request_ids) active_request_ids = set(self.requests.keys()) - set(waiting_request_ids) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 6308dc414b1..94f0a1a1f14 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1990,6 +1990,10 @@ def _add_rl_args(parser): help='Algorithm for distributing packed bins across ranks. ' 'fifo: first-in-first-out sequential distribution, ' 'round-robin: distribute bins cyclically across ranks for better load balancing') + group.add_argument('-–rl-persist-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, + default=False, + help='If set, do not call `delete_cuda_graphs` when the inference engine is suspended. ' + 'Use only when all training and inference cudagraphs and the KV cache fit on device.') return parser def _add_training_args(parser): From c8903d9bf35e2771a7f4c07f014929a07de25991 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Mon, 8 Dec 2025 08:45:33 -0800 Subject: [PATCH 05/18] Fix typo in args --- megatron/training/arguments.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 94f0a1a1f14..f1e9fdfa8f6 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1990,7 +1990,7 @@ def _add_rl_args(parser): help='Algorithm for distributing packed bins across ranks. ' 'fifo: first-in-first-out sequential distribution, ' 'round-robin: distribute bins cyclically across ranks for better load balancing') - group.add_argument('-–rl-persist-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, + group.add_argument('--rl-persist-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, default=False, help='If set, do not call `delete_cuda_graphs` when the inference engine is suspended. ' 'Use only when all training and inference cudagraphs and the KV cache fit on device.') From 4185aac5504de6089eb9b92dce0b5e9cfdd555c7 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Mon, 8 Dec 2025 09:44:22 -0800 Subject: [PATCH 06/18] persist cudagraphs --- megatron/core/inference/contexts/dynamic_context.py | 2 ++ megatron/core/inference/engines/dynamic_engine.py | 7 ++++--- megatron/rl/inference/megatron.py | 1 + 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 9b855effb7c..5d76818832f 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -270,6 +270,7 @@ def __init__( cuda_graph_mixed_prefill_count: Optional[int] = 16, metrics_writer: Optional['WandbModule'] = None, num_request_metadata: Optional[int] = None, + persist_cuda_graphs: Optional[bool] = False, ): super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits) @@ -360,6 +361,7 @@ def __init__( # Unified memory. self.unified_memory_level = unified_memory_level + self.persist_cuda_graphs = persist_cuda_graphs if unified_memory_level > 0: try: self.unified_memory_mempool = create_unified_mempool() diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 180fa628a1f..a854aeaec7a 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -167,6 +167,7 @@ def __init__( self.enable_chunked_prefill = enable_chunked_prefill self.inference_logging_step_interval = inference_logging_step_interval self.unified_memory_level = context.unified_memory_level + self.persist_cuda_graphs = context.persist_cuda_graphs if enable_cuda_graph is not None: self.cuda_graph_impl = "local" if enable_cuda_graph else "none" @@ -552,9 +553,9 @@ def suspend(self): self.context.deallocate_all_tensors() # Delete cuda graphs when not using unified memory at all (level 0) and - # `rl-persist-cuda-graphs` is not passed. For UVM levels 1 and 2, the context's tensors + # `--rl-persist-cuda-graphs` is not passed. For UVM levels 1 and 2, the context's tensors # maintain static memory addresses, so the cuda graphs are re-used. - if self.unified_memory_level == 0 and not args.rl_persist_cuda_graphs: + if self.unified_memory_level == 0 and not self.persist_cuda_graphs: delete_cuda_graphs() # Maintain references to requests before reset. @@ -596,7 +597,7 @@ def resume(self): # 0). For levels 1 and 2, the context's tensors maintain static # memory addresses, so the cuda graphs are re-used. capture_time = time.time() - if self.unified_memory_level == 0: + if self.unified_memory_level == 0 and not self.persist_cuda_graphs: self.create_cuda_graphs() capture_time = time.time() - capture_time diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py index ad22bd14ac9..86fb1ece025 100644 --- a/megatron/rl/inference/megatron.py +++ b/megatron/rl/inference/megatron.py @@ -136,6 +136,7 @@ def get_dynamic_inference_engine(args: Namespace, model: MegatronModule, inferen use_flashinfer_fused_rope=None, unified_memory_level=args.inference_dynamic_batching_unified_memory_level, metrics_writer=metrics_writer, + persist_cuda_graphs=args.rl_persist_cuda_graphs ) inference_wrapped_model = GPTInferenceWrapper(model, args, inference_context) From b9663033d608743d7ce092fad3f78a72f28cd89d Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 9 Dec 2025 12:56:25 -0800 Subject: [PATCH 07/18] Update to rl_training_cuda_graphs --- megatron/core/inference/engines/dynamic_engine.py | 2 +- megatron/rl/inference/megatron.py | 2 +- megatron/rl/rl_utils.py | 8 ++++++++ megatron/training/arguments.py | 4 ++-- 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index a854aeaec7a..3959ff5a529 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -553,7 +553,7 @@ def suspend(self): self.context.deallocate_all_tensors() # Delete cuda graphs when not using unified memory at all (level 0) and - # `--rl-persist-cuda-graphs` is not passed. For UVM levels 1 and 2, the context's tensors + # `--rl-training-cuda-graphs` is not passed. For UVM levels 1 and 2, the context's tensors # maintain static memory addresses, so the cuda graphs are re-used. if self.unified_memory_level == 0 and not self.persist_cuda_graphs: delete_cuda_graphs() diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py index 86fb1ece025..ce1129fa7b4 100644 --- a/megatron/rl/inference/megatron.py +++ b/megatron/rl/inference/megatron.py @@ -136,7 +136,7 @@ def get_dynamic_inference_engine(args: Namespace, model: MegatronModule, inferen use_flashinfer_fused_rope=None, unified_memory_level=args.inference_dynamic_batching_unified_memory_level, metrics_writer=metrics_writer, - persist_cuda_graphs=args.rl_persist_cuda_graphs + persist_cuda_graphs=args.rl_training_cuda_graphs ) inference_wrapped_model = GPTInferenceWrapper(model, args, inference_context) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index b16795fa9c0..af0940bfaf2 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -2275,6 +2275,10 @@ def megatron_rl_inference_mode( with nvtx_range("offload-optimizer-before-inference"): optimizer.offload_to_cpu() + # TODO: Remove this if statement once a change to `toggle_cuda_graphs` makes it safe to. + if cuda_graph_impl != "none" and not args.rl_training_cuda_graphs: + toggle_cuda_graphs(lang_module, cuda_graph_impl, reset_cuda_graphs=reset_cuda_graphs) + inference_interface = get_inference_interface(args, loop, model) with nvtx_range("onload-kv-cache-before-inference"): @@ -2320,6 +2324,10 @@ def megatron_rl_inference_mode( elif remove_kv_cache_during_training: inference_interface._inference_engine.context.memory_buffer = None + # TODO: Remove this if statement once a change to `toggle_cuda_graphs` makes it safe to. + if cuda_graph_impl != "none" and not args.rl_training_cuda_graphs: + toggle_cuda_graphs(lang_module, 'none', reset_cuda_graphs=reset_cuda_graphs) + if offload_optimizer_during_inference: with nvtx_range("onload-optimizer-after-inference"): optimizer.restore_from_cpu() diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index f1e9fdfa8f6..b76ab43e8bd 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1990,9 +1990,9 @@ def _add_rl_args(parser): help='Algorithm for distributing packed bins across ranks. ' 'fifo: first-in-first-out sequential distribution, ' 'round-robin: distribute bins cyclically across ranks for better load balancing') - group.add_argument('--rl-persist-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, + group.add_argument('--rl-training-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='If set, do not call `delete_cuda_graphs` when the inference engine is suspended. ' + help='If set, do not call `delete_cuda_graphs` or `toggle_cuda_graphs` when the inference engine is suspended. ' 'Use only when all training and inference cudagraphs and the KV cache fit on device.') return parser From 587ee8ca0cc4d35e931aa1c8dcc707faec45b08d Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Mon, 15 Dec 2025 09:01:27 -0800 Subject: [PATCH 08/18] Fix RNG state tracking for cudagraphs and grad mode --- .../abstract_model_inference_wrapper.py | 2 +- .../text_generation_controller.py | 2 +- megatron/core/transformer/cuda_graphs.py | 186 +++++++++++++++++- 3 files changed, 186 insertions(+), 4 deletions(-) diff --git a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py index 4c3d02b1fc5..95d476a9f83 100644 --- a/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py +++ b/megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py @@ -348,7 +348,7 @@ def forward_pass_with_pipeline_parallel_large_input_batch( # NOTE: Only returns the logits on the last pipeline stage return logits - @torch.no_grad() + @torch.inference_mode() def run_one_forward_step( self, inference_input: Dict[str, Any], recv_buffer_seq_len: Optional[int] = None ) -> torch.Tensor: diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index e430916a3e7..2ff0dcd579e 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -550,7 +550,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor) context = self.inference_wrapped_model.inference_context active_request_count = context.total_request_count - context.paused_request_count - with torch.no_grad(): + with torch.inference_mode(): logits = self.inference_wrapped_model.run_one_forward_step( {"tokens": input_ids, "position_ids": position_ids, "attention_mask": None} ) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 390b826e5da..594af2ef68e 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -536,6 +536,37 @@ def backward(ctx, *grads): ) return None, None, *output_grads +def _clone_nested_tensors(value: Any) -> Any: + """Recursively clone tensors inside nested containers.""" + if torch.is_tensor(value): + return value.clone() + if isinstance(value, (tuple, list)): + return type(value)(_clone_nested_tensors(v) for v in value) + if isinstance(value, dict): + return {k: _clone_nested_tensors(v) for k, v in value.items()} + return value + +def _ensure_generator_state_is_cudagraph_safe(gen: torch.Generator) -> torch.Generator: + """Make generator state safe for CUDA graph capture/replay. + + Generator state tensors can become inference tensors if created under `torch.inference_mode()`. + CUDA graph capture may later attempt in-place updates on that state; this fails for inference + tensors. Fix the generator *in-place* (preserving identity) by cloning its state outside + inference mode and setting it back. + """ + with torch.inference_mode(mode=False): + if hasattr(gen, "graphsafe_get_state"): + state = gen.graphsafe_get_state() + else: + state = gen.get_state() + + cloned_state = _clone_nested_tensors(state) + if hasattr(gen, "graphsafe_set_state"): + gen.graphsafe_set_state(cloned_state) + else: + gen.set_state(cloned_state) + + return gen class _CudaGraphRunner(torch.nn.Module): """Represents the execution of a cudagraphed module for a single microbatch. @@ -681,14 +712,149 @@ def create_fwd_graph(self, args, kwargs, clone_inputs=True): self.fwd_graph = torch.cuda.CUDAGraph() + # For cases with multiple active RNG states, e.g. TP. + rng_states = get_all_rng_states() + with torch.inference_mode(mode=False): + for gen in rng_states.values(): + self.fwd_graph.register_generator_state( + _ensure_generator_state_is_cudagraph_safe(gen)) + + # warmup again as case graph capture mode may execute a different codepath + for _ in range(self.num_warmup_steps): + with self.get_quantization_context(): + outputs = self.base_module.forward(*args, **kwargs) + if self.training and torch.is_grad_enabled(): + if isinstance(outputs, torch.Tensor): + outputs = (outputs,) + outputs = self.get_tensors(outputs) + grad_inputs = torch.autograd.grad( + outputs=tuple(o for o in outputs if o.requires_grad), + inputs=tuple(i for i in self.fwd_graph_input_surface if i.requires_grad), + grad_outputs=tuple( + torch.zeros_like(o) if o.requires_grad else None for o in outputs + ), + only_inputs=True, + allow_unused=True, + ) + + with self.get_quantization_context(): + torch.cuda.synchronize() + # Register default CUDA generators ourselves (fixed in-place to have normal tensors) + # before capture begins, to avoid inference-tensor state issues during capture. + with torch.inference_mode(mode=False): + for device_idx in range(torch.cuda.device_count()): + default_gen = torch.cuda.default_generators[device_idx] + self.fwd_graph.register_generator_state( + _ensure_generator_state_is_cudagraph_safe(default_gen) + ) + + with torch.cuda.graph( + self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local" + ): + outputs = self.base_module.forward(*args, **kwargs) + + # save cudagraph output buffer + if isinstance(outputs, torch.Tensor): + outputs = (outputs,) + self.fwd_graph_outputs = outputs + self.fwd_graph_output_surface = self.get_tensors(outputs) + + if self.training and torch.is_grad_enabled(): + assert ( + len(self.fwd_graph_output_surface) > 0 + ), """Tried graphing a moudule that returned no tensors in training mode, + however the graphed module must output at least one tensor, + so that a corresponding backward node may be registered in the autograd graph.""" + + # restore cached grads + for param in self.base_module.parameters(): + if hasattr(param, 'main_grad'): + saved_grad = save_main_grads.pop(0) + assert ( + param.main_grad.shape == saved_grad.shape + ), "Error restoring grads while cudagraphing!" + param.main_grad.copy_(saved_grad) + + if self.fp8_enabled or self.fp4_enabled: + restore_fp8_tensors([self.base_module], saved_fp8_tensors) + + # Unfreeze GC. + if FREEZE_GC: + gc.unfreeze() + + # gc.collect() drops references to unreachable tensors created during capture, + # returning their storage to the allocator to avoid a slowdown during replay. However, + # it forces expensive global garbage collection, so must be done only on the last layer + # per-device to avoid slowing down graph creation. + if self.is_last_layer: + gc.collect() + + def create_fwd_graph_old(self, args, kwargs, clone_inputs=True): + """Create a fwd cudagraph for this runner. Should be called inside + 'create_cudagraphs()'.""" + + # Freeze GC, to speed up capture time ~15-20x. + if FREEZE_GC: + gc.freeze() + + # save grads and other variables that may be affected by graph warmup + if self.training and torch.is_grad_enabled(): + save_main_grads = [ + param.main_grad.clone() + for param in self.base_module.parameters() + if hasattr(param, 'main_grad') + ] + + saved_fp8_tensors = None + + if self.fp8_enabled: + if is_te_min_version("1.13.0"): + saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp8_recipe) + else: + saved_fp8_tensors = save_fp8_tensors( + [self.base_module], self.fp8_recipe.amax_history_len + ) + elif self.fp4_enabled: + if is_te_min_version("2.7.0.dev0"): + saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp4_recipe) + else: + raise ValueError("FP4 requires TE >= 2.7.0.dev0 for NVFP4BlockScaling support.") + + if clone_inputs: + args, kwargs = self.zero_out_tensors(args, kwargs) + + input_tensors = self.get_tensors(args, kwargs) + self.fwd_graph_input_surface = input_tensors + tuple(self.base_module.parameters()) + + self.fwd_graph = torch.cuda.CUDAGraph() + # For cases with multiple active RNG states, e.g. TP. for _, state in get_all_rng_states().items(): + rng_states = get_all_rng_states() + with torch.inference_mode(mode=False): + for gen in rng_states.values(): + self.fwd_graph.register_generator_state( + _ensure_generator_state_is_cudagraph_safe(gen)) self.fwd_graph.register_generator_state(state) # warmup again as case graph capture mode may execute a different codepath for _ in range(self.num_warmup_steps): with self.get_quantization_context(): - outputs = self.base_module.forward(*args, **kwargs) + torch.cuda.synchronize() + # Register default CUDA generators ourselves (fixed in-place to have normal tensors) + # before capture begins, to avoid inference-tensor state issues during capture. + with torch.inference_mode(mode=False): + for device_idx in range(torch.cuda.device_count()): + default_gen = torch.cuda.default_generators[device_idx] + self.fwd_graph.register_generator_state( + _ensure_generator_state_is_cudagraph_safe(default_gen) + ) + + with torch.cuda.graph( + self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local" + ): + outputs = self.base_module.forward(*args, **kwargs) + # outputs = self.base_module.forward(*args, **kwargs) if self.training and torch.is_grad_enabled(): if isinstance(outputs, torch.Tensor): outputs = (outputs,) @@ -705,6 +871,15 @@ def create_fwd_graph(self, args, kwargs, clone_inputs=True): with self.get_quantization_context(): torch.cuda.synchronize() + # Register default CUDA generators ourselves (fixed in-place to have normal tensors) + # before capture begins, to avoid inference-tensor state issues during capture. + with torch.inference_mode(mode=False): + for device_idx in range(torch.cuda.device_count()): + default_gen = torch.cuda.default_generators[device_idx] + self.fwd_graph.register_generator_state( + _ensure_generator_state_is_cudagraph_safe(default_gen) + ) + with torch.cuda.graph( self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local" ): @@ -746,6 +921,8 @@ def create_fwd_graph(self, args, kwargs, clone_inputs=True): if self.is_last_layer: gc.collect() + + def create_bwd_graph(self, static_grad_outputs=None): """Create a bwd cudagraph for this runner. Should be called inside 'create_cudagraphs()'.""" @@ -758,7 +935,12 @@ def create_bwd_graph(self, static_grad_outputs=None): # For cases with multiple active RNG states, e.g. TP. for _, state in get_all_rng_states().items(): - self.bwd_graph.register_generator_state(state) + rng_states = get_all_rng_states() + with torch.inference_mode(mode=False): + for gen in rng_states.values(): + self.bwd_graph.register_generator_state( + _ensure_generator_state_is_cudagraph_safe(gen)) + # self.bwd_graph.register_generator_state(state) if static_grad_outputs is None: static_grad_outputs = tuple( From 484d9169fa4aca3fd1a446a733550360b83c13a4 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Mon, 15 Dec 2025 09:19:50 -0800 Subject: [PATCH 09/18] Include Peter's change on RNG state tracking safety for cudagraphs --- megatron/core/transformer/cuda_graphs.py | 90 ++++++++++++------------ 1 file changed, 44 insertions(+), 46 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 594af2ef68e..03fbd359226 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -175,6 +175,40 @@ def _determine_if_first_last_layer_of_this_vp_chunk(base_module): ) +def _clone_nested_tensors(value: Any) -> Any: + """Recursively clone tensors inside nested containers.""" + if torch.is_tensor(value): + return value.clone() + if isinstance(value, (tuple, list)): + return type(value)(_clone_nested_tensors(v) for v in value) + if isinstance(value, dict): + return {k: _clone_nested_tensors(v) for k, v in value.items()} + return value + + +def _ensure_generator_state_is_cudagraph_safe(gen: torch.Generator) -> torch.Generator: + """Make generator state safe for CUDA graph capture/replay. + + Generator state tensors can become inference tensors if created under `torch.inference_mode()`. + CUDA graph capture may later attempt in-place updates on that state; this fails for inference + tensors. Fix the generator *in-place* (preserving identity) by cloning its state outside + inference mode and setting it back. + """ + with torch.inference_mode(mode=False): + if hasattr(gen, "graphsafe_get_state"): + state = gen.graphsafe_get_state() + else: + state = gen.get_state() + + cloned_state = _clone_nested_tensors(state) + if hasattr(gen, "graphsafe_set_state"): + gen.graphsafe_set_state(cloned_state) + else: + gen.set_state(cloned_state) + + return gen + + class _CudagraphGlobalRecord: """A global datastructure that records of the ordering of all _CudaGraphRunner's first fwd or bwd passes. 'create_cudagraphs' will use this to create @@ -536,37 +570,6 @@ def backward(ctx, *grads): ) return None, None, *output_grads -def _clone_nested_tensors(value: Any) -> Any: - """Recursively clone tensors inside nested containers.""" - if torch.is_tensor(value): - return value.clone() - if isinstance(value, (tuple, list)): - return type(value)(_clone_nested_tensors(v) for v in value) - if isinstance(value, dict): - return {k: _clone_nested_tensors(v) for k, v in value.items()} - return value - -def _ensure_generator_state_is_cudagraph_safe(gen: torch.Generator) -> torch.Generator: - """Make generator state safe for CUDA graph capture/replay. - - Generator state tensors can become inference tensors if created under `torch.inference_mode()`. - CUDA graph capture may later attempt in-place updates on that state; this fails for inference - tensors. Fix the generator *in-place* (preserving identity) by cloning its state outside - inference mode and setting it back. - """ - with torch.inference_mode(mode=False): - if hasattr(gen, "graphsafe_get_state"): - state = gen.graphsafe_get_state() - else: - state = gen.get_state() - - cloned_state = _clone_nested_tensors(state) - if hasattr(gen, "graphsafe_set_state"): - gen.graphsafe_set_state(cloned_state) - else: - gen.set_state(cloned_state) - - return gen class _CudaGraphRunner(torch.nn.Module): """Represents the execution of a cudagraphed module for a single microbatch. @@ -717,7 +720,8 @@ def create_fwd_graph(self, args, kwargs, clone_inputs=True): with torch.inference_mode(mode=False): for gen in rng_states.values(): self.fwd_graph.register_generator_state( - _ensure_generator_state_is_cudagraph_safe(gen)) + _ensure_generator_state_is_cudagraph_safe(gen) + ) # warmup again as case graph capture mode may execute a different codepath for _ in range(self.num_warmup_steps): @@ -749,7 +753,7 @@ def create_fwd_graph(self, args, kwargs, clone_inputs=True): ) with torch.cuda.graph( - self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local" + self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local" ): outputs = self.base_module.forward(*args, **kwargs) @@ -761,7 +765,7 @@ def create_fwd_graph(self, args, kwargs, clone_inputs=True): if self.training and torch.is_grad_enabled(): assert ( - len(self.fwd_graph_output_surface) > 0 + len(self.fwd_graph_output_surface) > 0 ), """Tried graphing a moudule that returned no tensors in training mode, however the graphed module must output at least one tensor, so that a corresponding backward node may be registered in the autograd graph.""" @@ -771,7 +775,7 @@ def create_fwd_graph(self, args, kwargs, clone_inputs=True): if hasattr(param, 'main_grad'): saved_grad = save_main_grads.pop(0) assert ( - param.main_grad.shape == saved_grad.shape + param.main_grad.shape == saved_grad.shape ), "Error restoring grads while cudagraphing!" param.main_grad.copy_(saved_grad) @@ -834,7 +838,8 @@ def create_fwd_graph_old(self, args, kwargs, clone_inputs=True): with torch.inference_mode(mode=False): for gen in rng_states.values(): self.fwd_graph.register_generator_state( - _ensure_generator_state_is_cudagraph_safe(gen)) + _ensure_generator_state_is_cudagraph_safe(gen) + ) self.fwd_graph.register_generator_state(state) # warmup again as case graph capture mode may execute a different codepath @@ -851,10 +856,10 @@ def create_fwd_graph_old(self, args, kwargs, clone_inputs=True): ) with torch.cuda.graph( - self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local" + self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local" ): outputs = self.base_module.forward(*args, **kwargs) - # outputs = self.base_module.forward(*args, **kwargs) + if self.training and torch.is_grad_enabled(): if isinstance(outputs, torch.Tensor): outputs = (outputs,) @@ -921,8 +926,6 @@ def create_fwd_graph_old(self, args, kwargs, clone_inputs=True): if self.is_last_layer: gc.collect() - - def create_bwd_graph(self, static_grad_outputs=None): """Create a bwd cudagraph for this runner. Should be called inside 'create_cudagraphs()'.""" @@ -935,12 +938,7 @@ def create_bwd_graph(self, static_grad_outputs=None): # For cases with multiple active RNG states, e.g. TP. for _, state in get_all_rng_states().items(): - rng_states = get_all_rng_states() - with torch.inference_mode(mode=False): - for gen in rng_states.values(): - self.bwd_graph.register_generator_state( - _ensure_generator_state_is_cudagraph_safe(gen)) - # self.bwd_graph.register_generator_state(state) + self.bwd_graph.register_generator_state(state) if static_grad_outputs is None: static_grad_outputs = tuple( From b271f282db5dd2f579296ea0e341a628754c99b9 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Mon, 15 Dec 2025 09:19:50 -0800 Subject: [PATCH 10/18] Include Peter's change on RNG state tracking safety for cudagraphs --- megatron/core/transformer/cuda_graphs.py | 209 ++++------------------- 1 file changed, 37 insertions(+), 172 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 594af2ef68e..a962d31c7af 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -175,6 +175,40 @@ def _determine_if_first_last_layer_of_this_vp_chunk(base_module): ) +def _clone_nested_tensors(value: Any) -> Any: + """Recursively clone tensors inside nested containers.""" + if torch.is_tensor(value): + return value.clone() + if isinstance(value, (tuple, list)): + return type(value)(_clone_nested_tensors(v) for v in value) + if isinstance(value, dict): + return {k: _clone_nested_tensors(v) for k, v in value.items()} + return value + + +def _ensure_generator_state_is_cudagraph_safe(gen: torch.Generator) -> torch.Generator: + """Make generator state safe for CUDA graph capture/replay. + + Generator state tensors can become inference tensors if created under `torch.inference_mode()`. + CUDA graph capture may later attempt in-place updates on that state; this fails for inference + tensors. Fix the generator *in-place* (preserving identity) by cloning its state outside + inference mode and setting it back. + """ + with torch.inference_mode(mode=False): + if hasattr(gen, "graphsafe_get_state"): + state = gen.graphsafe_get_state() + else: + state = gen.get_state() + + cloned_state = _clone_nested_tensors(state) + if hasattr(gen, "graphsafe_set_state"): + gen.graphsafe_set_state(cloned_state) + else: + gen.set_state(cloned_state) + + return gen + + class _CudagraphGlobalRecord: """A global datastructure that records of the ordering of all _CudaGraphRunner's first fwd or bwd passes. 'create_cudagraphs' will use this to create @@ -536,37 +570,6 @@ def backward(ctx, *grads): ) return None, None, *output_grads -def _clone_nested_tensors(value: Any) -> Any: - """Recursively clone tensors inside nested containers.""" - if torch.is_tensor(value): - return value.clone() - if isinstance(value, (tuple, list)): - return type(value)(_clone_nested_tensors(v) for v in value) - if isinstance(value, dict): - return {k: _clone_nested_tensors(v) for k, v in value.items()} - return value - -def _ensure_generator_state_is_cudagraph_safe(gen: torch.Generator) -> torch.Generator: - """Make generator state safe for CUDA graph capture/replay. - - Generator state tensors can become inference tensors if created under `torch.inference_mode()`. - CUDA graph capture may later attempt in-place updates on that state; this fails for inference - tensors. Fix the generator *in-place* (preserving identity) by cloning its state outside - inference mode and setting it back. - """ - with torch.inference_mode(mode=False): - if hasattr(gen, "graphsafe_get_state"): - state = gen.graphsafe_get_state() - else: - state = gen.get_state() - - cloned_state = _clone_nested_tensors(state) - if hasattr(gen, "graphsafe_set_state"): - gen.graphsafe_set_state(cloned_state) - else: - gen.set_state(cloned_state) - - return gen class _CudaGraphRunner(torch.nn.Module): """Represents the execution of a cudagraphed module for a single microbatch. @@ -717,144 +720,13 @@ def create_fwd_graph(self, args, kwargs, clone_inputs=True): with torch.inference_mode(mode=False): for gen in rng_states.values(): self.fwd_graph.register_generator_state( - _ensure_generator_state_is_cudagraph_safe(gen)) - - # warmup again as case graph capture mode may execute a different codepath - for _ in range(self.num_warmup_steps): - with self.get_quantization_context(): - outputs = self.base_module.forward(*args, **kwargs) - if self.training and torch.is_grad_enabled(): - if isinstance(outputs, torch.Tensor): - outputs = (outputs,) - outputs = self.get_tensors(outputs) - grad_inputs = torch.autograd.grad( - outputs=tuple(o for o in outputs if o.requires_grad), - inputs=tuple(i for i in self.fwd_graph_input_surface if i.requires_grad), - grad_outputs=tuple( - torch.zeros_like(o) if o.requires_grad else None for o in outputs - ), - only_inputs=True, - allow_unused=True, - ) - - with self.get_quantization_context(): - torch.cuda.synchronize() - # Register default CUDA generators ourselves (fixed in-place to have normal tensors) - # before capture begins, to avoid inference-tensor state issues during capture. - with torch.inference_mode(mode=False): - for device_idx in range(torch.cuda.device_count()): - default_gen = torch.cuda.default_generators[device_idx] - self.fwd_graph.register_generator_state( - _ensure_generator_state_is_cudagraph_safe(default_gen) - ) - - with torch.cuda.graph( - self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local" - ): - outputs = self.base_module.forward(*args, **kwargs) - - # save cudagraph output buffer - if isinstance(outputs, torch.Tensor): - outputs = (outputs,) - self.fwd_graph_outputs = outputs - self.fwd_graph_output_surface = self.get_tensors(outputs) - - if self.training and torch.is_grad_enabled(): - assert ( - len(self.fwd_graph_output_surface) > 0 - ), """Tried graphing a moudule that returned no tensors in training mode, - however the graphed module must output at least one tensor, - so that a corresponding backward node may be registered in the autograd graph.""" - - # restore cached grads - for param in self.base_module.parameters(): - if hasattr(param, 'main_grad'): - saved_grad = save_main_grads.pop(0) - assert ( - param.main_grad.shape == saved_grad.shape - ), "Error restoring grads while cudagraphing!" - param.main_grad.copy_(saved_grad) - - if self.fp8_enabled or self.fp4_enabled: - restore_fp8_tensors([self.base_module], saved_fp8_tensors) - - # Unfreeze GC. - if FREEZE_GC: - gc.unfreeze() - - # gc.collect() drops references to unreachable tensors created during capture, - # returning their storage to the allocator to avoid a slowdown during replay. However, - # it forces expensive global garbage collection, so must be done only on the last layer - # per-device to avoid slowing down graph creation. - if self.is_last_layer: - gc.collect() - - def create_fwd_graph_old(self, args, kwargs, clone_inputs=True): - """Create a fwd cudagraph for this runner. Should be called inside - 'create_cudagraphs()'.""" - - # Freeze GC, to speed up capture time ~15-20x. - if FREEZE_GC: - gc.freeze() - - # save grads and other variables that may be affected by graph warmup - if self.training and torch.is_grad_enabled(): - save_main_grads = [ - param.main_grad.clone() - for param in self.base_module.parameters() - if hasattr(param, 'main_grad') - ] - - saved_fp8_tensors = None - - if self.fp8_enabled: - if is_te_min_version("1.13.0"): - saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp8_recipe) - else: - saved_fp8_tensors = save_fp8_tensors( - [self.base_module], self.fp8_recipe.amax_history_len + _ensure_generator_state_is_cudagraph_safe(gen) ) - elif self.fp4_enabled: - if is_te_min_version("2.7.0.dev0"): - saved_fp8_tensors = save_fp8_tensors([self.base_module], self.fp4_recipe) - else: - raise ValueError("FP4 requires TE >= 2.7.0.dev0 for NVFP4BlockScaling support.") - - if clone_inputs: - args, kwargs = self.zero_out_tensors(args, kwargs) - - input_tensors = self.get_tensors(args, kwargs) - self.fwd_graph_input_surface = input_tensors + tuple(self.base_module.parameters()) - - self.fwd_graph = torch.cuda.CUDAGraph() - - # For cases with multiple active RNG states, e.g. TP. - for _, state in get_all_rng_states().items(): - rng_states = get_all_rng_states() - with torch.inference_mode(mode=False): - for gen in rng_states.values(): - self.fwd_graph.register_generator_state( - _ensure_generator_state_is_cudagraph_safe(gen)) - self.fwd_graph.register_generator_state(state) # warmup again as case graph capture mode may execute a different codepath for _ in range(self.num_warmup_steps): with self.get_quantization_context(): - torch.cuda.synchronize() - # Register default CUDA generators ourselves (fixed in-place to have normal tensors) - # before capture begins, to avoid inference-tensor state issues during capture. - with torch.inference_mode(mode=False): - for device_idx in range(torch.cuda.device_count()): - default_gen = torch.cuda.default_generators[device_idx] - self.fwd_graph.register_generator_state( - _ensure_generator_state_is_cudagraph_safe(default_gen) - ) - - with torch.cuda.graph( - self.fwd_graph, pool=self.fwd_mempool, capture_error_mode="thread_local" - ): - outputs = self.base_module.forward(*args, **kwargs) - # outputs = self.base_module.forward(*args, **kwargs) + outputs = self.base_module.forward(*args, **kwargs) if self.training and torch.is_grad_enabled(): if isinstance(outputs, torch.Tensor): outputs = (outputs,) @@ -921,8 +793,6 @@ def create_fwd_graph_old(self, args, kwargs, clone_inputs=True): if self.is_last_layer: gc.collect() - - def create_bwd_graph(self, static_grad_outputs=None): """Create a bwd cudagraph for this runner. Should be called inside 'create_cudagraphs()'.""" @@ -935,12 +805,7 @@ def create_bwd_graph(self, static_grad_outputs=None): # For cases with multiple active RNG states, e.g. TP. for _, state in get_all_rng_states().items(): - rng_states = get_all_rng_states() - with torch.inference_mode(mode=False): - for gen in rng_states.values(): - self.bwd_graph.register_generator_state( - _ensure_generator_state_is_cudagraph_safe(gen)) - # self.bwd_graph.register_generator_state(state) + self.bwd_graph.register_generator_state(state) if static_grad_outputs is None: static_grad_outputs = tuple( From 075530f5fd9b93d3bd987a5a3705340bb86e4d29 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Mon, 12 Jan 2026 10:20:28 -0800 Subject: [PATCH 11/18] Flip training cudagraphs back on after refit MR --- megatron/rl/rl_utils.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index 248f2545e18..d1a946a9cfe 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -989,13 +989,7 @@ def prepare_data_for_update( nvtx_range = get_nvtx_range() runtime_state = get_rl_runtime_state() - # RL policy updates + logprob computations should run eagerly; only rollout generation - # (inference engine) should use CUDA graphs until training cuda-graphs MR goes in. - # In the single-model case this is naturally handled by `megatron_rl_inference_mode` - # toggling graphs on/off around inference. In the refit case (separate inference_model), - # we must explicitly keep the training model (this `model`) with CUDA graphs disabled, - # otherwise training/logprobs can get cudagraphed. - if args.cuda_graph_impl != "none": + if args.cuda_graph_impl != "none" and not args.rl_training_cuda_graphs: lang_module = ( model[0].module.module if hasattr(model[0].module, "module") else model[0].module ) From 3e6fafc4283d697fd35c3708eef275a55eebba9a Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene <34819528+tdene@users.noreply.github.com> Date: Mon, 12 Jan 2026 15:06:56 -0600 Subject: [PATCH 12/18] Fix RL sequence packing bin size (#2909) --- megatron/rl/sequence_packing_utils.py | 14 +- megatron/training/arguments.py | 2 +- .../golden_values_dev_dgx_h100.json | 352 ++++++++--------- .../golden_values_dev_dgx_h100.json | 372 +++++++++--------- .../golden_values_dev_dgx_h100.json | 366 ++++++++--------- 5 files changed, 557 insertions(+), 549 deletions(-) diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py index 1e9063f0947..ddff9555ccf 100644 --- a/megatron/rl/sequence_packing_utils.py +++ b/megatron/rl/sequence_packing_utils.py @@ -430,19 +430,25 @@ def create_packed_seq_params(packing_context: PackingContext): cached_packed_seq_params = [] packing_info = packing_context.packing_info bin_size = packing_context.bin_size + max_sequences_per_bin = packing_context.packer.max_sequences_per_bin device = packing_context.packed_trajs.device for bin_idx in range(len(packing_context.packed_trajs)): params = create_packed_seq_params_for_bin( packing_info=packing_info, bin_idx=bin_idx, bin_size=bin_size, + max_sequences_per_bin=max_sequences_per_bin, device=device, ) cached_packed_seq_params.append(params) return cached_packed_seq_params def create_packed_seq_params_for_bin( - packing_info: PackingInfo, bin_idx: int, bin_size: int, device: torch.device + packing_info: PackingInfo, + bin_idx: int, + bin_size: int, + max_sequences_per_bin: int, + device: torch.device ) -> Optional[PackedSeqParams]: """Create PackedSeqParams for a single bin to enable proper attention masking in TE. @@ -454,6 +460,7 @@ def create_packed_seq_params_for_bin( packing_info: PackingInfo object containing packing metadata from SequencePacker bin_idx: Index of the bin to create params for bin_size: Size of the bin (padded sequence length) + max_sequences_per_bin: Maximum number of sequences per bin device: Device to create tensors on Returns: @@ -476,8 +483,8 @@ def create_packed_seq_params_for_bin( # Pad cu_seqlens to bin_size by repeating the last value (creates zero-length ghost sequences) # This ensures a fixed tensor size for CUDA graph compatibility - if len(cu_seqlens) < bin_size: - out = cu_seqlens.new_full((bin_size,), bin_size) + if len(cu_seqlens) < max_sequences_per_bin: + out = cu_seqlens.new_full((max_sequences_per_bin,), bin_size) out[:len(cu_seqlens)] = cu_seqlens cu_seqlens = out @@ -1038,6 +1045,7 @@ def pack_all_trajectories(trajs, generation_masks, inference_logprobs, global_ad packing_info=packing_info, bin_idx=bin_idx, bin_size=bin_size, + max_sequences_per_bin=max_sequences_per_bin, device=packed_trajs.device, ) for bin_idx in range(len(packed_trajs)) ] diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 8df53e61863..ce9f155d68a 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2094,7 +2094,7 @@ def _add_rl_args(parser): help='If set, calculate the intra-group similarity of rollouts.') group.add_argument('--rl-use-sequence-packing', action=argparse.BooleanOptionalAction, type=bool, default=False, help='Enable sequence packing') - group.add_argument('--rl-sequence-packing-max-sequences-per-bin', type=int, default=50, + group.add_argument('--rl-sequence-packing-max-sequences-per-bin', type=int, default=32, help='Maximum number of sequences that can be packed into a single bin. ') group.add_argument('--rl-sequence-packing-algo', type=str, default='fifo', choices=['fifo', 'round-robin'], diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json index e58cb5d3349..78e76d55735 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json @@ -4,31 +4,31 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 0.0, - "2": -0.04564, + "1": 0.00762, + "2": 0.0, "3": 0.0, "4": 0.0, "5": 0.0, - "6": 0.0, - "7": 0.0, + "6": 0.05397, + "7": 0.01964, "8": 0.0, - "9": 0.04101, + "9": 0.0, "10": 0.0, "11": 0.0, "12": 0.0, "13": 0.0, - "14": 0.05164, + "14": 0.0, "15": 0.0, - "16": 0.0, + "16": 0.02209, "17": 0.0, - "18": 0.03448, - "19": 0.00346, + "18": 0.0, + "19": 0.0, "20": 0.0, "21": 0.0, "22": 0.0, "23": 0.0, - "24": 0.05792, - "25": 0.03686, + "24": 0.0, + "25": 0.0, "26": 0.0, "27": 0.0, "28": 0.0, @@ -49,11 +49,11 @@ "43": 0.0, "44": 0.0, "45": 0.0, - "46": 0.05118, + "46": 0.0, "47": 0.0, "48": 0.0, "49": 0.0, - "50": 0.0 + "50": 0.04447 } }, "num-zeros": { @@ -61,31 +61,31 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 583687296.0, - "2": 70.0, + "1": 58.0, + "2": 583687296.0, "3": 583687296.0, "4": 583687296.0, "5": 583687296.0, - "6": 583687296.0, - "7": 583687296.0, + "6": 35.0, + "7": 52.0, "8": 583687296.0, - "9": 19.0, + "9": 583687296.0, "10": 583687296.0, "11": 583687296.0, "12": 583687296.0, "13": 583687296.0, - "14": 20.0, + "14": 583687296.0, "15": 583687296.0, - "16": 583687296.0, + "16": 52.0, "17": 583687296.0, - "18": 53.0, - "19": 54.0, + "18": 583687296.0, + "19": 583687296.0, "20": 583687296.0, "21": 583687296.0, "22": 583687296.0, "23": 583687296.0, - "24": 40.0, - "25": 53.0, + "24": 583687296.0, + "25": 583687296.0, "26": 583687296.0, "27": 583687296.0, "28": 583687296.0, @@ -106,11 +106,11 @@ "43": 583687296.0, "44": 583687296.0, "45": 583687296.0, - "46": 30.0, + "46": 583687296.0, "47": 583687296.0, "48": 583687296.0, "49": 583687296.0, - "50": 583687296.0 + "50": 45.0 } }, "mem-allocated-bytes": { @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 99613442048.0, - "2": 99615326208.0, - "3": 99615236096.0, - "4": 99615236096.0, - "5": 99615219712.0, - "6": 99615203328.0, - "7": 99615203328.0, - "8": 99615211520.0, - "9": 99615178752.0, - "10": 99615154176.0, - "11": 99615105024.0, - "12": 99615105024.0, - "13": 99615105024.0, - "14": 99615105024.0, - "15": 99615113216.0, - "16": 99615113216.0, - "17": 99615113216.0, - "18": 99615121408.0, - "19": 99615113216.0, - "20": 99615121408.0, - "21": 99615121408.0, - "22": 99615113216.0, - "23": 99615121408.0, - "24": 99615113216.0, - "25": 99615113216.0, - "26": 99615113216.0, - "27": 99615113216.0, - "28": 99615121408.0, - "29": 99615121408.0, - "30": 99615121408.0, - "31": 99615121408.0, - "32": 99615121408.0, - "33": 99615121408.0, - "34": 99615121408.0, - "35": 99615121408.0, - "36": 99615129600.0, - "37": 99615121408.0, - "38": 99615129600.0, - "39": 99615121408.0, - "40": 99615129600.0, - "41": 99615121408.0, - "42": 99615129600.0, - "43": 99615129600.0, - "44": 99615129600.0, - "45": 99615129600.0, - "46": 99615121408.0, - "47": 99615121408.0, - "48": 99615129600.0, - "49": 99615129600.0, - "50": 99615121408.0 + "1": 55289954304.0, + "2": 55292747776.0, + "3": 55292731392.0, + "4": 55292891136.0, + "5": 55292878848.0, + "6": 55292878848.0, + "7": 55292878848.0, + "8": 55292788736.0, + "9": 55292788736.0, + "10": 55292788736.0, + "11": 55292792832.0, + "12": 55292792832.0, + "13": 55292792832.0, + "14": 55292792832.0, + "15": 55292792832.0, + "16": 55292796928.0, + "17": 55292796928.0, + "18": 55292801024.0, + "19": 55292805120.0, + "20": 55292801024.0, + "21": 55292801024.0, + "22": 55292796928.0, + "23": 55292801024.0, + "24": 55292796928.0, + "25": 55292801024.0, + "26": 55292796928.0, + "27": 55292796928.0, + "28": 55292801024.0, + "29": 55292801024.0, + "30": 55292805120.0, + "31": 55292805120.0, + "32": 55292805120.0, + "33": 55292805120.0, + "34": 55292805120.0, + "35": 55292805120.0, + "36": 55292805120.0, + "37": 55292801024.0, + "38": 55292801024.0, + "39": 55292801024.0, + "40": 55292805120.0, + "41": 55292805120.0, + "42": 55292805120.0, + "43": 55292801024.0, + "44": 55292796928.0, + "45": 55292801024.0, + "46": 55292801024.0, + "47": 55292801024.0, + "48": 55292801024.0, + "49": 55292805120.0, + "50": 55292805120.0 } }, "mem-max-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 99613450240.0, - "2": 101703827456.0, - "3": 101704925184.0, - "4": 101704925184.0, - "5": 101704925184.0, - "6": 101704925184.0, - "7": 101704925184.0, - "8": 101708570624.0, - "9": 101708570624.0, - "10": 101708570624.0, - "11": 101708570624.0, - "12": 101708570624.0, - "13": 101708570624.0, - "14": 101708570624.0, - "15": 101708570624.0, - "16": 101708570624.0, - "17": 101708570624.0, - "18": 101708570624.0, - "19": 101708570624.0, - "20": 101708570624.0, - "21": 101708570624.0, - "22": 101708570624.0, - "23": 101708570624.0, - "24": 101708570624.0, - "25": 101708570624.0, - "26": 101708570624.0, - "27": 101708570624.0, - "28": 101708570624.0, - "29": 101708570624.0, - "30": 101708570624.0, - "31": 101708570624.0, - "32": 101708570624.0, - "33": 101708570624.0, - "34": 101708570624.0, - "35": 101708570624.0, - "36": 101708570624.0, - "37": 101708570624.0, - "38": 101708570624.0, - "39": 101708570624.0, - "40": 101708570624.0, - "41": 101708570624.0, - "42": 101708570624.0, - "43": 101708570624.0, - "44": 101708570624.0, - "45": 101708570624.0, - "46": 101708570624.0, - "47": 101708570624.0, - "48": 101708570624.0, - "49": 101708570624.0, - "50": 101708570624.0 + "1": 55289958400.0, + "2": 57103880192.0, + "3": 57104392192.0, + "4": 57104416768.0, + "5": 57104416768.0, + "6": 57104416768.0, + "7": 57104416768.0, + "8": 57104416768.0, + "9": 57104416768.0, + "10": 57104416768.0, + "11": 57104416768.0, + "12": 57104416768.0, + "13": 57104416768.0, + "14": 57104416768.0, + "15": 57104416768.0, + "16": 57104416768.0, + "17": 57104416768.0, + "18": 57104416768.0, + "19": 57104416768.0, + "20": 57104416768.0, + "21": 57104416768.0, + "22": 57104416768.0, + "23": 57104416768.0, + "24": 57104416768.0, + "25": 57104416768.0, + "26": 57104416768.0, + "27": 57104416768.0, + "28": 57104416768.0, + "29": 57104416768.0, + "30": 57104416768.0, + "31": 57104416768.0, + "32": 57104416768.0, + "33": 57104416768.0, + "34": 57104416768.0, + "35": 57104416768.0, + "36": 57104416768.0, + "37": 57104416768.0, + "38": 57104416768.0, + "39": 57104416768.0, + "40": 57104416768.0, + "41": 57104416768.0, + "42": 57104416768.0, + "43": 57104416768.0, + "44": 57104416768.0, + "45": 57104416768.0, + "46": 57104416768.0, + "47": 57104416768.0, + "48": 57104416768.0, + "49": 57104416768.0, + "50": 57104416768.0 } }, "iteration-time": { @@ -232,56 +232,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 37.07577, - "2": 5.51337, - "3": 4.10557, - "4": 3.55106, - "5": 3.45444, - "6": 3.48579, - "7": 3.39066, - "8": 3.49615, - "9": 3.63661, - "10": 3.5452, - "11": 3.52083, - "12": 3.44924, - "13": 3.34272, - "14": 3.39596, - "15": 3.42629, - "16": 3.31287, - "17": 3.32152, - "18": 3.39771, - "19": 3.42125, - "20": 3.61702, - "21": 3.45153, - "22": 3.35039, - "23": 3.39949, - "24": 3.32904, - "25": 3.36768, - "26": 3.34536, - "27": 3.30363, - "28": 3.36734, - "29": 3.41942, - "30": 3.38079, - "31": 3.35877, - "32": 3.34474, - "33": 3.27045, - "34": 3.18637, - "35": 3.24522, - "36": 3.34784, - "37": 3.33885, - "38": 3.37193, - "39": 3.31138, - "40": 3.25321, - "41": 3.21574, - "42": 3.24275, - "43": 3.27418, - "44": 3.30596, - "45": 3.30984, - "46": 3.36254, - "47": 3.43668, - "48": 3.27358, - "49": 3.25891, - "50": 3.34573 + "1": 38.24908, + "2": 4.52458, + "3": 3.69393, + "4": 3.38577, + "5": 3.41862, + "6": 3.27421, + "7": 3.32023, + "8": 3.83723, + "9": 4.07373, + "10": 3.47799, + "11": 3.27499, + "12": 3.37017, + "13": 3.3918, + "14": 3.25114, + "15": 3.29905, + "16": 3.29943, + "17": 3.50383, + "18": 3.56844, + "19": 3.30276, + "20": 3.34553, + "21": 3.29165, + "22": 3.30348, + "23": 3.33814, + "24": 3.31525, + "25": 3.29337, + "26": 3.26119, + "27": 3.5167, + "28": 3.2312, + "29": 3.45063, + "30": 3.3088, + "31": 3.32522, + "32": 3.28154, + "33": 3.23551, + "34": 3.20003, + "35": 3.25844, + "36": 3.67071, + "37": 3.1881, + "38": 3.30757, + "39": 3.32895, + "40": 3.29602, + "41": 3.25522, + "42": 3.28932, + "43": 3.32204, + "44": 3.26419, + "45": 3.75371, + "46": 3.23126, + "47": 3.25929, + "48": 3.19512, + "49": 3.32815, + "50": 3.25617 } } } diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/golden_values_dev_dgx_h100.json index 4206fac0d0d..821e71269b6 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1_pp1_dp8_583m_throughputtest_github/golden_values_dev_dgx_h100.json @@ -4,22 +4,22 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 0.05412, - "2": 0.04523, - "3": 0.09444, - "4": 0.04451, - "5": 0.05201, + "1": 0.0, + "2": 0.0, + "3": 0.01527, + "4": 0.0, + "5": 0.0, "6": 0.0, "7": 0.0, - "8": 0.04699, + "8": 0.0, "9": 0.0, "10": 0.0, "11": 0.0, - "12": 0.0, - "13": 0.0, - "14": 0.03773, + "12": 0.04095, + "13": 0.04612, + "14": 0.0, "15": 0.0, - "16": 0.0, + "16": 0.03061, "17": 0.0, "18": 0.0, "19": 0.0, @@ -28,7 +28,7 @@ "22": 0.0, "23": 0.0, "24": 0.0, - "25": 0.0, + "25": 0.04155, "26": 0.0, "27": 0.0, "28": 0.0, @@ -36,24 +36,24 @@ "30": 0.0, "31": 0.0, "32": 0.0, - "33": 0.0, + "33": 0.04325, "34": 0.0, "35": 0.0, "36": 0.0, - "37": 0.04296, + "37": 0.0, "38": 0.0, - "39": 0.0, + "39": 0.02738, "40": 0.0, "41": 0.0, "42": 0.0, "43": 0.0, - "44": 0.0, + "44": 0.00715, "45": 0.0, "46": 0.0, - "47": 0.05684, - "48": 0.04259, + "47": 0.0, + "48": 0.04492, "49": 0.0, - "50": 0.02801 + "50": 0.0 } }, "num-zeros": { @@ -61,22 +61,22 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 24.0, - "2": 44.0, - "3": 20.0, - "4": 30.0, - "5": 37.0, + "1": 583687296.0, + "2": 583687296.0, + "3": 55.0, + "4": 583687296.0, + "5": 583687296.0, "6": 583687296.0, "7": 583687296.0, - "8": 53.0, + "8": 583687296.0, "9": 583687296.0, "10": 583687296.0, "11": 583687296.0, - "12": 583687296.0, - "13": 583687296.0, - "14": 50.0, + "12": 38.0, + "13": 37.0, + "14": 583687296.0, "15": 583687296.0, - "16": 583687296.0, + "16": 49.0, "17": 583687296.0, "18": 583687296.0, "19": 583687296.0, @@ -85,7 +85,7 @@ "22": 583687296.0, "23": 583687296.0, "24": 583687296.0, - "25": 583687296.0, + "25": 28.0, "26": 583687296.0, "27": 583687296.0, "28": 583687296.0, @@ -93,24 +93,24 @@ "30": 583687296.0, "31": 583687296.0, "32": 583687296.0, - "33": 583687296.0, + "33": 20.0, "34": 583687296.0, "35": 583687296.0, "36": 583687296.0, - "37": 46.0, + "37": 583687296.0, "38": 583687296.0, - "39": 583687296.0, + "39": 33.0, "40": 583687296.0, "41": 583687296.0, "42": 583687296.0, "43": 583687296.0, - "44": 583687296.0, + "44": 55.0, "45": 583687296.0, "46": 583687296.0, - "47": 33.0, - "48": 19.0, + "47": 583687296.0, + "48": 42.0, "49": 583687296.0, - "50": 41.0 + "50": 583687296.0 } }, "mem-allocated-bytes": { @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 99614597120.0, - "2": 99614261248.0, - "3": 99614236672.0, - "4": 99614228480.0, - "5": 99614220288.0, - "6": 99614212096.0, - "7": 99614212096.0, - "8": 99614212096.0, - "9": 99614146560.0, - "10": 99614146560.0, - "11": 99614146560.0, - "12": 99614146560.0, - "13": 99614146560.0, - "14": 99614146560.0, - "15": 99614154752.0, - "16": 99614154752.0, - "17": 99614154752.0, - "18": 99614154752.0, - "19": 99614154752.0, - "20": 99614154752.0, - "21": 99614154752.0, - "22": 99614154752.0, - "23": 99614162944.0, - "24": 99614162944.0, - "25": 99614162944.0, - "26": 99614162944.0, - "27": 99614162944.0, - "28": 99614162944.0, - "29": 99614162944.0, - "30": 99614162944.0, - "31": 99614162944.0, - "32": 99614171136.0, - "33": 99614171136.0, - "34": 99614162944.0, - "35": 99614162944.0, - "36": 99614162944.0, - "37": 99614162944.0, - "38": 99614154752.0, - "39": 99614162944.0, - "40": 99614162944.0, - "41": 99614162944.0, - "42": 99614154752.0, - "43": 99614154752.0, - "44": 99614154752.0, - "45": 99614154752.0, - "46": 99614154752.0, - "47": 99614154752.0, - "48": 99614154752.0, - "49": 99614154752.0, - "50": 99614162944.0 + "1": 55288418304.0, + "2": 55292735488.0, + "3": 55292968960.0, + "4": 55292956672.0, + "5": 55292952576.0, + "6": 55292948480.0, + "7": 55292948480.0, + "8": 55292899328.0, + "9": 55292899328.0, + "10": 55292891136.0, + "11": 55292887040.0, + "12": 55292882944.0, + "13": 55292878848.0, + "14": 55292862464.0, + "15": 55292862464.0, + "16": 55292854272.0, + "17": 55292858368.0, + "18": 55292854272.0, + "19": 55292858368.0, + "20": 55292854272.0, + "21": 55292854272.0, + "22": 55292850176.0, + "23": 55292850176.0, + "24": 55292846080.0, + "25": 55292854272.0, + "26": 55292858368.0, + "27": 55292858368.0, + "28": 55292858368.0, + "29": 55292862464.0, + "30": 55292862464.0, + "31": 55292862464.0, + "32": 55292862464.0, + "33": 55292862464.0, + "34": 55292866560.0, + "35": 55292866560.0, + "36": 55292866560.0, + "37": 55292870656.0, + "38": 55292870656.0, + "39": 55292870656.0, + "40": 55292866560.0, + "41": 55292854272.0, + "42": 55292854272.0, + "43": 55292854272.0, + "44": 55292850176.0, + "45": 55292850176.0, + "46": 55292854272.0, + "47": 55292854272.0, + "48": 55292854272.0, + "49": 55292850176.0, + "50": 55292850176.0 } }, "mem-max-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 99614605312.0, - "2": 101701984256.0, - "3": 101701984256.0, - "4": 101701984256.0, - "5": 101701984256.0, - "6": 101701984256.0, - "7": 101701984256.0, - "8": 101701984256.0, - "9": 101701984256.0, - "10": 101705539584.0, - "11": 101705539584.0, - "12": 101705539584.0, - "13": 101705547776.0, - "14": 101705547776.0, - "15": 101705547776.0, - "16": 101705547776.0, - "17": 101705547776.0, - "18": 101705547776.0, - "19": 101705547776.0, - "20": 101705547776.0, - "21": 101705547776.0, - "22": 101705547776.0, - "23": 101705555968.0, - "24": 101705555968.0, - "25": 101705555968.0, - "26": 101705555968.0, - "27": 101705555968.0, - "28": 101705564160.0, - "29": 101705564160.0, - "30": 101705564160.0, - "31": 101705564160.0, - "32": 101705564160.0, - "33": 101705564160.0, - "34": 101705564160.0, - "35": 101705564160.0, - "36": 101705564160.0, - "37": 101705564160.0, - "38": 101705564160.0, - "39": 101705564160.0, - "40": 101705564160.0, - "41": 101705564160.0, - "42": 101705564160.0, - "43": 101705564160.0, - "44": 101705564160.0, - "45": 101705564160.0, - "46": 101705564160.0, - "47": 101705564160.0, - "48": 101705564160.0, - "49": 101705564160.0, - "50": 101705564160.0 + "1": 55288422400.0, + "2": 57104965632.0, + "3": 57104965632.0, + "4": 57104965632.0, + "5": 57104965632.0, + "6": 57104965632.0, + "7": 57104965632.0, + "8": 57104965632.0, + "9": 57104965632.0, + "10": 57104965632.0, + "11": 57104965632.0, + "12": 57104965632.0, + "13": 57104965632.0, + "14": 57104965632.0, + "15": 57104965632.0, + "16": 57104965632.0, + "17": 57104965632.0, + "18": 57104965632.0, + "19": 57104965632.0, + "20": 57104965632.0, + "21": 57104965632.0, + "22": 57104965632.0, + "23": 57104965632.0, + "24": 57104965632.0, + "25": 57104965632.0, + "26": 57104965632.0, + "27": 57104965632.0, + "28": 57104965632.0, + "29": 57104965632.0, + "30": 57104965632.0, + "31": 57104965632.0, + "32": 57104965632.0, + "33": 57104965632.0, + "34": 57104965632.0, + "35": 57104965632.0, + "36": 57104965632.0, + "37": 57104965632.0, + "38": 57104965632.0, + "39": 57104965632.0, + "40": 57104965632.0, + "41": 57104965632.0, + "42": 57104965632.0, + "43": 57104965632.0, + "44": 57104965632.0, + "45": 57104965632.0, + "46": 57104965632.0, + "47": 57104965632.0, + "48": 57104965632.0, + "49": 57104965632.0, + "50": 57104965632.0 } }, "iteration-time": { @@ -232,56 +232,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 130.25253, - "2": 9.88948, - "3": 8.72032, - "4": 8.5427, - "5": 8.26483, - "6": 8.59126, - "7": 8.02799, - "8": 8.21142, - "9": 8.57808, - "10": 8.03187, - "11": 8.04941, - "12": 8.01158, - "13": 8.18497, - "14": 8.13065, - "15": 8.12456, - "16": 8.0261, - "17": 8.24415, - "18": 8.12356, - "19": 8.01872, - "20": 7.96605, - "21": 8.02618, - "22": 7.98249, - "23": 8.03059, - "24": 7.87244, - "25": 7.92321, - "26": 7.99325, - "27": 8.03815, - "28": 8.0646, - "29": 8.03226, - "30": 7.92917, - "31": 8.0803, - "32": 7.9272, - "33": 7.93803, - "34": 7.9555, - "35": 8.10923, - "36": 8.01863, - "37": 7.97726, - "38": 7.86783, - "39": 7.89458, - "40": 7.92858, - "41": 7.9655, - "42": 8.11402, - "43": 7.92667, - "44": 8.10251, - "45": 7.84423, - "46": 8.02262, - "47": 7.90143, - "48": 8.11201, - "49": 8.26159, - "50": 8.02742 + "1": 112.4651, + "2": 8.9818, + "3": 7.9959, + "4": 7.79497, + "5": 7.75717, + "6": 7.71573, + "7": 7.88261, + "8": 8.18461, + "9": 7.98469, + "10": 8.03914, + "11": 7.78529, + "12": 7.60637, + "13": 8.00128, + "14": 7.70029, + "15": 7.55303, + "16": 7.64575, + "17": 7.5524, + "18": 7.63808, + "19": 7.63046, + "20": 7.67546, + "21": 7.59065, + "22": 7.8971, + "23": 7.67556, + "24": 7.60639, + "25": 8.10731, + "26": 7.838, + "27": 7.88628, + "28": 7.76089, + "29": 7.97056, + "30": 7.69011, + "31": 7.51965, + "32": 7.68428, + "33": 7.63171, + "34": 7.67041, + "35": 9.81602, + "36": 7.77524, + "37": 8.84437, + "38": 7.68084, + "39": 7.73586, + "40": 7.59679, + "41": 7.69328, + "42": 7.62829, + "43": 7.83834, + "44": 7.89707, + "45": 7.91019, + "46": 7.95723, + "47": 7.66212, + "48": 7.85804, + "49": 7.73712, + "50": 7.54183 } } } diff --git a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json index 3fee233b8c1..03e8f2c649a 100644 --- a/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/gpt/gpt_grpo_tp1tp2_pp1_dp8_583m_throughputtest/golden_values_dev_dgx_h100.json @@ -4,29 +4,29 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 0.0, + "1": -0.05654, "2": 0.0, - "3": 0.0, + "3": 0.05082, "4": 0.0, "5": 0.0, - "6": 0.03951, - "7": 0.0, + "6": 0.02971, + "7": 0.0929, "8": 0.0, "9": 0.0, "10": 0.0, - "11": -0.02226, + "11": 0.0, "12": 0.0, "13": 0.0, "14": 0.0, - "15": 0.0, - "16": 0.0, + "15": 0.06444, + "16": 0.05003, "17": 0.0, "18": 0.0, "19": 0.0, - "20": 0.0, + "20": 0.04942, "21": 0.0, "22": 0.0, - "23": 0.04844, + "23": 0.0, "24": 0.0, "25": 0.0, "26": 0.0, @@ -42,15 +42,15 @@ "36": 0.0, "37": 0.0, "38": 0.0, - "39": 0.0, + "39": 0.02196, "40": 0.0, - "41": 0.05432, - "42": 0.03022, - "43": 0.05067, + "41": 0.0, + "42": 0.0, + "43": 0.0, "44": 0.0, - "45": 0.05642, - "46": 0.02496, - "47": 0.0377, + "45": 0.0, + "46": 0.0, + "47": 0.0, "48": 0.0, "49": 0.0, "50": 0.0 @@ -61,29 +61,29 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 583687296.0, + "1": 37.0, "2": 583687296.0, - "3": 583687296.0, + "3": 29.0, "4": 583687296.0, "5": 583687296.0, - "6": 39.0, - "7": 583687296.0, + "6": 56.0, + "7": 30.0, "8": 583687296.0, "9": 583687296.0, "10": 583687296.0, - "11": 50.0, + "11": 583687296.0, "12": 583687296.0, "13": 583687296.0, "14": 583687296.0, - "15": 583687296.0, - "16": 583687296.0, + "15": 19.0, + "16": 41.0, "17": 583687296.0, "18": 583687296.0, "19": 583687296.0, - "20": 583687296.0, + "20": 31.0, "21": 583687296.0, "22": 583687296.0, - "23": 22.0, + "23": 583687296.0, "24": 583687296.0, "25": 583687296.0, "26": 583687296.0, @@ -99,15 +99,15 @@ "36": 583687296.0, "37": 583687296.0, "38": 583687296.0, - "39": 583687296.0, + "39": 40.0, "40": 583687296.0, - "41": 7.0, - "42": 37.0, - "43": 27.0, + "41": 583687296.0, + "42": 583687296.0, + "43": 583687296.0, "44": 583687296.0, - "45": 18.0, - "46": 22.0, - "47": 70.0, + "45": 583687296.0, + "46": 583687296.0, + "47": 583687296.0, "48": 583687296.0, "49": 583687296.0, "50": 583687296.0 @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 57332711424.0, - "2": 57337454592.0, - "3": 57337208832.0, - "4": 57338224640.0, - "5": 57338208256.0, - "6": 57338191872.0, - "7": 57338175488.0, - "8": 57338163200.0, - "9": 57338003456.0, - "10": 57338003456.0, - "11": 57338007552.0, - "12": 57338011648.0, - "13": 57338019840.0, - "14": 57338015744.0, - "15": 57338019840.0, - "16": 57338023936.0, - "17": 57338019840.0, - "18": 57338015744.0, - "19": 57338015744.0, - "20": 57338003456.0, - "21": 57338007552.0, - "22": 57338007552.0, - "23": 57338011648.0, - "24": 57338015744.0, - "25": 57338015744.0, - "26": 57338011648.0, - "27": 57338015744.0, - "28": 57338011648.0, - "29": 57338019840.0, - "30": 57338015744.0, - "31": 57338011648.0, - "32": 57338003456.0, - "33": 57338011648.0, - "34": 57338007552.0, - "35": 57338011648.0, - "36": 57338011648.0, - "37": 57338015744.0, - "38": 57338015744.0, - "39": 57338011648.0, - "40": 57338011648.0, - "41": 57338015744.0, - "42": 57338023936.0, - "43": 57338019840.0, - "44": 57338015744.0, - "45": 57338011648.0, - "46": 57338011648.0, - "47": 57338011648.0, - "48": 57338015744.0, - "49": 57338011648.0, - "50": 57338011648.0 + "1": 57332219904.0, + "2": 57336553472.0, + "3": 57337122816.0, + "4": 57336414208.0, + "5": 57337241600.0, + "6": 57337225216.0, + "7": 57337208832.0, + "8": 57337208832.0, + "9": 57337192448.0, + "10": 57337020416.0, + "11": 57336926208.0, + "12": 57336930304.0, + "13": 57336934400.0, + "14": 57336938496.0, + "15": 57336938496.0, + "16": 57336942592.0, + "17": 57336946688.0, + "18": 57336942592.0, + "19": 57336946688.0, + "20": 57336946688.0, + "21": 57336942592.0, + "22": 57336938496.0, + "23": 57336938496.0, + "24": 57336942592.0, + "25": 57336946688.0, + "26": 57336946688.0, + "27": 57336942592.0, + "28": 57336942592.0, + "29": 57336938496.0, + "30": 57336938496.0, + "31": 57336938496.0, + "32": 57336938496.0, + "33": 57336942592.0, + "34": 57336946688.0, + "35": 57336946688.0, + "36": 57336946688.0, + "37": 57336950784.0, + "38": 57336954880.0, + "39": 57336950784.0, + "40": 57336946688.0, + "41": 57336946688.0, + "42": 57336938496.0, + "43": 57336942592.0, + "44": 57336946688.0, + "45": 57336946688.0, + "46": 57336950784.0, + "47": 57336950784.0, + "48": 57336954880.0, + "49": 57336954880.0, + "50": 57336954880.0 } }, "mem-max-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 57332711424.0, - "2": 59426697216.0, - "3": 59427717120.0, - "4": 59429486592.0, - "5": 59429486592.0, - "6": 59429486592.0, - "7": 59429486592.0, - "8": 59429486592.0, - "9": 59429920768.0, - "10": 59429920768.0, - "11": 59429920768.0, - "12": 59429920768.0, - "13": 59429937152.0, - "14": 59429937152.0, - "15": 59429937152.0, - "16": 59429941248.0, - "17": 59429941248.0, - "18": 59429941248.0, - "19": 59429941248.0, - "20": 59429941248.0, - "21": 59429941248.0, - "22": 59429941248.0, - "23": 59429941248.0, - "24": 59429941248.0, - "25": 59429941248.0, - "26": 59429941248.0, - "27": 59429941248.0, - "28": 59429941248.0, - "29": 59429941248.0, - "30": 59429941248.0, - "31": 59429941248.0, - "32": 59429941248.0, - "33": 59429941248.0, - "34": 59429941248.0, - "35": 59429941248.0, - "36": 59429941248.0, - "37": 59429941248.0, - "38": 59429941248.0, - "39": 59429941248.0, - "40": 59429941248.0, - "41": 59429941248.0, - "42": 59429941248.0, - "43": 59429941248.0, - "44": 59429941248.0, - "45": 59429941248.0, - "46": 59429941248.0, - "47": 59429941248.0, - "48": 59429941248.0, - "49": 59429941248.0, - "50": 59429941248.0 + "1": 57332224000.0, + "2": 59149262848.0, + "3": 59150077952.0, + "4": 59150077952.0, + "5": 59150077952.0, + "6": 59150077952.0, + "7": 59151114240.0, + "8": 59151114240.0, + "9": 59151114240.0, + "10": 59151114240.0, + "11": 59151114240.0, + "12": 59151114240.0, + "13": 59151114240.0, + "14": 59151114240.0, + "15": 59151114240.0, + "16": 59151114240.0, + "17": 59151114240.0, + "18": 59151114240.0, + "19": 59151114240.0, + "20": 59151114240.0, + "21": 59151114240.0, + "22": 59151114240.0, + "23": 59151114240.0, + "24": 59151114240.0, + "25": 59151114240.0, + "26": 59151114240.0, + "27": 59151114240.0, + "28": 59151114240.0, + "29": 59151114240.0, + "30": 59151114240.0, + "31": 59151114240.0, + "32": 59151114240.0, + "33": 59151114240.0, + "34": 59151114240.0, + "35": 59151114240.0, + "36": 59151114240.0, + "37": 59151114240.0, + "38": 59151114240.0, + "39": 59151114240.0, + "40": 59151114240.0, + "41": 59151114240.0, + "42": 59151114240.0, + "43": 59151114240.0, + "44": 59151114240.0, + "45": 59151114240.0, + "46": 59151114240.0, + "47": 59151114240.0, + "48": 59151114240.0, + "49": 59151114240.0, + "50": 59151114240.0 } }, "iteration-time": { @@ -232,56 +232,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 83.94299, - "2": 7.51381, - "3": 5.78991, - "4": 5.04681, - "5": 5.19935, - "6": 5.86003, - "7": 5.37343, - "8": 6.22227, - "9": 5.79141, - "10": 5.00793, - "11": 5.08468, - "12": 4.91629, - "13": 5.00142, - "14": 5.12995, - "15": 4.99047, - "16": 5.27209, - "17": 4.97418, - "18": 4.89528, - "19": 4.89114, - "20": 5.02499, - "21": 4.90696, - "22": 4.83721, - "23": 4.88444, - "24": 4.87204, - "25": 4.94429, - "26": 5.167, - "27": 4.89225, - "28": 4.99395, - "29": 4.90509, - "30": 4.95357, - "31": 4.9475, - "32": 5.37096, - "33": 4.85755, - "34": 4.95423, - "35": 4.88925, - "36": 4.92784, - "37": 4.86931, - "38": 5.23366, - "39": 4.96287, - "40": 4.88792, - "41": 4.93789, - "42": 4.91727, - "43": 4.90454, - "44": 5.39987, - "45": 4.97505, - "46": 4.83257, - "47": 4.9187, - "48": 4.87401, - "49": 4.90853, - "50": 5.384 + "1": 52.61195, + "2": 7.49893, + "3": 5.44917, + "4": 5.25399, + "5": 5.19843, + "6": 5.38782, + "7": 5.46312, + "8": 5.51632, + "9": 5.4408, + "10": 5.49874, + "11": 5.36303, + "12": 5.69186, + "13": 5.34955, + "14": 5.21328, + "15": 5.48714, + "16": 5.15125, + "17": 5.23431, + "18": 5.41636, + "19": 5.23728, + "20": 5.29286, + "21": 5.38276, + "22": 5.3511, + "23": 5.44691, + "24": 5.92416, + "25": 5.45103, + "26": 5.46551, + "27": 6.16072, + "28": 5.39877, + "29": 5.33127, + "30": 5.16681, + "31": 5.2505, + "32": 5.31096, + "33": 5.22232, + "34": 5.2931, + "35": 5.17914, + "36": 5.1828, + "37": 5.1557, + "38": 5.23643, + "39": 5.14327, + "40": 5.50636, + "41": 5.25893, + "42": 5.26202, + "43": 5.21742, + "44": 5.70532, + "45": 5.2714, + "46": 5.42226, + "47": 5.26848, + "48": 5.1815, + "49": 5.10326, + "50": 5.20865 } } -} \ No newline at end of file +} From c0882a41783ac3def7346823648d1127a8b542b6 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Mon, 12 Jan 2026 13:33:16 -0800 Subject: [PATCH 13/18] If empty bin, give it a default PackedSeqParams so the signature will match for cudagraphs --- train_rl.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/train_rl.py b/train_rl.py index d767e30401b..a76af7c65bf 100644 --- a/train_rl.py +++ b/train_rl.py @@ -25,6 +25,8 @@ from megatron.training.arguments import core_transformer_config_from_args from model_provider import model_provider +from megatron.rl.sequence_packing_utils import get_default_packed_seq_params + stimer = StragglerDetector() import logging @@ -255,6 +257,13 @@ def forward_step(data_iterator, model: GPTModel, loss_only: bool = False): # Common logic for both paths model_to_use = model[0] if isinstance(model, list) else model + if packed_seq_params is None: + print(f"WHAT DOES THIS INPUT DATA EVEN LOOK LIKE? tokens {tokens} | position_ids: {position_ids}") + packed_seq_params = get_default_packed_seq_params( + seq_length=tokens.shape[1], + device=tokens.device, + ) + # Clear RoPE cache to avoid inference tensor errors try: for module in model_to_use.modules(): From 3d7144206268b4a7bf2c97fc355711c55af5b23a Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 13 Jan 2026 10:56:31 -0800 Subject: [PATCH 14/18] runs fast need refactor --- megatron/rl/sequence_packing_utils.py | 8 ++++++-- train_rl.py | 2 ++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py index ddff9555ccf..17db4ccafc3 100644 --- a/megatron/rl/sequence_packing_utils.py +++ b/megatron/rl/sequence_packing_utils.py @@ -412,8 +412,12 @@ def get_default_packed_seq_params(seq_length: int, device: torch.device) -> Pack Returns: PackedSeqParams configured as a single unpacked sequence. """ - # Single sequence spanning the full length = no actual packing - cu_seqlens = torch.full((seq_length,), seq_length, dtype=torch.int32, device=device) + # print("CREATING FAKE SEQPACK") + # Pad with all zero-length sequences indicates to the attention kernel that + # cu_seqlens = torch.full((2,), seq_length, dtype=torch.int32, device=device) + # cu_seqlens[0] = 0 + # cu_seqlens = torch.tensor([0, seq_length], dtype=torch.int32, device=device) + cu_seqlens = torch.full((50,), seq_length, dtype=torch.int32, device=device) cu_seqlens[0] = 0 return PackedSeqParams( diff --git a/train_rl.py b/train_rl.py index a76af7c65bf..b05bfe122f7 100644 --- a/train_rl.py +++ b/train_rl.py @@ -257,6 +257,7 @@ def forward_step(data_iterator, model: GPTModel, loss_only: bool = False): # Common logic for both paths model_to_use = model[0] if isinstance(model, list) else model + if packed_seq_params is None: print(f"WHAT DOES THIS INPUT DATA EVEN LOOK LIKE? tokens {tokens} | position_ids: {position_ids}") packed_seq_params = get_default_packed_seq_params( @@ -264,6 +265,7 @@ def forward_step(data_iterator, model: GPTModel, loss_only: bool = False): device=tokens.device, ) + # Clear RoPE cache to avoid inference tensor errors try: for module in model_to_use.modules(): From 67b9f3ba9f1bd465068452638323dde943a8ac70 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 13 Jan 2026 11:11:24 -0800 Subject: [PATCH 15/18] Cleanup --- megatron/rl/sequence_packing_utils.py | 11 +++++------ train_rl.py | 3 --- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py index 17db4ccafc3..4599d3a0934 100644 --- a/megatron/rl/sequence_packing_utils.py +++ b/megatron/rl/sequence_packing_utils.py @@ -412,12 +412,11 @@ def get_default_packed_seq_params(seq_length: int, device: torch.device) -> Pack Returns: PackedSeqParams configured as a single unpacked sequence. """ - # print("CREATING FAKE SEQPACK") - # Pad with all zero-length sequences indicates to the attention kernel that - # cu_seqlens = torch.full((2,), seq_length, dtype=torch.int32, device=device) - # cu_seqlens[0] = 0 - # cu_seqlens = torch.tensor([0, seq_length], dtype=torch.int32, device=device) - cu_seqlens = torch.full((50,), seq_length, dtype=torch.int32, device=device) + + args = get_args() + + # Pad to the maximum number of sequences in the bin for the attention kernel. + cu_seqlens = torch.full((args.rl_sequence_packing_max_sequences_per_bin,), seq_length, dtype=torch.int32, device=device) cu_seqlens[0] = 0 return PackedSeqParams( diff --git a/train_rl.py b/train_rl.py index b05bfe122f7..299843bcff3 100644 --- a/train_rl.py +++ b/train_rl.py @@ -257,15 +257,12 @@ def forward_step(data_iterator, model: GPTModel, loss_only: bool = False): # Common logic for both paths model_to_use = model[0] if isinstance(model, list) else model - if packed_seq_params is None: - print(f"WHAT DOES THIS INPUT DATA EVEN LOOK LIKE? tokens {tokens} | position_ids: {position_ids}") packed_seq_params = get_default_packed_seq_params( seq_length=tokens.shape[1], device=tokens.device, ) - # Clear RoPE cache to avoid inference tensor errors try: for module in model_to_use.modules(): From db1bae47116c91ae10d7a94f3ee32fe402538f25 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Tue, 13 Jan 2026 14:04:09 -0800 Subject: [PATCH 16/18] Maximize --rl-sequence-packing-max-sequences-per-bin for now --- megatron/training/arguments.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index ce9f155d68a..8df53e61863 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2094,7 +2094,7 @@ def _add_rl_args(parser): help='If set, calculate the intra-group similarity of rollouts.') group.add_argument('--rl-use-sequence-packing', action=argparse.BooleanOptionalAction, type=bool, default=False, help='Enable sequence packing') - group.add_argument('--rl-sequence-packing-max-sequences-per-bin', type=int, default=32, + group.add_argument('--rl-sequence-packing-max-sequences-per-bin', type=int, default=50, help='Maximum number of sequences that can be packed into a single bin. ') group.add_argument('--rl-sequence-packing-algo', type=str, default='fifo', choices=['fifo', 'round-robin'], From 1c488ffec51e22281b2554320925e0129bee64e1 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Wed, 14 Jan 2026 07:35:06 -0800 Subject: [PATCH 17/18] Scope the self.training check so logprobs will not trigger cudagraph capture --- megatron/core/transformer/cuda_graphs.py | 2 +- megatron/rl/rl_utils.py | 10 +++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index edc1161d7bc..82f319e4ccb 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -1372,7 +1372,7 @@ def __call__(self, megatron_module, args, kwargs): # Now replay the graph out = runner.replay_graph_capture(self.is_first_microbatch, args, kwargs) - elif self.training and torch.is_grad_enabled(): + elif self.training: # Training mode runner = self.get_cudagraph_runner(megatron_module, args, kwargs) # check if a layer is frozen during training. diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index e48ed9c0a8a..20c6c9eeaff 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -1102,6 +1102,11 @@ def prepare_data_for_update( ) def logprobs_forward_step(data_iterator, model): + + # Avoid self.training checks which will trigger cudagraph capture; this path reuses + # the forward pass from training after it has been captured on the 1st iteration. + model.eval() + if args.rl_use_sequence_packing: # When using sequence packing, the data iterator returns a tuple with a single element, the bin index. bin_tensor = next(data_iterator)[0] @@ -1117,7 +1122,7 @@ def logprobs_forward_step(data_iterator, model): b_trajs = b_trajs.cuda() b_posids = b_posids.cuda() - return ( + logprobs = ( get_logprobs( model, b_trajs, @@ -1129,6 +1134,9 @@ def logprobs_forward_step(data_iterator, model): None, ) + model.train() + return logprobs + dtype = ( torch.bfloat16 if args.bf16 else (torch.float16 if args.fp16 else torch.float32) ) From d37d0ce8e6f0485f2e1001ceeab0dc62563bc9e6 Mon Sep 17 00:00:00 2001 From: Helen Ngo Date: Wed, 14 Jan 2026 09:03:00 -0800 Subject: [PATCH 18/18] Add an error if we ever pass a set of tensors through _clone_nested_tensors --- megatron/core/transformer/cuda_graphs.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 82f319e4ccb..b732aba6fc1 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -185,6 +185,10 @@ def _clone_nested_tensors(value: Any) -> Any: return type(value)(_clone_nested_tensors(v) for v in value) if isinstance(value, dict): return {k: _clone_nested_tensors(v) for k, v in value.items()} + if isinstance(value, set): + raise TypeError( + "Sets of tensors are unsupported in cudagraph helpers; use list/tuple instead" + ) return value