Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / context_parallel / context_parallel.py: 87%

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"""Unified Context Parallel: Pure Ulysses, Pure Colossal AI, and Hybrid CP.""" 

16from functools import partial 

17from typing import Optional 

18 

19from hyper_parallel.core.dtensor.device_mesh import DeviceMesh 

20from hyper_parallel.core.dtensor.dtensor import DTensor 

21from hyper_parallel.core.tensor_parallel.style import ParallelStyle 

22from hyper_parallel.core.dtensor.placement_types import Shard, Replicate, StridedShard 

23from hyper_parallel.platform import get_platform 

24 

25platform = get_platform() 

26Module = platform.Module 

27Tensor = platform.Tensor 

28 

29_OUTPUT_LAYOUT_STACK_ATTR = "_context_parallel_output_layout_stack" 

30_OUTPUT_LOCAL = "local" 

31_OUTPUT_CP = "cp" 

32_OUTPUT_NON_CP = "non_cp" 

33 

34 

35# --------------------------------------------------------------------------- 

36# DTensor boundary helpers 

37# --------------------------------------------------------------------------- 

38 

39def _same_mesh_rank_list(lhs: DeviceMesh, rhs: DeviceMesh) -> bool: 

40 """Return whether two meshes describe the same participant ranks.""" 

41 return tuple(lhs.rank_list) == tuple(rhs.rank_list) 

42 

43 

44def _same_mesh(lhs: DeviceMesh, rhs: DeviceMesh) -> bool: 

45 """Return whether two meshes represent the same named layout.""" 

46 return lhs.to_hash() == rhs.to_hash() 

47 

48 

49def _is_foreign_dtensor(tensor: Tensor, submesh: DeviceMesh) -> bool: 

50 """Return whether *tensor* is a DTensor whose mesh differs from *submesh*.""" 

51 return isinstance(tensor, DTensor) and not _same_mesh(tensor.device_mesh, submesh) 

52 

53 

54def _is_cp_composed_dtensor(tensor: Tensor, cp_mesh: DeviceMesh) -> bool: 

55 """Return whether *tensor* already lives on a ``CP + non-CP`` mesh.""" 

56 if not isinstance(tensor, DTensor): 

57 return False 

58 mesh = tensor.device_mesh 

59 if mesh.ndim <= cp_mesh.ndim or not mesh.mesh_dim_names: 

60 return False 

61 cp_names = tuple(cp_mesh.mesh_dim_names or ()) 

62 if len(cp_names) != cp_mesh.ndim or not all(name in mesh.mesh_dim_names for name in cp_names): 

63 return False 

64 try: 

65 return _same_mesh_rank_list(mesh[cp_names], cp_mesh) 

66 except (KeyError, RuntimeError, ValueError): 

67 return False 

68 

69 

70def _compose_cp_non_cp_mesh(cp_mesh: DeviceMesh, non_cp_mesh: DeviceMesh) -> DeviceMesh: 

71 """Build a ``CP + non-CP`` mesh used while executing CP attention.""" 

72 meshes = (cp_mesh, non_cp_mesh) 

73 root_mesh = cp_mesh.root_mesh or cp_mesh 

74 root_dim_names = tuple(root_mesh.mesh_dim_names or ()) 

75 if root_dim_names and all( 

76 mesh.mesh_dim_names and all(name in root_dim_names for name in mesh.mesh_dim_names) 

77 for mesh in meshes 

78 ): 

79 meshes = tuple( 

80 sorted(meshes, key=lambda mesh: root_dim_names.index(mesh.mesh_dim_names[0])) 

81 ) 

82 return DeviceMesh.concatenate(meshes) 

83 

84 

85def _compose_cp_non_cp_placements( 

86 composed_mesh: DeviceMesh, 

87 cp_mesh: DeviceMesh, 

88 cp_placements, 

89 non_cp_mesh: DeviceMesh, 

90 non_cp_placements, 

91) -> tuple: 

92 """Return placements aligned to the actual composed mesh dimension order.""" 

93 placement_by_name = {} 

94 for name, placement in zip(cp_mesh.mesh_dim_names or (), cp_placements): 

95 placement_by_name[name] = placement 

96 for name, placement in zip(non_cp_mesh.mesh_dim_names or (), non_cp_placements): 

97 placement_by_name[name] = placement 

98 placements = [placement_by_name[name] for name in composed_mesh.mesh_dim_names] 

99 for mesh_idx, placement in enumerate(placements): 

100 if not placement.is_shard(): 

101 continue 

102 split_factor = 1 

103 for right_mesh_idx in range(mesh_idx + 1, len(placements)): 

104 if placements[right_mesh_idx].is_shard(placement.dim): 

105 split_factor *= composed_mesh.mesh_shape[right_mesh_idx] 

106 if split_factor > 1: 

107 placements[mesh_idx] = StridedShard(placement.dim, split_factor) 

108 return tuple(placements) 

109 

110 

111def _compose_cp_mesh(cp_mesh: DeviceMesh, non_cp_mesh: DeviceMesh) -> DeviceMesh: 

112 """Compatibility wrapper for CP + non-CP mesh composition.""" 

113 return _compose_cp_non_cp_mesh(cp_mesh, non_cp_mesh) 

114 

115 

116def _composed_placements( 

117 composed_mesh: DeviceMesh, 

118 cp_mesh: DeviceMesh, 

119 cp_placements, 

120 non_cp_mesh: DeviceMesh, 

121 non_cp_placements, 

122) -> tuple: 

123 """Compatibility wrapper for CP + non-CP placements composition.""" 

124 return _compose_cp_non_cp_placements( 

125 composed_mesh, 

126 cp_mesh, 

127 cp_placements, 

128 non_cp_mesh, 

129 non_cp_placements, 

130 ) 

131 

132 

133def _non_cp_placements_from_composed( 

134 tensor: "DTensor", 

135 cp_mesh: DeviceMesh, 

136) -> tuple: 

137 """Extract non-CP placements from a composed DTensor by mesh dimension name.""" 

138 cp_names = set(cp_mesh.mesh_dim_names or ()) 

139 return tuple( 

140 placement 

141 for name, placement in zip(tensor.device_mesh.mesh_dim_names, tensor.placements) 

142 if name not in cp_names 

143 ) 

144 

145 

146def _cp_mesh_from_composed(composed_mesh: DeviceMesh, non_cp_mesh: DeviceMesh) -> DeviceMesh: 

147 """Return the CP portion of a composed mesh.""" 

148 non_cp_names = set(non_cp_mesh.mesh_dim_names or ()) 

149 cp_names = tuple(name for name in composed_mesh.mesh_dim_names if name not in non_cp_names) 

150 return composed_mesh[cp_names] 

151 

152 

153def _split_cp_composed_layout(tensor: "DTensor", cp_mesh: DeviceMesh): 

154 """Return ``(non_cp_mesh, placements, composed_mesh)`` for a composed DTensor.""" 

155 if not _is_cp_composed_dtensor(tensor, cp_mesh): 

156 return None 

157 mesh = tensor.device_mesh 

158 dim_names = mesh.mesh_dim_names 

159 if dim_names is None: 

160 raise ValueError("dim_names must not be None") 

161 cp_names = set(cp_mesh.mesh_dim_names or ()) 

162 non_cp_names = tuple(name for name in dim_names if name not in cp_names) 

163 if len(non_cp_names) == 0: 

164 return None 

165 return mesh[non_cp_names], _non_cp_placements_from_composed(tensor, cp_mesh), mesh 

166 

167 

168def _non_cp_dtensor_layout(tensor: Tensor, cp_mesh: DeviceMesh, seq_dim: int): 

169 """Return metadata for dropping CP and keeping the incoming non-CP layout.""" 

170 del seq_dim 

171 if not isinstance(tensor, DTensor): 

172 return None 

173 composed_layout = _split_cp_composed_layout(tensor, cp_mesh) 

174 if composed_layout is not None: 

175 return composed_layout 

176 if not _is_foreign_dtensor(tensor, cp_mesh): 

177 return None 

178 non_cp_mesh = tensor.device_mesh 

179 return non_cp_mesh, tuple(tensor.placements), _compose_cp_non_cp_mesh(cp_mesh, non_cp_mesh) 

180 

181 

182def _localize_foreign_dtensor(tensor: Tensor, cp_mesh: DeviceMesh, seq_dim: int): 

183 """Convert a safe non-CP DTensor to its local tensor at the CP boundary.""" 

184 del seq_dim 

185 if _is_cp_composed_dtensor(tensor, cp_mesh) or not _is_foreign_dtensor(tensor, cp_mesh): 

186 return tensor 

187 return tensor.to_local() 

188 

189 

190def _cp_boundary_dim_size(tensor: Tensor, submesh: DeviceMesh, dim: int) -> int: 

191 """Return the dimension size CP should validate at its local boundary.""" 

192 if _is_foreign_dtensor(tensor, submesh): 

193 return tensor.to_local().shape[dim] 

194 return tensor.shape[dim] 

195 

196 

197def _to_cp_dtensor( 

198 tensor: Tensor, 

199 cp_mesh: DeviceMesh, 

200 src_placements, 

201 dst_placements, 

202 seq_dim: int, 

203) -> "DTensor": 

204 """Wrap a local/non-CP tensor as a CP DTensor and redistribute it.""" 

205 if _is_cp_composed_dtensor(tensor, cp_mesh): 

206 non_cp_mesh, non_cp_placements, composed_mesh = _split_cp_composed_layout(tensor, cp_mesh) 

207 return tensor.redistribute( 

208 composed_mesh, 

209 _composed_placements(composed_mesh, cp_mesh, dst_placements, non_cp_mesh, non_cp_placements), 

210 ) 

211 layout = _non_cp_dtensor_layout(tensor, cp_mesh, seq_dim) 

212 if layout is not None: 

213 non_cp_mesh, non_cp_placements, composed_mesh = layout 

214 src_composed_placements = _composed_placements( 

215 composed_mesh, cp_mesh, src_placements, non_cp_mesh, non_cp_placements 

216 ) 

217 dst_composed_placements = _composed_placements( 

218 composed_mesh, cp_mesh, dst_placements, non_cp_mesh, non_cp_placements 

219 ) 

220 return DTensor.from_local( 

221 tensor.to_local(), 

222 composed_mesh, 

223 src_composed_placements, 

224 ).redistribute(composed_mesh, dst_composed_placements) 

225 tensor = _localize_foreign_dtensor(tensor, cp_mesh, seq_dim) 

226 if isinstance(tensor, DTensor): 

227 return tensor.redistribute(cp_mesh, dst_placements) 

228 cp_tensor = DTensor.from_local(tensor, cp_mesh, src_placements) 

229 if not hasattr(cp_tensor, "redistribute"): 

230 return cp_tensor 

231 return cp_tensor.redistribute(cp_mesh, dst_placements) 

232 

233 

234def _output_layout_from_q(tensor: Tensor, cp_mesh: DeviceMesh, seq_dim: int): 

235 """Return the output-layout policy implied by the Q input.""" 

236 if not isinstance(tensor, DTensor): 

237 return _OUTPUT_LOCAL, None 

238 non_cp_layout = _non_cp_dtensor_layout(tensor, cp_mesh, seq_dim) 

239 if non_cp_layout is not None: 

240 return _OUTPUT_NON_CP, non_cp_layout 

241 return _OUTPUT_CP, tensor.device_mesh 

242 

243 

244def _push_output_layout(module, layout) -> None: 

245 """Push the current forward's output-layout policy for the matching post-hook.""" 

246 if module is None: 

247 return 

248 stack = getattr(module, _OUTPUT_LAYOUT_STACK_ATTR, None) 

249 if stack is None: 

250 stack = [] 

251 setattr(module, _OUTPUT_LAYOUT_STACK_ATTR, stack) 

252 stack.append(layout) 

253 

254 

255def _pop_output_layout(module): 

256 """Pop the output-layout policy recorded by the matching pre-hook.""" 

257 if module is None: 

258 return _OUTPUT_LOCAL, None 

259 stack = getattr(module, _OUTPUT_LAYOUT_STACK_ATTR, None) 

260 if not stack: 

261 return _OUTPUT_LOCAL, None 

262 layout = stack.pop() 

263 if not stack: 

264 delattr(module, _OUTPUT_LAYOUT_STACK_ATTR) 

265 return layout 

266 

267 

268def _drop_cp_from_output(output, layout, cp_placements): 

269 """Drop CP metadata from the current local shard and keep the non-CP layout.""" 

270 if layout is None or not isinstance(output, (Tensor, DTensor)): 

271 return output 

272 non_cp_mesh, non_cp_placements, composed_mesh = layout 

273 cp_mesh = _cp_mesh_from_composed(composed_mesh, non_cp_mesh) 

274 composed_placements = _composed_placements( 

275 composed_mesh, 

276 cp_mesh, 

277 cp_placements, 

278 non_cp_mesh, 

279 non_cp_placements, 

280 ) 

281 if isinstance(output, DTensor): 

282 if not _same_mesh(output.device_mesh, composed_mesh): 

283 output = DTensor.from_local(output.to_local(), composed_mesh, composed_placements) 

284 else: 

285 output = DTensor.from_local(output, composed_mesh, composed_placements) 

286 if tuple(output.placements) != composed_placements: 

287 output = output.redistribute(composed_mesh, composed_placements) 

288 return DTensor.from_local(output.to_local(), non_cp_mesh, non_cp_placements) 

289 

290 

291def _wrap_cp_output_dtensor(output, device_mesh: DeviceMesh, placements): 

292 """Wrap/redistribute CP output as a CP DTensor.""" 

293 if not isinstance(output, (Tensor, DTensor)): 

294 return output 

295 if isinstance(output, DTensor): 

296 if _same_mesh(output.device_mesh, device_mesh): 

297 if tuple(output.placements) == tuple(placements): 

298 return output 

299 return output.redistribute(device_mesh, placements) 

300 output = output.to_local() 

301 return DTensor.from_local(output, device_mesh, placements) 

302 

303 

304def _finalize_output(output, output_layout, cp_mesh, cp_placements, use_local_output: bool): 

305 """Apply CP boundary output policy. 

306 

307 With ``use_local_output=True`` the public output is always local. Otherwise 

308 the output mirrors the Q input boundary: local input returns local output, 

309 CP-DTensor input returns a CP DTensor, and non-CP DTensor input drops CP 

310 back to that non-CP mesh. 

311 """ 

312 if use_local_output: 

313 return output.to_local() if isinstance(output, DTensor) else output 

314 output_kind, layout = output_layout 

315 if output_kind == _OUTPUT_NON_CP: 

316 return _drop_cp_from_output(output, layout, cp_placements) 

317 if output_kind == _OUTPUT_CP: 

318 return _wrap_cp_output_dtensor(output, cp_mesh, cp_placements) 

319 return output.to_local() if isinstance(output, DTensor) else output 

320 

321 

322def _finalize_colossal_output(output, output_layout, co_submesh, seq_dim: int, use_local_output: bool): 

323 """Apply Colossal output policy: local tensor or DTensor output.""" 

324 return _finalize_output(output, output_layout, co_submesh, (Shard(seq_dim),), use_local_output) 

325 

326 

327def _finalize_ata_output(output_dtensor, output_layout, ds_submesh, seq_dim: int, use_local_output: bool): 

328 """Apply Ulysses/Hybrid output policy after reverse ATA.""" 

329 return _finalize_output(output_dtensor, output_layout, ds_submesh, (Shard(seq_dim),), use_local_output) 

330 

331 

332def _is_tensor_or_dtensor(value) -> bool: 

333 """Return whether *value* can participate in CP tensor conversion.""" 

334 return isinstance(value, (Tensor, DTensor)) 

335 

336 

337# --------------------------------------------------------------------------- 

338# Low-level communication primitives 

339# --------------------------------------------------------------------------- 

340 

341def _ensure_1d(device_mesh: DeviceMesh) -> DeviceMesh: 

342 """Return a 1-D DeviceMesh (flatten if multi-dimensional).""" 

343 if device_mesh.ndim == 1: 

344 return device_mesh 

345 ranks = list(device_mesh.rank_list) 

346 return DeviceMesh(device_mesh.device_type, ranks, mesh_dim_names=("cp",)) 

347 

348 

349def _build_2d_mesh(device_mesh: DeviceMesh, ds: int, co: int) -> DeviceMesh: 

350 """Build or validate a 2-D ``(co × ds)`` DeviceMesh for Hybrid CP. 

351 

352 If *device_mesh* is already 2-D it is returned as-is (must have 

353 ``mesh_dim_names`` set). Otherwise the ranks of the 1-D mesh are tiled 

354 into *co* rows of *ds* adjacent ranks each. 

355 """ 

356 if device_mesh.ndim == 2: 

357 if not device_mesh.mesh_dim_names: 

358 raise ValueError( 

359 "2-D device_mesh for Hybrid CP must have mesh_dim_names=(\"co\", \"ds\")." 

360 ) 

361 return device_mesh 

362 ranks = list(device_mesh.rank_list) 

363 return DeviceMesh( 

364 device_mesh.device_type, 

365 [ranks[i * ds:(i + 1) * ds] for i in range(co)], 

366 mesh_dim_names=("co", "ds"), 

367 ) 

368 

369 

370def _build_hybrid_cp_mesh(cp_mesh: DeviceMesh, ds: int, co: int) -> DeviceMesh: 

371 """Build the Hybrid CP ``(co, ds)`` mesh, preserving root layout when possible.""" 

372 if cp_mesh.ndim == 2: 

373 if not cp_mesh.mesh_dim_names: 

374 raise ValueError( 

375 "2-D device_mesh for Hybrid CP must have mesh_dim_names=(\"co\", \"ds\")." 

376 ) 

377 return cp_mesh 

378 if cp_mesh.ndim != 1: 

379 raise ValueError(f"Hybrid CP expects a 1-D or 2-D CP mesh, got {cp_mesh.ndim}D.") 

380 if cp_mesh.root_mesh is not None and cp_mesh.mesh_dim_names == ("cp",): 

381 return cp_mesh._unflatten("cp", (co, ds), ("co", "ds")) # pylint: disable=protected-access 

382 return _build_2d_mesh(cp_mesh, ds, co) 

383 

384 

385def _scatter_seq_to_head( 

386 tensor: Tensor, 

387 submesh: DeviceMesh, 

388 seq_dim: int, 

389 head_dim: int, 

390 submesh_size: int, 

391) -> "DTensor": 

392 """All-to-all: ``Shard(seq_dim)`` → ``Shard(head_dim)``. Returns DTensor.""" 

393 head_size = _cp_boundary_dim_size(tensor, submesh, head_dim) 

394 if head_size % submesh_size != 0: 

395 raise ValueError( 

396 f"num_heads ({head_size}) must be divisible by " 

397 f"ulysses_degree ({submesh_size})." 

398 ) 

399 return _to_cp_dtensor(tensor, submesh, (Shard(seq_dim),), (Shard(head_dim),), seq_dim) 

400 

401 

402def _gather_head_to_seq( 

403 tensor: Tensor, 

404 submesh: DeviceMesh, 

405 seq_dim: int, 

406 head_dim: int, 

407) -> "DTensor": 

408 """Reverse all-to-all: ``Shard(head_dim)`` → ``Shard(seq_dim)``. Returns DTensor.""" 

409 return _to_cp_dtensor(tensor, submesh, (Shard(head_dim),), (Shard(seq_dim),), seq_dim) 

410 

411 

412def _gather_seq( 

413 tensor: Tensor, 

414 submesh: DeviceMesh, 

415 seq_dim: int, 

416) -> "DTensor": 

417 """All-gather: ``Shard(seq_dim)`` → ``Replicate``. Returns DTensor.""" 

418 return _to_cp_dtensor(tensor, submesh, (Shard(seq_dim),), (Replicate(),), seq_dim) 

419 

420 

421# --------------------------------------------------------------------------- 

422# Unified ContextParallel 

423# --------------------------------------------------------------------------- 

424 

425class ContextParallel(ParallelStyle): 

426 """Unified Context Parallel for core-attention modules. 

427 

428 Three modes controlled by ``ulysses_degree``: 

429 

430 +-----------------+--------------------+------------------------------------------+ 

431 | Mode | ``ulysses_degree`` | Mechanism | 

432 +=================+====================+==========================================+ 

433 | Pure Ulysses | ``None`` (default) | seq→head A2A before attn; | 

434 | | (= cp_size) | head→seq A2A after. | 

435 | | | Requires ``num_heads % cp_size == 0``. | 

436 +-----------------+--------------------+------------------------------------------+ 

437 | Pure Colossal AI| ``1`` | Q stays as local Shard(seq); | 

438 | | | K/V all-gathered (Replicate). | 

439 | | | No head-count constraint. | 

440 +-----------------+--------------------+------------------------------------------+ 

441 | Hybrid | ``1 < k < cp_size``| Q/K/V seq→head A2A on Ulysses sub-mesh | 

442 | | | (size ``k``); K/V then all-gathered on | 

443 | | | Colossal sub-mesh (size ``cp_size // k``)| 

444 | | | Requires ``num_heads % k == 0``. | 

445 +-----------------+--------------------+------------------------------------------+ 

446 

447 Args: 

448 seq_dim: Sequence dimension index. 1 for BSHD, 2 for BNSD. 

449 head_dim: Head dimension index. 2 for BSHD, 1 for BNSD. 

450 ulysses_degree: Ulysses sub-mesh size (see table above). 

451 qkv_indices: Positional-argument indices for (Q, K, V). 

452 qkv_kwarg_names: Keyword-argument names for (Q, K, V). 

453 use_local_output: Return local tensors after CP when True. When False, 

454 keep CP DTensor outputs, or drop the CP axis from 

455 composed CP+TP outputs and keep the non-CP layout. 

456 load_balance: Enable Head-Tail Q-exchange load balancing. 

457 Only valid with Pure Colossal AI (``ulysses_degree=1``). 

458 

459 **Important**: When ``load_balance=True``, ``q.shape[seq_dim]`` 

460 inside ``forward()`` returns ``S / 2`` (global shape / 2) 

461 rather than the true global ``S``. This is because 

462 ``DTensor.shape`` returns ``local_tensor_size * mesh_size``, 

463 and each sub-FA call wraps a half-sized Q shard 

464 (``S / (2 * cp_size)`` tokens) with a ``co_submesh`` of 

465 size ``cp_size``, giving a DTensor global shape of 

466 ``S / (2 * cp_size) * cp_size = S / 2``. 

467 K/V are always Replicate so ``k.shape[seq_dim]`` always 

468 returns the true ``S``. **When building the attention mask, 

469 use ``k.shape[seq_dim]`` (not ``q.shape[seq_dim]``) to 

470 obtain the correct global sequence length.** 

471 """ 

472 

473 def __init__( 

474 self, 

475 seq_dim: int = 1, 

476 head_dim: int = 2, 

477 ulysses_degree: Optional[int] = None, 

478 qkv_indices: tuple = (0, 1, 2), 

479 qkv_kwarg_names: tuple = (), 

480 use_local_output: bool = False, 

481 load_balance: bool = False, 

482 ): 

483 if load_balance and ulysses_degree != 1: 

484 raise ValueError( 

485 "load_balance=True requires ulysses_degree=1 (Pure Colossal AI mode)." 

486 ) 

487 self.seq_dim = seq_dim 

488 self.head_dim = head_dim 

489 self.ulysses_degree = ulysses_degree 

490 self.qkv_indices = qkv_indices 

491 self.qkv_kwarg_names = qkv_kwarg_names 

492 self.use_local_output = use_local_output 

493 self.load_balance = load_balance 

494 

495 # ------------------------------------------------------------------ 

496 # ParallelStyle interface 

497 # ------------------------------------------------------------------ 

498 

499 def apply(self, module: Module, device_mesh: DeviceMesh) -> Module: 

500 """Register forward hooks on *module* and return it. 

501 

502 Args: 

503 module: attention submodule to parallelise. 

504 device_mesh: CP device mesh (1-D or 2-D). 

505 """ 

506 cp_size = device_mesh.mesh.numel() 

507 ds = self.ulysses_degree if self.ulysses_degree is not None else cp_size 

508 if cp_size % ds != 0: 

509 raise ValueError( 

510 f"cp_size ({cp_size}) must be divisible by ulysses_degree ({ds})." 

511 ) 

512 co = cp_size // ds 

513 

514 if ds == 1: 

515 # Pure Colossal AI 

516 co_submesh = _ensure_1d(device_mesh) 

517 if self.load_balance: 

518 self._apply_lb_colossal(module, co_submesh) 

519 else: 

520 module.register_forward_pre_hook( 

521 partial(self._pre_hook_colossal, co_submesh=co_submesh), 

522 with_kwargs=True, 

523 ) 

524 module.register_forward_hook( 

525 partial(self._post_hook_colossal, co_submesh=co_submesh) 

526 ) 

527 elif co == 1: 

528 # Pure Ulysses 

529 ds_submesh = _ensure_1d(device_mesh) 

530 module.register_forward_pre_hook( 

531 partial(self._pre_hook_ulysses, ds_submesh=ds_submesh, ds_size=ds), 

532 with_kwargs=True, 

533 ) 

534 module.register_forward_hook( 

535 partial(self._post_hook_ata, ds_submesh=ds_submesh) 

536 ) 

537 else: 

538 # Hybrid 

539 hybrid_cp_mesh = _build_hybrid_cp_mesh(device_mesh, ds, co) 

540 dim_names = hybrid_cp_mesh.mesh_dim_names 

541 if dim_names is None: 

542 raise ValueError("2-D mesh must have mesh_dim_names (guaranteed by _build_2d_mesh)") 

543 ds_submesh = hybrid_cp_mesh[dim_names[1]] 

544 module.register_forward_pre_hook( 

545 partial( 

546 self._pre_hook_hybrid, 

547 hybrid_cp_mesh=hybrid_cp_mesh, 

548 ds_submesh=ds_submesh, 

549 ds_size=ds, 

550 ), 

551 with_kwargs=True, 

552 ) 

553 module.register_forward_hook( 

554 partial( 

555 self._post_hook_hybrid, 

556 hybrid_cp_mesh=hybrid_cp_mesh, 

557 ds_submesh=ds_submesh, 

558 ) 

559 ) 

560 

561 return module 

562 

563 # ------------------------------------------------------------------ 

564 # Pre-hooks 

565 # ------------------------------------------------------------------ 

566 

567 def _record_q_output_tp_layout(self, module, args, kwargs, submesh) -> None: 

568 """Remember how the CP output should cross the public boundary.""" 

569 output_layout = (_OUTPUT_LOCAL, None) 

570 q_idx = self.qkv_indices[0] 

571 if q_idx < len(args): 

572 output_layout = _output_layout_from_q(args[q_idx], submesh, self.seq_dim) 

573 if output_layout[0] == _OUTPUT_LOCAL and self.qkv_kwarg_names: 

574 q_name = self.qkv_kwarg_names[0] 

575 if q_name in kwargs: 

576 output_layout = _output_layout_from_q(kwargs[q_name], submesh, self.seq_dim) 

577 _push_output_layout(module, output_layout) 

578 

579 def _pre_hook_colossal(self, module, args, kwargs, co_submesh): # pylint: disable=unused-argument 

580 """Wrap Q as ``DTensor(co_submesh, Shard(seq))``; all-gather K/V.""" 

581 new_args = list(args) 

582 new_kwargs = dict(kwargs) 

583 

584 self._record_q_output_tp_layout(module, new_args, new_kwargs, co_submesh) 

585 q_idx = self.qkv_indices[0] 

586 if q_idx < len(new_args) and _is_tensor_or_dtensor(new_args[q_idx]): 

587 new_args[q_idx] = _to_cp_dtensor( 

588 new_args[q_idx], 

589 co_submesh, 

590 (Shard(self.seq_dim),), 

591 (Shard(self.seq_dim),), 

592 self.seq_dim, 

593 ) 

594 for idx in self.qkv_indices[1:]: 

595 if idx < len(new_args) and _is_tensor_or_dtensor(new_args[idx]): 

596 new_args[idx] = _gather_seq(new_args[idx], co_submesh, self.seq_dim) 

597 

598 if self.qkv_kwarg_names: 

599 q_name = self.qkv_kwarg_names[0] 

600 if q_name in new_kwargs and _is_tensor_or_dtensor(new_kwargs[q_name]): 

601 new_kwargs[q_name] = _to_cp_dtensor( 

602 new_kwargs[q_name], 

603 co_submesh, 

604 (Shard(self.seq_dim),), 

605 (Shard(self.seq_dim),), 

606 self.seq_dim, 

607 ) 

608 for name in self.qkv_kwarg_names[1:]: 

609 if name in new_kwargs and _is_tensor_or_dtensor(new_kwargs[name]): 

610 new_kwargs[name] = _gather_seq(new_kwargs[name], co_submesh, self.seq_dim) 

611 

612 return tuple(new_args), new_kwargs 

613 

614 def _pre_hook_ulysses(self, module, args, kwargs, ds_submesh, ds_size): # pylint: disable=unused-argument 

615 """Seq→head all-to-all for Q, K, and V.""" 

616 new_args = list(args) 

617 new_kwargs = dict(kwargs) 

618 

619 self._record_q_output_tp_layout(module, new_args, new_kwargs, ds_submesh) 

620 for idx in self.qkv_indices: 

621 if idx < len(new_args) and _is_tensor_or_dtensor(new_args[idx]): 

622 new_args[idx] = _scatter_seq_to_head( 

623 new_args[idx], ds_submesh, self.seq_dim, self.head_dim, ds_size 

624 ) 

625 

626 for name in self.qkv_kwarg_names: 

627 if name in new_kwargs and _is_tensor_or_dtensor(new_kwargs[name]): 

628 new_kwargs[name] = _scatter_seq_to_head( 

629 new_kwargs[name], ds_submesh, self.seq_dim, self.head_dim, ds_size 

630 ) 

631 

632 return tuple(new_args), new_kwargs 

633 

634 def _hybrid_cp_layout_from_input(self, t, hybrid_cp_mesh, cp_placements): 

635 """Return the Hybrid execution mesh and placements for one input tensor.""" 

636 layout = _non_cp_dtensor_layout(t, hybrid_cp_mesh, self.seq_dim) 

637 if layout is None: 

638 return hybrid_cp_mesh, tuple(cp_placements) 

639 non_cp_mesh, non_cp_placements, composed_mesh = layout 

640 return composed_mesh, _compose_cp_non_cp_placements( 

641 composed_mesh, 

642 hybrid_cp_mesh, 

643 cp_placements, 

644 non_cp_mesh, 

645 non_cp_placements, 

646 ) 

647 

648 def _ata_scatter_to_hybrid(self, t, ds_submesh, hybrid_cp_mesh, ds_size): 

649 """ATA scatter on ds, then wrap as Hybrid CP or Hybrid CP + non-CP DTensor. 

650 

651 Args: 

652 t: Plain local tensor to scatter. 

653 ds_submesh: 1-D Ulysses sub-mesh. 

654 hybrid_cp_mesh: 2-D mesh (co × ds). 

655 ds_size: Ulysses degree (world size on ds_submesh). 

656 

657 Returns: 

658 DTensor with CP placements ``(Shard(seq_dim), Shard(head_dim))``, 

659 preserving any incoming non-CP placements. 

660 """ 

661 out_mesh, out_placements = self._hybrid_cp_layout_from_input( 

662 t, 

663 hybrid_cp_mesh, 

664 (Shard(self.seq_dim), Shard(self.head_dim)), 

665 ) 

666 t = _localize_foreign_dtensor(t, ds_submesh, self.seq_dim) 

667 if isinstance(t, DTensor): 

668 t = t.redistribute(ds_submesh, (Shard(self.seq_dim),)).to_local() 

669 if t.shape[self.head_dim] % ds_size != 0: 

670 raise ValueError( 

671 f"num_heads ({t.shape[self.head_dim]}) must be divisible by " 

672 f"ulysses_degree ({ds_size})." 

673 ) 

674 local = ( 

675 DTensor.from_local(t, ds_submesh, (Shard(self.seq_dim),)) 

676 .redistribute(ds_submesh, (Shard(self.head_dim),)) 

677 .to_local() 

678 ) 

679 return DTensor.from_local(local, out_mesh, out_placements) 

680 

681 def _hybrid_kv_gather_placements(self, t, hybrid_cp_mesh): 

682 """Return K/V placements after co gather for a Hybrid input DTensor.""" 

683 out_mesh, out_placements = self._hybrid_cp_layout_from_input( 

684 t, 

685 hybrid_cp_mesh, 

686 (Replicate(), Shard(self.head_dim)), 

687 ) 

688 return out_mesh, out_placements 

689 

690 def _pre_hook_hybrid( # pylint: disable=unused-argument 

691 self, module, args, kwargs, hybrid_cp_mesh, ds_submesh, ds_size 

692 ): 

693 """Hybrid: seq→head ATA on ds-submesh, then all-gather K/V on co-submesh. 

694 

695 After this hook, CP placements on ``hybrid_cp_mesh`` are: 

696 Q → ``(Shard(seq_dim), Shard(head_dim))`` 

697 K/V → ``(Replicate(), Shard(head_dim))`` 

698 Incoming non-CP placements are preserved on a composed mesh. 

699 """ 

700 new_args = list(args) 

701 new_kwargs = dict(kwargs) 

702 

703 self._record_q_output_tp_layout(module, new_args, new_kwargs, hybrid_cp_mesh) 

704 

705 # Step 1: ATA on ds_submesh for all of Q/K/V; wrap as 2-D DTensor 

706 for idx in self.qkv_indices: 

707 if idx < len(new_args) and self._needs_hybrid_ata(new_args[idx], hybrid_cp_mesh): 

708 new_args[idx] = self._ata_scatter_to_hybrid( 

709 new_args[idx], ds_submesh, hybrid_cp_mesh, ds_size 

710 ) 

711 

712 # Step 2: all-gather K/V on co-dim (Shard(seq)→Replicate) 

713 for idx in self.qkv_indices[1:]: 

714 if idx < len(new_args) and isinstance(new_args[idx], DTensor): 

715 out_mesh, out_placements = self._hybrid_kv_gather_placements( 

716 new_args[idx], 

717 hybrid_cp_mesh, 

718 ) 

719 new_args[idx] = new_args[idx].redistribute(out_mesh, out_placements) 

720 

721 # Same for kwargs 

722 for name in self.qkv_kwarg_names: 

723 if name in new_kwargs and self._needs_hybrid_ata(new_kwargs[name], hybrid_cp_mesh): 

724 new_kwargs[name] = self._ata_scatter_to_hybrid( 

725 new_kwargs[name], ds_submesh, hybrid_cp_mesh, ds_size 

726 ) 

727 for name in self.qkv_kwarg_names[1:]: 

728 if name in new_kwargs and isinstance(new_kwargs[name], DTensor): 

729 out_mesh, out_placements = self._hybrid_kv_gather_placements( 

730 new_kwargs[name], 

731 hybrid_cp_mesh, 

732 ) 

733 new_kwargs[name] = new_kwargs[name].redistribute(out_mesh, out_placements) 

734 

735 return tuple(new_args), new_kwargs 

736 

737 @staticmethod 

738 def _needs_hybrid_ata(value, hybrid_cp_mesh) -> bool: 

739 """Return whether Hybrid CP should run ATA before entering attention.""" 

740 if not _is_tensor_or_dtensor(value): 

741 return False 

742 return not isinstance(value, DTensor) or _is_foreign_dtensor(value, hybrid_cp_mesh) 

743 

744 # ------------------------------------------------------------------ 

745 # Post-hooks 

746 # ------------------------------------------------------------------ 

747 

748 def _post_hook_ata(self, module, inputs, outputs, ds_submesh): # pylint: disable=unused-argument 

749 """Reverse all-to-all: head→seq on ds-submesh; returns local tensor. 

750 

751 Handles both Ulysses (1-D DTensor or plain tensor) and Hybrid 

752 (2-D DTensor — ``to_local()`` first to project onto the 1-D ds-submesh). 

753 """ 

754 output_layout = _pop_output_layout(module) 

755 

756 def _process(out): 

757 if isinstance(out, (Tensor, DTensor)): 

758 if isinstance(out, DTensor) and not _is_cp_composed_dtensor(out, ds_submesh): 

759 out = out.to_local() 

760 seq_dtensor = _gather_head_to_seq( 

761 out, ds_submesh, self.seq_dim, self.head_dim 

762 ) 

763 return _finalize_ata_output( 

764 seq_dtensor, output_layout, ds_submesh, self.seq_dim, self.use_local_output 

765 ) 

766 return out 

767 

768 if isinstance(outputs, (tuple, list)): 

769 return type(outputs)(_process(item) for item in outputs) 

770 return _process(outputs) 

771 

772 def _post_hook_hybrid(self, module, inputs, outputs, hybrid_cp_mesh, ds_submesh): # pylint: disable=unused-argument 

773 """Hybrid reverse ATA, then drop the full ``(co, ds)`` CP mesh if needed.""" 

774 output_layout = _pop_output_layout(module) 

775 

776 def _process(out): 

777 if not isinstance(out, (Tensor, DTensor)): 

778 return out 

779 out_local = out.to_local() if isinstance(out, DTensor) else out 

780 seq_dtensor = _gather_head_to_seq( 

781 out_local, 

782 ds_submesh, 

783 self.seq_dim, 

784 self.head_dim, 

785 ) 

786 seq_local = seq_dtensor.to_local() 

787 if self.use_local_output: 

788 return seq_local 

789 

790 output_kind, layout = output_layout 

791 if output_kind == _OUTPUT_NON_CP: 

792 non_cp_mesh, non_cp_placements, _ = layout 

793 return DTensor.from_local(seq_local, non_cp_mesh, non_cp_placements) 

794 if output_kind == _OUTPUT_CP: 

795 return DTensor.from_local( 

796 seq_local, 

797 hybrid_cp_mesh, 

798 (Shard(self.seq_dim), Replicate()), 

799 ) 

800 return seq_local 

801 

802 if isinstance(outputs, (tuple, list)): 

803 return type(outputs)(_process(item) for item in outputs) 

804 return _process(outputs)