Diff Coverage

            ep_size=ep_size,
            padded_num_embeddings=self.padded_num_embeddings,
        )
        fused = self._fuse(hidden_states, embeddings)
        token_mask = local_token_mask(token_mask, hidden_states, tp_rank=tp_rank, tp_size=tp_size)
        if token_mask is None:
            return fused
        return torch.where(token_mask.to(torch.bool).unsqueeze(-1).unsqueeze(-1), fused, hidden_states)

                int(getattr(mesh_context, "tp_size", 1)) > 1 or int(getattr(mesh_context, "cp_size", 1)) > 1
        ):
            raise ValueError("OmniParallelBatch requires an initialized process group for TP or CP")
        self.device = device
        self.sequence_parallel = bool(mesh_context.sequence_parallel)
        self.runtime_input_adapter = runtime_input_adapter
        token_pad_values = {field: OMNI_CP_PAD_VALUES.get(field, 0) for field in OMNI_CP_TOKEN_FIELDS}
        cp_size = int(mesh_context.cp_size)
        tp_parts = int(mesh_context.tp_size) if mesh_context.sequence_parallel else 1
        alignment = lcm(2 * cp_size, cp_size * tp_parts)
        self.cp_sharder = CPBatchSharder(
            self.parallel_context, token_pad_values=token_pad_values,
            required_global_alignment=alignment,
        )
        self.tp_broadcaster = TPBatchBroadcaster(self.parallel_context, device)

        """
        self.parallel_context = parallel_context
        self.token_pad_values = token_pad_values
        if required_global_alignment is not None and required_global_alignment < 1:
            raise ValueError("required_global_alignment must be positive")
        self.required_global_alignment = required_global_alignment

    def shard(
            self,

    @staticmethod
    def _extend_cu_seq_lens(cp_batch: dict[str, Any], sequence_length: int, pad_length: int) -> None:
        """Keep packed sequence boundaries aligned with token padding."""
        if "cu_seq_lens" not in cp_batch:
            return
        boundaries = cp_batch["cu_seq_lens"].reshape(-1)
        if int(boundaries[0]) != 0 or int(boundaries[-1]) != sequence_length:
            raise ValueError("cu_seq_lens must cover the complete CP input sequence")
        if pad_length:
            boundaries = torch.cat((boundaries, boundaries[-1:] + pad_length))
        cp_batch["cu_seq_lens"] = boundaries

    every other boundary as a descendant.
    """
    boundaries = [fqn for fqn, spec in plan.modules.items() if spec.is_boundary]
    if module_fqn == "":
        return [f for f in boundaries if f]
    prefix = module_fqn + "."
    return [f[len(prefix):] for f in boundaries if f.startswith(prefix)]

                                     key_heads: int, value_heads: int) -> None:
        """Check projected head counts before entering a collective."""
        actual = {"query": query_heads, "key": key_heads, "value": value_heads}
        if actual != self.declared_heads or any(count % self.cp_mesh.size() for count in actual.values()):
            raise ValueError(
                "[HP-ENCODER-005] Ulysses runtime head count mismatch; phase=runtime; "
                f"module_fqn={module_fqn}; declared_head_counts={self.declared_heads}; "
                f"runtime_head_counts={actual}; encoder_cp_size={self.cp_mesh.size()}"
            )

        """Validate locally, build other inputs, then agree on the visual plan."""
        prepared = self.runtime.preflight_batch(batch)
        generated = {} if self.inner is None else self.inner.build_runtime_inputs(batch=batch, context=context)
        if not isinstance(generated, Mapping):
            raise TypeError("RuntimeInputAdapter.build_runtime_inputs must return a mapping")
        inputs = dict(generated)
        if "encoder_parallel_plan" in batch or "encoder_parallel_plan" in inputs or set(batch).intersection(inputs):
            raise ValueError("[HP-DATA-001] runtime inputs cannot replace Omni model inputs")
        inputs["encoder_parallel_plan"] = self.runtime.finalize_plan(*prepared)

    encoder_cp_size: int = 1

    def __post_init__(self) -> None:
        if not isinstance(self.enabled, bool):
            raise ValueError("[HP-ENCODER-002] enabled must be a boolean")
        if (isinstance(self.encoder_cp_size, bool) or not isinstance(self.encoder_cp_size, int)
                or self.encoder_cp_size < 1):
            raise ValueError(
                "[HP-ENCODER-002] accelerator.encoder_parallel.encoder_cp_size "

        Args:
            raw: User's optional encoder-parallel configuration mapping.
        """
        if raw is None:
            return cls()
        if not isinstance(raw, Mapping):
            raise ValueError("[HP-ENCODER-002] encoder_parallel must be a mapping")
        unknown = set(raw) - {field.name for field in fields(cls)}
        if unknown:
            raise ValueError(f"[HP-ENCODER-002] unknown encoder_parallel keys: {sorted(unknown)}")
        return cls(**raw)

        return FeatureExchangeMode.PARENT_VARIABLE_A2A
    if mesh_context.tp_size == 1:
        return FeatureExchangeMode.LOCAL
    if mesh_context.sequence_parallel:
        return FeatureExchangeMode.TP_VARIABLE_A2A
    return FeatureExchangeMode.TP_VARIABLE_ALL_GATHER


def cat_or_empty(rows: Sequence[torch.Tensor], *, anchor: torch.Tensor,

        feature_width: Width of an empty feature tensor.
    """
    if rows:
        return torch.cat(tuple(rows), dim=0)
    return anchor.sum().reshape(1, 1).expand(0, feature_width)


def _ordered_source_records(buckets: Sequence[Sequence[EncoderWorkItem]], encoder_cp_size: int
                            ) -> dict[int, tuple[EncoderWorkItem, ...]]:

                         source_records: Mapping[int, tuple[EncoderWorkItem, ...]],
                         prototype: torch.Tensor, parent_rank: int,
                         exchange_mesh: Any) -> tuple[torch.Tensor, tuple[EncoderWorkItem, ...]]:
    """Gather each producer's complete records to every decoder TP rank."""
    group_size = exchange_mesh.size()
    local_items = source_records.get(parent_rank, ())
    local_rows = cat_or_empty([local_by_id[item.item_id] for item in local_items],
                              anchor=prototype, feature_width=int(prototype.shape[-1]))
    splits = [sum(item.output_rows for item in source_records.get(rank, ())) for rank in range(group_size)]
    rows = differentiable_variable_all_gather(local_rows, output_splits=splits,
                                               group=exchange_mesh.get_group())
    received_items = tuple(item for rank in range(group_size) for item in source_records.get(rank, ()))
    return rows, received_items


def _route_variable_a2a(*, mode: FeatureExchangeMode,
                        local_by_id: Mapping[int, torch.Tensor],

                        destinations: Mapping[int, tuple[int, ...]],
                        encoder_cp_size: int, exchange_mesh: Any
                        ) -> tuple[torch.Tensor, tuple[EncoderWorkItem, ...]]:
    """Send each full record only to its exact decoder token owners."""
    group_size = exchange_mesh.size()
    local_rank = exchange_mesh.get_local_rank()
    if group_size != len(source_records) * encoder_cp_size:
        raise ValueError(f"[HP-ENCODER-009] {mode.value} exchange size disagrees with encoder parent")
    width = int(prototype.shape[-1])
    send_chunks = []
    input_splits = []
    for destination_rank in range(group_size):
        items = [item for item in source_records.get(parent_rank, ())
                 if destination_rank in destinations[item.item_id]]
        chunk = cat_or_empty([local_by_id[item.item_id] for item in items],
                             anchor=prototype, feature_width=width)
        send_chunks.append(chunk)
        input_splits.append(chunk.shape[0])
    packed = cat_or_empty(send_chunks, anchor=prototype, feature_width=width)
    output_splits = []
    received_items = []
    for source_rank in range(group_size):
        items = [item for item in source_records.get(source_rank, ())
                 if local_rank in destinations[item.item_id]]
        output_splits.append(sum(item.output_rows for item in items))
        received_items.extend(items)
    rows = differentiable_all_to_all_single(packed, input_splits=input_splits,
                                            output_splits=output_splits, group=exchange_mesh.get_group())
    return rows, tuple(received_items)


def route_features(*, mode: FeatureExchangeMode, producer_records: Sequence[EncoderFeature],
                   prototype: torch.Tensor, parent_rank: int, destinations: Mapping[int, tuple[int, ...]],

        received_items = tuple(item for item, _ in producer_records)
        rows = cat_or_empty([feature for _, feature in producer_records], anchor=prototype, feature_width=width)
        return rows, received_items

    if mode is FeatureExchangeMode.TP_VARIABLE_ALL_GATHER:
        return _route_tp_all_gather(local_by_id=local_by_id, source_records=source_records,
                                    prototype=prototype, parent_rank=parent_rank,
                                    exchange_mesh=exchange_mesh)
    if mode in (FeatureExchangeMode.TP_VARIABLE_A2A, FeatureExchangeMode.PARENT_VARIABLE_A2A):
        return _route_variable_a2a(
            mode=mode, local_by_id=local_by_id, source_records=source_records,
            prototype=prototype, parent_rank=parent_rank, destinations=destinations,
            encoder_cp_size=encoder_cp_size, exchange_mesh=exchange_mesh,
        )
    raise ValueError(f"[HP-ENCODER-009] unsupported feature exchange mode: {mode!r}")


def unpack_records(rows: torch.Tensor, received_items: Sequence[EncoderWorkItem]
                   ) -> tuple[EncoderFeature, ...]:

        stop = offset + item.output_rows
        records.append((item, rows[offset:stop]))
        offset = stop
    if offset != rows.shape[0]:
        raise ValueError(
            "[HP-ENCODER-009] transported rows disagree with planned images; "
            f"rows={rows.shape[0]}; described={offset}"
        )
    return tuple(sorted(records, key=lambda record: (

        target_names: Names assigned to the derived axes.
    """
    source = root_mesh[source_dims]
    if math.prod(target_sizes) != math.prod(source.mesh_shape) or len(target_sizes) != len(target_names):
        raise ValueError("[HP-ENCODER-002] encoder mesh shape does not match decoder parent")
    flat = source.flatten("encoder_parent")
    return flat._unflatten("encoder_parent", target_sizes, target_names)  # pylint: disable=protected-access

        cp_rank = int(mesh_context.cp_rank if cp_rank is None else cp_rank)
        tp_rank = int(mesh_context.tp_rank if tp_rank is None else tp_rank)
        sp = bool(mesh_context.sequence_parallel)
        if not 0 <= cp_rank < cp_size or not 0 <= tp_rank < tp_size:
            raise ValueError(
                "[HP-ENCODER-004] decoder CP/TP coordinate is outside topology; "
                f"cp_rank={cp_rank}; cp_size={cp_size}; tp_rank={tp_rank}; tp_size={tp_size}"
            )
        if cp_local_length < 1 or (sp and cp_local_length % tp_size):

            value: Auxiliary tensor replicated across decoder TP ranks.
            sequence_dim: Dimension containing CP-local token rows.
        """
        if value.shape[sequence_dim] != self.cp_local_length:
            raise ValueError(
                "[HP-ENCODER-004] CP-local token tensor length mismatch; "
                f"actual={value.shape[sequence_dim]}; expected={self.cp_local_length}; "
                f"cp_rank={self.cp_rank}; tp_rank={self.tp_rank}"
            )

        encoder_dp_size: Number of visual work buckets.
        encoder_cp_size: Sequence-sharding degree of each bucket.
    """
    if encoder_dp_size < 1 or encoder_cp_size < 1:
        raise ValueError("[HP-ENCODER-002] encoder DP and CP sizes must be positive")
    buckets: list[list[EncoderWorkItem]] = [[] for _ in range(encoder_dp_size)]
    loads = [0] * encoder_dp_size
    for item in sorted(items, key=lambda value: (-value.sequence_length ** 2, value.batch_index, value.item_id)):
        bucket = min(range(encoder_dp_size), key=lambda rank: (loads[rank], rank))

        encoder_cp_rank: Sequence-shard coordinate within the bucket.
        encoder_cp_size: Number of shards per image.
    """
    if not 0 <= encoder_cp_rank < encoder_cp_size:
        raise ValueError("[HP-ENCODER-002] invalid encoder CP rank")
    segments = []
    offset = 0
    for item in items:
        if item.sequence_length < 1:
            raise ValueError("[HP-ENCODER-013] image patch count must be positive")
        local_length = (item.sequence_length + encoder_cp_size - 1) // encoder_cp_size
        start = encoder_cp_rank * local_length
        valid = max(0, min(local_length, item.sequence_length - start))
        segments.append(EncoderSegment(item, start, valid, local_length, offset))

        layouts: Decoder token interval for every rank in the parent group.
    """
    lengths = {layout.global_sequence_length for layout in layouts.values()}
    if len(lengths) != 1:
        raise ValueError("[HP-ENCODER-004] inconsistent decoder global lengths")
    global_length = next(iter(lengths))
    result = {}
    for item in items:
        start = item.decoder_token_start

            )
        owners = tuple(rank for rank, layout in layouts.items()
                       if max(start, layout.global_start) < min(stop, layout.global_stop))
        if not owners:
            raise ValueError(f"[HP-ENCODER-016] image span has no decoder token owner; item_id={item.item_id}")
        result[item.item_id] = owners
    return result

        bridge: Model-bound validator for vision-specific restrictions.
    """
    bridge.validate_training(options)
    if options.get("compile_enabled"):
        raise NotImplementedError("[HP-ENCODER-011] v1 requires compile.enabled=false")
    if options.get("validate_placement"):
        raise NotImplementedError("[HP-ENCODER-012] v1 requires model.validate_placement=false")


@dataclass
class EncoderParallelPlan:

            real = segment is not None and not segment.is_dummy and self.runtime.encoder_cp_rank == 0
            if real:
                feature = bridge.project(full_features[slot], segment.item)
                if feature.ndim != 2 or feature.shape[0] != segment.item.output_rows:
                    raise ValueError(
                        "[HP-ENCODER-017] aligner output does not match planned rows; "
                        f"item_id={segment.item.item_id}; actual_shape={tuple(feature.shape)}; "
                        f"expected_rows={segment.item.output_rows}"
                    )

                        f"expected_rows={segment.item.output_rows}"
                    )
                records.append((segment.item, feature))
            else:
                feature = bridge.project(hidden, None)
                anchors.append(feature.sum() * 0.0)
            prototype = feature
        if prototype is None:
            raise RuntimeError("[HP-ENCODER-009] active plan must have an aligner slot")
        anchor = torch.stack(anchors).sum() if anchors else prototype.sum() * 0.0
        return tuple(records), prototype, anchor

    def encode_and_route(self, image_inputs: Mapping[str, Any]

            {**execution_options, "loss_parallel": mesh_context.loss_parallel}, self.bridge,
        )
        parent_size = mesh_context.cp_size * mesh_context.tp_size
        if mesh_context.pp_size != 1:
            raise NotImplementedError("[HP-ENCODER-003] v1 requires pp_size=1")
        if parent_size % config.encoder_cp_size:
            raise ValueError("[HP-ENCODER-002] encoder_cp_size must divide cp_size*tp_size")
        validate_ulysses_head_divisibility(encoder_cp_size=config.encoder_cp_size,
                                           head_counts=self.bridge.head_counts,
                                           attention_pattern=spec.attention_pattern, model=model)
        self.expected = validate_attention_replacement(model, spec)

        """
        marker = torch.tensor(int(bool(items)), device=self.model.device if hasattr(self.model, "device")
                              else next(self.model.parameters()).device)
        if dist.is_initialized() and self.parent_mesh.size() > 1:
            signature = tuple(
                (item.item_id, item.batch_index, item.sequence_length, item.output_rows,
                 item.decoder_token_start, item.decoder_token_length, item.payload_start, item.payload_stop,
                 tuple(sorted(item.metadata.items())))
                for item in items

                 item.decoder_token_start, item.decoder_token_length, item.payload_start, item.payload_stop,
                 tuple(sorted(item.metadata.items())))
                for item in items
            )
            digest = hashlib.sha256(repr((layout.cp_local_length, signature)).encode()).digest()
            digest_value = int.from_bytes(digest[:8], "big") & ((1 << 63) - 1)
            local_digest = torch.tensor(digest_value, dtype=torch.int64, device=marker.device)
            peer_digests = [torch.empty_like(local_digest) for _ in range(self.parent_mesh.size())]
            dist.all_gather(peer_digests, local_digest, group=self.parent_mesh.get_group())
            if any(int(peer.item()) != digest_value for peer in peer_digests):
                raise ValueError("[HP-ENCODER-007] image metadata differs across decoder parent ranks")
        if dist.is_initialized() and self.global_mesh.size() > 1:
            dist.all_reduce(marker, op=dist.ReduceOp.MAX, group=self.global_mesh.get_group())
        active = bool(marker.item())
        buckets = assign_lpt(items, self.encoder_dp_size, self.config.encoder_cp_size)
        layouts = {rank: DecoderSequenceLayout.build(mesh_context=self.mesh_context,
                                                     cp_local_length=layout.cp_local_length,

        segments = build_segments(buckets[self.encoder_dp_rank], self.encoder_cp_rank,
                                  self.config.encoder_cp_size)
        slot_count = torch.tensor(len(buckets[self.encoder_dp_rank]), device=marker.device)
        if dist.is_initialized() and self.global_mesh.size() > 1:
            dist.all_reduce(slot_count, op=dist.ReduceOp.MAX, group=self.global_mesh.get_group())
        return EncoderParallelPlan(
            runtime=self,
            active=active, decoder_layout=layout, buckets=buckets,
            destinations=destinations, local_segments=segments,

def build_encoder_parallel_runtime(*, config: EncoderParallelConfig, mesh_context: Any,
                                   model: nn.Module | None, model_adapter_spec: Any,
                                   execution_options: Mapping[str, Any] | None) -> EncoderParallelRuntime | None:
    """Create the model-neutral runtime only when the feature is enabled."""
    if not config.enabled:
        return None
    if model is None or model_adapter_spec is None or model_adapter_spec.encoder_parallel is None:
        raise ValueError("[HP-ENCODER-001] enabled encoder parallel requires a model adapter provider")
    spec = model_adapter_spec.encoder_parallel()
    if not isinstance(spec, EncoderParallelAdapterSpec):
        raise TypeError("[HP-ENCODER-001] encoder_parallel provider returned an invalid spec")
    return EncoderParallelRuntime(config=config, mesh_context=mesh_context, model=model,
                                  spec=spec, execution_options=execution_options or {})

    if entry.low_precision_dtype_scheme is not None and (
        not isinstance(entry.low_precision_dtype_scheme, str)
        or not entry.low_precision_dtype_scheme.strip()
    ):
        raise ValueError(
            f"plan_overrides match={entry.match!r}: "
            "low_precision_dtype_scheme must be a non-empty string"
        )
    if entry.low_precision_dtype_scheme is not None and entry.when != "low_precision":

            f"plan_overrides match={entry.match!r} sets "
            "low_precision_dtype_scheme but is not a low-precision entry"
        )
    if entry.when == "low_precision" and effective_config is None:
        raise ValueError(
            f"plan_overrides match={entry.match!r} enables a low-precision "
            "replacement but training.low_precision is missing"
        )
    if entry.module_type is None:
        raise ValueError(
            f"plan_overrides match={entry.match!r} uses replace_module but omits module_type"
        )
    if not isinstance(entry.replace_module, Target):
        raise TypeError(
            f"plan_overrides match={entry.match!r} replace_module must be a YAML Target"
        )

from collections import defaultdict
from typing import Any, Dict

from hyper_parallel.data.batching import calculate_num_micro_batches
from hyper_parallel.data.batching.get_batch import OmniParallelBatch
from hyper_parallel.data.omni import OmniDataTransform
from hyper_parallel.distributed.encoder_parallel import build_encoder_parallel_runtime
from hyper_parallel.distributed.encoder_parallel.batch_input import EncoderParallelInputAdapter
from hyper_parallel.models.registry import resolve_model_adapter
from hyper_parallel.trainer.runtime.loss_aggregation import count_loss_token
from hyper_parallel.trainer.runtime.logging import create_logger
from hyper_parallel.trainer.runtime.memory import print_device_mem_info
from hyper_parallel.trainer.runtime.device import synchronize  # pylint: disable=syntax-error

        """Build the DataLoader-to-LLM batch runtime."""
        config = self.base.config
        if config.dataloader.get_batch is None:
            raise ValueError("dataloader.get_batch must define a batching runtime target")
        if getattr(config.dataloader.get_batch, "encoder_parallel", None) is not None:
            raise ValueError("encoder_parallel belongs under accelerator, not dataloader.get_batch")
        encoder_runtime = build_encoder_parallel_runtime(
            config=config.accelerator.encoder_parallel,
            mesh_context=self.base.mesh,
            model=self.base.model,
            model_adapter_spec=resolve_model_adapter(self.base.model),

            device=self.base.device,
            data_config=getattr(config.dataset, "data_config", {}),
            pp_shared_data=bool(getattr(config.dataloader, "pp_shared_data", False)),
        )
        if encoder_runtime is not None:
            if not isinstance(get_batch, OmniParallelBatch):
                raise ValueError("[HP-ENCODER-001] encoder parallel requires OmniParallelBatch")
            get_batch.runtime_input_adapter = EncoderParallelInputAdapter(
                runtime=encoder_runtime,
                inner=get_batch.runtime_input_adapter,
            )
        self.base.get_batch = get_batch