Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / source_shard.py: 67%
125 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-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# pylint: disable=missing-apache-license-header
16"""source_shard: build_source_shard_info (05 §6.7.1).
18source_shard_info is read from the ShardingPlan (rather than from DTensors — under
19production the parameters have already been unwrapped by
20_local_params_context, and only the plan retains the complete placement
21information).
23Each entry records the parameter's **complete source layout**: the placements
24on every source (non-FSDP) mesh axis, together with the source sub-mesh.
25``fully_shard`` later prefixes its own FSDP shard/replicate dimensions, so the
26final sharded parameter carries the full distributed layout (FSDP dims +
27source dims) for downstream consumers such as distributed checkpointing.
28"""
29# pylint: disable=forbidden-backend-import
31from __future__ import annotations
33from collections.abc import Mapping
34from typing import Any, Dict, List, Optional, Tuple, TypeAlias
36from torch.nn import Module as ModuleClass
37from torch.nn import Parameter as ParameterClass
38from torch.utils.weak import WeakIdKeyDictionary
40import hyper_parallel.core.fully_shard.utils as fully_shard_utils
41from hyper_parallel import DeviceMesh, DTensor, Replicate
42from hyper_parallel.core.dtensor.placement_types import Partial, Placement, Shard
43from hyper_parallel.core.fully_shard.hsdp_utils import get_managed_modules_parameters
44from hyper_parallel.distributed.plan import ShardingPlan
45from hyper_parallel.distributed.recipe_spec import EP, resolve_placements
47# Mesh dimensions whose sharding semantics belong to FSDP (or to other
48# parallelism concerns), never to the source layout recorded here:
49# - dp/dp_shard/dp_replicate/cp: FSDP shards weights over the whole DP+CP
50# domain (cp is folded into the fsdp shard axis at mesh build time), and
51# FSDP's reduce-scatter already covers gradient sync on these axes;
52# - pp: different stages hold different parameters (distinguished by FQN),
53# a weight is never distributed along pp;
54# - fsdp_*/edp_*: the FSDP child-mesh axes themselves.
55# Declaring a Shard placement for a weight on one of these axes would
56# double-count the sharding (Phase A shard + FSDP shard on the same ranks),
57# so build_source_shard_info rejects it fail-fast.
58FSDP_OWNED_DIMS = frozenset({
59 "dp", "dp_replicate", "dp_shard", "cp", "pp",
60 "fsdp_replicate", "fsdp_shard",
61 "edp", "edp_replicate", "edp_shard",
62})
65def _source_dim_names(mesh, fallback_dim_names):
66 """Return the non-FSDP axis names of a source mesh.
68 Falls back to the plan's axis names when the mesh carries no dimension
69 names (e.g. a mock or a plain 1-D TP mesh).
70 """
71 dim_names = getattr(mesh, "mesh_dim_names", None) if mesh is not None else None
72 if dim_names is None:
73 dim_names = tuple(fallback_dim_names or ())
74 return tuple(name for name in dim_names if name not in FSDP_OWNED_DIMS)
77def _source_sub_mesh(mesh, source_dim_names):
78 """Slice the source sub-mesh covering exactly ``source_dim_names``."""
79 if mesh is None:
80 return None
81 dim_names = getattr(mesh, "mesh_dim_names", None)
82 if dim_names is None or tuple(dim_names) == tuple(source_dim_names):
83 return mesh
84 if not source_dim_names:
85 return None
86 if len(source_dim_names) == 1:
87 return mesh[source_dim_names[0]]
88 return mesh[source_dim_names]
91def _check_fsdp_owned_axes(full_fqn, named_placement):
92 """Reject weight sharding declared on axes owned by FSDP/PP."""
93 for axis, placement in named_placement.items():
94 if axis in FSDP_OWNED_DIMS and not placement.is_replicate():
95 raise ValueError(
96 f"{full_fqn}: placement {placement} on axis {axis!r} conflicts with "
97 "FSDP ownership of that axis (weights are sharded over the DP+CP "
98 "domain by fully_shard itself); declare Replicate or remove the axis"
99 )
102def _normalize_tied_pair_placements(info, tied_pairs):
103 """Normalize tied parameters to the finer source placement."""
104 if not tied_pairs:
105 return
106 for a, b in tied_pairs:
107 if a in info and b in info:
108 pa, _ = info[a]
109 pb, _ = info[b]
110 if pa != pb and len(pa) == len(pb):
111 norm = tuple(
112 x if isinstance(x, Shard) else y
113 for x, y in zip(pa, pb)
114 )
115 info[a] = (norm, info[a][1])
116 info[b] = (norm, info[b][1])
119def build_source_shard_info(
120 plan: ShardingPlan,
121 dense_source_mesh: Any,
122 *,
123 expert_source_mesh: Any = None,
124 tied_pairs: Optional[List[Tuple[str, str]]] = None,
125) -> Dict[str, Tuple[Tuple[Placement, ...], Any]]:
126 """Build complete source-layout metadata for dense and routed-expert parameters.
128 Args:
129 plan: The finalized ShardingPlan.
130 dense_source_mesh: Dense-region source mesh. May be the full dense
131 FSDP mesh (e.g. ``(fsdp_replicate, fsdp_shard, tp)``) — FSDP-owned
132 axes are stripped; or a plain TP child mesh (used as-is).
133 expert_source_mesh: Full expert mesh (e.g. ``(edp_shard, ep)``) or the
134 EP child mesh; required when any spec marks routed experts.
135 tied_pairs: parameter pairs with shared storage (defaults to
136 plan.tied_pairs). Both ends of a tied pair must map to the same
137 placements — axis by axis, the finer sharding wins (Shard takes
138 precedence over Replicate), guaranteeing consistent TP all-reduce /
139 reduce-scatter semantics on both ends.
141 Returns:
142 ``{param_fqn: (placements_tuple, source_sub_mesh)}`` where
143 ``placements_tuple`` holds one placement per source (non-FSDP) mesh
144 axis, missing axes filled with ``Replicate()``.
145 """
146 dense_source_dims = _source_dim_names(dense_source_mesh, plan.mesh_dim_names)
147 dense_mesh = _source_sub_mesh(dense_source_mesh, dense_source_dims)
148 expert_source_dims = _source_dim_names(expert_source_mesh, ("ep",))
149 expert_mesh = _source_sub_mesh(expert_source_mesh, expert_source_dims)
151 info = {}
152 for fqn, spec in plan.modules.items():
153 for param_name, named_placement in spec.params.items():
154 full_fqn = f"{fqn}.{param_name}"
155 _check_fsdp_owned_axes(full_fqn, named_placement)
156 if (spec._ep_size > 0 # pylint: disable=protected-access
157 and isinstance((named_placement or {}).get(EP), Shard)):
158 if expert_source_mesh is None:
159 raise ValueError(
160 "Virtual-EP parameter metadata requires an EP source mesh"
161 )
162 placements = tuple(resolve_placements(named_placement, expert_source_dims))
163 info[full_fqn] = (placements, expert_mesh)
164 else:
165 placements = tuple(resolve_placements(named_placement, dense_source_dims))
166 info[full_fqn] = (placements, dense_mesh)
168 _normalize_tied_pair_placements(
169 info,
170 tied_pairs if tied_pairs is not None else plan.tied_pairs,
171 )
172 return info
175# ────────────────────────────────────────────────────────────────────────────
176# FSDP2 source-layout discovery (moved verbatim from components/distributed/
177# fsdp2.py in stage 4f, 05 §15.2.4 row 406; ``self`` became an explicit
178# ``manager`` parameter — FSDP2Manager calls these free functions).
179# ────────────────────────────────────────────────────────────────────────────
181SourceShardInfoByFQN: TypeAlias = Mapping[ # pylint: disable=invalid-name
182 str, tuple[tuple[Placement, ...], DeviceMesh]
183]
186SourceShardInfoByParam: TypeAlias = dict[ # pylint: disable=invalid-name
187 ParameterClass,
188 "fully_shard_utils.SourceShardMetaInfo",
189]
192def _build_dtensor_source_shard_info(
193 manager,
194 model: ModuleClass,
195) -> SourceShardInfoByParam | None:
196 """Build source metadata from DTensor parameters when TP is active."""
197 if manager.mesh_context.tp_size <= 1:
198 return None
200 parameters = list(model.parameters())
201 if not parameters or not all(
202 isinstance(parameter, DTensor) for parameter in parameters
203 ):
204 raise ValueError(
205 "TP is enabled but source_shard_info is missing and "
206 "not all model parameters are DTensors"
207 )
208 return {
209 parameter: fully_shard_utils.SourceShardMetaInfo( # pylint: disable=no-member
210 mesh=parameter.device_mesh,
211 placements=tuple(parameter.placements),
212 origin_is_dtensor=True,
213 )
214 for parameter in parameters
215 }
218def _build_parameter_source_shard_info(
219 parameter_fqn: str,
220 placements: tuple[Placement, ...],
221 source_mesh: DeviceMesh,
222) -> fully_shard_utils.SourceShardMetaInfo:
223 """Validate and build source metadata for one model parameter."""
224 if any(isinstance(placement, Partial) for placement in placements):
225 raise ValueError(
226 "source_shard_info does not support Partial placement: "
227 f"{parameter_fqn}"
228 )
229 return fully_shard_utils.SourceShardMetaInfo( # pylint: disable=no-member
230 mesh=source_mesh,
231 placements=tuple(placements),
232 origin_is_dtensor=False,
233 )
236def _record_parameter_source_shard_info(
237 metadata_by_parameter: SourceShardInfoByParam,
238 parameter: ParameterClass,
239 parameter_fqn: str,
240 parameter_source_shard_info: fully_shard_utils.SourceShardMetaInfo,
241) -> None:
242 """Record one layout while rejecting conflicting tied aliases."""
243 previous_source_shard_info = metadata_by_parameter.get(parameter)
244 if previous_source_shard_info is not None and (
245 previous_source_shard_info.mesh is not parameter_source_shard_info.mesh
246 or previous_source_shard_info.placements
247 != parameter_source_shard_info.placements
248 ):
249 raise ValueError(
250 "Tied parameter aliases have conflicting source layouts; "
251 f"conflict found at {parameter_fqn}"
252 )
253 metadata_by_parameter[parameter] = parameter_source_shard_info
256def _build_source_shard_info_by_param(
257 manager,
258 model: ModuleClass,
259 source_shard_info: SourceShardInfoByFQN | None,
260) -> SourceShardInfoByParam | None:
261 """Resolve stable FQNs to parameter identities after model rewrites.
263 Activation-checkpoint wrappers may add an internal child-module prefix,
264 but the supported wrapper contract removes that prefix from
265 ``named_modules``/``named_parameters``. Therefore planner metadata
266 produced before checkpointing remains resolvable here, while FSDP still
267 receives the final parameter objects that it will manage.
269 Args:
270 model: Model after sharding and checkpoint rewrites, before layer
271 compile is installed.
272 source_shard_info: Source-layout metadata produced by the TP planner.
274 Returns:
275 Parameter-keyed metadata for later ``fully_shard`` calls, or
276 ``None`` when tensor parallelism is disabled.
278 Raises:
279 ValueError: If TP metadata is missing, references an unknown FQN,
280 contains ``Partial``, or gives tied aliases conflicting layouts.
281 """
282 if source_shard_info is None:
283 return _build_dtensor_source_shard_info(manager, model)
284 if manager.mesh_context.tp_size > 1 and not source_shard_info:
285 raise ValueError("TP is enabled but source_shard_info is empty")
287 parameters_by_fqn = dict(model.named_parameters(remove_duplicate=False))
288 metadata_by_parameter: SourceShardInfoByParam = {}
289 for parameter_fqn, (placements, source_mesh) in source_shard_info.items():
290 parameter = parameters_by_fqn.get(parameter_fqn)
291 if parameter is None:
292 raise ValueError(
293 f"source_shard_info contains unknown parameter FQN: {parameter_fqn}"
294 )
295 parameter_source_shard_info = _build_parameter_source_shard_info(
296 parameter_fqn,
297 placements,
298 source_mesh,
299 )
300 _record_parameter_source_shard_info(
301 metadata_by_parameter,
302 parameter,
303 parameter_fqn,
304 parameter_source_shard_info,
305 )
306 return metadata_by_parameter
309def _get_default_source_shard_info(manager) -> fully_shard_utils.SourceShardMetaInfo:
310 """Build all-Replicate source metadata covering every non-FSDP mesh axis.
312 Used to complete source_shard_info for parameters the plan does not mention;
313 the axis set must match the per-parameter entries so that all parameters
314 in one ``fully_shard`` unit share the same source-mesh dimensionality.
315 """
316 world_mesh = manager.mesh_context.fsdp_non_moe_mesh or manager.mesh_context.device_mesh
317 if world_mesh is None or world_mesh.mesh_dim_names is None:
318 raise ValueError("TP metadata completion requires a named world mesh")
319 source_dim_names = tuple(
320 dim_name
321 for dim_name in world_mesh.mesh_dim_names
322 if dim_name not in FSDP_OWNED_DIMS
323 )
324 if not source_dim_names:
325 raise ValueError(
326 "TP is enabled but the world mesh has no non-FSDP source dimension"
327 )
328 source_mesh = (
329 world_mesh
330 if tuple(world_mesh.mesh_dim_names) == source_dim_names
331 else world_mesh[source_dim_names[0]]
332 if len(source_dim_names) == 1
333 else world_mesh[source_dim_names]
334 )
335 return fully_shard_utils.SourceShardMetaInfo( # pylint: disable=no-member
336 mesh=source_mesh,
337 placements=tuple(Replicate() for _ in source_dim_names),
338 origin_is_dtensor=False,
339 )
342def _build_managed_source_shard_info(
343 manager,
344 owner: ModuleClass,
345 metadata_by_parameter: SourceShardInfoByParam | None,
346 ignored_params: set[ParameterClass] | None = None,
347) -> SourceShardInfoByParam | None:
348 """Select and complete metadata for parameters managed by one wrap call.
350 Metadata must use the same parameter selection as ``fully_shard``. In
351 particular, nested FSDP units mark their parameters initialized before the
352 parent is wrapped, so ownership assigned from the pre-wrap module tree can
353 over-approximate the actual parameter set. Deriving this set through the
354 shared helper keeps source-layout metadata complete and exact.
355 """
356 if metadata_by_parameter is None:
357 return None
359 default_source_shard_info = None
360 if (
361 manager.mesh_context.fsdp_non_moe_mesh is not None
362 or manager.mesh_context.device_mesh is not None
363 ):
364 default_source_shard_info = _get_default_source_shard_info(manager)
365 managed_source_shard_info = {}
366 for parameter in get_managed_modules_parameters((owner,), ignored_params):
367 source_shard_info = metadata_by_parameter.get(parameter)
368 if source_shard_info is None:
369 if default_source_shard_info is None:
370 raise ValueError(
371 "source_shard_info is present but does not cover an FSDP-managed parameter"
372 )
373 source_shard_info = default_source_shard_info
374 managed_source_shard_info[parameter] = source_shard_info
375 return managed_source_shard_info
378def _source_infos_for_fully_shard(
379 managed_source_shard_info: SourceShardInfoByParam | None,
380) -> Mapping[ParameterClass, fully_shard_utils.SourceShardMetaInfo] | None:
381 """Return non-owning metadata only for plain source-layout parameters.
383 Validate-mode parameters remain native DTensors. The FSDP
384 state derives their source layouts from each parameter directly and
385 rejects duplicate explicit metadata.
387 The adapter keeps its construction-local mapping for post-wrap gradient
388 configuration. FSDP receives weak identity keys so its scheduler and state
389 cannot keep the replaced original parameters and their storage alive.
390 """
391 if managed_source_shard_info is None:
392 return None
393 if managed_source_shard_info and all(
394 source_shard_info.origin_is_dtensor
395 for source_shard_info in managed_source_shard_info.values()
396 ):
397 return None
398 return WeakIdKeyDictionary(managed_source_shard_info)