modeling_residualnet_v2.py

# coding: utf-8
from typing import Optional, List, Tuple
import torch
import torch.nn.functional as F
from torch import nn

from transformers.models.phi3.configuration_phi3 import Phi3Config
from transformers.models.phi3.modeling_phi3 import (
     Phi3PreTrainedModel,
     Phi3RMSNorm,
     Phi3MLP,
     Phi3Attention,
     # Phi3SdpaAttention,
     Phi3RotaryEmbedding,
)
#from models.phi3_config import Phi3Config
#from models.phi3 import (
#    Phi3PreTrainedModel,
#    Phi3RMSNorm,
#    Phi3MLP,
#    # Phi3SdpaAttention,
#    Phi3Attention,
#    Phi3RotaryEmbedding,
#)

from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
from transformers.modeling_outputs import BaseModelOutput, CausalLMOutputWithPast
from transformers.generation.utils import GenerationMixin


class ResidualNetV2Config(Phi3Config):
    model_type = "ResidualNetV2Config"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)


# ==============
# helpers (差分/マスク, 形状ユーティリティ, optional RoPE)
# ==============

def first_order_diff(x: torch.Tensor) -> torch.Tensor:
    # (B,L,H) -> (B,L-1,H)
    return x[:, 1:, :] - x[:, :-1, :]

def second_order_diff(x: torch.Tensor) -> torch.Tensor:
    # (B,L,H) -> (B,L-2,H): x_{t+2} - 2 x_{t+1} + x_t
    return x[:, 2:, :] - 2.0 * x[:, 1:-1, :] + x[:, :-2, :]

def mask_and(*tensors: torch.Tensor) -> torch.Tensor:
    # AND を丁寧に（float/bool いずれでも可）
    out = tensors[0].bool()
    for t in tensors[1:]:
        out = out & t.bool()
    return out.to(tensors[0].dtype)

def build_mask_for_diff(mask2d: Optional[torch.Tensor], order: int) -> Optional[torch.Tensor]:
    if mask2d is None:
        return None
    if order == 0:
        return mask2d
    elif order == 1:
        return mask_and(mask2d[:, 1:], mask2d[:, :-1])
    elif order == 2:
        return mask_and(mask2d[:, 2:], mask2d[:, 1:-1], mask2d[:, :-2])
    else:
        raise ValueError("order must be 0,1,2")

def shape_qkv(x: torch.Tensor, num_heads: int) -> torch.Tensor:
    # (B,L,H) -> (B,heads,L,head_dim)
    B, L, H = x.shape
    head_dim = H // num_heads
    x = x.view(B, L, num_heads, head_dim).transpose(1, 2)  # (B,heads,L,head_dim)
    return x

def unshape_ctx(x: torch.Tensor) -> torch.Tensor:
    # (B,heads,L,head_dim) -> (B,L,H)
    B, nH, L, d = x.shape
    return x.transpose(1, 2).contiguous().view(B, L, nH * d)

# RoPE helper（必要な場合のみ使用）
def rotate_half(x):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    rotary_dim = cos.shape[-1]
    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_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


# ==============
# Self-Block （Phi-3 部品そのまま）
# ==============

class Phi3SelfBlock(nn.Module):
    """
    PreNorm -> Self-Attn(SDPA+RoPE) -> resid -> PreNorm -> MLP -> resid
    """
    def __init__(self, config: ResidualNetV2Config, layer_idx: int, rotary_emb: Phi3RotaryEmbedding):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.in_norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # self.attn = Phi3SdpaAttention(config, layer_idx=layer_idx)
        self.attn = Phi3Attention(config, layer_idx=layer_idx)
        self.dropout_attn = nn.Dropout(config.resid_pdrop)
        self.ff_norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = Phi3MLP(config)
        self.dropout_mlp = nn.Dropout(config.resid_pdrop)
        self.rotary_emb = rotary_emb

    def _prepare_4d_mask(self, mask2d, bsz, seqlen, hidden_states):
        if mask2d is None:
            return None
        return _prepare_4d_causal_attention_mask(
            mask2d, (bsz, seqlen), hidden_states, past_key_values_length=0, sliding_window=self.config.sliding_window
        )

    def forward(
        self,
        hidden_states: torch.Tensor,         # (B,L,H)
        mask2d: Optional[torch.Tensor],      # (B,L)
        position_ids: Optional[torch.Tensor],# unused here; we re-make per len
        output_attentions: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        x = self.in_norm(hidden_states)
        B, L, _ = x.shape
        pos = torch.arange(L, device=x.device).unsqueeze(0).expand(B, -1)
        attn_mask = self._prepare_4d_mask(mask2d, B, L, x)
        position_embeddings = self.rotary_emb(hidden_states, pos)
        
        attn_out, attn_weights = self.attn(
            hidden_states=x,
            attention_mask=attn_mask,
            position_ids=pos,
            position_embeddings=position_embeddings,
            past_key_value=None,
            output_attentions=output_attentions,
            use_cache=False,
        )
        x = hidden_states + self.dropout_attn(attn_out)

        h = self.ff_norm(x)
        h = self.mlp(h)
        x = x + self.dropout_mlp(h)
        return x, (attn_weights if output_attentions else None)


# ==============
# Cross-Attention （シンプル実装／GQA対応。既定で RoPE 無し）
# ==============

class SimpleCrossAttention(nn.Module):
    """
    Query: x_q (B,Lq,H)   Key/Value: x_kv (B,Lk,H)
    - num_heads / num_key_value_heads は Phi-3 と同じ設定に合わせる（GQA対応）
    - 既定: RoPE 適用なし（decoder-encoder cross は相対位置の意味付けが曖昧なため）。
      use_rope_in_cross_attn=True で RoPE を適用可能。
    """
    def __init__(self, config: ResidualNetV2Config, rotary_emb: Phi3RotaryEmbedding, use_rope_in_cross_attn: bool = False):
        super().__init__()
        self.config = config
        self.rotary_emb = rotary_emb
        self.use_rope = use_rope_in_cross_attn

        H = config.hidden_size
        nH = config.num_attention_heads
        nKV = getattr(config, "num_key_value_heads", nH)
        self.nH = nH
        self.nKV = nKV
        self.groups = nH // nKV
        self.head_dim = H // nH

        self.q_proj = nn.Linear(H, H, bias=False)
        self.k_proj = nn.Linear(H, nKV * self.head_dim, bias=False)
        self.v_proj = nn.Linear(H, nKV * self.head_dim, bias=False)
        self.o_proj = nn.Linear(H, H, bias=False)
        self.dropout = nn.Dropout(config.attention_dropout)

        # 追加正規化（安定のため）
        self.q_norm = Phi3RMSNorm(H, eps=config.rms_norm_eps)
        self.kv_norm = Phi3RMSNorm(H, eps=config.rms_norm_eps)

    def _kv_repeat(self, x: torch.Tensor) -> torch.Tensor:
        # (B, nKV, Lk, d) -> (B, nH, Lk, d) へ繰り返し
        if self.nKV == self.nH:
            return x
        return x.repeat_interleave(self.groups, dim=1)

    def _make_attn_mask_bool(self, enc_mask2d: Optional[torch.Tensor], Lq: int, Lk: int, B: int) -> Optional[torch.Tensor]:
        # enc_mask2d: (B,Lk) in {0,1}  -> broadcastable bool mask of shape (B,1,Lq,Lk)
        if enc_mask2d is None:
            return None
        m = (~enc_mask2d.bool()).unsqueeze(1).unsqueeze(2)  # True=mask
        return m.expand(B, 1, Lq, Lk)

    def forward(
        self,
        x_q: torch.Tensor,             # (B,Lq,H)
        x_kv: torch.Tensor,            # (B,Lk,H)
        enc_mask2d: Optional[torch.Tensor] = None,  # (B,Lk)
    ) -> torch.Tensor:
        B, Lq, H = x_q.shape
        Lk = x_kv.size(1)

        q = self.q_proj(self.q_norm(x_q))                      # (B,Lq,H)
        k = self.k_proj(self.kv_norm(x_kv))                    # (B,Lk, nKV*Hd)
        v = self.v_proj(self.kv_norm(x_kv))                    # (B,Lk, nKV*Hd)

        q = shape_qkv(q, self.nH)                              # (B,nH,Lq,Hd)
        k = k.view(B, Lk, self.nKV, self.head_dim).transpose(1, 2)  # (B,nKV,Lk,Hd)
        v = v.view(B, Lk, self.nKV, self.head_dim).transpose(1, 2)  # (B,nKV,Lk,Hd)
        k = self._kv_repeat(k)                                 # (B,nH,Lk,Hd)
        v = self._kv_repeat(v)                                 # (B,nH,Lk,Hd)

        if self.use_rope:
            # 参考実装：各系列長に合わせて cos/sin を取り、q/k に適用
            # Phi3RotaryEmbedding の forward は (x, position_ids) -> (cos, sin) を返す前提
            pos_q = torch.arange(Lq, device=x_q.device).unsqueeze(0).expand(B, -1)
            pos_k = torch.arange(Lk, device=x_q.device).unsqueeze(0).expand(B, -1)
            # ダミーの [B,L,head_dim] を渡して cos/sin を得る（実装に依存するため try/except）
            try:
                dummy_q = torch.zeros(B, Lq, self.head_dim, device=x_q.device, dtype=x_q.dtype)
                dummy_k = torch.zeros(B, Lk, self.head_dim, device=x_q.device, dtype=x_q.dtype)
                cos_q, sin_q = self.rotary_emb(dummy_q, pos_q)
                cos_k, sin_k = self.rotary_emb(dummy_k, pos_k)
                q, k = apply_rotary_pos_emb(q, k, cos_q, sin_k, position_ids=None, unsqueeze_dim=2)
            except Exception:
                # RoPE 未対応環境では静かにスキップ（Self-Attn 側で RoPE が効いていれば全体としては相対位置信号を保持）
                pass

        # scaled dot-product attention
        attn_mask = self._make_attn_mask_bool(enc_mask2d, Lq, Lk, B)  # True=mask
        y = F.scaled_dot_product_attention(
            q, k, v, attn_mask=attn_mask, dropout_p=self.dropout.p if self.training else 0.0, is_causal=False
        )  # (B,nH,Lq,Hd)

        y = unshape_ctx(y)                                           # (B,Lq,H)
        y = self.o_proj(y)
        return y


# ==============
# v2: 3本の枝を「各3層」通してから、0階に Cross-Attn(←1階) → Cross-Attn(←2階)
# ==============

class ResidualNetV2Model(Phi3PreTrainedModel):
    """
    1) embedding -> 0/1/2階差分 3枝
    2) 各枝: [SelfBlock] x 3 （同一枝内で3層）
    3) x0 に CrossAttn(x1_final) → residual、続けて CrossAttn(x2_final) → residual
    出力は x0（原系列長 L）
    """
    def __init__(self, config: ResidualNetV2Config, use_rope_in_cross_attn: bool = False):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.rotary_emb = Phi3RotaryEmbedding(config=config)
        self.norm_out = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # 各枝 3 層
        self.branch0 = nn.ModuleList([Phi3SelfBlock(config, layer_idx=i, rotary_emb=self.rotary_emb) for i in range(3)])
        self.branch1 = nn.ModuleList([Phi3SelfBlock(config, layer_idx=100+i, rotary_emb=self.rotary_emb) for i in range(3)])
        self.branch2 = nn.ModuleList([Phi3SelfBlock(config, layer_idx=200+i, rotary_emb=self.rotary_emb) for i in range(3)])

        # Cross-Attn × 2 （0階 <- 1階, 0階 <- 2階）
        self.cross01_norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.cross01 = SimpleCrossAttention(config, self.rotary_emb, use_rope_in_cross_attn)
        self.cross12_norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.cross02 = SimpleCrossAttention(config, self.rotary_emb, use_rope_in_cross_attn)

        self.dropout = nn.Dropout(config.resid_pdrop)
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,   # (B,L)
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> BaseModelOutput:
        output_attentions = False if output_attentions is None else output_attentions
        output_hidden_states = False if output_hidden_states is None else output_hidden_states
        return_dict = True if return_dict is None else return_dict

        if inputs_embeds is None:
            x0 = self.embed_tokens(input_ids)  # (B,L,H)  0階
        else:
            x0 = inputs_embeds
        B, L, H = x0.shape

        m0 = attention_mask if attention_mask is not None else torch.ones(B, L, device=x0.device, dtype=torch.long)
        # 1階/2階差分
        x1 = first_order_diff(x0)      # (B,L-1,H)
        x2 = second_order_diff(x0)     # (B,L-2,H)
        m1 = build_mask_for_diff(m0, 1)
        m2 = build_mask_for_diff(m0, 2)

        # 各枝 3 層
        for blk in self.branch0:
            x0, _ = blk(x0, m0, None, output_attentions=False)
        for blk in self.branch1:
            x1, _ = blk(x1, m1, None, output_attentions=False)
        for blk in self.branch2:
            x2, _ = blk(x2, m2, None, output_attentions=False)

        # CrossAttn: x0 <- x1
        x0 = x0 + self.dropout(self.cross01(self.cross01_norm(x0), x1, enc_mask2d=m1))
        # CrossAttn: x0 <- x2
        x0 = x0 + self.dropout(self.cross02(self.cross12_norm(x0), x2, enc_mask2d=m2))

        x_out = self.norm_out(x0)

        if not return_dict:
            return (x_out,)

        return BaseModelOutput(
            last_hidden_state=x_out,
            hidden_states=None,
            attentions=None,
        )


class ResidualNetV2ForCausalLM(Phi3PreTrainedModel, GenerationMixin):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: ResidualNetV2Config, use_rope_in_cross_attn: bool = False):
        super().__init__(config)
        self.model = ResidualNetV2Model(config, use_rope_in_cross_attn=use_rope_in_cross_attn)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        # weight tying
        self.lm_head.weight = self.model.embed_tokens.weight
        self.post_init()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        inputs_embeds=None,
        labels: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=kwargs.get("output_attentions", False),
            output_hidden_states=kwargs.get("output_hidden_states", False),
            return_dict=True,
        )
        logits = self.lm_head(out.last_hidden_state).float()

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, self.vocab_size), shift_labels.view(-1))

        return CausalLMOutputWithPast(
            loss=loss, logits=logits, past_key_values=None, hidden_states=None, attentions=None
        )

    @property
    def base_model(self):
        return self.model


# ==============
# v3: 「各枝1層 + x0<-x1 Cross + x0<-x2 Cross」を1ブロックとして **3回** 反復
# ==============

class ResidualNetV3Model(Phi3PreTrainedModel):
    """
    1 block = { 3枝: SelfBlock各1層 → x0<-x1 Cross → x0<-x2 Cross }
    これを 3 回繰り返す（早期融合 + 反復洗練）
    """
    def __init__(self, config: ResidualNetV2Config, use_rope_in_cross_attn: bool = False):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.rotary_emb = Phi3RotaryEmbedding(config=config)
        self.norm_out = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.dropout = nn.Dropout(config.resid_pdrop)

        # 3 ブロック分の層を用意（枝それぞれ + CrossAttn×2）
        self.blocks_branch0 = nn.ModuleList([Phi3SelfBlock(config, layer_idx=10+i, rotary_emb=self.rotary_emb) for i in range(3)])
        self.blocks_branch1 = nn.ModuleList([Phi3SelfBlock(config, layer_idx=110+i, rotary_emb=self.rotary_emb) for i in range(3)])
        self.blocks_branch2 = nn.ModuleList([Phi3SelfBlock(config, layer_idx=210+i, rotary_emb=self.rotary_emb) for i in range(3)])

        self.cross_norm_01 = nn.ModuleList([Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) for _ in range(3)])
        self.cross_01 = nn.ModuleList([SimpleCrossAttention(config, self.rotary_emb, use_rope_in_cross_attn) for _ in range(3)])

        self.cross_norm_02 = nn.ModuleList([Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) for _ in range(3)])
        self.cross_02 = nn.ModuleList([SimpleCrossAttention(config, self.rotary_emb, use_rope_in_cross_attn) for _ in range(3)])

        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> BaseModelOutput:
        output_attentions = False if output_attentions is None else output_attentions
        output_hidden_states = False if output_hidden_states is None else output_hidden_states
        return_dict = True if return_dict is None else return_dict

        if inputs_embeds is None:
            x0 = self.embed_tokens(input_ids)
        else:
            x0 = inputs_embeds
        B, L, H = x0.shape
        m0 = attention_mask if attention_mask is not None else torch.ones(B, L, device=x0.device, dtype=torch.long)

        # 初回の差分（1,2階）は x0 から
        def mk_x1x2(x0, m0):
            return first_order_diff(x0), second_order_diff(x0), build_mask_for_diff(m0, 1), build_mask_for_diff(m0, 2)

        x1, x2, m1, m2 = mk_x1x2(x0, m0)

        # 3 ブロック反復
        for i in range(3):
            # 各枝 1 層
            x0, _ = self.blocks_branch0[i](x0, m0, None, output_attentions=False)
            x1, _ = self.blocks_branch1[i](x1, m1, None, output_attentions=False)
            x2, _ = self.blocks_branch2[i](x2, m2, None, output_attentions=False)

            # Cross: x0 <- x1, ついで x0 <- x2
            x0 = x0 + self.dropout(self.cross_01[i](self.cross_norm_01[i](x0), x1, enc_mask2d=m1))
            x0 = x0 + self.dropout(self.cross_02[i](self.cross_norm_02[i](x0), x2, enc_mask2d=m2))

            # 次ブロック用に 1/2階差分を「最新の x0」から再計算する手もある
            # （Down(2) 的な早期融合→再分解の設計に合わせたい場合）
            # ここでは「枝連鎖の継続」を優先し、x1/x2 は枝内の連続層として進める。
            # もし再分解を望むなら下記を有効化：
            # x1, x2, m1, m2 = mk_x1x2(x0, m0)

        x_out = self.norm_out(x0)

        if not return_dict:
            return (x_out,)

        return BaseModelOutput(
            last_hidden_state=x_out,
            hidden_states=None,
            attentions=None,
        )


class ResidualNetV3ForCausalLM(Phi3PreTrainedModel, GenerationMixin):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: ResidualNetV2Config, use_rope_in_cross_attn: bool = False):
        super().__init__(config)
        self.model = ResidualNetV3Model(config, use_rope_in_cross_attn=use_rope_in_cross_attn)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.lm_head.weight = self.model.embed_tokens.weight
        self.post_init()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        inputs_embeds=None,
        labels: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=kwargs.get("output_attentions", False),
            output_hidden_states=kwargs.get("output_hidden_states", False),
            return_dict=True,
        )
        logits = self.lm_head(out.last_hidden_state).float()

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, self.vocab_size), shift_labels.view(-1))

        return CausalLMOutputWithPast(
            loss=loss, logits=logits, past_key_values=None, hidden_states=None, attentions=None
        )

    @property
    def base_model(self):
        return self.model