Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Configurable :class:`Module` with declarative init and tensor parallelism.

17Layers inherit from :class:`Module`, declare ``Config`` with ``param_init`` and

18:class:`~hyper_parallel.dmodule.sharding.ShardingConfig`, then call

19:meth:`Module.init_states` and :meth:`Module.parallelize`.

20"""

21from __future__ import annotations

23__all__ = ["Module"]

25import inspect

26import logging

27from collections.abc import Callable

28from dataclasses import dataclass

29from typing import Any

31from hyper_parallel import get_platform

32from hyper_parallel.core.dtensor.device_mesh import DeviceMesh

33from hyper_parallel.core.dtensor.dtensor import DTensor, distribute_tensor

34from hyper_parallel.core.dtensor.placement_types import Placement

35from hyper_parallel.config.configurable import Configurable, enforce_module_config_slots

36from hyper_parallel.dmodule.sharding import ShardingConfig, resolve_placements

38logger = logging.getLogger(__name__)

40_FORWARD_POSITIONAL_KINDS = (

41 inspect.Parameter.POSITIONAL_ONLY,

42 inspect.Parameter.POSITIONAL_OR_KEYWORD,

43)

44platform = get_platform()

45_PlatformModule = platform.Module

47_created_classes: dict[type, type] = {}

50def _get_attr_by_path(obj: Any, path: str) -> Any:

51 parts = path.split(".")

52 for part in parts[:-1]:

53 obj = getattr(obj, part)

54 return getattr(obj, parts[-1])

57def _set_param_by_path(module: _PlatformModule, path: str, param: Any) -> None:

58 parts = path.split(".")

59 if len(parts) == 1:

60 module.register_parameter(parts[0], param)

61 return

62 parent = module

63 for part in parts[:-1]:

64 parent = getattr(parent, part)

65 parent.register_parameter(parts[-1], param)

68def _placements_equal(

69 left: tuple[Placement, ...] | list[Placement],

70 right: list[Placement],

71) -> bool:

72 if len(left) != len(right):

73 return False

74 return all(a == b for a, b in zip(left, right))

77class Module(_PlatformModule, Configurable):

78 """Declarative distributed layer base class.

80 Combines the platform :class:`~hyper_parallel.platform.platform.Module` with

81 :class:`~hyper_parallel.config.configurable.Configurable`. Each subclass

82 defines ``Config`` (hyperparameters, ``param_init``, ``sharding_config``) and

83 ``__init__(self, config)``.

85 Example::

87 class Linear(Module):

88 class Config(Module.Config):

89 in_features: int = 128

90 out_features: int = 64

91 param_init = {"weight": nn.init.kaiming_uniform_}

92 sharding_config = ShardingConfig(

93 state_shardings={"weight": {MeshAxisName.TP: Shard(0)}},

94 )

96 def __init__(self, config: Config):

97 super().__init__()

98 self.weight = nn.Parameter(torch.empty(config.out_features, config.in_features))

100 def forward(self, x):

101 return F.linear(x, self.weight, None)

102

103 layer = Linear.Config().build()

104 layer.init_states()

105 layer.parallelize(tp_mesh)

106 """

107

108 _param_init: dict[str, Callable] | None = None

109 _sharding_config: ShardingConfig | None = None

110 _pos_arg_list: list[str] | None = None

111 _parallelized: bool = False

112

113 @dataclass(kw_only=True, slots=True)

114 class Config(Configurable.Config):

115 """Per-layer config: initializers and optional sharding spec.

116

117 Args:

118 param_init: Map parameter name (for example ``"weight"``) to an

119 init callable applied in :meth:`Module.init_states`.

120 sharding_config: Optional :class:`~hyper_parallel.dmodule.sharding.ShardingConfig`

121 consumed by :meth:`Module.parallelize`.

122

123 Example::

124

125 Linear.Config(

126 in_features=128,

127 out_features=64,

128 param_init={"weight": nn.init.kaiming_uniform_},

129 sharding_config=ShardingConfig(

130 state_shardings={"weight": {MeshAxisName.TP: Shard(0)}},

131 ),

132 ).build()

133 """

134

135 param_init: dict[str, Callable] | None = None

136 sharding_config: ShardingConfig | None = None

137

138 def build(self, **kwargs):

139 """Build the layer and attach ``param_init`` / ``sharding_config``.

140

141 Args:

142 **kwargs: Runtime-only ``__init__`` arguments (must not overlap

143 config field names).

144

145 Returns:

146 Instance of the outer :class:`Module` subclass.

147

148 Example::

149

150 layer = Linear.Config(in_features=128, out_features=64).build()

151 """

152 instance = Configurable.Config.build(self, **kwargs)

153 if self.param_init is not None:

154 instance._param_init = self.param_init

155 if self.sharding_config is not None:

156 instance._sharding_config = self.sharding_config

157 return instance

158

159 def __init_subclass__(cls, **kwargs):

160 super().__init_subclass__(**kwargs)

161 if "Config" in cls.__dict__:

162 enforce_module_config_slots(cls.__dict__["Config"], cls.__name__)

163

164 @property

165 def sharding_config(self) -> ShardingConfig | None:

166 """Sharding spec used by :meth:`parallelize` (from Config at build time)."""

167 return self._sharding_config

168

169 def init_states(

170 self,

171 *,

172 buffer_device: Any | None = None,

173 ) -> None:

174 """Initialize parameters and buffers for this subtree.

175

176 Recurses into child :class:`Module` instances, then applies

177 ``param_init`` on local parameters and calls :meth:`_init_self_buffers`.

178

179 Args:

180 buffer_device: Optional device for buffer creation (subclasses may

181 use this in :meth:`_init_self_buffers`).

182

183 Example::

184

185 model = model_cfg.build()

186 model.init_states(buffer_device=torch.device("npu", 0))

187 """

188 queue = list(self.children())

189 while queue:

190 child = queue.pop(0)

191 if isinstance(child, Module):

192 child.init_states(buffer_device=buffer_device)

193 else:

194 queue.extend(child.children())

195

196 self._init_self_parameters()

197

198 dtensor_meta = {

199 name: (buf.device_mesh, buf.placements)

200 for name, buf in self._buffers.items()

201 if isinstance(buf, DTensor)

202 }

203 self._init_self_buffers(buffer_device=buffer_device)

204 for name, (mesh, placements) in dtensor_meta.items():

205 new_buf = self._buffers.get(name)

206 if new_buf is None or isinstance(new_buf, DTensor):

207 continue

208 persistent = name not in self._non_persistent_buffers_set

209 self.register_buffer(

210 name,

211 distribute_tensor(new_buf, mesh, list(placements)),

212 persistent=persistent,

213 )

214

215 def _init_self_parameters(self) -> None:

216 for name, param in self.named_parameters(recurse=False):

217 self._init_param(name, param)

218

219 def _init_param(self, name: str, param: Any) -> None:

220 """Apply ``param_init`` initializer for a single parameter."""

221 if self._param_init is None:

222 raise ValueError(

223 f"No param_init found for parameter '{name}' in "

224 f"{type(self).__name__}. Set param_init on this module's Config."

225 )

226 if name not in self._param_init:

227 raise ValueError(

228 f"No initializer for parameter '{name}' in {type(self).__name__}. "

229 f"Available: {list(self._param_init.keys())}"

230 )

231 self._param_init[name](param)

232

233 def _init_self_buffers(self, *, buffer_device: Any | None = None) -> None:

234 pass

235

236 def _cache_pos_arg_names(self) -> list[str]:

237 """Return positional argument names of :meth:`forward` (cached)."""

238 if self._pos_arg_list is not None:

239 return self._pos_arg_list

240 sig = inspect.signature(type(self).forward)

241 self._pos_arg_list = [

242 p.name

243 for p in sig.parameters.values()

244 if p.kind in _FORWARD_POSITIONAL_KINDS and p.name != "self"

245 ]

246 return self._pos_arg_list

247

248 def parallelize(self, tp_mesh: DeviceMesh) -> None:

249 """Apply ``sharding_config`` and wrap ``forward`` with redistribution.

250

251 For layers with a :class:`~hyper_parallel.dmodule.sharding.ShardingConfig`:

252 shards parameters, then runs

253 ``redistribute inputs -> forward -> redistribute outputs``.

254 Always recurses into child :class:`Module` nodes. Callable at most once

255 per instance.

256

257 Args:

258 tp_mesh: Device mesh whose ``mesh_dim_names`` align with

259 :class:`~hyper_parallel.dmodule.types.MeshAxisName` keys.

260

261 Raises:

262 ValueError: If already parallelized or mesh has no axis names.

263 NotImplementedError: If ``sharding_config.local_map`` is set (M9).

264

265 Example::

266

267 mesh = init_device_mesh("npu", mesh_shape=(2,), mesh_dim_names=("tp",))

268 model = model_cfg.build()

269 model.init_states()

270 model.parallelize(mesh)

271 """

272 if self._parallelized:

273 raise ValueError(

274 f"{type(self).__name__} has already been parallelized. "

275 "Module.parallelize() must be called at most once per instance."

276 )

277 self._parallelized = True

278

279 sc = self.sharding_config

280 if sc is None:

281 for child in self.children():

282 if isinstance(child, Module):

283 child.parallelize(tp_mesh)

284 return

285

286 if sc.local_map is not None:

287 raise NotImplementedError("local_map will be added in M9")

288

289 mesh_axis_names = tuple(tp_mesh.mesh_dim_names or ())

290 if not mesh_axis_names:

291 raise ValueError("DeviceMesh must have mesh_dim_names for parallelize()")

292

293 self._shard_states(tp_mesh, sc, mesh_axis_names)

294 _ = self._cache_pos_arg_names()

295 unbound_forward = type(self).forward

296

297 def forward_with_redistribution(*args, **kwargs):

298 args, kwargs = self._redistribute_inputs(tp_mesh, mesh_axis_names, sc, args, kwargs)

299 outputs = unbound_forward(self, *args, **kwargs)

300 return self._redistribute_outputs(tp_mesh, mesh_axis_names, sc, outputs)

301

302 self.forward = forward_with_redistribution # type: ignore[method-assign]

303

304 for child in self.children():

305 if isinstance(child, Module):

306 child.parallelize(tp_mesh)

307

308 def _shard_states(

309 self,

310 tp_mesh: DeviceMesh,

311 sharding_config: ShardingConfig,

312 mesh_axis_names: tuple[str, ...],

313 ) -> None:

314 """Shard parameters listed in ``sharding_config.state_shardings``."""

315 for path, named_placements in sharding_config.state_shardings.items():

316 param = _get_attr_by_path(self, path)

317 placements = resolve_placements(named_placements, mesh_axis_names)

318 if isinstance(param, DTensor):

319 if not _placements_equal(tuple(param.placements), placements):

320 raise ValueError(

321 f"{type(self).__name__}.{path} is already a DTensor with "

322 f"placements {param.placements}, but sharding_config expects "

323 f"{placements}."

324 )

325 continue

326 tensor = param.data if hasattr(param, "data") else param

327 new_local = distribute_tensor(tensor, tp_mesh, placements)

328 requires_grad = getattr(param, "requires_grad", True)

329 _set_param_by_path(

330 self,

331 path,

332 platform.Parameter(new_local, requires_grad=requires_grad),

333 )

334

335 def _redistribute_inputs(

336 self,

337 tp_mesh: DeviceMesh,

338 mesh_axis_names: tuple[str, ...],

339 sharding_config: ShardingConfig,

340 args: tuple,

341 kwargs: dict,

342 ) -> tuple[tuple, dict]:

343 """Redistribute forward inputs per ``in_src_shardings`` / ``in_dst_shardings``."""

344 if (

345 sharding_config.in_dst_shardings is None

346 and sharding_config.in_src_shardings is None

347 ):

348 return args, kwargs

349

350 pos_arg_names = [

351 name for name in self._cache_pos_arg_names() if name not in kwargs

352 ]

353 new_kwargs = dict(zip(pos_arg_names, args))

354 new_kwargs.update(kwargs)

355

356 in_dst_shardings = sharding_config.in_dst_shardings or {}

357 in_src_shardings = sharding_config.in_src_shardings or {}

358

359 for name, value in new_kwargs.items():

360 if not platform.is_tensor(value) and not isinstance(value, DTensor):

361 continue

362 src_named = in_src_shardings.get(name)

363 dst_named = in_dst_shardings.get(name)

364 if src_named is None and dst_named is None:

365 continue

366

367 if not isinstance(value, DTensor):

368 if src_named is not None:

369 layout = resolve_placements(src_named, mesh_axis_names)

370 value = DTensor.from_local(value, tp_mesh, layout)

371 elif dst_named is not None:

372 layout = resolve_placements(dst_named, mesh_axis_names)

373 value = DTensor.from_local(value, tp_mesh, layout)

374

375 if dst_named is not None and isinstance(value, DTensor):

376 desired = resolve_placements(dst_named, mesh_axis_names)

377 if not _placements_equal(tuple(value.placements), desired):

378 value = value.redistribute(tp_mesh, desired)

379

380 new_kwargs[name] = value

381

382 new_args = tuple(new_kwargs.pop(name) for name in pos_arg_names)

383 return new_args, new_kwargs

384

385 def _redistribute_outputs(

386 self,

387 tp_mesh: DeviceMesh,

388 mesh_axis_names: tuple[str, ...],

389 sharding_config: ShardingConfig,

390 outputs: Any,

391 ) -> Any:

392 """Redistribute forward outputs per ``out_dst_shardings``."""

393 out_named = sharding_config.out_dst_shardings

394 if out_named is None:

395 return outputs

396 if not isinstance(outputs, DTensor):

397 return outputs

398 desired = resolve_placements(out_named, mesh_axis_names)

399 if not _placements_equal(tuple(outputs.placements), desired):

400 outputs = outputs.redistribute(tp_mesh, desired)

401 return outputs

402

403 @classmethod

404 def from_nn_module(cls, nn_module_cls: type) -> type["Module"]:

405 """Wrap a platform ``nn.*`` class as a :class:`Module` subclass.

406

407 Args:

408 nn_module_cls: Source class (for example ``torch.nn.Conv2d``).

409

410 Returns:

411 New type mixing *nn_module_cls* and :class:`Module`. Results are

412 cached per source class.

413

414 Example::

415

416 Conv2d = Module.from_nn_module(torch.nn.Conv2d)

417 layer = Conv2d(3, 16, kernel_size=3)

418 layer.init_states()

419 """

420 if nn_module_cls in _created_classes:

421 return _created_classes[nn_module_cls]

422

423 attrs: dict[str, Any] = {}

424 if hasattr(nn_module_cls, "reset_parameters"):

425

426 def _init_self_parameters(self: Any) -> None:

427 self.reset_parameters()

428

429 attrs["_init_self_parameters"] = _init_self_parameters

430

431 name = f"Module({nn_module_cls.__name__})"

432 new_cls = type(name, (nn_module_cls, Module), attrs)

433 new_cls.__module__ = __name__

434 new_cls.__qualname__ = name

435 _created_classes[nn_module_cls] = new_cls

436 return new_cls

437

438

439def _get_torch_nn_container(kind: str) -> type:

440 """Return ``torch.nn.{ModuleList,ModuleDict,Sequential}`` (requires PyTorch, lazy)."""

441 # pylint: disable=C0415

442 try:

443 import torch.nn as nn

444 except ImportError as exc:

445 raise NotImplementedError(

446 f"{kind} container wrappers require PyTorch (torch.nn) in M1"

447 ) from exc

448

449 mapping = {

450 "ModuleList": nn.ModuleList,

451 "ModuleDict": nn.ModuleDict,

452 "Sequential": nn.Sequential,

453 }

454 if kind not in mapping:

455 raise ValueError(f"Unknown container kind: {kind}")

456 return mapping[kind]

457

458

459_LAZY_CONTAINER_NAMES = frozenset({"ModuleList", "ModuleDict", "Sequential"})

460

461

462def __getattr__(name: str):

463 """Lazy wrappers for ``ModuleList``, ``ModuleDict``, and ``Sequential``.

464

465 Example::

466

467 from hyper_parallel.dmodule import module as hp_module

468

469 layers = hp_module.ModuleList([Linear.Config().build() for _ in range(4)])

470 """

471 if name in _LAZY_CONTAINER_NAMES:

472 cls = Module.from_nn_module(_get_torch_nn_container(name))

473 globals()[name] = cls

474 return cls

475 raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / dmodule / module.py: 54%

200 statements