diff --git a/src/prime_rl/trainer/compile_compat.py b/src/prime_rl/trainer/compile_compat.py new file mode 100644 index 00000000000..282b72eafc3 --- /dev/null +++ b/src/prime_rl/trainer/compile_compat.py @@ -0,0 +1,47 @@ +"""PyTorch 2.13 compatibility for runtime autotuning under checkpointing. + +Modes opting out of compiler internals must not observe runtime +autotuning allocations or benchmark executions during checkpoint recomputation. +""" + +from contextlib import contextmanager + +import torch +from torch.utils._python_dispatch import _get_current_dispatch_mode_stack, _pop_mode, _push_mode + + +@contextmanager +def _disable_modes_for_compile_internals(): + modes = _get_current_dispatch_mode_stack() + retained = [mode for mode in modes if not mode.ignore_compile_internals()] + if len(retained) == len(modes): + yield + return + + for _ in modes: + _pop_mode() + for mode in retained: + _push_mode(mode) + try: + yield + finally: + for _ in retained: + _pop_mode() + for mode in modes: + _push_mode(mode) + + +def apply_autotune_checkpoint_patch() -> None: + """Apply the tested runtime backport once; refuse unvalidated Torch versions.""" + if torch.__version__.split("+")[0] != "2.13.0": + raise RuntimeError(f"Autotuning checkpoint backport requires Torch 2.13.0; got {torch.__version__}") + + from torch._inductor.runtime.triton_heuristics import CachingAutotuner + + for name in ("bench", "autotune_to_one_config", "coordinate_descent_tuning"): + original = getattr(CachingAutotuner, name) + if getattr(original, "_prime_rl_checkpoint_backport", False): + continue + wrapped = _disable_modes_for_compile_internals()(original) + wrapped._prime_rl_checkpoint_backport = True + setattr(CachingAutotuner, name, wrapped) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 1e37dbcc8ce..6197a0b1d67 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -16,7 +16,6 @@ from huggingface_hub import snapshot_download from jaxtyping import Int from torch import Tensor -from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageReader from torch.distributed.checkpoint.state_dict_loader import load as dcp_load from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, OffloadPolicy, fully_shard @@ -842,18 +841,17 @@ def mark_dynamic_int_args() -> None: def apply_compile(model: nn.Module, compile_config: CompileConfig): + from prime_rl.trainer.compile_compat import apply_autotune_checkpoint_patch + + apply_autotune_checkpoint_patch() + get_logger().info(f"Autotuning checkpoint dispatch backport active (torch={torch.__version__})") torch._dynamo.config.capture_scalar_outputs = True mark_dynamic_int_args() language_model = get_language_model(model) for layer_id in range(len(language_model.layers)): layer = language_model.layers[layer_id] - if isinstance(layer, CheckpointWrapper) and any( - isinstance(module, FSDPModule) for module in layer._checkpoint_wrapped_module.modules() - ): - # A nested FSDP unit's hooks always break the graph, and a break inside a compiled checkpoint - # sends the block to eager, so keep AC eager around the compiled block. pytorch/pytorch#196626 - # removes the fp32 router's nested unit by letting it join the block's FSDP unit. - layer = layer._checkpoint_wrapped_module + # Compile the checkpoint wrapper as one region: an eager fallback during + # recomputation must not change the operations observed by an eager SAC context. # Doing it in-place avoids mangled fqn which can break checkpoint loading layer.compile(fullgraph=compile_config.fullgraph, mode=compile_config.mode) get_logger().info( @@ -1016,7 +1014,7 @@ def setup_model( override_attr=config.vlm.vision_encoder_attr if config.vlm is not None else None, ) - # the right order is AC -> FSDP -> Compile: compile needs to see which blocks hold nested FSDP units + # Keep the checkpoint wrapper inside the compiled region and apply FSDP before tracing. if config.ac is not None: apply_ac(model, config.ac)