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

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

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

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

37 Args:

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

40 Returns:

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

42 """

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

44 return 0

45 if group is None:

46 return torch.distributed.get_rank()

47 if hasattr(group, "rank"):

48 return group.rank()

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

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

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

55 Args:

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

58 Returns:

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

61 Raises:

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

63 """

64 try:

65 handle = getattr(torch, device_type)

66 except AttributeError as e:

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

68 return handle

71def get_world_size() -> int:

72 """Return the Torch distributed world size.

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

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

76 raises before ``init_process_group``.

77 """

78 return dist.get_world_size()

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

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

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

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

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

88def ensure_contiguous(x):

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

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

92 trace time, so the copy is emitted unconditionally.

93 """

94 if torch.compiler.is_compiling():

95 return x.contiguous()

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

97 return x.contiguous()

98 return x

100

101# ---------------------------------------------------------------------------

102# Process group helpers

103# ---------------------------------------------------------------------------

104

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

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

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

108 if group_key in EXISTING_COMM_GROUPS:

109 return EXISTING_COMM_GROUPS[group_key]

110 return None

111

112

113def split_group(parent_pg=None,

114 split_ranks: Optional[list] = None,

115 timeout=None,

116 pg_options=None,

117 group_desc=None,

118 ):

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

120

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

122 """

123 del parent_pg, timeout, group_desc

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

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

126

127 split_group_pg = None

128 for split_rank in split_ranks:

129 dist_group = get_created_group(split_rank)

130 if dist_group is None:

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

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

133 if dist.get_rank() in split_rank:

134 split_group_pg = dist_group

135

136 return split_group_pg

137

138

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

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

141 if not dist.is_initialized():

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

143

144

145# ---------------------------------------------------------------------------

146# Differentiable collectives

147# ---------------------------------------------------------------------------

148

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

150_OP_MAP = {

151 'sum': dist.ReduceOp.SUM,

152 'prod': dist.ReduceOp.PRODUCT,

153 'max': dist.ReduceOp.MAX,

154 'min': dist.ReduceOp.MIN,

155 # convert tensor elements to int32 and use MIN

156 'all': dist.ReduceOp.MIN,

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

158 'avg': dist.ReduceOp.SUM,

159}

160

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

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

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

164else:

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

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

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

168

169

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

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

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

173

174

175def _is_tracing_or_compiling() -> bool:

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

177

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

179 so we catch it and return False.

180 """

181 if torch.compiler.is_compiling():

182 return True

183 try:

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

185 except Exception: # pylint: disable=W0718

186 return False

187

188

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

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

191 op_str = {

192 dist.ReduceOp.SUM: 'sum',

193 dist.ReduceOp.PRODUCT: 'product',

194 dist.ReduceOp.MAX: 'max',

195 dist.ReduceOp.MIN: 'min',

196 }.get(reduce_op, 'sum')

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

198 op_str = 'avg'

199 return op_str

200

201

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

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

204

205 @staticmethod

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

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

208 del ctx

209 return tensor

210

211 @staticmethod

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

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

214 del ctx

215 return grad_output.contiguous()

216

217

218class _AsyncA2ALazyBwd(torch.autograd.Function):

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

220

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

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

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

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

225 stream and blocks Attention launches.

226

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

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

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

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

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

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

233 """

234

235 @staticmethod

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

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

238 ctx.input_splits = input_splits

239 ctx.output_splits = output_splits

240 ctx.group = group

241 # pylint: disable=C0415

242 from torch.distributed._functional_collectives import all_to_all_single

243 return all_to_all_single(

244 input_tensor, output_splits, input_splits, group,

245 )

246

247 @staticmethod

248 def backward(ctx, grad_output):

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

250 # pylint: disable=C0415

251 from torch.distributed._functional_collectives import all_to_all_single

252 grad_input = all_to_all_single(

253 grad_output, ctx.input_splits, ctx.output_splits, ctx.group,

254 )

255 return grad_input, None, None, None

256

257

258class _TorchP2PExchangeFunction(torch.autograd.Function):

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

260

261 @staticmethod

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

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

264 ctx.peer_rank = peer_rank

265 ctx.group = group

266 send_buf = tensor.contiguous()

267 recv_buf = torch.empty_like(send_buf)

268 reqs = dist.batch_isend_irecv([

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

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

271 ])

272 for req in reqs:

273 req.wait()

274 return recv_buf

275

276 @staticmethod

277 def backward(ctx, grad_output: Tensor):

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

279 send_buf = grad_output.contiguous()

280 recv_buf = torch.empty_like(send_buf)

281 reqs = dist.batch_isend_irecv([

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

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

284 ])

285 for req in reqs:

286 req.wait()

287 return recv_buf, None, None

288

289

290class _TorchDifferentiableVariableAllGather(torch.autograd.Function):

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

292

293 @staticmethod

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

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

296 if input_tensor.ndim == 0:

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

298 splits = tuple(output_splits)

299 if not splits:

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

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

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

303

304 group_rank = dist.get_rank(group=group)

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

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

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

308 raise ValueError(

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

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

311 f"output_splits={splits!r}"

312 )

313

314 input_tensor = input_tensor.contiguous()

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

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

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

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

319 else:

320 max_rows = max(splits)

321 if max_rows == 0:

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

323 else:

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

325 if input_tensor.shape[0] > 0:

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

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

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

329 gathered = [

330 output[:rows].contiguous()

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

332 ]

333

334 ctx.output_splits = splits

335 ctx.group = group

336 ctx.group_rank = group_rank

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

338

339 @staticmethod

340 def backward(ctx, grad_output):

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

342 output_rows = ctx.output_splits[ctx.group_rank]

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

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

345 return output, None, None

346

347 grad_output = grad_output.contiguous()

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

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

350 reduce_scatter_tensor_uneven(

351 output,

352 grad_output,

353 input_split_sizes=list(ctx.output_splits),

354 op=dist.ReduceOp.SUM,

355 group=ctx.group,

356 )

357 else:

358 reduced = grad_output.clone()

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

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

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

362 return output, None, None

363

364

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

366 rank_list=None) -> Tensor:

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

368

369 Args:

370 data: Local shard to gather.

371 group: Process group to gather over.

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

373 concat_dim: Dimension the gathered shards are concatenated along.

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

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

376 callers can request mesh order without changing the group.

377

378 Returns:

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

380 """

381 del concat_size

382 data = ensure_contiguous(data)

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

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

385 # collectives (fc.all_gather_single) for torch.compile / make_fx.

386 if rank_list is not None:

387 group_ranks = dist.get_process_group_ranks(group)

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

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

390 output = [

391 _TorchContiguousGrad.apply(tensor)

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

393 ]

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

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

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

397 if not _is_tracing_or_compiling():

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

399 output = [

400 _TorchContiguousGrad.apply(tensor)

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

402 ]

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

404 # Compile: use functional collective API for torch.compile compatibility.

405 # all_gather_single only reliably supports gather_dim=0; for other dims we

406 # gather on dim 0 then reshape to avoid _maybe_view_chunk_cat view() bugs

407 # in Dynamo fake tensor tracing.

408 group_size = dist.get_world_size(group)

409 import torch.distributed._functional_collectives as fc # pylint: disable=C0415

410 gathered = fc.all_gather_single(data, 0, group)

411 if hasattr(gathered, 'wait'):

412 gathered = gathered.wait()

413 if concat_dim == 0:

414 return gathered

415 # Reshape from [group_size * d0, d1, ...] to [d0 * group_size, d1, ...]

416 # by moving the group_size factor from dim 0 to concat_dim.

417 orig_shape = list(data.shape)

418 gathered_shape = [orig_shape[0] * group_size] + orig_shape[1:]

419 gathered = gathered.view(gathered_shape)

420 # Move group_size from dim 0 to concat_dim via unflatten + movedim + flatten

421 unflattened = gathered.unflatten(0, (group_size, orig_shape[0]))

422 moved = torch.movedim(unflattened, 0, concat_dim)

423 result_shape = list(orig_shape)

424 result_shape[concat_dim] *= group_size

425 return moved.reshape(result_shape).contiguous()

426

427

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

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

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

431 out_total = sum(output_splits)

432 output = torch.empty(

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

434 dtype=input_tensor.dtype, device=input_tensor.device,

435 )

436 return dist_func.all_to_all_single(

437 output, input_tensor,

438 output_split_sizes=output_splits,

439 input_split_sizes=input_splits,

440 group=group,

441 )

442

443

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

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

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

447

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

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

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

451 backward needs the same treatment.

452

453 Args:

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

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

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

457 group: Process group.

458

459 Returns:

460 ``AsyncCollectiveTensor`` of shape

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

462 """

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

464

465

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

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

468 input_data = ensure_contiguous(input_data)

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

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

471

472

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

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

475 data = ensure_contiguous(data)

476 reduce_op = resolve_reduce_op(op)

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

478 if not _is_tracing_or_compiling():

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

480 # Compile: use functional collective API for torch.compile compatibility.

481 # fc.all_reduce expects a string reduceOp and a ProcessGroup.

482 import torch.distributed._functional_collectives as fc # pylint: disable=C0415

483 return fc.all_reduce(data, _reduce_op_to_str(reduce_op), group)

484

485

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

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

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

489 data = ensure_contiguous(data)

490 reduce_op = resolve_reduce_op(op)

491

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

493 if not _is_tracing_or_compiling():

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

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

496 output_tensor = dist_func.reduce_scatter(

497 output_tensor, input_tuple, op=reduce_op, group=group

498 )

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

500 if op == 'avg':

501 output_tensor = output_tensor / dev_num

502 return output_tensor

503

504 # Compile: use functional collective API for torch.compile compatibility.

505 # fc.reduce_scatter_single expects a string reduceOp, scatter_dim, and group.

506 import torch.distributed._functional_collectives as fc # pylint: disable=C0415

507 result = fc.reduce_scatter_single(data, _reduce_op_to_str(reduce_op), axis, group)

508 if hasattr(result, 'wait'):

509 result = result.wait()

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

511 if op == 'avg':

512 result = result / dev_num

513 return result

514

515

516def differentiable_variable_all_gather(

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

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

519 return _TorchDifferentiableVariableAllGather.apply(

520 input_tensor, tuple(output_splits), group

521 )

522

523

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

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

526

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

528 """

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

530 wait_tensor(tensor)

531 return tensor

532

533

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

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

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

537 return tensor

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

539

540

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

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

543

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

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

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

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

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

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

550 """

551 output = torch.empty(

552 [input_tensor.shape[0]],

553 dtype=input_tensor.dtype, device=input_tensor.device,

554 )

555 handle = dist.all_to_all_single(

556 output, input_tensor, group=group, async_op=True,

557 )

558 if handle is not None:

559 handle.wait()

560 return output

561

562

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

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

565

566 Same async-with-explicit-wait rationale as

567 :func:`exchange_splits_via_all_to_all`.

568 """

569 output = torch.empty(

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

571 dtype=input_tensor.dtype,

572 device=input_tensor.device,

573 )

574 handle = dist.all_gather_into_tensor(

575 output, input_tensor, group=group, async_op=True,

576 )

577 if handle is not None:

578 handle.wait()

579 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: 45%

241 statements