Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / compile_adapter.py: 32%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Adapters for using PyTorch-native selective checkpointing under compile.""" 

16 

17__all__ = ["create_native_selective_checkpoint_contexts"] 

18 

19from functools import partial 

20from typing import Any, Callable, Tuple 

21 

22from hyper_parallel.core.activation_memory.policy import CheckpointPolicy 

23 

24 

25_SUPPORTED_POLICY_NAMES = ( 

26 "MUST_SAVE", 

27 "PREFER_SAVE", 

28 "MUST_RECOMPUTE", 

29 "PREFER_RECOMPUTE", 

30) 

31 

32 

33def _to_torch_checkpoint_policy(policy: Any) -> Any: 

34 """Convert supported HyperParallel policies to the native Torch enum.""" 

35 from torch.utils import checkpoint as torch_checkpoint # pylint: disable=C0415 

36 

37 torch_policy_cls = torch_checkpoint.CheckpointPolicy 

38 supported_native = { 

39 getattr(torch_policy_cls, name) for name in _SUPPORTED_POLICY_NAMES 

40 } 

41 if isinstance(policy, torch_policy_cls): 

42 if policy in supported_native: 

43 return policy 

44 raise ValueError( 

45 f"Torch checkpoint policy {policy.name} is not supported by " 

46 "HyperParallel compile mode. Only SAVE and RECOMPUTE policies are supported." 

47 ) 

48 if isinstance(policy, CheckpointPolicy): 

49 if policy.name in _SUPPORTED_POLICY_NAMES: 

50 return getattr(torch_policy_cls, policy.name) 

51 raise ValueError( 

52 f"HyperParallel checkpoint policy {policy.name} is not supported in compile mode. " 

53 "Only SAVE and RECOMPUTE policies are supported." 

54 ) 

55 raise TypeError( 

56 "Selective checkpoint policy_fn must return a HyperParallel or Torch " 

57 f"CheckpointPolicy, but got {type(policy).__name__}." 

58 ) 

59 

60 

61def _torch_policy_adapter( 

62 policy_fn: Callable, torch_context: Any, op: Any, *args: Any, **kwargs: Any 

63) -> Any: 

64 """Pass native Torch inputs through and adapt only the policy result.""" 

65 return _to_torch_checkpoint_policy(policy_fn(torch_context, op, *args, **kwargs)) 

66 

67 

68def create_native_selective_checkpoint_contexts(policy_fn: Callable) -> Tuple[Any, Any]: 

69 """Create Torch-native selective-checkpoint contexts for compile capture.""" 

70 if not callable(policy_fn): 

71 raise TypeError("policy_fn must be callable in HyperParallel compile mode.") 

72 from torch.utils import checkpoint as torch_checkpoint # pylint: disable=C0415 

73 

74 return torch_checkpoint.create_selective_checkpoint_contexts( 

75 partial(_torch_policy_adapter, policy_fn) 

76 )