Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/core/fully_shard/state_dict

3# Licensed under the Apache License, Version 2.0 (the "License");

4# you may not use this file except in compliance with the License.

5# You may obtain a copy of the License at

7# http://www.apache.org/licenses/LICENSE-2.0

9# Unless required by applicable law or agreed to in writing, software

10# distributed under the License is distributed on an "AS IS" BASIS,

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

12# See the License for the specific language governing permissions and

13# limitations under the License.

14# ============================================================================

15"""State dict utilities for fully_shard (torch-specific).

17Provides model state dict get/set (``get_model_state_dict``,

18``set_model_state_dict``) and optimizer state dict get/set

19(``get_optim_state_dict``, ``set_optim_state_dict``) that handle HSDP

20DTensor parameters, FQN-based keys, full/local/CPU conversion,

21broadcast-from-rank0, flatten/unflatten, and DCP load template construction.

23The optimizer state dict functions do **not** call PyTorch DCP's

24``get_optimizer_state_dict``/``set_optimizer_state_dict``.

25"""

26from __future__ import annotations

28import logging

29from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

31import torch

32import torch.distributed as dist

33from torch import nn

34from torch.distributed.checkpoint.state_dict import StateDictOptions

36from hyper_parallel.core.dtensor.dtensor import DTensor, distribute_tensor

37from hyper_parallel.core.dtensor.device_mesh import DeviceMesh

38from hyper_parallel.core.dtensor.placement_types import Replicate

40logger = logging.getLogger(__name__)

43# =====================================================================

44# Model state dict helpers

45# =====================================================================

47def _gather_full_state_dict(

48 state_dict: dict[str, Any], cpu_offload: bool

49) -> dict[str, Any]:

50 """All-gather every DTensor shard into a full tensor.

52 Args:

53 state_dict: Model state dict with DTensor or plain tensor values.

54 cpu_offload: If True, only rank-0 keeps the result on CPU;

55 other ranks return an empty dict to save memory.

56 """

57 is_rank0 = (not dist.is_initialized()) or (dist.get_rank() == 0)

59 gathered: dict[str, Any] = {}

60 for key, val in state_dict.items():

61 if isinstance(val, DTensor):

62 val = val.full_tensor()

63 if cpu_offload:

64 if not is_rank0:

65 del val

66 continue

67 if isinstance(val, torch.Tensor):

68 val = val.cpu()

69 gathered[key] = val

71 if cpu_offload and not is_rank0:

72 return {}

73 return gathered

76def _offload_sharded_state_dict(

77 state_dict: dict[str, Any],

78) -> dict[str, Any]:

79 """Move each shard to CPU without all-gathering.

81 Args:

82 state_dict: Model state dict with DTensor or plain tensor values.

83 """

84 offloaded: dict[str, Any] = {}

85 for key, val in state_dict.items():

86 if isinstance(val, DTensor):

87 val = DTensor.from_local(

88 val.to_local().cpu(), val.device_mesh, val.layout.alias_placements,

89 )

90 elif isinstance(val, torch.Tensor):

91 val = val.cpu()

92 offloaded[key] = val

93 return offloaded

96def get_model_state_dict(

97 model: nn.Module,

98 *,

99 options: Optional[StateDictOptions] = None,

100) -> dict[str, Any]:

101 """Return the model state dict with configurable gathering and offloading.

102

103 Behaviour matrix:

104

105 +-----------------+-------------+--------------------------------------+

106 | full_state_dict | cpu_offload | result |

107 +=================+=============+======================================+

108 | False | False | DTensor (sharded, as-is) |

109 +-----------------+-------------+--------------------------------------+

110 | False | True | DTensor local shard offloaded to CPU |

111 +-----------------+-------------+--------------------------------------+

112 | True | False | full Tensor on **every** rank |

113 +-----------------+-------------+--------------------------------------+

114 | True | True | full Tensor on CPU, **rank 0 only** |

115 +-----------------+-------------+--------------------------------------+

116

117 Args:

118 model: The model whose state dict to retrieve.

119 options: Controls full_state_dict, cpu_offload,

120 ignore_frozen_params, and broadcast_from_rank0 flags.

121

122 Raises:

123 ValueError: If ``broadcast_from_rank0`` is True while

124 ``full_state_dict`` is False.

125 NotImplementedError: If ``broadcast_from_rank0`` is True.

126 ``broadcast_from_rank0`` requires a cross-rank tensor broadcast,

127 which is not available. Use ``full_state_dict=True`` with a full

128 tensor on every rank instead.

129 """

130 options = options or StateDictOptions()

131

132 if options.broadcast_from_rank0 and not options.full_state_dict:

133 raise ValueError(

134 "full_state_dict must be True when broadcast_from_rank0 is True."

135 )

136 if options.broadcast_from_rank0:

137 raise NotImplementedError(

138 "broadcast_from_rank0=True is not supported. "

139 "broadcast_from_rank0 requires a cross-rank tensor broadcast, "

140 "which is not available. Use full_state_dict=True with a full "

141 "tensor on every rank instead."

142 )

143

144 state_dict: dict[str, Any] = model.state_dict()

145

146 if options.ignore_frozen_params:

147 frozen_keys = {

148 name for name, p in model.named_parameters()

149 if not p.requires_grad

150 }

151 for key in frozen_keys:

152 state_dict.pop(key, None)

153

154 if options.full_state_dict:

155 return _gather_full_state_dict(state_dict, options.cpu_offload)

156

157 if options.cpu_offload:

158 return _offload_sharded_state_dict(state_dict)

159

160 return state_dict

161

162

163def _scatter_model_state_dict(

164 model: nn.Module,

165 state_dict: dict[str, Any],

166 cpu_offload: bool,

167 strict: bool,

168) -> dict[str, Any]:

169 """Scatter full tensors into DTensor shards matching model's layout.

170

171 Inverse of _gather_full_state_dict. Uses distribute_tensor (no communication,

172 assumes every rank holds the same global tensor).

173

174 When ``cpu_offload`` is True the input tensors are expected to live on CPU

175 (e.g. produced by ``get_model_state_dict(full_state_dict=True, cpu_offload=True)``).

176 The scattered local shard is moved onto the target parameter's device so the

177 subsequent ``load_state_dict(assign=True)`` does not leave the parameter on

178 the wrong device.

179

180 Args:

181 model: The model whose layout the scattered tensors must match.

182 state_dict: Input state dict whose values are plain (global) tensors.

183 cpu_offload: Whether the input tensors live on CPU and must be moved

184 back onto the target parameter's device after scattering.

185 strict: When True, keys present in ``state_dict`` but absent from the

186 model are treated as errors (unexpected keys), mirroring the

187 PyTorch ``load_state_dict(strict=True)`` semantics of the

188 ``full_state_dict=False`` passthrough path.

189

190 Raises:

191 ValueError: If ``strict`` is True and ``state_dict`` contains keys that

192 do not exist in the model (unexpected keys).

193 """

194 target_state_dict = model.state_dict()

195 scattered: dict[str, Any] = {}

196 unexpected_keys: list[str] = []

197 for key, val in state_dict.items():

198 target = target_state_dict.get(key)

199 if target is None:

200 unexpected_keys.append(key)

201 continue

202 if isinstance(target, DTensor):

203 if isinstance(val, DTensor):

204 scattered[key] = val

205 else:

206 # Slice a plain (global) tensor into a DTensor shard.

207 placements = (

208 target.layout.alias_placements if target.layout else target.placements

209 )

210 scattered[key] = distribute_tensor(val, target.device_mesh, placements)

211 else:

212 scattered[key] = val

213

214 # When the input came from CPU (cpu_offload=True round-trip), move the

215 # scattered local shard onto the target parameter's device. assign=True

216 # would otherwise bind a CPU DTensor to an on-device parameter.

217 if cpu_offload and isinstance(scattered[key], DTensor) and isinstance(target, DTensor):

218 dt = scattered[key]

219 target_device = target._local_tensor.device # pylint: disable=protected-access

220 scattered[key] = DTensor.from_local(

221 dt.to_local().to(target_device),

222 dt.device_mesh,

223 dt.layout.alias_placements if dt.layout else dt.placements,

224 )

225

226 if strict and unexpected_keys:

227 raise ValueError(

228 f"Unexpected key(s) in state_dict: {unexpected_keys}. "

229 f"To allow loading a state_dict with extra keys, pass strict=False."

230 )

231 return scattered

232

233

234def set_model_state_dict(

235 model: nn.Module,

236 state_dict: dict[str, Any],

237 *,

238 options: Optional[StateDictOptions] = None,

239 assign: bool = True,

240) -> None:

241 """Load state dict into model, scattering full tensors to DTensor shards.

242

243 Behaviour matrix:

244

245 +-----------------+-------------+--------------------------------+

246 | full_state_dict | cpu_offload | behaviour |

247 +=================+=============+================================+

248 | True | False | scatter full -> DTensor on dev |

249 +-----------------+-------------+--------------------------------+

250 | True | True | scatter full -> DTensor on dev |

251 +-----------------+-------------+--------------------------------+

252 | False (any) | * | load sharded DTensors as-is |

253 +-----------------+-------------+--------------------------------+

254

255 Note:

256 ``ignore_frozen_params`` is intentionally a no-op on the setter path,

257 matching the distributed checkpoint contract: frozen parameter

258 filtering is a getter-only feature. Callers that strip frozen keys

259 from the input should also pass ``strict=False``.

260

261 Args:

262 model: The model to load state into.

263 state_dict: State dict to load. Values may be plain (global) tensors

264 when ``full_state_dict=True`` or sharded DTensors otherwise.

265 options: Controls full_state_dict, cpu_offload, strict and

266 broadcast_from_rank0 flags.

267 assign: Replace parameters when True; use False when an optimizer already

268 owns the final parameters and their identities must remain stable.

269

270 Raises:

271 ValueError: If ``broadcast_from_rank0`` is True while

272 ``full_state_dict`` is False.

273 NotImplementedError: If ``broadcast_from_rank0`` is True.

274 ``broadcast_from_rank0`` requires a cross-rank tensor broadcast,

275 which is not available. Use ``full_state_dict=True`` with a full

276 tensor on every rank instead.

277 """

278 options = options or StateDictOptions()

279

280 if options.broadcast_from_rank0 and not options.full_state_dict:

281 raise ValueError(

282 "full_state_dict must be True when broadcast_from_rank0 is True."

283 )

284 if options.broadcast_from_rank0:

285 raise NotImplementedError(

286 "broadcast_from_rank0=True is not supported. "

287 "broadcast_from_rank0 requires a cross-rank tensor broadcast, "

288 "which is not available. Use full_state_dict=True with a full "

289 "tensor on every rank instead."

290 )

291

292 # Scatter full tensors into DTensor shards matching the model's layout.

293 if options.full_state_dict:

294 scattered = _scatter_model_state_dict(

295 model, state_dict, options.cpu_offload, options.strict,

296 )

297 else:

298 scattered = state_dict

299

300 model.load_state_dict(scattered, strict=options.strict, assign=assign)

301

302

303# =====================================================================

304# Optimizer state dict helpers

305# =====================================================================

306

307_SCALAR_STATE_KEYS = {"step"}

308

309

310class UnsupportedConfigurationError(RuntimeError):

311 """Raised when the requested configuration is not supported."""

312

313

314def _is_dtensor_param(param: Any) -> bool:

315 return isinstance(param, DTensor)

316

317

318def _is_replicate_group_root(

319 mesh: DeviceMesh,

320 placements: Sequence,

321) -> bool:

322 """Return True if current rank is the root (coordinate-0) in every

323 Replicate dimension of the given mesh+placements.

324

325 For HSDP with mesh (replicate, shard) and placements (Replicate(), Shard(0)),

326 only the rank whose replicate-dim coordinate is 0 returns True — even if

327 its global rank is not 0 (e.g. PP stage-1 ranks in a 3-D mesh).

328

329 If there are no Replicate dimensions, falls back to ``dist.get_rank() == 0``.

330 """

331 coord = mesh.get_coordinate()

332 if coord is None:

333 return False

334 replicate_dims = [i for i, p in enumerate(placements) if isinstance(p, Replicate)]

335 if not replicate_dims:

336 return dist.get_rank() == 0

337 return all(coord[d] == 0 for d in replicate_dims)

338

339

340def _param_to_fqn(model: nn.Module) -> Dict[nn.Parameter, str]:

341 """Build a mapping from parameter object to its fully-qualified name."""

342 param_to_name: Dict[nn.Parameter, str] = {}

343 for name, param in model.named_parameters():

344 if param in param_to_name:

345 raise ValueError(

346 f"Parameter {name} shares the same object as "

347 f"{param_to_name[param]}. Duplicate parameter objects are not supported."

348 )

349 param_to_name[param] = name

350 return param_to_name

351

352

353def _build_id_to_fqn(

354 optimizer: torch.optim.Optimizer,

355 model: nn.Module,

356) -> Tuple[Dict[int, str], Dict[str, int]]:

357 """Build mappings between optimizer saved IDs and FQNs.

358

359 Returns:

360 Tuple of (saved_id_to_fqn, fqn_to_saved_id).

361 """

362 param_to_name = _param_to_fqn(model)

363 raw_sd = optimizer.state_dict()

364 saved_id_to_fqn: Dict[int, str] = {}

365 saved_id_to_param: Dict[int, nn.Parameter] = {}

366

367 for runtime_group, saved_group in zip(optimizer.param_groups, raw_sd["param_groups"]):

368 for parameter, saved_id in zip(runtime_group["params"], saved_group["params"]):

369 if saved_id in saved_id_to_param and saved_id_to_param[saved_id] is not parameter:

370 raise ValueError(

371 f"saved_id {saved_id} maps to different parameters; "

372 f"this is not supported."

373 )

374 saved_id_to_param[saved_id] = parameter

375

376 for saved_id, param in saved_id_to_param.items():

377 if param not in param_to_name:

378 raise ValueError(

379 f"Parameter with saved_id={saved_id} not found in "

380 f"model.named_parameters(). Ensure the model matches the optimizer."

381 )

382 saved_id_to_fqn[saved_id] = param_to_name[param]

383

384 fqn_to_saved_id = {v: k for k, v in saved_id_to_fqn.items()}

385 return saved_id_to_fqn, fqn_to_saved_id

386

387

388def _get_param_dtensor_info(

389 param: nn.Parameter,

390) -> Optional[Tuple[DeviceMesh, Sequence]]:

391 """Return (device_mesh, placements) if param is DTensor, else None."""

392 if _is_dtensor_param(param):

393 return param.device_mesh, param.placements

394 return None

395

396

397def _convert_state_tensor(

398 tensor: torch.Tensor,

399 param: nn.Parameter,

400 full_state_dict: bool,

401 cpu_offload: bool,

402 is_rank0: bool,

403) -> Optional[torch.Tensor]:

404 """Convert a single optimizer state tensor according to options.

405

406 For scalar states (step), only cpu_offload applies.

407 For tensor states sharing the parameter shape, full/local conversion applies.

408

409 When the optimizer state value is itself a DTensor (e.g. when

410 ``optimizer.step()`` ran without ``SkipDTensorDispatch`` and the

411 parameter is a DTensor), we extract the local shard via

412 ``.to_local()`` or gather via ``.full_tensor()`` directly rather

413 than wrapping it again.

414

415 Args:

416 tensor: The optimizer state tensor (plain Tensor or DTensor).

417 param: The corresponding model parameter (may be DTensor).

418 full_state_dict: If True, gather to full tensor.

419 cpu_offload: If True, move to CPU.

420 is_rank0: Whether current rank is 0 (used for full+cpu to drop non-rank0 data).

421

422 Returns:

423 Converted tensor, or None if this rank should not keep the result.

424 """

425 dtensor_info = _get_param_dtensor_info(param)

426 is_dtensor_value = isinstance(tensor, DTensor)

427

428 if full_state_dict and dtensor_info is not None:

429 if is_dtensor_value:

430 full_t = tensor.full_tensor()

431 else:

432 mesh, placements = dtensor_info

433 full_t = DTensor.from_local(tensor, mesh, placements).full_tensor()

434 if cpu_offload:

435 if not is_rank0:

436 return None

437 return full_t.cpu()

438 return full_t

439

440 if is_dtensor_value:

441 local_t = tensor.to_local()

442 if cpu_offload:

443 return local_t.cpu()

444 return local_t

445

446 if cpu_offload:

447 return tensor.cpu()

448

449 return tensor

450

451

452def _is_scalar_state(key: str) -> bool:

453 return key in _SCALAR_STATE_KEYS

454

455

456def _convert_state_scalar(

457 tensor: torch.Tensor,

458 cpu_offload: bool,

459) -> torch.Tensor:

460 if isinstance(tensor, DTensor):

461 tensor = tensor.to_local()

462 if cpu_offload:

463 return tensor.cpu()

464 return tensor

465

466

467def _determine_is_root(

468 full_state_dict: bool,

469 cpu_offload: bool,

470 dtensor_info: Optional[Tuple[Any, Any]],

471) -> bool:

472 """Determine if current rank is the replicate-group root."""

473 if full_state_dict and cpu_offload and dtensor_info is not None:

474 return _is_replicate_group_root(dtensor_info[0], dtensor_info[1])

475 if full_state_dict and cpu_offload:

476 return (not dist.is_initialized()) or (dist.get_rank() == 0)

477 return True

478

479

480def _convert_state_entries(

481 state: Dict[str, Any],

482 param: nn.Parameter,

483 full_state_dict: bool,

484 cpu_offload: bool,

485) -> Dict[str, Any]:

486 """Convert optimizer state entries for a single parameter."""

487 dtensor_info = _get_param_dtensor_info(param)

488 is_root = _determine_is_root(full_state_dict, cpu_offload, dtensor_info)

489 skip_non_root = full_state_dict and cpu_offload and not is_root

490

491 converted: Dict[str, Any] = {}

492 for key, value in state.items():

493 if isinstance(value, torch.Tensor):

494 if _is_scalar_state(key):

495 result = _convert_state_scalar(value, cpu_offload)

496 if not skip_non_root:

497 converted[key] = result

498 else:

499 result = _convert_state_tensor(

500 value, param, full_state_dict, cpu_offload, is_root,

501 )

502 if result is not None:

503 converted[key] = result

504 else:

505 if not skip_non_root:

506 converted[key] = value

507 return converted

508

509

510def _build_result_param_groups(

511 optimizer: torch.optim.Optimizer,

512 raw_sd: Dict[str, Any],

513 saved_id_to_fqn: Dict[int, str],

514) -> List[Dict[str, Any]]:

515 """Build result param_groups with FQN keys from saved id-based groups."""

516 del optimizer

517 result_param_groups: List[Dict[str, Any]] = []

518 for saved_group in raw_sd["param_groups"]:

519 pg: Dict[str, Any] = {}

520 for k, v in saved_group.items():

521 if k == "params":

522 pg["params"] = [saved_id_to_fqn[sid] for sid in v]

523 else:

524 pg[k] = v

525 result_param_groups.append(pg)

526 return result_param_groups

527

528

529def _check_chained_optimizer(optimizer: Any) -> None:

530 if type(optimizer).__name__ == "ChainedOptimizer":

531 raise ValueError(

532 "ChainedOptimizer is not supported by get_optim_state_dict / "

533 "set_optim_state_dict. Use ChainedOptimizer.state_dict() and "

534 "ChainedOptimizer.load_state_dict() instead, which handle "

535 "multi-optimizer merging internally."

536 )

537

538

539def get_optim_state_dict(

540 model: nn.Module,

541 optimizer: torch.optim.Optimizer,

542 *,

543 options: Any = None,

544) -> Dict[str, Any]:

545 """Return the optimizer state dict with FQN keys and configurable conversion.

546

547 Args:

548 model: The model whose parameters are optimized.

549 optimizer: The optimizer instance.

550 options: Options object with fields:

551 full_state_dict (bool): Gather DTensor shards to full tensors.

552 cpu_offload (bool): Move tensors to CPU.

553 flatten_optimizer_state_dict (bool): Flatten to single-level keys.

554 strict (bool): Not used in get; reserved for set.

555 broadcast_from_rank0 (bool): Not used in get; reserved for set.

556

557 Returns:

558 Optimizer state dict with FQN-based keys.

559 """

560 _check_chained_optimizer(optimizer)

561

562 full_state_dict = getattr(options, "full_state_dict", False)

563 cpu_offload = getattr(options, "cpu_offload", False)

564 flatten = getattr(options, "flatten_optimizer_state_dict", False)

565

566 raw_sd = optimizer.state_dict()

567 saved_id_to_fqn, _ = _build_id_to_fqn(optimizer, model)

568

569 param_by_id: Dict[int, nn.Parameter] = {}

570 for runtime_group, saved_group in zip(optimizer.param_groups, raw_sd["param_groups"]):

571 for parameter, saved_id in zip(runtime_group["params"], saved_group["params"]):

572 param_by_id[saved_id] = parameter

573

574 result_state: Dict[str, Dict[str, Any]] = {}

575 for saved_id, state in raw_sd["state"].items():

576 fqn = saved_id_to_fqn[saved_id]

577 param = param_by_id[saved_id]

578 converted = _convert_state_entries(

579 state, param, full_state_dict, cpu_offload,

580 )

581 if converted:

582 result_state[fqn] = converted

583

584 result_param_groups = _build_result_param_groups(

585 optimizer, raw_sd, saved_id_to_fqn,

586 )

587

588 result: Dict[str, Any] = {

589 "state": result_state,

590 "param_groups": result_param_groups,

591 }

592

593 if flatten:

594 result = _flatten_optim_state_dict(result)

595

596 return result

597

598

599def _check_strict_fqns(

600 source_fqns: set,

601 target_fqns: set,

602 strict: bool,

603) -> None:

604 """Validate FQN compatibility under strict mode."""

605 if not strict:

606 return

607 extra_fqns = source_fqns - target_fqns

608 if extra_fqns:

609 raise ValueError(

610 f"strict=True but checkpoint contains FQNs not in target "

611 f"optimizer: {sorted(extra_fqns)}"

612 )

613

614

615def _load_state_values(

616 source_state: Dict[str, Any],

617 param: nn.Parameter,

618 full_state_dict: bool,

619 cpu_offload: bool,

620 stores_dtensor: bool,

621) -> Dict[str, Any]:

622 """Convert and load state values for a single FQN."""

623 converted: Dict[str, Any] = {}

624 for key, value in source_state.items():

625 if isinstance(value, torch.Tensor):

626 if _is_scalar_state(key):

627 converted[key] = _convert_input_scalar_to_target(

628 value, param, cpu_offload, stores_dtensor,

629 )

630 else:

631 converted[key] = _convert_input_tensor_to_target(

632 value, param, full_state_dict, cpu_offload, stores_dtensor,

633 )

634 else:

635 converted[key] = value

636 return converted

637

638

639def _merge_param_groups(

640 source_param_groups: List[Dict[str, Any]],

641 target_raw_sd: Dict[str, Any],

642 target_fqn_to_saved_id: Dict[str, int],

643) -> None:

644 """Merge source param_group fields into target param_groups in-place."""

645 if not source_param_groups:

646 return

647 target_saved_id_to_fqn = {v: k for k, v in target_fqn_to_saved_id.items()}

648

649 target_pg_by_fqns: Dict[frozenset, Dict[str, Any]] = {}

650 for saved_group in target_raw_sd["param_groups"]:

651 fqns_in_group = frozenset(

652 target_saved_id_to_fqn.get(sid, "")

653 for sid in saved_group.get("params", [])

654 )

655 target_pg_by_fqns[fqns_in_group] = saved_group

656

657 for source_pg in source_param_groups:

658 source_fqn_set = frozenset(source_pg.get("params", []))

659 matched_target_pg = None

660 for target_fqn_set, target_pg in target_pg_by_fqns.items():

661 if source_fqn_set & target_fqn_set:

662 matched_target_pg = target_pg

663 break

664 if matched_target_pg is None:

665 continue

666 for k, v in source_pg.items():

667 if k == "params":

668 continue

669 matched_target_pg[k] = v

670

671

672def _build_target_fqn_mappings(

673 optimizer: torch.optim.Optimizer,

674 model: nn.Module,

675 param_by_fqn: Dict[str, nn.Parameter],

676) -> Tuple[Dict[str, int], set]:

677 """Build target FQN→saved_id mapping and identify DTensor state IDs.

678

679 Returns:

680 Tuple of (fqn_to_saved_id, dtensor_state_ids).

681 """

682 del param_by_fqn

683 target_raw_sd = optimizer.state_dict()

684

685 target_id_to_param: Dict[int, nn.Parameter] = {}

686 for runtime_group, saved_group in zip(optimizer.param_groups, target_raw_sd["param_groups"]):

687 for parameter, saved_id in zip(runtime_group["params"], saved_group["params"]):

688 target_id_to_param[saved_id] = parameter

689

690 fqn_to_saved_id: Dict[str, int] = {}

691 for saved_id, param in target_id_to_param.items():

692 for name, p in model.named_parameters():

693 if p is param:

694 fqn_to_saved_id[name] = saved_id

695 break

696

697 dtensor_state_ids: set = set()

698 for saved_id, state in target_raw_sd["state"].items():

699 for value in state.values():

700 if isinstance(value, DTensor):

701 dtensor_state_ids.add(saved_id)

702 break

703

704 return fqn_to_saved_id, dtensor_state_ids

705

706

707def _load_fqn_states(

708 optim_state_dict: Dict[str, Any],

709 target_raw_sd: Dict[str, Any],

710 target_fqn_to_saved_id: Dict[str, int],

711 param_by_fqn: Dict[str, nn.Parameter],

712 dtensor_state_ids: set,

713 full_state_dict: bool,

714 cpu_offload: bool,

715 strict: bool,

716) -> Dict[int, Dict[str, Any]]:

717 """Load source FQN states into the target state dict format.

718

719 Returns:

720 Updated state dict with saved_id keys.

721 """

722 new_state: Dict[int, Dict[str, Any]] = dict(target_raw_sd["state"])

723

724 for fqn, source_state in optim_state_dict.get("state", {}).items():

725 if fqn not in target_fqn_to_saved_id:

726 if strict:

727 raise ValueError(

728 f"strict=True but FQN '{fqn}' not found in target optimizer."

729 )

730 continue

731

732 target_saved_id = target_fqn_to_saved_id[fqn]

733 param = param_by_fqn[fqn]

734

735 if target_saved_id not in new_state:

736 new_state[target_saved_id] = {}

737

738 new_state[target_saved_id].update(

739 _load_state_values(

740 source_state, param, full_state_dict, cpu_offload,

741 target_saved_id in dtensor_state_ids,

742 )

743 )

744

745 return new_state

746

747

748def set_optim_state_dict(

749 model: nn.Module,

750 optimizer: torch.optim.Optimizer,

751 optim_state_dict: Dict[str, Any],

752 *,

753 options: Any = None,

754) -> None:

755 """Load an optimizer state dict with FQN keys into the optimizer.

756

757 Args:

758 model: The model whose parameters are optimized.

759 optimizer: The target optimizer instance.

760 optim_state_dict: State dict with FQN-based keys (possibly flatten format).

761 options: Options object with fields:

762 full_state_dict (bool): Input contains full (non-sharded) tensors.

763 cpu_offload (bool): Input tensors are on CPU.

764 flatten_optimizer_state_dict (bool): Input is in flatten format.

765 strict (bool): If True, reject FQNs in the checkpoint that do not

766 exist in the target optimizer (extra FQNs). Missing FQNs are

767 always allowed because the target optimizer may have

768 untrained/empty state for some parameters. If False, extra

769 FQNs in the checkpoint are silently ignored.

770 broadcast_from_rank0 (bool): If True and full_state_dict+cpu_offload,

771 broadcast from the replicate-group root to all ranks within

772 the same replicate subgroup. In PP+HSDP, each PP stage

773 broadcasts independently within its own HSDP subgroup.

774 """

775 full_state_dict = getattr(options, "full_state_dict", False)

776 cpu_offload = getattr(options, "cpu_offload", False)

777 flatten = getattr(options, "flatten_optimizer_state_dict", False)

778 strict = getattr(options, "strict", True)

779 broadcast_from_rank0 = getattr(options, "broadcast_from_rank0", False)

780

781 _check_chained_optimizer(optimizer)

782

783 if flatten:

784 optim_state_dict = _unflatten_optim_state_dict(optim_state_dict, model, strict=strict)

785

786 if full_state_dict and cpu_offload and not broadcast_from_rank0:

787 has_any_state = bool(optim_state_dict.get("state", {}))

788 if not has_any_state:

789 raise ValueError(

790 "Received empty state dict with full_state_dict=True "

791 "and cpu_offload=True but broadcast_from_rank0=False. "

792 "Set broadcast_from_rank0=True to allow the replicate-group "

793 "root to broadcast."

794 )

795

796 if broadcast_from_rank0:

797 optim_state_dict = _broadcast_state_from_rank0(

798 optim_state_dict, model, full_state_dict, cpu_offload,

799 )

800

801 param_by_fqn: Dict[str, nn.Parameter] = {}

802 for name, param in model.named_parameters():

803 param_by_fqn[name] = param

804

805 target_fqn_to_saved_id, dtensor_state_ids = _build_target_fqn_mappings(

806 optimizer, model, param_by_fqn,

807 )

808

809 source_fqns = set(optim_state_dict.get("state", {}).keys())

810 target_fqns = set(target_fqn_to_saved_id.keys())

811 _check_strict_fqns(source_fqns, target_fqns, strict)

812

813 target_raw_sd = optimizer.state_dict()

814 new_state = _load_fqn_states(

815 optim_state_dict, target_raw_sd, target_fqn_to_saved_id,

816 param_by_fqn, dtensor_state_ids, full_state_dict, cpu_offload, strict,

817 )

818

819 target_raw_sd["state"] = new_state

820

821 _merge_param_groups(

822 optim_state_dict.get("param_groups", []),

823 target_raw_sd,

824 target_fqn_to_saved_id,

825 )

826

827 optimizer.load_state_dict(target_raw_sd)

828

829

830def _optimizer_tensor_result(tensor: DTensor, stores_dtensor: bool) -> Union[torch.Tensor, DTensor]:

831 """Keep global layout only when the target optimizer stores DTensor values."""

832 return tensor if stores_dtensor else tensor.to_local()

833

834

835def _convert_input_tensor_to_target(

836 tensor: torch.Tensor,

837 param: nn.Parameter,

838 full_state_dict: bool,

839 cpu_offload: bool,

840 stores_dtensor: bool = False,

841) -> Union[torch.Tensor, DTensor]:

842 """Convert an incoming state tensor to the target optimizer's format.

843

844 For DTensor params whose optimizer stores DTensor state values (i.e. when

845 ``optimizer.step()`` ran without ``SkipDTensorDispatch``), we return a

846 DTensor so that ``optimizer.load_state_dict()`` receives the correct type.

847

848 For DTensor params whose optimizer stores plain local-shard tensors (i.e.

849 when ``optimizer.step()`` ran inside ``SkipDTensorDispatch``), we return

850 a plain local-shard tensor.

851

852 For plain (non-DTensor) params, we return a plain tensor on the correct

853 device.

854

855 Args:

856 tensor: Incoming state tensor (plain Tensor, possibly on CPU).

857 param: The corresponding model parameter (may be DTensor).

858 full_state_dict: If True, input tensor represents the full (non-sharded) value.

859 cpu_offload: If True, input tensor is on CPU and may need device transfer.

860 stores_dtensor: If True, the target optimizer stores DTensor state values

861 for this parameter (detected from existing ``optimizer.state_dict()``).

862

863 Returns:

864 Tensor or DTensor matching the target optimizer's expected format.

865 """

866 dtensor_info = _get_param_dtensor_info(param)

867

868 if not dtensor_info:

869 target_device = param.data.device

870 result = tensor

871 if cpu_offload and result.device != target_device:

872 result = result.to(target_device)

873 return result

874

875 mesh, placements = dtensor_info

876 target_device = param.to_local().device

877

878 if isinstance(tensor, DTensor):

879 if tensor.device_mesh == mesh and tensor.placements == placements:

880 return _optimizer_tensor_result(tensor, stores_dtensor)

881 redistributed = tensor.redistribute(mesh, placements)

882 return _optimizer_tensor_result(redistributed, stores_dtensor)

883

884 if full_state_dict:

885 if cpu_offload:

886 tensor = tensor.to(target_device)

887 dt = distribute_tensor(tensor, mesh, placements)

888 return _optimizer_tensor_result(dt, stores_dtensor)

889

890 if cpu_offload and tensor.device != target_device:

891 tensor = tensor.to(target_device)

892

893 if stores_dtensor:

894 if tensor.shape == param.to_local().shape:

895 return DTensor.from_local(tensor, mesh, placements)

896 dt = distribute_tensor(tensor, mesh, placements)

897 return dt

898

899 if tensor.shape == param.to_local().shape:

900 return tensor

901

902 dt = distribute_tensor(tensor, mesh, placements)

903 return dt.to_local()

904

905

906def _convert_input_scalar_to_target(

907 tensor: torch.Tensor,

908 param: nn.Parameter,

909 cpu_offload: bool,

910 stores_dtensor: bool = False,

911) -> Union[torch.Tensor, DTensor]:

912 """Convert an incoming scalar state tensor (e.g. step) to the target format.

913

914 Scalar states are not sharded — they are replicated on all ranks. When the

915 target optimizer stores DTensor values (``stores_dtensor=True``), we wrap

916 the scalar as a replicated DTensor; otherwise we return a plain tensor on

917 the correct device.

918

919 Args:

920 tensor: Incoming scalar tensor (plain Tensor, possibly on CPU).

921 param: The corresponding model parameter (may be DTensor).

922 cpu_offload: If True, input tensor is on CPU and may need device transfer.

923 stores_dtensor: If True, the target optimizer stores DTensor state values.

924

925 Returns:

926 Plain tensor or replicated DTensor matching the target's expected format.

927 """

928 dtensor_info = _get_param_dtensor_info(param)

929

930 if not dtensor_info or not stores_dtensor:

931 target_device = param.data.device if not dtensor_info else param.to_local().device

932 result = tensor

933 if cpu_offload and result.device != target_device:

934 result = result.to(target_device)

935 return result

936

937 mesh, _ = dtensor_info

938 target_device = param.to_local().device

939

940 scalar_placements = [Replicate()] * mesh.ndim

941 if cpu_offload and tensor.device != target_device:

942 tensor = tensor.to(target_device)

943 return DTensor.from_local(tensor, mesh, scalar_placements)

944

945

946def _get_broadcast_groups(

947 model: nn.Module,

948) -> List[Dict[str, Any]]:

949 """Identify broadcast groups from the model's DTensor parameters.

950

951 A broadcast group is a set of ranks that share the same model

952 parameters (same FQNs). In HSDP, ranks within the same replicate

953 subgroup share parameters. In PP, each PP stage is a different

954 broadcast group because different stages own different parameters.

955

956 Returns:

957 List of dicts, each with keys:

958 fqns: list of FQN strings owned by this group

959 pg: ProcessGroup for broadcast within the group

960 src_rank: global rank of the replicate-group root

961 is_root: True if current rank is the root of this group

962 """

963 param_by_fqn: Dict[str, nn.Parameter] = {}

964 for name, param in model.named_parameters():

965 param_by_fqn[name] = param

966

967 model_fqns = list(param_by_fqn.keys())

968

969 if not model_fqns:

970 return []

971

972 first_param = param_by_fqn[model_fqns[0]]

973 dtensor_info = _get_param_dtensor_info(first_param)

974

975 if dtensor_info is not None:

976 mesh, placements = dtensor_info

977 is_root = _is_replicate_group_root(mesh, placements)

978 replicate_dims = [i for i, p in enumerate(placements) if isinstance(p, Replicate)]

979 if replicate_dims:

980 pg = mesh.get_group(replicate_dims[0])

981 coord = mesh.get_coordinate()

982 if coord is not None:

983 root_coord = list(coord)

984 for d in replicate_dims:

985 root_coord[d] = 0

986 src_rank = int(mesh.mesh[tuple(root_coord)])

987 else:

988 src_rank = 0

989 else:

990 pg = dist.group.WORLD

991 src_rank = 0

992 else:

993 is_root = dist.get_rank() == 0

994 pg = dist.group.WORLD

995 src_rank = 0

996

997 return [{

998 "fqns": model_fqns,

999 "pg": pg,

1000 "src_rank": src_rank,

1001 "is_root": is_root,

1002 "param_by_fqn": param_by_fqn,

1003 }]

1004

1005

1006def _broadcast_fqn_list(

1007 is_root: bool,

1008 fqns: List[str],

1009 pg: dist.ProcessGroup,

1010 src_rank: int,

1011) -> List[str]:

1012 """Broadcast FQN list from root to all ranks in the replicate group."""

1013 fqn_list = fqns if is_root else []

1014 obj = [fqn_list]

1015 dist.broadcast_object_list(obj, src=src_rank, group=pg)

1016 return obj[0]

1017

1018

1019def _build_state_schema(

1020 state: Dict[str, Any],

1021) -> Dict[str, Any]:

1022 """Build schema dict describing the types and shapes of state entries."""

1023 schema: Dict[str, Any] = {}

1024 for key, value in state.items():

1025 if isinstance(value, torch.Tensor):

1026 schema[key] = {

1027 "shape": tuple(value.shape),

1028 "dtype": str(value.dtype),

1029 "is_scalar": _is_scalar_state(key),

1030 }

1031 else:

1032 schema[key] = {"type": type(value).__name__, "value": value}

1033 return schema

1034

1035

1036def _broadcast_schema_per_fqn(

1037 fqn_list: List[str],

1038 optim_state_dict: Dict[str, Any],

1039 is_root: bool,

1040 pg: dist.ProcessGroup,

1041 src_rank: int,

1042) -> Dict[str, Dict[str, Any]]:

1043 """Broadcast schema for each FQN from root to all ranks."""

1044 fqn_schema: Dict[str, Dict[str, Any]] = {}

1045 for fqn in fqn_list:

1046 if is_root and fqn in optim_state_dict.get("state", {}):

1047 schema = _build_state_schema(optim_state_dict["state"][fqn])

1048 else:

1049 schema = {}

1050 schema_list = [schema] if is_root else [None]

1051 dist.broadcast_object_list(schema_list, src=src_rank, group=pg)

1052 fqn_schema[fqn] = schema_list[0]

1053 return fqn_schema

1054

1055

1056def _get_local_device() -> torch.device:

1057 """Get the device for the current local rank."""

1058 local_rank = dist.get_rank() if dist.is_initialized() else 0

1059 if torch.npu.is_available():

1060 return torch.device(f"npu:{local_rank}")

1061 return torch.device(f"cuda:{local_rank}")

1062

1063

1064def _broadcast_scalar_entry(

1065 key: str,

1066 fqn: str,

1067 info: Dict[str, Any],

1068 optim_state_dict: Dict[str, Any],

1069 is_root: bool,

1070 src_rank: int,

1071 pg: dist.ProcessGroup,

1072 cpu_offload: bool,

1073) -> torch.Tensor:

1074 """Broadcast a scalar state entry and return the result tensor."""

1075 dtype_str = info.get("dtype", "torch.float32")

1076 dtype = _resolve_dtype(dtype_str)

1077 device = _get_local_device()

1078

1079 if is_root and fqn in optim_state_dict.get("state", {}):

1080 scalar_val = optim_state_dict["state"][fqn][key]

1081 t = scalar_val.clone() if isinstance(scalar_val, torch.Tensor) else torch.tensor(scalar_val)

1082 if t.dim() == 0:

1083 t = t.reshape(1)

1084 else:

1085 t = torch.zeros(1, dtype=dtype)

1086

1087 if t.device.type == "cpu":

1088 t = t.to(device)

1089 dist.broadcast(t, src=src_rank, group=pg)

1090 return t.reshape(()).cpu() if cpu_offload else t.reshape(())

1091

1092

1093def _broadcast_tensor_entry(

1094 key: str,

1095 fqn: str,

1096 info: Dict[str, Any],

1097 optim_state_dict: Dict[str, Any],

1098 is_root: bool,

1099 src_rank: int,

1100 pg: dist.ProcessGroup,

1101 full_state_dict: bool,

1102 cpu_offload: bool,

1103) -> torch.Tensor:

1104 """Broadcast a tensor state entry and return the result tensor."""

1105 shape = info.get("shape", ())

1106 dtype_str = info.get("dtype", "torch.float32")

1107 dtype = _resolve_dtype(dtype_str)

1108 device = _get_local_device()

1109

1110 if is_root and fqn in optim_state_dict.get("state", {}):

1111 src_tensor = optim_state_dict["state"][fqn][key]

1112 if src_tensor.device.type == "cpu":

1113 src_tensor = src_tensor.to(device)

1114 else:

1115 src_tensor = torch.zeros(shape, dtype=dtype, device=device)

1116

1117 dist.broadcast(src_tensor, src=src_rank, group=pg)

1118

1119 if full_state_dict and cpu_offload:

1120 return src_tensor.cpu() if is_root else src_tensor

1121 return src_tensor.cpu() if cpu_offload else src_tensor

1122

1123

1124def _broadcast_tensor_data_per_fqn(

1125 fqn_list: List[str],

1126 fqn_schema: Dict[str, Dict[str, Any]],

1127 optim_state_dict: Dict[str, Any],

1128 is_root: bool,

1129 src_rank: int,

1130 pg: dist.ProcessGroup,

1131 full_state_dict: bool,

1132 cpu_offload: bool,

1133) -> Dict[str, Dict[str, Any]]:

1134 """Broadcast tensor data for each FQN from root to all ranks."""

1135 result_state: Dict[str, Dict[str, Any]] = {}

1136 for fqn in fqn_list:

1137 schema = fqn_schema[fqn]

1138 if not schema:

1139 continue

1140 result_state[fqn] = {}

1141 for key, info in schema.items():

1142 is_scalar = info.get("is_scalar", False)

1143 if is_scalar:

1144 result_state[fqn][key] = _broadcast_scalar_entry(

1145 key, fqn, info, optim_state_dict,

1146 is_root, src_rank, pg, cpu_offload,

1147 )

1148 else:

1149 result_state[fqn][key] = _broadcast_tensor_entry(

1150 key, fqn, info, optim_state_dict,

1151 is_root, src_rank, pg, full_state_dict, cpu_offload,

1152 )

1153 return result_state

1154

1155

1156def _broadcast_state_from_rank0(

1157 optim_state_dict: Dict[str, Any],

1158 model: nn.Module,

1159 full_state_dict: bool,

1160 cpu_offload: bool,

1161) -> Dict[str, Any]:

1162 """Broadcast full optimizer state from replicate-group root to all ranks.

1163

1164 Only meaningful when full_state_dict=True and cpu_offload=True.

1165 The root rank (coordinate-0 in every Replicate dimension) has the full

1166 data; other ranks in the same replicate subgroup receive it via

1167 broadcast on the parameter's replicate-dim process group.

1168

1169 For non-DTensor parameters (pure FSDP), falls back to global rank 0

1170 as the broadcast source.

1171

1172 **MPMD / Pipeline Parallelism**: each rank uses its own model's FQN

1173 list, which only contains the parameters of its own PP stage. The

1174 broadcast is scoped to the replicate-dim process group (HSDP

1175 subgroup), so different PP stages broadcast independently without

1176 cross-stage interference. This avoids the problem where global

1177 rank 0's FQN list does not include other stages' parameters.

1178 """

1179 if not dist.is_initialized():

1180 return optim_state_dict

1181

1182 groups = _get_broadcast_groups(model)

1183

1184 if not groups:

1185 return {

1186 "state": {},

1187 "param_groups": optim_state_dict.get("param_groups", []),

1188 }

1189

1190 group = groups[0]

1191 pg = group["pg"]

1192 src_rank = group["src_rank"]

1193 is_root = group["is_root"]

1194 fqns = group["fqns"]

1195

1196 fqn_list = _broadcast_fqn_list(is_root, fqns, pg, src_rank)

1197 fqn_schema = _broadcast_schema_per_fqn(

1198 fqn_list, optim_state_dict, is_root, pg, src_rank,

1199 )

1200 result_state = _broadcast_tensor_data_per_fqn(

1201 fqn_list, fqn_schema, optim_state_dict,

1202 is_root, src_rank, pg, full_state_dict, cpu_offload,

1203 )

1204

1205 return {

1206 "state": result_state,

1207 "param_groups": optim_state_dict.get("param_groups", []),

1208 }

1209

1210

1211def _flatten_optim_state_dict(

1212 state_dict: Dict[str, Any],

1213) -> Dict[str, Any]:

1214 """Flatten nested state dict to single-level dot-separated keys.

1215

1216 Format:

1217 state.<FQN>.<state_key> = value

1218 param_group.<FQN>.<group_field> = value

1219

1220 **MPMD / Pipeline Parallelism caveat**: in PP, different ranks may

1221 own parameters with the same local FQN (e.g. ``layers.0.net1.weight``

1222 exists on both stage 0 and stage 1). The flatten format does not

1223 include a PP-rank or stage identifier, so these FQNs would collide

1224 when combined into a single flat dict. The recommended PP checkpoint

1225 pattern is for each rank to independently save/load its own stages'

1226 optimizer state dicts using ``no_dist=True`` (see P6 test pattern),

1227 rather than using a coordinated DCP save with flatten format across

1228 all ranks.

1229 """

1230 flat: Dict[str, Any] = {}

1231

1232 for fqn, state in state_dict.get("state", {}).items():

1233 for key, value in state.items():

1234 flat[f"state.{fqn}.{key}"] = value

1235

1236 for group_idx, group in enumerate(state_dict.get("param_groups", [])):

1237 params_list = group.get("params", [])

1238 if not params_list:

1239 raise UnsupportedConfigurationError(

1240 f"Cannot flatten param_groups[{group_idx}]: empty 'params' list. "

1241 f"Empty param groups cannot be represented in flatten format "

1242 f"because there is no FQN to prefix the group fields. "

1243 f"Use non-flatten format or provide a stable group_name."

1244 )

1245 for fqn in params_list:

1246 for key, value in group.items():

1247 if key == "params":

1248 continue

1249 flat[f"param_group.{fqn}.{key}"] = value

1250

1251 return flat

1252

1253

1254def _check_inconsistent_pg_fields(

1255 existing: Dict[str, Any],

1256 incoming: Dict[str, Any],

1257 fqn: str,

1258 strict: bool,

1259) -> None:

1260 """Check and raise/warn on inconsistent param_group fields."""

1261 common_keys = set(existing.keys()) & set(incoming.keys())

1262 for k in common_keys:

1263 if existing[k] != incoming[k]:

1264 if strict:

1265 raise ValueError(

1266 f"strict=True but param_group field '{k}' is "

1267 f"inconsistent within the same group: "

1268 f"existing={existing[k]!r}, "

1269 f"incoming(from {fqn})={incoming[k]!r}"

1270 )

1271 logger.warning(

1272 "param_group field '%s' is inconsistent within the "

1273 "same group: existing=%r, incoming(from %s)=%r. "

1274 "Keeping existing value (strict=False).",

1275 k, existing[k], fqn, incoming[k],

1276 )

1277

1278

1279def _parse_unflatten_entries(

1280 flat_dict: Dict[str, Any],

1281 known_fqns_sorted: List[str],

1282) -> Tuple[Dict[str, Dict[str, Any]], Dict[str, Dict[str, Any]]]:

1283 """Parse flat dict entries into state and param_group_fields dicts.

1284

1285 Returns:

1286 Tuple of (state, param_group_fields).

1287 """

1288 state: Dict[str, Dict[str, Any]] = {}

1289 param_group_fields: Dict[str, Dict[str, Any]] = {}

1290

1291 for key, value in flat_dict.items():

1292 if key.startswith("state."):

1293 remainder = key[len("state."):]

1294 fqn = _match_fqn_from_remainder(remainder, known_fqns_sorted)

1295 if fqn is None:

1296 raise ValueError(

1297 f"Cannot match FQN from flat key '{key}'. "

1298 f"Known FQNs: {known_fqns_sorted}"

1299 )

1300 state_key = remainder[len(fqn) + 1:]

1301 state.setdefault(fqn, {})[state_key] = value

1302

1303 elif key.startswith("param_group."):

1304 remainder = key[len("param_group."):]

1305 fqn = _match_fqn_from_remainder(remainder, known_fqns_sorted)

1306 if fqn is None:

1307 raise ValueError(

1308 f"Cannot match FQN from flat key '{key}'. "

1309 f"Known FQNs: {known_fqns_sorted}"

1310 )

1311 field_name = remainder[len(fqn) + 1:]

1312 param_group_fields.setdefault(fqn, {})[field_name] = value

1313

1314 return state, param_group_fields

1315

1316

1317def _assemble_param_groups(

1318 known_fqns: List[str],

1319 state: Dict[str, Dict[str, Any]],

1320 param_group_fields: Dict[str, Dict[str, Any]],

1321 strict: bool,

1322) -> List[Dict[str, Any]]:

1323 """Assemble param_groups list from FQN-based state and fields."""

1324 param_groups: List[Dict[str, Any]] = []

1325 current_group: Dict[str, Any] = {"params": []}

1326

1327 for fqn in known_fqns:

1328 if fqn in param_group_fields or fqn in state:

1329 fields = param_group_fields.get(fqn, {})

1330 if not current_group["params"]:

1331 current_group["params"].append(fqn)

1332 current_group.update(fields)

1333 else:

1334 existing_fields = {

1335 k: v for k, v in current_group.items() if k != "params"

1336 }

1337 if existing_fields == fields or not fields:

1338 current_group["params"].append(fqn)

1339 else:

1340 _check_inconsistent_pg_fields(existing_fields, fields, fqn, strict)

1341 current_group["params"].append(fqn)

1342

1343 if current_group["params"]:

1344 param_groups.append(current_group)

1345

1346 if not param_groups:

1347 raise UnsupportedConfigurationError(

1348 "Cannot unflatten: empty param_group in flatten format. "

1349 "Provide a stable group_name or use non-flatten format."

1350 )

1351

1352 return param_groups

1353

1354

1355def _unflatten_optim_state_dict(

1356 flat_dict: Dict[str, Any],

1357 model: nn.Module,

1358 strict: bool = True,

1359) -> Dict[str, Any]:

1360 """Unflatten a flat state dict back to nested format.

1361

1362 Uses longest-prefix matching against model FQNs to handle FQNs

1363 containing dots.

1364

1365 Args:

1366 flat_dict: The flattened state dict.

1367 model: The model whose FQNs are used for prefix matching.

1368 strict: If True, raise ValueError when param_group fields are

1369 inconsistent within the same group. If False, keep the

1370 current group's values and log a warning.

1371

1372 Returns:

1373 The nested state dict with FQN keys.

1374 """

1375 known_fqns = [name for name, _ in model.named_parameters()]

1376 known_fqns_sorted = sorted(known_fqns, key=len, reverse=True)

1377

1378 state, param_group_fields = _parse_unflatten_entries(flat_dict, known_fqns_sorted)

1379 param_groups = _assemble_param_groups(known_fqns, state, param_group_fields, strict)

1380

1381 return {

1382 "state": state,

1383 "param_groups": param_groups,

1384 }

1385

1386

1387def _build_optim_state_dict_load_template(

1388 model: nn.Module,

1389 optimizer: torch.optim.Optimizer,

1390 storage_reader: Any,

1391 options: Any = None,

1392) -> Dict[str, Any]:

1393 """Build an empty optimizer state dict template for DCP load.

1394

1395 Reads checkpoint metadata from ``storage_reader`` to determine

1396 key, shape, and dtype for each tensor. This is needed when loading

1397 into a new optimizer that has no state yet (optimizer.step() has not

1398 been called), so ``optimizer.state_dict()`` returns empty state.

1399

1400 The template uses the checkpoint metadata to infer the correct keys

1401 and shapes, then creates zero-filled tensors so that ``dcp.load``

1402 has a valid target to write into.

1403

1404 Args:

1405 model: The model whose parameters are optimized.

1406 optimizer: The target optimizer (may have empty state).

1407 storage_reader: DCP StorageReader with ``load_metadata()`` method.

1408 options: Options for full_state_dict, cpu_offload,

1409 flatten_optimizer_state_dict, etc.

1410

1411 Returns:

1412 State dict template with empty tensors of correct shape/dtype.

1413 """

1414 full_state_dict = getattr(options, "full_state_dict", False)

1415 cpu_offload = getattr(options, "cpu_offload", False)

1416 flatten = getattr(options, "flatten_optimizer_state_dict", False)

1417

1418 try:

1419 metadata = storage_reader.load_metadata()

1420 except FileNotFoundError:

1421 rank = dist.get_rank() if dist.is_initialized() else 0

1422 metadata = storage_reader.load_metadata(rank=rank)

1423

1424 param_by_fqn: Dict[str, nn.Parameter] = {}

1425 for name, param in model.named_parameters():

1426 param_by_fqn[name] = param

1427

1428 if flatten:

1429 return _build_flatten_template_from_metadata(

1430 metadata, param_by_fqn, full_state_dict, cpu_offload,

1431 )

1432

1433 return _build_nested_template_from_metadata(

1434 metadata, optimizer, param_by_fqn, full_state_dict, cpu_offload,

1435 )

1436

1437

1438def _resolve_dtype(dtype_str: str) -> torch.dtype:

1439 """Resolve a dtype string like 'torch.float32' to torch.dtype."""

1440 try:

1441 return getattr(torch, dtype_str.replace("torch.", ""))

1442 except AttributeError:

1443 return torch.float32

1444

1445

1446def _match_fqn_from_remainder(

1447 remainder: str,

1448 known_fqns_sorted: List[str],

1449) -> Optional[str]:

1450 """Match a remainder string against known FQNs using longest-prefix match."""

1451 for fqn in known_fqns_sorted:

1452 if remainder == fqn or remainder.startswith(fqn + "."):

1453 return fqn

1454 return None

1455

1456

1457def _build_saved_id_to_fqn(

1458 optimizer: torch.optim.Optimizer,

1459 param_by_fqn: Dict[str, nn.Parameter],

1460) -> Dict[int, str]:

1461 """Build mapping from optimizer saved IDs to FQNs."""

1462 raw_sd = optimizer.state_dict()

1463 saved_id_to_fqn: Dict[int, str] = {}

1464 for runtime_group, saved_group in zip(optimizer.param_groups, raw_sd["param_groups"]):

1465 for parameter, saved_id in zip(runtime_group["params"], saved_group["params"]):

1466 for name, p in param_by_fqn.items():

1467 if p is parameter:

1468 saved_id_to_fqn[saved_id] = name

1469 break

1470 return saved_id_to_fqn

1471

1472

1473def _parse_nested_state_entry(

1474 meta_key: str,

1475 meta_val: Any,

1476 param_by_fqn: Dict[str, nn.Parameter],

1477 full_state_dict: bool,

1478 cpu_offload: bool,

1479) -> Optional[Tuple[str, str, torch.Tensor]]:

1480 """Parse a state.* metadata entry and return (fqn, state_key, tensor) or None."""

1481 # DCP re-exports these state helpers; defer metadata to avoid that package import cycle.

1482 from hyper_parallel.core.distributed_checkpoint.metadata import ( # pylint: disable=import-outside-toplevel

1483 TensorStorageMetadata,

1484 )

1485

1486 remainder = meta_key[len("state."):]

1487 fqns_sorted = sorted(param_by_fqn.keys(), key=len, reverse=True)

1488 matched_fqn = _match_fqn_from_remainder(remainder, fqns_sorted)

1489 if matched_fqn is None:

1490 return None

1491

1492 state_key = remainder[len(matched_fqn) + 1:] if len(remainder) > len(matched_fqn) else None

1493 if state_key is None:

1494 return None

1495

1496 if not isinstance(meta_val, TensorStorageMetadata):

1497 return None

1498

1499 dtype = _resolve_dtype(meta_val.properties.dtype)

1500 global_shape = meta_val.size

1501 param = param_by_fqn.get(matched_fqn)

1502

1503 if _is_scalar_state(state_key):

1504 device = torch.device("cpu")

1505 t = torch.zeros(global_shape, dtype=dtype, device=device)

1506 else:

1507 t = _create_empty_state_tensor(param, global_shape, dtype, full_state_dict, cpu_offload)

1508

1509 return matched_fqn, state_key, t

1510

1511

1512def _parse_nested_param_group_entry(

1513 meta_key: str,

1514 meta_val: Any,

1515) -> Optional[Tuple[int, str, torch.Tensor]]:

1516 """Parse a param_group.* metadata entry and return (group_idx, field_name, tensor) or None."""

1517 # DCP re-exports these state helpers; defer metadata to avoid that package import cycle.

1518 from hyper_parallel.core.distributed_checkpoint.metadata import ( # pylint: disable=import-outside-toplevel

1519 TensorStorageMetadata,

1520 )

1521

1522 remainder = meta_key[len("param_group."):]

1523 parts = remainder.split(".", 1)

1524 if len(parts) < 2:

1525 return None

1526 try:

1527 group_idx = int(parts[0])

1528 except ValueError:

1529 return None

1530 field_name = parts[1]

1531

1532 if not isinstance(meta_val, TensorStorageMetadata):

1533 return None

1534

1535 dtype = _resolve_dtype(meta_val.properties.dtype)

1536 return group_idx, field_name, torch.zeros(meta_val.size, dtype=dtype)

1537

1538

1539def _build_nested_template_from_metadata(

1540 metadata: Any,

1541 optimizer: torch.optim.Optimizer,

1542 param_by_fqn: Dict[str, nn.Parameter],

1543 full_state_dict: bool,

1544 cpu_offload: bool,

1545) -> Dict[str, Any]:

1546 """Build nested (non-flatten) template from checkpoint metadata.

1547

1548 Non-flatten checkpoint keys look like:

1549 state.<FQN>.<state_key> -> TensorStorageMetadata

1550 param_group.<group_idx>.<field> -> TensorStorageMetadata or BytesStorageMetadata

1551

1552 We reconstruct:

1553 {

1554 "state": { "<FQN>": { "<state_key>": empty_tensor, ... }, ... },

1555 "param_groups": [ { "params": [...], "<field>": value, ... }, ... ]

1556 }

1557 """

1558 state: Dict[str, Dict[str, Any]] = {}

1559 param_groups_raw: Dict[int, Dict[str, Any]] = {}

1560 param_groups_fqns: Dict[int, List[str]] = {}

1561

1562 saved_id_to_fqn = _build_saved_id_to_fqn(optimizer, param_by_fqn)

1563 raw_sd = optimizer.state_dict()

1564

1565 for group_idx, saved_group in enumerate(raw_sd["param_groups"]):

1566 param_groups_raw[group_idx] = {

1567 k: v for k, v in saved_group.items() if k != "params"

1568 }

1569 param_groups_fqns[group_idx] = [

1570 saved_id_to_fqn[sid] for sid in saved_group["params"]

1571 ]

1572

1573 for meta_key, meta_val in metadata.state_dict_metadata.items():

1574 if meta_key.startswith("state."):

1575 result = _parse_nested_state_entry(

1576 meta_key, meta_val, param_by_fqn, full_state_dict, cpu_offload,

1577 )

1578 if result is not None:

1579 fqn, state_key, t = result

1580 state.setdefault(fqn, {})[state_key] = t

1581 elif meta_key.startswith("param_group."):

1582 result = _parse_nested_param_group_entry(meta_key, meta_val)

1583 if result is not None:

1584 group_idx, field_name, t = result

1585 param_groups_raw.setdefault(group_idx, {})[field_name] = t

1586

1587 result_param_groups: List[Dict[str, Any]] = []

1588 for group_idx in sorted(param_groups_fqns.keys()):

1589 pg: Dict[str, Any] = {"params": param_groups_fqns[group_idx]}

1590 pg.update(param_groups_raw.get(group_idx, {}))

1591 result_param_groups.append(pg)

1592

1593 return {

1594 "state": state,

1595 "param_groups": result_param_groups,

1596 }

1597

1598

1599def _build_flatten_template_from_metadata(

1600 metadata: Any,

1601 param_by_fqn: Dict[str, nn.Parameter],

1602 full_state_dict: bool,

1603 cpu_offload: bool,

1604) -> Dict[str, Any]:

1605 """Build flatten template from checkpoint metadata.

1606

1607 Flatten checkpoint keys look like:

1608 state.<FQN>.<state_key> -> TensorStorageMetadata

1609 param_group.<FQN>.<field> -> TensorStorageMetadata or BytesStorageMetadata

1610 """

1611 # DCP re-exports these state helpers; defer metadata to avoid that package import cycle.

1612 from hyper_parallel.core.distributed_checkpoint.metadata import ( # pylint: disable=import-outside-toplevel

1613 TensorStorageMetadata,

1614 )

1615

1616 flat: Dict[str, Any] = {}

1617 fqns_sorted = sorted(param_by_fqn.keys(), key=len, reverse=True)

1618

1619 for meta_key, meta_val in metadata.state_dict_metadata.items():

1620 if meta_key.startswith("state."):

1621 remainder = meta_key[len("state."):]

1622 matched_fqn = _match_fqn_from_remainder(remainder, fqns_sorted)

1623 if matched_fqn is None:

1624 continue

1625

1626 state_key = remainder[len(matched_fqn) + 1:] if len(remainder) > len(matched_fqn) else None

1627 if state_key is None:

1628 continue

1629

1630 if isinstance(meta_val, TensorStorageMetadata):

1631 dtype = _resolve_dtype(meta_val.properties.dtype)

1632 global_shape = meta_val.size

1633 param = param_by_fqn.get(matched_fqn)

1634

1635 if _is_scalar_state(state_key):

1636 t = torch.zeros(global_shape, dtype=dtype)

1637 else:

1638 t = _create_empty_state_tensor(

1639 param, global_shape, dtype,

1640 full_state_dict, cpu_offload,

1641 )

1642 flat[meta_key] = t

1643

1644 elif meta_key.startswith("param_group."):

1645 if isinstance(meta_val, TensorStorageMetadata):

1646 dtype = _resolve_dtype(meta_val.properties.dtype)

1647 flat[meta_key] = torch.zeros(meta_val.size, dtype=dtype)

1648

1649 return flat

1650

1651

1652def _create_empty_state_tensor(

1653 param: Optional[nn.Parameter],

1654 global_shape: tuple,

1655 dtype: torch.dtype,

1656 full_state_dict: bool,

1657 cpu_offload: bool,

1658) -> torch.Tensor:

1659 """Create an empty state tensor with appropriate shape and device.

1660

1661 For DTensor params, determines whether to use full or local shape

1662 based on full_state_dict. For plain params, uses the global shape.

1663

1664 Args:

1665 param: The model parameter (may be DTensor or None).

1666 global_shape: Global shape from checkpoint metadata.

1667 dtype: Tensor dtype.

1668 full_state_dict: If True, use global shape; otherwise use local.

1669 cpu_offload: If True, place on CPU.

1670

1671 Returns:

1672 Empty tensor with correct shape, dtype, and device.

1673 """

1674 dtensor_info = _get_param_dtensor_info(param) if param is not None else None

1675

1676 if dtensor_info is not None:

1677 if full_state_dict:

1678 shape = global_shape

1679 device = torch.device("cpu") if cpu_offload else param.to_local().device

1680 else:

1681 shape = param.to_local().shape

1682 device = torch.device("cpu") if cpu_offload else param.to_local().device

1683 else:

1684 shape = global_shape

1685 device = torch.device("cpu") if cpu_offload else (

1686 param.data.device if param is not None else torch.device("cpu")

1687 )

1688

1689 return torch.zeros(shape, dtype=dtype, device=device)

1690

1691

1692def _infer_state_keys(optimizer: torch.optim.Optimizer) -> List[str]:

1693 """Infer expected state keys from optimizer defaults.

1694

1695 Supports Adam, AdamW, and SGD. Returns keys that would appear

1696 after at least one optimizer.step() call.

1697 """

1698 opt_cls = type(optimizer)

1699 opt_name = opt_cls.__name__.lower()

1700

1701 if "adam" in opt_name:

1702 keys = ["step", "exp_avg", "exp_avg_sq"]

1703 if optimizer.defaults.get("amsgrad", False):

1704 keys.append("max_exp_avg_sq")

1705 return keys

1706

1707 if "sgd" in opt_name:

1708 keys = []

1709 if optimizer.defaults.get("momentum", 0) != 0:

1710 keys.append("momentum_buffer")

1711 return keys

1712

1713 logger.warning(

1714 "Cannot infer state keys for optimizer type '%s'. "

1715 "DCP load template may be incomplete.",

1716 opt_cls.__name__,

1717 )

1718 return []

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / state_dict_utils.py: 44%

680 statements