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