Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / encoder_parallel / mesh.py: 92%

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"""Derive visual encoder coordinates from the existing decoder mesh.""" 

16 

17import math 

18from typing import Any 

19 

20 

21def derive_reshaped_submesh(root_mesh: Any, *, source_dims: tuple[str, ...], 

22 target_sizes: tuple[int, ...], target_names: tuple[str, ...]) -> Any: 

23 """Reshape a named submesh while preserving its physical rank order. 

24 

25 Args: 

26 root_mesh: Existing decoder mesh. 

27 source_dims: Decoder axes forming one visual parent group. 

28 target_sizes: Encoder-DP and Encoder-CP degrees. 

29 target_names: Names assigned to the derived axes. 

30 """ 

31 source = root_mesh[source_dims] 

32 if math.prod(target_sizes) != math.prod(source.mesh_shape) or len(target_sizes) != len(target_names): 

33 raise ValueError("encoder mesh shape does not match decoder parent") 

34 flat = source.flatten("encoder_parent") 

35 return flat._unflatten("encoder_parent", target_sizes, target_names) # pylint: disable=protected-access 

36 

37 

38def encoder_coordinates(parent_rank: int, encoder_cp_size: int) -> tuple[int, int]: 

39 """Return Encoder-DP and Encoder-CP coordinates in parent rank order. 

40 

41 Args: 

42 parent_rank: Linear rank within the decoder CP-by-TP parent. 

43 encoder_cp_size: Number of sequence shards per visual bucket. 

44 """ 

45 return divmod(parent_rank, encoder_cp_size) 

46 

47 

48def decoder_coordinates(parent_rank: int, decoder_tp_size: int) -> tuple[int, int]: 

49 """Return decoder CP and TP coordinates in parent rank order. 

50 

51 Args: 

52 parent_rank: Linear rank within the decoder CP-by-TP parent. 

53 decoder_tp_size: Number of decoder TP ranks per CP rank. 

54 """ 

55 return divmod(parent_rank, decoder_tp_size)