Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / plan / offline_mesh.py: 37%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# Licensed under the Apache License, Version 2.0 

3# ============================================================================ 

4"""Offline mesh: a static stand-in for a real distributed ``DeviceMesh``. 

5 

6``ShardingPlanner.plan()`` consumes a mesh through a narrow contract — 

7``mesh_dim_names`` + ``mesh_shape`` for axis validation and planning 

8(``getattr`` guarded, see ``_validate_dtensor_axes`` / 

9``_build_mesh_dim_names`` in ``components.distributed.sharding_planner``). 

10A real ``DeviceMesh`` requires an initialized 

11process group — on the codegen host (CPU, no ``torchrun``) the default 

12``hccl`` backend fails at construction. The offline mesh provides the same 

13*shape* interface over the parallel dims projected from the YAML, so the 

14plan can be derived (and frozen into meta) before any training process 

15group exists. 

16 

17The plan coordinate system still excludes DP placements, but the offline mesh 

18must carry the runtime DP domain so planner checks that depend on the dense 

19region (notably extended EP for MoE) see the same topology as native trainer: 

20 

21- ``dp`` is declared with the data-parallel size derived from ``WORLD_SIZE``; 

22- ``tp`` / ``cp`` axes are declared whenever the config declares them 

23 (present even at size 1 so the planner's validation can see them); 

24- ``ep`` is **never** declared as a mesh axis. An explicit ``"ep"`` axis 

25 selects old-style EP in the planner (``_mark_hf_native_moe``), which 

26 leaves a TP key on expert weights. The HF-native extended expert-parallel 

27 path derives the EP group from the dense region and must NOT see an ``"ep"`` 

28 axis. ``ep_size`` still reaches 

29 the planner through ``plan(..., ep_size=...)``, so EP is planned; it 

30 just is not a coordinate of the frozen plan; 

31- ``pp`` is declared as an axis but stays outside the plan coordinate system; 

32 it never takes part in ``mesh_dim_names`` (the planner rejects pipeline 

33 parallelism with extended expert parallelism, and the plan coordinate system 

34 is a single dp slice). 

35""" 

36from __future__ import annotations 

37 

38from dataclasses import dataclass 

39 

40from hyper_parallel.codegen.spec.types import ParallelDimsSpec 

41 

42 

43@dataclass(frozen=True) 

44class OfflineMesh: 

45 """Static mesh with the surface ``ShardingPlanner.plan()`` needs. 

46 

47 Immutable and JSON-safe: it carries only axis names + sizes, never a 

48 process group. 

49 """ 

50 

51 mesh_dim_names: tuple[str, ...] = () 

52 mesh_shape: tuple[int, ...] = () 

53 

54 def __post_init__(self) -> None: 

55 if len(self.mesh_dim_names) != len(self.mesh_shape): 

56 raise ValueError( 

57 f"mesh_dim_names {self.mesh_dim_names!r} and mesh_shape " 

58 f"{self.mesh_shape!r} must have the same length" 

59 ) 

60 

61 

62def build_offline_mesh(parallel_dims: ParallelDimsSpec) -> OfflineMesh: 

63 """Build the static mesh for the projected parallel dims. 

64 

65 Axis rules (see the module docstring): 

66 

67 - ``dp`` / ``cp`` / ``tp`` present when the config declares them (even at size 1); 

68 - ``ep`` never named — naming it would select old-style EP and leave a 

69 TP key on fused expert weights; ``ep_size`` is passed to ``plan()`` 

70 instead so the planner derives the expert mesh; 

71 - ``pp`` declared when the config sets ``pp_size > 1``; 

72 """ 

73 names: list[str] = [] 

74 sizes: list[int] = [] 

75 if parallel_dims.dp_size >= 1: 

76 names.append("dp") 

77 sizes.append(parallel_dims.dp_size) 

78 if parallel_dims.cp_size >= 1: 

79 names.append("cp") 

80 sizes.append(parallel_dims.cp_size) 

81 if parallel_dims.tp_size >= 1: 

82 names.append("tp") 

83 sizes.append(parallel_dims.tp_size) 

84 if parallel_dims.pp_size > 1: 

85 names.append("pp") 

86 sizes.append(parallel_dims.pp_size) 

87 return OfflineMesh(mesh_dim_names=tuple(names), mesh_shape=tuple(sizes)) 

88 

89 

90__all__ = ["OfflineMesh", "build_offline_mesh"]