Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / fully_shard / state_dict_utils.py: 99%
79 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""State dict utilities for fully_shard (torch-specific)."""
16from typing import Any, Optional
18import torch
19import torch.distributed as dist
20from torch import nn
21from torch.distributed.checkpoint.state_dict import StateDictOptions
23from hyper_parallel.core.dtensor.dtensor import DTensor, distribute_tensor
26def _gather_full_state_dict(
27 state_dict: dict[str, Any], cpu_offload: bool
28) -> dict[str, Any]:
29 """All-gather every DTensor shard into a full tensor.
31 Args:
32 state_dict: Model state dict with DTensor or plain tensor values.
33 cpu_offload: If True, only rank-0 keeps the result on CPU;
34 other ranks return an empty dict to save memory.
35 """
36 is_rank0 = (not dist.is_initialized()) or (dist.get_rank() == 0)
38 gathered: dict[str, Any] = {}
39 for key, val in state_dict.items():
40 if isinstance(val, DTensor):
41 val = val.full_tensor()
42 if cpu_offload:
43 if not is_rank0:
44 del val
45 continue
46 if isinstance(val, torch.Tensor):
47 val = val.cpu()
48 gathered[key] = val
50 if cpu_offload and not is_rank0:
51 return {}
52 return gathered
55def _offload_sharded_state_dict(
56 state_dict: dict[str, Any],
57) -> dict[str, Any]:
58 """Move each shard to CPU without all-gathering.
60 Args:
61 state_dict: Model state dict with DTensor or plain tensor values.
62 """
63 offloaded: dict[str, Any] = {}
64 for key, val in state_dict.items():
65 if isinstance(val, DTensor):
66 val = DTensor.from_local(
67 val.to_local().cpu(), val.device_mesh, val.layout.alias_placements,
68 )
69 elif isinstance(val, torch.Tensor):
70 val = val.cpu()
71 offloaded[key] = val
72 return offloaded
75def get_model_state_dict(
76 model: nn.Module,
77 *,
78 options: Optional[StateDictOptions] = None,
79) -> dict[str, Any]:
80 """Return the model state dict with configurable gathering and offloading.
82 Behaviour matrix:
84 +-----------------+-------------+--------------------------------------+
85 | full_state_dict | cpu_offload | result |
86 +=================+=============+======================================+
87 | False | False | DTensor (sharded, as-is) |
88 +-----------------+-------------+--------------------------------------+
89 | False | True | DTensor local shard offloaded to CPU |
90 +-----------------+-------------+--------------------------------------+
91 | True | False | full Tensor on **every** rank |
92 +-----------------+-------------+--------------------------------------+
93 | True | True | full Tensor on CPU, **rank 0 only** |
94 +-----------------+-------------+--------------------------------------+
96 Args:
97 model: The model whose state dict to retrieve.
98 options: Controls full_state_dict, cpu_offload,
99 ignore_frozen_params, and broadcast_from_rank0 flags.
101 Raises:
102 ValueError: If ``broadcast_from_rank0`` is True while
103 ``full_state_dict`` is False.
104 NotImplementedError: If ``broadcast_from_rank0`` is True.
105 ``broadcast_from_rank0`` requires a cross-rank tensor broadcast,
106 which is not available. Use ``full_state_dict=True`` with a full
107 tensor on every rank instead.
108 """
109 options = options or StateDictOptions()
111 if options.broadcast_from_rank0 and not options.full_state_dict:
112 raise ValueError(
113 "full_state_dict must be True when broadcast_from_rank0 is True."
114 )
115 if options.broadcast_from_rank0:
116 raise NotImplementedError(
117 "broadcast_from_rank0=True is not supported. "
118 "broadcast_from_rank0 requires a cross-rank tensor broadcast, "
119 "which is not available. Use full_state_dict=True with a full "
120 "tensor on every rank instead."
121 )
123 state_dict: dict[str, Any] = model.state_dict()
125 if options.ignore_frozen_params:
126 frozen_keys = {
127 name for name, p in model.named_parameters()
128 if not p.requires_grad
129 }
130 for key in frozen_keys:
131 state_dict.pop(key, None)
133 if options.full_state_dict:
134 return _gather_full_state_dict(state_dict, options.cpu_offload)
136 if options.cpu_offload:
137 return _offload_sharded_state_dict(state_dict)
139 return state_dict
142def _scatter_model_state_dict(
143 model: nn.Module,
144 state_dict: dict[str, Any],
145 cpu_offload: bool,
146 strict: bool,
147) -> dict[str, Any]:
148 """Scatter full tensors into DTensor shards matching model's layout.
150 Inverse of _gather_full_state_dict. Uses distribute_tensor (no communication,
151 assumes every rank holds the same global tensor).
153 When ``cpu_offload`` is True the input tensors are expected to live on CPU
154 (e.g. produced by ``get_model_state_dict(full_state_dict=True, cpu_offload=True)``).
155 The scattered local shard is moved onto the target parameter's device so the
156 subsequent ``load_state_dict(assign=True)`` does not leave the parameter on
157 the wrong device.
159 Args:
160 model: The model whose layout the scattered tensors must match.
161 state_dict: Input state dict whose values are plain (global) tensors.
162 cpu_offload: Whether the input tensors live on CPU and must be moved
163 back onto the target parameter's device after scattering.
164 strict: When True, keys present in ``state_dict`` but absent from the
165 model are treated as errors (unexpected keys), mirroring the
166 PyTorch ``load_state_dict(strict=True)`` semantics of the
167 ``full_state_dict=False`` passthrough path.
169 Raises:
170 ValueError: If ``strict`` is True and ``state_dict`` contains keys that
171 do not exist in the model (unexpected keys).
172 """
173 target_state_dict = model.state_dict()
174 scattered: dict[str, Any] = {}
175 unexpected_keys: list[str] = []
176 for key, val in state_dict.items():
177 target = target_state_dict.get(key)
178 if target is None:
179 unexpected_keys.append(key)
180 continue
181 if isinstance(target, DTensor):
182 if isinstance(val, DTensor):
183 scattered[key] = val
184 else:
185 # Slice a plain (global) tensor into a DTensor shard.
186 placements = (
187 target.layout.alias_placements if target.layout else target.placements
188 )
189 scattered[key] = distribute_tensor(val, target.device_mesh, placements)
190 else:
191 scattered[key] = val
193 # When the input came from CPU (cpu_offload=True round-trip), move the
194 # scattered local shard onto the target parameter's device. assign=True
195 # would otherwise bind a CPU DTensor to an on-device parameter.
196 if cpu_offload and isinstance(scattered[key], DTensor) and isinstance(target, DTensor):
197 dt = scattered[key]
198 target_device = target._local_tensor.device # pylint: disable=protected-access
199 scattered[key] = DTensor.from_local(
200 dt.to_local().to(target_device),
201 dt.device_mesh,
202 dt.layout.alias_placements if dt.layout else dt.placements,
203 )
205 if strict and unexpected_keys:
206 raise ValueError(
207 f"Unexpected key(s) in state_dict: {unexpected_keys}. "
208 f"To allow loading a state_dict with extra keys, pass strict=False."
209 )
210 return scattered
213def set_model_state_dict(
214 model: nn.Module,
215 state_dict: dict[str, Any],
216 *,
217 options: Optional[StateDictOptions] = None,
218) -> None:
219 """Load state dict into model, scattering full tensors to DTensor shards.
221 Behaviour matrix:
223 +-----------------+-------------+--------------------------------+
224 | full_state_dict | cpu_offload | behaviour |
225 +=================+=============+================================+
226 | True | False | scatter full -> DTensor on dev |
227 +-----------------+-------------+--------------------------------+
228 | True | True | scatter full -> DTensor on dev |
229 +-----------------+-------------+--------------------------------+
230 | False (any) | * | load sharded DTensors as-is |
231 +-----------------+-------------+--------------------------------+
233 Note:
234 ``ignore_frozen_params`` is intentionally a no-op on the setter path,
235 mirroring torch.distributed.checkpoint.state_dict: frozen parameters
236 filtering is a getter-only feature. Callers that strip frozen keys
237 from the input should also pass ``strict=False``.
239 Args:
240 model: The model to load state into.
241 state_dict: State dict to load. Values may be plain (global) tensors
242 when ``full_state_dict=True`` or sharded DTensors otherwise.
243 options: Controls full_state_dict, cpu_offload, strict and
244 broadcast_from_rank0 flags.
246 Raises:
247 ValueError: If ``broadcast_from_rank0`` is True while
248 ``full_state_dict`` is False.
249 NotImplementedError: If ``broadcast_from_rank0`` is True.
250 ``broadcast_from_rank0`` requires a cross-rank tensor broadcast,
251 which is not available. Use ``full_state_dict=True`` with a full
252 tensor on every rank instead.
253 """
254 options = options or StateDictOptions()
256 if options.broadcast_from_rank0 and not options.full_state_dict:
257 raise ValueError(
258 "full_state_dict must be True when broadcast_from_rank0 is True."
259 )
260 if options.broadcast_from_rank0:
261 raise NotImplementedError(
262 "broadcast_from_rank0=True is not supported. "
263 "broadcast_from_rank0 requires a cross-rank tensor broadcast, "
264 "which is not available. Use full_state_dict=True with a full "
265 "tensor on every rank instead."
266 )
268 # Scatter full tensors into DTensor shards matching the model's layout.
269 if options.full_state_dict:
270 scattered = _scatter_model_state_dict(
271 model, state_dict, options.cpu_offload, options.strict,
272 )
273 else:
274 scattered = state_dict
276 model.load_state_dict(scattered, strict=options.strict, assign=True)