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

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"""HSDP parameter""" 

16from typing import List, Callable, Optional, cast, Tuple 

17import itertools 

18import mindspore as ms 

19from mindspore import nn 

20from mindspore.common.api import _no_grad 

21from mindspore import ops, Parameter 

22import mindspore.mint.distributed as dist 

23from mindspore.ops.function.comm_func import CommHandle 

24from hyper_parallel.core.fully_shard.utils import ( 

25 MixedPrecisionPolicy, 

26 CPUOffloadPolicy, 

27 OffloadPolicy, 

28 FSDPMeshInfo, 

29 HSDPMeshInfo, 

30) 

31from hyper_parallel.core.dtensor.dtensor import DTensor 

32from hyper_parallel.core.dtensor.layout import Layout 

33from hyper_parallel.core.fully_shard.hsdp_param import HSDPParamV2 

34from hyper_parallel.core.fully_shard.hsdp_utils import ( 

35 ShardedState, 

36 FullyShardParamMode, 

37 apply_gradient_scaling_factor, 

38 unwrap_dtensor_param, 

39) 

40from hyper_parallel.core.dtensor.placement_types import Shard, StridedShard 

41from hyper_parallel.core.fully_shard.hsdp_utils import ParamModuleInfo 

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

43from hyper_parallel.platform.mindspore.utils import normalize_runtime_device 

44from hyper_parallel.platform.mindspore.fully_shard.pack_utils import ( 

45 build_rs_plan, 

46 pack_for_reduce_scatter, 

47 unpack_from_all_gather, 

48) 

49 

50 

51def _pack_for_reduce_scatter(local_tensor: ms.Tensor, shard_dim: int, world_size: int) -> ms.Tensor: 

52 """Pack one local gradient into the row-major reduce-scatter layout. 

53 

54 MindSpore currently aligns with the torch non-comm-fusion V1 path: 

55 

56 - shard on dim 0: identity flatten 

57 - shard on non-dim0: chunk on shard dim, then concatenate on dim 0 

58 """ 

59 if world_size <= 1 or shard_dim == 0: 

60 return local_tensor 

61 chunks = ms.mint.chunk(local_tensor, world_size, dim=shard_dim) 

62 return ms.mint.cat(chunks, dim=0).contiguous() 

63 

64 

65def _to_dtype_if_needed( 

66 tensor: ms.Tensor, dtype: Optional[ms.Type] 

67) -> ms.Tensor: 

68 """Cast tensor to the given dtype if it differs from current dtype.""" 

69 if isinstance(dtype, ms.Type) and tensor.dtype != dtype: 

70 return tensor.to(dtype) 

71 return tensor 

72 

73 

74def make_contiguous_strides_for(shape, row_major=True): 

75 """ 

76 Compute strides for a contiguous tensor of the given shape. 

77 

78 Args: 

79 shape (tuple of int): The shape of the tensor. Each dimension must be a non-negative integer. 

80 row_major (bool):  

81 - If True (default), returns C-style (row-major) strides: last dimension changes fastest. 

82 - If False, returns strides where the last two dimensions are Fortran-style  

83 (i.e., for batched matrix operations in BLAS/LAPACK): second-to-last dim changes fastest. 

84 

85 Returns: 

86 tuple of int: The computed strides. 

87 

88 Examples: 

89 >>> make_contiguous_strides_for((2, 3, 4)) 

90 (12, 4, 1) 

91 >>> make_contiguous_strides_for((2, 3, 4), row_major=False) 

92 (12, 1, 3) 

93 >>> make_contiguous_strides_for((5,)) 

94 (1,) 

95 >>> make_contiguous_strides_for((5,), row_major=False) 

96 (1,) 

97 >>> make_contiguous_strides_for(()) 

98 () 

99 """ 

100 if not isinstance(shape, (tuple, list)): 

101 raise TypeError("shape must be a tuple or list of non-negative integers") 

102 

103 # Validate shape elements 

104 for dim in shape: 

105 if not isinstance(dim, int) or dim < 0: 

106 raise ValueError("All dimensions in shape must be non-negative integers") 

107 

108 if not shape: 

109 return () 

110 

111 # Compute C-style (row-major) strides: stride[i] = product(shape[i+1:]) 

112 strides = [] 

113 multiplier = 1 

114 # Traverse shape in reverse order 

115 for size in reversed(shape): 

116 strides.append(multiplier) 

117 multiplier *= max(size, 1) # handle size=0 gracefully (treat as 1 for stride calc) 

118 

119 # Reverse to get correct order 

120 c_strides = tuple(reversed(strides)) 

121 

122 if row_major: 

123 return c_strides 

124 # For column-major: only affect last two dimensions 

125 if len(shape) < 2: 

126 return c_strides 

127 # In Fortran-style for matrices: 

128 # stride of last dim = 1 

129 # stride of second-to-last dim = shape[-1] 

130 # But note: in batched case (..., M, N), we want strides (..., N, 1) → wait! 

131 # However, the original PyTorch logic returns: result[:-2] + (1, max(shape[-2], 1)) 

132 # Let's follow that exactly: 

133 # Example: shape=(B, M, N) → c_strides=(M*N, N, 1) 

134 # col-major → (M*N, 1, M) 

135 # So: keep all but last two, then (1, shape[-2]) 

136 return c_strides[:-2] + (1, max(shape[-2], 1)) 

137 

138 

139class MindSporeHSDPParamV2(HSDPParamV2): 

140 """ 

141 MindSpore HSDP parameter. 

142 """ 

143 

144 def __init__( 

145 self, 

146 param: Parameter, 

147 module_info: ParamModuleInfo, 

148 mesh_info: FSDPMeshInfo, 

149 shard_placement_fn: Optional[Callable[[Parameter], Optional[Shard]]] = None, 

150 mp_policy: Optional[MixedPrecisionPolicy] = None, 

151 offload_policy: Optional[OffloadPolicy] = None, 

152 device: Optional[str] = None, 

153 param_mode: Optional[FullyShardParamMode] = None, 

154 enable_fsdp_shard: bool = True, 

155 ): 

156 self._module_info: ParamModuleInfo = module_info 

157 self.mesh_info = mesh_info 

158 self.mp_policy = mp_policy 

159 self.device = device 

160 if param_mode is None: 

161 raise AssertionError("param_mode must be resolved before MindSporeHSDPParamV2 initialization.") 

162 self.param_mode = param_mode 

163 self.enable_fsdp_shard = enable_fsdp_shard 

164 self.offload_to_cpu: bool = isinstance(offload_policy, CPUOffloadPolicy) 

165 self.pin_memory = ( 

166 self.offload_to_cpu and cast(CPUOffloadPolicy, offload_policy).pin_memory 

167 ) 

168 self._orig_param_hooks: List[Callable] = [] 

169 self.grad_offload_event: Optional[ms.runtime.Event] = None 

170 dtensor_payload = unwrap_dtensor_param(param) 

171 self._orig_param_is_dtensor = dtensor_payload is not None 

172 self._orig_dtensor_mesh = dtensor_payload.device_mesh if dtensor_payload is not None else None 

173 self._orig_dtensor_placements = ( 

174 tuple(dtensor_payload.placements) if dtensor_payload is not None else None 

175 ) 

176 self._spmd_shard_mesh_dim = getattr(self.mesh_info, "shard_mesh_dim", None) 

177 self._spmd_replicate_mesh_dim = getattr(self.mesh_info, "replicate_mesh_dim", None) 

178 self._init_sharded_param(param, shard_placement_fn) 

179 self._init_group_infos() 

180 self._save_backward_hooks(param) 

181 self.all_gather_outputs: List[ms.Tensor] = [] 

182 self.unsharded_accumulated_grad = None 

183 self._unsharded_param: Optional[Parameter] = None 

184 self._param_fqn: Optional[str] = None 

185 # Communication attributes for prefetch pattern 

186 self.prefetch_handle: Optional[CommHandle] = None 

187 self._reduce_scatter_output = None 

188 self.reduce_scatter_handle: Optional[CommHandle] = None 

189 self._all_reduce_output = None 

190 self.all_reduce_handle: Optional[CommHandle] = None 

191 self._accumulated_allreduced_grad = True 

192 self._post_load_hook_handle = ( 

193 module_info.module.register_load_state_dict_post_hook( 

194 lambda *args, **kwargs: self.reset_sharded_param() 

195 ) 

196 ) 

197 self.gradient_scaling_factor = None 

198 

199 @property 

200 def accumulated_allreduced_grad(self) -> bool: 

201 return self._accumulated_allreduced_grad 

202 

203 @accumulated_allreduced_grad.setter 

204 def accumulated_allreduced_grad(self, value: bool) -> None: 

205 self._accumulated_allreduced_grad = value 

206 

207 @property 

208 def uses_param_shard(self) -> bool: 

209 """Whether FSDP sharding is enabled for this parameter.""" 

210 return self.enable_fsdp_shard 

211 

212 @property 

213 def is_dtensor_compat_mode(self) -> bool: 

214 """Whether this parameter uses DTensor compatibility mode.""" 

215 return self.param_mode == FullyShardParamMode.DTENSOR_COMPAT 

216 

217 def _get_data_parallel_shard_placement(self, placements: list, shard_placement: Shard): 

218 """Return the explicit fully_shard placement on the unified SPMD mesh.""" 

219 split_factor = 1 

220 shard_mesh_dim = getattr(self, "_spmd_shard_mesh_dim", None) 

221 for mesh_idx, placement in enumerate(placements): 

222 if mesh_idx == shard_mesh_dim: 

223 continue 

224 if placement.is_shard(shard_placement.dim): 

225 split_factor *= self._spmd_mesh.mesh_shape[mesh_idx] 

226 if split_factor > 1: 

227 return StridedShard(shard_placement.dim, split_factor=split_factor) 

228 return shard_placement 

229 

230 def _release_full_param_storage_if_safe(self, param_data: ms.Tensor) -> None: 

231 """Release the temporary full-parameter storage once the sharded param is installed. 

232 

233 Skip storage reclamation only for meta tensors. Both plain Tensor inputs and DTensor local 

234 tensors should drop their original storage after the sharded Parameter has been installed 

235 onto the owning modules. 

236 """ 

237 if param_data.is_meta: 

238 return 

239 storage = param_data.untyped_storage() 

240 if storage.size() != 0: 

241 storage.resize_(0) 

242 

243 def _iter_backward_hooks(self, param: Parameter) -> List[Callable]: 

244 """Return backward hooks registered on a MindSpore Tensor/Parameter.""" 

245 hooks_getter = getattr(param, "hooks", None) 

246 if callable(hooks_getter): 

247 try: 

248 return list(hooks_getter()) 

249 except (AttributeError, RuntimeError, TypeError, ValueError): 

250 pass 

251 

252 backward_hooks = getattr(param, "_backward_hooks", None) 

253 if backward_hooks is None: 

254 return [] 

255 if hasattr(backward_hooks, "values"): 

256 return list(backward_hooks.values()) 

257 return list(backward_hooks) 

258 

259 def _save_backward_hooks(self, param: Parameter) -> None: 

260 """Save user-registered parameter backward hooks for later parameter swaps.""" 

261 if not hasattr(self, "_orig_param_hooks"): 

262 self._orig_param_hooks = [] 

263 if not hasattr(self, "_saved_hook_ids"): 

264 self._saved_hook_ids = set() 

265 

266 for hook_func in self._iter_backward_hooks(param): 

267 hook_func_id = id(hook_func) 

268 if hook_func_id not in self._saved_hook_ids: 

269 self._orig_param_hooks.append(hook_func) 

270 self._saved_hook_ids.add(hook_func_id) 

271 

272 def _migrate_backward_hooks(self, new_param: Parameter) -> None: 

273 """Migrate saved user backward hooks to the active sharded/unsharded parameter.""" 

274 if not getattr(self, "_orig_param_hooks", None): 

275 return 

276 if hasattr(new_param, "migrate_backward_hooks_run_once"): 

277 return 

278 register_hook = getattr(new_param, "register_hook", None) 

279 if not callable(register_hook): 

280 return 

281 

282 for hook_func in self._orig_param_hooks: 

283 try: 

284 if getattr(new_param, "requires_grad", False): 

285 register_hook(hook_func) 

286 except (RuntimeError, TypeError, ValueError): 

287 pass 

288 new_param.migrate_backward_hooks_run_once = True 

289 

290 @_no_grad() 

291 def _init_sharded_param( 

292 self, 

293 param: Parameter, 

294 shard_placement_fn: Optional[Callable], 

295 ) -> None: 

296 param_device = normalize_runtime_device(param.device) 

297 if param_device not in ("meta", self.device): 

298 raise AssertionError( 

299 f"Expects the parameter to already be moved to device {self.device} but got {param.device}" 

300 ) 

301 hsdp_placement = shard_placement_fn(param) if shard_placement_fn else None 

302 if hsdp_placement is None: 

303 hsdp_placement = Shard(0) 

304 elif hsdp_placement.dim < 0: 

305 # if dim is negative, add the number of dimensions of the parameter 

306 hsdp_placement = Shard(hsdp_placement.dim + param.ndim) 

307 

308 if not isinstance(hsdp_placement, Shard): 

309 raise AssertionError( 

310 f"Expected Shard, got {type(hsdp_placement)}: {hsdp_placement}" 

311 ) 

312 

313 self.hsdp_placement = hsdp_placement 

314 base_placements = list(self._get_base_spmd_placements()) 

315 self._spmd_placements = self._apply_data_parallel_placements(base_placements, hsdp_placement) 

316 param_data = unwrap_dtensor_param(param).to_local() if self._orig_param_is_dtensor else param 

317 

318 shard_dim = hsdp_placement.dim 

319 self._orig_size = param_data.shape 

320 self._contiguous_orig_stride = make_contiguous_strides_for(self._orig_size) 

321 

322 if self.uses_param_shard and isinstance(self.mesh_info, FSDPMeshInfo): # FSDP or HSDP 

323 shard_rank = self.mesh_info.shard_mesh_rank 

324 shard_world_size = self.mesh_info.shard_mesh_size 

325 else: # DDP 

326 shard_rank = 0 

327 shard_world_size = 1 

328 

329 self.is_sharded = bool(self.uses_param_shard and shard_world_size > 1) 

330 

331 if param_data.shape[shard_dim] % shard_world_size != 0: 

332 raise NotImplementedError( 

333 f"Uneven sharding on dim {shard_dim} not supported: " 

334 f"shape={param_data.shape}, world_size={shard_world_size}" 

335 ) 

336 chunks = ms.mint.chunk(param_data, shard_world_size, dim=shard_dim) 

337 sharded_param = chunks[shard_rank].clone().contiguous() 

338 self.sharded_size = sharded_param.shape 

339 self.contiguous_sharded_stride = make_contiguous_strides_for(self.sharded_size) 

340 self._sharded_param_data = sharded_param.view(-1) 

341 

342 self._sharding_spec = Layout.from_device_mesh(self._spmd_mesh) 

343 self._sharding_spec.set_placements(self._spmd_placements) 

344 self._sharding_spec.placement_to_tensor_map(param.ndim) 

345 

346 shard_dtensor = DTensor.from_local(sharded_param, self._spmd_mesh, self._spmd_placements) 

347 self.sharded_param = Parameter(shard_dtensor, name=param.name) 

348 set_requires_grad_if_needed(param, self.sharded_param) 

349 self.sharded_param.grad = None 

350 

351 self._setattr_on_modules(self.sharded_param) 

352 self._release_full_param_storage_if_safe(param_data) 

353 self.sharded_param._hsdp_param_initialized = True 

354 self.sharded_state = ShardedState.SHARDED 

355 self.param_dtype = None 

356 

357 def init_dtype_attrs(self, mp_policy: MixedPrecisionPolicy): 

358 param_dtype, reduce_dtype = (mp_policy.param_dtype, mp_policy.reduce_dtype) 

359 self.orig_dtype = self.sharded_param.dtype 

360 if reduce_dtype == param_dtype: 

361 reduce_dtype = None 

362 if param_dtype == self.orig_dtype: 

363 param_dtype = None 

364 self.param_dtype = param_dtype 

365 self.reduce_dtype = reduce_dtype 

366 

367 def init_all_gather_outputs( 

368 self, 

369 all_gather_input_numels: list[int], 

370 all_gather_input_dtypes: list[ms.Type], 

371 world_size: int, 

372 device: str, 

373 force_recreate: bool = False, 

374 ): 

375 if not force_recreate and len(self.all_gather_outputs) > 0: 

376 return # already initialized 

377 self.all_gather_outputs = [ 

378 ms.mint.empty([numel * world_size], dtype=dtype, device=device.split(':')[0]) 

379 for numel, dtype in zip(all_gather_input_numels, all_gather_input_dtypes) 

380 ] 

381 

382 def init_unsharded_param(self): 

383 """ 

384 Initialize unsharded parameter from all-gather outputs. 

385 

386 This reconstructs the full parameter after all-gather by unpacking the 

387 gathered flat buffer back to the original tensor layout. 

388 """ 

389 unsharded_param = self._get_unsharded_param_from_all_gather_output() 

390 if self._unsharded_param is not None: 

391 # Keep the Parameter identity stable across forward-reshard-backward 

392 # cycles so backward hooks continue to read gradients from the same 

393 # object that participated in the forward graph. 

394 if self._orig_param_is_dtensor: 

395 self._unsharded_param.set_data(unsharded_param) 

396 else: 

397 self._unsharded_param.data = unsharded_param 

398 set_requires_grad_if_needed(self.sharded_param, self._unsharded_param) 

399 self._unsharded_param.grad = None 

400 return 

401 if self._orig_param_is_dtensor: 

402 self._unsharded_param = Parameter( 

403 unsharded_param, 

404 name=self.sharded_param.name, 

405 requires_grad=self.sharded_param.requires_grad, 

406 ) 

407 return 

408 # For MindSpore, if use `Parameter(tensor)`, Parameter will create a new Tensor instead of a view. 

409 # Here we need to share storage, so we use the `.data = tensor` approach to create shared storage. 

410 self._unsharded_param = Parameter( 

411 [], 

412 name=self.sharded_param.name, 

413 requires_grad=False, 

414 ) 

415 self._unsharded_param.data = unsharded_param 

416 if self.sharded_param.requires_grad: 

417 self._unsharded_param.requires_grad = True 

418 

419 def _get_unsharded_param_from_all_gather_output(self): 

420 """Reconstruct the full local parameter view from the packed all-gather output.""" 

421 if len(self.all_gather_outputs) != 1: 

422 raise AssertionError( 

423 f"Expected 1 all_gather_output, got {len(self.all_gather_outputs)}" 

424 ) 

425 unsharded_tensor = self.all_gather_outputs[0] 

426 plan = build_rs_plan( 

427 self, 

428 self._sharded_local_tensor, 

429 self.shard_world_size if self.is_sharded else 1, 

430 ) 

431 unsharded_param = unpack_from_all_gather(unsharded_tensor, plan) 

432 if getattr(self, "_orig_param_is_dtensor", False): 

433 unsharded_param = DTensor.from_local( 

434 unsharded_param, 

435 self._orig_dtensor_mesh, 

436 self._orig_dtensor_placements, 

437 ) 

438 return unsharded_param 

439 

440 def to_sharded(self) -> None: 

441 if not self.uses_param_shard and self._unsharded_param is not None: 

442 # Replicate params keep the same local shape across shard/unshard, 

443 # so persist forward-time state updates before switching objects. 

444 src = self._unsharded_param.to_local() if isinstance(self._unsharded_param, DTensor) \ 

445 else self._unsharded_param 

446 dst = self.sharded_param.to_local() if isinstance(self.sharded_param, DTensor) else self.sharded_param 

447 copy_without_bumping_version(dst, src) 

448 self._setattr_on_modules(self.sharded_param) 

449 self.free_unsharded_param() 

450 self.sharded_state = ShardedState.SHARDED 

451 

452 def to_unsharded(self) -> None: 

453 set_requires_grad_if_needed(self.sharded_param, self._unsharded_param) 

454 self._setattr_on_modules(self._unsharded_param) 

455 self.sharded_state = ShardedState.UNSHARDED 

456 

457 def _setattr_on_modules(self, param: Parameter) -> None: 

458 if getattr(self._module_info.module.__setattr__, "__func__", None) is nn.Cell.__setattr__: 

459 # fast path 

460 self._module_info.module._params[self._module_info.param_name] = param 

461 else: 

462 # slow path 

463 setattr(self._module_info.module, self._module_info.param_name, param) 

464 if hasattr(self, "sharded_param"): 

465 self._save_backward_hooks(self.sharded_param) 

466 self._migrate_backward_hooks(param) 

467 

468 # Iterate through all modules that share this parameter to prevent pointer desync. 

469 for shared_module, shared_param_name in zip( 

470 self._module_info.shared_modules, self._module_info.shared_param_names 

471 ): 

472 if getattr(shared_module.__setattr__, "__func__", None) is nn.Cell.__setattr__: 

473 shared_module._params[shared_param_name] = param 

474 else: 

475 setattr(shared_module, shared_param_name, param) 

476 

477 def to_sharded_dtensor(self, tensor: ms.Tensor) -> DTensor: 

478 """ 

479 Converts a local tensor representing either the sharded parameter or 

480 sharded gradient to DTensor. 

481 """ 

482 return DTensor.from_local( 

483 tensor, 

484 self._sharding_spec.mesh, 

485 self._sharding_spec.placements 

486 ) 

487 

488 def _to_local_unsharded_grad(self, grad): 

489 """Normalize a pending gradient to the local tensor expected by fully_shard collectives.""" 

490 return self._normalize_unsharded_grad_to_local(grad, reduce_partial_dtensor=False) 

491 

492 def to_accumulated_grad_if_needed(self) -> None: 

493 if self._unsharded_param.grad is None: 

494 return 

495 unsharded_grad = self._unsharded_param.grad 

496 self._unsharded_param.grad = None 

497 if self.reduce_dtype is not None and unsharded_grad.dtype != self.reduce_dtype: 

498 unsharded_grad = unsharded_grad.to(self.reduce_dtype) 

499 if self.unsharded_accumulated_grad is None: 

500 self.unsharded_accumulated_grad = unsharded_grad 

501 else: 

502 self.unsharded_accumulated_grad += unsharded_grad 

503 

504 def accumulate_unsharded_grad_if_needed(self) -> None: 

505 if ( 

506 self.unsharded_accumulated_grad is not None 

507 and self.unsharded_param.grad is not None 

508 ): 

509 # need to handle the gradient 

510 self.unsharded_accumulated_grad += self._to_local_unsharded_grad(self.unsharded_param.grad) 

511 self.unsharded_param.grad = None 

512 

513 def alloc_all_gather_outputs(self) -> None: 

514 for tensor in self.all_gather_outputs: 

515 expected_size = tensor.numel() * tensor.itemsize 

516 

517 storage = tensor.untyped_storage() 

518 if storage.size() != expected_size: 

519 storage.resize_(expected_size) 

520 

521 def free_unsharded_param(self) -> None: 

522 for tensor in itertools.chain( 

523 self.all_gather_outputs 

524 ): 

525 storage = tensor.untyped_storage() 

526 if storage.size() != 0: 

527 storage.resize_(0) 

528 

529 @property 

530 def all_gather_inputs(self) -> list[ms.Tensor]: 

531 self._assert_in_states(ShardedState.SHARDED) 

532 sharded_param_data = self._sharded_param_data 

533 if self.offload_to_cpu: 

534 sharded_param_data = sharded_param_data.to( 

535 self.device, non_blocking=True 

536 ) 

537 if self.param_dtype is not None and self.param_dtype != sharded_param_data.dtype: 

538 return [sharded_param_data.to(self.param_dtype)] 

539 return [sharded_param_data] 

540 

541 @property 

542 def unsharded_param(self) -> Parameter: 

543 """Return the full unsharded parameter after all-gather.""" 

544 return self._unsharded_param 

545 

546 @property 

547 def unsharded_grad_data(self) -> ms.Tensor: 

548 """ 

549 Get the unsharded gradient data as a local tensor. 

550 """ 

551 grad = self.unsharded_param.grad 

552 if grad is None: 

553 raise AssertionError("Expects unsharded_param.grad to not be None") 

554 return self._to_local_unsharded_grad(grad) 

555 

556 @property 

557 def unsharded_accumulated_grad_data(self) -> ms.Tensor: 

558 """ 

559 Get the unsharded accumulated gradient data as a local tensor. 

560 """ 

561 grad = self.unsharded_accumulated_grad 

562 return grad 

563 

564 @property 

565 def _sharded_local_tensor(self) -> ms.Tensor: 

566 """Return the underlying local tensor of the sharded DTensor parameter.""" 

567 return cast(DTensor, self.sharded_param)._local_tensor 

568 

569 def _sharded_param_storage_dtype(self) -> Optional[ms.Type]: 

570 """Return the dtype of the sharded parameter's on-device storage.""" 

571 if not hasattr(self.sharded_param, "dtype"): 

572 return None 

573 dtype = self.sharded_param.dtype 

574 if isinstance(dtype, ms.Type): 

575 return dtype 

576 return None 

577 

578 @property 

579 def shard_world_size(self) -> int: 

580 """Get the world size for shard dimension.""" 

581 if isinstance(self.mesh_info, FSDPMeshInfo): 

582 return self.mesh_info.shard_mesh_size 

583 return 1 

584 

585 @property 

586 def replicate_world_size(self) -> int: 

587 """Get the world size for replicate dimension (HSDP only).""" 

588 if isinstance(self.mesh_info, HSDPMeshInfo): 

589 return self.mesh_info.replicate_mesh_size 

590 return 1 

591 

592 def _assert_in_states(self, *states: ShardedState) -> None: 

593 """Assert current state is one of expected states.""" 

594 if self.sharded_state not in states: 

595 raise AssertionError( 

596 f"Expected sharded_state in {states}, got {self.sharded_state}" 

597 ) 

598 

599 def _is_same_sharded_local_tensor(self, local_tensor: ms.Tensor) -> bool: 

600 """Whether the cached flat shard view already points to the ``local_tensor`` storage.""" 

601 if not isinstance(self._sharded_param_data, ms.Tensor): 

602 return False 

603 cached_storage = self._sharded_param_data.untyped_storage() 

604 local_storage = local_tensor.untyped_storage() 

605 # when sharding param with shape (1, ...) over 2 ranks 

606 # local_tensor on rank 1 can be size 0, data_ptr() can be 0 

607 return ( 

608 cached_storage.data_ptr() > 0 

609 and cached_storage.data_ptr() == local_storage.data_ptr() 

610 ) 

611 

612 def _validate_sharded_local_tensor_shape(self, local_tensor: ms.Tensor) -> None: 

613 """Validate that a replaced local tensor still matches the expected shard shape.""" 

614 if local_tensor.shape != self.sharded_size: 

615 raise AssertionError( 

616 f"Expected sharded_size to be {self.sharded_size}, got {local_tensor.shape}" 

617 ) 

618 

619 def _pin_sharded_local_tensor_if_needed(self, local_tensor: ms.Tensor) -> Tuple[ms.Tensor, bool]: 

620 """Pin the local tensor memory when CPU offload requires it.""" 

621 if self.pin_memory and not local_tensor.is_pinned(): 

622 return local_tensor.to("cpu").pin_memory(), True 

623 return local_tensor, False 

624 

625 def _assert_sharded_param_is_dtensor(self) -> None: 

626 """Assert that ``self.sharded_param`` is backed by a DTensor.""" 

627 if not isinstance(self.sharded_param, DTensor): 

628 raise AssertionError(f"Expected DTensor, got {type(self.sharded_param)}") 

629 

630 def _refresh_sharded_local_tensor_view( 

631 self, 

632 local_tensor: ms.Tensor, 

633 shard_dim: int, 

634 length: int, 

635 ) -> None: 

636 """Refresh ``self.sharded_param`` to point to a local tensor view.""" 

637 # Only change the local tensor object if needed 

638 with _no_grad(): 

639 local_view = local_tensor.narrow(dim=shard_dim, start=0, length=length) 

640 set_requires_grad_if_needed(self.sharded_param, local_view) 

641 self.sharded_param._local_tensor = local_view 

642 if not self.sharded_param._local_tensor.is_contiguous(): 

643 raise AssertionError( 

644 "Expected sharded_param._local_tensor to be contiguous" 

645 ) 

646 

647 def reset_sharded_param(self) -> None: 

648 """Reset the sharded param after ``load_state_dict``.""" 

649 module_info = self._module_info 

650 new_param = getattr(module_info.module, module_info.param_name) 

651 if new_param is not self.sharded_param: 

652 if isinstance(new_param, DTensor): 

653 self.sharded_param = new_param 

654 if not getattr(self.sharded_param, "_hsdp_param_initialized", None): 

655 # reset _hsdp_param_initialized flag. 

656 self.sharded_param._hsdp_param_initialized = True 

657 elif isinstance(new_param, ms.Tensor): 

658 # if new_param is Tensor, don't re-ref 'self.sharded_param' 

659 # just update self.sharded_param._local_tensor and self.sharded_param_data. 

660 pass 

661 

662 local_tensor = new_param._local_tensor if isinstance(new_param, DTensor) else new_param 

663 if local_tensor.is_meta: 

664 return 

665 # local_tensor can be padded twice 

666 # 1st time in fully_shard(model) 

667 # 2nd time in model(input) lazy_init 

668 # 2nd time should be no-op if parameters remain unchanged 

669 # 2nd time shouldn't be no-op if people call model.load_state_dict(...) before lazy_init 

670 # this makes it possible for trainer to call `sd = model.state_dict()` before the training loop 

671 # and use `sd` without calling .state_dict() per iteration 

672 same_local_tensor = self._is_same_sharded_local_tensor(local_tensor) 

673 shard_dim = self.hsdp_placement.dim 

674 length = local_tensor.shape[shard_dim] if local_tensor.numel() > 0 else 0 

675 if not same_local_tensor: 

676 self._validate_sharded_local_tensor_shape(local_tensor) 

677 local_tensor, pinned_local_tensor = self._pin_sharded_local_tensor_if_needed(local_tensor) 

678 updated_local_tensor = not same_local_tensor or pinned_local_tensor 

679 if not same_local_tensor: 

680 self._sharded_param_data = local_tensor.view(-1) 

681 self._assert_sharded_param_is_dtensor() 

682 if updated_local_tensor: 

683 self._refresh_sharded_local_tensor_view(local_tensor, shard_dim, length) 

684 self._sharding_spec = cast(DTensor, self.sharded_param).layout 

685 

686 @_no_grad() 

687 def _get_unsharded_param_data(self, async_op: bool = False) -> Tuple[ms.Tensor, Optional[CommHandle]]: 

688 """ 

689 Perform all-gather to get unsharded parameter data. 

690 

691 Args: 

692 async_op: Whether to execute asynchronously. 

693 

694 Returns: 

695 (unsharded_param, handle): Unsharded parameter data and communication handle. 

696 """ 

697 # Optimizer steps may refresh the underlying local tensor storage. Re-sync 

698 # the cached flat shard view before reading all_gather_inputs for the next 

699 # unshard cycle. 

700 self.reset_sharded_param() 

701 all_gather_input = self.all_gather_inputs[0] 

702 

703 # If parameter is not sharded (below threshold), no communication needed 

704 if not self.is_sharded: 

705 self.init_all_gather_outputs( 

706 all_gather_input_numels=[all_gather_input.numel()], 

707 all_gather_input_dtypes=[all_gather_input.dtype], 

708 world_size=1, 

709 device=all_gather_input.device.split(':')[0], 

710 ) 

711 self.alloc_all_gather_outputs() 

712 copy_without_bumping_version(self.all_gather_outputs[0], all_gather_input) 

713 return self.all_gather_outputs[0], None 

714 

715 # Initialize output buffer 

716 self.init_all_gather_outputs( 

717 all_gather_input_numels=[all_gather_input.numel()], 

718 all_gather_input_dtypes=[all_gather_input.dtype], 

719 world_size=self.shard_world_size, 

720 device=self._sharded_param_data.device.split(':')[0], 

721 ) 

722 self.alloc_all_gather_outputs() 

723 

724 # Get communication group 

725 shard_group = self.mesh_info.shard_process_group if isinstance(self.mesh_info, FSDPMeshInfo) else None 

726 

727 if shard_group is None or self.shard_world_size <= 1: 

728 # No communication needed, just copy 

729 copy_without_bumping_version(self.all_gather_outputs[0], all_gather_input) 

730 return self.all_gather_outputs[0], None 

731 

732 # Execute all_gather_into_tensor 

733 handle = dist.all_gather_into_tensor( 

734 self.all_gather_outputs[0], 

735 all_gather_input, 

736 group=shard_group, 

737 async_op=async_op, 

738 ) 

739 

740 return self.all_gather_outputs[0], handle 

741 

742 def unshard(self, async_op: bool = False) -> None: 

743 if self.prefetch_handle is not None: 

744 # Already triggered by HSDPState.prefetch(), so return directly. 

745 return # no-op 

746 

747 _, handle = self._get_unsharded_param_data(async_op=async_op) 

748 self.prefetch_handle = handle 

749 

750 def wait_for_unshard(self) -> None: 

751 self._assert_in_states(ShardedState.SHARDED) 

752 

753 if self.prefetch_handle is not None: 

754 self.prefetch_handle.wait() 

755 self.prefetch_handle = None 

756 

757 self.init_unsharded_param() 

758 self.to_unsharded() 

759 

760 def shard(self) -> None: 

761 """ 

762 Transition parameter from unsharded back to sharded state. 

763 """ 

764 self._assert_in_states(ShardedState.UNSHARDED) 

765 self.to_sharded() 

766 

767 def reduce_scatter_output(self): 

768 """Return cached reduce-scatter output after waiting pending async work.""" 

769 if self.reduce_scatter_handle is not None: 

770 self.reduce_scatter_handle.wait() 

771 self.reduce_scatter_handle = None 

772 return self._reduce_scatter_output 

773 

774 def clear_reduce_scatter_output(self): 

775 """Clear cached reduce-scatter output.""" 

776 self._reduce_scatter_output = None 

777 

778 def reduce_scatter_grad( 

779 self, 

780 async_op: bool = True, 

781 dtype: Optional[ms.Type] = None, 

782 reduce_op: Optional[ops.ReduceOp] = ops.ReduceOp.AVG, 

783 output_buffer: Optional[ms.Tensor] = None, 

784 ) -> Tuple[ms.Tensor, Optional[CommHandle]]: 

785 """ 

786 Perform reduce-scatter on gradient to reduce and shard the full gradient. 

787 

788 Args: 

789 async_op: Whether to execute asynchronously. 

790 dtype: reduce dtype. 

791 reduce_op: do reduce-scatter avg or sum. 

792 output_buffer: Optional pre-allocated output for fused all-reduce groups. 

793 

794 Returns: 

795 (sharded_grad, handle): Sharded gradient and communication handle. 

796 """ 

797 self._assert_in_states(ShardedState.UNSHARDED) 

798 

799 # Choose gradient source based on use_accumulated_grad flag 

800 if self.unsharded_accumulated_grad is not None: 

801 grad = self.unsharded_accumulated_grad_data 

802 else: 

803 grad = self.unsharded_grad_data