Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/core/dtensor/

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"""Torch implementations backing ``hyper_parallel.core.dtensor``.

17The DTensor core used to reach these routines through the platform abstraction

18layer. Only the torch backend is supported now, so the nontrivial torch

19implementations live here and every caller calls them directly.

21Collectives and the small rank/device/module helpers that other ``core``

22modules also needed were lifted into :mod:`hyper_parallel.core.utils`. The

23collectives stay re-exported here because the debug tracer patches these

24attributes on this module, so ``_utils.<name>`` keeps working for that path;

25callers elsewhere should import from :mod:`hyper_parallel.core.utils`.

26"""

27# pylint: disable=C9006,C9007

28from contextlib import contextmanager

29from typing import Any, Optional, Union

31import torch

32from torch import nn

33from torch._C._distributed_c10d import ProcessGroup

34from torch.distributed.distributed_c10d import _get_default_group

36import torch.distributed as dist

38from hyper_parallel.core.utils.communication import (

39 EXISTING_COMM_GROUPS,

40 differentiable_all_gather_concat,

41 differentiable_all_to_all,

42 differentiable_all_to_all_single,

43 differentiable_all_to_all_single_async,

44 differentiable_all_reduce,

45 differentiable_reduce_scatter,

46 differentiable_variable_all_gather,

47)

48from hyper_parallel.core.utils.communication import get_device_handle

49from hyper_parallel.core.shard.utils import get_op_name

52# ---------------------------------------------------------------------------

53# Small shared helpers

54# ---------------------------------------------------------------------------

56def tensor_type_cast(input_data, cast_type):

57 """Cast tensor to specified data type."""

58 type_mapping = {

59 'float32': torch.float32,

60 'float16': torch.float16,

61 'int64': torch.int64,

62 'int32': torch.int32

63 }

64 if cast_type not in type_mapping:

65 raise ValueError(f"Unknown cast type: {cast_type}. Supported types: {list(type_mapping.keys())}")

66 return input_data.to(type_mapping[cast_type])

69# ---------------------------------------------------------------------------

70# Device / process group helpers

71# ---------------------------------------------------------------------------

73def device_count(device_handle):

74 """Return the number of available devices for *device_handle*."""

75 return device_handle.device_count()

78def device_type():

79 """Return the current device type string ("npu" for NPU, "cuda" for GPU)."""

80 device_handle = get_device_handle()

81 if device_handle == getattr(torch, "npu", None):

82 return "npu"

83 return "cuda"

86def device(device_idx=None):

87 """Return a :class:`torch.device` for the current device type."""

88 current_device_type = device_type()

89 if device_idx is None:

90 return torch.device(current_device_type)

91 return torch.device(f"{current_device_type}:{device_idx:d}")

94def get_rng_state(device=None, device_handle=None): # pylint: disable=W0621

95 """Get the random number generator state."""

96 if device_handle is None:

97 return torch.get_rng_state()

98 if device is None:

99 return device_handle.get_rng_state()

100 return device_handle.get_rng_state(device)

101

102

103def set_rng_state(state, device=None, device_handle=None): # pylint: disable=W0621

104 """Set the random number generator state."""

105 if device_handle is None:

106 return torch.set_rng_state(state)

107 if device is None:

108 return device_handle.set_rng_state(state)

109 return device_handle.set_rng_state(state, device)

110

111

112def get_created_group(rank_list: Union[list[int], tuple[int]]):

113 """Return an existing process group by rank list, or ``None``."""

114 group_key = str(tuple(sorted(rank_list)))

115 if group_key in EXISTING_COMM_GROUPS:

116 return EXISTING_COMM_GROUPS[group_key]

117 return None

118

119

120def create_group(rank_list):

121 """Create or retrieve a communication group with the specified ranks.

122

123 If a group with the same rank list already exists, returns the existing

124 group instead of creating a new one.

125 """

126 group_key = str(tuple(sorted(rank_list)))

127 if group_key in EXISTING_COMM_GROUPS:

128 return EXISTING_COMM_GROUPS[group_key]

129

130 normalized_rank_list = tuple(sorted(rank_list))

131 world_rank_list = tuple(range(dist.get_world_size()))

132 if normalized_rank_list == world_rank_list:

133 group = _get_default_group()

134 else:

135 group = create_sub_groups(rank_list)[normalized_rank_list]

136

137 EXISTING_COMM_GROUPS[group_key] = group

138 return group

139

140

141def split_group(parent_pg: Optional[ProcessGroup] = None,

142 split_ranks: Optional[list] = None,

143 timeout: Optional[Any] = None,

144 pg_options: Optional[Any] = None,

145 group_desc: Optional[str] = None,

146 ) -> Optional[ProcessGroup]:

147 """Create split groups for every rank list in *split_ranks*.

148

149 Returns the split process group relative to the current rank id.

150 """

151 del parent_pg, timeout, group_desc

152 if split_ranks is None or len(split_ranks) == 0:

153 raise ValueError("split_ranks cannot be None or empty")

154

155 split_group_pg = None

156 for split_rank in split_ranks:

157 dist_group = get_created_group(split_rank)

158 if dist_group is None:

159 dist_group = dist.new_group(ranks=split_rank, pg_options=pg_options)

160 EXISTING_COMM_GROUPS[str(tuple(sorted(split_rank)))] = dist_group

161 if dist.get_rank() in split_rank:

162 split_group_pg = dist_group

163

164 return split_group_pg

165

166

167def init_process_group(*args, **kwargs):

168 """Initialize the default torch distributed process group."""

169 if not dist.is_initialized():

170 dist.init_process_group(*args, **kwargs)

171

172

173# ---------------------------------------------------------------------------

174# Sub-group construction

175# ---------------------------------------------------------------------------

176

177def _validate_intra_step(normalized_template: list[int], template_len: int) -> int:

178 """Verify a normalized template has a uniform step between its members.

179

180 The step between adjacent members defines the stride of the groups built

181 from this template, so a template such as ``[0, 1, 3]`` (steps 1 then 2)

182 has no single well-defined stride and is rejected.

183

184 Args:

185 normalized_template: Template rebased so its first member is ``0``.

186 template_len: Number of members; ``normalized_template`` must hold at

187 least two for a step to exist.

188

189 Returns:

190 int: The step shared by every adjacent pair of members.

191

192 Raises:

193 ValueError: If any adjacent pair differs in step from the first pair.

194 """

195 intra_step = normalized_template[1] - normalized_template[0]

196 for i in range(1, template_len - 1):

197 diff = normalized_template[i + 1] - normalized_template[i]

198 if diff != intra_step:

199 msg = (

200 f"Template must have consistent intra-group step. "

201 f"Found {normalized_template[i+1]} - {normalized_template[i]} = {diff}, "

202 f"expected {intra_step}"

203 )

204 raise ValueError(msg)

205 return intra_step

206

207

208def _compute_group_starts(world_size: int, block_size: int, inter_step: int) -> list[int]:

209 """Compute every block start position that fits inside ``world_size``.

210

211 Start positions are spaced ``inter_step`` apart, and a start is kept only

212 when the whole block fits, i.e. the last rank it can reach stays within

213 ``world_size``. Ranks past the final full block are left uncovered.

214

215 Args:

216 world_size: Total number of processes.

217 block_size: Ranks spanned by one block; used for the fit check.

218 inter_step: Distance between adjacent candidate starts.

219

220 Returns:

221 list[int]: Ascending start positions of the blocks that fit.

222 """

223 return [s for s in range(0, world_size, inter_step) if s + block_size <= world_size]

224

225

226def _build_groups_for_blocks(

227 group_starts: list[int],

228 block_size: int,

229 template_span_int: int,

230 normalized_template: list[int],

231 template_len: int,

232 world_size: int,

233) -> list[list[int]]:

234 """Expand each block start into the groups the template can place there.

235

236 Within a block the template may be slid to offsets ``0 .. block_size -

237 template_span_int``; an offset's group is kept only when every rank it

238 produces lands inside ``world_size``. Groups are appended in start order,

239 then offset order, so every process building the same template sees the

240 same sequence before the caller sorts.

241

242 Args:

243 group_starts: Block start positions, as returned by

244 :func:`_compute_group_starts`.

245 block_size: Ranks spanned by one block.

246 template_span_int: Distance from the template's first member to its

247 last; bounds how far the template can slide within a block.

248 normalized_template: Template rebased so its first member is ``0``.

249 template_len: Number of members in the template.

250 world_size: Total number of processes.

251

252 Returns:

253 list[list[int]]: One rank list per (block, offset) pair that fits.

254 """

255 all_groups = []

256 for start_block in group_starts:

257 max_offset = block_size - template_span_int

258 for offset in range(0, max_offset):

259 group = [start_block + offset + normalized_template[i] for i in range(template_len)]

260 if all(0 <= r < world_size for r in group):

261 all_groups.append(group)

262 return all_groups

263

264

265def _report_template_stage(verbose: bool, my_rank: int, message: str) -> None:

266 """Print one ``generate_groups_from_template`` stage when *verbose* is set.

267

268 Each generated group must be identical across processes, so a template bug

269 shows up as a mismatch in the printed rank lists rather than as a crash.

270 This helper exists so those stage prints stay a single ``verbose`` check

271 each, keeping the caller's local-variable count down.

272

273 Args:

274 verbose: Whether the caller asked for debug output.

275 my_rank: Current process rank, prefixed to every line.

276 message: Stage-specific text; the caller formats it.

277 """

278 if verbose:

279 print(f"Rank {my_rank}: {message}")

280

281

282def generate_groups_from_template(

283 template: Union[list[int], tuple[int, ...]],

284 world_size: int,

285 my_rank: int,

286 verbose: bool = False

287) -> list[list[int]]:

288 """

289 Auto-generate all communication groups from a template (supports any valid starting template).

290

291 Args:

292 template: Template group, e.g. [0,1], [0,2,4,6] or [1,3,5,7]

293 world_size: Total number of processes

294 my_rank: Current process rank (for debug output)

295 verbose: Whether to print debug info

296

297 Returns:

298 Full rank list, e.g.:

299 - template [0,1] + world_size=8 -> [[0,1], [2,3], [4,5], [6,7]]

300 - template [0,2,4,6] + world_size=8 -> [[0,2,4,6], [1,3,5,7]]

301 - template [1,3,5,7] + world_size=8 -> [[0,2,4,6], [1,3,5,7]]

302

303 Algorithm:

304 1. Template normalization: convert any starting template to 0-based

305 2. Analyze pattern (intra-step, template span)

306 3. Iterate by blocks, generate valid sub-groups per block

307 4. Ensure each rank appears in exactly one group

308 """

309 # convert template to int list and sort (rank_list may come from numpy/tensor as float)

310 template = sorted([int(x) for x in list(template)])

311 world_size = int(world_size)

312 my_rank = int(my_rank)

313 template_len = len(template)

314

315 _report_template_stage(verbose, my_rank, f"Original Template = {template}, World size = {world_size}")

316

317 if template_len == 1:

318 return [[i] for i in range(world_size)]

319

320 if template_len < 2:

321 raise ValueError(f"Template must have at least 2 ranks, got {template}")

322

323 # 1. Template normalization: convert to 0-based template

324 normalized_template = [x - template[0] for x in template]

325 _report_template_stage(verbose, my_rank, f"Normalized Template = {normalized_template}")

326

327 # 2. Analyze normalized template core params

328 # intra-step: spacing between elements in template

329 intra_step = _validate_intra_step(normalized_template, template_len)

330 # block size: ranks per block; it is also the spacing between adjacent blocks

331 block_size = intra_step * template_len

332 template_span = normalized_template[-1] - normalized_template[0]

333 _report_template_stage(

334 verbose, my_rank,

335 f"Template analysis - intra_step={intra_step}, template_span={template_span}, block_size={block_size}"

336 )

337

338 # 3. Compute all valid block start positions

339 group_starts = _compute_group_starts(world_size, block_size, block_size)

340 _report_template_stage(verbose, my_rank, f"Possible block starts: {group_starts}")

341

342 # 4. Generate all valid sub-groups for each block

343 all_groups = _build_groups_for_blocks(

344 group_starts, block_size, template_span,

345 normalized_template, template_len, world_size

346 )

347

348 # 5. Validate: ensure each rank appears exactly once

349 all_ranks = [rank for group in all_groups for rank in group]

350 unique_ranks = set(all_ranks)

351 if len(all_ranks) != len(unique_ranks):

352 raise ValueError("Duplicate ranks found! Some ranks appear in multiple groups.")

353

354 # 6. Sort: ensure all processes generate groups in same order

355 all_groups.sort(key=lambda x: (x[0], x[1] if len(x) > 1 else 0))

356

357 _report_template_stage(

358 verbose, my_rank,

359 f"Generated {len(all_groups)} groups, covering {len(unique_ranks)} unique ranks\n"

360 f"Final group list: {all_groups}"

361 )

362

363 return all_groups

364

365

366def create_sub_groups(

367 rank_list: Union[list[int], tuple[int, ...]],

368 verbose: bool = False

369) -> dict[tuple, dist.ProcessGroup]:

370 """

371 Create sub-communication groups, supports template auto-expansion.

372

373 Args:

374 rank_list: One of:

375 1. Full group list, e.g. [[0,1], [2,3], [4,5], [6,7]]

376 2. Template group, e.g. [0,1] or [0,2], will auto-expand

377 verbose: Whether to print debug info

378

379 Returns:

380 Dict, key is tuple of group ranks, value is ProcessGroup

381 """

382 my_rank = dist.get_rank()

383 world_size = dist.get_world_size()

384 template = list(rank_list)

385 full_rank_list = generate_groups_from_template(template, world_size, my_rank, verbose=verbose)

386

387 if verbose:

388 print(f"Rank {my_rank}: Full rank list to create: {full_rank_list}")

389

390 # validate full group list format

391 for i, group in enumerate(full_rank_list):

392 if not isinstance(group, (list, tuple)):

393 raise ValueError(f"Group {i} must be a list or tuple, got {type(group)}")

394 if len(group) == 0:

395 raise ValueError(f"Group {i} is empty")

396 if len(group) != len(set(group)):

397 raise ValueError(f"Group {i} contains duplicate ranks")

398 for rank in group:

399 if not isinstance(rank, int):

400 raise ValueError(f"Rank must be integer, got {type(rank)} in group {i}")

401

402 # sort by first element to ensure all processes create groups in same order

403 sorted_groups = sorted(full_rank_list, key=lambda x: x[0])

404

405 if verbose:

406 print(f"Rank {my_rank}: Sorted groups for creation: {sorted_groups}")

407

408 # create all groups and collect groups current process belongs to

409 group_dict = {}

410 for group_ranks in sorted_groups:

411 # ensure ranks are ordered so each process passes same order

412 sorted_ranks = sorted(group_ranks)

413

414 if verbose:

415 print(f"Rank {my_rank}: Creating group with ranks {sorted_ranks}")

416

417 # key: all processes participate in each group creation

418 group = dist.new_group(ranks=sorted_ranks)

419 EXISTING_COMM_GROUPS[str(tuple(sorted_ranks))] = group

420

421 # only save when current process is in the group

422 if my_rank in sorted_ranks:

423 group_dict[tuple(sorted_ranks)] = group

424

425 if verbose:

426 print(f"Rank {my_rank}: Created {len(group_dict)} groups I belong to")

427

428 return group_dict

429

430

431# ---------------------------------------------------------------------------

432# Unsupported legacy redistribution hooks

433# ---------------------------------------------------------------------------

434

435def get_tensor_transform():

436 """Legacy tensor transform hook — not available on torch."""

437 raise NotImplementedError("Unsupported get_tensor_transform for torch platform")

438

439

440def construct_strided_slice(x, begin, end, stride):

441 """Legacy strided-slice hook — not available on torch."""

442 raise NotImplementedError("Unsupported construct_strided_slice for torch platform")

443

444

445# ---------------------------------------------------------------------------

446# Weight initialization

447# ---------------------------------------------------------------------------

448

449@contextmanager

450def init_on_device(device, include_buffers=False): # pylint: disable=W0621

451 """Monkey-patch ``nn.Module`` so that every parameter (and optionally every

452 buffer) is placed on *device* at registration time.

453

454 Args:

455 device (torch.device): Target device.

456 include_buffers (bool): Also redirect buffers to *device*.

457 """

458 orig_register_parameter = nn.Module.register_parameter

459 orig_register_buffer = nn.Module.register_buffer

460

461 # pylint: disable=W0212

462 def _register_parameter(module, name, param):

463 orig_register_parameter(module, name, param)

464 if param is None or param.device == device:

465 return

466 # Rebuild with data only, then restore instance attributes via __dict__:

467 # forwarding them to __new__ crashes subclasses with a narrow signature.

468 new_param = type(param)(param.to(device))

469 new_param.__dict__.update(param.__dict__)

470 new_param.requires_grad = param.requires_grad

471 module._parameters[name] = new_param

472

473 # pylint: disable=W0212

474 def _register_buffer(module, name, buffer, persistent=True):

475 orig_register_buffer(module, name, buffer, persistent=persistent)

476 if buffer is not None:

477 module._buffers[name] = module._buffers[name].to(device)

478

479 try:

480 nn.Module.register_parameter = _register_parameter

481 if include_buffers:

482 nn.Module.register_buffer = _register_buffer

483 yield

484 finally:

485 nn.Module.register_parameter = orig_register_parameter

486 if include_buffers:

487 nn.Module.register_buffer = orig_register_buffer

488

489

490__all__ = [

491 # Tensor/device/rng helpers.

492 "tensor_type_cast",

493 "device_count",

494 "device_type",

495 "device",

496 "get_rng_state",

497 "set_rng_state",

498 "get_tensor_transform",

499 "construct_strided_slice",

500 "init_on_device",

501 # Process-group construction.

502 "get_created_group",

503 "create_group",

504 "split_group",

505 "init_process_group",

506 "generate_groups_from_template",

507 "create_sub_groups",

508 # Re-exports kept for the debug tracer, which patches these attributes on

509 # this module; the canonical home is

510 # :mod:`hyper_parallel.core.utils.communication`.

511 "differentiable_all_gather_concat",

512 "differentiable_all_to_all",

513 "differentiable_all_to_all_single",

514 "differentiable_all_to_all_single_async",

515 "differentiable_all_reduce",

516 "differentiable_reduce_scatter",

517 "differentiable_variable_all_gather",

518 # Small shared helpers re-exported for callers that reach them here.

519 "get_device_handle",

520 "get_op_name",

521]

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / _utils.py: 35%

180 statements