From 2248c6915242be497d74ca37bee91810f5f20ad0 Mon Sep 17 00:00:00 2001 From: zitai-wang <2531131993@qq.com> Date: Wed, 26 Aug 2026 14:35:21 +0800 Subject: [PATCH 01/10] feat(checkpoint): support packed-section TP mapping --- areno/engine/checkpoints/common.py | 83 +++++++++++++++++++++++++++++- 1 file changed, 82 insertions(+), 1 deletion(-) 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: From 20ef04900975082d7cf522c31f1367dee088979a Mon Sep 17 00:00:00 2001 From: zitai-wang <2531131993@qq.com> Date: Wed, 26 Aug 2026 14:36:24 +0800 Subject: [PATCH 02/10] feat(models): add Phi-4 text-only adapter --- areno/engine/layers/attention.py | 23 +- areno/engine/runtime/decode_graph.py | 13 + areno/models/__init__.py | 8 + areno/models/phi4mm/__init__.py | 21 ++ areno/models/phi4mm/checkpoint.py | 162 ++++++++++ areno/models/phi4mm/model.py | 467 +++++++++++++++++++++++++++ 6 files changed, 691 insertions(+), 3 deletions(-) create mode 100644 areno/models/phi4mm/__init__.py create mode 100644 areno/models/phi4mm/checkpoint.py create mode 100644 areno/models/phi4mm/model.py 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/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..9fcd4d9a --- /dev/null +++ b/areno/models/phi4mm/checkpoint.py @@ -0,0 +1,162 @@ +"""Strict text-only checkpoint mapping for Phi-4-Multimodal.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path + +from torch import nn + +from areno.engine.checkpoints.common import ( + CheckpointSpec, + LayerSpec, + PackedSectionColumnSpec, + ParallelTensorSpec, + ReplicatedTensorSpec, + TopLevelSpec, + load_checkpoint_weights, + save_checkpoint_weights, +) +from areno.engine.checkpoints.io import SafetensorsIndex +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) -> 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() + required = _required_base_keys(num_hidden_layers) + 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 base-language tensors.""" + + model.config.validate_tp(get_tp_context().world_size) + audit = audit_phi4mm_checkpoint(model_path, len(model.layers)) + load_checkpoint_weights(model, str(model_path), CHECKPOINT_SPEC) + 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 save_phi4mm_weights( + model: nn.Module, + output_path: str | Path, + source_path: str | Path | None, +) -> str | None: + """Save only Phi-4 base-language weights in 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, + copy_passthrough=False, + ) diff --git a/areno/models/phi4mm/model.py b/areno/models/phi4mm/model.py new file mode 100644 index 00000000..45118cf3 --- /dev/null +++ b/areno/models/phi4mm/model.py @@ -0,0 +1,467 @@ +"""Phi-4-Multimodal language-backbone adapter. + +PR1 intentionally supports the checkpoint's text path only. The vision and +audio towers and their modality-specific LoRA adapters are not runtime model +components here. +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch +from torch import nn + +from areno.accel.ops import is_cuda_graph_capturing +from areno.engine.config import ModelConfig, _parse_dtype +from areno.engine.layers.attention import CausalSelfAttention +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 ( + 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 + + +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)) + + 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) + + 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): + """Text-only Phi-4 transformer body.""" + + def __init__(self, config: ModelConfig): + super().__init__() + self.config = config + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size, dtype=config.dtype) + 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) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None = None, + train_meta: TrainMeta | None = None, + infer_meta: InferMeta | 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) + hidden_states = self.embed_tokens(input_ids) + 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) + + +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, + ) -> 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) + 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.""" + del num_slots + 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) + + @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 + + 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, + ) + + 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) From 0c13a01b41c17782a768bc3601cc2dfca8f2eb59 Mon Sep 17 00:00:00 2001 From: zitai-wang <2531131993@qq.com> Date: Wed, 26 Aug 2026 14:36:31 +0800 Subject: [PATCH 03/10] test(models): add Phi-4 adapter and checkpoint coverage --- tests/test_phi4mm_adapter_cpu.py | 434 ++++++++++++++++++++++++++++ tests/test_phi4mm_checkpoint_cpu.py | 275 ++++++++++++++++++ 2 files changed, 709 insertions(+) create mode 100644 tests/test_phi4mm_adapter_cpu.py create mode 100644 tests/test_phi4mm_checkpoint_cpu.py diff --git a/tests/test_phi4mm_adapter_cpu.py b/tests/test_phi4mm_adapter_cpu.py new file mode 100644 index 00000000..f4c6fcc5 --- /dev/null +++ b/tests/test_phi4mm_adapter_cpu.py @@ -0,0 +1,434 @@ +from __future__ import annotations + +import json +import math + +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), + }, + }, + ) + + +@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"], + ) From 0b29173c787f8c9ca56e520ef1ed4d85f89107df Mon Sep 17 00:00:00 2001 From: zitai-wang <2531131993@qq.com> Date: Wed, 26 Aug 2026 16:43:41 +0800 Subject: [PATCH 04/10] feat(models): add Phi-4 multimodal vision support --- areno/models/phi4mm/checkpoint.py | 148 ++++++++++++- areno/models/phi4mm/model.py | 339 +++++++++++++++++++++++++++++- areno/models/phi4mm/vision.py | 253 ++++++++++++++++++++++ 3 files changed, 724 insertions(+), 16 deletions(-) create mode 100644 areno/models/phi4mm/vision.py diff --git a/areno/models/phi4mm/checkpoint.py b/areno/models/phi4mm/checkpoint.py index 9fcd4d9a..dae80858 100644 --- a/areno/models/phi4mm/checkpoint.py +++ b/areno/models/phi4mm/checkpoint.py @@ -1,4 +1,4 @@ -"""Strict text-only checkpoint mapping for Phi-4-Multimodal.""" +"""Strict checkpoint mapping for the supported Phi-4-Multimodal paths.""" from __future__ import annotations @@ -6,19 +6,26 @@ 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 +from areno.engine.checkpoints.io import SafetensorsIndex, _copy_row from areno.engine.parallel.context import get_tp_context TOP_LEVEL_SPEC = TopLevelSpec( @@ -94,7 +101,12 @@ def _required_base_keys(num_hidden_layers: int) -> set[str]: return required -def audit_phi4mm_checkpoint(model_path: str | Path, num_hidden_layers: int) -> Phi4MMCheckpointAudit: +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) @@ -102,7 +114,8 @@ def audit_phi4mm_checkpoint(model_path: str | Path, num_hidden_layers: int) -> P checkpoint_keys = set(index.weight_map) finally: index.close() - required = _required_base_keys(num_hidden_layers) + 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]) @@ -135,22 +148,107 @@ def audit_phi4mm_checkpoint(model_path: str | Path, num_hidden_layers: int) -> P def load_phi4mm_weights(model: nn.Module, model_path: str | Path) -> Phi4MMCheckpointAudit: - """Audit and load the supported Phi-4 base-language tensors.""" + """Audit and load the supported Phi-4 language and vision tensors.""" model.config.validate_tp(get_tp_context().world_size) - audit = audit_phi4mm_checkpoint(model_path, len(model.layers)) + 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 only Phi-4 base-language weights in official HF key layout.""" + """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( @@ -158,5 +256,41 @@ def save_phi4mm_weights( 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 index 45118cf3..6691a965 100644 --- a/areno/models/phi4mm/model.py +++ b/areno/models/phi4mm/model.py @@ -1,9 +1,4 @@ -"""Phi-4-Multimodal language-backbone adapter. - -PR1 intentionally supports the checkpoint's text path only. The vision and -audio towers and their modality-specific LoRA adapters are not runtime model -components here. -""" +"""Phi-4-Multimodal language and vision adapter.""" from __future__ import annotations @@ -12,21 +7,180 @@ from typing import Any import torch +import torch.nn.functional as F from torch import nn from areno.accel.ops 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: @@ -182,6 +336,16 @@ 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, @@ -214,6 +378,18 @@ def __init__(self, config: ModelConfig, layer_idx: int): 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, @@ -231,14 +407,25 @@ def forward( class Phi4MMModel(nn.Module): - """Text-only Phi-4 transformer body.""" + """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, @@ -246,10 +433,15 @@ def forward( 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) @@ -266,6 +458,125 @@ def forward( ) 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.""" @@ -300,10 +611,11 @@ def forward( 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) + 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) @@ -311,11 +623,17 @@ 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.""" - del num_slots 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: @@ -425,6 +743,7 @@ def config_from_hf(self, hf_config: dict[str, Any]) -> ModelConfig: 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, @@ -449,6 +768,8 @@ def config_from_hf(self, hf_config: dict[str, Any]) -> ModelConfig: 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: 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) From 7976d5a1e1d6773a766efcd4903fdfbec2510922 Mon Sep 17 00:00:00 2001 From: zitai-wang <2531131993@qq.com> Date: Wed, 26 Aug 2026 16:43:50 +0800 Subject: [PATCH 05/10] fix(runtime): preserve vision state across chunked prefill --- areno/api/multimodal.py | 4 ++-- areno/engine/data/rollout_state.py | 35 ++++++++++++++++++++++++++++-- 2 files changed, 35 insertions(+), 4 deletions(-) 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/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.""" From 2856e5d599b09782c85b8327a8c37f85eef34811 Mon Sep 17 00:00:00 2001 From: zitai-wang <2531131993@qq.com> Date: Wed, 26 Aug 2026 16:44:01 +0800 Subject: [PATCH 06/10] test(models): add Phi-4 vision coverage --- tests/test_phi4mm_vision_cpu.py | 495 ++++++++++++++++++++++++++++++++ 1 file changed, 495 insertions(+) create mode 100644 tests/test_phi4mm_vision_cpu.py 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) From 3508d5a08b3953afc08e2ca11d0bc79d45537e78 Mon Sep 17 00:00:00 2001 From: zitai-wang <2531131993@qq.com> Date: Wed, 26 Aug 2026 17:54:00 +0800 Subject: [PATCH 07/10] feat(models): add Phi-4 audio encoder and speech adapter --- areno/models/phi4mm/audio.py | 299 +++++++++++++++++++++++ areno/models/phi4mm/checkpoint.py | 142 ++++++----- areno/models/phi4mm/model.py | 379 ++++++++++++++++++++++-------- areno/models/phi4mm/vision.py | 10 +- 4 files changed, 671 insertions(+), 159 deletions(-) create mode 100644 areno/models/phi4mm/audio.py diff --git a/areno/models/phi4mm/audio.py b/areno/models/phi4mm/audio.py new file mode 100644 index 00000000..f79887e5 --- /dev/null +++ b/areno/models/phi4mm/audio.py @@ -0,0 +1,299 @@ +"""Native Phi-4-Multimodal Conformer audio encoder and projector.""" + +from __future__ import annotations + +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 Phi4MMAudioConfig: + input_size: int = 80 + attention_dim: int = 1024 + attention_heads: int = 16 + linear_units: int = 1536 + num_blocks: int = 24 + kernel_size: int = 3 + time_reduction: int = 8 + relative_attention_max_distance: int = 500 + + @classmethod + def from_dict(cls, values: dict[str, Any]) -> Phi4MMAudioConfig: + relative = values.get("relative_attention_bias_args") or {} + normalized = dict(values) + normalized["relative_attention_max_distance"] = int(relative.get("t5_bias_max_distance", 500)) + return cls(**{name: normalized[name] for name in cls.__dataclass_fields__ if name in normalized}) + + +class _Swish(nn.Module): + def forward(self, value: torch.Tensor) -> torch.Tensor: + return value * torch.sigmoid(value) + + +class _GLULinear(nn.Module): + def __init__(self, input_dim: int, output_dim: int, dtype: torch.dtype): + super().__init__() + self.linear = nn.Linear(input_dim, output_dim * 2, dtype=dtype) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + left, gate = self.linear(value).chunk(2, dim=-1) + return left * (gate * torch.sigmoid(gate)) + + +class _FeedForward(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int, dtype: torch.dtype): + super().__init__() + self.layer_norm = nn.LayerNorm(hidden_size, dtype=dtype) + self.net = nn.Sequential( + _GLULinear(hidden_size, intermediate_size, dtype), + nn.Dropout(0.0), + nn.Linear(intermediate_size, hidden_size, dtype=dtype), + nn.Dropout(0.0), + ) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + return self.net(self.layer_norm(value)) + + +class _GLUPointWiseConv(nn.Module): + def __init__(self, hidden_size: int, dtype: torch.dtype): + super().__init__() + self.output_dim = hidden_size + self.ext_pw_conv_1d = nn.Conv1d(hidden_size, hidden_size * 2, 1, padding=0, dtype=dtype) + self.b1 = nn.Parameter(torch.zeros(1, hidden_size, 1, dtype=dtype)) + self.b2 = nn.Parameter(torch.zeros(1, hidden_size, 1, dtype=dtype)) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + value = self.ext_pw_conv_1d(value.transpose(1, 2)) + left = value[:, : self.output_dim] + self.b1 + gate = value[:, self.output_dim :] + self.b2 + return (left * (gate * torch.sigmoid(gate))).transpose(1, 2) + + +class _DepthWiseSeperableConv1d(nn.Module): + # Keep the official misspelling in this private class because it defines + # checkpoint-compatible attribute names. + def __init__(self, hidden_size: int, kernel_size: int, dtype: torch.dtype): + super().__init__() + self.dw_conv = nn.Conv1d( + hidden_size, + hidden_size, + kernel_size, + padding=kernel_size - 1, + groups=hidden_size, + dtype=dtype, + ) + self.pw_conv = nn.Conv1d(hidden_size, hidden_size, 1, dtype=dtype) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + return self.pw_conv(self.dw_conv(value)) + + +class _ConvModule(nn.Module): + def __init__(self, hidden_size: int, kernel_size: int, dtype: torch.dtype): + super().__init__() + self.kernel_size = kernel_size + self.layer_norm = nn.LayerNorm(hidden_size, dtype=dtype) + self.glu = _GLUPointWiseConv(hidden_size, dtype) + self.dw_sep_conv_1d = _DepthWiseSeperableConv1d(hidden_size, kernel_size, dtype) + self.ext_pw_conv_1d = nn.Conv1d(hidden_size, hidden_size, 1, dtype=dtype) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + value = self.glu(self.layer_norm(value)).transpose(1, 2) + value = self.dw_sep_conv_1d(value) + value = value[:, :, : -(self.kernel_size - 1)] + value = value * torch.sigmoid(value) + return self.ext_pw_conv_1d(value).transpose(1, 2) + + +class _T5RelativeAttentionLogitBias(nn.Module): + def __init__(self, num_heads: int, max_distance: int, dtype: torch.dtype): + super().__init__() + self.max_distance = max_distance + self.bias_values = nn.Embedding(max_distance * 2, num_heads, dtype=dtype) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + length = value.shape[1] + positions = torch.arange(length, device=value.device, dtype=torch.long) + relative = positions[None, :] - positions[:, None] + indices = relative.clamp(-self.max_distance, self.max_distance - 1) + self.max_distance + return self.bias_values(indices).permute(2, 0, 1).unsqueeze(0) + + +class _MultiHeadedAttention(nn.Module): + def __init__(self, hidden_size: int, num_heads: int, dtype: torch.dtype): + super().__init__() + if hidden_size % num_heads: + raise ValueError("Phi4MM audio attention dimension must be divisible by its head count") + self.h = num_heads + self.d_k = hidden_size // num_heads + self.inv_sqrt_d_k = self.d_k**-0.5 + self.linear_q = nn.Linear(hidden_size, hidden_size, dtype=dtype) + self.linear_k = nn.Linear(hidden_size, hidden_size, dtype=dtype) + self.linear_v = nn.Linear(hidden_size, hidden_size, dtype=dtype) + self.linear_out = nn.Linear(hidden_size, hidden_size, dtype=dtype) + + def forward( + self, + value: torch.Tensor, + mask: torch.Tensor | None, + relative_attention_bias: torch.Tensor, + ) -> torch.Tensor: + batch = value.shape[0] + query = self.linear_q(value).view(batch, -1, self.h, self.d_k).transpose(1, 2) * self.inv_sqrt_d_k + key = self.linear_k(value).view(batch, -1, self.h, self.d_k).transpose(1, 2) + projected = self.linear_v(value).view(batch, -1, self.h, self.d_k).transpose(1, 2) + scores = torch.matmul(query, key.transpose(-2, -1)) + relative_attention_bias + if mask is not None: + invalid = ~mask.unsqueeze(1) + attention = torch.softmax(scores.masked_fill(invalid, -torch.inf), dim=-1).masked_fill(invalid, 0.0) + else: + attention = torch.softmax(scores, dim=-1) + output = torch.matmul(attention.to(projected.dtype), projected) + output = output.transpose(1, 2).contiguous().view(batch, -1, self.h * self.d_k) + return self.linear_out(output) + + +class _ConformerEncoderLayer(nn.Module): + def __init__(self, config: Phi4MMAudioConfig, dtype: torch.dtype): + super().__init__() + self.feed_forward_in = _FeedForward(config.attention_dim, config.linear_units, dtype) + self.self_attn = _MultiHeadedAttention(config.attention_dim, config.attention_heads, dtype) + self.conv = _ConvModule(config.attention_dim, config.kernel_size, dtype) + self.feed_forward_out = _FeedForward(config.attention_dim, config.linear_units, dtype) + self.layer_norm_att = nn.LayerNorm(config.attention_dim, dtype=dtype) + self.layer_norm = nn.LayerNorm(config.attention_dim, dtype=dtype) + + def forward( + self, + value: torch.Tensor, + mask: torch.Tensor | None, + relative_attention_bias: torch.Tensor, + ) -> torch.Tensor: + value = value + 0.5 * self.feed_forward_in(value) + normalized = self.layer_norm_att(value) + value = value + self.self_attn(normalized, mask, relative_attention_bias) + value = value + self.conv(value) + value = value + 0.5 * self.feed_forward_out(value) + return self.layer_norm(value) + + +class _MeanVarianceNormLayer(nn.Module): + def __init__(self, input_size: int, dtype: torch.dtype): + super().__init__() + self.register_buffer("global_mean", torch.zeros(input_size, dtype=dtype)) + self.register_buffer("global_invstd", torch.ones(input_size, dtype=dtype)) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + return (value - self.global_mean) * self.global_invstd + + +class _NemoConvSubsampling(nn.Module): + def __init__(self, config: Phi4MMAudioConfig, dtype: torch.dtype): + super().__init__() + if config.time_reduction != 8 or config.input_size != 80: + raise ValueError("Phi4MM audio requires 80-bin features and time_reduction=8") + hidden = config.attention_dim + self.subsampling_factor = config.time_reduction + self.conv = nn.Sequential( + nn.Conv2d(1, hidden, 3, stride=2, padding=1, dtype=dtype), + nn.ReLU(), + nn.Conv2d(hidden, hidden, 3, stride=2, padding=1, groups=hidden, dtype=dtype), + nn.Conv2d(hidden, hidden, 1, dtype=dtype), + nn.ReLU(), + nn.Conv2d(hidden, hidden, 3, stride=2, padding=1, groups=hidden, dtype=dtype), + nn.Conv2d(hidden, hidden, 1, dtype=dtype), + nn.ReLU(), + ) + self.out = nn.Linear(hidden * 10, hidden, dtype=dtype) + + def forward( + self, value: torch.Tensor, attention_mask: torch.Tensor | None + ) -> tuple[torch.Tensor, torch.Tensor | None]: + value = self.conv(value.unsqueeze(1)) + batch, channels, time, frequency = value.shape + value = self.out(value.transpose(1, 2).reshape(batch, time, channels * frequency)) + if attention_mask is None: + return value, None + lengths = torch.ceil(attention_mask.sum(1) / self.subsampling_factor).to(dtype=torch.long) + valid = torch.arange(time, device=value.device).unsqueeze(0) < lengths.unsqueeze(1) + return value, valid.unsqueeze(1) + + +class _ConformerEncoder(nn.Module): + def __init__(self, config: Phi4MMAudioConfig, dtype: torch.dtype): + super().__init__() + self.embed = _NemoConvSubsampling(config, dtype) + self.relative_attention_bias_layer = _T5RelativeAttentionLogitBias( + config.attention_heads, config.relative_attention_max_distance, dtype + ) + self.encoders = nn.ModuleList([_ConformerEncoderLayer(config, dtype) for _ in range(config.num_blocks)]) + self.encoder_embedding = _MeanVarianceNormLayer(config.input_size, dtype) + + def _encode_chunk(self, value: torch.Tensor, padding_mask: torch.Tensor | None) -> torch.Tensor: + relative_bias = self.relative_attention_bias_layer(value) + # chunk_size=-1 in the released checkpoint denotes full-utterance + # attention; padding still masks invalid keys for variable-length input. + attention_mask = None + if padding_mask is not None: + attention_mask = padding_mask.expand(-1, value.shape[1], -1) + for layer in self.encoders: + value = layer(value, attention_mask, relative_bias) + return value + + def forward( + self, value: torch.Tensor, attention_mask: torch.Tensor | None + ) -> tuple[torch.Tensor, torch.Tensor | None]: + value = self.encoder_embedding(value) + value, padding_mask = self.embed(value, attention_mask) + original_length = value.shape[1] + if original_length <= 500: + return self._encode_chunk(value, padding_mask), padding_mask + padded_length = ((original_length + 499) // 500) * 500 + value = F.pad(value, (0, 0, 0, padded_length - original_length)) + chunks = value.reshape(-1, 500, value.shape[-1]) + chunk_masks = None + if padding_mask is not None: + valid = F.pad(padding_mask.squeeze(1), (0, padded_length - original_length), value=False) + chunk_masks = valid.reshape(-1, 1, 500) + output = self._encode_chunk(chunks, chunk_masks) + return output.reshape(value.shape[0], padded_length, -1)[:, :original_length], padding_mask + + +class Phi4MMAudioEmbedding(nn.Module): + """Checkpoint-compatible Phi-4 audio tower with both official projectors.""" + + def __init__( + self, + config: Phi4MMAudioConfig, + language_hidden_size: int, + dtype: torch.dtype, + ): + super().__init__() + self.encoder = _ConformerEncoder(config, dtype) + self.audio_projection = nn.ModuleDict( + { + mode: nn.Sequential( + nn.Linear(config.attention_dim, language_hidden_size, dtype=dtype), + nn.GELU(), + nn.Linear(language_hidden_size, language_hidden_size, dtype=dtype), + ) + for mode in ("speech", "vision") + } + ) + + def forward( + self, + input_audio_embeds: torch.Tensor, + audio_attention_mask: torch.Tensor | None, + projection_mode: str = "speech", + ) -> torch.Tensor: + if projection_mode not in self.audio_projection: + raise ValueError(f"unsupported Phi4MM audio projection mode {projection_mode!r}") + dtype = self.audio_projection[projection_mode][0].weight.dtype + features, _ = self.encoder(input_audio_embeds.to(dtype=dtype), audio_attention_mask) + return self.audio_projection[projection_mode](features) diff --git a/areno/models/phi4mm/checkpoint.py b/areno/models/phi4mm/checkpoint.py index dae80858..e8a5dbf2 100644 --- a/areno/models/phi4mm/checkpoint.py +++ b/areno/models/phi4mm/checkpoint.py @@ -148,51 +148,53 @@ def audit_phi4mm_checkpoint( def load_phi4mm_weights(model: nn.Module, model_path: str | Path) -> Phi4MMCheckpointAudit: - """Audit and load the supported Phi-4 language and vision tensors.""" + """Audit and load the supported Phi-4 language and multimodal 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) + multimodal_keys = _multimodal_checkpoint_keys(model) + lora_keys = _lora_checkpoint_keys(model) + audit = audit_phi4mm_checkpoint(model_path, len(model.layers), multimodal_keys, 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 multimodal_keys: + _load_multimodal_weights(model, model_path, multimodal_keys) + if lora_keys: + _load_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]: +def _multimodal_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()} + names = {name for name, _ in extended.named_parameters()} + names.update(name for name, _ in extended.named_buffers()) + return {f"model.embed_tokens_extend.{name}" for name in names} -def _vision_lora_checkpoint_keys(model: nn.Module) -> set[str]: +def _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 + if any(f".lora_{side}.{adapter}.weight" in name for side in "AB" for adapter in ("vision", "speech")) } @torch.no_grad() -def _load_vision_weights(model: nn.Module, model_path: str | Path, keys: set[str] | None = None) -> None: +def _load_multimodal_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) + expected = keys if keys is not None else _multimodal_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]}") + raise KeyError(f"missing Phi4MM multimodal weight {missing[0]}") index.prefetch(sorted(expected)) - for name, parameter in extended.named_parameters(): + for name, parameter in list(extended.named_parameters()) + list(extended.named_buffers()): key = f"model.embed_tokens_extend.{name}" source = index.get_tensor(key) if tuple(source.shape) != tuple(parameter.shape): @@ -205,40 +207,45 @@ def _load_vision_weights(model: nn.Module, model_path: str | Path, keys: set[str @torch.no_grad() -def _load_vision_lora_weights(model: nn.Module, model_path: str | Path) -> None: +def _load_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) - ) + for adapter in ("vision", "speech"): + 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), + ): + if adapter not in module.lora_A: + continue + lora_a = f"{prefix}.{name}.lora_A.{adapter}.weight" + lora_b = f"{prefix}.{name}.lora_B.{adapter}.weight" + module.lora_A[adapter].weight.copy_(index.get_tensor(lora_a).to(dtype=module.weight.dtype)) + copy_merged_column( + module.lora_B[adapter].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), + ): + if adapter not in module.lora_A: + continue + lora_a = f"{prefix}.{name}.lora_A.{adapter}.weight" + lora_b = f"{prefix}.{name}.lora_B.{adapter}.weight" + _copy_row( + module.lora_A[adapter].weight, + index.get_tensor(lora_a), + context.rank, + context.world_size, + ) + module.lora_B[adapter].weight.copy_( + index.get_tensor(lora_b).to(device=module.weight.device, dtype=module.weight.dtype) + ) finally: index.close() @@ -248,7 +255,7 @@ def save_phi4mm_weights( output_path: str | Path, source_path: str | Path | None, ) -> str | None: - """Save Phi-4 language and vision weights in the official HF key layout.""" + """Save Phi-4 language and multimodal weights in the official HF key layout.""" model.config.validate_tp(get_tp_context().world_size) return save_checkpoint_weights( @@ -256,19 +263,21 @@ def save_phi4mm_weights( 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), + extra_tensors_fn=lambda tensors: _save_multimodal_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.""" +def _save_multimodal_weights(tensors: CheckpointTensorStore | PolicyTensorStore, model: nn.Module) -> None: + """Stage replicated modality tensors and TP-aware 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 name, buffer in extended.named_buffers(): + tensors[f"model.embed_tokens_extend.{name}"] = rank0_tensor(buffer) for layer_idx, layer in enumerate(model.layers): prefix = f"model.layers.{layer_idx}" @@ -276,21 +285,32 @@ def _save_vision_weights(tensors: CheckpointTensorStore | PolicyTensorStore, mod ("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: + if not hasattr(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 adapter in module.lora_A: + tensors[f"{prefix}.{name}.lora_A.{adapter}.weight"] = rank0_tensor(module.lora_A[adapter].weight) + tensors[f"{prefix}.{name}.lora_B.{adapter}.weight"] = gather_tensor_parallel_split_column_tensor( + module.lora_B[adapter].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: + if not hasattr(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) + for adapter in module.lora_A: + tensors[f"{prefix}.{name}.lora_A.{adapter}.weight"] = gather_tensor_parallel_tensor( + module.lora_A[adapter].weight, + dim=1, + ) + tensors[f"{prefix}.{name}.lora_B.{adapter}.weight"] = rank0_tensor(module.lora_B[adapter].weight) + + +# Retain the PR2 private names for downstream tests and integrations that imported +# them before audio support generalized the checkpoint path. +_vision_checkpoint_keys = _multimodal_checkpoint_keys +_vision_lora_checkpoint_keys = _lora_checkpoint_keys +_load_vision_weights = _load_multimodal_weights +_load_vision_lora_weights = _load_lora_weights +_save_vision_weights = _save_multimodal_weights diff --git a/areno/models/phi4mm/model.py b/areno/models/phi4mm/model.py index 6691a965..3f464b5c 100644 --- a/areno/models/phi4mm/model.py +++ b/areno/models/phi4mm/model.py @@ -1,4 +1,4 @@ -"""Phi-4-Multimodal language and vision adapter.""" +"""Native Phi-4-Multimodal language, vision, and audio adapter.""" from __future__ import annotations @@ -28,9 +28,11 @@ 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.audio import Phi4MMAudioConfig, Phi4MMAudioEmbedding from areno.models.phi4mm.vision import Phi4MMExtendedEmbedding, Phi4MMVisionConfig _IMAGE_SPECIAL_TOKEN_ID = 200010 +_AUDIO_SPECIAL_TOKEN_ID = 200011 def _phi4mm_vision_config(hf_config: dict[str, Any]) -> dict[str, Any] | None: @@ -73,6 +75,48 @@ def _phi4mm_vision_config(hf_config: dict[str, Any]) -> dict[str, Any] | None: return config +def _phi4mm_audio_config(hf_config: dict[str, Any]) -> dict[str, Any] | None: + embedding = hf_config.get("embd_layer") + audio_embedding = embedding.get("audio_embd_layer") if isinstance(embedding, dict) else None + processor = hf_config.get("audio_processor") + if not isinstance(audio_embedding, dict) or not isinstance(processor, dict): + return None + required_embedding = { + "embedding_cls": "audio", + "projection_cls": "mlp", + "compression_rate": 8, + "downsample_rate": 1, + "use_qformer": False, + "use_conv_downsample": False, + } + for key, expected in required_embedding.items(): + actual = audio_embedding.get(key) + if actual != expected: + raise ValueError(f"Phi4MM audio requires embd_layer.audio_embd_layer.{key}={expected!r}, got {actual!r}") + if processor.get("name") != "cascades" or not isinstance(processor.get("config"), dict): + raise ValueError("Phi4MM audio requires the cascades processor configuration") + values = dict(processor["config"]) + required = { + "input_layer": "nemo_conv", + "input_size": 80, + "attention_dim": 1024, + "attention_heads": 16, + "num_blocks": 24, + "time_reduction": 8, + "causal": True, + "activation": "swish", + "conv_activation": "swish", + "conv_glu_type": "swish", + "batch_norm": False, + } + for key, expected in required.items(): + if values.get(key) != expected: + raise ValueError( + f"Phi4MM audio requires audio_processor.config.{key}={expected!r}, got {values.get(key)!r}" + ) + return values + + 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: @@ -103,84 +147,101 @@ def _feature_tensor( if value is None: return None tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + if tensor.numel() == 0: + return None 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: +def _lora_config(config: ModelConfig, adapter: str) -> tuple[int, float, float] | None: + values = (config.hf_text_config or {}).get(f"{adapter}_lora") + modality_config = config.vision_config if adapter == "vision" else config.audio_config + if modality_config is None: return None if not isinstance(values, dict): - raise ValueError("Phi4MM vision support requires a vision_lora config") + raise ValueError(f"Phi4MM {adapter} support requires a {adapter}_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)") + raise ValueError(f"Phi4MM {adapter}_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.lora_scales: dict[str, float] = {} + self.lora_dropouts: dict[str, float] = {} self.vision_lora_mask: torch.Tensor | None = None + self.speech_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) + for adapter in ("vision", "speech"): + lora = _lora_config(config, adapter) + if lora is None: + continue + rank, self.lora_scales[adapter], self.lora_dropouts[adapter] = lora + self.lora_A[adapter] = nn.Linear(in_features, rank, bias=False) + self.lora_B[adapter] = 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 + self.lora_A[adapter].weight, False, sequence_parallel=False, tp_grad_allreduce=True ) - mark_tensor_parallel_parameter(self.lora_B["vision"].weight, True, sequence_parallel=True) + mark_tensor_parallel_parameter(self.lora_B[adapter].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: + if all(getattr(self, f"{adapter}_lora_mask") is None for adapter 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) + for adapter in self.lora_A: + mask = getattr(self, f"{adapter}_lora_mask") + if mask is None: + continue + dropped = F.dropout(full_input, p=self.lora_dropouts[adapter], training=self.training) + delta = self.lora_B[adapter](self.lora_A[adapter](dropped)) * self.lora_scales[adapter] + output = output + delta * mask.to(device=delta.device, dtype=delta.dtype).unsqueeze(-1) + return output 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.lora_scales: dict[str, float] = {} + self.lora_dropouts: dict[str, float] = {} self.vision_lora_mask: torch.Tensor | None = None + self.speech_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) + for adapter in ("vision", "speech"): + lora = _lora_config(config, adapter) + if lora is None: + continue + rank, self.lora_scales[adapter], self.lora_dropouts[adapter] = lora + self.lora_A[adapter] = nn.Linear(self.local_in_features, rank, bias=False) + self.lora_B[adapter] = nn.Linear(rank, out_features, bias=False) + mark_tensor_parallel_parameter(self.lora_A[adapter].weight, True, sequence_parallel=True) mark_tensor_parallel_parameter( - self.lora_B["vision"].weight, False, sequence_parallel=False, tp_grad_allreduce=True + self.lora_B[adapter].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 + for adapter in self.lora_A: + mask = getattr(self, f"{adapter}_lora_mask") + if mask is None: + continue + dropped = F.dropout(x, p=self.lora_dropouts[adapter], training=self.training) + latent = all_reduce(self.lora_A[adapter](dropped)) + delta = self.lora_B[adapter](latent) * self.lora_scales[adapter] + delta = delta * mask.to(device=delta.device, dtype=delta.dtype).unsqueeze(-1) + if is_sequence_parallel_active(): + delta = scatter_to_sequence_parallel_region(delta) + output = output + delta + return output def _require_bool(hf_config: dict[str, Any], key: str, expected: bool) -> None: @@ -378,18 +439,17 @@ def __init__(self, config: ModelConfig, layer_idx: int): 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: + if config.vision_config is not None or config.audio_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 set_lora_masks(self, vision: torch.Tensor | None, speech: torch.Tensor | None) -> None: + for module in (self.self_attn.qkv_proj, self.self_attn.o_proj, self.mlp.gate_up_proj, self.mlp.down_proj): + if hasattr(module, "vision_lora_mask"): + module.vision_lora_mask = vision + module.speech_lora_mask = speech def forward( self, @@ -407,25 +467,29 @@ def forward( class Phi4MMModel(nn.Module): - """Phi-4 transformer body with an optional native vision embedding path.""" + """Phi-4 transformer body with optional native vision and audio paths.""" 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 + self.embed_tokens_extend = None + if config.vision_config is not None or config.audio_config is not None: + vision_config = ( + Phi4MMVisionConfig.from_dict(config.vision_config) if config.vision_config is not None else None ) - if config.vision_config is not None - else None - ) + self.embed_tokens_extend = Phi4MMExtendedEmbedding(vision_config, config.hidden_size, config.dtype) + if config.audio_config is not None: + self.embed_tokens_extend.audio_embed = Phi4MMAudioEmbedding( + Phi4MMAudioConfig.from_dict(config.audio_config), config.hidden_size, config.dtype + ) 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) + self.register_buffer("speech_lora_slots", torch.empty(0, dtype=torch.bool), persistent=False) def forward( self, @@ -437,9 +501,9 @@ def forward( ) -> 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) + vision_lora_mask, speech_lora_mask = self._lora_masks(input_ids, features, train_meta, infer_meta) for layer in self.layers: - layer.set_vision_lora_mask(vision_lora_mask) + layer.set_lora_masks(vision_lora_mask, speech_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) @@ -458,52 +522,98 @@ def forward( ) return self.norm(hidden_states) - def _vision_lora_mask( + def _lora_masks( 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: + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: if self.embed_tokens_extend is None: - return None + return None, 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) + if infer_meta.recurrent_slots is None: + raise ValueError("Phi4MM multimodal decode requires recurrent modality slots") + vision = ( + self.vision_lora_slots.index_select(0, infer_meta.recurrent_slots).view_as(input_ids) + if self.vision_lora_slots.numel() + else None + ) + speech = ( + self.speech_lora_slots.index_select(0, infer_meta.recurrent_slots).view_as(input_ids) + if self.speech_lora_slots.numel() + else None + ) + return vision, speech image_mask = self._image_token_mask(input_ids, features) - explicit_modes = None + audio_mask = self._audio_token_mask(input_ids, features) + explicit_image_modes = None + explicit_audio_modes = None if isinstance(features, dict) and features.get("image_sequence_mask") is not None: - explicit_modes = torch.as_tensor( + explicit_image_modes = torch.as_tensor( features["image_sequence_mask"], device=input_ids.device, dtype=torch.bool ).reshape(-1) + if isinstance(features, dict) and features.get("audio_sequence_mask") is not None: + explicit_audio_modes = torch.as_tensor( + features["audio_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) + image_modes = explicit_image_modes if explicit_image_modes is not None else image_mask.any(dim=1) + audio_modes = explicit_audio_modes if explicit_audio_modes is not None else audio_mask.any(dim=1) + if int(image_modes.numel()) != int(input_ids.shape[0]) or int(audio_modes.numel()) != int( + input_ids.shape[0] + ): + raise ValueError("Phi4MM modality sequence masks must contain one value per input row") + # Official VISION_SPEECH mode uses the vision adapter and vision audio projector. + speech_modes = audio_modes & ~image_modes + return image_modes[:, None].expand_as(input_ids), speech_modes[:, None].expand_as(input_ids) else: - flat = image_mask.reshape(-1) - mask = torch.zeros_like(flat) - modes = [] + flat_image = image_mask.reshape(-1) + flat_audio = audio_mask.reshape(-1) + vision_mask = torch.zeros_like(flat_image) + speech_mask = torch.zeros_like(flat_audio) + vision_modes = [] + speech_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: + if explicit_image_modes is not None and int(explicit_image_modes.numel()) != sequence_count: raise ValueError("Phi4MM image_sequence_mask must contain one value per packed sequence") + if explicit_audio_modes is not None and int(explicit_audio_modes.numel()) != sequence_count: + raise ValueError("Phi4MM audio_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) + image_mode = ( + bool(explicit_image_modes[sequence_idx]) + if explicit_image_modes is not None + else bool(flat_image[start:end].any()) + ) + audio_mode = ( + bool(explicit_audio_modes[sequence_idx]) + if explicit_audio_modes is not None + else bool(flat_audio[start:end].any()) + ) + speech_mode = audio_mode and not image_mode + vision_modes.append(image_mode) + speech_modes.append(speech_mode) + vision_mask[start:end] = image_mode + speech_mask[start:end] = speech_mode 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 + self.vision_lora_slots.index_copy_( + 0, + infer_meta.recurrent_slots, + torch.tensor(vision_modes, device=self.vision_lora_slots.device, dtype=torch.bool), + ) + if infer_meta is not None and infer_meta.recurrent_slots is not None and self.speech_lora_slots.numel() > 0: + self.speech_lora_slots.index_copy_( + 0, + infer_meta.recurrent_slots, + torch.tensor(speech_modes, device=self.speech_lora_slots.device, dtype=torch.bool), + ) + return vision_mask.view_as(input_ids), speech_mask.view_as(input_ids) def _image_token_mask( self, @@ -516,6 +626,28 @@ def _image_token_mask( ) return input_ids == int(self.config.image_token_id or _IMAGE_SPECIAL_TOKEN_ID) + 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: + """Backward-compatible view of the vision half of modality LoRA state.""" + + return self._lora_masks(input_ids, features, train_meta, infer_meta)[0] + + def _audio_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("audio_token_mask") is not None: + return torch.as_tensor(features["audio_token_mask"], device=input_ids.device, dtype=torch.bool).view_as( + input_ids + ) + return input_ids == int(self.config.audio_token_id or _AUDIO_SPECIAL_TOKEN_ID) + @torch._dynamo.disable def _apply_multimodal_features( self, @@ -526,29 +658,34 @@ def _apply_multimodal_features( if features is None: return hidden_states if self.embed_tokens_extend is None: - raise ValueError("Phi4MM image features require a configured vision tower") + raise ValueError("Phi4MM multimodal features require a configured modality 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) + for modality, embeds in ( + ("image", self._project_image_feature_rows(row, hidden_states.device)), + ("audio", self._project_audio_feature_rows(row, hidden_states.device)), + ): + if embeds is None: + continue + mask = None if row.get(f"{modality}_feature_rows") is not None else row.get(f"{modality}_token_mask") + if mask is None: + default_id = self.config.image_token_id if modality == "image" else self.config.audio_token_id + fallback_id = _IMAGE_SPECIAL_TOKEN_ID if modality == "image" else _AUDIO_SPECIAL_TOKEN_ID + token_id = int(row.get(f"{modality}_token_id", default_id or fallback_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(f"Phi4MM {modality}_token_mask must match the input token row") + if int(mask.sum().item()) != int(embeds.shape[0]): + raise ValueError( + f"Phi4MM {modality} token count does not match projected embeddings: " + f"tokens={int(mask.sum().item())} embeds={int(embeds.shape[0])}" + ) + output[row_idx, mask] = 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: @@ -570,16 +707,60 @@ def _project_image_feature(self, features: dict[str, Any], device: torch.device) 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) + image_embed = getattr(self.embed_tokens_extend, "image_embed", None) + if image_embed is None: + raise ValueError("Phi4MM image features require a configured vision tower") + image_embeds = 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:] + def _project_audio_feature_rows(self, features: dict[str, Any], device: torch.device) -> torch.Tensor | None: + rows = features.get("audio_feature_rows") + if rows is not None: + pieces = [self._project_audio_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_audio_feature(features, device) + + def _project_audio_feature(self, features: dict[str, Any], device: torch.device) -> torch.Tensor | None: + existing = _feature_tensor(features, "audio_embeds", device, self.config.dtype) + if existing is not None: + return existing + inputs = _feature_tensor(features, "input_audio_embeds", device, self.config.dtype) + if inputs is None: + return None + audio_embed = getattr(self.embed_tokens_extend, "audio_embed", None) + if audio_embed is None: + raise ValueError("Phi4MM audio features require a configured audio tower") + attention_mask = _feature_tensor(features, "audio_attention_mask", device, torch.bool) + sizes = _feature_tensor(features, "audio_embed_sizes", device, torch.long) + if sizes is None: + raise ValueError("Phi4MM processor output requires audio_embed_sizes") + if inputs.ndim == 2: + inputs = inputs.unsqueeze(0) + if attention_mask is not None and attention_mask.ndim == 1: + attention_mask = attention_mask.unsqueeze(0) + input_mode = features.get("input_mode", 2) + if isinstance(input_mode, torch.Tensor): + input_mode = int(input_mode.reshape(-1)[0].item()) + projection_mode = "vision" if int(input_mode) in (1, 3) else "speech" + projected = audio_embed(inputs, attention_mask, projection_mode) + sizes_list = sizes.reshape(-1).detach().to(device="cpu", dtype=torch.long).tolist() + if len(sizes_list) != int(projected.shape[0]): + raise ValueError("Phi4MM audio_embed_sizes must contain one entry per audio segment") + merged = torch.cat([projected[index, : int(size)] for index, size in enumerate(sizes_list)], dim=0) + offset = int(features.get("audio_token_offset", 0) or 0) + count = features.get("audio_token_count") + if count is not None: + return merged[offset : offset + int(count)] + return merged[offset:] + class Phi4MMForCausalLM(nn.Module): - """Text-only Phi-4 causal LM with a truly tied vocab-parallel head.""" + """Phi-4-Multimodal causal LM with a truly tied vocab-parallel head.""" def __init__(self, config: ModelConfig): super().__init__() @@ -629,11 +810,14 @@ def set_kv_caches( 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) + self.model.speech_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) + if self.model.speech_lora_slots.numel() > 0: + self.model.speech_lora_slots.index_fill_(0, slots, False) @torch.no_grad() def prepare_infer_weights(self) -> None: @@ -744,6 +928,7 @@ def config_from_hf(self, hf_config: dict[str, Any]) -> ModelConfig: text_config["rope_scaling"] = rope_scaling text_config["original_max_position_embeddings"] = original_max_position_embeddings vision_config = _phi4mm_vision_config(hf_config) + audio_config = _phi4mm_audio_config(hf_config) return ModelConfig( model_type=self.name, @@ -769,7 +954,9 @@ def config_from_hf(self, hf_config: dict[str, Any]) -> ModelConfig: sequence_parallel=bool(hf_config.get("sequence_parallel", True)), hf_text_config=text_config, vision_config=vision_config, + audio_config=audio_config, image_token_id=_IMAGE_SPECIAL_TOKEN_ID if vision_config is not None else None, + audio_token_id=_AUDIO_SPECIAL_TOKEN_ID if audio_config is not None else None, ) def build(self, config: ModelConfig) -> nn.Module: diff --git a/areno/models/phi4mm/vision.py b/areno/models/phi4mm/vision.py index 0cf10f9e..dfdf21e1 100644 --- a/areno/models/phi4mm/vision.py +++ b/areno/models/phi4mm/vision.py @@ -248,6 +248,12 @@ def forward( class Phi4MMExtendedEmbedding(nn.Module): """Checkpoint-compatible container for Phi multimodal embedding modules.""" - def __init__(self, config: Phi4MMVisionConfig, language_hidden_size: int, dtype: torch.dtype): + def __init__( + self, + config: Phi4MMVisionConfig | None, + language_hidden_size: int, + dtype: torch.dtype, + ): super().__init__() - self.image_embed = Phi4MMImageEmbedding(config, language_hidden_size, dtype) + if config is not None: + self.image_embed = Phi4MMImageEmbedding(config, language_hidden_size, dtype) From f9189d0b58729db545588e2d17184d9616f65ee0 Mon Sep 17 00:00:00 2001 From: zitai-wang <2531131993@qq.com> Date: Wed, 26 Aug 2026 17:54:07 +0800 Subject: [PATCH 08/10] feat(runtime): support Phi-4 audio preprocessing and state --- areno/api/multimodal.py | 103 +++++++++++++++++++++++++++++ areno/engine/data/rollout_state.py | 74 ++++++++++++++++++++- 2 files changed, 176 insertions(+), 1 deletion(-) diff --git a/areno/api/multimodal.py b/areno/api/multimodal.py index b27a6c0f..a98182a0 100644 --- a/areno/api/multimodal.py +++ b/areno/api/multimodal.py @@ -7,6 +7,7 @@ import threading from collections.abc import Mapping, Sequence from typing import Any +from urllib.request import urlopen import torch @@ -119,6 +120,9 @@ def encode_processor_messages( """Use a native multimodal processor to load media and expand soft-token slots.""" normalized = _normalize_multimodal_messages(messages) + identity = f"{type(processor).__module__}.{type(processor).__name__}".lower() + if "phi4mm" in identity: + return _encode_phi4mm_messages(processor, normalized, tools=tools) kwargs: dict[str, Any] = { "tokenize": True, "add_generation_prompt": True, @@ -144,6 +148,94 @@ def encode_processor_messages( return tokens, features or None +def _encode_phi4mm_messages( + processor: Any, + messages: list[dict[str, Any]], + *, + tools: Any = None, +) -> tuple[list[int], dict[str, Any] | None]: + """Bridge structured API messages to the released Phi-4 processor API.""" + + images = [] + audios = [] + rendered_messages = [] + for message in messages: + rendered = dict(message) + content = rendered.get("content") + if isinstance(content, list): + pieces = [] + for part in content: + if not isinstance(part, dict): + pieces.append(str(part)) + continue + kind = str(part.get("type", "")) + if kind == "text": + pieces.append(str(part.get("text", ""))) + elif kind == "image": + images.append(_load_phi4mm_image(part.get("url"))) + pieces.append(f"<|image_{len(images)}|>") + elif kind == "audio": + audios.append(_load_phi4mm_audio(part.get("url"))) + pieces.append(f"<|audio_{len(audios)}|>") + else: + raise ValueError(f"Phi4MM does not support multimodal content type {kind!r}") + rendered["content"] = "".join(pieces) + rendered_messages.append(rendered) + template_kwargs: dict[str, Any] = {"tokenize": False, "add_generation_prompt": True} + if tools: + template_kwargs["tools"] = tools + prompt = apply_chat_template_with_options(processor.tokenizer, rendered_messages, **template_kwargs) + if prompt.endswith("<|endoftext|>"): + prompt = prompt.removesuffix("<|endoftext|>") + encoded = processor( + text=prompt, + images=images or None, + audios=audios or None, + return_tensors=getattr(processor, "_areno_return_tensors", "pt"), + ) + input_ids = encoded["input_ids"] + tokens = normalize_token_ids(input_ids[0].tolist()) + features = { + key: value for key, value in encoded.items() if key not in {"input_ids", "attention_mask", "token_type_ids"} + } + token_ids = modality_token_ids(processor) + features["modality_token_ids"] = token_ids + features["image_token_id"] = token_ids["image"] + features["audio_token_id"] = token_ids["audio"] + return tokens, features + + +def _load_phi4mm_image(reference: Any) -> Any: + if not isinstance(reference, str) or not reference: + raise ValueError("Phi4MM image content requires a URL or data URI") + if reference.startswith("data:"): + return _load_base64_image(reference) + try: + from PIL import Image + except ImportError as exc: + raise ValueError("Phi4MM image input requires Pillow") from exc + if reference.startswith(("http://", "https://")): + with urlopen(reference, timeout=30) as response: # noqa: S310 + return Image.open(io.BytesIO(response.read())).convert("RGB") + return Image.open(reference).convert("RGB") + + +def _load_phi4mm_audio(reference: Any) -> tuple[Any, int]: + if not isinstance(reference, str) or not reference: + raise ValueError("Phi4MM audio content requires a URL or data URI") + try: + import soundfile + except ImportError as exc: + raise ValueError("Phi4MM audio input requires soundfile") from exc + if reference.startswith("data:"): + _, _, payload = reference.partition(",") + return soundfile.read(io.BytesIO(base64.b64decode(payload))) + if reference.startswith(("http://", "https://")): + with urlopen(reference, timeout=30) as response: # noqa: S310 + return soundfile.read(io.BytesIO(response.read())) + return soundfile.read(reference) + + def _ensure_gemma4_torchvision_video_fps(processor: Any) -> None: """Backfill FPS metadata omitted by torchvision for some browser videos.""" @@ -194,6 +286,17 @@ def modality_token_ids(processor: Any) -> dict[str, int]: value = getattr(processor, f"{modality}_token_id", None) if isinstance(value, int) and value >= 0: result[modality] = int(value) + tokenizer = getattr(processor, "tokenizer", None) + if tokenizer is not None: + for modality, token in (("image", "<|endoftext10|>"), ("audio", "<|endoftext11|>")): + if modality in result: + continue + identity = f"{type(processor).__module__}.{type(processor).__name__}".lower() + if "phi4mm" not in identity: + continue + token_id = tokenizer.convert_tokens_to_ids(token) + if isinstance(token_id, int) and token_id >= 0: + result[modality] = token_id return result diff --git a/areno/engine/data/rollout_state.py b/areno/engine/data/rollout_state.py index 1fe71065..1e6f82d9 100644 --- a/areno/engine/data/rollout_state.py +++ b/areno/engine/data/rollout_state.py @@ -112,8 +112,10 @@ def build_prefill_payload(self) -> dict | None: mrope_position_parts: list[torch.Tensor] = [] has_mrope_positions = False feature_mask: list[bool] = [] + audio_feature_mask: list[bool] = [] image_features: list[dict] = [] image_sequence_modes: list[bool] = [] + audio_sequence_modes: list[bool] = [] cu_seqlens = [0] sample_indices: list[int] = [] block_table: list[list[int]] = [] @@ -156,8 +158,10 @@ def build_prefill_payload(self) -> dict | None: position_ids, mrope_position_parts if has_mrope_positions else None, feature_mask, + audio_feature_mask, image_features, image_sequence_modes, + audio_sequence_modes, cu_seqlens, sample_indices, block_table, @@ -177,7 +181,11 @@ def build_prefill_payload(self) -> dict | None: chunk_len, ) feature_mask.extend(local_mask) + audio_feature_mask.extend( + _prompt_modality_mask(self.prompt_features[seq_id], prompt, "audio")[cursor : cursor + chunk_len] + ) image_sequence_modes.append(_prompt_has_image(self.prompt_features[seq_id], prompt)) + audio_sequence_modes.append(_prompt_has_audio(self.prompt_features[seq_id], prompt)) if local_features is not None: image_features.append(local_features) local_mrope_positions = _slice_prompt_mrope_positions( @@ -221,8 +229,10 @@ def build_prefill_payload(self) -> dict | None: position_ids, mrope_position_parts if has_mrope_positions else None, feature_mask, + audio_feature_mask, image_features, image_sequence_modes, + audio_sequence_modes, cu_seqlens, sample_indices, block_table, @@ -238,8 +248,10 @@ def _prefill_payload( position_ids: list[int], mrope_position_parts: list[torch.Tensor] | None, feature_mask: list[bool], + audio_feature_mask: list[bool], image_features: list[dict], image_sequence_modes: list[bool], + audio_sequence_modes: list[bool], cu_seqlens: list[int], sample_indices: list[int], block_table: list[list[int]], @@ -261,12 +273,20 @@ 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 any(image_sequence_modes) or mrope_position_parts is not None: + if ( + any(feature_mask) + or image_features + or any(image_sequence_modes) + or any(audio_sequence_modes) + or mrope_position_parts is not None + ): payload["features"] = _prefill_multimodal_features( feature_mask, image_features, mrope_position_parts, image_sequence_modes, + audio_sequence_modes, + audio_feature_mask, ) return payload @@ -344,6 +364,10 @@ def _slice_prompt_image_features( "target_sizes", "pixel_values_videos", "input_features", + "input_audio_embeds", + "audio_embeds", + "audio_embed_sizes", + "audio_attention_mask", "multimodal_feature_rows", ) ): @@ -377,6 +401,11 @@ def _slice_prompt_image_features( "image_token_count": local_count, } ) + if modality_token_ids: + payload_features["image_token_offset"] = modality_offsets.get("image", 0) + payload_features["image_token_count"] = modality_counts.get("image", 0) + payload_features["audio_token_offset"] = modality_offsets.get("audio", 0) + payload_features["audio_token_count"] = modality_counts.get("audio", 0) for key in ( "pixel_values", "input_image_embeds", @@ -388,6 +417,11 @@ def _slice_prompt_image_features( "downsample_mode", "processor_expanded_image_tokens", "image_token_id", + "input_audio_embeds", + "audio_embed_sizes", + "audio_attention_mask", + "audio_token_id", + "input_mode", ): if features.get(key) is not None: payload_features[key] = features[key] @@ -429,10 +463,14 @@ def _prefill_multimodal_features( image_features: list[dict], mrope_position_parts: list[torch.Tensor] | None = None, image_sequence_modes: list[bool] | None = None, + audio_sequence_modes: list[bool] | None = None, + audio_feature_mask: 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 audio_sequence_modes is not None and any(audio_sequence_modes): + features["audio_sequence_mask"] = torch.tensor(audio_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: @@ -445,8 +483,13 @@ def _prefill_multimodal_features( { "image_token_mask": torch.tensor(feature_mask, dtype=torch.bool), "image_feature_rows": image_features, + "audio_feature_rows": image_features, } ) + if audio_feature_mask is not None: + audio_mask = torch.tensor(audio_feature_mask, dtype=torch.bool) + features["audio_token_mask"] = audio_mask + features["image_token_mask"] &= ~audio_mask return features @@ -487,6 +530,35 @@ def _prompt_has_image(features: dict | None, prompt: list[int]) -> bool: return image_token_id is not None and any(int(token) == int(image_token_id) for token in prompt) +def _prompt_has_audio(features: dict | None, prompt: list[int]) -> bool: + if features is None: + return False + mask = features.get("audio_token_mask") + if mask is not None: + return bool(torch.as_tensor(mask, dtype=torch.bool).any()) + audio_token_id = features.get("audio_token_id") + if audio_token_id is None: + audio_token_id = (features.get("modality_token_ids") or {}).get("audio") + return audio_token_id is not None and any(int(token) == int(audio_token_id) for token in prompt) + + +def _prompt_modality_mask(features: dict | None, prompt: list[int], modality: str) -> list[bool]: + if features is None: + return [False] * len(prompt) + mask = features.get(f"{modality}_token_mask") + if mask is not None: + values = [bool(item) for item in torch.as_tensor(mask).reshape(-1).tolist()] + if len(values) != len(prompt): + raise ValueError(f"{modality}_token_mask length must match prompt length") + return values + token_id = features.get(f"{modality}_token_id") + if token_id is None: + token_id = (features.get("modality_token_ids") or {}).get(modality) + if token_id is None: + return [False] * len(prompt) + return [int(token) == int(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.""" From 3874d7cbe7414d495174a25ea0d713faae99f6ab Mon Sep 17 00:00:00 2001 From: zitai-wang <2531131993@qq.com> Date: Wed, 26 Aug 2026 17:54:15 +0800 Subject: [PATCH 09/10] test(models): add Phi-4 audio coverage --- tests/test_phi4mm_audio_cpu.py | 365 +++++++++++++++++++++++++++++++++ 1 file changed, 365 insertions(+) create mode 100644 tests/test_phi4mm_audio_cpu.py diff --git a/tests/test_phi4mm_audio_cpu.py b/tests/test_phi4mm_audio_cpu.py new file mode 100644 index 00000000..065af833 --- /dev/null +++ b/tests/test_phi4mm_audio_cpu.py @@ -0,0 +1,365 @@ +from __future__ import annotations + +import pytest +import torch +from safetensors.torch import save_file + +from areno.api.multimodal import encode_processor_messages, modality_token_ids +from areno.engine.config import ModelConfig +from areno.engine.data.rollout_state import InferenceBatchState, payload_to_infer_meta +from areno.engine.parallel.context import TPContext, get_tp_context, set_tp_context +from areno.models.phi4mm.audio import Phi4MMAudioConfig, Phi4MMAudioEmbedding + + +@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_audio_config() -> Phi4MMAudioConfig: + return Phi4MMAudioConfig( + input_size=80, + attention_dim=8, + attention_heads=2, + linear_units=12, + num_blocks=2, + kernel_size=3, + time_reduction=8, + relative_attention_max_distance=8, + ) + + +def _tiny_model_config() -> ModelConfig: + return ModelConfig( + model_type="phi4mm", + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=4, + rms_norm_eps=1e-5, + rope_theta=10_000.0, + max_position_embeddings=32, + tie_word_embeddings=True, + qkv_bias=False, + qk_norm=False, + dtype=torch.float32, + hidden_act="silu", + partial_rotary_factor=0.5, + sequence_parallel=False, + attn_backend="native", + audio_config={ + "input_size": 80, + "attention_dim": 8, + "attention_heads": 2, + "linear_units": 12, + "num_blocks": 1, + "kernel_size": 3, + "time_reduction": 8, + "relative_attention_bias_args": {"t5_bias_max_distance": 8}, + }, + audio_token_id=200011, + hf_text_config={ + "original_max_position_embeddings": 16, + "rope_scaling": {"type": "longrope", "short_factor": (1.0,), "long_factor": (2.0,)}, + "speech_lora": {"r": 4, "lora_alpha": 8, "dp": 0.0}, + }, + ) + + +def _official_audio_sections() -> dict: + return { + "embd_layer": { + "audio_embd_layer": { + "embedding_cls": "audio", + "projection_cls": "mlp", + "compression_rate": 8, + "downsample_rate": 1, + "use_qformer": False, + "use_conv_downsample": False, + } + }, + "audio_processor": { + "name": "cascades", + "config": { + "input_layer": "nemo_conv", + "input_size": 80, + "attention_dim": 1024, + "attention_heads": 16, + "num_blocks": 24, + "time_reduction": 8, + "causal": True, + "activation": "swish", + "conv_activation": "swish", + "conv_glu_type": "swish", + "batch_norm": False, + }, + }, + } + + +def test_phi4mm_audio_config_accepts_only_official_semantics(): + from areno.models.phi4mm.model import _phi4mm_audio_config + + config = _official_audio_sections() + assert _phi4mm_audio_config(config)["num_blocks"] == 24 + + config["embd_layer"]["audio_embd_layer"]["compression_rate"] = 4 + with pytest.raises(ValueError, match="compression_rate=8"): + _phi4mm_audio_config(config) + + +def test_phi4mm_audio_encoder_reduces_time_and_respects_padding(): + module = Phi4MMAudioEmbedding(_tiny_audio_config(), language_hidden_size=16, dtype=torch.float32).eval() + inputs = torch.randn(2, 25, 80) + mask = torch.tensor([[True] * 25, [True] * 17 + [False] * 8]) + + output = module(inputs, mask) + + assert output.shape == (2, 4, 16) + assert torch.isfinite(output).all() + + +def test_phi4mm_audio_checkpoint_names_match_official_layout(): + module = Phi4MMAudioEmbedding(_tiny_audio_config(), language_hidden_size=16, dtype=torch.float32) + keys = set(module.state_dict()) + + assert "encoder.embed.conv.0.weight" in keys + assert "encoder.encoder_embedding.global_mean" in keys + assert "encoder.encoders.0.self_attn.linear_q.weight" in keys + assert "encoder.encoders.0.conv.glu.b1" in keys + assert "audio_projection.speech.2.weight" in keys + assert "audio_projection.vision.2.weight" in keys + + +def test_phi4mm_audio_relative_bias_clamps_distant_positions(): + from areno.models.phi4mm.audio import _T5RelativeAttentionLogitBias + + bias = _T5RelativeAttentionLogitBias(num_heads=2, max_distance=3, dtype=torch.float32) + output = bias(torch.zeros(1, 6, 8)) + + assert output.shape == (1, 2, 6, 6) + torch.testing.assert_close(output[:, :, 0, 2], output[:, :, 0, 5]) + torch.testing.assert_close(output[:, :, 3, 0], output[:, :, 5, 0]) + + +def test_phi4mm_audio_merge_replaces_only_audio_placeholders(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(_tiny_model_config()).float() + input_ids = torch.tensor([[1, 200011, 200011, 2]]) + hidden = torch.randn(1, 4, 16) + audio_embeds = torch.full((2, 16), 3.0) + + merged = model.model._apply_multimodal_features( + hidden, + input_ids, + {"audio_embeds": audio_embeds, "audio_token_id": 200011}, + ) + + torch.testing.assert_close(merged[0, 1:3], audio_embeds) + torch.testing.assert_close(merged[0, [0, 3]], hidden[0, [0, 3]]) + + +def test_phi4mm_audio_merge_rejects_placeholder_count_mismatch(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(_tiny_model_config()).float() + input_ids = torch.tensor([[1, 200011, 2]]) + + with pytest.raises(ValueError, match="audio token count does not match"): + model.model._apply_multimodal_features( + torch.randn(1, 3, 16), + input_ids, + {"audio_embeds": torch.ones(2, 16), "audio_token_id": 200011}, + ) + + +def test_phi4mm_speech_lora_state_survives_chunked_prefill_and_decode(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + features = { + "audio_embeds": torch.ones(2, 16), + "audio_token_id": 200011, + "modality_token_ids": {"audio": 200011}, + } + state = InferenceBatchState( + [[200011, 200011, 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"]["audio_sequence_mask"].tolist() == [True] + assert second["features"]["audio_sequence_mask"].tolist() == [True] + + model = Phi4MMAdapter().build(_tiny_model_config()).float() + model.model.speech_lora_slots = torch.zeros(1, dtype=torch.bool) + _, speech = model.model._lora_masks( + second["input_ids"].unsqueeze(0), + second["features"], + None, + payload_to_infer_meta(second, torch.device("cpu")), + ) + assert speech.tolist() == [[True, True]] + assert model.model.speech_lora_slots.tolist() == [True] + + +def test_phi4mm_multi_audio_merge_preserves_segment_order(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(_tiny_model_config()).float() + first = torch.full((2, 16), 1.0) + second = torch.full((3, 16), 2.0) + input_ids = torch.tensor([[1, 200011, 200011, 2, 200011, 200011, 200011, 3]]) + hidden = torch.randn(1, input_ids.shape[1], 16) + features = { + "audio_feature_rows": [ + {"audio_embeds": first, "audio_token_count": 2}, + {"audio_embeds": second, "audio_token_count": 3}, + ], + "audio_token_id": 200011, + } + + 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_text_audio_batch_uses_speech_lora_per_row(): + pytest.importorskip("triton") + from areno.models.phi4mm.model import Phi4MMAdapter + + model = Phi4MMAdapter().build(_tiny_model_config()).float() + input_ids = torch.tensor([[200011, 1, 2], [3, 4, 5]]) + _, speech = model.model._lora_masks( + input_ids, + [{"audio_embeds": torch.ones(1, 16), "audio_token_id": 200011}, None], + None, + None, + ) + + assert speech.tolist() == [[True, True, True], [False, False, False]] + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_phi4mm_speech_lora_tp_mapping_reconstructs_full_weights(tmp_path, tp_size): + pytest.importorskip("triton") + from areno.models.phi4mm.checkpoint import _load_vision_lora_weights + from areno.models.phi4mm.model import Phi4MMAdapter + + prefix = "model.layers.0" + tensors = { + f"{prefix}.self_attn.qkv_proj.lora_A.speech.weight": torch.arange(4 * 16).view(4, 16).float(), + f"{prefix}.self_attn.qkv_proj.lora_B.speech.weight": torch.arange(48 * 4).view(48, 4).float(), + f"{prefix}.self_attn.o_proj.lora_A.speech.weight": torch.arange(4 * 16).view(4, 16).float(), + f"{prefix}.self_attn.o_proj.lora_B.speech.weight": torch.arange(16 * 4).view(16, 4).float(), + f"{prefix}.mlp.gate_up_proj.lora_A.speech.weight": torch.arange(4 * 16).view(4, 16).float(), + f"{prefix}.mlp.gate_up_proj.lora_B.speech.weight": torch.arange(64 * 4).view(64, 4).float(), + f"{prefix}.mlp.down_proj.lora_A.speech.weight": torch.arange(4 * 32).view(4, 32).float(), + f"{prefix}.mlp.down_proj.lora_B.speech.weight": torch.arange(16 * 4).view(16, 4).float(), + } + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + 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(_tiny_model_config()).float() + _load_vision_lora_weights(model, checkpoint) + qkv = model.layers[0].self_attn.qkv_proj + expected = torch.cat( + [ + part.chunk(tp_size)[rank] + for part in tensors[f"{prefix}.self_attn.qkv_proj.lora_B.speech.weight"].split((16, 16, 16)) + ] + ) + torch.testing.assert_close(qkv.lora_B["speech"].weight, expected) + torch.testing.assert_close( + model.layers[0].self_attn.o_proj.lora_A["speech"].weight, + tensors[f"{prefix}.self_attn.o_proj.lora_A.speech.weight"].chunk(tp_size, dim=1)[rank], + ) + finally: + set_tp_context(previous) + + +def test_phi4mm_processor_token_fallback_finds_audio_token(): + class Tokenizer: + def convert_tokens_to_ids(self, token): + return {"<|endoftext10|>": 200010, "<|endoftext11|>": 200011}[token] + + class Phi4MMProcessor: + tokenizer = Tokenizer() + + Phi4MMProcessor.__module__ = "transformers_modules.phi4mm" + processor = Phi4MMProcessor() + + assert modality_token_ids(processor) == {"image": 200010, "audio": 200011} + + +def test_phi4mm_processor_bridge_numbers_audio_placeholders(monkeypatch): + class Tokenizer: + def convert_tokens_to_ids(self, token): + return {"<|endoftext10|>": 200010, "<|endoftext11|>": 200011}[token] + + def apply_chat_template(self, messages, *, tokenize, add_generation_prompt): + assert tokenize is False + assert add_generation_prompt is True + assert messages[0]["content"] == "<|audio_1|><|audio_2|>Compare" + return "rendered<|endoftext|>" + + class Phi4MMProcessor: + tokenizer = Tokenizer() + + def __call__(self, *, text, images, audios, return_tensors): + assert text == "rendered" + assert images is None + assert audios == ["first", "second"] + assert return_tensors == "pt" + return { + "input_ids": torch.tensor([[1, 200011, 200011, 2]]), + "attention_mask": torch.ones(1, 4), + "input_audio_embeds": torch.zeros(2, 8, 80), + "audio_embed_sizes": torch.tensor([1, 1]), + } + + Phi4MMProcessor.__module__ = "transformers_modules.phi4mm" + monkeypatch.setattr( + "areno.api.multimodal._load_phi4mm_audio", + lambda reference: {"one.wav": "first", "two.wav": "second"}[reference], + ) + tokens, features = encode_processor_messages( + Phi4MMProcessor(), + [ + { + "role": "user", + "content": [ + {"type": "audio", "url": "one.wav"}, + {"type": "audio", "url": "two.wav"}, + {"type": "text", "text": "Compare"}, + ], + } + ], + ) + + assert tokens == [1, 200011, 200011, 2] + assert features["audio_token_id"] == 200011 + assert features["input_audio_embeds"].shape == (2, 8, 80) From 2c586a9d1b5c9ea9df929b11266532c8f73d79cc Mon Sep 17 00:00:00 2001 From: zitai-wang <2531131993@qq.com> Date: Thu, 27 Aug 2026 10:47:57 +0800 Subject: [PATCH 10/10] fix(runtime): make CPU model imports Triton-optional --- areno/accel/ops.py | 44 +------------------------------- areno/accel/utils.py | 41 +++++++++++++++++++++++++++++ areno/engine/layers/mlp.py | 3 ++- areno/engine/layers/norm.py | 6 +++-- areno/models/phi4mm/model.py | 2 +- tests/test_phi4mm_adapter_cpu.py | 14 ++++++++++ 6 files changed, 63 insertions(+), 47 deletions(-) create mode 100644 areno/accel/utils.py 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/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/models/phi4mm/model.py b/areno/models/phi4mm/model.py index 3f464b5c..0bb51ca2 100644 --- a/areno/models/phi4mm/model.py +++ b/areno/models/phi4mm/model.py @@ -10,7 +10,7 @@ import torch.nn.functional as F from torch import nn -from areno.accel.ops import is_cuda_graph_capturing +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 diff --git a/tests/test_phi4mm_adapter_cpu.py b/tests/test_phi4mm_adapter_cpu.py index f4c6fcc5..c2fb7e0f 100644 --- a/tests/test_phi4mm_adapter_cpu.py +++ b/tests/test_phi4mm_adapter_cpu.py @@ -2,6 +2,8 @@ import json import math +import subprocess +import sys import pytest import torch @@ -95,6 +97,18 @@ def _tiny_model_config() -> ModelConfig: ) +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):