diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 10fb335addd..4d4d76186f3 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -279,6 +279,7 @@ def __init__( cuda_graph_mixed_prefill_count: Optional[int] = 16, metrics_writer: Optional['WandbModule'] = None, request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None, + persist_cuda_graphs: Optional[bool] = False, ): super().__init__(materialize_only_last_token_logits=materialize_only_last_token_logits) @@ -400,6 +401,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 c7698b8a4bb..cede0ca06dc 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -184,6 +184,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" @@ -566,10 +567,10 @@ 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 when not using unified memory at all (level 0) and + # `--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() # Maintain references to requests before reset. @@ -611,7 +612,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/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index a1cb764b01d..b732aba6fc1 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/megatron/core/transformer/cuda_graphs.py @@ -101,6 +101,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 @@ -176,6 +177,44 @@ 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()} + if isinstance(value, set): + raise TypeError( + "Sets of tensors are unsupported in cudagraph helpers; use list/tuple instead" + ) + 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 @@ -683,8 +722,12 @@ 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. - for _, state in get_all_rng_states().items(): - self.fwd_graph.register_generator_state(state) + 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): @@ -706,6 +749,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" ): diff --git a/megatron/rl/inference/megatron.py b/megatron/rl/inference/megatron.py index 126da94e3a6..73ab5024a64 100644 --- a/megatron/rl/inference/megatron.py +++ b/megatron/rl/inference/megatron.py @@ -166,6 +166,7 @@ def get_dynamic_inference_engine( cuda_graph_max_tokens=args.inference_dynamic_batching_cuda_graph_max_tokens, cuda_graph_mixed_prefill_count=args.inference_dynamic_batching_cuda_graph_mixed_prefill_count, metrics_writer=metrics_writer, + persist_cuda_graphs=args.rl_training_cuda_graphs ) inference_wrapped_model = GPTInferenceWrapper(model, args, inference_context, pg_collection=pg_collection) diff --git a/megatron/rl/rl_utils.py b/megatron/rl/rl_utils.py index cdb8049a100..20c6c9eeaff 100644 --- a/megatron/rl/rl_utils.py +++ b/megatron/rl/rl_utils.py @@ -987,13 +987,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 ) @@ -1108,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] @@ -1123,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, @@ -1135,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) ) @@ -1600,7 +1602,7 @@ def megatron_rl_inference_mode( 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": + 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) @@ -1649,7 +1651,7 @@ def megatron_rl_inference_mode( 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": + 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 this is a separate RL inference model, prefetch weights back to CPU so they don't consume diff --git a/megatron/rl/sequence_packing_utils.py b/megatron/rl/sequence_packing_utils.py index ddff9555ccf..4599d3a0934 100644 --- a/megatron/rl/sequence_packing_utils.py +++ b/megatron/rl/sequence_packing_utils.py @@ -412,8 +412,11 @@ 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) + + 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/megatron/training/arguments.py b/megatron/training/arguments.py index 262c4ce79b2..63637a274da 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2127,13 +2127,17 @@ 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'], 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-training-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, + default=False, + 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.') group.add_argument('--rl-inference-tensor-model-parallel-size', type=int, default=None, help='Degree of tensor model parallelism for inference for RL.') group.add_argument( diff --git a/train_rl.py b/train_rl.py index d767e30401b..299843bcff3 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,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: + 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():