Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/core/utils/communication.py: 55%

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"""Communication helpers shared by the ``core`` modules.

17Everything here is collective-related and therefore Torch-only:

18differentiable collective wrappers (plus the autograd ``Function``s backing

19them), the process-group cache, and the P2P exchange. The synchronous

20collectives each caller needs are thin enough to stay local; what is shared

21here is the part with nontrivial autograd/stream-lifetime semantics, which

22must stay identical across dtensor, expert_parallel and context_parallel.

23"""

24# pylint: disable=C9006,C9007

25from typing import Any, Optional, Sequence, Union

27import torch

28import torch.distributed as dist

29import torch.distributed.nn.functional as dist_func

30from torch import Tensor

31from torch._C._distributed_c10d import ProcessGroup

32from torch.distributed import _functional_collectives

35def get_group_local_rank(group: Optional[ProcessGroup] = None) -> int:

36 """Return the current rank's index within ``group``, or 0 before init.

38 Args:

39 group: The group to index into. ``None`` means the default group.

41 Returns:

42 int: The rank of the current process inside ``group``.

43 """

44 if not torch.distributed.is_available() or not torch.distributed.is_initialized():

45 return 0

46 if group is None:

47 return torch.distributed.get_rank()

48 if hasattr(group, "rank"):

49 return group.rank()

50 return torch.distributed.get_group_rank(group, torch.distributed.get_rank())

53def get_device_handle(device_type: str = "npu"):

54 """Return the ``torch`` device module for ``device_type``.

56 Args:

57 device_type: Device backend name, e.g. ``"npu"`` or ``"cuda"``.

59 Returns:

60 The matching ``torch.<device_type>`` module.

62 Raises:

63 RuntimeError: If torch exposes no module for ``device_type``.

64 """

65 try:

66 handle = getattr(torch, device_type)

67 except AttributeError as e:

68 raise RuntimeError(f"Failed to resolve device handle: 'torch.{device_type}'.") from e

69 return handle

72def get_world_size() -> int:

73 """Return the Torch distributed world size.

75 Kept as a thin wrapper so single-process call sites and white-box unit

76 tests share one patchable seam; ``torch.distributed.get_world_size`` itself

77 raises before ``init_process_group``.

78 """

79 return dist.get_world_size()

82# Process-group cache keyed by the string form of the sorted rank tuple, e.g.

83# ``"(0, 1, 2, 3)"``. Lives here rather than in a module-specific ``utils`` so

84# that every sub-group creator in ``core`` shares one map. Values are

85# framework-native groups, hence ``Any``.

86EXISTING_COMM_GROUPS: dict[str, Any] = {}

89def ensure_contiguous(x):

90 """Return a contiguous copy of *x* if not already contiguous.

92 Under ``torch.compile`` the contiguity of a graph input is not knowable at

93 trace time, so the copy is emitted unconditionally.

94 """

95 if torch.compiler.is_compiling():

96 return x.contiguous()

97 if not x.is_contiguous() or x.storage_offset() != 0:

98 return x.contiguous()

99 return x

100

101

102# ---------------------------------------------------------------------------

103# Process group helpers

104# ---------------------------------------------------------------------------

105

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

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

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

109 if group_key in EXISTING_COMM_GROUPS:

110 return EXISTING_COMM_GROUPS[group_key]

111 return None

112

113

114def split_group(parent_pg=None,

115 split_ranks: Optional[list] = None,

116 timeout=None,

117 pg_options=None,

118 group_desc=None,

119 ):

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

121

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

123 """

124 del parent_pg, timeout, group_desc

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

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

127

128 split_group_pg = None

129 for split_rank in split_ranks:

130 dist_group = get_created_group(split_rank)

131 if dist_group is None:

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

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

134 if dist.get_rank() in split_rank:

135 split_group_pg = dist_group

136

137 return split_group_pg

138

139

140def init_process_group(*args, **kwargs) -> None:

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

142 if not dist.is_initialized():

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

144

145

146# ---------------------------------------------------------------------------

147# Differentiable collectives

148# ---------------------------------------------------------------------------

149

150# Mapping from string op names to torch.distributed.ReduceOp

151_OP_MAP = {

152 'sum': dist.ReduceOp.SUM,

153 'prod': dist.ReduceOp.PRODUCT,

154 'max': dist.ReduceOp.MAX,

155 'min': dist.ReduceOp.MIN,

156 # convert tensor elements to int32 and use MIN

157 'all': dist.ReduceOp.MIN,

158 # 'avg' is typically handled by SUM followed by division in current implementation logic

159 'avg': dist.ReduceOp.SUM,

160}

161

162# Try to add AVG for 'mean' if supported by current torch version

163if hasattr(dist.ReduceOp, "AVG"):

164 _OP_MAP['mean'] = dist.ReduceOp.AVG

165else:

166 # Fallback for older torch versions if necessary, though this might require manual division upstream

167 # Assuming standard behavior where 'mean' implies native AVG support or upstream handling

168 _OP_MAP['mean'] = dist.ReduceOp.SUM

169

170

171def resolve_reduce_op(op: Union[str, Any]) -> Any:

172 """Resolve a string op name (or pass through an already-resolved ``ReduceOp``)."""

173 return _OP_MAP.get(op, dist.ReduceOp.SUM) if isinstance(op, str) else op

174

175

176def _is_tracing_or_compiling() -> bool:

177 """Return True if we are inside a make_fx trace or torch.compile.

178

179 ``TracingContext.get()`` raises ``RuntimeError`` when no trace is active,

180 so we catch it and return False.

181 """

182 if torch.compiler.is_compiling():

183 return True

184 try:

185 return torch._guards.TracingContext.get() is not None # pylint: disable=protected-access

186 except Exception: # pylint: disable=W0718

187 return False

188

189

190def _reduce_op_to_str(reduce_op: Any) -> str:

191 """Map a ``dist.ReduceOp`` to the string name ``_functional_collectives`` expects."""

192 op_str = {

193 dist.ReduceOp.SUM: 'sum',

194 dist.ReduceOp.PRODUCT: 'product',

195 dist.ReduceOp.MAX: 'max',

196 dist.ReduceOp.MIN: 'min',

197 }.get(reduce_op, 'sum')

198 if hasattr(dist.ReduceOp, 'AVG') and reduce_op == dist.ReduceOp.AVG:

199 op_str = 'avg'

200 return op_str

201

202

203class _TorchContiguousGrad(torch.autograd.Function): # pylint: disable=abstract-method

204 """Autograd identity that materializes gradients before upstream collectives."""

205

206 @staticmethod

207 def forward(ctx: Any, tensor: Tensor) -> Tensor: # pylint: disable=arguments-differ

208 """Return the input unchanged in the forward pass."""

209 del ctx

210 return tensor

211

212 @staticmethod

213 def backward(ctx: Any, grad_output: Tensor) -> Tensor: # pylint: disable=arguments-differ

214 """Return a contiguous gradient to the preceding autograd node."""

215 del ctx

216 return grad_output.contiguous()

217

218

219class _AsyncA2ALazyBwd(torch.autograd.Function):

220 """All-to-all whose forward AND backward return ``AsyncCollectiveTensor``.

221

222 PyTorch's stock ``all_to_all_single_autograd`` calls ``wait_tensor`` in

223 its backward eagerly, and the autograd engine binds backward stream

224 context to the forward stream — so even if the BWD thread is wrapped

225 in a side-stream context, that wait still lands on the FWD main

226 stream and blocks Attention launches.

227

228 This Function bypasses the engine's binding by calling the

229 non-autograd functional op in both directions and returning ACT.

230 The wait is deferred to the next consumer's first non-view access

231 (e.g. the indexing backward of ``_unpermute``), giving the FWD

232 thread a small Python window to enqueue its Attention kernels onto

233 the main stream **before** the wait lands there.

234 """

235

236 @staticmethod

237 def forward(ctx, input_tensor, output_splits, input_splits, group): # pylint: disable=arguments-differ

238 """Perform the forward all-to-all single collective, saving splits and group for backward."""

239 ctx.input_splits = input_splits

240 ctx.output_splits = output_splits

241 ctx.group = group

242 # pylint: disable=C0415

243 from torch.distributed._functional_collectives import all_to_all_single

244 return all_to_all_single(

245 input_tensor, output_splits, input_splits, group,

246 )

247

248 @staticmethod

249 def backward(ctx, grad_output):

250 """Compute the backward pass by performing the inverse all-to-all with swapped splits."""

251 # pylint: disable=C0415

252 from torch.distributed._functional_collectives import all_to_all_single

253 grad_input = all_to_all_single(

254 grad_output, ctx.input_splits, ctx.output_splits, ctx.group,

255 )

256 return grad_input, None, None, None

257

258

259class _TorchP2PExchangeFunction(torch.autograd.Function):

260 """Symmetric bidirectional P2P: send local tensor to peer, receive peer's tensor."""

261

262 @staticmethod

263 def forward(ctx, tensor: Tensor, peer_rank: int, group) -> Tensor: # pylint: disable=arguments-differ

264 """Perform symmetric bidirectional P2P exchange with ``peer_rank``."""

265 ctx.peer_rank = peer_rank

266 ctx.group = group

267 send_buf = tensor.contiguous()

268 recv_buf = torch.empty_like(send_buf)

269 reqs = dist.batch_isend_irecv([

270 dist.P2POp(dist.isend, send_buf, peer_rank, group),

271 dist.P2POp(dist.irecv, recv_buf, peer_rank, group),

272 ])

273 for req in reqs:

274 req.wait()

275 return recv_buf

276

277 @staticmethod

278 def backward(ctx, grad_output: Tensor):

279 """Perform symmetric P2P exchange for the backward gradient pass."""

280 send_buf = grad_output.contiguous()

281 recv_buf = torch.empty_like(send_buf)

282 reqs = dist.batch_isend_irecv([

283 dist.P2POp(dist.isend, send_buf, ctx.peer_rank, ctx.group),

284 dist.P2POp(dist.irecv, recv_buf, ctx.peer_rank, ctx.group),

285 ])

286 for req in reqs:

287 req.wait()

288 return recv_buf, None, None

289

290

291class _TorchDifferentiableVariableAllGather(torch.autograd.Function):

292 """Variable dim-zero all-gather with an uneven reduce-scatter backward."""

293

294 @staticmethod

295 def forward(ctx, input_tensor, output_splits, group): # pylint: disable=arguments-differ

296 """Gather each rank's true row count without replicating inputs for A2A."""

297 if input_tensor.ndim == 0:

298 raise ValueError("variable all-gather input must have at least one dimension")

299 splits = tuple(output_splits)

300 if not splits:

301 raise ValueError("output_splits must contain at least one group rank")

302 if any(not isinstance(rows, int) or isinstance(rows, bool) or rows < 0 for rows in splits):

303 raise ValueError(f"output_splits must contain non-negative integers, got {splits!r}")

304

305 group_rank = dist.get_rank(group=group)

306 if group_rank < 0 or group_rank >= len(splits):

307 raise ValueError(f"group rank must be in [0, {len(splits)}), got {group_rank}")

308 if input_tensor.shape[0] != splits[group_rank]:

309 raise ValueError(

310 "variable all-gather local rows must match output_splits at the group rank, "

311 f"got local_rows={input_tensor.shape[0]}, group_rank={group_rank}, "

312 f"output_splits={splits!r}"

313 )

314

315 input_tensor = input_tensor.contiguous()

316 feature_shape = tuple(input_tensor.shape[1:])

317 if input_tensor.device.type == "npu":

318 gathered = [input_tensor.new_empty((rows, *feature_shape)) for rows in splits]

319 dist.all_gather(gathered, input_tensor, group=group)

320 else:

321 max_rows = max(splits)

322 if max_rows == 0:

323 gathered = [input_tensor.new_empty((0, *feature_shape)) for _ in splits]

324 else:

325 padded = input_tensor.new_zeros((max_rows, *feature_shape))

326 if input_tensor.shape[0] > 0:

327 padded[:input_tensor.shape[0]].copy_(input_tensor)

328 padded_outputs = [torch.empty_like(padded) for _ in splits]

329 dist.all_gather(padded_outputs, padded, group=group)

330 gathered = [

331 output[:rows].contiguous()

332 for output, rows in zip(padded_outputs, splits)

333 ]

334

335 ctx.output_splits = splits

336 ctx.group = group

337 ctx.group_rank = group_rank

338 return torch.cat(gathered, dim=0)

339

340 @staticmethod

341 def backward(ctx, grad_output):

342 """Sum replicated output gradients and return this rank's uneven shard."""

343 output_rows = ctx.output_splits[ctx.group_rank]

344 output = grad_output.new_empty((output_rows, *grad_output.shape[1:]))

345 if sum(ctx.output_splits) == 0:

346 return output, None, None

347

348 grad_output = grad_output.contiguous()

349 if grad_output.device.type == "npu":

350 from torch_npu.distributed import reduce_scatter_tensor_uneven # pylint: disable=C0415

351 reduce_scatter_tensor_uneven(

352 output,

353 grad_output,

354 input_split_sizes=list(ctx.output_splits),

355 op=dist.ReduceOp.SUM,

356 group=ctx.group,

357 )

358 else:

359 reduced = grad_output.clone()

360 dist.all_reduce(reduced, op=dist.ReduceOp.SUM, group=ctx.group)

361 start = sum(ctx.output_splits[:ctx.group_rank])

362 output.copy_(reduced.narrow(0, start, output_rows))

363 return output, None, None

364

365

366def differentiable_all_gather_concat(data: Tensor, group, concat_size: int, concat_dim: int,

367 rank_list=None) -> Tensor:

368 """Autograd-aware all-gather whose results are concatenated along ``concat_dim``.

369

370 Args:

371 data: Local shard to gather.

372 group: Process group to gather over.

373 concat_size: Unused; kept for call-site compatibility.

374 concat_dim: Dimension the gathered shards are concatenated along.

375 rank_list: Optional global-rank order for the gathered shards. When it

376 differs from the group's own rank order the shards are permuted, so

377 callers can request mesh order without changing the group.

378

379 Returns:

380 The gathered shards concatenated along ``concat_dim``.

381 """

382 del concat_size

383 data = ensure_contiguous(data)

384 # Eager path: use legacy autograd API (dist_func.all_gather) to preserve

385 # original numerical behavior. Compile/graph-trace path: use functional

386 # collectives for torch.compile / make_fx.

387 if rank_list is not None:

388 group_ranks = dist.get_process_group_ranks(group)

389 if tuple(rank_list) != tuple(group_ranks):

390 # Rank subset gathering (not common path) — same for both modes.

391 output = [

392 _TorchContiguousGrad.apply(tensor)

393 for tensor in dist_func.all_gather(data, group=group)

394 ]

395 rank_to_idx = {int(rank): idx for idx, rank in enumerate(group_ranks)}

396 output = [output[rank_to_idx[int(rank)]] for rank in rank_list]

397 return torch.cat(output, dim=concat_dim)

398 if not _is_tracing_or_compiling():

399 # Eager: original implementation using torch.distributed.nn.functional.

400 output = [

401 _TorchContiguousGrad.apply(tensor)

402 for tensor in dist_func.all_gather(data, group=group)

403 ]

404 return torch.cat(output, dim=concat_dim)

405 # Torch 2.14 functional collectives provide autograd and tracing-aware waits.

406 if not -data.ndim <= concat_dim < data.ndim:

407 raise IndexError("All-gather concatenation dimension is out of range")

408 concat_dim %= data.ndim

409 return _functional_collectives.all_gather_single(data, concat_dim, group)

410

411

412def differentiable_all_to_all_single(input_tensor: Tensor, input_splits: Sequence[int],

413 output_splits: Sequence[int], group) -> Tensor:

414 """Variable-split all-to-all with autograd support for EP token dispatch/combine."""

415 out_total = sum(output_splits)

416 output = torch.empty(

417 out_total, *input_tensor.shape[1:],

418 dtype=input_tensor.dtype, device=input_tensor.device,

419 )

420 return dist_func.all_to_all_single(

421 output, input_tensor,

422 output_split_sizes=output_splits,

423 input_split_sizes=input_splits,

424 group=group,

425 )

426

427

428def differentiable_all_to_all_single_async(input_tensor: Tensor, input_splits: Sequence[int],

429 output_splits: Sequence[int], group) -> Tensor:

430 """Truly-async variant of :func:`differentiable_all_to_all_single`.

431

432 Both forward AND backward return ``AsyncCollectiveTensor``, so the

433 ``wait_tensor`` op is queued lazily — only when a downstream kernel

434 actually reads the result. See :class:`_AsyncA2ALazyBwd` for why the

435 backward needs the same treatment.

436

437 Args:

438 input_tensor: Input tensor, split along dim 0 by ``input_splits``.

439 input_splits: ``list[int]`` — rows sent to each rank.

440 output_splits: ``list[int]`` — rows received from each rank.

441 group: Process group.

442

443 Returns:

444 ``AsyncCollectiveTensor`` of shape

445 ``[sum(output_splits), *input_tensor.shape[1:]]``.

446 """

447 return _AsyncA2ALazyBwd.apply(input_tensor, output_splits, input_splits, group)

448

449

450def differentiable_all_to_all(input_data: Tensor, output_shape: Sequence[int], group) -> Tensor:

451 """Autograd-aware all-to-all producing a tensor of ``output_shape``."""

452 input_data = ensure_contiguous(input_data)

453 output_tensor = torch.empty(output_shape, device=input_data.device, dtype=input_data.dtype)

454 return dist_func.all_to_all_single(output_tensor, input_data, group=group)

455

456

457def differentiable_all_reduce(data: Tensor, op: Union[str, Any], group) -> Tensor:

458 """Autograd-aware all-reduce with string or ``ReduceOp`` *op*."""

459 data = ensure_contiguous(data)

460 reduce_op = resolve_reduce_op(op)

461 # Eager: use legacy autograd API. Compile/graph-trace: use functional collectives.

462 if not _is_tracing_or_compiling():

463 return dist_func.all_reduce(data, op=reduce_op, group=group)

464 return _functional_collectives.all_reduce(data, _reduce_op_to_str(reduce_op), group)

465

466

467def differentiable_reduce_scatter(data: Tensor, dev_num: int, axis: int,

468 op: Union[str, Any], group) -> Tensor:

469 """Autograd-aware reduce-scatter splitting ``axis`` into ``dev_num`` parts."""

470 data = ensure_contiguous(data)

471 reduce_op = resolve_reduce_op(op)

472

473 # Eager: use legacy autograd API. Compile/graph-trace: use functional collectives.

474 if not _is_tracing_or_compiling():

475 input_tuple = torch.chunk(data, dev_num, dim=axis)

476 output_tensor = torch.empty(input_tuple[0].shape, device=data.device, dtype=data.dtype)

477 output_tensor = dist_func.reduce_scatter(

478 output_tensor, input_tuple, op=reduce_op, group=group

479 )

480 # 'avg' maps to SUM in _OP_MAP, so the division stays manual.

481 if op == 'avg':

482 output_tensor = output_tensor / dev_num

483 return output_tensor

484

485 result = _functional_collectives.reduce_scatter_single(data, _reduce_op_to_str(reduce_op), axis, group)

486 # 'avg' maps to SUM in _OP_MAP, so the division stays manual.

487 if op == 'avg':

488 result = result / dev_num

489 return result

490

491

492def differentiable_variable_all_gather(

493 input_tensor: Tensor, output_splits: Sequence[int], group) -> Tensor:

494 """Gather variable dim-zero shards on HCCL or Gloo with autograd support."""

495 return _TorchDifferentiableVariableAllGather.apply(

496 input_tensor, tuple(output_splits), group

497 )

498

499

500def wait_async_tensor(tensor: Tensor) -> Tensor:

501 """Wait for an async collective tensor to become materialised.

502

503 Idempotent — calling on an already-waited tensor is a no-op.

504 """

505 from torch.distributed._functional_collectives import wait_tensor # pylint: disable=C0415

506 wait_tensor(tensor)

507 return tensor

508

509

510def p2p_exchange(tensor: Tensor, peer_rank: int, group=None) -> Tensor:

511 """Symmetric bidirectional P2P exchange with ``peer_rank``."""

512 if peer_rank == dist.get_rank(group):

513 return tensor

514 return _TorchP2PExchangeFunction.apply(tensor, peer_rank, group)

515

516

517def exchange_splits_via_all_to_all(input_tensor: Tensor, group) -> Tensor:

518 """All-to-all a per-rank split vector, returning the received counts.

519

520 Used by the token dispatchers to learn how many tokens every peer sends

521 them; the exchange carries no gradient, so plain ``dist`` collectives are

522 used. Runs asynchronously with an explicit ``handle.wait()`` rather than

523 ``async_op=False``: the implicit cross-stream sync is NCCL-only, and on

524 HCCL the compute stream may read the output before the collective write is

525 visible, producing garbage values that blow up downstream allocations.

526 """

527 output = torch.empty(

528 [input_tensor.shape[0]],

529 dtype=input_tensor.dtype, device=input_tensor.device,

530 )

531 handle = dist.all_to_all_single(

532 output, input_tensor, group=group, async_op=True,

533 )

534 if handle is not None:

535 handle.wait()

536 return output

537

538

539def gather_counts_via_all_gather(input_tensor: Tensor, group, world_size: int) -> Tensor:

540 """All-gather a per-rank count vector into an ``[world_size, *shape]`` view.

541

542 Same async-with-explicit-wait rationale as

543 :func:`exchange_splits_via_all_to_all`.

544 """

545 output = torch.empty(

546 [world_size * input_tensor.shape[0]],

547 dtype=input_tensor.dtype,

548 device=input_tensor.device,

549 )

550 handle = dist.all_gather_into_tensor(

551 output, input_tensor, group=group, async_op=True,

552 )

553 if handle is not None:

554 handle.wait()

555 return output.view(world_size, input_tensor.shape[0])

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / utils / communication.py: 55%

227 statements