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) -> None:

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

241

242 Behaviour matrix:

243

244 +-----------------+-------------+--------------------------------+

245 | full_state_dict | cpu_offload | behaviour |

246 +=================+=============+================================+

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

248 +-----------------+-------------+--------------------------------+

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

250 +-----------------+-------------+--------------------------------+

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

252 +-----------------+-------------+--------------------------------+

253

254 Note:

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

256 mirroring torch.distributed.checkpoint.state_dict: frozen parameters

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

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

259

260 Args:

261 model: The model to load state into.

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

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

264 options: Controls full_state_dict, cpu_offload, strict and

265 broadcast_from_rank0 flags.

266

267 Raises:

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

269 ``full_state_dict`` is False.

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

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

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

273 tensor on every rank instead.

274 """

275 options = options or StateDictOptions()

276

277 if options.broadcast_from_rank0 and not options.full_state_dict:

278 raise ValueError(

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

280 )

281 if options.broadcast_from_rank0:

282 raise NotImplementedError(

283 "broadcast_from_rank0=True is not supported. "

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

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

286 "tensor on every rank instead."

287 )

288

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

290 if options.full_state_dict:

291 scattered = _scatter_model_state_dict(

292 model, state_dict, options.cpu_offload, options.strict,

293 )

294 else:

295 scattered = state_dict

296

297 model.load_state_dict(scattered, strict=options.strict, assign=True)

298

299

300# =====================================================================

301# Optimizer state dict helpers

302# =====================================================================

303

304_SCALAR_STATE_KEYS = {"step"}

305

306

307class UnsupportedConfigurationError(RuntimeError):

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

309

310

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

312 return isinstance(param, DTensor)

313

314

315def _is_replicate_group_root(

316 mesh: DeviceMesh,

317 placements: Sequence,

318) -> bool:

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

320 Replicate dimension of the given mesh+placements.

321

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

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

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

325

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

327 """

328 coord = mesh.get_coordinate()

329 if coord is None:

330 return False

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

332 if not replicate_dims:

333 return dist.get_rank() == 0

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

335

336

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

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

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

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

341 if param in param_to_name:

342 raise ValueError(

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

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

345 )

346 param_to_name[param] = name

347 return param_to_name

348

349

350def _build_id_to_fqn(

351 optimizer: torch.optim.Optimizer,

352 model: nn.Module,

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

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

355

356 Returns:

357 Tuple of (saved_id_to_fqn, fqn_to_saved_id).

358 """

359 param_to_name = _param_to_fqn(model)

360 raw_sd = optimizer.state_dict()

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

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

363

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

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

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

367 raise ValueError(

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

369 f"this is not supported."

370 )

371 saved_id_to_param[saved_id] = parameter

372

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

374 if param not in param_to_name:

375 raise ValueError(

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

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

378 )

379 saved_id_to_fqn[saved_id] = param_to_name[param]

380

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

382 return saved_id_to_fqn, fqn_to_saved_id

383

384

385def _get_param_dtensor_info(

386 param: nn.Parameter,

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

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

389 if _is_dtensor_param(param):

390 return param.device_mesh, param.placements

391 return None

392

393

394def _convert_state_tensor(

395 tensor: torch.Tensor,

396 param: nn.Parameter,

397 full_state_dict: bool,

398 cpu_offload: bool,

399 is_rank0: bool,

400) -> Optional[torch.Tensor]:

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

402

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

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

405

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

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

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

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

410 than wrapping it again.

411

412 Args:

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

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

415 full_state_dict: If True, gather to full tensor.

416 cpu_offload: If True, move to CPU.

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

418

419 Returns:

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

421 """

422 dtensor_info = _get_param_dtensor_info(param)

423 is_dtensor_value = isinstance(tensor, DTensor)

424

425 if full_state_dict and dtensor_info is not None:

426 if is_dtensor_value:

427 full_t = tensor.full_tensor()

428 else:

429 mesh, placements = dtensor_info

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

431 if cpu_offload:

432 if not is_rank0:

433 return None

434 return full_t.cpu()

435 return full_t

436

437 if is_dtensor_value:

438 local_t = tensor.to_local()

439 if cpu_offload:

440 return local_t.cpu()

441 return local_t

442

443 if cpu_offload:

444 return tensor.cpu()

445

446 return tensor

447

448

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

450 return key in _SCALAR_STATE_KEYS

451

452

453def _convert_state_scalar(

454 tensor: torch.Tensor,

455 cpu_offload: bool,

456) -> torch.Tensor:

457 if isinstance(tensor, DTensor):

458 tensor = tensor.to_local()

459 if cpu_offload:

460 return tensor.cpu()

461 return tensor

462

463

464def _determine_is_root(

465 full_state_dict: bool,

466 cpu_offload: bool,

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

468) -> bool:

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

470 if full_state_dict and cpu_offload and dtensor_info is not None:

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

472 if full_state_dict and cpu_offload:

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

474 return True

475

476

477def _convert_state_entries(

478 state: Dict[str, Any],

479 param: nn.Parameter,

480 full_state_dict: bool,

481 cpu_offload: bool,

482) -> Dict[str, Any]:

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

484 dtensor_info = _get_param_dtensor_info(param)

485 is_root = _determine_is_root(full_state_dict, cpu_offload, dtensor_info)

486 skip_non_root = full_state_dict and cpu_offload and not is_root

487

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

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

490 if isinstance(value, torch.Tensor):

491 if _is_scalar_state(key):

492 result = _convert_state_scalar(value, cpu_offload)

493 if not skip_non_root:

494 converted[key] = result

495 else:

496 result = _convert_state_tensor(

497 value, param, full_state_dict, cpu_offload, is_root,

498 )

499 if result is not None:

500 converted[key] = result

501 else:

502 if not skip_non_root:

503 converted[key] = value

504 return converted

505

506

507def _build_result_param_groups(

508 optimizer: torch.optim.Optimizer,

509 raw_sd: Dict[str, Any],

510 saved_id_to_fqn: Dict[int, str],

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

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

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

514 for saved_group in raw_sd["param_groups"]:

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

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

517 if k == "params":

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

519 else:

520 pg[k] = v

521 result_param_groups.append(pg)

522 return result_param_groups

523

524

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

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

527 raise ValueError(

528 "ChainedOptimizer is not supported by get_optim_state_dict / "

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

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

531 "multi-optimizer merging internally."

532 )

533

534

535def get_optim_state_dict(

536 model: nn.Module,

537 optimizer: torch.optim.Optimizer,

538 *,

539 options: Any = None,

540) -> Dict[str, Any]:

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

542

543 Args:

544 model: The model whose parameters are optimized.

545 optimizer: The optimizer instance.

546 options: Options object with fields:

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

548 cpu_offload (bool): Move tensors to CPU.

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

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

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

552

553 Returns:

554 Optimizer state dict with FQN-based keys.

555 """

556 _check_chained_optimizer(optimizer)

557

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

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

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

561

562 raw_sd = optimizer.state_dict()

563 saved_id_to_fqn, _ = _build_id_to_fqn(optimizer, model)

564

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

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

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

568 param_by_id[saved_id] = parameter

569

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

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

572 fqn = saved_id_to_fqn[saved_id]

573 param = param_by_id[saved_id]

574 converted = _convert_state_entries(

575 state, param, full_state_dict, cpu_offload,

576 )

577 if converted:

578 result_state[fqn] = converted

579

580 result_param_groups = _build_result_param_groups(

581 optimizer, raw_sd, saved_id_to_fqn,

582 )

583

584 result: Dict[str, Any] = {

585 "state": result_state,

586 "param_groups": result_param_groups,

587 }

588

589 if flatten:

590 result = _flatten_optim_state_dict(result)

591

592 return result

593

594

595def _check_strict_fqns(

596 source_fqns: set,

597 target_fqns: set,

598 strict: bool,

599) -> None:

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

601 if not strict:

602 return

603 extra_fqns = source_fqns - target_fqns

604 if extra_fqns:

605 raise ValueError(

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

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

608 )

609

610

611def _load_state_values(

612 source_state: Dict[str, Any],

613 param: nn.Parameter,

614 full_state_dict: bool,

615 cpu_offload: bool,

616 stores_dtensor: bool,

617) -> Dict[str, Any]:

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

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

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

621 if isinstance(value, torch.Tensor):

622 if _is_scalar_state(key):

623 converted[key] = _convert_input_scalar_to_target(

624 value, param, cpu_offload, stores_dtensor,

625 )

626 else:

627 converted[key] = _convert_input_tensor_to_target(

628 value, param, full_state_dict, cpu_offload, stores_dtensor,

629 )

630 else:

631 converted[key] = value

632 return converted

633

634

635def _merge_param_groups(

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

637 target_raw_sd: Dict[str, Any],

638 target_fqn_to_saved_id: Dict[str, int],

639) -> None:

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

641 if not source_param_groups:

642 return

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

644

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

646 for saved_group in target_raw_sd["param_groups"]:

647 fqns_in_group = frozenset(

648 target_saved_id_to_fqn.get(sid, "")

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

650 )

651 target_pg_by_fqns[fqns_in_group] = saved_group

652

653 for source_pg in source_param_groups:

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

655 matched_target_pg = None

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

657 if source_fqn_set & target_fqn_set:

658 matched_target_pg = target_pg

659 break

660 if matched_target_pg is None:

661 continue

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

663 if k == "params":

664 continue

665 matched_target_pg[k] = v

666

667

668def _build_target_fqn_mappings(

669 optimizer: torch.optim.Optimizer,

670 model: nn.Module,

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

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

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

674

675 Returns:

676 Tuple of (fqn_to_saved_id, dtensor_state_ids).

677 """

678 target_raw_sd = optimizer.state_dict()

679

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

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

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

683 target_id_to_param[saved_id] = parameter

684

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

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

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

688 if p is param:

689 fqn_to_saved_id[name] = saved_id

690 break

691

692 dtensor_state_ids: set = set()

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

694 for value in state.values():

695 if isinstance(value, DTensor):

696 dtensor_state_ids.add(saved_id)

697 break

698

699 return fqn_to_saved_id, dtensor_state_ids

700

701

702def _load_fqn_states(

703 optim_state_dict: Dict[str, Any],

704 target_raw_sd: Dict[str, Any],

705 target_fqn_to_saved_id: Dict[str, int],

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

707 dtensor_state_ids: set,

708 full_state_dict: bool,

709 cpu_offload: bool,

710 strict: bool,

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

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

713

714 Returns:

715 Updated state dict with saved_id keys.

716 """

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

718

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

720 if fqn not in target_fqn_to_saved_id:

721 if strict:

722 raise ValueError(

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

724 )

725 continue

726

727 target_saved_id = target_fqn_to_saved_id[fqn]

728 param = param_by_fqn[fqn]

729

730 if target_saved_id not in new_state:

731 new_state[target_saved_id] = {}

732

733 new_state[target_saved_id].update(

734 _load_state_values(

735 source_state, param, full_state_dict, cpu_offload,

736 target_saved_id in dtensor_state_ids,

737 )

738 )

739

740 return new_state

741

742

743def set_optim_state_dict(

744 model: nn.Module,

745 optimizer: torch.optim.Optimizer,

746 optim_state_dict: Dict[str, Any],

747 *,

748 options: Any = None,

749) -> None:

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

751

752 Args:

753 model: The model whose parameters are optimized.

754 optimizer: The target optimizer instance.

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

756 options: Options object with fields:

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

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

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

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

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

762 always allowed because the target optimizer may have

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

764 FQNs in the checkpoint are silently ignored.

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

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

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

768 broadcasts independently within its own HSDP subgroup.

769 """

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

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

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

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

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

775

776 _check_chained_optimizer(optimizer)

777

778 if flatten:

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

780

781 if full_state_dict and cpu_offload and not broadcast_from_rank0:

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

783 if not has_any_state:

784 raise ValueError(

785 "Received empty state dict with full_state_dict=True "

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

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

788 "root to broadcast."

789 )

790

791 if broadcast_from_rank0:

792 optim_state_dict = _broadcast_state_from_rank0(

793 optim_state_dict, model, full_state_dict, cpu_offload,

794 )

795

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

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

798 param_by_fqn[name] = param

799

800 target_fqn_to_saved_id, dtensor_state_ids = _build_target_fqn_mappings(

801 optimizer, model, param_by_fqn,

802 )

803

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

805 target_fqns = set(target_fqn_to_saved_id.keys())

806 _check_strict_fqns(source_fqns, target_fqns, strict)

807

808 target_raw_sd = optimizer.state_dict()

809 new_state = _load_fqn_states(

810 optim_state_dict, target_raw_sd, target_fqn_to_saved_id,

811 param_by_fqn, dtensor_state_ids, full_state_dict, cpu_offload, strict,

812 )

813

814 target_raw_sd["state"] = new_state

815

816 _merge_param_groups(

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

818 target_raw_sd,

819 target_fqn_to_saved_id,

820 )

821

822 optimizer.load_state_dict(target_raw_sd)

823

824

825def _convert_input_tensor_to_target(

826 tensor: torch.Tensor,

827 param: nn.Parameter,

828 full_state_dict: bool,

829 cpu_offload: bool,

830 stores_dtensor: bool = False,

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

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

833

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

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

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

837

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

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

840 a plain local-shard tensor.

841

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

843 device.

844

845 Args:

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

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

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

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

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

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

852

853 Returns:

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

855 """

856 dtensor_info = _get_param_dtensor_info(param)

857

858 if not dtensor_info:

859 target_device = param.data.device

860 result = tensor

861 if cpu_offload and result.device != target_device:

862 result = result.to(target_device)

863 return result

864

865 mesh, placements = dtensor_info

866 target_device = param.to_local().device

867

868 if isinstance(tensor, DTensor):

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

870 if stores_dtensor:

871 return tensor

872 return tensor.to_local()

873 redistributed = tensor.redistribute(mesh, placements)

874 if stores_dtensor:

875 return redistributed

876 return redistributed.to_local()

877

878 if full_state_dict:

879 if cpu_offload:

880 tensor = tensor.to(target_device)

881 dt = distribute_tensor(tensor, mesh, placements)

882 if stores_dtensor:

883 return dt

884 return dt.to_local()

885

886 if cpu_offload and tensor.device != target_device:

887 tensor = tensor.to(target_device)

888

889 if stores_dtensor:

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

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

892 dt = distribute_tensor(tensor, mesh, placements)

893 return dt

894

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

896 return tensor

897

898 dt = distribute_tensor(tensor, mesh, placements)

899 return dt.to_local()

900

901

902def _convert_input_scalar_to_target(

903 tensor: torch.Tensor,

904 param: nn.Parameter,

905 cpu_offload: bool,

906 stores_dtensor: bool = False,

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

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

909

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

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

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

913 the correct device.

914

915 Args:

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

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

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

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

920

921 Returns:

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

923 """

924 dtensor_info = _get_param_dtensor_info(param)

925

926 if not dtensor_info or not stores_dtensor:

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

928 result = tensor

929 if cpu_offload and result.device != target_device:

930 result = result.to(target_device)

931 return result

932

933 mesh, _ = dtensor_info

934 target_device = param.to_local().device

935

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

937 if cpu_offload and tensor.device != target_device:

938 tensor = tensor.to(target_device)

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

940

941

942def _get_broadcast_groups(

943 model: nn.Module,

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

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

946

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

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

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

950 broadcast group because different stages own different parameters.

951

952 Returns:

953 List of dicts, each with keys:

954 fqns: list of FQN strings owned by this group

955 pg: ProcessGroup for broadcast within the group

956 src_rank: global rank of the replicate-group root

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

958 """

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

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

961 param_by_fqn[name] = param

962

963 model_fqns = list(param_by_fqn.keys())

964

965 if not model_fqns:

966 return []

967

968 first_param = param_by_fqn[model_fqns[0]]

969 dtensor_info = _get_param_dtensor_info(first_param)

970

971 if dtensor_info is not None:

972 mesh, placements = dtensor_info

973 is_root = _is_replicate_group_root(mesh, placements)

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

975 if replicate_dims:

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

977 coord = mesh.get_coordinate()

978 if coord is not None:

979 root_coord = list(coord)

980 for d in replicate_dims:

981 root_coord[d] = 0

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

983 else:

984 src_rank = 0

985 else:

986 pg = dist.group.WORLD

987 src_rank = 0

988 else:

989 is_root = dist.get_rank() == 0

990 pg = dist.group.WORLD

991 src_rank = 0

992

993 return [{

994 "fqns": model_fqns,

995 "pg": pg,

996 "src_rank": src_rank,

997 "is_root": is_root,

998 "param_by_fqn": param_by_fqn,

999 }]

1000

1001

1002def _broadcast_fqn_list(

1003 is_root: bool,

1004 fqns: List[str],

1005 pg: dist.ProcessGroup,

1006 src_rank: int,

1007) -> List[str]:

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

1009 fqn_list = fqns if is_root else []

1010 obj = [fqn_list]

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

1012 return obj[0]

1013

1014

1015def _build_state_schema(

1016 state: Dict[str, Any],

1017) -> Dict[str, Any]:

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

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

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

1021 if isinstance(value, torch.Tensor):

1022 schema[key] = {

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

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

1025 "is_scalar": _is_scalar_state(key),

1026 }

1027 else:

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

1029 return schema

1030

1031

1032def _broadcast_schema_per_fqn(

1033 fqn_list: List[str],

1034 optim_state_dict: Dict[str, Any],

1035 is_root: bool,

1036 pg: dist.ProcessGroup,

1037 src_rank: int,

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

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

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

1041 for fqn in fqn_list:

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

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

1044 else:

1045 schema = {}

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

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

1048 fqn_schema[fqn] = schema_list[0]

1049 return fqn_schema

1050

1051

1052def _get_local_device() -> torch.device:

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

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

1055 if torch.npu.is_available():

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

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

1058

1059

1060def _broadcast_scalar_entry(

1061 key: str,

1062 fqn: str,

1063 info: Dict[str, Any],

1064 optim_state_dict: Dict[str, Any],

1065 is_root: bool,

1066 src_rank: int,

1067 pg: dist.ProcessGroup,

1068 cpu_offload: bool,

1069) -> torch.Tensor:

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

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

1072 dtype = _resolve_dtype(dtype_str)

1073 device = _get_local_device()

1074

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

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

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

1078 if t.dim() == 0:

1079 t = t.reshape(1)

1080 else:

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

1082

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

1084 t = t.to(device)

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

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

1087

1088

1089def _broadcast_tensor_entry(

1090 key: str,

1091 fqn: str,

1092 info: Dict[str, Any],

1093 optim_state_dict: Dict[str, Any],

1094 is_root: bool,

1095 src_rank: int,

1096 pg: dist.ProcessGroup,

1097 full_state_dict: bool,

1098 cpu_offload: bool,

1099) -> torch.Tensor:

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

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

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

1103 dtype = _resolve_dtype(dtype_str)

1104 device = _get_local_device()

1105

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

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

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

1109 src_tensor = src_tensor.to(device)

1110 else:

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

1112

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

1114

1115 if full_state_dict and cpu_offload:

1116 return src_tensor.cpu() if is_root else src_tensor

1117 return src_tensor.cpu() if cpu_offload else src_tensor

1118

1119

1120def _broadcast_tensor_data_per_fqn(

1121 fqn_list: List[str],

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

1123 optim_state_dict: Dict[str, Any],

1124 is_root: bool,

1125 src_rank: int,

1126 pg: dist.ProcessGroup,

1127 full_state_dict: bool,

1128 cpu_offload: bool,

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

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

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

1132 for fqn in fqn_list:

1133 schema = fqn_schema[fqn]

1134 if not schema:

1135 continue

1136 result_state[fqn] = {}

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

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

1139 if is_scalar:

1140 result_state[fqn][key] = _broadcast_scalar_entry(

1141 key, fqn, info, optim_state_dict,

1142 is_root, src_rank, pg, cpu_offload,

1143 )

1144 else:

1145 result_state[fqn][key] = _broadcast_tensor_entry(

1146 key, fqn, info, optim_state_dict,

1147 is_root, src_rank, pg, full_state_dict, cpu_offload,

1148 )

1149 return result_state

1150

1151

1152def _broadcast_state_from_rank0(

1153 optim_state_dict: Dict[str, Any],

1154 model: nn.Module,

1155 full_state_dict: bool,

1156 cpu_offload: bool,

1157) -> Dict[str, Any]:

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

1159

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

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

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

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

1164

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

1166 as the broadcast source.

1167

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

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

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

1171 subgroup), so different PP stages broadcast independently without

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

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

1174 """

1175 if not dist.is_initialized():

1176 return optim_state_dict

1177

1178 groups = _get_broadcast_groups(model)

1179

1180 if not groups:

1181 return {

1182 "state": {},

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

1184 }

1185

1186 group = groups[0]

1187 pg = group["pg"]

1188 src_rank = group["src_rank"]

1189 is_root = group["is_root"]

1190 fqns = group["fqns"]

1191

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

1193 fqn_schema = _broadcast_schema_per_fqn(

1194 fqn_list, optim_state_dict, is_root, pg, src_rank,

1195 )

1196 result_state = _broadcast_tensor_data_per_fqn(

1197 fqn_list, fqn_schema, optim_state_dict,

1198 is_root, src_rank, pg, full_state_dict, cpu_offload,

1199 )

1200

1201 return {

1202 "state": result_state,

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

1204 }

1205

1206

1207def _flatten_optim_state_dict(

1208 state_dict: Dict[str, Any],

1209) -> Dict[str, Any]:

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

1211

1212 Format:

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

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

1215

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

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

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

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

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

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

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

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

1224 all ranks.

1225 """

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

1227

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

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

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

1231

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

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

1234 if not params_list:

1235 raise UnsupportedConfigurationError(

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

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

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

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

1240 )

1241 for fqn in params_list:

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

1243 if key == "params":

1244 continue

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

1246

1247 return flat

1248

1249

1250def _check_inconsistent_pg_fields(

1251 existing: Dict[str, Any],

1252 incoming: Dict[str, Any],

1253 fqn: str,

1254 strict: bool,

1255) -> None:

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

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

1258 for k in common_keys:

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

1260 if strict:

1261 raise ValueError(

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

1263 f"inconsistent within the same group: "

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

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

1266 )

1267 logger.warning(

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

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

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

1271 k, existing[k], fqn, incoming[k],

1272 )

1273

1274

1275def _parse_unflatten_entries(

1276 flat_dict: Dict[str, Any],

1277 known_fqns_sorted: List[str],

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

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

1280

1281 Returns:

1282 Tuple of (state, param_group_fields).

1283 """

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

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

1286

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

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

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

1290 fqn = _match_fqn_from_remainder(remainder, known_fqns_sorted)

1291 if fqn is None:

1292 raise ValueError(

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

1294 f"Known FQNs: {known_fqns_sorted}"

1295 )

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

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

1298

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

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

1301 fqn = _match_fqn_from_remainder(remainder, known_fqns_sorted)

1302 if fqn is None:

1303 raise ValueError(

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

1305 f"Known FQNs: {known_fqns_sorted}"

1306 )

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

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

1309

1310 return state, param_group_fields

1311

1312

1313def _assemble_param_groups(

1314 known_fqns: List[str],

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

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

1317 strict: bool,

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

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

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

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

1322

1323 for fqn in known_fqns:

1324 if fqn in param_group_fields or fqn in state:

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

1326 if not current_group["params"]:

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

1328 current_group.update(fields)

1329 else:

1330 existing_fields = {

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

1332 }

1333 if existing_fields == fields or not fields:

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

1335 else:

1336 _check_inconsistent_pg_fields(existing_fields, fields, fqn, strict)

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

1338

1339 if current_group["params"]:

1340 param_groups.append(current_group)

1341

1342 if not param_groups:

1343 raise UnsupportedConfigurationError(

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

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

1346 )

1347

1348 return param_groups

1349

1350

1351def _unflatten_optim_state_dict(

1352 flat_dict: Dict[str, Any],

1353 model: nn.Module,

1354 strict: bool = True,

1355) -> Dict[str, Any]:

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

1357

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

1359 containing dots.

1360

1361 Args:

1362 flat_dict: The flattened state dict.

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

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

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

1366 current group's values and log a warning.

1367

1368 Returns:

1369 The nested state dict with FQN keys.

1370 """

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

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

1373

1374 state, param_group_fields = _parse_unflatten_entries(flat_dict, known_fqns_sorted)

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

1376

1377 return {

1378 "state": state,

1379 "param_groups": param_groups,

1380 }

1381

1382

1383def _build_optim_state_dict_load_template(

1384 model: nn.Module,

1385 optimizer: torch.optim.Optimizer,

1386 storage_reader: Any,

1387 options: Any = None,

1388) -> Dict[str, Any]:

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

1390

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

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

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

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

1395

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

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

1398 has a valid target to write into.

1399

1400 Args:

1401 model: The model whose parameters are optimized.

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

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

1404 options: Options for full_state_dict, cpu_offload,

1405 flatten_optimizer_state_dict, etc.

1406

1407 Returns:

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

1409 """

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

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

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

1413

1414 try:

1415 metadata = storage_reader.load_metadata()

1416 except FileNotFoundError:

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

1418 metadata = storage_reader.load_metadata(rank=rank)

1419

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

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

1422 param_by_fqn[name] = param

1423

1424 if flatten:

1425 return _build_flatten_template_from_metadata(

1426 metadata, param_by_fqn, full_state_dict, cpu_offload,

1427 )

1428

1429 return _build_nested_template_from_metadata(

1430 metadata, optimizer, param_by_fqn, full_state_dict, cpu_offload,

1431 )

1432

1433

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

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

1436 try:

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

1438 except AttributeError:

1439 return torch.float32

1440

1441

1442def _match_fqn_from_remainder(

1443 remainder: str,

1444 known_fqns_sorted: List[str],

1445) -> Optional[str]:

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

1447 for fqn in known_fqns_sorted:

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

1449 return fqn

1450 return None

1451

1452

1453def _build_saved_id_to_fqn(

1454 optimizer: torch.optim.Optimizer,

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

1456) -> Dict[int, str]:

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

1458 raw_sd = optimizer.state_dict()

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

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

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

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

1463 if p is parameter:

1464 saved_id_to_fqn[saved_id] = name

1465 break

1466 return saved_id_to_fqn

1467

1468

1469def _parse_nested_state_entry(

1470 meta_key: str,

1471 meta_val: Any,

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

1473 full_state_dict: bool,

1474 cpu_offload: bool,

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

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

1477 from hyper_parallel.core.distributed_checkpoint.metadata import (

1478 TensorStorageMetadata,

1479 )

1480

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

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

1483 matched_fqn = _match_fqn_from_remainder(remainder, fqns_sorted)

1484 if matched_fqn is None:

1485 return None

1486

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

1488 if state_key is None:

1489 return None

1490

1491 if not isinstance(meta_val, TensorStorageMetadata):

1492 return None

1493

1494 dtype = _resolve_dtype(meta_val.properties.dtype)

1495 global_shape = meta_val.size

1496 param = param_by_fqn.get(matched_fqn)

1497

1498 if _is_scalar_state(state_key):

1499 device = torch.device("cpu")

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

1501 else:

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

1503

1504 return matched_fqn, state_key, t

1505

1506

1507def _parse_nested_param_group_entry(

1508 meta_key: str,

1509 meta_val: Any,

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

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

1512 from hyper_parallel.core.distributed_checkpoint.metadata import (

1513 TensorStorageMetadata,

1514 )

1515

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

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

1518 if len(parts) < 2:

1519 return None

1520 try:

1521 group_idx = int(parts[0])

1522 except ValueError:

1523 return None

1524 field_name = parts[1]

1525

1526 if not isinstance(meta_val, TensorStorageMetadata):

1527 return None

1528

1529 dtype = _resolve_dtype(meta_val.properties.dtype)

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

1531

1532

1533def _build_nested_template_from_metadata(

1534 metadata: Any,

1535 optimizer: torch.optim.Optimizer,

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

1537 full_state_dict: bool,

1538 cpu_offload: bool,

1539) -> Dict[str, Any]:

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

1541

1542 Non-flatten checkpoint keys look like:

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

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

1545

1546 We reconstruct:

1547 {

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

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

1550 }

1551 """

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

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

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

1555

1556 saved_id_to_fqn = _build_saved_id_to_fqn(optimizer, param_by_fqn)

1557 raw_sd = optimizer.state_dict()

1558

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

1560 param_groups_raw[group_idx] = {

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

1562 }

1563 param_groups_fqns[group_idx] = [

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

1565 ]

1566

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

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

1569 result = _parse_nested_state_entry(

1570 meta_key, meta_val, param_by_fqn, full_state_dict, cpu_offload,

1571 )

1572 if result is not None:

1573 fqn, state_key, t = result

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

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

1576 result = _parse_nested_param_group_entry(meta_key, meta_val)

1577 if result is not None:

1578 group_idx, field_name, t = result

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

1580

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

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

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

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

1585 result_param_groups.append(pg)

1586

1587 return {

1588 "state": state,

1589 "param_groups": result_param_groups,

1590 }

1591

1592

1593def _build_flatten_template_from_metadata(

1594 metadata: Any,

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

1596 full_state_dict: bool,

1597 cpu_offload: bool,

1598) -> Dict[str, Any]:

1599 """Build flatten template from checkpoint metadata.

1600

1601 Flatten checkpoint keys look like:

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

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

1604 """

1605 from hyper_parallel.core.distributed_checkpoint.metadata import (

1606 TensorStorageMetadata,

1607 )

1608

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

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

1611

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

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

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

1615 matched_fqn = _match_fqn_from_remainder(remainder, fqns_sorted)

1616 if matched_fqn is None:

1617 continue

1618

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

1620 if state_key is None:

1621 continue

1622

1623 if isinstance(meta_val, TensorStorageMetadata):

1624 dtype = _resolve_dtype(meta_val.properties.dtype)

1625 global_shape = meta_val.size

1626 param = param_by_fqn.get(matched_fqn)

1627

1628 if _is_scalar_state(state_key):

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

1630 else:

1631 t = _create_empty_state_tensor(

1632 param, global_shape, dtype,

1633 full_state_dict, cpu_offload,

1634 )

1635 flat[meta_key] = t

1636

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

1638 if isinstance(meta_val, TensorStorageMetadata):

1639 dtype = _resolve_dtype(meta_val.properties.dtype)

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

1641

1642 return flat

1643

1644

1645def _create_empty_state_tensor(

1646 param: Optional[nn.Parameter],

1647 global_shape: tuple,

1648 dtype: torch.dtype,

1649 full_state_dict: bool,

1650 cpu_offload: bool,

1651) -> torch.Tensor:

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

1653

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

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

1656

1657 Args:

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

1659 global_shape: Global shape from checkpoint metadata.

1660 dtype: Tensor dtype.

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

1662 cpu_offload: If True, place on CPU.

1663

1664 Returns:

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

1666 """

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

1668

1669 if dtensor_info is not None:

1670 if full_state_dict:

1671 shape = global_shape

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

1673 else:

1674 shape = param.to_local().shape

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

1676 else:

1677 shape = global_shape

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

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

1680 )

1681

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

1683

1684

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

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

1687

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

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

1690 """

1691 opt_cls = type(optimizer)

1692 opt_name = opt_cls.__name__.lower()

1693

1694 if "adam" in opt_name:

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

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

1697 keys.append("max_exp_avg_sq")

1698 return keys

1699

1700 if "sgd" in opt_name:

1701 keys = []

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

1703 keys.append("momentum_buffer")

1704 return keys

1705

1706 logger.warning(

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

1708 "DCP load template may be incomplete.",

1709 opt_cls.__name__,

1710 )

1711 return []

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

682 statements