Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / fully_shard / param_group.py: 83%

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"""MindSpore HSDP parameter group with fused communication.""" 

16 

17from __future__ import annotations 

18 

19import math 

20from dataclasses import dataclass, field 

21from typing import Any, List, NamedTuple, Optional 

22 

23import mindspore as ms 

24from mindspore import ops 

25from mindspore.common.api import _no_grad 

26import mindspore.mint.distributed as dist 

27from mindspore.ops.function.comm_func import CommHandle 

28 

29from hyper_parallel.core.fully_shard.hsdp_utils import apply_gradient_scaling_factor 

30from hyper_parallel.core.fully_shard.utils import DDPMeshInfo, FSDPMeshInfo, HSDPMeshInfo, MixedPrecisionPolicy 

31from hyper_parallel.platform.mindspore.fully_shard._version_utils import copy_without_bumping_version 

32from hyper_parallel.platform.mindspore.fully_shard.pack_utils import build_rs_plan, pack_for_reduce_scatter 

33from hyper_parallel.platform.mindspore.fully_shard.param import MindSporeHSDPParamV2 

34 

35 

36def _normalize_device(device: Any) -> str: 

37 if isinstance(device, str): 

38 return device.split(":", 1)[0] 

39 return str(device).split(":", 1)[0] 

40 

41 

42def _shape_numel(shape) -> int: 

43 return math.prod(int(dim) for dim in shape) 

44 

45 

46def get_all_gather_metadata(hsdp_params): 

47 """Collect metadata required for fused all-gather.""" 

48 param_input_dtypes = [] 

49 param_input_numels = [] 

50 inp_split_sizes = [] 

51 total_input_numel = 0 

52 first_dtype = None 

53 

54 for hsdp_param in hsdp_params: 

55 inputs = hsdp_param.all_gather_inputs 

56 if first_dtype is None: 

57 first_dtype = inputs[0].dtype 

58 elif first_dtype != inputs[0].dtype: 

59 raise ValueError("All parameters in the group must have a uniform dtype.") 

60 param_dtypes = [t.dtype for t in inputs] 

61 param_numels = [t.numel() for t in inputs] 

62 param_input_dtypes.append(param_dtypes) 

63 param_input_numels.append(param_numels) 

64 inp_split_sizes.extend(param_numels) 

65 total_input_numel += sum(param_numels) 

66 

67 return AllGatherMetadata( 

68 param_input_dtypes, 

69 param_input_numels, 

70 first_dtype, 

71 inp_split_sizes, 

72 total_input_numel, 

73 ) 

74 

75 

76@dataclass 

77class AllGatherMetadata: 

78 """Metadata describing the fused all-gather buffer layout.""" 

79 

80 param_input_dtypes: list[list[Any]] 

81 param_input_numels: list[list[int]] 

82 dtype: Any 

83 inp_split_sizes: list[int] 

84 total_input_numel: int 

85 hash_key: int = field(init=False) 

86 

87 def __post_init__(self): 

88 self.hash_key = hash( 

89 ( 

90 tuple(tuple(d) for d in self.param_input_dtypes), 

91 tuple(tuple(n) for n in self.param_input_numels), 

92 self.dtype, 

93 tuple(self.inp_split_sizes), 

94 self.total_input_numel, 

95 ) 

96 ) 

97 

98 

99class AllGatherResult(NamedTuple): 

100 """Result of a fused all-gather operation.""" 

101 

102 all_gather_output: Optional[ms.Tensor] 

103 metadata: Optional[AllGatherMetadata] 

104 handle: Optional[CommHandle] 

105 

106 

107@dataclass 

108class CommContext: 

109 """Global communication context for pipelined fused reductions.""" 

110 

111 comm_handle: Optional[CommHandle] = None 

112 all_reduce_handle: Optional[CommHandle] = None 

113 pre_param_group = None 

114 all_reduce_param_group = None 

115 

116 

117comm_ctx = CommContext() 

118 

119 

120def get_comm_ctx(): 

121 """Return the global communication context singleton.""" 

122 return comm_ctx 

123 

124 

125@dataclass 

126class ReplicateBucket: 

127 """One fused all-reduce bucket sharing the same replicate process group.""" 

128 

129 key: int 

130 group: Any 

131 group_size: int 

132 param_indices: list[int] 

133 flat_numel: int 

134 buffer: Optional[ms.Tensor] = None 

135 

136 

137@dataclass 

138class PendingBucketAllReduce: 

139 """One in-flight async all-reduce launched for a replicate bucket.""" 

140 

141 bucket_key: int 

142 handle: Any 

143 

144 

145class AllGatherMetadataCache: 

146 """Cache for all-gather metadata across iterations.""" 

147 

148 _cache: dict[int, AllGatherMetadata] = {} 

149 

150 @classmethod 

151 def get_metadata(cls, hsdp_params, fn): 

152 """Retrieve or compute all-gather metadata, caching the result by parameter identity.""" 

153 param_key = tuple((id(p), getattr(p, "version", 0)) for p in hsdp_params) 

154 key = hash(param_key) 

155 if key in cls._cache: 

156 return cls._cache[key] 

157 metadata = fn(hsdp_params) 

158 cls._cache[key] = metadata 

159 return metadata 

160 

161 

162@_no_grad() 

163def all_gather_copy_in(all_gather_inputs, all_gather_output, inp_split_sizes, all_gather_input_numel, rank): 

164 """Copy per-parameter local shards into one fused rank-local all-gather slice.""" 

165 all_gather_input = all_gather_output.narrow(0, all_gather_input_numel * rank, all_gather_input_numel) 

166 offset = 0 

167 for src, size in zip(all_gather_inputs, inp_split_sizes): 

168 src_flat = src.view(-1) 

169 all_gather_input.narrow(0, offset, size).copy_(src_flat) 

170 offset += size 

171 return all_gather_input, all_gather_output 

172 

173 

174@_no_grad() 

175def split_with_sizes_copy(all_gather_output, split_sizes, dim, out): 

176 """Copy split views from a fused all-gather output into pre-allocated outputs.""" 

177 if dim != 1: 

178 raise NotImplementedError("split_with_sizes_copy currently only supports dim=1") 

179 offset = 0 

180 for dst, size in zip(out, split_sizes): 

181 src = all_gather_output.narrow(dim, offset, size) 

182 copy_without_bumping_version(dst, src) 

183 offset += size 

184 

185 

186@_no_grad() 

187def reduce_scatter_copy_in( 

188 hsdp_params: List[MindSporeHSDPParamV2], 

189 unsharded_grads: List[ms.Tensor], 

190 reduce_scatter_input: ms.Tensor, 

191 world_size: int, 

192) -> None: 

193 """Pack all unsharded gradients into one fused reduce-scatter input buffer.""" 

194 if len(hsdp_params) != len(unsharded_grads): 

195 raise AssertionError( 

196 "reduce_scatter_copy_in expects one hsdp_param per unsharded_grad, but got " 

197 f"{len(hsdp_params)} params and {len(unsharded_grads)} grads" 

198 ) 

199 packed_rows = reduce_scatter_input.view(world_size, -1) 

200 col_offset = 0 

201 for hsdp_param, grad in zip(hsdp_params, unsharded_grads): 

202 grad = grad.contiguous() 

203 plan = build_rs_plan(hsdp_param, grad, world_size) 

204 packed_grad = pack_for_reduce_scatter(grad, plan) 

205 next_col_offset = col_offset + packed_grad.shape[1] 

206 for row_idx in range(world_size): 

207 packed_rows[row_idx].narrow(0, col_offset, packed_grad.shape[1]).copy_( 

208 packed_grad[row_idx].view(-1) 

209 ) 

210 col_offset = next_col_offset 

211 if col_offset != packed_rows.shape[1]: 

212 raise AssertionError( 

213 "reduce_scatter_copy_in packed an unexpected number of elements: " 

214 f"{col_offset} != {packed_rows.shape[1]}" 

215 ) 

216 

217 

218class HSDPParamGroup: 

219 """Group HSDP parameters within a module for fused collectives.""" 

220 

221 def __init__( 

222 self, 

223 hsdp_params, 

224 mesh_info: FSDPMeshInfo, 

225 device: Optional[str] = None, 

226 mp_policy: Optional[MixedPrecisionPolicy] = None, 

227 enable_zero_copy_param_buffer: bool = False, 

228 ): 

229 self.mesh_info = mesh_info 

230 self.device = device 

231 self.hsdp_params = hsdp_params 

232 self.enable_zero_copy_param_buffer = enable_zero_copy_param_buffer 

233 if isinstance(self.mesh_info, (FSDPMeshInfo, HSDPMeshInfo)): 

234 self.shard_rank = self.mesh_info.shard_mesh_rank 

235 self.shard_world_size = self.mesh_info.shard_mesh_size 

236 else: 

237 self.shard_rank = 0 

238 self.shard_world_size = 1 

239 self.shard_group = self.mesh_info.shard_process_group 

240 self.replicate_group = None 

241 if isinstance(self.mesh_info, (HSDPMeshInfo, DDPMeshInfo)): 

242 self.replicate_group = self.mesh_info.replicate_process_group 

243 elif isinstance(self.mesh_info, FSDPMeshInfo): 

244 self.replicate_group = self._infer_layout_replicate_group() 

245 self.ag_output: Optional[ms.Tensor] = None 

246 self.metadata_cache = None 

247 self.mp_policy = mp_policy 

248 self._result = None 

249 self._reduce_output = None 

250 self._reduce_op = None 

251 self._reduce_hsdp_params = None 

252 self._active_replicate_buckets: dict[int, ReplicateBucket] = {} 

253 self._active_param_flat_offsets: list[int] = [] 

254 self._pending_all_reduce_handles: list[PendingBucketAllReduce] = [] 

255 self._flat_param_buffer: Optional[ms.Tensor] = None 

256 self._flat_cast_buffer: Optional[ms.Tensor] = None 

257 self._init_mp_dtypes() 

258 if self.enable_zero_copy_param_buffer: 

259 self._init_flat_param_buffer() 

260 self.gradient_scaling_factor = None 

261 

262 def _infer_layout_replicate_group(self): 

263 replicate_groups = [] 

264 for hsdp_param in self.hsdp_params: 

265 group_info = getattr(hsdp_param, "unsharded_group_info", None) 

266 group = getattr(group_info, "group", None) 

267 if group is None or getattr(hsdp_param, "replicate_world_size", 1) <= 1: 

268 continue 

269 replicate_groups.append(group) 

270 if not replicate_groups: 

271 return None 

272 return replicate_groups[0] 

273 

274 @staticmethod 

275 def _build_active_replicate_buckets(hsdp_params): 

276 buckets: dict[int, ReplicateBucket] = {} 

277 for idx, hsdp_param in enumerate(hsdp_params): 

278 group_info = getattr(hsdp_param, "unsharded_group_info", None) 

279 group = getattr(group_info, "group", None) 

280 group_size = getattr(group_info, "rank_size", getattr(hsdp_param, "replicate_world_size", 1)) 

281 if group is None or group_size <= 1: 

282 continue 

283 key = id(group) 

284 if key not in buckets: 

285 buckets[key] = ReplicateBucket( 

286 key=key, 

287 group=group, 

288 group_size=group_size, 

289 param_indices=[], 

290 flat_numel=0, 

291 ) 

292 buckets[key].param_indices.append(idx) 

293 buckets[key].flat_numel += _shape_numel(hsdp_param.sharded_size) 

294 return buckets 

295 

296 def _init_flat_param_buffer(self): 

297 """Rebase local shards into one flat buffer when storage semantics allow it.""" 

298 if not self.enable_zero_copy_param_buffer: 

299 return 

300 if self.shard_world_size <= 1 or len(self.hsdp_params) == 0: 

301 return 

302 if any(p.offload_to_cpu or str(p.sharded_param.device) == "meta" for p in self.hsdp_params): 

303 return 

304 

305 total_numel = sum(hsdp_param._sharded_param_data.numel() for hsdp_param in self.hsdp_params) 

306 orig_dtype = self.hsdp_params[0]._sharded_param_data.dtype 

307 flat_buffer = ms.mint.empty((total_numel,), dtype=orig_dtype, device=_normalize_device(self.device)) 

308 

309 offset = 0 

310 original_locals = [] 

311 try: 

312 for hsdp_param in self.hsdp_params: 

313 original_locals.append((hsdp_param, hsdp_param._sharded_param_data, hsdp_param._sharded_local_tensor)) 

314 numel = hsdp_param._sharded_param_data.numel() 

315 flat_slice = flat_buffer.narrow(0, offset, numel) 

316 flat_slice.copy_(hsdp_param._sharded_param_data) 

317 hsdp_param._sharded_param_data = flat_slice 

318 new_local = flat_slice.view(hsdp_param.sharded_size) 

319 req_grad = hsdp_param.sharded_param.requires_grad 

320 hsdp_param.sharded_param.set_data(new_local) 

321 hsdp_param.sharded_param._local_tensor = new_local 

322 if req_grad: 

323 new_local.requires_grad_(True) 

324 hsdp_param.sharded_param.requires_grad_(True) 

325 offset += numel 

326 except Exception: # pylint: disable=W0718 

327 for hsdp_param, orig_flat, orig_local in original_locals: 

328 hsdp_param._sharded_param_data = orig_flat 

329 hsdp_param.sharded_param.set_data(orig_local) 

330 hsdp_param.sharded_param._local_tensor = orig_local 

331 self._flat_param_buffer = None 

332 self._flat_cast_buffer = None 

333 return 

334 

335 self._flat_param_buffer = flat_buffer 

336 has_param_dtype = any(p.param_dtype is not None for p in self.hsdp_params) 

337 if has_param_dtype: 

338 cast_dtype = next(p.param_dtype for p in self.hsdp_params if p.param_dtype is not None) 

339 self._flat_cast_buffer = ms.mint.empty( 

340 (total_numel,), dtype=cast_dtype, device=_normalize_device(self.device) 

341 ) 

342 

343 def _is_flat_buffer_valid(self): 

344 """Check if flat buffer still backs the params' sharded data.""" 

345 if self._flat_param_buffer is None or len(self.hsdp_params) == 0: 

346 return False 

347 first_param = self.hsdp_params[0] 

348 return ( 

349 first_param._sharded_param_data.untyped_storage().data_ptr() 

350 == self._flat_param_buffer.untyped_storage().data_ptr() 

351 ) 

352 

353 def _allocate_bucket_buffers_if_needed(self, device, dtype): 

354 normalized_device = _normalize_device(device) 

355 for bucket in self._active_replicate_buckets.values(): 

356 if bucket.flat_numel == 0: 

357 continue 

358 needs_new_buffer = ( 

359 bucket.buffer is None 

360 or bucket.buffer.numel() != bucket.flat_numel 

361 or bucket.buffer.dtype != dtype 

362 ) 

363 if needs_new_buffer: 

364 bucket.buffer = ms.mint.empty((bucket.flat_numel,), dtype=dtype, device=normalized_device) 

365 

366 def _pack_bucket_from_reduce_output(self, bucket: ReplicateBucket) -> ms.Tensor: 

367 if bucket.buffer is None: 

368 raise AssertionError("Bucket buffer must be allocated before packing from reduce output") 

369 if self._reduce_output is None or self._reduce_hsdp_params is None: 

370 raise AssertionError("Bucket packing requires an active fused reduce output") 

371 dst_offset = 0 

372 for idx in bucket.param_indices: 

373 hsdp_param = self._reduce_hsdp_params[idx] 

374 src_offset = self._active_param_flat_offsets[idx] 

375 numel = _shape_numel(hsdp_param.sharded_size) 

376 bucket.buffer.narrow(0, dst_offset, numel).copy_( 

377 self._reduce_output.narrow(0, src_offset, numel) 

378 ) 

379 dst_offset += numel 

380 return bucket.buffer 

381 

382 def _unpack_bucket_to_reduce_output(self, bucket: ReplicateBucket) -> None: 

383 if bucket.buffer is None: 

384 raise AssertionError("Bucket buffer must exist before unpacking to reduce output") 

385 if self._reduce_output is None or self._reduce_hsdp_params is None: 

386 raise AssertionError("Bucket unpack requires an active fused reduce output") 

387 src_offset = 0 

388 for idx in bucket.param_indices: 

389 hsdp_param = self._reduce_hsdp_params[idx] 

390 dst_offset = self._active_param_flat_offsets[idx] 

391 numel = _shape_numel(hsdp_param.sharded_size) 

392 self._reduce_output.narrow(0, dst_offset, numel).copy_( 

393 bucket.buffer.narrow(0, src_offset, numel) 

394 ) 

395 src_offset += numel 

396 

397 def unshard(self, async_op: bool = False): 

398 """Trigger fused all-gather for all parameters in this group.""" 

399 if self._result is not None: 

400 return 

401 if self.shard_world_size == 1: 

402 self._result = AllGatherResult(None, None, None) 

403 return 

404 self.foreach_all_gather(async_op=async_op) 

405 

406 def _init_mp_dtypes(self): 

407 for hsdp_param in self.hsdp_params: 

408 hsdp_param.init_dtype_attrs(self.mp_policy) 

409 trainable_params: list[MindSporeHSDPParamV2] = [ 

410 p for p in self.hsdp_params if p.sharded_param.requires_grad 

411 ] 

412 orig_dtypes = {p.orig_dtype for p in trainable_params} 

413 reduce_dtypes = {p.reduce_dtype for p in trainable_params} 

414 if len(trainable_params) > 0 and len(orig_dtypes) != 1: 

415 raise AssertionError( 

416 f"hsdp expects uniform original parameter dtype but got {orig_dtypes}" 

417 ) 

418 self._orig_dtype = next(iter(orig_dtypes)) if trainable_params else None 

419 if len(trainable_params) > 0 and len(reduce_dtypes) != 1: 

420 raise AssertionError( 

421 f"hsdp expects uniform reduce dtype but got {reduce_dtypes}" 

422 ) 

423 self._reduce_dtype = next(iter(reduce_dtypes)) if trainable_params else None 

424 

425 def wait_for_unshard(self): 

426 """Wait for fused all-gather and materialize per-parameter unsharded views.""" 

427 if self._result is None: 

428 return 

429 if self.shard_world_size == 1: 

430 for hsdp_param in self.hsdp_params: 

431 all_gather_input = hsdp_param.all_gather_inputs[0] 

432 hsdp_param.init_all_gather_outputs( 

433 [all_gather_input.numel()], 

434 [all_gather_input.dtype], 

435 self.shard_world_size, 

436 _normalize_device(self.device), 

437 ) 

438 hsdp_param.alloc_all_gather_outputs() 

439 copy_without_bumping_version(hsdp_param.all_gather_outputs[0], all_gather_input) 

440 self._result = None 

441 else: 

442 self.foreach_all_gather_copy_out() 

443 for hsdp_param in self.hsdp_params: 

444 hsdp_param.init_unsharded_param() 

445 hsdp_param.to_unsharded() 

446 

447 def alloc_all_gather_output(self, total_output_numel, dtype): 

448 """Allocate or resize the fused all-gather output buffer to the specified size and dtype.""" 

449 normalized_device = _normalize_device(self.device) 

450 if self.ag_output is None or self.ag_output.dtype != dtype: 

451 self.ag_output = ms.mint.empty((total_output_numel,), dtype=dtype, device=normalized_device) 

452 return 

453 storage = self.ag_output.untyped_storage() 

454 expected_size = total_output_numel * self.ag_output.itemsize 

455 if storage.size() != expected_size: 

456 storage.resize_(expected_size) 

457 

458 def free_all_gather_output(self): 

459 """Release the fused all-gather output buffer by resizing its storage to zero.""" 

460 if self.ag_output is None: 

461 return 

462 storage = self.ag_output.untyped_storage() 

463 if storage.size() != 0: 

464 storage.resize_(0) 

465 

466 @_no_grad() 

467 def foreach_all_gather(self, async_op=False): 

468 """Perform one fused all-gather across all parameters in the group.""" 

469 if self.metadata_cache is None: 

470 self.metadata_cache = AllGatherMetadataCache() 

471 metadata = self.metadata_cache.get_metadata(self.hsdp_params, get_all_gather_metadata) 

472 if metadata.total_input_numel == 0: 

473 return 

474 world_size = self.shard_world_size 

475 rank = self.shard_rank 

476 total_output_numel = metadata.total_input_numel * world_size 

477 self.alloc_all_gather_output(total_output_numel, metadata.dtype) 

478 for hsdp_param in self.hsdp_params: 

479 hsdp_param.reset_sharded_param() 

480 if self.enable_zero_copy_param_buffer and not self._is_flat_buffer_valid(): 

481 self._init_flat_param_buffer() 

482 

483 use_flat_buffer = ( 

484 self.enable_zero_copy_param_buffer 

485 and self._flat_param_buffer is not None 

486 and self._is_flat_buffer_valid() 

487 ) 

488 if use_flat_buffer: 

489 if self._flat_cast_buffer is not None: 

490 self._flat_cast_buffer.copy_(self._flat_param_buffer) 

491 all_gather_input = self._flat_cast_buffer 

492 else: 

493 all_gather_input = self._flat_param_buffer 

494 else: 

495 all_gather_inputs = [] 

496 for hsdp_param in self.hsdp_params: 

497 all_gather_inputs.extend(hsdp_param.all_gather_inputs) 

498 if len(all_gather_inputs) == 0: 

499 return 

500 all_gather_input, _ = all_gather_copy_in( 

501 all_gather_inputs, 

502 self.ag_output, 

503 metadata.inp_split_sizes, 

504 metadata.total_input_numel, 

505 rank, 

506 ) 

507 handle = dist.all_gather_into_tensor(self.ag_output, all_gather_input, self.shard_group, async_op) 

508 self._result = AllGatherResult(self.ag_output, metadata, handle) 

509 

510 @_no_grad() 

511 def foreach_all_gather_copy_out(self): 

512 """Scatter one fused all-gather result back into per-parameter buffers.""" 

513 ag_output, metadata, handle = self._result 

514 if handle is not None: 

515 handle.wait() 

516 world_size = self.shard_world_size 

517 split_with_sizes_out = [] 

518 for input_numels, input_dtypes, hsdp_param in zip( 

519 metadata.param_input_numels, metadata.param_input_dtypes, self.hsdp_params 

520 ): 

521 hsdp_param.init_all_gather_outputs( 

522 input_numels, 

523 input_dtypes, 

524 world_size, 

525 _normalize_device(ag_output.device), 

526 ) 

527 hsdp_param.alloc_all_gather_outputs() 

528 split_with_sizes_out.extend(hsdp_param.all_gather_outputs) 

529 ag_output = ag_output.view(world_size, -1) 

530 out = [t.view(world_size, -1) for t in split_with_sizes_out] 

531 split_with_sizes_copy(ag_output, metadata.inp_split_sizes, dim=1, out=out) 

532 self._result = None 

533 self.free_all_gather_output() 

534 

535 @_no_grad() 

536 def foreach_reduce( 

537 self, 

538 reduce_scatter_reduce_op: Optional[ops.ReduceOp] = ops.ReduceOp.SUM, 

539 async_op: bool = True, 

540 ) -> Optional[ms.Tensor]: 

541 """Perform fused reduce-scatter and optional bucketed all-reduce.""" 

542 hsdp_params: List[MindSporeHSDPParamV2] = [] 

543 unsharded_grads: List[ms.Tensor] = [] 

544 for hsdp_param in self.hsdp_params: 

545 if not hasattr(hsdp_param, "_unsharded_param"): 

546 continue 

547 if hsdp_param.unsharded_accumulated_grad is not None: 

548 hsdp_params.append(hsdp_param) 

549 unsharded_grads.append(hsdp_param.unsharded_accumulated_grad_data) 

550 elif hsdp_param._unsharded_param.grad is not None: 

551 hsdp_params.append(hsdp_param) 

552 unsharded_grads.append(hsdp_param.unsharded_grad_data) 

553 if not hsdp_params: 

554 return None 

555 grad_dtypes = {g.dtype for g in unsharded_grads} 

556 if len(grad_dtypes) != 1: 

557 raise ValueError( 

558 f"FSDP reduce-scatter expects uniform grad dtype but got {grad_dtypes}" 

559 ) 

560 grad_dtype = unsharded_grads[0].dtype 

561 reduce_dtype = self._reduce_dtype or grad_dtype 

562 world_size = self.shard_world_size 

563 reduce_scatter_input_numel = sum(s.numel() for s in unsharded_grads) 

564 reduce_scatter_output_numel = reduce_scatter_input_numel // world_size 

565 device = _normalize_device(unsharded_grads[0].device) 

566 reduce_scatter_input = ms.mint.empty((reduce_scatter_input_numel,), dtype=reduce_dtype, device=device) 

567 reduce_scatter_copy_in(hsdp_params, unsharded_grads, reduce_scatter_input, world_size) 

568 # Captured here, consumed once in _apply_reduced_grad after all collectives 

569 # complete. Async paths cross method boundaries, so the field is unavoidable. 

570 reduce_output = ms.mint.empty((reduce_scatter_output_numel,), dtype=reduce_dtype, device=device) 

571 self._reduce_op = reduce_scatter_reduce_op 

572 self._reduce_hsdp_params = hsdp_params 

573 self._active_param_flat_offsets = [] 

574 flat_offset = 0 

575 for hsdp_param in hsdp_params: 

576 self._active_param_flat_offsets.append(flat_offset) 

577 flat_offset += _shape_numel(hsdp_param.sharded_size) 

578 self._active_replicate_buckets = self._build_active_replicate_buckets(hsdp_params) 

579 self._allocate_bucket_buffers_if_needed(reduce_output.device, reduce_output.dtype) 

580 self._pending_all_reduce_handles = [] 

581 if self.shard_group is None or world_size <= 1: 

582 comm_ctx.comm_handle = None 

583 self._reduce_output = reduce_scatter_input 

584 if async_op: 

585 comm_ctx.pre_param_group = self 

586 else: 

587 self.apply_fusion_reduced_grad() 

588 return self._reduce_output 

589 apply_gradient_scaling_factor(reduce_scatter_input, self.gradient_scaling_factor) 

590 rs_handle = dist.reduce_scatter_tensor( 

591 output=reduce_output, 

592 input=reduce_scatter_input, 

593 group=self.shard_group, 

594 op=reduce_scatter_reduce_op, 

595 async_op=async_op, 

596 ) 

597 comm_ctx.comm_handle = rs_handle 

598 self._reduce_output = reduce_output 

599 if async_op: 

600 comm_ctx.pre_param_group = self 

601 else: 

602 self.apply_fusion_reduced_grad() 

603 return reduce_output 

604 

605 def wait_reduce_scatter_and_issue_all_reduce(self): 

606 """Wait for reduce-scatter and issue async all-reduces for active buckets.""" 

607 if comm_ctx.comm_handle is not None: 

608 comm_ctx.comm_handle.wait() 

609 comm_ctx.comm_handle = None 

610 if not self._active_replicate_buckets: 

611 self._apply_reduced_grad() 

612 return 

613 self._pending_all_reduce_handles = [] 

614 for bucket in self._active_replicate_buckets.values(): 

615 packed = self._pack_bucket_from_reduce_output(bucket) 

616 ar_handle = dist.all_reduce( 

617 packed, 

618 group=bucket.group, 

619 op=self._reduce_op, 

620 async_op=True, 

621 ) 

622 self._pending_all_reduce_handles.append( 

623 PendingBucketAllReduce(bucket_key=bucket.key, handle=ar_handle) 

624 ) 

625 comm_ctx.all_reduce_param_group = self 

626 

627 def wait_all_reduce_and_apply_grad(self): 

628 """Wait for pending bucket all-reduces and apply reduced grads.""" 

629 for pending in self._pending_all_reduce_handles: 

630 bucket = self._active_replicate_buckets[pending.bucket_key] 

631 pending.handle.wait() 

632 self._unpack_bucket_to_reduce_output(bucket) 

633 self._pending_all_reduce_handles = [] 

634 comm_ctx.all_reduce_handle = None 

635 self._apply_reduced_grad() 

636 

637 def apply_fusion_reduced_grad(self): 

638 """Synchronous fallback: wait, all-reduce buckets, then apply grads.""" 

639 if comm_ctx.comm_handle is not None: 

640 comm_ctx.comm_handle.wait() 

641 comm_ctx.comm_handle = None 

642 for bucket in self._active_replicate_buckets.values(): 

643 packed = self._pack_bucket_from_reduce_output(bucket) 

644 dist.all_reduce( 

645 packed, 

646 group=bucket.group, 

647 op=self._reduce_op, 

648 ) 

649 self._unpack_bucket_to_reduce_output(bucket) 

650 self._apply_reduced_grad() 

651 

652 def _apply_reduced_grad(self): 

653 """Write reduced gradients from the fused output buffer back to params.""" 

654 flat_grad_offset = 0 

655 if self._reduce_hsdp_params is None or self._reduce_output is None: 

656 return 

657 # All collectives have completed; scale once on the fused buffer right 

658 # before slicing it into per-parameter sharded grads. 

659 for hsdp_param in self._reduce_hsdp_params: 

660 shard_numel = _shape_numel(hsdp_param.sharded_size) 

661 new_sharded_grad = self._reduce_output.narrow(0, flat_grad_offset, shard_numel) 

662 hsdp_param.apply_reduced_grad(new_sharded_grad, self._orig_dtype) 

663 flat_grad_offset += shard_numel 

664 self._reduce_output = None 

665 self._reduce_hsdp_params = None 

666 self._active_param_flat_offsets = [] 

667 self._active_replicate_buckets = {} 

668 self._pending_all_reduce_handles = [] 

669 

670 

671class AllReduceParamGroup: 

672 """Groups HSDP parameters by replicate group for fused async all-reduce.""" 

673 

674 ALIGNMENT_BYTES = 512 

675 

676 @staticmethod 

677 def _resolve_reduce_dtype( 

678 reduce_dtype: Any, 

679 hsdp_params: List[MindSporeHSDPParamV2], 

680 orig_dtypes: List[Any], 

681 ) -> Any: 

682 """Resolve None reduce_dtype to match ``reduce_scatter_grad``'s ``dtype or grad.dtype``.""" 

683 if reduce_dtype is not None: 

684 return reduce_dtype 

685 for hsdp_param in hsdp_params: 

686 if getattr(hsdp_param, "unsharded_accumulated_grad", None) is not None: 

687 return hsdp_param.unsharded_accumulated_grad_data.dtype 

688 unsharded_param = getattr(hsdp_param, "unsharded_param", None) 

689 if unsharded_param is not None and getattr(unsharded_param, "grad", None) is not None: 

690 return hsdp_param.unsharded_grad_data.dtype 

691 return orig_dtypes[0] if orig_dtypes else None 

692 

693 def __init__( 

694 self, 

695 replicate_group, 

696 hsdp_params: List[MindSporeHSDPParamV2], 

697 orig_dtypes: List[Any], 

698 reduce_dtype: Any, 

699 reduce_op: ops.ReduceOp, 

700 mp_policy: Optional[MixedPrecisionPolicy] = None, 

701 replicate_world_size: Optional[int] = None, 

702 ): 

703 self.replicate_group = replicate_group 

704 self.hsdp_params = hsdp_params 

705 self.orig_dtypes = orig_dtypes 

706 self.reduce_dtype = self._resolve_reduce_dtype(reduce_dtype, hsdp_params, orig_dtypes) 

707 self.reduce_op = reduce_op 

708 self.mp_policy = mp_policy 

709 if replicate_world_size is not None: 

710 self.replicate_world_size = replicate_world_size 

711 elif replicate_group is not None and hasattr(replicate_group, "rank_size"): 

712 self.replicate_world_size = replicate_group.rank_size 

713 elif hsdp_params: 

714 self.replicate_world_size = hsdp_params[0].unsharded_group_info.rank_size 

715 else: 

716 self.replicate_world_size = 1 

717 self.fused_buffer: Optional[ms.Tensor] = None 

718 self.param_offsets: List[int] = [] 

719 self.param_numels: List[int] = [] 

720 self.all_reduce_handle: Optional[CommHandle] = None 

721 

722 def compute_aligned_layout(self) -> int: 

723 """Compute fused buffer layout with trailing 512-byte alignment.""" 

724 self.param_offsets = [] 

725 self.param_numels = [] 

726 element_size = int(ms.Tensor([], dtype=self.reduce_dtype).itemsize) 

727 current_offset = 0 

728 for hsdp_param in self.hsdp_params: 

729 numel = _shape_numel(hsdp_param.sharded_size) 

730 self.param_numels.append(numel) 

731 self.param_offsets.append(current_offset) 

732 current_offset += numel 

733 total_bytes = current_offset * element_size 

734 aligned_total_bytes = ( 

735 (total_bytes + self.ALIGNMENT_BYTES - 1) // self.ALIGNMENT_BYTES 

736 ) * self.ALIGNMENT_BYTES 

737 return aligned_total_bytes // element_size 

738 

739 def allocate_fused_buffer(self, device: Any) -> None: 

740 """Allocate the fused buffer with computed layout.""" 

741 total_numel = self.compute_aligned_layout() 

742 normalized_device = _normalize_device(device) 

743 self.fused_buffer = ms.mint.empty( 

744 (total_numel,), dtype=self.reduce_dtype, device=normalized_device 

745 ) 

746 self.fused_buffer.zero_() 

747 

748 def get_param_buffer_view(self, idx: int) -> ms.Tensor: 

749 """Return a flat view for reduce_scatter output of parameter idx.""" 

750 if self.fused_buffer is None: 

751 raise RuntimeError("Fused buffer not allocated. Call allocate_fused_buffer first.") 

752 offset = self.param_offsets[idx]