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