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