Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / build.py: 100%

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"""Device-free orchestration for snapshot-driven code generation.""" 

16 

17from typing import Any, Optional 

18 

19from hyper_parallel.codegen.artifact import ArtifactLayout 

20from hyper_parallel.codegen.model_build_overlay import use_codegen_model_build_overlay 

21from hyper_parallel.codegen.planner_overlay import use_codegen_planner_overlay 

22from hyper_parallel.codegen.session import CodegenBuildSession 

23from hyper_parallel.distributed.mesh import ( 

24 DistributedSetup, 

25 LogicalDeviceMesh, 

26 MeshContext, 

27 ParallelTopologySpec, 

28 infer_minimum_world_size, 

29 resolve_parallel_layout, 

30) 

31from hyper_parallel.trainer.config import normalize_distributed_setup_overrides 

32 

33 

34def generate_from_native_build( 

35 config: Any, 

36 *, 

37 artifact_dir: Optional[str] = None, 

38) -> ArtifactLayout: 

39 """Observe the configured model on meta and publish a fresh bundle. 

40 

41 Args: 

42 config: Resolved TrainerConfig. Only model and planning settings are 

43 consumed; data, optimizer, scheduler, and training loops are never built. 

44 artifact_dir: Optional explicit bundle output directory. 

45 

46 Returns: 

47 The freshly generated artifact layout. 

48 

49 Raises: 

50 RuntimeError: If the native model does not publish a complete artifact. 

51 """ 

52 return _build_and_publish(config, artifact_dir) 

53 

54 

55def _build_and_publish(config: Any, artifact_dir: Optional[str]) -> ArtifactLayout: 

56 """Build and publish from a process-group-free planning topology.""" 

57 distributed_setup = create_offline_distributed_setup(config) 

58 normalize_distributed_setup_overrides(distributed_setup, config) 

59 

60 session = CodegenBuildSession() 

61 with use_codegen_model_build_overlay(session, artifact_dir) as result, use_codegen_planner_overlay(): 

62 config.model.build( 

63 distributed_setup=distributed_setup, 

64 peft_config=config.peft, 

65 activation_checkpoint=config.activation_checkpoint.mode, 

66 activation_checkpoint_selection=config.activation_checkpoint.selection, 

67 swap_inputs=getattr(config.activation_checkpoint, "swap_inputs", False), 

68 allow_cache_entry_mutation=getattr( 

69 config.activation_checkpoint, 

70 "allow_cache_entry_mutation", 

71 False, 

72 ), 

73 activation_swap=config.activation_swap, 

74 compile_config=config.compile, 

75 model_init_dtype=config.model_init_dtype, 

76 ) 

77 

78 if result.artifact_layout is None: 

79 failure = getattr(result.snapshot, "failure", None) 

80 if failure is not None: 

81 raise RuntimeError( 

82 f"codegen {failure.stage} failed: {failure.reason}" 

83 ) 

84 raise RuntimeError("native model build completed without publishing a codegen artifact") 

85 return result.artifact_layout 

86 

87 

88def create_offline_distributed_setup(config: Any) -> DistributedSetup: 

89 """Project YAML parallel settings into a process-group-free setup. 

90 

91 Args: 

92 config: Trainer configuration describing compute and FSDP parallel degrees. 

93 

94 Returns: 

95 A setup with named logical dense and expert meshes, without communication groups. 

96 """ 

97 topology_spec = ParallelTopologySpec.from_config(config) 

98 layout = resolve_parallel_layout( 

99 topology_spec, 

100 infer_minimum_world_size(topology_spec), 

101 ) 

102 topology = LogicalDeviceMesh( 

103 mesh_dim_names=("dp", "cp", "tp"), 

104 mesh_shape=(layout.dp_size, topology_spec.cp_size, topology_spec.tp_size), 

105 ) 

106 mesh_context = MeshContext.from_layout(layout) 

107 mesh_context.device_mesh = topology 

108 if topology_spec.ep_size > 1: 

109 mesh_context.fsdp_moe_mesh = LogicalDeviceMesh( 

110 mesh_dim_names=("edp", "ep"), 

111 mesh_shape=(topology.size() // topology_spec.ep_size, topology_spec.ep_size), 

112 ) 

113 return DistributedSetup.from_config( 

114 config, 

115 mesh_context=mesh_context, 

116 strategy_enabled=True, 

117 ) 

118 

119 

120__all__ = ["create_offline_distributed_setup", "generate_from_native_build"]