Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / tracer / dynamic_shapes.py: 96%

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"""Selective symbolic inputs and runtime guards for make_fx joint graphs.""" 

16 

17from __future__ import annotations 

18 

19from typing import Any 

20 

21import torch 

22from torch._dynamo.source import LocalSource 

23from torch._subclasses import FakeTensorMode 

24from torch.fx.experimental.symbolic_shapes import DimDynamic, ShapeEnv, StatelessSymbolicContext 

25from torch.utils import _pytree as pytree 

26 

27 

28def normalize_dynamic_arg_dims(dynamic_arg_dims: dict[str, int | list[int]] | None) -> dict[str, list[int]] | None: 

29 """Validate and copy the public path-to-dimensions mapping.""" 

30 if dynamic_arg_dims is None: 

31 return None 

32 if not isinstance(dynamic_arg_dims, dict): 

33 raise ValueError("dynamic_arg_dims must be a mapping of input paths to dimensions") 

34 normalized = {} 

35 for path, dims in dynamic_arg_dims.items(): 

36 if not isinstance(path, str) or not path or any(not part for part in path.split(".")): 

37 raise ValueError("dynamic_arg_dims keys must be non-empty dotted input paths") 

38 dims = [dims] if isinstance(dims, int) else dims 

39 if not isinstance(dims, (list, tuple)) or any( 

40 isinstance(dim, bool) or not isinstance(dim, int) for dim in dims 

41 ): 

42 raise ValueError(f"dynamic_arg_dims[{path!r}] must be an int or a list of ints") 

43 normalized[path] = list(dims) 

44 return normalized 

45 

46 

47def _resolve_input(inputs: dict[str, Any], path: str) -> torch.Tensor: 

48 """Follow a dotted path through dicts, sequences and attributes to a tensor.""" 

49 value = inputs 

50 try: 

51 for part in path.split("."): 

52 if isinstance(value, dict): 

53 value = value[part] 

54 elif isinstance(value, (list, tuple)): 

55 index = int(part) 

56 if index < 0: 

57 raise IndexError(index) 

58 value = value[index] 

59 else: 

60 value = getattr(value, part) 

61 except (KeyError, IndexError, ValueError, AttributeError) as exc: 

62 raise ValueError(f"dynamic_arg_dims input path {path!r} does not exist") from exc 

63 if not isinstance(value, torch.Tensor): 

64 raise ValueError(f"dynamic_arg_dims input {path!r} must be a tensor, got {type(value).__name__}") 

65 return value 

66 

67 

68def build_symbolic_inputs( 

69 state_flat: list[Any], 

70 inputs: dict[str, Any], 

71 dynamic_arg_dims: dict[str, int | list[int]] | None, 

72) -> tuple[FakeTensorMode, tuple[Any, ...]]: 

73 """Fakeify state statically and selected user dimensions symbolically. 

74 

75 Args: 

76 state_flat: Flat model parameters and buffers. 

77 inputs: Keyword inputs, including registered pytree containers. 

78 dynamic_arg_dims: Explicit paths/dimensions; None symbolizes all user tensor dimensions. 

79 

80 Returns: 

81 Fake mode and flattened state/user inputs sharing its ShapeEnv. 

82 """ 

83 mapping = normalize_dynamic_arg_dims(dynamic_arg_dims) 

84 user_flat, _ = pytree.tree_flatten(inputs) 

85 tensor_ids = {id(value) for value in user_flat if isinstance(value, torch.Tensor)} 

86 dynamic_dims = {} 

87 if mapping is None: 

88 dynamic_dims = {id(value): set(range(value.ndim)) for value in user_flat if isinstance(value, torch.Tensor)} 

89 else: 

90 for path, dims in mapping.items(): 

91 tensor = _resolve_input(inputs, path) 

92 if id(tensor) not in tensor_ids: 

93 raise ValueError(f"dynamic_arg_dims input {path!r} must be a registered pytree tensor leaf") 

94 selected = dynamic_dims.setdefault(id(tensor), set()) 

95 for dim in dims: 

96 if not -tensor.ndim <= dim < tensor.ndim: 

97 raise ValueError(f"dynamic_arg_dims[{path!r}] dimension {dim} is invalid for rank {tensor.ndim}") 

98 selected.add(dim % tensor.ndim) 

99 

100 # Equal sample sizes do not imply equality across future batches. 

101 fake_mode = FakeTensorMode(allow_non_fake_inputs=True, shape_env=ShapeEnv(duck_shape=False)) 

102 fake_args = [fake_mode.from_tensor(value, static_shapes=True) for value in state_flat] 

103 for index, value in enumerate(user_flat): 

104 if not isinstance(value, torch.Tensor): 

105 fake_args.append(value) 

106 continue 

107 if value.layout != torch.strided: 

108 raise ValueError("Dynamic graph inputs must use torch.strided layout") 

109 dims = dynamic_dims.get(id(value), set()) 

110 context = StatelessSymbolicContext( 

111 dynamic_sizes=[DimDynamic.DYNAMIC if dim in dims else DimDynamic.STATIC for dim in range(value.ndim)] 

112 ) 

113 fake_args.append(fake_mode.from_tensor(value, source=LocalSource(f"input_{index}"), symbolic_context=context)) 

114 return fake_mode, tuple(fake_args) 

115 

116 

117class InputGuards: 

118 """Check input metadata, constants and traced shape branches before execution. 

119 

120 The graph is called directly, so Dynamo will not check its guards for us. 

121 Only user tensors participate: FSDP may legitimately shard live model state. 

122 No real input tensors or storage are retained here. 

123 """ 

124 

125 def __init__(self, inputs: list[Any], fake_inputs: tuple[Any, ...], shape_env: ShapeEnv) -> None: 

126 self.shape_env = shape_env 

127 self.fake_tensors = [value for value in fake_inputs if isinstance(value, torch.Tensor)] 

128 self.leaves = [self._metadata(value) for value in inputs] 

129 self.aliases = self._aliases(inputs) 

130 self.expression = None 

131 self.refresh() 

132 

133 @staticmethod 

134 def _metadata(value: Any) -> tuple: 

135 if isinstance(value, torch.Tensor): 

136 return (torch.Tensor, value.ndim, value.dtype, value.device, value.layout, value.requires_grad) 

137 return (type(value), value) 

138 

139 @staticmethod 

140 def _aliases(inputs: list[Any]) -> list[int]: 

141 seen = {} 

142 return [ 

143 seen.setdefault(id(value), index) 

144 for index, value in enumerate(inputs) if isinstance(value, torch.Tensor) 

145 ] 

146 

147 def refresh(self) -> None: 

148 """Include constraints introduced by tracing and graph passes.""" 

149 self.expression = self.shape_env.produce_guards_expression(self.fake_tensors, ignore_static=False) 

150 

151 def validate(self, inputs: list[Any]) -> None: 

152 """Reject incompatible inputs before any graph or collective is executed.""" 

153 if [self._metadata(value) for value in inputs] != self.leaves: 

154 raise ValueError("Graph input rank, dtype, device, layout, requires_grad or Python constant changed") 

155 if self._aliases(inputs) != self.aliases: 

156 raise ValueError("Graph input tensor aliasing changed since compilation") 

157 tensors = [value for value in inputs if isinstance(value, torch.Tensor)] 

158 if self.expression and not self.shape_env.evaluate_guards_expression(self.expression, tensors): 

159 raise ValueError( 

160 "Graph input shape/stride violates the traced shape guards. Check unmarked static dimensions, " 

161 "related input sizes and shape-dependent branches. Compile with representative dynamic sizes > 1; " 

162 "sizes 0/1 and Python shape branches may specialize the graph." 

163 )