Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / compile / trainer.py: 97%

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""" 

16Graph Trainer - Graph-mode Trainer 

17 

18Users provide model code and parallel configuration. 

19Framework automatically handles all parallel logic. 

20 

21The trainer composes ``GraphCompiler``: compilation and graph execution 

22(``compile`` / ``forward_backward``) are delegated to the compiler, while the 

23trainer owns the training policy -- optimizer lifecycle, grad clip, batch 

24device placement, and the ``train`` loop. Model / device / compiled-graph 

25state lives on the compiler (``trainer._compiler``). 

26""" 

27 

28from __future__ import annotations 

29 

30__all__ = ["GraphTrainer"] 

31 

32import logging 

33from typing import Any, Callable, Iterable, Iterator, List, Optional 

34 

35import torch 

36import torch.distributed as dist 

37 

38from .compiler import GraphCompiler 

39from .pass_config import PassConfig 

40from .graph_parallel_plan import GraphParallelPlan 

41 

42_LOG = logging.getLogger(__name__) 

43 

44 

45class GraphTrainer: 

46 """ 

47 Graph-mode Trainer 

48 

49 Users provide model code and parallel configuration. 

50 Framework automatically handles all parallel logic. 

51 

52 Holds a ``GraphCompiler`` and delegates compilation and graph execution 

53 to it; the optimizer and the training loop stay here. 

54 """ 

55 

56 def __init__( 

57 self, 

58 model: torch.nn.Module, 

59 train_fn: Callable, 

60 pass_config: Optional[PassConfig] = None, 

61 parallel_plan: Optional[GraphParallelPlan] = None, 

62 trainer_config: Optional[Any] = None, 

63 optimizer_config: Optional[dict] = None, 

64 device: Optional[torch.device] = None, 

65 mesh_context: Optional[Any] = None, 

66 dynamic: Optional[bool] = None, 

67 dynamic_arg_dims: dict[str, int | list[int]] | None = None, 

68 ) -> None: 

69 """ 

70 Args: 

71 model: Model to train 

72 train_fn: Training function signature: (model, input, label) -> loss 

73 pass_config: Parallel configuration. When omitted, GraphTrainer 

74 uses ``trainer_config`` to project topology intent onto a 

75 graph-mode ``PassConfig``; if ``trainer_config`` is also 

76 omitted, it falls back to ``PassConfig()``. 

77 parallel_plan: GraphParallelPlan declaring which modules to shard 

78 (optional; enables declarative sharding) 

79 trainer_config: Optional eager trainer config used only to infer a 

80 default graph-mode ``PassConfig`` when ``pass_config`` is not 

81 provided. 

82 optimizer_config: Optimizer configuration 

83 device: Device to place the model and run training on. Defaults to 

84 the NPU device when available, otherwise CPU. 

85 mesh_context: Optional automodel ``MeshContext`` carrying a 

86 pre-built TP/FSDP mesh (forwarded to ``GraphCompiler``). When 

87 provided, the TP group is reused as-is (boundary forwards 

88 already hold the group object) and only the FSDP shard 

89 sub-mesh is registered under ``"fsdp"``. Use this to feed an 

90 automodel TP-sharded model into the graph-mode FSDP pass. 

91 dynamic: Enable symbolic user input dimensions (forwarded to GraphCompiler). 

92 dynamic_arg_dims: Optional input paths to dynamic dimensions (forwarded to GraphCompiler). 

93 """ 

94 # The compiler owns the model, device, config, and the compiled 

95 # graph; the trainer keeps only the optimizer / loop policy. 

96 self._compiler = GraphCompiler( 

97 model=model, 

98 train_fn=train_fn, 

99 pass_config=pass_config, 

100 parallel_plan=parallel_plan, 

101 trainer_config=trainer_config, 

102 device=device, 

103 mesh_context=mesh_context, 

104 dynamic=dynamic, 

105 dynamic_arg_dims=dynamic_arg_dims, 

106 ) 

107 self.optimizer_config = optimizer_config or {} 

108 self.optimizer = None 

109 

110 def compile(self, **inputs: Any) -> None: 

111 """ 

112 Compile model into parallel graph 

113 

114 Users can explicitly call this, or it will be automatically compiled 

115 at first train_step 

116 

117 Args: 

118 **inputs: Model inputs, forwarded to ``train_fn`` as keyword 

119 arguments and used to trace the joint graph 

120 """ 

121 self._compiler.compile(**inputs) 

122 self._init_optimizer() 

123 

124 def train_step(self, **inputs: Any) -> Any: 

125 """ 

126 Execute one training step 

127 

128 Compiles lazily on the first step. ``compile`` also initializes the 

129 optimizer: the compiler's own lazy path in ``forward_backward`` cannot 

130 do that, because the optimizer belongs to the trainer. 

131 

132 Args: 

133 **inputs: Model inputs, forwarded to ``train_fn`` as keyword 

134 arguments 

135 

136 Returns: 

137 loss: Loss value 

138 """ 

139 if self._compiler.is_compiled is False: 

140 self.compile(**inputs) 

141 

142 return self._compiler.forward_backward(**inputs) 

143 

144 @property 

145 def pass_config(self) -> PassConfig: 

146 """Return the compiler-owned graph pass configuration.""" 

147 return self._compiler.pass_config 

148 

149 def to(self, device: torch.device) -> "GraphTrainer": 

150 """Move the model to ``device`` and remember it for batch placement. 

151 

152 Args: 

153 device: Target device for the model and subsequent batches 

154 

155 Returns: 

156 self, so calls can be chained 

157 """ 

158 self._compiler.to(device) 

159 return self 

160 

161 def optimizer_step(self) -> None: 

162 """Optimizer update""" 

163 if self.optimizer is None: 

164 return 

165 

166 if self.optimizer_config.get("grad_clip"): 

167 torch.nn.utils.clip_grad_norm_( 

168 self._compiler.model.parameters(), self.optimizer_config["grad_clip"] 

169 ) 

170 

171 self.optimizer.step() 

172 self.optimizer.zero_grad() 

173 

174 def _place_on_device(self, inputs: dict) -> dict: 

175 """Move a dict of model inputs onto the compiler's device.""" 

176 device = self._compiler.device 

177 if device is None: 

178 return inputs 

179 return { 

180 key: value.to(device) if isinstance(value, torch.Tensor) else value 

181 for key, value in inputs.items() 

182 } 

183 

184 def train( 

185 self, 

186 data_iterable: Iterable, 

187 max_steps: Optional[int] = None, 

188 log_interval: Optional[int] = None, 

189 log_fn: Optional[Callable[[int, Any], None]] = None, 

190 ) -> List[Any]: 

191 """Run the full training loop over ``data_iterable``. 

192 

193 The data iterator must yield dicts of model inputs (the keyword 

194 arguments ``train_fn`` and ``train_step`` consume, e.g. 

195 ``{"input_ids": ..., "labels": ...}``). Each batch is moved onto the 

196 compiler's device, then ``train_step(**inputs)`` + 

197 ``optimizer_step`` are driven. The graph is compiled lazily on the 

198 first batch via ``train_step``. 

199 

200 Args: 

201 data_iterable: An iterable / iterator of model-input dicts. 

202 max_steps: Stop after this many steps. Runs the whole iterator when 

203 ``None``. 

204 log_interval: Log a loss every ``log_interval`` steps (requires 

205 ``log_fn`` or a rank-0 printer). 

206 log_fn: Callback ``log_fn(step, loss)`` for progress reporting. When 

207 ``None`` the loss is printed to stdout on ``log_interval``. 

208 

209 Returns: 

210 List of per-step losses. 

211 """ 

212 if not isinstance(data_iterable, Iterator): 

213 data_iterable = iter(data_iterable) 

214 

215 losses: List[Any] = [] 

216 rank = dist.get_rank() if dist.is_initialized() else 0 

217 for step, batch in enumerate(data_iterable): 

218 if max_steps is not None and step >= max_steps: 

219 break 

220 

221 inputs = self._place_on_device(batch) 

222 

223 loss, _ = self.train_step(**inputs) 

224 self.optimizer_step() 

225 losses.append(loss) 

226 

227 if rank == 0 and log_interval and (step + 1) % log_interval == 0: 

228 if log_fn is not None: 

229 log_fn(step + 1, loss) 

230 else: 

231 _LOG.info("Step %s | Loss: %.4f", step + 1, loss.item()) 

232 

233 return losses 

234 

235 def _init_optimizer(self): 

236 """Initialize optimizer on the model's (FSDP-sharded) parameters. 

237 

238 FSDPPass shards the compiler's model parameters in place during 

239 compile, so ``model.parameters()`` already yields the local shards 

240 and the optimizer needs no FSDP awareness. 

241 

242 When ``torch_npu`` is installed but no NPU is available (e.g. a 

243 CPU-only CI run), Adam's automatic foreach/fused kernel selection 

244 probes ``torch_npu.npu.current_device()`` via ``_lazy_init()``, 

245 which raises even though the parameters live on CPU. Disable the 

246 probe explicitly in that case so single-process CPU runs (and the 

247 UT suite) succeed; NPU runs keep the default foreach/fused path. 

248 """ 

249 optimizer_class = torch.optim.Adam 

250 kwargs = {"lr": self.optimizer_config.get("lr", 1e-4)} 

251 if hasattr(torch, "npu") and not torch.npu.is_available(): 

252 kwargs["foreach"] = False 

253 self.optimizer = optimizer_class(self._compiler.model.parameters(), **kwargs)