diff --git a/areno/accel/ops.py b/areno/accel/ops.py index 6c85d7f1..76b278ea 100644 --- a/areno/accel/ops.py +++ b/areno/accel/ops.py @@ -12,11 +12,8 @@ from __future__ import annotations -import logging from typing import Any -import torch - from areno.accel.activations import areno_gelu_tanh_and_mul, areno_silu_and_mul from areno.accel.attention import ( areno_causal_attention, @@ -28,46 +25,7 @@ from areno.accel.kernels.fused_moe import is_available as fused_moe_is_available from areno.accel.kernels.group_rmsnorm import rms_norm_gate_fwd from areno.accel.kernels.seg_la import SegLaMeta, seg_la_fwd - -logger = logging.getLogger(__name__) -# Process-wide set of message keys already emitted by log_once/warn_once. -_LOGGED: set[str] = set() - - -def log_once(key: str, message: str, *, level: int = logging.DEBUG) -> None: - """Log ``message`` at most once per process for the given ``key``.""" - - if key in _LOGGED: - return - logger.log(level, message) - _LOGGED.add(key) - - -def warn_once(key: str, message: str) -> None: - """Emit a warning at most once per process for the given ``key``.""" - - log_once(key, message, level=logging.WARNING) - - -@torch._dynamo.disable -def is_cuda_graph_capturing(tensor: torch.Tensor) -> bool: - """True if the tensor lives on CUDA and we are inside a graph capture.""" - - return tensor.is_cuda and torch.cuda.is_current_stream_capturing() - - -@torch._dynamo.disable -def can_use_cuda_kernel(tensor: torch.Tensor, name: str, *, allow_sm121: bool = False) -> bool: - """Decide whether to take the fused kernel path for ``tensor``. - - Returns False only on non-CUDA tensors. ``name`` and ``allow_sm121`` are - kept for compatibility with existing call sites. - """ - - if not tensor.is_cuda: - return False - return True - +from areno.accel.utils import can_use_cuda_kernel, is_cuda_graph_capturing, log_once, warn_once __all__ = [ "Any", diff --git a/areno/accel/utils.py b/areno/accel/utils.py new file mode 100644 index 00000000..a04a5696 --- /dev/null +++ b/areno/accel/utils.py @@ -0,0 +1,41 @@ +"""Lightweight acceleration helpers that do not import optional kernels.""" + +from __future__ import annotations + +import logging + +import torch + +logger = logging.getLogger(__name__) +_LOGGED: set[str] = set() + + +def log_once(key: str, message: str, *, level: int = logging.DEBUG) -> None: + """Log ``message`` at most once per process for the given ``key``.""" + + if key in _LOGGED: + return + logger.log(level, message) + _LOGGED.add(key) + + +def warn_once(key: str, message: str) -> None: + """Emit a warning at most once per process for the given ``key``.""" + + log_once(key, message, level=logging.WARNING) + + +@torch._dynamo.disable +def is_cuda_graph_capturing(tensor: torch.Tensor) -> bool: + """True if the tensor lives on CUDA and we are inside a graph capture.""" + + return tensor.is_cuda and torch.cuda.is_current_stream_capturing() + + +@torch._dynamo.disable +def can_use_cuda_kernel(tensor: torch.Tensor, name: str, *, allow_sm121: bool = False) -> bool: + """Return whether a fused CUDA kernel can run for ``tensor``.""" + + if not tensor.is_cuda: + return False + return True diff --git a/areno/api/multimodal.py b/areno/api/multimodal.py index 5c76f9d2..b27a6c0f 100644 --- a/areno/api/multimodal.py +++ b/areno/api/multimodal.py @@ -430,7 +430,7 @@ def _image_processor_from_processor(processor: Any): def _image_token_id(tokenizer: Any, processor: Any) -> int | None: for obj in (processor, tokenizer): - for attr in ("image_token_id", "image_token_index"): + for attr in ("image_token_id", "image_token_index", "special_image_token_id"): value = getattr(obj, attr, None) if isinstance(value, int): return int(value) @@ -443,7 +443,7 @@ def _image_token_id(tokenizer: Any, processor: Any) -> int | None: return int(token_id) convert = getattr(tokenizer, "convert_tokens_to_ids", None) if callable(convert): - for token in ("<|image_pad|>", "<|image|>", ""): + for token in ("<|image_pad|>", "<|image|>", "", "<|endoftext10|>"): token_id = convert(token) if isinstance(token_id, int) and token_id >= 0: return int(token_id) diff --git a/areno/engine/checkpoints/common.py b/areno/engine/checkpoints/common.py index 16358c49..a17a085e 100644 --- a/areno/engine/checkpoints/common.py +++ b/areno/engine/checkpoints/common.py @@ -98,6 +98,16 @@ class MergedColumnSpec: keys: tuple[str, ...] +@dataclass(frozen=True, slots=True) +class PackedSectionColumnSpec: + """One HF tensor whose semantic row sections are TP-sharded separately.""" + + key: str + tensor_attr: str + global_sizes_attr: str + local_sizes_attr: str + + @dataclass(frozen=True, slots=True) class KSharedQKVColumnSpec: """QKV load spec for checkpoints where later layers may share K/V.""" @@ -403,6 +413,8 @@ def save_checkpoint_weights( source_path: str | None, spec: CheckpointSpec, extra_tensors_fn: Callable[[CheckpointTensorStore], None] | None = None, + *, + copy_passthrough: bool = True, ) -> str | None: """Save a tensor-parallel model as a HF sharded safetensors checkpoint.""" @@ -430,7 +442,7 @@ def save_checkpoint_weights( writer.write(tensors, "extra-tensors") tensors.clear() saved_path = writer.finish() - if saved_path is not None and source_path is not None: + if copy_passthrough and saved_path is not None and source_path is not None: copy_source_passthrough_weights( source_path, saved_path, protected_prefix=_protected_prefix_from_top_level(spec.top_level) ) @@ -599,6 +611,9 @@ def load_layer_op( if isinstance(op, MergedColumnSpec): load_merged_column_spec(module, index, prefix, op, rank, world_size) return + if isinstance(op, PackedSectionColumnSpec): + load_packed_section_column_spec(module, index, prefix, op, rank, world_size) + return if isinstance(op, KSharedQKVColumnSpec): load_k_shared_qkv_column_spec(module, index, prefix, op, rank, world_size) return @@ -642,6 +657,9 @@ def save_layer_op( if isinstance(op, SplitColumnSpec): save_split_column_spec(tensors, module, prefix, op) return + if isinstance(op, PackedSectionColumnSpec): + save_packed_section_column_spec(tensors, module, prefix, op) + return if isinstance(op, RangedSplitColumnSpec): save_ranged_split_column_spec(tensors, module, prefix, op) return @@ -751,6 +769,47 @@ def load_merged_column_spec( copy_merged_column_from_index(dst, index, tensor_keys, rank, world_size) +def load_packed_section_column_spec( + module: nn.Module, + index: SafetensorsIndex, + prefix: str, + spec: PackedSectionColumnSpec, + rank: int, + world_size: int, +) -> None: + """Shard each row section of one packed HF tensor independently.""" + + dst = attr_path(module, spec.tensor_attr) + global_sizes = tuple(int(size) for size in attr_path(module, spec.global_sizes_attr)) + local_sizes = tuple(int(size) for size in attr_path(module, spec.local_sizes_attr)) + if len(global_sizes) != len(local_sizes): + raise ValueError("packed-section global and local size counts differ") + ranges = tuple(_shard_range(size, rank, world_size) for size in global_sizes) + expected_local_sizes = tuple(end - start for start, end in ranges) + if local_sizes != expected_local_sizes: + raise ValueError(f"packed-section local sizes {local_sizes} do not match TP shard sizes {expected_local_sizes}") + if dst.shape[0] != sum(local_sizes): + raise ValueError(f"packed-section destination has {dst.shape[0]} rows, expected {sum(local_sizes)}") + + tensor_key = key(spec.key, prefix) + filename = index.weight_map.get(tensor_key) + if filename is None: + raise KeyError(f"missing HF weight {tensor_key}") + with safe_open(index.model_path / filename, framework="pt", device="cpu") as handle: + source = handle.get_slice(tensor_key) + source_shape = tuple(source.get_shape()) + expected_shape = (sum(global_sizes), *dst.shape[1:]) + if source_shape != expected_shape: + raise ValueError(f"checkpoint tensor {tensor_key} has shape {source_shape}, expected {expected_shape}") + source_offset = 0 + destination_offset = 0 + for global_size, local_size, (start, end) in zip(global_sizes, local_sizes, ranges, strict=True): + shard = source[source_offset + start : source_offset + end] + dst[destination_offset : destination_offset + local_size].copy_(shard.to(dtype=dst.dtype)) + source_offset += global_size + destination_offset += local_size + + def load_k_shared_qkv_column_spec( module: nn.Module, index: SafetensorsIndex, prefix: str, spec: KSharedQKVColumnSpec, rank: int, world_size: int ) -> None: @@ -810,6 +869,28 @@ def save_split_column_spec( tensors[key(template, prefix)] = tensor +def save_packed_section_column_spec( + tensors: dict[str, torch.Tensor | None], + module: nn.Module, + prefix: str, + spec: PackedSectionColumnSpec, +) -> None: + """Gather local packed sections into their original single HF tensor.""" + + tensor = attr_path(module, spec.tensor_attr) + global_sizes = tuple(int(size) for size in attr_path(module, spec.global_sizes_attr)) + local_sizes = [int(size) for size in attr_path(module, spec.local_sizes_attr)] + world_size = get_tp_context().world_size + if any( + global_size != local_size * world_size + for global_size, local_size in zip(global_sizes, local_sizes, strict=True) + ): + raise ValueError("packed-section sizes are incompatible with the tensor-parallel world size") + if tensor.shape[0] != sum(local_sizes): + raise ValueError(f"packed-section source has {tensor.shape[0]} rows, expected {sum(local_sizes)}") + tensors[key(spec.key, prefix)] = gather_tensor_parallel_split_column_tensor(tensor, local_sizes) + + def save_ranged_split_column_spec( tensors: dict[str, torch.Tensor | None], module: nn.Module, prefix: str, spec: RangedSplitColumnSpec ) -> None: diff --git a/areno/engine/data/rollout_state.py b/areno/engine/data/rollout_state.py index c378bd7a..1fe71065 100644 --- a/areno/engine/data/rollout_state.py +++ b/areno/engine/data/rollout_state.py @@ -113,6 +113,7 @@ def build_prefill_payload(self) -> dict | None: has_mrope_positions = False feature_mask: list[bool] = [] image_features: list[dict] = [] + image_sequence_modes: list[bool] = [] cu_seqlens = [0] sample_indices: list[int] = [] block_table: list[list[int]] = [] @@ -156,6 +157,7 @@ def build_prefill_payload(self) -> dict | None: mrope_position_parts if has_mrope_positions else None, feature_mask, image_features, + image_sequence_modes, cu_seqlens, sample_indices, block_table, @@ -175,6 +177,7 @@ def build_prefill_payload(self) -> dict | None: chunk_len, ) feature_mask.extend(local_mask) + image_sequence_modes.append(_prompt_has_image(self.prompt_features[seq_id], prompt)) if local_features is not None: image_features.append(local_features) local_mrope_positions = _slice_prompt_mrope_positions( @@ -219,6 +222,7 @@ def build_prefill_payload(self) -> dict | None: mrope_position_parts if has_mrope_positions else None, feature_mask, image_features, + image_sequence_modes, cu_seqlens, sample_indices, block_table, @@ -235,6 +239,7 @@ def _prefill_payload( mrope_position_parts: list[torch.Tensor] | None, feature_mask: list[bool], image_features: list[dict], + image_sequence_modes: list[bool], cu_seqlens: list[int], sample_indices: list[int], block_table: list[list[int]], @@ -256,8 +261,13 @@ def _prefill_payload( "cache_block_offsets": torch.tensor(cache_block_offsets, dtype=torch.long), "recurrent_slots": torch.tensor(recurrent_slots, dtype=torch.long), } - if any(feature_mask) or image_features or mrope_position_parts is not None: - payload["features"] = _prefill_multimodal_features(feature_mask, image_features, mrope_position_parts) + if any(feature_mask) or image_features or any(image_sequence_modes) or mrope_position_parts is not None: + payload["features"] = _prefill_multimodal_features( + feature_mask, + image_features, + mrope_position_parts, + image_sequence_modes, + ) return payload def ensure_decode_blocks(self, seq_ids: list[int], next_positions: list[int]) -> None: @@ -327,6 +337,9 @@ def _slice_prompt_image_features( key in features for key in ( "pixel_values", + "input_image_embeds", + "image_sizes", + "image_attention_mask", "image_grid_thw", "target_sizes", "pixel_values_videos", @@ -366,6 +379,9 @@ def _slice_prompt_image_features( ) for key in ( "pixel_values", + "input_image_embeds", + "image_sizes", + "image_attention_mask", "image_grid_thw", "target_sizes", "num_patches_per_image", @@ -412,8 +428,11 @@ def _prefill_multimodal_features( feature_mask: list[bool], image_features: list[dict], mrope_position_parts: list[torch.Tensor] | None = None, + image_sequence_modes: list[bool] | None = None, ) -> dict: features = {} + if image_sequence_modes is not None and any(image_sequence_modes): + features["image_sequence_mask"] = torch.tensor(image_sequence_modes, dtype=torch.bool) if mrope_position_parts is not None: features["mrope_position_ids"] = torch.cat(mrope_position_parts, dim=1).to(dtype=torch.long) if not image_features: @@ -456,6 +475,18 @@ def _prompt_image_mask(features: dict, prompt: list[int]) -> list[bool]: return [int(token) in values for token in prompt] +def _prompt_has_image(features: dict | None, prompt: list[int]) -> bool: + if features is None: + return False + mask = features.get("image_token_mask") + if mask is not None: + return bool(torch.as_tensor(mask, dtype=torch.bool).any()) + image_token_id = features.get("image_token_id") + if image_token_id is None: + image_token_id = (features.get("modality_token_ids") or {}).get("image") + return image_token_id is not None and any(int(token) == int(image_token_id) for token in prompt) + + def payload_to_infer_meta(payload: dict, device: torch.device) -> InferMeta: """Move a scheduler payload to device and expose it as model metadata.""" diff --git a/areno/engine/layers/attention.py b/areno/engine/layers/attention.py index 0f1b05bb..20ef88cb 100644 --- a/areno/engine/layers/attention.py +++ b/areno/engine/layers/attention.py @@ -32,7 +32,7 @@ class CausalSelfAttention(nn.Module): reduces across ranks to reassemble the full hidden state. """ - def __init__(self, config: ModelConfig, layer_idx: int): + def __init__(self, config: ModelConfig, layer_idx: int, *, rotary_embedding: nn.Module | None = None): super().__init__() ctx = get_tp_context() self.layer_idx = layer_idx @@ -58,7 +58,11 @@ def __init__(self, config: ModelConfig, layer_idx: int): # Row-parallel output projection: input is already sharded along # head dimension, output is all-reduced across ranks. self.o_proj = RowParallelLinear(self.num_heads * self.head_dim, config.hidden_size, bias=False) - self.rope = RotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta) + self.rope = ( + rotary_embedding + if rotary_embedding is not None + else RotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta) + ) # Optional per-head QK normalization (used by some recent models). self.q_norm = RMSNorm(config.head_dim, config.rms_norm_eps) if config.qk_norm else None self.k_norm = RMSNorm(config.head_dim, config.rms_norm_eps) if config.qk_norm else None @@ -93,7 +97,7 @@ def forward( k = self.k_norm(k) # Rotary embedding is applied on the head dim using position-indexed # cos/sin tables; positions are broadcast across heads. - q, k = self.rope(q, k, position_ids) + q, k = self.apply_rotary(q, k, position_ids, train_meta, infer_meta) # Presence of infer_meta selects the paged KV-cache backend; otherwise # we run the training-mode FlashAttention (padded or varlen packed). @@ -101,6 +105,19 @@ def forward( return self.forward_infer(q, k, v, infer_meta) return self.forward_train(q, k, v, train_meta) + def apply_rotary( + self, + q: torch.Tensor, + k: torch.Tensor, + position_ids: torch.Tensor, + train_meta: TrainMeta | None, + infer_meta: InferMeta | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Apply the model's rotary embedding, with a model override hook.""" + + del train_meta, infer_meta + return self.rope(q, k, position_ids) + def forward_train( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, train_meta: TrainMeta | None ) -> torch.Tensor: diff --git a/areno/engine/layers/mlp.py b/areno/engine/layers/mlp.py index 6b785671..65c92209 100644 --- a/areno/engine/layers/mlp.py +++ b/areno/engine/layers/mlp.py @@ -10,7 +10,8 @@ import torch from torch import nn -from areno.accel.ops import areno_silu_and_mul, log_once +from areno.accel.activations import areno_silu_and_mul +from areno.accel.utils import log_once from areno.engine.config import ModelConfig from areno.engine.layers.linear import MergedColumnParallelLinear, RowParallelLinear diff --git a/areno/engine/layers/norm.py b/areno/engine/layers/norm.py index a640a96a..e765773c 100644 --- a/areno/engine/layers/norm.py +++ b/areno/engine/layers/norm.py @@ -13,7 +13,7 @@ from torch import nn from areno.accel import areno_rmsnorm -from areno.accel.ops import can_use_cuda_kernel, log_once, rms_norm_gate_fwd +from areno.accel.utils import can_use_cuda_kernel, log_once from areno.engine.layers.linear import mark_tensor_parallel_parameter @@ -90,8 +90,10 @@ def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: # Reshape last dim into (groups_per_rank, group_width) for the kernel. x = x.view(*shape[:-1], self.groups_per_rank, self.group_width) gate = gate.view(*shape[:-1], self.groups_per_rank, self.group_width) - if rms_norm_gate_fwd is None or not can_use_cuda_kernel(x, "fused group RMSNorm sigmoid gate kernel"): + if not can_use_cuda_kernel(x, "fused group RMSNorm sigmoid gate kernel"): raise RuntimeError("ARENO group RMSNorm sigmoid gate requires the fused CUDA kernel") + from areno.accel.kernels.group_rmsnorm import rms_norm_gate_fwd + log_once("group_rmsnorm_sigmoid_gate", "using fused group RMSNorm sigmoid gate kernel") # Flatten the leading dims into a single batch so the kernel only # sees a 3D (B, groups, width) tensor. diff --git a/areno/engine/runtime/decode_graph.py b/areno/engine/runtime/decode_graph.py index 79f0d0eb..339e29d4 100644 --- a/areno/engine/runtime/decode_graph.py +++ b/areno/engine/runtime/decode_graph.py @@ -81,6 +81,7 @@ def __init__( """Allocate static input buffers and the `InferMeta` baked into capture.""" self.model = model + self.decode_cache_length_limit = getattr(model, "decode_cache_length_limit", None) self.bucket = bucket self.scratch_block = scratch_block self.scratch_recurrent_slot = scratch_recurrent_slot @@ -156,6 +157,7 @@ def replay_tensors( actual = int(input_ids.numel()) if actual > self.bucket: raise ValueError(f"decode payload has {actual} tokens, graph bucket is {self.bucket}") + _validate_decode_cache_length(cache_seqlens, actual, self.decode_cache_length_limit) # Copy the live values into the captured-stable buffers. The graph # was recorded against these buffer addresses so `copy_` here is what @@ -186,3 +188,14 @@ def replay_tensors( self.graph.replay() assert self.logits_shard is not None return self.logits_shard + + +def _validate_decode_cache_length( + cache_seqlens: torch.Tensor, + actual: int, + limit: int | None, +) -> None: + if limit is not None and actual and int(cache_seqlens[:actual].max().item()) >= limit: + raise ValueError( + "cached decode cannot cross the model's rotary-factor boundary; run a full long-context prefill" + ) diff --git a/areno/models/__init__.py b/areno/models/__init__.py index 83540786..cb74c51e 100644 --- a/areno/models/__init__.py +++ b/areno/models/__init__.py @@ -25,6 +25,13 @@ def _register_qwen35() -> None: register_adapter(Qwen35Adapter()) +def _register_phi4mm() -> None: + from areno.models.phi4mm import Phi4MMAdapter + from areno.models.registry import register_adapter + + register_adapter(Phi4MMAdapter()) + + def _register_bailing() -> None: from areno.models.bailing import BailingMoeLinearV2Adapter from areno.models.registry import register_adapter @@ -71,6 +78,7 @@ def _register_olmo2() -> None: "llama": _register_llama, "qwen3": _register_qwen3, "qwen3_5": _register_qwen35, + "phi4mm": _register_phi4mm, "bailing": _register_bailing, "bailing_v3": _register_bailing_v3, "gemma4": _register_gemma4, diff --git a/areno/models/phi4mm/__init__.py b/areno/models/phi4mm/__init__.py new file mode 100644 index 00000000..3cf01600 --- /dev/null +++ b/areno/models/phi4mm/__init__.py @@ -0,0 +1,21 @@ +"""Phi-4-Multimodal language-backbone adapter.""" + +from __future__ import annotations + +from areno.models.phi4mm.model import ( + Phi4MMAdapter, + Phi4MMAttention, + Phi4MMDecoderLayer, + Phi4MMForCausalLM, + Phi4MMLongRoPEScaledRotaryEmbedding, + Phi4MMModel, +) + +__all__ = [ + "Phi4MMAdapter", + "Phi4MMAttention", + "Phi4MMDecoderLayer", + "Phi4MMForCausalLM", + "Phi4MMLongRoPEScaledRotaryEmbedding", + "Phi4MMModel", +] diff --git a/areno/models/phi4mm/checkpoint.py b/areno/models/phi4mm/checkpoint.py new file mode 100644 index 00000000..dae80858 --- /dev/null +++ b/areno/models/phi4mm/checkpoint.py @@ -0,0 +1,296 @@ +"""Strict checkpoint mapping for the supported Phi-4-Multimodal paths.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path + +import torch +from torch import nn + +from areno.engine.checkpoints.common import ( + CheckpointSpec, + CheckpointTensorStore, + LayerSpec, + PackedSectionColumnSpec, + ParallelTensorSpec, + PolicyTensorStore, + ReplicatedTensorSpec, + TopLevelSpec, + copy_merged_column, + gather_tensor_parallel_split_column_tensor, + gather_tensor_parallel_tensor, + load_checkpoint_weights, + rank0_tensor, + save_checkpoint_weights, +) +from areno.engine.checkpoints.io import SafetensorsIndex, _copy_row +from areno.engine.parallel.context import get_tp_context + +TOP_LEVEL_SPEC = TopLevelSpec( + embedding_key="model.embed_tokens.weight", + embedding_attr="model.embed_tokens", + norm_key="model.norm.weight", + norm_attr="model.norm.weight", +) +LAYER_NORM_SPECS = ( + ReplicatedTensorSpec("{prefix}.input_layernorm.weight", "input_layernorm.weight"), + ReplicatedTensorSpec("{prefix}.post_attention_layernorm.weight", "post_attention_layernorm.weight"), +) +QKV_SPEC = PackedSectionColumnSpec( + key="{prefix}.self_attn.qkv_proj.base_layer.weight", + tensor_attr="self_attn.qkv_proj.weight", + global_sizes_attr="self_attn.qkv_proj.out_features", + local_sizes_attr="self_attn.qkv_proj.local_out_features", +) +ATTN_OUT_SPEC = ParallelTensorSpec( + "{prefix}.self_attn.o_proj.base_layer.weight", + "self_attn.o_proj.weight", + 1, +) +GATE_UP_SPEC = PackedSectionColumnSpec( + key="{prefix}.mlp.gate_up_proj.base_layer.weight", + tensor_attr="mlp.gate_up_proj.weight", + global_sizes_attr="mlp.gate_up_proj.out_features", + local_sizes_attr="mlp.gate_up_proj.local_out_features", +) +MLP_DOWN_SPEC = ParallelTensorSpec( + "{prefix}.mlp.down_proj.base_layer.weight", + "mlp.down_proj.weight", + 1, +) +LAYER_SPEC = LayerSpec( + prefix="model.layers.{layer}", + replicated=LAYER_NORM_SPECS, + load_ops=(QKV_SPEC, ATTN_OUT_SPEC, GATE_UP_SPEC, MLP_DOWN_SPEC), + save_ops=(QKV_SPEC, ATTN_OUT_SPEC, GATE_UP_SPEC, MLP_DOWN_SPEC), +) +CHECKPOINT_SPEC = CheckpointSpec(top_level=TOP_LEVEL_SPEC, layer=LAYER_SPEC) + +_LAYER_BASE_SUFFIXES = ( + "input_layernorm.weight", + "post_attention_layernorm.weight", + "self_attn.qkv_proj.base_layer.weight", + "self_attn.o_proj.base_layer.weight", + "mlp.gate_up_proj.base_layer.weight", + "mlp.down_proj.base_layer.weight", +) +_LORA_PATTERN = re.compile( + r"^model\.layers\.(\d+)\." + r"(?:self_attn\.(?:qkv_proj|o_proj)|mlp\.(?:gate_up_proj|down_proj))\." + r"lora_[AB]\.(vision|speech)\.weight$" +) + + +@dataclass(frozen=True, slots=True) +class Phi4MMCheckpointAudit: + total: int + consumed: int + vision_lora_skipped: int + speech_lora_skipped: int + vision_skipped: int + audio_skipped: int + unknown: int + + +def _required_base_keys(num_hidden_layers: int) -> set[str]: + required = {"model.embed_tokens.weight", "model.norm.weight"} + for layer in range(num_hidden_layers): + required.update(f"model.layers.{layer}.{suffix}" for suffix in _LAYER_BASE_SUFFIXES) + return required + + +def audit_phi4mm_checkpoint( + model_path: str | Path, + num_hidden_layers: int, + vision_keys: set[str] | None = None, + vision_lora_keys: set[str] | None = None, +) -> Phi4MMCheckpointAudit: + """Classify every checkpoint key and reject missing or unknown tensors.""" + + index = SafetensorsIndex(model_path, progress=False) + try: + checkpoint_keys = set(index.weight_map) + finally: + index.close() + base_keys = _required_base_keys(num_hidden_layers) + required = base_keys | (vision_keys or set()) | (vision_lora_keys or set()) + missing = sorted(required - checkpoint_keys) + if missing: + preview = ", ".join(missing[:5]) + raise ValueError(f"Phi4MM checkpoint is missing {len(missing)} required base-language tensors: {preview}") + + counts = {"vision_lora": 0, "speech_lora": 0, "vision": 0, "audio": 0} + unknown = [] + for tensor_key in checkpoint_keys - required: + lora_match = _LORA_PATTERN.fullmatch(tensor_key) + if lora_match is not None and int(lora_match.group(1)) < num_hidden_layers: + counts[f"{lora_match.group(2)}_lora"] += 1 + elif tensor_key.startswith("model.embed_tokens_extend.image_embed."): + counts["vision"] += 1 + elif tensor_key.startswith("model.embed_tokens_extend.audio_embed."): + counts["audio"] += 1 + else: + unknown.append(tensor_key) + if unknown: + preview = ", ".join(sorted(unknown)[:5]) + raise ValueError(f"Phi4MM checkpoint contains {len(unknown)} unknown tensors: {preview}") + return Phi4MMCheckpointAudit( + total=len(checkpoint_keys), + consumed=len(required), + vision_lora_skipped=counts["vision_lora"], + speech_lora_skipped=counts["speech_lora"], + vision_skipped=counts["vision"], + audio_skipped=counts["audio"], + unknown=0, + ) + + +def load_phi4mm_weights(model: nn.Module, model_path: str | Path) -> Phi4MMCheckpointAudit: + """Audit and load the supported Phi-4 language and vision tensors.""" + + model.config.validate_tp(get_tp_context().world_size) + vision_keys = _vision_checkpoint_keys(model) + vision_lora_keys = _vision_lora_checkpoint_keys(model) + audit = audit_phi4mm_checkpoint(model_path, len(model.layers), vision_keys, vision_lora_keys) + load_checkpoint_weights(model, str(model_path), CHECKPOINT_SPEC) + if vision_keys: + _load_vision_weights(model, model_path, vision_keys) + if vision_lora_keys: + _load_vision_lora_weights(model, model_path) + if model.lm_head.weight is not model.model.embed_tokens.weight: + raise RuntimeError("Phi4MM embedding and LM head weight tying was lost during checkpoint loading") + return audit + + +def _vision_checkpoint_keys(model: nn.Module) -> set[str]: + extended = getattr(model.model, "embed_tokens_extend", None) + if extended is None: + return set() + return {f"model.embed_tokens_extend.{name}" for name, _ in extended.named_parameters()} + + +def _vision_lora_checkpoint_keys(model: nn.Module) -> set[str]: + return { + f"model.layers.{layer_idx}.{name}" + for layer_idx, layer in enumerate(model.layers) + for name, _ in layer.named_parameters() + if ".lora_A.vision.weight" in name or ".lora_B.vision.weight" in name + } + + +@torch.no_grad() +def _load_vision_weights(model: nn.Module, model_path: str | Path, keys: set[str] | None = None) -> None: + extended = getattr(model.model, "embed_tokens_extend", None) + if extended is None: + return + expected = keys if keys is not None else _vision_checkpoint_keys(model) + index = SafetensorsIndex(model_path) + try: + missing = sorted(expected - set(index.weight_map)) + if missing: + raise KeyError(f"missing Phi4MM vision weight {missing[0]}") + index.prefetch(sorted(expected)) + for name, parameter in extended.named_parameters(): + key = f"model.embed_tokens_extend.{name}" + source = index.get_tensor(key) + if tuple(source.shape) != tuple(parameter.shape): + raise ValueError( + f"checkpoint tensor {key} shape {tuple(source.shape)} does not match {tuple(parameter.shape)}" + ) + parameter.copy_(source.to(device=parameter.device, dtype=parameter.dtype)) + finally: + index.close() + + +@torch.no_grad() +def _load_vision_lora_weights(model: nn.Module, model_path: str | Path) -> None: + context = get_tp_context() + index = SafetensorsIndex(model_path) + try: + for layer_idx, layer in enumerate(model.layers): + prefix = f"model.layers.{layer_idx}" + for name, module, sections in ( + ("self_attn.qkv_proj", layer.self_attn.qkv_proj, layer.self_attn.qkv_proj.out_features), + ("mlp.gate_up_proj", layer.mlp.gate_up_proj, layer.mlp.gate_up_proj.out_features), + ): + lora_a = f"{prefix}.{name}.lora_A.vision.weight" + lora_b = f"{prefix}.{name}.lora_B.vision.weight" + module.lora_A["vision"].weight.copy_(index.get_tensor(lora_a).to(dtype=module.weight.dtype)) + copy_merged_column( + module.lora_B["vision"].weight, + list(index.get_tensor(lora_b).split(tuple(sections), dim=0)), + context.rank, + context.world_size, + ) + for name, module in ( + ("self_attn.o_proj", layer.self_attn.o_proj), + ("mlp.down_proj", layer.mlp.down_proj), + ): + lora_a = f"{prefix}.{name}.lora_A.vision.weight" + lora_b = f"{prefix}.{name}.lora_B.vision.weight" + _copy_row( + module.lora_A["vision"].weight, + index.get_tensor(lora_a), + context.rank, + context.world_size, + ) + module.lora_B["vision"].weight.copy_( + index.get_tensor(lora_b).to(device=module.weight.device, dtype=module.weight.dtype) + ) + finally: + index.close() + + +def save_phi4mm_weights( + model: nn.Module, + output_path: str | Path, + source_path: str | Path | None, +) -> str | None: + """Save Phi-4 language and vision weights in the official HF key layout.""" + + model.config.validate_tp(get_tp_context().world_size) + return save_checkpoint_weights( + model, + str(output_path), + None if source_path is None else str(source_path), + CHECKPOINT_SPEC, + extra_tensors_fn=lambda tensors: _save_vision_weights(tensors, model), + copy_passthrough=False, + ) + + +def _save_vision_weights(tensors: CheckpointTensorStore | PolicyTensorStore, model: nn.Module) -> None: + """Stage replicated vision tensors and TP-aware Vision LoRA tensors.""" + + extended = getattr(model.model, "embed_tokens_extend", None) + if extended is None: + return + for name, parameter in extended.named_parameters(): + tensors[f"model.embed_tokens_extend.{name}"] = rank0_tensor(parameter) + + for layer_idx, layer in enumerate(model.layers): + prefix = f"model.layers.{layer_idx}" + for name, module in ( + ("self_attn.qkv_proj", layer.self_attn.qkv_proj), + ("mlp.gate_up_proj", layer.mlp.gate_up_proj), + ): + if not hasattr(module, "lora_A") or "vision" not in module.lora_A: + continue + tensors[f"{prefix}.{name}.lora_A.vision.weight"] = rank0_tensor(module.lora_A["vision"].weight) + tensors[f"{prefix}.{name}.lora_B.vision.weight"] = gather_tensor_parallel_split_column_tensor( + module.lora_B["vision"].weight, + list(module.local_out_features), + ) + for name, module in ( + ("self_attn.o_proj", layer.self_attn.o_proj), + ("mlp.down_proj", layer.mlp.down_proj), + ): + if not hasattr(module, "lora_A") or "vision" not in module.lora_A: + continue + tensors[f"{prefix}.{name}.lora_A.vision.weight"] = gather_tensor_parallel_tensor( + module.lora_A["vision"].weight, + dim=1, + ) + tensors[f"{prefix}.{name}.lora_B.vision.weight"] = rank0_tensor(module.lora_B["vision"].weight) diff --git a/areno/models/phi4mm/model.py b/areno/models/phi4mm/model.py new file mode 100644 index 00000000..f77431e5 --- /dev/null +++ b/areno/models/phi4mm/model.py @@ -0,0 +1,788 @@ +"""Phi-4-Multimodal language and vision adapter.""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from areno.accel.utils import is_cuda_graph_capturing +from areno.engine.config import ModelConfig, _parse_dtype +from areno.engine.layers.attention import CausalSelfAttention +from areno.engine.layers.linear import MergedColumnParallelLinear, RowParallelLinear, mark_tensor_parallel_parameter +from areno.engine.layers.mlp import GatedMLP +from areno.engine.layers.norm import RMSNorm +from areno.engine.layers.vocab import VocabParallelEmbedding, VocabParallelLMHead +from areno.engine.parallel.collectives import ( + all_reduce, + copy_to_tensor_parallel_region, + gather_from_sequence_parallel_region, + is_sequence_parallel_active, + scatter_to_sequence_parallel_region, + sequence_parallel_region, +) +from areno.engine.runtime.metadata import InferMeta, TrainMeta +from areno.engine.runtime.recompute import checkpoint_layer +from areno.models.base import CausalLMOutput, ModelAdapter +from areno.models.phi4mm.vision import Phi4MMExtendedEmbedding, Phi4MMVisionConfig + +_IMAGE_SPECIAL_TOKEN_ID = 200010 + + +def _phi4mm_vision_config(hf_config: dict[str, Any]) -> dict[str, Any] | None: + embedding = hf_config.get("embd_layer") + if not isinstance(embedding, dict): + return None + image = embedding.get("image_embd_layer") + if not isinstance(image, dict): + return None + required = { + "embedding_cls": "tune_image", + "image_token_compression_cls": "avg_pool_2d", + "projection_cls": "mlp", + "use_hd_transform": True, + "with_learnable_separator": True, + "hd_transform_order": "sub_glb", + } + for key, expected in required.items(): + actual = image.get(key) + if actual != expected: + raise ValueError(f"Phi4MM vision requires embd_layer.image_embd_layer.{key}={expected!r}, got {actual!r}") + config = { + "hidden_size": 1152, + "intermediate_size": 4304, + "num_hidden_layers": 27, + "num_attention_heads": 16, + "num_channels": 3, + "image_size": 448, + "patch_size": 14, + "layer_norm_eps": 1e-6, + "attention_dropout": 0.0, + "hidden_act": "gelu_pytorch_tanh", + "feature_layer": -2, + "crop_size": int(image.get("crop_size", 448)), + "hd_transform_order": str(image["hd_transform_order"]), + } + override = hf_config.get("vision_config") + if isinstance(override, dict): + config.update(override) + return config + + +def _features_by_row(features: dict[str, Any] | list[dict[str, Any] | None], batch: int) -> list[dict[str, Any] | None]: + if isinstance(features, list): + if len(features) != batch: + raise ValueError(f"Phi4MM multimodal features batch mismatch: got {len(features)} rows for batch {batch}") + return features + if not isinstance(features, dict): + raise TypeError("Phi4MM multimodal features must be a dict or batch-aligned list") + if batch == 1: + return [features] + rows = [] + for row_idx in range(batch): + row = {} + for key, value in features.items(): + if isinstance(value, torch.Tensor) and value.ndim > 0 and int(value.shape[0]) == batch: + row[key] = value[row_idx] + elif isinstance(value, list) and len(value) == batch: + row[key] = value[row_idx] + else: + row[key] = value + rows.append(row) + return rows + + +def _feature_tensor( + features: dict[str, Any], key: str, device: torch.device, dtype: torch.dtype | None = None +) -> torch.Tensor | None: + value = features.get(key) + if value is None: + return None + tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + return tensor.to(device=device, dtype=dtype) + + +def _vision_lora_config(config: ModelConfig) -> tuple[int, float, float] | None: + values = (config.hf_text_config or {}).get("vision_lora") + if config.vision_config is None: + return None + if not isinstance(values, dict): + raise ValueError("Phi4MM vision support requires a vision_lora config") + rank = int(values["r"]) + alpha = float(values["lora_alpha"]) + dropout = float(values.get("dp", 0.0)) + if rank <= 0 or alpha <= 0 or not 0.0 <= dropout < 1.0: + raise ValueError("Phi4MM vision_lora requires positive r/alpha and dp in [0, 1)") + return rank, alpha / rank, dropout + + +class _Phi4MMColumnLoRA(MergedColumnParallelLinear): + def __init__(self, in_features: int, out_features: tuple[int, ...], config: ModelConfig): + super().__init__(in_features, out_features, bias=False) + lora = _vision_lora_config(config) + self.vision_lora_scale = 0.0 + self.vision_lora_dropout = 0.0 + self.vision_lora_mask: torch.Tensor | None = None + self.lora_A = nn.ModuleDict() + self.lora_B = nn.ModuleDict() + if lora is not None: + rank, self.vision_lora_scale, self.vision_lora_dropout = lora + self.lora_A["vision"] = nn.Linear(in_features, rank, bias=False) + self.lora_B["vision"] = nn.Linear(rank, sum(self.local_out_features), bias=False) + mark_tensor_parallel_parameter( + self.lora_A["vision"].weight, False, sequence_parallel=False, tp_grad_allreduce=True + ) + mark_tensor_parallel_parameter(self.lora_B["vision"].weight, True, sequence_parallel=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output = super().forward(x) + if self.vision_lora_mask is None or "vision" not in self.lora_A: + return output + full_input = ( + gather_from_sequence_parallel_region(x) + if is_sequence_parallel_active() + else copy_to_tensor_parallel_region(x) + ) + dropped = F.dropout(full_input, p=self.vision_lora_dropout, training=self.training) + delta = self.lora_B["vision"](self.lora_A["vision"](dropped)) * self.vision_lora_scale + return output + delta * self.vision_lora_mask.to(device=delta.device, dtype=delta.dtype).unsqueeze(-1) + + +class _Phi4MMRowLoRA(RowParallelLinear): + def __init__(self, in_features: int, out_features: int, config: ModelConfig): + super().__init__(in_features, out_features, bias=False) + lora = _vision_lora_config(config) + self.vision_lora_scale = 0.0 + self.vision_lora_dropout = 0.0 + self.vision_lora_mask: torch.Tensor | None = None + self.lora_A = nn.ModuleDict() + self.lora_B = nn.ModuleDict() + if lora is not None: + rank, self.vision_lora_scale, self.vision_lora_dropout = lora + self.lora_A["vision"] = nn.Linear(self.local_in_features, rank, bias=False) + self.lora_B["vision"] = nn.Linear(rank, out_features, bias=False) + mark_tensor_parallel_parameter(self.lora_A["vision"].weight, True, sequence_parallel=True) + mark_tensor_parallel_parameter( + self.lora_B["vision"].weight, False, sequence_parallel=False, tp_grad_allreduce=True + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output = super().forward(x) + if self.vision_lora_mask is None or "vision" not in self.lora_A: + return output + dropped = F.dropout(x, p=self.vision_lora_dropout, training=self.training) + latent = all_reduce(self.lora_A["vision"](dropped)) + delta = self.lora_B["vision"](latent) * self.vision_lora_scale + delta = delta * self.vision_lora_mask.to(device=delta.device, dtype=delta.dtype).unsqueeze(-1) + if is_sequence_parallel_active(): + delta = scatter_to_sequence_parallel_region(delta) + return output + delta + + +def _require_bool(hf_config: dict[str, Any], key: str, expected: bool) -> None: + value = bool(hf_config.get(key, expected)) + if value is not expected: + raise ValueError(f"Phi4MM requires {key}={expected}, got {value}") + + +def _validated_longrope(hf_config: dict[str, Any], rotary_dim: int) -> dict[str, Any]: + rope = hf_config.get("rope_scaling") + if not isinstance(rope, dict): + raise ValueError("Phi4MM requires a rope_scaling mapping") + if set(rope) != {"type", "short_factor", "long_factor"}: + raise ValueError("Phi4MM rope_scaling must contain exactly: type, short_factor, long_factor") + if rope["type"] != "longrope": + raise ValueError(f"Phi4MM only supports rope_scaling.type='longrope', got {rope['type']!r}") + + expected_factors = rotary_dim // 2 + normalized = {"type": "longrope"} + for key in ("short_factor", "long_factor"): + factors = rope[key] + if not isinstance(factors, list) or len(factors) != expected_factors: + raise ValueError(f"Phi4MM {key} must contain {expected_factors} values") + if any(isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0 for value in factors): + raise ValueError(f"Phi4MM {key} values must be positive numbers") + normalized[key] = tuple(float(value) for value in factors) + return normalized + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + first, second = x.chunk(2, dim=-1) + return torch.cat((-second, first), dim=-1) + + +class Phi4MMLongRoPEScaledRotaryEmbedding(nn.Module): + """Official Phi-4 partial LongRoPE math without per-layer position caches.""" + + def __init__(self, config: ModelConfig): + super().__init__() + if config.hf_text_config is None: + raise ValueError("Phi4MM requires the validated HF text config") + self.dim = int(config.head_dim * config.partial_rotary_factor) + if self.dim <= 0 or self.dim % 2: + raise ValueError("Phi4MM rotary dimension must be a positive even number") + rope_scaling = config.hf_text_config["rope_scaling"] + expected_factors = self.dim // 2 + short_factor = rope_scaling["short_factor"] + long_factor = rope_scaling["long_factor"] + if len(short_factor) != expected_factors or len(long_factor) != expected_factors: + raise ValueError(f"Phi4MM short_factor and long_factor must contain {expected_factors} values") + + self.max_position_embeddings = int(config.max_position_embeddings) + self.original_max_position_embeddings = int(config.hf_text_config["original_max_position_embeddings"]) + inv_freq_shape = torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim + base_freq = config.rope_theta**inv_freq_shape + self.register_buffer( + "short_inv_freq", 1.0 / (torch.tensor(short_factor, dtype=torch.float32) * base_freq), persistent=False + ) + self.register_buffer( + "long_inv_freq", 1.0 / (torch.tensor(long_factor, dtype=torch.float32) * base_freq), persistent=False + ) + scale = self.max_position_embeddings / self.original_max_position_embeddings + self.scaling_factor = ( + 1.0 if scale <= 1.0 else math.sqrt(1.0 + math.log(scale) / math.log(self.original_max_position_embeddings)) + ) + + def _apply(self, fn): + super()._apply(fn) + # Long-context phases must remain FP32 even when model weights are cast. + self.short_inv_freq = self.short_inv_freq.float() + self.long_inv_freq = self.long_inv_freq.float() + return self + + @torch.no_grad() + def cos_sin( + self, + x: torch.Tensor, + position_ids: torch.Tensor, + sequence_length: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if sequence_length is None: + sequence_length = int(torch.max(position_ids).item()) + 1 + inv_freq = ( + self.long_inv_freq if sequence_length > self.original_max_position_embeddings else self.short_inv_freq + ) + expanded_inv_freq = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + expanded_positions = position_ids[:, None, :].float() + device_type = x.device.type if x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (expanded_inv_freq @ expanded_positions).transpose(1, 2) + embedding = torch.cat((freqs, freqs), dim=-1) + cos = embedding.cos() * self.scaling_factor + sin = embedding.sin() * self.scaling_factor + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + position_ids: torch.Tensor, + sequence_length: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + cos, sin = self.cos_sin(q, position_ids, sequence_length) + cos = cos.unsqueeze(2) + sin = sin.unsqueeze(2) + q_rot, q_pass = q[..., : self.dim], q[..., self.dim :] + k_rot, k_pass = k[..., : self.dim], k[..., self.dim :] + q_embed = torch.cat((q_rot * cos + _rotate_half(q_rot) * sin, q_pass), dim=-1) + k_embed = torch.cat((k_rot * cos + _rotate_half(k_rot) * sin, k_pass), dim=-1) + return q_embed, k_embed + + +def _phi4mm_longrope_sequence_length( + position_ids: torch.Tensor, + train_meta: TrainMeta | None, + infer_meta: InferMeta | None, + original_max_position_embeddings: int, +) -> int: + if infer_meta is not None and infer_meta.mode == "decode": + if infer_meta.cache_seqlens is None: + raise ValueError("Phi4MM decode requires cache_seqlens for LongRoPE selection") + sequence_length = int(infer_meta.cache_seqlens.max().item()) + 1 + if sequence_length > original_max_position_embeddings: + raise ValueError( + "Phi4MM cached decode cannot cross the LongRoPE boundary because cached keys may use short factors; " + "run a full long-context prefill" + ) + return sequence_length + + if infer_meta is not None: + sequence_length = int(position_ids.max().item()) + 1 + if sequence_length > original_max_position_embeddings: + if infer_meta.cu_seqlens is None: + raise ValueError("Phi4MM prefill requires cu_seqlens for LongRoPE boundary validation") + starts = infer_meta.cu_seqlens[:-1].to(dtype=torch.long) + flat_positions = position_ids.reshape(-1) + if bool(torch.any(flat_positions[starts] != 0)): + raise ValueError( + "Phi4MM chunked prefill cannot cross the LongRoPE boundary because cached keys use short factors; " + "increase the prefill token budget and run a full prefill" + ) + return sequence_length + + if train_meta is not None and train_meta.max_seqlen is not None: + return int(train_meta.max_seqlen) + return int(position_ids.shape[-1]) + + +class Phi4MMAttention(CausalSelfAttention): + """AReno GQA attention with a Phi-owned rotary implementation.""" + + def __init__(self, config: ModelConfig, layer_idx: int): + if config.qk_norm: + raise ValueError("Phi4MMAttention requires qk_norm=False") + super().__init__(config, layer_idx, rotary_embedding=Phi4MMLongRoPEScaledRotaryEmbedding(config)) + self.qkv_proj = _Phi4MMColumnLoRA( + config.hidden_size, + ( + config.num_attention_heads * config.head_dim, + config.num_key_value_heads * config.head_dim, + config.num_key_value_heads * config.head_dim, + ), + config, + ) + self.o_proj = _Phi4MMRowLoRA(config.num_attention_heads * config.head_dim, config.hidden_size, config) + + def apply_rotary( + self, + q: torch.Tensor, + k: torch.Tensor, + position_ids: torch.Tensor, + train_meta: TrainMeta | None, + infer_meta: InferMeta | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if infer_meta is not None and infer_meta.mode == "decode" and is_cuda_graph_capturing(q): + # DecodeGraph validates its dynamic cache lengths before replay. + # Capture itself always records the supported short-factor path. + sequence_length = self.rope.original_max_position_embeddings + else: + sequence_length = _phi4mm_longrope_sequence_length( + position_ids, + train_meta, + infer_meta, + self.rope.original_max_position_embeddings, + ) + return self.rope(q, k, position_ids, sequence_length) + + +class Phi4MMDecoderLayer(nn.Module): + """Phi-4 pre-norm decoder block composed from AReno shared layers.""" + + def __init__(self, config: ModelConfig, layer_idx: int): + super().__init__() + self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps) + self.self_attn = Phi4MMAttention(config, layer_idx) + self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps) + self.mlp = GatedMLP(config) + if config.vision_config is not None: + self.mlp.gate_up_proj = _Phi4MMColumnLoRA( + config.hidden_size, (config.intermediate_size, config.intermediate_size), config + ) + self.mlp.down_proj = _Phi4MMRowLoRA(config.intermediate_size, config.hidden_size, config) + + def set_vision_lora_mask(self, mask: torch.Tensor | None) -> None: + self.self_attn.qkv_proj.vision_lora_mask = mask + self.self_attn.o_proj.vision_lora_mask = mask + if hasattr(self.mlp.gate_up_proj, "vision_lora_mask"): + self.mlp.gate_up_proj.vision_lora_mask = mask + self.mlp.down_proj.vision_lora_mask = mask + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + train_meta: TrainMeta | None = None, + infer_meta: InferMeta | None = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = residual + self.self_attn(hidden_states, position_ids, train_meta, infer_meta) + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + return residual + self.mlp(hidden_states) + + +class Phi4MMModel(nn.Module): + """Phi-4 transformer body with an optional native vision embedding path.""" + + def __init__(self, config: ModelConfig): + super().__init__() + self.config = config + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size, dtype=config.dtype) + self.embed_tokens_extend = ( + Phi4MMExtendedEmbedding( + Phi4MMVisionConfig.from_dict(config.vision_config), config.hidden_size, config.dtype + ) + if config.vision_config is not None + else None + ) + if self.embed_tokens_extend is not None: + for parameter in self.embed_tokens_extend.parameters(): + mark_tensor_parallel_parameter(parameter, False, sequence_parallel=False, tp_grad_allreduce=True) + self.layers = nn.ModuleList([Phi4MMDecoderLayer(config, index) for index in range(config.num_hidden_layers)]) + self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps) + self.register_buffer("vision_lora_slots", torch.empty(0, dtype=torch.bool), persistent=False) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None = None, + train_meta: TrainMeta | None = None, + infer_meta: InferMeta | None = None, + features: dict[str, Any] | list[dict[str, Any] | None] | None = None, + ) -> torch.Tensor: + if position_ids is None: + position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0).expand_as(input_ids) + vision_lora_mask = self._vision_lora_mask(input_ids, features, train_meta, infer_meta) + for layer in self.layers: + layer.set_vision_lora_mask(vision_lora_mask) + hidden_states = self.embed_tokens(input_ids) + hidden_states = self._apply_multimodal_features(hidden_states, input_ids, features) + use_sequence_parallel = bool(train_meta is not None and train_meta.sequence_parallel) + if use_sequence_parallel: + hidden_states = scatter_to_sequence_parallel_region(hidden_states) + with sequence_parallel_region(use_sequence_parallel): + for layer in self.layers: + hidden_states = checkpoint_layer( + layer, + hidden_states, + position_ids, + train_meta, + infer_meta, + train_meta=train_meta, + infer_meta=infer_meta, + ) + return self.norm(hidden_states) + + def _vision_lora_mask( + self, + input_ids: torch.Tensor, + features: dict[str, Any] | list[dict[str, Any] | None] | None, + train_meta: TrainMeta | None, + infer_meta: InferMeta | None, + ) -> torch.Tensor | None: + if self.embed_tokens_extend is None: + return None + if infer_meta is not None and infer_meta.mode == "decode": + if infer_meta.recurrent_slots is None or self.vision_lora_slots.numel() == 0: + raise ValueError("Phi4MM vision decode requires recurrent modality slots") + return self.vision_lora_slots.index_select(0, infer_meta.recurrent_slots).view_as(input_ids) + image_mask = self._image_token_mask(input_ids, features) + explicit_modes = None + if isinstance(features, dict) and features.get("image_sequence_mask") is not None: + explicit_modes = torch.as_tensor( + features["image_sequence_mask"], device=input_ids.device, dtype=torch.bool + ).reshape(-1) + sequence_offsets = None + if infer_meta is not None and infer_meta.cu_seqlens is not None: + sequence_offsets = infer_meta.cu_seqlens + elif train_meta is not None and train_meta.cu_seqlens is not None: + sequence_offsets = train_meta.cu_seqlens + if sequence_offsets is None: + row_modes = explicit_modes if explicit_modes is not None else image_mask.any(dim=1) + if int(row_modes.numel()) != int(input_ids.shape[0]): + raise ValueError("Phi4MM image_sequence_mask must contain one value per input row") + mask = row_modes[:, None].expand_as(input_ids) + else: + flat = image_mask.reshape(-1) + mask = torch.zeros_like(flat) + modes = [] + offsets = sequence_offsets.detach().to(device="cpu", dtype=torch.long).tolist() + sequence_count = len(offsets) - 1 + if explicit_modes is not None and int(explicit_modes.numel()) != sequence_count: + raise ValueError("Phi4MM image_sequence_mask must contain one value per packed sequence") + for sequence_idx, (start, end) in enumerate(zip(offsets[:-1], offsets[1:], strict=True)): + mode = bool(explicit_modes[sequence_idx]) if explicit_modes is not None else bool(flat[start:end].any()) + modes.append(mode) + mask[start:end] = mode + mask = mask.view_as(input_ids) + if infer_meta is not None and infer_meta.recurrent_slots is not None and self.vision_lora_slots.numel() > 0: + mode_tensor = torch.tensor(modes, device=self.vision_lora_slots.device, dtype=torch.bool) + self.vision_lora_slots.index_copy_(0, infer_meta.recurrent_slots, mode_tensor) + return mask + + def _image_token_mask( + self, + input_ids: torch.Tensor, + features: dict[str, Any] | list[dict[str, Any] | None] | None, + ) -> torch.Tensor: + if isinstance(features, dict) and features.get("image_token_mask") is not None: + return torch.as_tensor(features["image_token_mask"], device=input_ids.device, dtype=torch.bool).view_as( + input_ids + ) + return input_ids == int(self.config.image_token_id or _IMAGE_SPECIAL_TOKEN_ID) + + @torch._dynamo.disable + def _apply_multimodal_features( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + features: dict[str, Any] | list[dict[str, Any] | None] | None, + ) -> torch.Tensor: + if features is None: + return hidden_states + if self.embed_tokens_extend is None: + raise ValueError("Phi4MM image features require a configured vision tower") + rows = _features_by_row(features, int(input_ids.shape[0])) + output = hidden_states.clone() + for row_idx, row in enumerate(rows): + if row is None: + continue + image_embeds = self._project_image_feature_rows(row, hidden_states.device) + if image_embeds is None: + continue + mask = row.get("image_token_mask") + if mask is None: + token_id = int(row.get("image_token_id", self.config.image_token_id or _IMAGE_SPECIAL_TOKEN_ID)) + mask = input_ids[row_idx] == token_id + else: + mask = torch.as_tensor(mask, device=input_ids.device, dtype=torch.bool).reshape(-1) + if mask.shape != input_ids[row_idx].shape: + raise ValueError("Phi4MM image_token_mask must match the input token row") + if int(mask.sum().item()) != int(image_embeds.shape[0]): + raise ValueError( + "Phi4MM image token count does not match projected embeddings: " + f"tokens={int(mask.sum().item())} embeds={int(image_embeds.shape[0])}" + ) + output[row_idx, mask] = image_embeds.to(device=output.device, dtype=output.dtype) + return output + + def _project_image_feature_rows(self, features: dict[str, Any], device: torch.device) -> torch.Tensor | None: + rows = features.get("image_feature_rows") + if rows is not None: + pieces = [self._project_image_feature(dict(row), device) for row in rows if row is not None] + pieces = [piece for piece in pieces if piece is not None] + return torch.cat(pieces, dim=0) if pieces else None + return self._project_image_feature(features, device) + + def _project_image_feature(self, features: dict[str, Any], device: torch.device) -> torch.Tensor | None: + existing = _feature_tensor(features, "image_embeds", device, self.config.dtype) + if existing is not None: + return existing + pixels = _feature_tensor(features, "input_image_embeds", device, self.config.dtype) + if pixels is None: + return None + sizes = _feature_tensor(features, "image_sizes", device, torch.long) + mask = _feature_tensor(features, "image_attention_mask", device, torch.bool) + if sizes is None or mask is None: + raise ValueError("Phi4MM processor output requires image_sizes and image_attention_mask") + image_embeds = self.embed_tokens_extend.image_embed(pixels, sizes, mask) + offset = int(features.get("image_token_offset", 0) or 0) + count = features.get("image_token_count") + if count is not None: + return image_embeds[offset : offset + int(count)] + return image_embeds[offset:] + + +class Phi4MMForCausalLM(nn.Module): + """Text-only Phi-4 causal LM with a truly tied vocab-parallel head.""" + + def __init__(self, config: ModelConfig): + super().__init__() + if not config.tie_word_embeddings: + raise ValueError("Phi4MMForCausalLM requires tied word embeddings") + self.config = config + self.decode_cache_length_limit = int(config.hf_text_config["original_max_position_embeddings"]) + self.model = Phi4MMModel(config) + self.lm_head = VocabParallelLMHead(config.hidden_size, config.vocab_size, dtype=config.dtype) + self._tie_word_embeddings() + + def _tie_word_embeddings(self) -> None: + embedding = self.model.embed_tokens + if (self.lm_head.vocab_start, self.lm_head.vocab_end) != (embedding.vocab_start, embedding.vocab_end): + raise ValueError("Phi4MM embedding and LM head use different TP vocabulary ranges") + if self.lm_head.weight.shape != embedding.weight.shape: + raise ValueError("Phi4MM embedding and LM head local weight shapes differ") + self.lm_head.weight = embedding.weight + + @property + def layers(self) -> nn.ModuleList: + """Expose decoder layers to the shared checkpoint machinery.""" + + return self.model.layers + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None = None, + train_meta: TrainMeta | None = None, + infer_meta: InferMeta | None = None, + features: dict[str, Any] | list[dict[str, Any] | None] | None = None, + ) -> CausalLMOutput: + use_sequence_parallel = bool(train_meta is not None and train_meta.sequence_parallel) + with sequence_parallel_region(use_sequence_parallel): + hidden_states = self.model(input_ids, position_ids, train_meta, infer_meta, features) + logits_shard = self.lm_head(hidden_states) + return CausalLMOutput(logits_shard=logits_shard, hidden_states=hidden_states) + + def set_kv_caches( + self, kv_caches: list[tuple[torch.Tensor, torch.Tensor]], *, num_slots: int | None = None + ) -> None: + """Bind one paged KV-cache pair to each decoder layer.""" + if len(kv_caches) != len(self.layers): + raise ValueError(f"expected {len(self.layers)} layer caches, got {len(kv_caches)}") + for layer, (k_cache, v_cache) in zip(self.layers, kv_caches, strict=True): + layer.self_attn.set_kv_cache(k_cache, v_cache) + slot_count = int(num_slots) if num_slots is not None else (int(kv_caches[0][0].shape[0]) if kv_caches else 0) + self.model.vision_lora_slots = torch.zeros(slot_count, device=next(self.parameters()).device, dtype=torch.bool) + + @torch.no_grad() + def reset_recurrent_cache_slots(self, slots: torch.Tensor) -> None: + if self.model.vision_lora_slots.numel() > 0: + self.model.vision_lora_slots.index_fill_(0, slots, False) + + @torch.no_grad() + def prepare_infer_weights(self) -> None: + return None + + @torch.no_grad() + def clear_infer_weights(self) -> None: + return None + + @torch.no_grad() + def offload_train_weights(self) -> None: + return None + + @torch.no_grad() + def onload_train_weights(self, device: torch.device) -> None: + del device + return None + + @torch.no_grad() + def finalize_router_expert_bias(self, tp_group, dp_group) -> None: + del tp_group, dp_group + return None + + def allocate_kv_caches( + self, num_blocks: int, block_size: int, device: torch.device + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Allocate the standard paged GQA cache layout for every layer.""" + caches = [] + for layer in self.layers: + attention = layer.self_attn + shape = (num_blocks, block_size, attention.local_kv_heads, attention.head_dim) + caches.append( + ( + torch.empty(shape, device=device, dtype=self.config.dtype), + torch.empty(shape, device=device, dtype=self.config.dtype), + ) + ) + return caches + + def clear_kv_caches(self) -> None: + for layer in self.layers: + layer.self_attn.clear_kv_cache() + + @torch.no_grad() + def reset_kv_caches(self) -> None: + return None + + @torch.no_grad() + def offload_kv_caches(self) -> None: + for layer in self.layers: + attention = layer.self_attn + if attention.k_cache.numel() > 0: + attention.k_cache = attention.k_cache.to(device="cpu") + if attention.v_cache.numel() > 0: + attention.v_cache = attention.v_cache.to(device="cpu") + attention.infer_backend = None + + @torch.no_grad() + def onload_kv_caches(self, device: torch.device) -> bool: + found = False + for layer in self.layers: + attention = layer.self_attn + if attention.k_cache.numel() > 0: + found = True + if attention.k_cache.device != device: + attention.k_cache = attention.k_cache.to(device=device) + if attention.v_cache.numel() > 0 and attention.v_cache.device != device: + attention.v_cache = attention.v_cache.to(device=device) + return found + + +class Phi4MMAdapter(ModelAdapter): + """Translate the official Phi-4-Multimodal config into AReno semantics.""" + + name = "phi4mm" + + def match_hf_config(self, hf_config: dict[str, Any]) -> bool: + return str(hf_config.get("model_type", "")).lower() == self.name + + def config_from_hf(self, hf_config: dict[str, Any]) -> ModelConfig: + hidden_size = int(hf_config["hidden_size"]) + num_attention_heads = int(hf_config["num_attention_heads"]) + if hidden_size % num_attention_heads != 0: + raise ValueError("Phi4MM hidden_size must be divisible by num_attention_heads") + head_dim = hidden_size // num_attention_heads + partial_rotary_factor = float(hf_config.get("partial_rotary_factor", 1.0)) + if not 0.0 < partial_rotary_factor <= 1.0: + raise ValueError("Phi4MM partial_rotary_factor must be in (0, 1]") + rotary_dim = int(head_dim * partial_rotary_factor) + if rotary_dim <= 0 or rotary_dim % 2 != 0: + raise ValueError("Phi4MM rotary dimension must be a positive even number") + + if str(hf_config.get("hidden_act", "silu")) != "silu": + raise ValueError("Phi4MM language backbone requires hidden_act='silu'") + _require_bool(hf_config, "attention_bias", False) + _require_bool(hf_config, "mlp_bias", False) + _require_bool(hf_config, "lm_head_bias", False) + _require_bool(hf_config, "tie_word_embeddings", True) + + original_max_position_embeddings = int(hf_config.get("original_max_position_embeddings", 4096)) + max_position_embeddings = int(hf_config.get("max_position_embeddings", original_max_position_embeddings)) + if original_max_position_embeddings <= 0 or max_position_embeddings < original_max_position_embeddings: + raise ValueError("Phi4MM max_position_embeddings must be at least original_max_position_embeddings > 0") + rope_scaling = _validated_longrope(hf_config, rotary_dim) + + # Preserve the validated LongRoPE fields for the Phi-specific rotary implementation. + text_config = dict(hf_config) + text_config["rope_scaling"] = rope_scaling + text_config["original_max_position_embeddings"] = original_max_position_embeddings + vision_config = _phi4mm_vision_config(hf_config) + + return ModelConfig( + model_type=self.name, + checkpoint_prefix="model", + vocab_size=int(hf_config["vocab_size"]), + pad_token_id=int(hf_config.get("pad_token_id", 0) or 0), + hidden_size=hidden_size, + intermediate_size=int(hf_config["intermediate_size"]), + num_hidden_layers=int(hf_config["num_hidden_layers"]), + num_attention_heads=num_attention_heads, + num_key_value_heads=int(hf_config.get("num_key_value_heads", num_attention_heads)), + head_dim=head_dim, + rms_norm_eps=float(hf_config.get("rms_norm_eps", 1e-5)), + rope_theta=float(hf_config.get("rope_theta", 10_000.0)), + max_position_embeddings=max_position_embeddings, + tie_word_embeddings=True, + qkv_bias=False, + qk_norm=False, + dtype=_parse_dtype(hf_config.get("torch_dtype") or hf_config.get("dtype")), + hidden_act="silu", + sliding_window=hf_config.get("sliding_window"), + partial_rotary_factor=partial_rotary_factor, + sequence_parallel=bool(hf_config.get("sequence_parallel", True)), + hf_text_config=text_config, + vision_config=vision_config, + image_token_id=_IMAGE_SPECIAL_TOKEN_ID if vision_config is not None else None, + ) + + def build(self, config: ModelConfig) -> nn.Module: + if config.model_type != self.name: + raise ValueError(f"Phi4MMAdapter cannot build model_type={config.model_type!r}") + return Phi4MMForCausalLM(config) + + def load_weights(self, model: nn.Module, model_path: str | Path) -> None: + from areno.models.phi4mm.checkpoint import load_phi4mm_weights + + load_phi4mm_weights(model, model_path) + + def save_weights(self, model: nn.Module, output_path: str | Path, source_path: str | Path | None) -> str | None: + from areno.models.phi4mm.checkpoint import save_phi4mm_weights + + return save_phi4mm_weights(model, output_path, source_path) diff --git a/areno/models/phi4mm/vision.py b/areno/models/phi4mm/vision.py new file mode 100644 index 00000000..0cf10f9e --- /dev/null +++ b/areno/models/phi4mm/vision.py @@ -0,0 +1,253 @@ +"""Native Phi-4-Multimodal SigLIP vision tower and HD projector.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + + +@dataclass(frozen=True, slots=True) +class Phi4MMVisionConfig: + hidden_size: int = 1152 + intermediate_size: int = 4304 + num_hidden_layers: int = 27 + num_attention_heads: int = 16 + num_channels: int = 3 + image_size: int = 448 + patch_size: int = 14 + layer_norm_eps: float = 1e-6 + attention_dropout: float = 0.0 + hidden_act: str = "gelu_pytorch_tanh" + feature_layer: int = -2 + crop_size: int = 448 + hd_transform_order: str = "sub_glb" + + @classmethod + def from_dict(cls, values: dict[str, Any]) -> Phi4MMVisionConfig: + return cls(**{name: values[name] for name in cls.__dataclass_fields__ if name in values}) + + +class Phi4MMVisionEmbeddings(nn.Module): + def __init__(self, config: Phi4MMVisionConfig, dtype: torch.dtype): + super().__init__() + self.patch_size = config.patch_size + self.num_patches_per_side = config.image_size // config.patch_size + self.patch_embedding = nn.Conv2d( + config.num_channels, + config.hidden_size, + kernel_size=config.patch_size, + stride=config.patch_size, + dtype=dtype, + ) + self.position_embedding = nn.Embedding( + self.num_patches_per_side**2, + config.hidden_size, + dtype=dtype, + ) + + def forward(self, pixel_values: torch.Tensor, patch_attention_mask: torch.Tensor) -> torch.Tensor: + embeddings = ( + self.patch_embedding(pixel_values.to(dtype=self.patch_embedding.weight.dtype)).flatten(2).transpose(1, 2) + ) + batch, patch_height, patch_width = patch_attention_mask.shape + boundaries = torch.arange( + 1 / self.num_patches_per_side, + 1.0, + 1 / self.num_patches_per_side, + device="cpu", + ) + position_ids = torch.zeros((batch, patch_height * patch_width), dtype=torch.long, device="cpu") + for row, mask in enumerate(patch_attention_mask.detach().to(device="cpu", dtype=torch.bool)): + valid_height = int(mask[:, 0].sum().item()) + valid_width = int(mask[0].sum().item()) + if valid_height <= 0 or valid_width <= 0: + raise ValueError("Phi4MM image attention mask must contain at least one valid patch") + height_coords = torch.arange(valid_height, dtype=torch.float32) / valid_height + width_coords = torch.arange(valid_width, dtype=torch.float32) / valid_width + height_buckets = torch.bucketize(height_coords, boundaries, right=True) + width_buckets = torch.bucketize(width_coords, boundaries, right=True) + ids = (height_buckets[:, None] * self.num_patches_per_side + width_buckets).flatten() + position_ids[row, mask.reshape(-1)] = ids + return embeddings + self.position_embedding(position_ids.to(self.position_embedding.weight.device)) + + +class Phi4MMVisionAttention(nn.Module): + def __init__(self, config: Phi4MMVisionConfig, dtype: torch.dtype): + super().__init__() + if config.hidden_size % config.num_attention_heads: + raise ValueError("Phi4MM vision hidden_size must be divisible by num_attention_heads") + self.num_heads = config.num_attention_heads + self.head_dim = config.hidden_size // config.num_attention_heads + self.scale = self.head_dim**-0.5 + self.dropout = config.attention_dropout + self.k_proj = nn.Linear(config.hidden_size, config.hidden_size, dtype=dtype) + self.v_proj = nn.Linear(config.hidden_size, config.hidden_size, dtype=dtype) + self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, dtype=dtype) + self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, dtype=dtype) + + def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor: + batch, seqlen, hidden_size = hidden_states.shape + query = self.q_proj(hidden_states).view(batch, seqlen, self.num_heads, self.head_dim).transpose(1, 2) + key = self.k_proj(hidden_states).view(batch, seqlen, self.num_heads, self.head_dim).transpose(1, 2) + value = self.v_proj(hidden_states).view(batch, seqlen, self.num_heads, self.head_dim).transpose(1, 2) + scores = torch.matmul(query, key.transpose(-2, -1)) * self.scale + if attention_mask is not None: + scores = scores.masked_fill(~attention_mask[:, None, None, :], torch.finfo(scores.dtype).min) + probabilities = F.softmax(scores, dim=-1, dtype=torch.float32).to(dtype=query.dtype) + probabilities = F.dropout(probabilities, p=self.dropout, training=self.training) + output = torch.matmul(probabilities, value).transpose(1, 2).reshape(batch, seqlen, hidden_size) + return self.out_proj(output) + + +class Phi4MMVisionMLP(nn.Module): + def __init__(self, config: Phi4MMVisionConfig, dtype: torch.dtype): + super().__init__() + if config.hidden_act != "gelu_pytorch_tanh": + raise ValueError(f"unsupported Phi4MM vision activation {config.hidden_act!r}") + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, dtype=dtype) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, dtype=dtype) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.fc2(F.gelu(self.fc1(hidden_states), approximate="tanh")) + + +class Phi4MMVisionEncoderLayer(nn.Module): + def __init__(self, config: Phi4MMVisionConfig, dtype: torch.dtype): + super().__init__() + self.self_attn = Phi4MMVisionAttention(config, dtype) + self.layer_norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, dtype=dtype) + self.mlp = Phi4MMVisionMLP(config, dtype) + self.layer_norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, dtype=dtype) + + def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor: + hidden_states = hidden_states + self.self_attn(self.layer_norm1(hidden_states), attention_mask) + return hidden_states + self.mlp(self.layer_norm2(hidden_states)) + + +class Phi4MMVisionEncoder(nn.Module): + def __init__(self, config: Phi4MMVisionConfig, dtype: torch.dtype): + super().__init__() + self.layers = nn.ModuleList([Phi4MMVisionEncoderLayer(config, dtype) for _ in range(config.num_hidden_layers)]) + + +class Phi4MMVisionPoolingHead(nn.Module): + def __init__(self, config: Phi4MMVisionConfig, dtype: torch.dtype): + super().__init__() + self.probe = nn.Parameter(torch.empty(1, 1, config.hidden_size, dtype=dtype)) + self.attention = nn.MultiheadAttention( + config.hidden_size, + config.num_attention_heads, + batch_first=True, + dtype=dtype, + ) + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, dtype=dtype) + self.mlp = Phi4MMVisionMLP(config, dtype) + + +class Phi4MMVisionTransformer(nn.Module): + """SigLIP NaViT module with checkpoint-compatible parameter names.""" + + def __init__(self, config: Phi4MMVisionConfig, dtype: torch.dtype): + super().__init__() + self.config = config + self.embeddings = Phi4MMVisionEmbeddings(config, dtype) + self.encoder = Phi4MMVisionEncoder(config, dtype) + self.post_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, dtype=dtype) + self.head = Phi4MMVisionPoolingHead(config, dtype) + + def patch_features(self, pixel_values: torch.Tensor, patch_attention_mask: torch.Tensor) -> torch.Tensor: + hidden_states = self.embeddings(pixel_values, patch_attention_mask) + flat_mask = patch_attention_mask.reshape(patch_attention_mask.shape[0], -1).to(dtype=torch.bool) + attention_mask = None if bool(flat_mask.all()) else flat_mask + hidden_states_by_layer = [hidden_states] + for layer in self.encoder.layers: + hidden_states = layer(hidden_states, attention_mask) + hidden_states_by_layer.append(hidden_states) + return hidden_states_by_layer[self.config.feature_layer] + + +class Phi4MMImageEmbedding(nn.Module): + """Project processor crops into the language model's expanded image slots.""" + + def __init__(self, config: Phi4MMVisionConfig, language_hidden_size: int, dtype: torch.dtype): + super().__init__() + self.config = config + self.img_processor = Phi4MMVisionTransformer(config, dtype) + self.glb_GN = nn.Parameter(torch.zeros(1, 1, config.hidden_size, dtype=dtype)) + self.sub_GN = nn.Parameter(torch.zeros(1, 1, 1, config.hidden_size, dtype=dtype)) + self.img_projection = nn.Sequential( + nn.Linear(config.hidden_size, language_hidden_size, dtype=dtype), + nn.GELU(), + nn.Linear(language_hidden_size, language_hidden_size, dtype=dtype), + ) + + def forward( + self, + input_image_embeds: torch.Tensor, + image_sizes: torch.Tensor, + image_attention_mask: torch.Tensor, + ) -> torch.Tensor: + if input_image_embeds.ndim != 5: + raise ValueError("Phi4MM input_image_embeds must have shape (images, crops, 3, H, W)") + if image_sizes.numel() == 0 or image_attention_mask.numel() == 0: + raise ValueError("Phi4MM vision inputs require image_sizes and image_attention_mask") + image_count, max_crops = input_image_embeds.shape[:2] + masks = image_attention_mask.to(device=input_image_embeds.device, dtype=torch.bool) + features = self.img_processor.patch_features(input_image_embeds.flatten(0, 1), masks.flatten(0, 1)) + side = math.isqrt(int(features.shape[1])) + if side * side != int(features.shape[1]): + raise ValueError("Phi4MM vision patch count must form a square grid") + features = features.view(image_count * max_crops, side, side, self.config.hidden_size) + features = F.avg_pool2d(features.permute(0, 3, 1, 2), kernel_size=2, stride=2).permute(0, 2, 3, 1) + pooled_side = int(features.shape[1]) + features = features.reshape(image_count, max_crops, pooled_side, pooled_side, self.config.hidden_size) + + projected = [] + sizes = image_sizes.reshape(-1, 2).detach().to(device="cpu", dtype=torch.long) + for image_idx, (height_value, width_value) in enumerate(sizes.tolist()): + crop_rows = int(height_value) // self.config.crop_size + crop_cols = int(width_value) // self.config.crop_size + local_crop_count = crop_rows * crop_cols + if local_crop_count + 1 > max_crops: + raise ValueError("Phi4MM image size requires more crops than input_image_embeds provides") + + global_image = features[image_idx, 0:1] + global_separators = self.sub_GN.expand(1, pooled_side, 1, -1) + global_image = torch.cat((global_image, global_separators), dim=2).reshape(1, -1, self.config.hidden_size) + + local_image = features[image_idx, 1 : local_crop_count + 1] + local_image = ( + local_image.reshape(crop_rows, crop_cols, pooled_side, pooled_side, self.config.hidden_size) + .permute(0, 2, 1, 3, 4) + .reshape(1, crop_rows * pooled_side, crop_cols * pooled_side, self.config.hidden_size) + ) + local_mask = masks[image_idx, 1 : local_crop_count + 1, 0::2, 0::2] + local_mask = ( + local_mask.reshape(crop_rows, crop_cols, pooled_side, pooled_side) + .permute(0, 2, 1, 3) + .reshape(crop_rows * pooled_side, crop_cols * pooled_side) + ) + useful_height = int(local_mask[:, 0].sum().item()) + useful_width = int(local_mask[0].sum().item()) + local_image = local_image[:, :useful_height, :useful_width] + local_separators = self.sub_GN.expand(1, useful_height, 1, -1) + local_image = torch.cat((local_image, local_separators), dim=2).reshape(1, -1, self.config.hidden_size) + + if self.config.hd_transform_order != "sub_glb": + raise ValueError(f"unsupported Phi4MM hd_transform_order {self.config.hd_transform_order!r}") + image_features = torch.cat((local_image, self.glb_GN, global_image), dim=1) + projected.append(self.img_projection(image_features)) + return torch.cat(projected, dim=1).squeeze(0) + + +class Phi4MMExtendedEmbedding(nn.Module): + """Checkpoint-compatible container for Phi multimodal embedding modules.""" + + def __init__(self, config: Phi4MMVisionConfig, language_hidden_size: int, dtype: torch.dtype): + super().__init__() + self.image_embed = Phi4MMImageEmbedding(config, language_hidden_size, dtype) diff --git a/tests/test_phi4mm_adapter_cpu.py b/tests/test_phi4mm_adapter_cpu.py new file mode 100644 index 00000000..c2fb7e0f --- /dev/null +++ b/tests/test_phi4mm_adapter_cpu.py @@ -0,0 +1,448 @@ +from __future__ import annotations + +import json +import math +import subprocess +import sys + +import pytest +import torch +import torch.nn.functional as F + +import areno.models +from areno.engine.config import ModelConfig, OptimizerConfig +from areno.engine.layers import mlp, norm, vocab +from areno.engine.modeling import build_optimizer +from areno.engine.parallel.collectives import is_sequence_parallel_active +from areno.engine.parallel.context import TPContext, get_tp_context, set_tp_context +from areno.engine.runtime.decode_graph import _validate_decode_cache_length +from areno.engine.runtime.metadata import InferMeta, TrainMeta +from areno.models import registry +from areno.models.phi4mm import Phi4MMAdapter, Phi4MMForCausalLM +from areno.models.phi4mm.model import ( + Phi4MMLongRoPEScaledRotaryEmbedding, + _phi4mm_longrope_sequence_length, +) + + +@pytest.fixture(autouse=True) +def _isolate_tp_context(): + previous_context = get_tp_context() + set_tp_context(TPContext(rank=0, world_size=1, device=torch.device("cpu"), group=None)) + try: + yield + finally: + set_tp_context(previous_context) + + +def _phi4mm_config() -> dict: + return { + "model_type": "phi4mm", + "vocab_size": 200064, + "hidden_size": 3072, + "intermediate_size": 8192, + "num_hidden_layers": 32, + "num_attention_heads": 24, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-5, + "rope_theta": 10_000.0, + "max_position_embeddings": 131072, + "original_max_position_embeddings": 4096, + "partial_rotary_factor": 0.75, + "rope_scaling": { + "type": "longrope", + "short_factor": [1.0] * 48, + "long_factor": [float(index + 1) for index in range(48)], + }, + "sliding_window": 262144, + "hidden_act": "silu", + "attention_bias": False, + "mlp_bias": False, + "lm_head_bias": False, + "tie_word_embeddings": True, + "pad_token_id": 199999, + "torch_dtype": "bfloat16", + } + + +def _tiny_model_config() -> ModelConfig: + return ModelConfig( + model_type="phi4mm", + vocab_size=32, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=8, + num_key_value_heads=4, + head_dim=8, + rms_norm_eps=1e-5, + rope_theta=10_000.0, + max_position_embeddings=64, + tie_word_embeddings=True, + qkv_bias=False, + qk_norm=False, + dtype=torch.float32, + hidden_act="silu", + partial_rotary_factor=0.75, + sequence_parallel=False, + attn_backend="native", + hf_text_config={ + "original_max_position_embeddings": 32, + "rope_scaling": { + "type": "longrope", + "short_factor": (1.0, 1.0, 1.0), + "long_factor": (1.0, 2.0, 3.0), + }, + }, + ) + + +def test_phi4mm_import_does_not_require_triton(): + script = """ +import sys +sys.modules['triton'] = None +import areno.models.phi4mm +assert 'areno.accel.kernels.fused_moe' not in sys.modules +assert 'areno.accel.kernels.group_rmsnorm' not in sys.modules +assert 'areno.accel.kernels.seg_la' not in sys.modules +""" + subprocess.run([sys.executable, "-c", script], check=True) + + +@pytest.fixture +def cpu_reference_kernels(monkeypatch): + def embedding(input_ids, weight, vocab_start, vocab_end): + local_ids = input_ids - vocab_start + local_mask = (input_ids >= vocab_start) & (input_ids < vocab_end) + safe_ids = local_ids.masked_fill(~local_mask, 0) + return F.embedding(safe_ids, weight) * local_mask.unsqueeze(-1) + + def rms_norm(x, weight, eps): + normalized = x.float() * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps) + return normalized.to(dtype=x.dtype) * weight.to(dtype=x.dtype) + + def silu_and_mul(x): + gate, up = x.chunk(2, dim=-1) + return F.silu(gate) * up + + monkeypatch.setattr(vocab, "areno_vocab_embedding", embedding) + monkeypatch.setattr(norm, "_areno_rmsnorm_no_compile", rms_norm) + monkeypatch.setattr(mlp, "_areno_silu_and_mul_no_compile", silu_and_mul) + + +def test_phi4mm_config_translation_matches_official_language_backbone(): + config = Phi4MMAdapter().config_from_hf(_phi4mm_config()) + + assert config.model_type == "phi4mm" + assert config.vocab_size == 200064 + assert config.hidden_size == 3072 + assert config.intermediate_size == 8192 + assert config.num_hidden_layers == 32 + assert config.num_attention_heads == 24 + assert config.num_key_value_heads == 8 + assert config.head_dim == 128 + assert config.partial_rotary_factor == 0.75 + assert config.qk_norm is False + assert config.qkv_bias is False + assert config.tie_word_embeddings is True + assert config.dtype == torch.bfloat16 + assert config.hf_text_config is not None + assert config.hf_text_config["original_max_position_embeddings"] == 4096 + assert config.hf_text_config["rope_scaling"]["short_factor"] == (1.0,) * 48 + + +@pytest.mark.parametrize( + ("update", "message"), + [ + ({"tie_word_embeddings": False}, "tie_word_embeddings=True"), + ({"attention_bias": True}, "attention_bias=False"), + ({"hidden_act": "gelu"}, "hidden_act='silu'"), + ({"rope_scaling": {"type": "linear", "short_factor": [1.0] * 48, "long_factor": [1.0] * 48}}, "longrope"), + ({"rope_scaling": {"type": "longrope", "short_factor": [1.0] * 47, "long_factor": [1.0] * 48}}, "48 values"), + ], +) +def test_phi4mm_config_rejects_unsupported_language_semantics(update, message): + hf_config = _phi4mm_config() + hf_config.update(update) + + with pytest.raises(ValueError, match=message): + Phi4MMAdapter().config_from_hf(hf_config) + + +def test_phi4mm_registry_resolves_config(tmp_path, monkeypatch): + (tmp_path / "config.json").write_text(json.dumps(_phi4mm_config()), encoding="utf-8") + monkeypatch.setattr(registry, "_PLUGINS_LOADED", False) + monkeypatch.setattr(areno.models, "_REGISTERED_GROUPS", set()) + monkeypatch.setattr(registry, "_ADAPTERS", {}) + + config = registry.config_from_hf(tmp_path) + + assert config.model_type == "phi4mm" + assert isinstance(registry.adapter_from_hf(tmp_path), Phi4MMAdapter) + + +def test_phi4mm_tp_validation_rejects_non_divisible_kv_heads(): + config = Phi4MMAdapter().config_from_hf(_phi4mm_config()) + + config.validate_tp(1) + config.validate_tp(2) + config.validate_tp(4) + config.validate_tp(8) + with pytest.raises(ValueError, match="num_key_value_heads must be divisible"): + config.validate_tp(3) + with pytest.raises(ValueError, match="num_key_value_heads must be divisible"): + config.validate_tp(6) + + +def test_phi4mm_model_construction_has_expected_text_layers(): + config = _tiny_model_config() + model = Phi4MMAdapter().build(config) + + assert isinstance(model, Phi4MMForCausalLM) + assert len(model.model.layers) == 2 + assert model.model.embed_tokens.weight.shape == (32, 64) + assert model.lm_head.weight.shape == (32, 64) + assert model.model.norm.eps == 1e-5 + for layer in model.model.layers: + assert layer.input_layernorm.eps == 1e-5 + assert layer.post_attention_layernorm.eps == 1e-5 + assert layer.self_attn.qkv_proj.out_features == (64, 32, 32) + assert layer.self_attn.qkv_proj.local_out_features == [64, 32, 32] + assert layer.self_attn.o_proj.weight.shape == (64, 64) + assert layer.mlp.gate_up_proj.out_features == (128, 128) + assert layer.mlp.gate_up_proj.weight.shape == (256, 64) + assert layer.mlp.down_proj.weight.shape == (64, 128) + + +def test_phi4mm_projection_biases_and_qk_norm_are_disabled(): + model = Phi4MMAdapter().build(_tiny_model_config()) + + assert not hasattr(model.lm_head, "bias") + for layer in model.model.layers: + assert layer.self_attn.qkv_proj.bias is None + assert layer.self_attn.o_proj.bias is None + assert layer.self_attn.q_norm is None + assert layer.self_attn.k_norm is None + assert layer.mlp.gate_up_proj.bias is None + assert layer.mlp.down_proj.bias is None + + +def test_phi4mm_embedding_and_lm_head_share_one_optimizer_parameter(): + model = Phi4MMAdapter().build(_tiny_model_config()) + + assert model.lm_head.weight is model.model.embed_tokens.weight + parameter_ids = [id(parameter) for parameter in model.parameters()] + assert len(parameter_ids) == len(set(parameter_ids)) + + optimizer = build_optimizer( + model.parameters(), + OptimizerConfig(), + type("Context", (), {"dp_rank": 0, "dp_size": 1, "dp_group": None})(), + ) + optimizer_parameter_ids = [id(parameter) for parameter in optimizer.model_params] + assert len(optimizer_parameter_ids) == len(set(optimizer_parameter_ids)) + assert optimizer_parameter_ids.count(id(model.model.embed_tokens.weight)) == 1 + + +def test_phi4mm_text_forward_shapes_and_causal_prefix(cpu_reference_kernels): + del cpu_reference_kernels + torch.manual_seed(0) + model = Phi4MMAdapter().build(_tiny_model_config()).eval() + input_ids = torch.tensor([[1, 2, 3], [1, 2, 4]]) + + output = model(input_ids) + + assert output.hidden_states is not None + assert output.logits_shard is not None + assert output.hidden_states.shape == (2, 3, 64) + assert output.logits_shard.shape == (2, 3, 32) + assert torch.isfinite(output.hidden_states).all() + assert torch.isfinite(output.logits_shard).all() + torch.testing.assert_close(output.logits_shard[0, :2], output.logits_shard[1, :2]) + + +def test_phi4mm_lm_head_runs_inside_sequence_parallel_region(cpu_reference_kernels, monkeypatch): + del cpu_reference_kernels + model = Phi4MMAdapter().build(_tiny_model_config()).eval() + original_forward = model.lm_head.forward + sequence_parallel_states = [] + + def record_sequence_parallel_state(hidden_states): + sequence_parallel_states.append(is_sequence_parallel_active()) + return original_forward(hidden_states) + + monkeypatch.setattr(model.lm_head, "forward", record_sequence_parallel_state) + model(torch.tensor([[1, 2, 3]]), train_meta=TrainMeta(sequence_parallel=True)) + + assert sequence_parallel_states == [True] + + +def test_phi4mm_kv_cache_lifecycle(): + model = Phi4MMAdapter().build(_tiny_model_config()) + caches = model.allocate_kv_caches(num_blocks=3, block_size=4, device=torch.device("cpu")) + + assert len(caches) == len(model.layers) + assert caches[0][0].shape == (3, 4, 4, 8) + assert caches[0][0].dtype == model.config.dtype + + model.set_kv_caches(caches) + assert model.layers[0].self_attn.k_cache is caches[0][0] + assert model.onload_kv_caches(torch.device("cpu")) + + model.offload_kv_caches() + assert model.layers[0].self_attn.infer_backend is None + model.clear_kv_caches() + assert model.layers[0].self_attn.k_cache.numel() == 0 + assert not model.onload_kv_caches(torch.device("cpu")) + + with pytest.raises(ValueError, match="expected 2 layer caches"): + model.set_kv_caches(caches[:1]) + + +def _official_longrope_reference( + x: torch.Tensor, + position_ids: torch.Tensor, + dim: int, + base: float, + factors: tuple[float, ...], + max_position_embeddings: int, + original_max_position_embeddings: int, +) -> tuple[torch.Tensor, torch.Tensor]: + ext_factors = torch.tensor(factors, dtype=torch.float32, device=x.device) + inv_freq_shape = torch.arange(0, dim, 2, dtype=torch.int64, device=x.device).float() / dim + inv_freq = 1.0 / (ext_factors * base**inv_freq_shape) + inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) + embedding = torch.cat((freqs, freqs), dim=-1) + scale = max_position_embeddings / original_max_position_embeddings + scaling_factor = ( + 1.0 if scale <= 1.0 else math.sqrt(1.0 + math.log(scale) / math.log(original_max_position_embeddings)) + ) + return (embedding.cos() * scaling_factor).to(x.dtype), (embedding.sin() * scaling_factor).to(x.dtype) + + +def test_phi4mm_official_config_builds_partial_longrope_without_position_caches(): + config = Phi4MMAdapter().config_from_hf(_phi4mm_config()) + + rope = Phi4MMLongRoPEScaledRotaryEmbedding(config) + + assert rope.dim == 96 + assert config.head_dim - rope.dim == 32 + assert rope.short_inv_freq.shape == (48,) + assert rope.long_inv_freq.shape == (48,) + assert all("cached" not in name for name, _ in rope.named_buffers()) + + +def test_phi4mm_longrope_keeps_inverse_frequencies_in_fp32_when_model_is_cast(): + rope = Phi4MMLongRoPEScaledRotaryEmbedding(_tiny_model_config()).to(dtype=torch.bfloat16) + + assert rope.short_inv_freq.dtype == torch.float32 + assert rope.long_inv_freq.dtype == torch.float32 + + +@pytest.mark.parametrize( + ("sequence_length", "positions", "factor_key"), + [ + (32, [0, 1, 7, 31], "short_factor"), + (64, [0, 1, 31, 32, 63], "long_factor"), + ], +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_phi4mm_longrope_cos_sin_matches_official_reference(sequence_length, positions, factor_key, dtype): + config = _tiny_model_config() + rope = Phi4MMLongRoPEScaledRotaryEmbedding(config) + x = torch.zeros(1, len(positions), 1, config.head_dim, dtype=dtype) + position_ids = torch.tensor([positions]) + + actual_cos, actual_sin = rope.cos_sin(x, position_ids, sequence_length) + expected_cos, expected_sin = _official_longrope_reference( + x, + position_ids, + rope.dim, + config.rope_theta, + config.hf_text_config["rope_scaling"][factor_key], + config.max_position_embeddings, + config.hf_text_config["original_max_position_embeddings"], + ) + + torch.testing.assert_close(actual_cos, expected_cos, rtol=0, atol=0) + torch.testing.assert_close(actual_sin, expected_sin, rtol=0, atol=0) + + +def test_phi4mm_longrope_preserves_non_rotary_head_dimensions_and_applies_scale(): + config = _tiny_model_config() + rope = Phi4MMLongRoPEScaledRotaryEmbedding(config) + q = torch.randn(1, 3, 2, config.head_dim) + k = torch.randn(1, 3, 1, config.head_dim) + position_ids = torch.tensor([[0, 1, 2]]) + + rotated_q, rotated_k = rope(q, k, position_ids, sequence_length=3) + cos, sin = rope.cos_sin(q, torch.tensor([[0]]), sequence_length=3) + + torch.testing.assert_close(rotated_q[..., rope.dim :], q[..., rope.dim :]) + torch.testing.assert_close(rotated_k[..., rope.dim :], k[..., rope.dim :]) + assert cos[0, 0, 0].item() == pytest.approx(rope.scaling_factor) + assert sin[0, 0, 0].item() == 0.0 + + +def test_phi4mm_longrope_full_long_prefill_selects_long_factors(): + attention = Phi4MMAdapter().build(_tiny_model_config()).model.layers[0].self_attn + positions = torch.arange(40).unsqueeze(0) + infer_meta = InferMeta(mode="prefill", cu_seqlens=torch.tensor([0, 40], dtype=torch.int32), max_seqlen=40) + q = torch.randn(1, 40, attention.local_heads, attention.head_dim) + k = torch.randn(1, 40, attention.local_kv_heads, attention.head_dim) + + sequence_length = _phi4mm_longrope_sequence_length(positions, None, infer_meta, 32) + actual_q, actual_k = attention.apply_rotary(q, k, positions, None, infer_meta) + expected_q, expected_k = attention.rope(q, k, positions, sequence_length=40) + + assert sequence_length == 40 + torch.testing.assert_close(actual_q, expected_q) + torch.testing.assert_close(actual_k, expected_k) + + +def test_phi4mm_longrope_rejects_chunked_prefill_crossing_boundary(): + positions = torch.arange(28, 40).unsqueeze(0) + infer_meta = InferMeta(mode="prefill", cu_seqlens=torch.tensor([0, 12], dtype=torch.int32), max_seqlen=12) + + with pytest.raises(ValueError, match="chunked prefill cannot cross"): + _phi4mm_longrope_sequence_length(positions, None, infer_meta, 32) + + +def test_phi4mm_longrope_rejects_cached_decode_crossing_boundary(): + below_boundary = InferMeta(mode="decode", cache_seqlens=torch.tensor([31], dtype=torch.int32)) + crossing_boundary = InferMeta(mode="decode", cache_seqlens=torch.tensor([32], dtype=torch.int32)) + + assert _phi4mm_longrope_sequence_length(torch.tensor([[31]]), None, below_boundary, 32) == 32 + with pytest.raises(ValueError, match="cached decode cannot cross"): + _phi4mm_longrope_sequence_length(torch.tensor([[32]]), None, crossing_boundary, 32) + + +def test_phi4mm_longrope_decode_graph_replay_enforces_same_boundary(): + _validate_decode_cache_length(torch.tensor([31, 99], dtype=torch.int32), actual=1, limit=32) + with pytest.raises(ValueError, match="rotary-factor boundary"): + _validate_decode_cache_length(torch.tensor([32], dtype=torch.int32), actual=1, limit=32) + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_phi4mm_tp_construction_uses_compatible_local_shards(tp_size): + old_context = get_tp_context() + try: + set_tp_context(TPContext(rank=0, world_size=tp_size, device=torch.device("cpu"), group=None)) + config = _tiny_model_config() + config.validate_tp(tp_size) + model = Phi4MMAdapter().build(config) + finally: + set_tp_context(old_context) + + attention = model.model.layers[0].self_attn + assert attention.local_heads == 8 // tp_size + assert attention.local_kv_heads == 4 // tp_size + assert attention.qkv_proj.local_out_features == [64 // tp_size, 32 // tp_size, 32 // tp_size] + assert attention.o_proj.weight.shape == (64, 64 // tp_size) + assert model.model.layers[0].mlp.gate_up_proj.weight.shape == (256 // tp_size, 64) + assert model.model.layers[0].mlp.down_proj.weight.shape == (64, 128 // tp_size) + assert model.model.embed_tokens.weight.shape == (32 // tp_size, 64) + assert model.lm_head.weight.shape == (32 // tp_size, 64) + assert model.lm_head.weight is model.model.embed_tokens.weight diff --git a/tests/test_phi4mm_checkpoint_cpu.py b/tests/test_phi4mm_checkpoint_cpu.py new file mode 100644 index 00000000..456547af --- /dev/null +++ b/tests/test_phi4mm_checkpoint_cpu.py @@ -0,0 +1,275 @@ +from __future__ import annotations + +import json + +import pytest +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from areno.engine.checkpoints.common import load_packed_section_column_spec, save_packed_section_column_spec +from areno.engine.checkpoints.io import PolicyTensorStore, SafetensorsIndex +from areno.engine.config import ModelConfig +from areno.engine.parallel.context import TPContext, get_tp_context, set_tp_context +from areno.models.phi4mm import Phi4MMAdapter +from areno.models.phi4mm.checkpoint import QKV_SPEC, audit_phi4mm_checkpoint + + +@pytest.fixture(autouse=True) +def _isolate_tp_context(): + previous_context = get_tp_context() + set_tp_context(TPContext(rank=0, world_size=1, device=torch.device("cpu"), group=None)) + try: + yield + finally: + set_tp_context(previous_context) + + +def _tiny_config() -> ModelConfig: + return ModelConfig( + model_type="phi4mm", + vocab_size=32, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=8, + num_key_value_heads=4, + head_dim=8, + rms_norm_eps=1e-5, + rope_theta=10_000.0, + max_position_embeddings=64, + tie_word_embeddings=True, + qkv_bias=False, + qk_norm=False, + dtype=torch.float32, + hidden_act="silu", + partial_rotary_factor=0.75, + sequence_parallel=False, + attn_backend="native", + hf_text_config={ + "original_max_position_embeddings": 32, + "rope_scaling": { + "type": "longrope", + "short_factor": (1.0, 1.0, 1.0), + "long_factor": (1.0, 2.0, 3.0), + }, + }, + ) + + +def _row_values(rows: int, columns: int, base: float) -> torch.Tensor: + return (base + torch.arange(rows, dtype=torch.float32)).unsqueeze(1).expand(rows, columns).clone() + + +def _column_values(rows: int, columns: int, base: float) -> torch.Tensor: + return (base + torch.arange(columns, dtype=torch.float32)).unsqueeze(0).expand(rows, columns).clone() + + +def _synthetic_weights(*, skipped: bool = False) -> dict[str, torch.Tensor]: + config = _tiny_config() + tensors = { + "model.embed_tokens.weight": torch.arange(config.vocab_size * config.hidden_size, dtype=torch.float32).view( + config.vocab_size, config.hidden_size + ), + "model.norm.weight": torch.arange(config.hidden_size, dtype=torch.float32) + 10, + } + for layer in range(config.num_hidden_layers): + prefix = f"model.layers.{layer}" + offset = layer * 10_000 + q = _row_values(64, 64, 1_000 + offset) + k = _row_values(32, 64, 2_000 + offset) + v = _row_values(32, 64, 3_000 + offset) + gate = _row_values(128, 64, 4_000 + offset) + up = _row_values(128, 64, 5_000 + offset) + tensors.update( + { + f"{prefix}.input_layernorm.weight": torch.arange(64, dtype=torch.float32) + 20 + offset, + f"{prefix}.post_attention_layernorm.weight": torch.arange(64, dtype=torch.float32) + 30 + offset, + f"{prefix}.self_attn.qkv_proj.base_layer.weight": torch.cat((q, k, v)), + f"{prefix}.self_attn.o_proj.base_layer.weight": _column_values(64, 64, 6_000 + offset), + f"{prefix}.mlp.gate_up_proj.base_layer.weight": torch.cat((gate, up)), + f"{prefix}.mlp.down_proj.base_layer.weight": _column_values(64, 128, 7_000 + offset), + } + ) + if skipped: + tensors.update( + { + "model.layers.0.self_attn.qkv_proj.lora_A.vision.weight": torch.ones(1), + "model.layers.0.self_attn.qkv_proj.lora_B.speech.weight": torch.ones(1), + "model.embed_tokens_extend.image_embed.img_projection.weight": torch.ones(1), + "model.embed_tokens_extend.audio_embed.audio_projection.weight": torch.ones(1), + } + ) + return tensors + + +def _write_checkpoint(path, tensors: dict[str, torch.Tensor]) -> None: + path.mkdir() + save_file(tensors, path / "model.safetensors") + (path / "config.json").write_text(json.dumps({"model_type": "phi4mm"}), encoding="utf-8") + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_phi4mm_checkpoint_loads_each_packed_section_independently(tmp_path, monkeypatch, tp_size): + monkeypatch.setenv("ARENO_CKPT_PROGRESS", "0") + tensors = _synthetic_weights() + checkpoint = tmp_path / "source" + _write_checkpoint(checkpoint, tensors) + old_context = get_tp_context() + try: + for rank in range(tp_size): + set_tp_context(TPContext(rank=rank, world_size=tp_size, device=torch.device("cpu"), group=None)) + model = Phi4MMAdapter().build(_tiny_config()) + Phi4MMAdapter().load_weights(model, checkpoint) + + layer = model.model.layers[0] + q, k, v = tensors["model.layers.0.self_attn.qkv_proj.base_layer.weight"].split((64, 32, 32)) + gate, up = tensors["model.layers.0.mlp.gate_up_proj.base_layer.weight"].split((128, 128)) + expected_qkv = torch.cat((q.chunk(tp_size)[rank], k.chunk(tp_size)[rank], v.chunk(tp_size)[rank])) + expected_gate_up = torch.cat((gate.chunk(tp_size)[rank], up.chunk(tp_size)[rank])) + + torch.testing.assert_close(layer.self_attn.qkv_proj.weight, expected_qkv) + torch.testing.assert_close(layer.mlp.gate_up_proj.weight, expected_gate_up) + torch.testing.assert_close( + layer.self_attn.o_proj.weight, + tensors["model.layers.0.self_attn.o_proj.base_layer.weight"].chunk(tp_size, dim=1)[rank], + ) + torch.testing.assert_close( + layer.mlp.down_proj.weight, + tensors["model.layers.0.mlp.down_proj.base_layer.weight"].chunk(tp_size, dim=1)[rank], + ) + torch.testing.assert_close( + model.model.embed_tokens.weight, + tensors["model.embed_tokens.weight"].chunk(tp_size)[rank], + ) + torch.testing.assert_close(layer.input_layernorm.weight, tensors["model.layers.0.input_layernorm.weight"]) + torch.testing.assert_close( + layer.post_attention_layernorm.weight, + tensors["model.layers.0.post_attention_layernorm.weight"], + ) + torch.testing.assert_close(model.model.norm.weight, tensors["model.norm.weight"]) + assert model.lm_head.weight is model.model.embed_tokens.weight + finally: + set_tp_context(old_context) + + +def test_phi4mm_checkpoint_audit_accepts_only_documented_skips(tmp_path): + checkpoint = tmp_path / "source" + _write_checkpoint(checkpoint, _synthetic_weights(skipped=True)) + + audit = audit_phi4mm_checkpoint(checkpoint, num_hidden_layers=2) + + assert audit.total == 18 + assert audit.consumed == 14 + assert audit.vision_lora_skipped == 1 + assert audit.speech_lora_skipped == 1 + assert audit.vision_skipped == 1 + assert audit.audio_skipped == 1 + assert audit.unknown == 0 + + +def test_phi4mm_checkpoint_audit_rejects_unknown_base_key(tmp_path): + tensors = _synthetic_weights() + tensors["model.layers.0.self_attn.foo.weight"] = torch.ones(1) + checkpoint = tmp_path / "source" + _write_checkpoint(checkpoint, tensors) + + with pytest.raises(ValueError, match="unknown tensors.*self_attn.foo.weight"): + audit_phi4mm_checkpoint(checkpoint, num_hidden_layers=2) + + +def test_phi4mm_checkpoint_audit_rejects_missing_required_key(tmp_path): + tensors = _synthetic_weights() + del tensors["model.layers.0.self_attn.qkv_proj.base_layer.weight"] + checkpoint = tmp_path / "source" + _write_checkpoint(checkpoint, tensors) + + with pytest.raises(ValueError, match="missing 1 required.*qkv_proj.base_layer.weight"): + audit_phi4mm_checkpoint(checkpoint, num_hidden_layers=2) + + +def test_phi4mm_checkpoint_rejects_wrong_packed_shape(tmp_path, monkeypatch): + monkeypatch.setenv("ARENO_CKPT_PROGRESS", "0") + tensors = _synthetic_weights() + tensors["model.layers.0.self_attn.qkv_proj.base_layer.weight"] = torch.zeros(127, 64) + checkpoint = tmp_path / "source" + _write_checkpoint(checkpoint, tensors) + + with pytest.raises(ValueError, match=r"shape \(127, 64\), expected \(128, 64\)"): + Phi4MMAdapter().load_weights(Phi4MMAdapter().build(_tiny_config()), checkpoint) + + +def test_packed_section_loader_rejects_non_divisible_sections(tmp_path): + checkpoint = tmp_path / "source" + _write_checkpoint(checkpoint, _synthetic_weights()) + model = Phi4MMAdapter().build(_tiny_config()) + index = SafetensorsIndex(checkpoint, progress=False) + try: + with pytest.raises(ValueError, match="cannot shard size 64 across 3 ranks"): + load_packed_section_column_spec(model.model.layers[0], index, "model.layers.0", QKV_SPEC, 0, 3) + finally: + index.close() + + +@pytest.mark.parametrize("tp_size", [2, 4]) +def test_packed_section_save_layout_is_inverse_of_section_sharding(tp_size): + tensors = _synthetic_weights() + full = tensors["model.layers.0.self_attn.qkv_proj.base_layer.weight"] + q, k, v = full.split((64, 32, 32)) + old_context = get_tp_context() + contributions = [] + try: + for rank in range(tp_size): + set_tp_context(TPContext(rank=rank, world_size=tp_size, device=torch.device("cpu"), group=None)) + layer = Phi4MMAdapter().build(_tiny_config()).model.layers[0] + layer.self_attn.qkv_proj.weight.data.copy_( + torch.cat((q.chunk(tp_size)[rank], k.chunk(tp_size)[rank], v.chunk(tp_size)[rank])) + ) + store = PolicyTensorStore() + save_packed_section_column_spec(store, layer, "model.layers.0", QKV_SPEC) + layout = store["model.layers.0.self_attn.qkv_proj.base_layer.weight"].policy_layout() + contribution = torch.empty(layout.numel, dtype=layout.dtype) + layout.read_chunk(0, contribution) + contributions.append(contribution) + finally: + set_tp_context(old_context) + + reconstructed = torch.stack(contributions).sum(dim=0).reshape_as(full) + torch.testing.assert_close(reconstructed, full, rtol=0, atol=0) + + +def test_phi4mm_text_only_checkpoint_load_save_reload_closes(tmp_path, monkeypatch): + monkeypatch.setenv("ARENO_CKPT_PROGRESS", "0") + source = tmp_path / "source" + output = tmp_path / "output" + tensors = _synthetic_weights(skipped=True) + _write_checkpoint(source, tensors) + first = Phi4MMAdapter().build(_tiny_config()) + Phi4MMAdapter().load_weights(first, source) + + saved_path = Phi4MMAdapter().save_weights(first, output, source) + second = Phi4MMAdapter().build(_tiny_config()) + Phi4MMAdapter().load_weights(second, output) + + assert saved_path == str(output) + assert (output / "config.json").exists() + assert second.lm_head.weight is second.model.embed_tokens.weight + for (first_name, first_parameter), (second_name, second_parameter) in zip( + first.named_parameters(), second.named_parameters(), strict=True + ): + assert first_name == second_name + torch.testing.assert_close(first_parameter, second_parameter, rtol=0, atol=0) + audit = audit_phi4mm_checkpoint(output, num_hidden_layers=2) + assert audit.total == audit.consumed == 14 + with open(output / "model.safetensors.index.json", encoding="utf-8") as handle: + saved_keys = set(json.load(handle)["weight_map"]) + assert not any("embed_tokens_extend" in key or ".lora_" in key for key in saved_keys) + with safe_open(output / "model-rank00000-00002-layer-00000.safetensors", framework="pt") as handle: + torch.testing.assert_close( + handle.get_tensor("model.layers.0.self_attn.qkv_proj.base_layer.weight"), + tensors["model.layers.0.self_attn.qkv_proj.base_layer.weight"], + ) + torch.testing.assert_close( + handle.get_tensor("model.layers.0.mlp.gate_up_proj.base_layer.weight"), + tensors["model.layers.0.mlp.gate_up_proj.base_layer.weight"], + ) diff --git a/tests/test_phi4mm_vision_cpu.py b/tests/test_phi4mm_vision_cpu.py new file mode 100644 index 00000000..0f3b1786 --- /dev/null +++ b/tests/test_phi4mm_vision_cpu.py @@ -0,0 +1,495 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from safetensors.torch import save_file + +from areno.api.multimodal import _image_token_id +from areno.engine.data.rollout_state import InferenceBatchState, _slice_prompt_image_features, payload_to_infer_meta +from areno.engine.parallel.context import TPContext, get_tp_context, set_tp_context + + +@pytest.fixture(autouse=True) +def _isolate_tp_context(): + previous_context = get_tp_context() + set_tp_context(TPContext(rank=0, world_size=1, device=torch.device("cpu"), group=None)) + try: + yield + finally: + set_tp_context(previous_context) + + +def _config() -> dict: + return { + "model_type": "phi4mm", + "vocab_size": 128, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "partial_rotary_factor": 0.5, + "original_max_position_embeddings": 16, + "max_position_embeddings": 32, + "rope_scaling": {"type": "longrope", "short_factor": [1.0], "long_factor": [2.0]}, + "hidden_act": "silu", + "attention_bias": False, + "mlp_bias": False, + "lm_head_bias": False, + "tie_word_embeddings": True, + "torch_dtype": "float32", + "vision_lora": {"r": 4, "lora_alpha": 8, "dp": 0.0}, + "embd_layer": { + "image_embd_layer": { + "embedding_cls": "tune_image", + "crop_size": 8, + "image_token_compression_cls": "avg_pool_2d", + "projection_cls": "mlp", + "use_hd_transform": True, + "with_learnable_separator": True, + "hd_transform_order": "sub_glb", + } + }, + "vision_config": { + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "image_size": 8, + "patch_size": 2, + "feature_layer": -2, + "crop_size": 8, + }, + } + + +def test_phi4mm_adapter_constructs_native_vision_path(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + config = Phi4MMAdapter().config_from_hf(_config()) + model = Phi4MMAdapter().build(config).float() + + assert config.image_token_id == 200010 + assert config.vision_config["hidden_size"] == 8 + assert model.model.embed_tokens_extend.image_embed.img_processor.encoder.layers.__len__() == 2 + + +def test_phi4mm_hd_projection_matches_expanded_image_token_count(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + features = { + "input_image_embeds": torch.zeros(1, 2, 3, 8, 8), + "image_sizes": torch.tensor([[8, 8]], dtype=torch.long), + "image_attention_mask": torch.ones(1, 2, 4, 4, dtype=torch.bool), + "image_token_id": 99, + } + + image_embeds = model.model._project_image_feature(features, torch.device("cpu")) + + assert image_embeds.shape == (13, 16) + + +def test_phi4mm_replaces_only_expanded_image_slots(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + input_ids = torch.tensor([[1, *([99] * 13), 2]], dtype=torch.long) + hidden = torch.randn(1, input_ids.shape[1], 16) + features = { + "input_image_embeds": torch.zeros(1, 2, 3, 8, 8), + "image_sizes": torch.tensor([[8, 8]], dtype=torch.long), + "image_attention_mask": torch.ones(1, 2, 4, 4, dtype=torch.bool), + "image_token_id": 99, + } + + replaced = model.model._apply_multimodal_features(hidden, input_ids, features) + + assert torch.equal(replaced[:, :1], hidden[:, :1]) + assert torch.equal(replaced[:, -1:], hidden[:, -1:]) + assert not torch.equal(replaced[:, 1:-1], hidden[:, 1:-1]) + + +def test_phi4mm_processor_token_fallback_uses_endoftext10(): + tokenizer = SimpleNamespace(convert_tokens_to_ids=lambda token: 200010 if token == "<|endoftext10|>" else -1) + + assert _image_token_id(tokenizer, object()) == 200010 + + +def test_phi4mm_rollout_chunk_keeps_processor_vision_fields(): + features = { + "input_image_embeds": torch.zeros(1, 2, 3, 8, 8), + "image_sizes": torch.tensor([[8, 8]], dtype=torch.long), + "image_attention_mask": torch.ones(1, 2, 4, 4, dtype=torch.bool), + "image_token_id": 99, + } + + mask, payload = _slice_prompt_image_features(features, [1, 99, 99, 2], 0, 4) + + assert mask == [False, True, True, False] + assert payload is not None + assert payload["input_image_embeds"] is features["input_image_embeds"] + assert payload["image_sizes"] is features["image_sizes"] + assert payload["image_attention_mask"] is features["image_attention_mask"] + assert payload["image_token_count"] == 2 + + +def test_phi4mm_chunked_prefill_keeps_vision_lora_active_after_image_chunk(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + features = { + "input_image_embeds": torch.zeros(1, 2, 3, 8, 8), + "image_sizes": torch.tensor([[8, 8]], dtype=torch.long), + "image_attention_mask": torch.ones(1, 2, 4, 4, dtype=torch.bool), + "image_token_id": 99, + } + state = InferenceBatchState( + [[99, 99, 1, 2]], + max_new_tokens=1, + max_prefill_tokens=2, + max_cache_len=8, + kv_block_size=2, + num_cache_blocks=4, + prompt_features=[features], + ) + first = state.build_prefill_payload() + second = state.build_prefill_payload() + + assert first["features"]["image_sequence_mask"].tolist() == [True] + assert second["input_ids"].tolist() == [1, 2] + assert second["features"]["image_sequence_mask"].tolist() == [True] + + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + model.model.vision_lora_slots = torch.zeros(1, dtype=torch.bool) + input_ids = second["input_ids"].unsqueeze(0) + infer_meta = payload_to_infer_meta(second, torch.device("cpu")) + mask = model.model._vision_lora_mask(input_ids, second["features"], None, infer_meta) + + assert mask.tolist() == [[True, True]] + assert model.model.vision_lora_slots.tolist() == [True] + + +def test_phi4mm_projects_multiple_images_with_different_crop_counts(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + features = { + "input_image_embeds": torch.zeros(2, 3, 3, 8, 8), + "image_sizes": torch.tensor([[8, 8], [16, 8]], dtype=torch.long), + "image_attention_mask": torch.ones(2, 3, 4, 4, dtype=torch.bool), + } + + projected = model.model._project_image_feature(features, torch.device("cpu")) + + assert projected.shape == (32, 16) + + +def test_phi4mm_multiple_image_features_follow_placeholder_order(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + first = torch.full((2, 16), 1.0) + second = torch.full((3, 16), 2.0) + input_ids = torch.tensor([[7, 99, 99, 8, 99, 99, 99, 9]]) + hidden = torch.randn(1, input_ids.shape[1], 16) + features = { + "image_feature_rows": [ + {"image_embeds": first, "image_token_count": 2}, + {"image_embeds": second, "image_token_count": 3}, + ], + "image_token_id": 99, + } + + merged = model.model._apply_multimodal_features(hidden, input_ids, features) + + torch.testing.assert_close(merged[0, 1:3], first) + torch.testing.assert_close(merged[0, 4:7], second) + torch.testing.assert_close(merged[0, [0, 3, 7]], hidden[0, [0, 3, 7]]) + + +def test_phi4mm_mixed_batch_keeps_image_features_and_lora_row_local(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + image_token = model.config.image_token_id + input_ids = torch.tensor([[image_token, 1, 2], [3, 4, 5]]) + hidden = torch.randn(2, 3, 16) + image_embeds = torch.full((1, 16), 4.0) + features = [{"image_embeds": image_embeds, "image_token_id": image_token}, None] + + merged = model.model._apply_multimodal_features(hidden, input_ids, features) + lora_mask = model.model._vision_lora_mask(input_ids, features, None, None) + + torch.testing.assert_close(merged[0, 0], image_embeds[0]) + torch.testing.assert_close(merged[0, 1:], hidden[0, 1:]) + torch.testing.assert_close(merged[1], hidden[1]) + assert lora_mask.tolist() == [[True, True, True], [False, False, False]] + + +def test_phi4mm_packed_batch_maps_vision_modes_to_recurrent_slots(): + pytest.importorskip("triton") + from areno.engine.runtime.metadata import InferMeta + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + model.model.vision_lora_slots = torch.zeros(2, dtype=torch.bool) + input_ids = torch.tensor([[1, 2, 3, 4]]) + features = {"image_sequence_mask": torch.tensor([True, False])} + infer_meta = InferMeta( + mode="prefill", + cu_seqlens=torch.tensor([0, 2, 4], dtype=torch.int32), + recurrent_slots=torch.tensor([1, 0]), + ) + + mask = model.model._vision_lora_mask(input_ids, features, None, infer_meta) + + assert mask.tolist() == [[True, True, False, False]] + assert model.model.vision_lora_slots.tolist() == [False, True] + + +def _lora_weights() -> dict[str, torch.Tensor]: + prefix = "model.layers.0" + return { + f"{prefix}.self_attn.qkv_proj.lora_A.vision.weight": torch.arange(4 * 16).view(4, 16).float(), + f"{prefix}.self_attn.qkv_proj.lora_B.vision.weight": torch.arange(48 * 4).view(48, 4).float(), + f"{prefix}.self_attn.o_proj.lora_A.vision.weight": torch.arange(4 * 16).view(4, 16).float() + 1_000, + f"{prefix}.self_attn.o_proj.lora_B.vision.weight": torch.arange(16 * 4).view(16, 4).float() + 2_000, + f"{prefix}.mlp.gate_up_proj.lora_A.vision.weight": torch.arange(4 * 16).view(4, 16).float() + 3_000, + f"{prefix}.mlp.gate_up_proj.lora_B.vision.weight": torch.arange(64 * 4).view(64, 4).float() + 4_000, + f"{prefix}.mlp.down_proj.lora_A.vision.weight": torch.arange(4 * 32).view(4, 32).float() + 5_000, + f"{prefix}.mlp.down_proj.lora_B.vision.weight": torch.arange(16 * 4).view(16, 4).float() + 6_000, + } + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_phi4mm_vision_lora_tp_mapping_shards_each_fused_section(tmp_path, tp_size): + pytest.importorskip("triton") + from areno.models.phi4mm.checkpoint import _load_vision_lora_weights + from areno.models.phi4mm.model import Phi4MMAdapter + + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + tensors = _lora_weights() + save_file(tensors, checkpoint / "model.safetensors") + previous = get_tp_context() + try: + for rank in range(tp_size): + set_tp_context(TPContext(rank=rank, world_size=tp_size, device=torch.device("cpu"), group=None)) + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + _load_vision_lora_weights(model, checkpoint) + layer = model.layers[0] + + q, k, v = tensors["model.layers.0.self_attn.qkv_proj.lora_B.vision.weight"].split((16, 16, 16)) + expected_qkv_b = torch.cat((q.chunk(tp_size)[rank], k.chunk(tp_size)[rank], v.chunk(tp_size)[rank])) + gate, up = tensors["model.layers.0.mlp.gate_up_proj.lora_B.vision.weight"].split((32, 32)) + expected_gate_b = torch.cat((gate.chunk(tp_size)[rank], up.chunk(tp_size)[rank])) + + torch.testing.assert_close(layer.self_attn.qkv_proj.lora_B["vision"].weight, expected_qkv_b) + torch.testing.assert_close(layer.mlp.gate_up_proj.lora_B["vision"].weight, expected_gate_b) + torch.testing.assert_close( + layer.self_attn.o_proj.lora_A["vision"].weight, + tensors["model.layers.0.self_attn.o_proj.lora_A.vision.weight"].chunk(tp_size, dim=1)[rank], + ) + torch.testing.assert_close( + layer.mlp.down_proj.lora_A["vision"].weight, + tensors["model.layers.0.mlp.down_proj.lora_A.vision.weight"].chunk(tp_size, dim=1)[rank], + ) + torch.testing.assert_close( + layer.self_attn.qkv_proj.lora_A["vision"].weight, + tensors["model.layers.0.self_attn.qkv_proj.lora_A.vision.weight"], + ) + torch.testing.assert_close( + layer.self_attn.o_proj.lora_B["vision"].weight, + tensors["model.layers.0.self_attn.o_proj.lora_B.vision.weight"], + ) + finally: + set_tp_context(previous) + + +def test_phi4mm_vision_lora_tp1_forward_matches_peft_formula(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + projection = model.layers[0].self_attn.qkv_proj + projection.weight.data.zero_() + projection.lora_A["vision"].weight.data.copy_(torch.arange(4 * 16).view(4, 16).float() / 100) + projection.lora_B["vision"].weight.data.copy_(torch.arange(48 * 4).view(48, 4).float() / 100) + projection.vision_lora_mask = torch.tensor([[True, False, True]]) + inputs = torch.arange(3 * 16).view(1, 3, 16).float() / 100 + + actual = projection(inputs) + expected = 2.0 * F.linear(F.linear(inputs, projection.lora_A["vision"].weight), projection.lora_B["vision"].weight) + expected[:, 1].zero_() + + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize("tp_size", [2, 4]) +def test_phi4mm_vision_lora_tp_shards_reconstruct_peft_formula(tmp_path, tp_size): + pytest.importorskip("triton") + from areno.models.phi4mm.checkpoint import _load_vision_lora_weights + from areno.models.phi4mm.model import Phi4MMAdapter + + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + tensors = {name: tensor / 10_000 for name, tensor in _lora_weights().items()} + save_file(tensors, checkpoint / "model.safetensors") + inputs = torch.arange(3 * 16).view(3, 16).float() / 100 + down_inputs = torch.arange(3 * 32).view(3, 32).float() / 100 + qkv_parts: list[tuple[torch.Tensor, ...]] = [] + gate_parts: list[tuple[torch.Tensor, ...]] = [] + o_latents = [] + down_latents = [] + previous = get_tp_context() + try: + for rank in range(tp_size): + set_tp_context(TPContext(rank=rank, world_size=tp_size, device=torch.device("cpu"), group=None)) + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + _load_vision_lora_weights(model, checkpoint) + layer = model.layers[0] + qkv_delta = F.linear( + F.linear(inputs, layer.self_attn.qkv_proj.lora_A["vision"].weight), + layer.self_attn.qkv_proj.lora_B["vision"].weight, + ) + gate_delta = F.linear( + F.linear(inputs, layer.mlp.gate_up_proj.lora_A["vision"].weight), + layer.mlp.gate_up_proj.lora_B["vision"].weight, + ) + qkv_parts.append(qkv_delta.split((16 // tp_size,) * 3, dim=-1)) + gate_parts.append(gate_delta.split((32 // tp_size,) * 2, dim=-1)) + o_latents.append( + F.linear(inputs.chunk(tp_size, dim=-1)[rank], layer.self_attn.o_proj.lora_A["vision"].weight) + ) + down_latents.append( + F.linear(down_inputs.chunk(tp_size, dim=-1)[rank], layer.mlp.down_proj.lora_A["vision"].weight) + ) + qkv_actual = torch.cat( + [torch.cat([parts[section] for parts in qkv_parts], dim=-1) for section in range(3)], dim=-1 + ) + gate_actual = torch.cat( + [torch.cat([parts[section] for parts in gate_parts], dim=-1) for section in range(2)], dim=-1 + ) + o_actual = F.linear(sum(o_latents), layer.self_attn.o_proj.lora_B["vision"].weight) + down_actual = F.linear(sum(down_latents), layer.mlp.down_proj.lora_B["vision"].weight) + finally: + set_tp_context(previous) + + prefix = "model.layers.0" + qkv_expected = F.linear( + F.linear(inputs, tensors[f"{prefix}.self_attn.qkv_proj.lora_A.vision.weight"]), + tensors[f"{prefix}.self_attn.qkv_proj.lora_B.vision.weight"], + ) + gate_expected = F.linear( + F.linear(inputs, tensors[f"{prefix}.mlp.gate_up_proj.lora_A.vision.weight"]), + tensors[f"{prefix}.mlp.gate_up_proj.lora_B.vision.weight"], + ) + o_expected = F.linear( + F.linear(inputs, tensors[f"{prefix}.self_attn.o_proj.lora_A.vision.weight"]), + tensors[f"{prefix}.self_attn.o_proj.lora_B.vision.weight"], + ) + down_expected = F.linear( + F.linear(down_inputs, tensors[f"{prefix}.mlp.down_proj.lora_A.vision.weight"]), + tensors[f"{prefix}.mlp.down_proj.lora_B.vision.weight"], + ) + torch.testing.assert_close(qkv_actual, qkv_expected, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(gate_actual, gate_expected, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(o_actual, o_expected, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(down_actual, down_expected, rtol=1e-5, atol=1e-5) + + +def test_phi4mm_vision_checkpoint_save_reload_closes(tmp_path, monkeypatch): + pytest.importorskip("triton") + from areno.engine.checkpoints.io import SafetensorsIndex + from areno.models.phi4mm.checkpoint import ( + _vision_checkpoint_keys, + _vision_lora_checkpoint_keys, + audit_phi4mm_checkpoint, + ) + from areno.models.phi4mm.model import Phi4MMAdapter + + monkeypatch.setenv("ARENO_CKPT_PROGRESS", "0") + torch.manual_seed(7) + adapter = Phi4MMAdapter() + config = adapter.config_from_hf(_config()) + first = adapter.build(config).float() + output = tmp_path / "output" + + saved_path = adapter.save_weights(first, output, None) + second = adapter.build(config).float() + adapter.load_weights(second, output) + + assert saved_path == str(output) + assert second.lm_head.weight is second.model.embed_tokens.weight + for (first_name, first_parameter), (second_name, second_parameter) in zip( + first.named_parameters(), second.named_parameters(), strict=True + ): + assert first_name == second_name + torch.testing.assert_close(first_parameter, second_parameter, rtol=0, atol=0) + + vision_keys = _vision_checkpoint_keys(first) + vision_lora_keys = _vision_lora_checkpoint_keys(first) + audit = audit_phi4mm_checkpoint(output, len(first.layers), vision_keys, vision_lora_keys) + assert audit.consumed == audit.total + assert audit.speech_lora_skipped == audit.audio_skipped == audit.unknown == 0 + assert "model.embed_tokens_extend.image_embed.sub_GN" in vision_keys + assert "model.embed_tokens_extend.image_embed.glb_GN" in vision_keys + assert len(vision_lora_keys) == 8 + + index = SafetensorsIndex(output, progress=False) + try: + saved_keys = set(index.weight_map) + finally: + index.close() + assert vision_keys | vision_lora_keys <= saved_keys + assert not any(".speech." in key or ".audio_embed." in key for key in saved_keys) + + +@pytest.mark.parametrize("tp_size", [2, 4]) +def test_phi4mm_vision_lora_save_layout_inverts_tp_sharding(tmp_path, tp_size): + pytest.importorskip("triton") + from areno.engine.checkpoints.io import PolicyTensorStore, policy_plan_scope + from areno.models.phi4mm.checkpoint import _load_vision_lora_weights, _save_vision_weights + from areno.models.phi4mm.model import Phi4MMAdapter + + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + expected = _lora_weights() + save_file(expected, checkpoint / "model.safetensors") + sharded_keys = ( + "model.layers.0.self_attn.qkv_proj.lora_B.vision.weight", + "model.layers.0.mlp.gate_up_proj.lora_B.vision.weight", + "model.layers.0.self_attn.o_proj.lora_A.vision.weight", + "model.layers.0.mlp.down_proj.lora_A.vision.weight", + ) + contributions = {key: [] for key in sharded_keys} + previous = get_tp_context() + try: + for rank in range(tp_size): + set_tp_context(TPContext(rank=rank, world_size=tp_size, device=torch.device("cpu"), group=None)) + model = Phi4MMAdapter().build(Phi4MMAdapter().config_from_hf(_config())).float() + _load_vision_lora_weights(model, checkpoint) + store = PolicyTensorStore() + with policy_plan_scope(): + _save_vision_weights(store, model) + for key in sharded_keys: + layout = store[key].policy_layout() + contribution = torch.empty(layout.numel, dtype=layout.dtype) + layout.read_chunk(0, contribution) + contributions[key].append(contribution) + finally: + set_tp_context(previous) + + for key in sharded_keys: + reconstructed = torch.stack(contributions[key]).sum(dim=0).reshape_as(expected[key]) + torch.testing.assert_close(reconstructed, expected[key], rtol=0, atol=0)