Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / encoder_parallel / runtime.py: 83%
156 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07: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"""Per-batch visual encoder planning and packed execution."""
17from __future__ import annotations
19import fnmatch
20import hashlib
21from dataclasses import dataclass
22from typing import Any, Mapping
24import torch # pylint: disable=forbidden-backend-import
25import torch.distributed as dist # pylint: disable=forbidden-backend-import
26from torch import nn # pylint: disable=forbidden-backend-import
28from hyper_parallel.core.dtensor._utils import differentiable_variable_all_gather
29from hyper_parallel.distributed.encoder_parallel.attention import (
30 EncoderAttentionContext,
31 activate_encoder_attention,
32)
33from hyper_parallel.distributed.encoder_parallel.config import EncoderParallelConfig
34from hyper_parallel.distributed.encoder_parallel.contracts import (
35 EncoderFeature,
36 EncoderLocalInputs,
37 EncoderParallelAdapterSpec,
38 EncoderSegment,
39 EncoderWorkItem,
40)
41from hyper_parallel.distributed.encoder_parallel.exchange import (
42 FeatureExchangeMode,
43 route_features,
44 select_feature_exchange,
45 unpack_records,
46)
47from hyper_parallel.distributed.encoder_parallel.mesh import derive_reshaped_submesh
48from hyper_parallel.distributed.encoder_parallel.planning import (
49 DecoderSequenceLayout,
50 assign_lpt,
51 build_destinations,
52 build_segments,
53)
56def _invalid_head_count(count: Any) -> bool:
57 """Reject booleans, nonintegers, and zero-sized QKV head axes."""
58 return isinstance(count, bool) or not isinstance(count, int) or count < 1
61def validate_ulysses_head_divisibility(*, encoder_cp_size: int, head_counts: Mapping[str, int],
62 attention_pattern: str, model: nn.Module) -> None:
63 """Check every Q/K/V head count before deriving any Encoder-CP group."""
64 required = ("query", "key", "value")
65 missing = [name for name in required if name not in head_counts]
66 invalid = {name: head_counts[name] for name in required if name in head_counts
67 and _invalid_head_count(head_counts[name])}
68 non_divisible = {name: head_counts[name] for name in required if name in head_counts
69 and name not in invalid and head_counts[name] % encoder_cp_size}
70 if missing or invalid or non_divisible:
71 matched = sorted(name for name, _ in model.named_modules()
72 if fnmatch.fnmatchcase(name, attention_pattern))
73 raise ValueError(
74 "Encoder-CP/Ulysses head divisibility check failed; "
75 "constraint=head_count % encoder_cp_size == 0 for every Q/K/V; "
76 "config_path=accelerator.encoder_parallel.encoder_cp_size; "
77 f"encoder_cp_size={encoder_cp_size}; head_counts={dict(head_counts)}; "
78 f"missing={missing}; invalid_values={invalid}; non_divisible={non_divisible}; "
79 f"attention_pattern={attention_pattern!r}; matched_modules={matched}"
80 )
83def validate_attention_replacement(model: nn.Module, spec: EncoderParallelAdapterSpec) -> frozenset[str]:
84 """Require an exact structural replacement for every visual attention.
86 Args:
87 model: Final model after module replacements.
88 spec: Model-owned attention matching and replacement contract.
89 """
90 expected = frozenset(name for name, _ in model.named_modules()
91 if fnmatch.fnmatchcase(name, spec.attention_pattern))
92 recorded = set(getattr(model, "_hp_replaced_module_fqns", ()) or ())
93 recorded_visual = {name for name in recorded if fnmatch.fnmatchcase(name, spec.attention_pattern)}
94 missing = sorted(expected - recorded_visual)
95 unexpected = sorted(recorded_visual - expected)
96 wrong_type = sorted(name for name in expected
97 if not isinstance(model.get_submodule(name), spec.replacement_attention_type))
98 if not expected or missing or unexpected or wrong_type:
99 raise RuntimeError(
100 "incomplete vision attention replacement: "
101 f"expected={len(expected)}; missing={missing}; unexpected={unexpected}; wrong_type={wrong_type}"
102 )
103 return expected
106def validate_encoder_execution_options(options: Mapping[str, Any], bridge: Any) -> None:
107 """Reject unsupported execution modes before creating Encoder groups.
109 Args:
110 options: Trainer microbatch, compile, and placement settings.
111 bridge: Model-bound validator for vision-specific restrictions.
112 """
113 bridge.validate_training(options)
114 if options.get("compile_enabled"):
115 raise NotImplementedError("v1 requires compile.enabled=false")
116 if options.get("validate_placement"):
117 raise NotImplementedError("v1 requires model.validate_placement=false")
120@dataclass
121class EncoderParallelPlan:
122 """One microbatch of derived visual work and feature routing."""
124 runtime: "EncoderParallelRuntime"
125 active: bool
126 decoder_layout: DecoderSequenceLayout
127 buckets: tuple[tuple[EncoderWorkItem, ...], ...]
128 destinations: dict[int, tuple[int, ...]]
129 local_segments: tuple[EncoderSegment, ...]
130 k_fsdp_max: int
132 def _gather_features(self, hidden: torch.Tensor) -> tuple[torch.Tensor, ...]:
133 """All-gather each local segment and trim dummy padding."""
134 gathered = []
135 offset = 0
136 for segment in self.local_segments:
137 local = hidden[offset:offset + segment.local_padded_length]
138 offset += segment.local_padded_length
139 full = differentiable_variable_all_gather(
140 local, output_splits=[segment.local_padded_length] * self.runtime.config.encoder_cp_size,
141 group=self.runtime.encoder_cp_mesh.get_group(),
142 ) if self.runtime.config.encoder_cp_size > 1 else local
143 gathered.append(full if segment.is_dummy else full[:segment.item.sequence_length])
144 return tuple(gathered)
146 def _execute_local(self, image_inputs: Mapping[str, Any] | EncoderLocalInputs
147 ) -> tuple[tuple[EncoderFeature, ...], torch.Tensor, torch.Tensor]:
148 """Run the local ViT pass and build routed feature records."""
149 bridge = self.runtime.bridge
150 context = EncoderAttentionContext(
151 self.local_segments,
152 self.runtime.encoder_cp_mesh,
153 bridge.head_counts,
154 self.runtime.expected,
155 )
156 with activate_encoder_attention(context):
157 hidden = bridge.encode_local(image_inputs, self.local_segments)
158 context.assert_exactly_once()
159 full_features = self._gather_features(hidden)
160 records = []
161 anchors = [feature.sum() * 0.0 for feature in full_features]
162 prototype = None
163 for slot in range(self.k_fsdp_max):
164 segment = self.local_segments[slot] if slot < len(self.local_segments) else None
165 real = segment is not None and not segment.is_dummy and self.runtime.encoder_cp_rank == 0
166 if real:
167 feature = bridge.project(full_features[slot], segment.item)
168 if feature.ndim != 2 or feature.shape[0] != segment.item.output_rows:
169 raise ValueError(
170 "aligner output does not match planned rows; "
171 f"item_id={segment.item.item_id}; actual_shape={tuple(feature.shape)}; "
172 f"expected_rows={segment.item.output_rows}"
173 )
174 records.append((segment.item, feature))
175 else:
176 feature = bridge.project(hidden, None)
177 anchors.append(feature.sum() * 0.0)
178 prototype = feature
179 if prototype is None:
180 raise RuntimeError("active plan must have an aligner slot")
181 anchor = torch.stack(anchors).sum() if anchors else prototype.sum() * 0.0
182 return tuple(records), prototype, anchor
184 def encode_and_route(self, image_inputs: Mapping[str, Any] | EncoderLocalInputs
185 ) -> tuple[list[EncoderFeature], torch.Tensor | None]:
186 """Run packed ViT and route aligned features to this decoder owner.
188 Args:
189 image_inputs: Prepared local patches, or a native image bundle for direct callers.
191 Returns:
192 Routed image records and a zero-valued gradient anchor, or no work.
193 """
194 if not self.active:
195 return [], None
196 records, prototype, local_anchor = self._execute_local(image_inputs)
197 mode = self.runtime.feature_exchange_mode
198 exchange_mesh = self.runtime.tp_mesh if mode in (
199 FeatureExchangeMode.TP_VARIABLE_ALL_GATHER, FeatureExchangeMode.TP_VARIABLE_A2A
200 ) else self.runtime.parent_mesh
201 rows, received_items = route_features(
202 mode=mode, producer_records=records, prototype=prototype,
203 parent_rank=self.runtime.parent_rank, destinations=self.destinations,
204 buckets=self.buckets, encoder_cp_size=self.runtime.config.encoder_cp_size, exchange_mesh=exchange_mesh,
205 )
206 unpacked = unpack_records(rows, received_items)
207 return list(unpacked), local_anchor + rows.sum() * 0.0
210class EncoderParallelRuntime:
211 """Own static topology and create one plan per Omni microbatch."""
213 def __init__(self, *, config: EncoderParallelConfig, mesh_context: Any,
214 model: nn.Module, spec: EncoderParallelAdapterSpec,
215 execution_options: Mapping[str, Any]) -> None:
216 """Validate the static contract and derive Encoder-CP groups."""
217 self.config = config
218 self.mesh_context = mesh_context
219 self.model = model
220 self.spec = spec
221 self.bridge = spec.bind(model)
222 validate_encoder_execution_options(
223 {**execution_options, "loss_parallel": mesh_context.loss_parallel}, self.bridge,
224 )
225 parent_size = mesh_context.cp_size * mesh_context.tp_size
226 if mesh_context.pp_size != 1:
227 raise NotImplementedError("v1 requires pp_size=1")
228 if parent_size % config.encoder_cp_size:
229 raise ValueError("encoder_cp_size must divide cp_size*tp_size")
230 validate_ulysses_head_divisibility(encoder_cp_size=config.encoder_cp_size,
231 head_counts=self.bridge.head_counts,
232 attention_pattern=spec.attention_pattern, model=model)
233 self.expected = validate_attention_replacement(model, spec)
234 self.parent_rank = mesh_context.cp_rank * mesh_context.tp_size + mesh_context.tp_rank
235 self.encoder_cp_rank = self.parent_rank % config.encoder_cp_size
236 self.encoder_dp_rank = self.parent_rank // config.encoder_cp_size
237 self.encoder_dp_size = parent_size // config.encoder_cp_size
238 self.parent_mesh = mesh_context.device_mesh[("cp", "tp")].flatten("encoder_parent")
239 encoder_mesh = derive_reshaped_submesh(
240 mesh_context.device_mesh, source_dims=("cp", "tp"),
241 target_sizes=(self.encoder_dp_size, config.encoder_cp_size),
242 target_names=("encoder_dp", "encoder_cp"),
243 )
244 self.encoder_cp_mesh = encoder_mesh["encoder_cp"]
245 self.tp_mesh = mesh_context.device_mesh["tp"]
246 self.global_mesh = mesh_context.device_mesh.flatten("encoder_global")
247 self.feature_exchange_mode = select_feature_exchange(mesh_context)
249 def preflight_batch(self, batch: Mapping[str, Any]
250 ) -> tuple[tuple[EncoderWorkItem, ...], DecoderSequenceLayout]:
251 """Validate local metadata and SP geometry before control collectives.
253 Args:
254 batch: Device batch with decoder tokens and optional image fields.
255 """
256 layout = DecoderSequenceLayout.build(
257 mesh_context=self.mesh_context, cp_local_length=int(batch["input_ids"].shape[1])
258 )
259 present = [batch.get(name) is not None for name in self.bridge.input_fields]
260 if any(present) and not all(present):
261 missing = [name for name in self.bridge.input_fields if batch.get(name) is None]
262 raise ValueError(f"incomplete image metadata bundle: missing={missing}")
263 return tuple(self.bridge.parse_images(batch)), layout
265 def finalize_plan(self, items: tuple[EncoderWorkItem, ...],
266 layout: DecoderSequenceLayout) -> EncoderParallelPlan:
267 """Agree on active mode, then assign images and exact decoder owners.
269 Args:
270 items: Locally validated image jobs.
271 layout: Locally validated decoder token ownership.
272 """
273 marker = torch.tensor(int(bool(items)), device=self.model.device if hasattr(self.model, "device")
274 else next(self.model.parameters()).device)
275 if dist.is_initialized() and self.parent_mesh.size() > 1:
276 signature = tuple(
277 (item.item_id, item.batch_index, item.sequence_length, item.output_rows,
278 item.decoder_token_start, item.decoder_token_length, item.payload_start, item.payload_stop,
279 tuple(sorted(item.metadata.items())))
280 for item in items
281 )
282 digest = hashlib.sha256(repr((layout.cp_local_length, signature)).encode()).digest()
283 digest_value = int.from_bytes(digest[:8], "big") & ((1 << 63) - 1)
284 local_digest = torch.tensor(digest_value, dtype=torch.int64, device=marker.device)
285 peer_digests = [torch.empty_like(local_digest) for _ in range(self.parent_mesh.size())]
286 dist.all_gather(peer_digests, local_digest, group=self.parent_mesh.get_group())
287 if any(int(peer.item()) != digest_value for peer in peer_digests):
288 raise ValueError("image metadata differs across decoder parent ranks")
289 if dist.is_initialized() and self.global_mesh.size() > 1:
290 dist.all_reduce(marker, op=dist.ReduceOp.MAX, group=self.global_mesh.get_group())
291 active = bool(marker.item())
292 buckets = assign_lpt(items, self.encoder_dp_size, self.config.encoder_cp_size)
293 layouts = {rank: DecoderSequenceLayout.build(mesh_context=self.mesh_context,
294 cp_local_length=layout.cp_local_length,
295 cp_rank=rank // self.mesh_context.tp_size,
296 tp_rank=rank % self.mesh_context.tp_size)
297 for rank in range(self.mesh_context.cp_size * self.mesh_context.tp_size)}
298 destinations = build_destinations(items, layouts)
299 segments = build_segments(buckets[self.encoder_dp_rank], self.encoder_cp_rank,
300 self.config.encoder_cp_size)
301 slot_count = torch.tensor(len(buckets[self.encoder_dp_rank]), device=marker.device)
302 if dist.is_initialized() and self.global_mesh.size() > 1:
303 dist.all_reduce(slot_count, op=dist.ReduceOp.MAX, group=self.global_mesh.get_group())
304 return EncoderParallelPlan(
305 runtime=self,
306 active=active, decoder_layout=layout, buckets=buckets,
307 destinations=destinations, local_segments=segments,
308 k_fsdp_max=max(1, int(slot_count.item())),
309 )
312def build_encoder_parallel_runtime(*, config: EncoderParallelConfig, mesh_context: Any,
313 model: nn.Module | None, model_adapter_spec: Any,
314 execution_options: Mapping[str, Any] | None) -> EncoderParallelRuntime | None:
315 """Create the model-neutral runtime only when the feature is enabled."""
316 if not config.enabled:
317 return None
318 if model is None or model_adapter_spec is None or model_adapter_spec.encoder_parallel is None:
319 raise ValueError("enabled encoder parallel requires a model adapter provider")
320 spec = model_adapter_spec.encoder_parallel()
321 if not isinstance(spec, EncoderParallelAdapterSpec):
322 raise TypeError("encoder_parallel provider returned an invalid spec")
323 return EncoderParallelRuntime(config=config, mesh_context=mesh_context, model=model,
324 spec=spec, execution_options=execution_options or {})