Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / parameter_sharding.py: 80%
174 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"""parameter_sharding: parameter-side apply machinery (05 §4.2 + Phase D).
17Owns parameter placement/localize/stack (merged from the legacy
18``components/distributed/sharding/apply.py`` — there is exactly one parameter
19sharding implementation), source-mesh resolution, runtime source_shard_info
20construction, and the tied-weights replication pass (same-rank shared
21storage, no cross-rank shard broadcast).
22"""
24import logging
25from contextlib import contextmanager
26from typing import Any, Dict, List, Tuple
28import torch
29from torch import nn
31from hyper_parallel.components.modules.module_family import (
32 _family_source_shard_info, get_family_build_info, _validate_applied_layout,
33)
34from hyper_parallel.core.dtensor.dtensor import DTensor, distribute_tensor
35from hyper_parallel.core.dtensor.placement_types import Shard
36from hyper_parallel.distributed.recipe_spec import (
37 EP,
38 PlacementMismatchError,
39 resolve_placements,
40)
41from hyper_parallel.distributed._builder.source_shard import (
42 build_source_shard_info,
43)
44from hyper_parallel.distributed.tensor_parallel.head_count import (
45 maybe_update_head_counts,
46)
48logger = logging.getLogger(__name__)
51# ────────────────────────────────────────────────────────────────────────────
52# Legacy path helpers (merged from components/distributed/sharding/apply.py)
53# ────────────────────────────────────────────────────────────────────────────
55def _get_attr_by_path(model, fqn):
56 """Fetch an attribute along a dotted FQN (numeric segments index into ModuleLists)."""
57 obj = model
58 for p in fqn.split("."):
59 obj = obj[int(p)] if p.isdigit() else getattr(obj, p)
60 return obj
63def _set_param_by_path(model: nn.Module, fqn: str, new_param) -> None:
64 """Locate the parent module along a dotted FQN and replace the leaf parameter.
66 object.__setattr__(model, dotted_name, ...) would only set a stray
67 attribute on model and would not replace the submodule parameter — you
68 must walk the path to the true parent module before assigning.
69 """
70 *path, leaf = fqn.split(".")
71 obj = model
72 for p in path:
73 obj = obj[int(p)] if p.isdigit() else getattr(obj, p)
74 if hasattr(obj, "register_parameter"):
75 obj.register_parameter(leaf, new_param)
76 else:
77 object.__setattr__(obj, leaf, new_param)
80def _resolve_module(model, fqn):
81 """Fetch a module by FQN (the last segment is NOT stripped; call sites pass module FQNs).
83 Same semantics as _get_attr_by_path — every call site (Phase A/B/C)
84 passes a module fully-qualified name (e.g. `model.layers.0.self_attn`),
85 not a parameter FQN, so no last-segment stripping is done (stripping
86 would incorrectly return the parent module). The empty FQN "" resolves
87 to the model itself (D-14: a root-level outer spec, e.g. a whole-LM
88 contract, 05 §13.4).
89 """
90 obj = model
91 if not fqn:
92 return obj
93 for p in fqn.split("."):
94 obj = obj[int(p)] if p.isdigit() else getattr(obj, p)
95 return obj
98def _local_params_context(model: nn.Module):
99 """One-shot unwrap at build time: replace DTensor parameters with their
100 _local_tensor (plain), zero-copy.
102 Called at the Phase C entry of apply_sharding_plan, before fully_shard;
103 the permanent unwrap is not restored. _local_tensor shares storage with
104 the original DTensor (same data_ptr).
106 Returns a {fqn: placements} snapshot of the placements before unwrapping
107 (diagnostic use only; the canonical source for source_shard_info is the
108 ShardingPlan, see build_source_shard_info).
109 """
110 source_shard_records = {}
111 for name, param in list(model.named_parameters()):
112 if isinstance(param, DTensor):
113 source_shard_records[name] = param.placements
114 _set_param_by_path(model, name, nn.Parameter(
115 param.to_local(), requires_grad=param.requires_grad))
116 return source_shard_records
119@contextmanager
120def _temp_local_params(module, exclude=()):
121 """Temporarily unwrap DTensor parameters inside a validate-mode local region (restored on exit).
123 Under production the parameters were already permanently unwrapped at build
124 time, so this context is unnecessary. A validate-mode local region (MoE
125 all-to-all / HF CP attention) computes on local tensors internally and
126 needs local parameters; after restoration the DTensor propagation chain is
127 unbroken.
129 exclude: relative FQN prefixes of nested-boundary subtrees whose
130 parameters must stay DTensors (D-14 invariant 3, 05 §13.3 — inner
131 validate islands dispatch via __torch_function__ and break if the outer
132 region unwraps their parameters).
133 """
134 excluded = tuple(e.rstrip(".") + "." for e in exclude)
135 saved = []
136 for name, param in list(module.named_parameters(recurse=True)):
137 if excluded and name.startswith(excluded):
138 continue
139 if isinstance(param, DTensor):
140 saved.append((name, param))
141 *path, leaf = name.split(".")
142 owner = module
143 for part in path:
144 owner = owner[int(part)] if part.isdigit() else getattr(owner, part)
145 owner._parameters[leaf] = param.to_local() # pylint: disable=protected-access
146 try:
147 yield
148 finally:
149 for name, param in saved:
150 _set_param_by_path(module, name, param)
153class _StackedExperts(nn.Module):
154 """Container for stacked per-expert weights: gate_proj/up_proj/down_proj
155 (or w1/w2/w3) become Parameters of shape [E, ...], sharded uniformly by
156 EP Shard(0) + TP (the D-08 ndim=3 rule)."""
159def _install_stacked_expert_holders(module: nn.Module, holders: Dict[str, Dict[str, nn.Parameter]]) -> None:
160 """Replace per-expert containers with their stacked parameter holders."""
161 for parent_path, params in holders.items():
162 holder = _StackedExperts()
163 for name, p in params.items():
164 holder.register_parameter(name, p)
165 *path, leaf = parent_path.split(".")
166 obj = module
167 for seg in path:
168 obj = obj[int(seg)] if seg.isdigit() else getattr(obj, seg)
169 setattr(obj, leaf, holder) # replace the original ModuleList (original expert params freed)
172def _stack_moe_experts(module: nn.Module, ep_stack: Dict[str, List[str]]) -> None:
173 """Per-expert parameters → stacked 3D parameters (stack is concat, values
174 exactly equal).
176 Executed before _shard_module_params in Phase A:
177 - fetch weights by source path → torch.stack(dim=0), register onto the
178 replaced experts holder;
179 - the original ModuleList is replaced as a whole (memory freed);
180 - v1 asserts no bias and consistent source-parameter shapes; meta tensors
181 work the same way (concat of metas).
183 ep_stack: {stacked relative path: [source parameter relative paths
184 (ordered by expert idx)]}, e.g.
185 {"experts.gate_proj": ["experts.0.gate_proj.weight", ...]}.
186 """
187 holders: Dict[str, Dict[str, nn.Parameter]] = {}
188 for stacked_path, sources in ep_stack.items():
189 parent_path, param_name = stacked_path.rsplit(".", 1)
190 tensors = []
191 requires_grad = True
192 for src in sources:
193 owner = _resolve_module(module, src.rsplit(".", 1)[0])
194 if getattr(owner, "bias", None) is not None:
195 raise NotImplementedError(
196 f"D-09 v1 does not support experts with bias "
197 f"({src.rsplit('.', 1)[0]}); use an EP-aware MoE module instead"
198 )
199 t = _get_attr_by_path(module, src)
200 tensors.append(t.data if hasattr(t, "data") else t)
201 requires_grad = getattr(t, "requires_grad", True)
202 stacked = torch.stack(tensors, dim=0)
203 holders.setdefault(parent_path, {})[param_name] = nn.Parameter(
204 stacked, requires_grad=requires_grad)
206 _install_stacked_expert_holders(module, holders)
209# ────────────────────────────────────────────────────────────────────────────
210# Source-mesh resolution + Phase A parameter sharding (05 §4.2)
211# ────────────────────────────────────────────────────────────────────────────
213def _resolve_parameter_source_meshes(plan, mesh_context, full_mesh, tp_mesh):
214 """Resolve dense TP and virtual-EP source meshes for one sharding plan."""
215 # Lazy import: applier imports parameter_sharding (_resolve_module) at
216 # module level — importing applier here at module level would cycle.
217 from hyper_parallel.distributed._builder.applier import ( # pylint: disable=C0415,R0401
218 _build_expert_mesh,
219 )
220 ep_size = next(
221 (
222 spec._ep_size # pylint: disable=protected-access
223 for spec in plan.modules.values()
224 if spec._ep_size > 0 # pylint: disable=protected-access
225 ),
226 0,
227 )
228 if mesh_context is not None:
229 expert_mesh = getattr(mesh_context, "fsdp_moe_mesh", None)
230 elif ep_size > 0:
231 expert_mesh = _build_expert_mesh(
232 full_mesh,
233 full_mesh.mesh_dim_names,
234 ep_size,
235 )
236 else:
237 expert_mesh = None
238 if ep_size > 0 and expert_mesh is None:
239 raise ValueError("Virtual-EP plan requires MeshContext.fsdp_moe_mesh")
240 if expert_mesh is not None:
241 logger.info(
242 "virtual-EP mesh: using %s for parameter sharding, explicit compute injection, "
243 "and FSDP source metadata",
244 dict(zip(tuple(expert_mesh.mesh_dim_names), tuple(expert_mesh.mesh_shape))),
245 )
246 # Pass the full region meshes: build_source_shard_info strips the FSDP-owned
247 # axes itself and records the complete source (non-FSDP) layout.
248 dense_source_mesh = (
249 mesh_context.fsdp_non_moe_mesh
250 if mesh_context is not None and getattr(mesh_context, "fsdp_non_moe_mesh", None) is not None
251 else tp_mesh
252 )
253 expert_source_mesh = (
254 expert_mesh
255 if expert_mesh is not None and "ep" in expert_mesh.mesh_dim_names
256 else None
257 )
258 return expert_mesh, dense_source_mesh, expert_source_mesh
261def _shard_planned_parameters(models, plan, mesh, expert_mesh, validate_mode):
262 """Shard dense and virtual-EP parameters, then update local metadata."""
263 for model in models:
264 for module_fqn, spec in plan.modules.items():
265 module = _resolve_module(model, module_fqn)
266 info = get_family_build_info(module)
267 if info is not None and info.layout is not None and info.layout.applied:
268 _validate_applied_layout(info.layout, spec, module_fqn)
269 continue
270 # Runtime planner metadata is internal to the sharding pipeline.
271 if spec._ep_stack: # pylint: disable=protected-access
272 _stack_moe_experts(
273 module,
274 spec._ep_stack, # pylint: disable=protected-access
275 )
276 if spec._ep_size > 0: # pylint: disable=protected-access
277 expert_params = {
278 name: placement
279 for name, placement in spec.params.items()
280 if isinstance((placement or {}).get(EP), Shard)
281 }
282 dense_params = {
283 name: placement
284 for name, placement in spec.params.items()
285 if not isinstance((placement or {}).get(EP), Shard)
286 }
287 _shard_module_params(
288 module,
289 expert_params,
290 expert_mesh,
291 expert_mesh.mesh_dim_names,
292 )
293 _shard_module_params(
294 module,
295 dense_params,
296 mesh,
297 plan.mesh_dim_names,
298 )
299 else:
300 _shard_module_params(
301 module,
302 spec.params,
303 mesh,
304 plan.mesh_dim_names,
305 )
306 if not validate_mode:
307 maybe_update_head_counts(
308 module,
309 spec,
310 module_fqn,
311 mesh,
312 plan.mesh_dim_names,
313 )
316def _build_runtime_source_shard_info(
317 models,
318 plan,
319 dense_source_mesh,
320 expert_source_mesh,
321 validate_mode,
322):
323 """Unwrap production parameters and build FSDP source-layout metadata."""
324 if validate_mode:
325 return None
326 source_shard_records = {}
327 existing = {}
328 for model in models:
329 source_shard_records.update(_local_params_context(model))
330 existing.update(_family_source_shard_info(model))
331 if not source_shard_records and not existing:
332 return None
333 if dense_source_mesh is None and expert_source_mesh is None:
334 return existing or None
335 result = build_source_shard_info(
336 plan,
337 dense_source_mesh,
338 expert_source_mesh=expert_source_mesh,
339 )
340 result.update(existing)
341 return result
344def _shard_module_params(module, param_specs, mesh, mesh_dim_names):
345 """distribute_tensor() converts parameters into DTensors.
347 - meta tensor -> DTensor: _local_tensor remains meta (zero-memory path);
348 - real tensor -> DTensor: physically split; each rank holds a local shard;
349 - already a DTensor: skipped if the placement matches, otherwise raises
350 PlacementMismatchError.
351 """
352 for param_path, named in param_specs.items():
353 param = _get_attr_by_path(module, param_path)
354 placements = tuple(resolve_placements(named, mesh_dim_names))
355 if not placements:
356 continue # no active DTensor axes (all size 1) -- no sharding needed
358 if isinstance(param, DTensor):
359 if tuple(param.placements) != placements:
360 raise PlacementMismatchError(
361 f"{type(module).__name__}.{param_path}",
362 placements, tuple(param.placements), "params",
363 )
364 continue
366 src = param.data if hasattr(param, "data") else param
367 dt = distribute_tensor(src, mesh, placements)
368 requires_grad = getattr(param, "requires_grad", True)
369 _set_param_by_path(module, param_path,
370 nn.Parameter(dt, requires_grad=requires_grad))
373# ────────────────────────────────────────────────────────────────────────────
374# Phase D: tied weights
375# ────────────────────────────────────────────────────────────────────────────
377def detect_tied_weights(model: Any) -> List[Tuple[str, str]]:
378 """Detect tied-weight pairs (embed_tokens.weight <-> lm_head.weight).
380 In PP scenarios cross-stage pairs cannot be detected; the user must
381 explicitly declare plan.tied_pairs.
383 Args:
384 model: The model to inspect (an HF-style ``nn.Module`` with a
385 ``config`` carrying ``tie_word_embeddings``).
387 Returns:
388 A list of ``(embed_fqn, lm_head_fqn)`` tied-parameter FQN pairs.
389 """
390 tied = []
391 if getattr(getattr(model, "config", None), "tie_word_embeddings", False):
392 embed_fqn = lm_head_fqn = None
393 # remove_duplicate=False: under the default dedup of named_parameters
394 # a tied parameter appears only once; duplicates must be explicitly
395 # retained to discover the FQNs of both ends.
396 for name, _ in model.named_parameters(remove_duplicate=False):
397 if name.endswith("embed_tokens.weight"):
398 embed_fqn = name
399 elif name.endswith("lm_head.weight"):
400 lm_head_fqn = name
401 if embed_fqn and lm_head_fqn:
402 tied.append((embed_fqn, lm_head_fqn))
403 return tied
406def _broadcast_tied_param(model, tied_pair):
407 """A tied-weight pair shares storage within this rank (end A's storage is authoritative; end B shares it).
409 Cross-rank broadcast would be **wrong**: a tied pair (embed/lm_head) is
410 usually Shard(0)-sharded on both ends, and each rank's local shard carries
411 a different vocab interval -- broadcasting rank0's shard to rank1 would
412 corrupt rank1's sharding. Tied semantics require that **within the same
413 rank** the two ends are the same physical parameter (shared gradients),
414 not cross-rank consistency (sharding is naturally consistent: same global
415 source, same placement).
416 """
417 fqn_a, fqn_b = tied_pair
418 try:
419 param_a = _get_attr_by_path(model, fqn_a)
420 param_b = _get_attr_by_path(model, fqn_b)
421 except AttributeError:
422 return
423 if param_a is None or param_b is None:
424 return
425 tensor_a = param_a.to_local() if isinstance(param_a, DTensor) else param_a.data
426 # B shares storage with A (a tied weight is the same physical parameter)
427 if isinstance(param_b, DTensor):
428 param_b._local_tensor = tensor_a # pylint: disable=protected-access
429 else:
430 param_b.data = tensor_a
433def _replicate_tied_weights(model, tied_pairs=None):
434 """Phase D: replicate tied weights across ranks."""
435 for tied_pair in (tied_pairs if tied_pairs is not None
436 else detect_tied_weights(model)):
437 _broadcast_tied_param(model, tied_pair)