Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / activation_checkpoint / native_compile.py: 0%

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 

17from functools import partial 

18from typing import Any, Callable, Tuple 

19 

20from hyper_parallel.core.activation_checkpoint.activation_checkpoint import CheckpointPolicy 

21 

22 

23_SUPPORTED_POLICY_NAMES = ( 

24 "MUST_SAVE", 

25 "PREFER_SAVE", 

26 "MUST_RECOMPUTE", 

27 "PREFER_RECOMPUTE", 

28) 

29 

30 

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

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

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

34 

35 torch_policy_cls = torch_checkpoint.CheckpointPolicy 

36 supported_native = { 

37 getattr(torch_policy_cls, name) for name in _SUPPORTED_POLICY_NAMES 

38 } 

39 if isinstance(policy, torch_policy_cls): 

40 if policy in supported_native: 

41 return policy 

42 raise ValueError( 

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

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

45 ) 

46 if isinstance(policy, CheckpointPolicy): 

47 if policy.name in _SUPPORTED_POLICY_NAMES: 

48 return getattr(torch_policy_cls, policy.name) 

49 raise ValueError( 

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

51 "Only SAVE and RECOMPUTE policies are supported." 

52 ) 

53 raise TypeError( 

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

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

56 ) 

57 

58 

59def _torch_policy_adapter( 

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

61) -> Any: 

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

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

64 

65 

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

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

68 if not callable(policy_fn): 

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

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

71 

72 return torch_checkpoint.create_selective_checkpoint_contexts( 

73 partial(_torch_policy_adapter, policy_fn) 

74 ) 

75 

76 

77__all__ = ["create_native_selective_checkpoint_contexts"]