1012 lines
44 KiB
Python
1012 lines
44 KiB
Python
|
# Copyright 2023 Baichuan Inc. All Rights Reserved.
|
|||
|
|
#
|
|||
|
|
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
|
|||
|
|
#
|
|||
|
|
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
|||
|
|
# and OPT implementations in this library. It has been modified from its
|
|||
|
|
# original forms to accommodate minor architectural differences compared
|
|||
|
|
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
|||
|
|
#
|
|||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|||
|
|
# you may not use this file except in compliance with the License.
|
|||
|
|
# You may obtain a copy of the License at
|
|||
|
|
#
|
|||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|||
|
|
#
|
|||
|
|
# Unless required by applicable law or agreed to in writing, software
|
|||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|||
|
|
# See the License for the specific language governing permissions and
|
|||
|
|
# limitations under the License.
|
|||
|
|
""" PyTorch omni model."""
|
|||
|
|
import os
|
|||
|
|
import time
|
|||
|
|
import json
|
|||
|
|
import math
|
|||
|
|
import numpy as np
|
|||
|
|
from typing import List, Optional, Tuple, Union, Any
|
|||
|
|
from threading import Thread
|
|||
|
|
from easydict import EasyDict
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
import torch.distributed
|
|||
|
|
import torch.utils.checkpoint
|
|||
|
|
from torch import nn
|
|||
|
|
from torch.nn import CrossEntropyLoss
|
|||
|
|
from torch.nn import functional as F
|
|||
|
|
import torch.distributed as dist
|
|||
|
|
from transformers import PreTrainedModel
|
|||
|
|
from transformers.activations import ACT2FN
|
|||
|
|
from dataclasses import dataclass
|
|||
|
|
from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
|
|||
|
|
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, ModelOutput
|
|||
|
|
from transformers.generation.utils import GenerationConfig
|
|||
|
|
from transformers.utils import logging
|
|||
|
|
# import for dynamic import not used in this file
|
|||
|
|
from .vector_quantize import VectorQuantize, EuclideanCodebook
|
|||
|
|
from .matcha_components import (
|
|||
|
|
SinusoidalPosEmb,
|
|||
|
|
Block1D,
|
|||
|
|
ResnetBlock1D,
|
|||
|
|
Downsample1D,
|
|||
|
|
TimestepEmbedding,
|
|||
|
|
Upsample1D,
|
|||
|
|
)
|
|||
|
|
from .matcha_transformer import BasicTransformerBlock
|
|||
|
|
from .flow_matching import ConditionalDecoder, ConditionalCFM
|
|||
|
|
|
|||
|
|
from .configuration_omni import OmniConfig
|
|||
|
|
from .audio_modeling_omni import (RMSNorm,
|
|||
|
|
OmniAudioEncoder,
|
|||
|
|
OmniAudioDecoder,
|
|||
|
|
OmniAudioVQBridgeTokenizer,
|
|||
|
|
OmniAudioFlowMatchingDecoder)
|
|||
|
|
from .visual_modeling_omni import OmniVisualEncoder, OmniVisualBridge
|
|||
|
|
from .processor_omni import OmniMMProcessor
|
|||
|
|
|
|||
|
|
# support model path contain point(.)
|
|||
|
|
try:
|
|||
|
|
# step1: copy relative imports to transformers_modules
|
|||
|
|
from .generation_utils import build_chat_input, TextIterStreamer
|
|||
|
|
from .sequence_parallel_utils import (
|
|||
|
|
create_attention_layer,
|
|||
|
|
get_sequence_parallel_size,
|
|||
|
|
get_sequence_parallel_chunk,
|
|||
|
|
)
|
|||
|
|
except ModuleNotFoundError:
|
|||
|
|
# step2: direct import from transformers_modules
|
|||
|
|
try: # bypass check_imports failure
|
|||
|
|
import sys
|
|||
|
|
sys.path.append(os.path.dirname(__file__))
|
|||
|
|
from generation_utils import build_chat_input, TextIterStreamer
|
|||
|
|
from sequence_parallel_utils import (
|
|||
|
|
create_attention_layer,
|
|||
|
|
get_sequence_parallel_size,
|
|||
|
|
get_sequence_parallel_chunk,
|
|||
|
|
)
|
|||
|
|
except Exception:
|
|||
|
|
raise
|
|||
|
|
|
|||
|
|
logger = logging.get_logger(__name__)
|
|||
|
|
|
|||
|
|
def get_slopes(n):
|
|||
|
|
def get_slopes_power_of_2(n):
|
|||
|
|
start = (2 ** (-2 ** -(math.log2(n) - 3)))
|
|||
|
|
ratio = start
|
|||
|
|
return [start * ratio ** i for i in range(n)]
|
|||
|
|
|
|||
|
|
if math.log2(n).is_integer():
|
|||
|
|
return get_slopes_power_of_2(
|
|||
|
|
n) # In the paper, we only train models that have 2^a heads for some a. This function has
|
|||
|
|
else: # some good properties that only occur when the input is a power of 2. To maintain that even
|
|||
|
|
closest_power_of_2 = 2 ** math.floor(
|
|||
|
|
math.log2(n)) # when the number of heads is not a power of 2, we use this workaround.
|
|||
|
|
return get_slopes_power_of_2(closest_power_of_2) + get_slopes(2 * closest_power_of_2)[0::2][
|
|||
|
|
:n - closest_power_of_2]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RotaryEmbedding(torch.nn.Module):
|
|||
|
|
def __init__(self, dim, max_position_embeddings=2048, base=5e6, device=None):
|
|||
|
|
super().__init__()
|
|||
|
|
# 修复RePE初始化精度问题 https://zhuanlan.zhihu.com/p/678963442
|
|||
|
|
# DeepSpeed 会 Hack torch.arange 强制在 GPU 上运行，这里使用原生的 torch.arange
|
|||
|
|
try:
|
|||
|
|
import deepspeed
|
|||
|
|
self.arange = deepspeed.runtime.zero.partition_parameters._orig_torch_arange
|
|||
|
|
except:
|
|||
|
|
self.arange = torch.arange
|
|||
|
|
|
|||
|
|
self.inv_freq = 1.0 / (base ** (self.arange(0, dim, 2).float().to(device) / dim))
|
|||
|
|
self.max_seq_len_cached = max_position_embeddings
|
|||
|
|
t = self.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
|
|||
|
|
freqs = torch.outer(t, self.inv_freq)
|
|||
|
|
emb = torch.cat((freqs, freqs), dim=-1)
|
|||
|
|
self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32)
|
|||
|
|
self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32)
|
|||
|
|
|
|||
|
|
def forward(self, x, seq_len=None):
|
|||
|
|
# x: [bs, num_attention_heads, seq_len, head_size]
|
|||
|
|
# This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
|
|||
|
|
if seq_len > self.max_seq_len_cached:
|
|||
|
|
self.max_seq_len_cached = seq_len
|
|||
|
|
t = self.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
|
|||
|
|
freqs = torch.outer(t, self.inv_freq)
|
|||
|
|
emb = torch.cat((freqs, freqs), dim=-1)
|
|||
|
|
self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32).to(x.device)
|
|||
|
|
self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32).to(x.device)
|
|||
|
|
return (
|
|||
|
|
self.cos_cached[:, :, :seq_len, ...].to(torch.float32).to(x.device),
|
|||
|
|
self.sin_cached[:, :, :seq_len, ...].to(torch.float32).to(x.device),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def rotate_half(x):
|
|||
|
|
"""Rotates half the hidden dims of the input."""
|
|||
|
|
x1 = x[..., : x.shape[-1] // 2]
|
|||
|
|
x2 = x[..., x.shape[-1] // 2:]
|
|||
|
|
return torch.cat((-x2, x1), dim=-1)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def apply_rotary_pos_emb(q, k, cos_, sin_, position_ids):
|
|||
|
|
cos = cos_.squeeze(1).squeeze(0) # [seq_len, dim]
|
|||
|
|
sin = sin_.squeeze(1).squeeze(0) # [seq_len, dim]
|
|||
|
|
cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
|
|||
|
|
sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
|
|||
|
|
q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
|
|||
|
|
k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
|
|||
|
|
return q_embed.to(q.dtype), k_embed.to(k.dtype)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MLP(nn.Module):
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
hidden_size: int,
|
|||
|
|
intermediate_size: int,
|
|||
|
|
hidden_act: str,
|
|||
|
|
):
|
|||
|
|
super().__init__()
|
|||
|
|
self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
|
|||
|
|
self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
|
|||
|
|
self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
|
|||
|
|
self.act_fn = ACT2FN[hidden_act]
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
|||
|
|
|
|||
|
|
# Copied from transformers.models.llama.modeling_llama.repeat_kv
|
|||
|
|
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
|||
|
|
"""
|
|||
|
|
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
|
|||
|
|
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
|
|||
|
|
"""
|
|||
|
|
batch, num_key_value_heads, slen, head_dim = hidden_states.shape
|
|||
|
|
if n_rep == 1:
|
|||
|
|
return hidden_states
|
|||
|
|
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
|
|||
|
|
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Attention(nn.Module):
|
|||
|
|
"""Multi-headed attention from 'Attention Is All You Need' paper"""
|
|||
|
|
def __init__(self, config: OmniConfig, is_sparse=False):
|
|||
|
|
super().__init__()
|
|||
|
|
self.config = config
|
|||
|
|
self.position_embedding_type = config.position_embedding_type.lower()
|
|||
|
|
self.num_kv_heads = config.num_key_value_heads
|
|||
|
|
self.head_dim = config.head_dim
|
|||
|
|
self.hidden_size = config.num_attention_heads * self.head_dim
|
|||
|
|
self.hidden_kv_size = self.num_kv_heads * self.head_dim
|
|||
|
|
|
|||
|
|
if is_sparse:
|
|||
|
|
self.num_heads = config.sparse_attention_heads
|
|||
|
|
assert self.num_kv_heads == config.num_attention_heads
|
|||
|
|
self.W_pack = nn.Linear(self.hidden_size, 3 * self.num_heads * self.head_dim, bias=config.attention_qkv_bias)
|
|||
|
|
self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
|
|||
|
|
else:
|
|||
|
|
self.num_heads = config.num_attention_heads
|
|||
|
|
if self.config.attention_qkv_pack:
|
|||
|
|
self.W_pack = nn.Linear(config.hidden_size, self.hidden_size + self.hidden_kv_size * 2, bias=config.attention_qkv_bias)
|
|||
|
|
else:
|
|||
|
|
self.q_proj = nn.Linear(config.hidden_size, self.hidden_size, bias=config.attention_qkv_bias)
|
|||
|
|
self.k_proj = nn.Linear(config.hidden_size, self.hidden_kv_size, bias=config.attention_qkv_bias)
|
|||
|
|
self.v_proj = nn.Linear(config.hidden_size, self.hidden_kv_size, bias=config.attention_qkv_bias)
|
|||
|
|
|
|||
|
|
self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
|
|||
|
|
|
|||
|
|
if self.position_embedding_type == 'rope':
|
|||
|
|
self.rotary_emb = RotaryEmbedding(
|
|||
|
|
dim=self.head_dim,
|
|||
|
|
max_position_embeddings=config.max_position_embeddings,
|
|||
|
|
base=config.get_rotary_base()
|
|||
|
|
)
|
|||
|
|
elif self.position_embedding_type == 'alibi':
|
|||
|
|
self.alibi_slopes = get_slopes(self.num_heads)
|
|||
|
|
self.attention = create_attention_layer(self.hidden_size, self.num_heads, self.head_dim)
|
|||
|
|
|
|||
|
|
def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
|
|||
|
|
return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
|
|||
|
|
|
|||
|
|
def _repeat_kv(self, hidden_states: torch.Tensor, num_heads: int) -> torch.Tensor:
|
|||
|
|
assert hidden_states.size(1) <= num_heads and num_heads % hidden_states.size(1) == 0
|
|||
|
|
return repeat_kv(hidden_states, num_heads // hidden_states.size(1))
|
|||
|
|
|
|||
|
|
def forward(
|
|||
|
|
self,
|
|||
|
|
hidden_states: torch.Tensor,
|
|||
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|||
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|||
|
|
seqlens: Optional[torch.IntTensor] = None,
|
|||
|
|
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
|||
|
|
output_attentions: bool = False,
|
|||
|
|
use_cache: bool = False,
|
|||
|
|
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
|
|||
|
|
bsz, q_len = hidden_states.shape[:2]
|
|||
|
|
|
|||
|
|
if self.config.attention_qkv_pack:
|
|||
|
|
proj = self.W_pack(hidden_states)
|
|||
|
|
query_states, key_states, value_states = proj.split([self.hidden_size, self.hidden_kv_size, self.hidden_kv_size], dim=-1)
|
|||
|
|
else:
|
|||
|
|
query_states = self.q_proj(hidden_states)
|
|||
|
|
key_states = self.k_proj(hidden_states)
|
|||
|
|
value_states = self.v_proj(hidden_states)
|
|||
|
|
|
|||
|
|
# (B, S, hidden_size) -> (B, num_heads, S, head_size)
|
|||
|
|
query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
|
|||
|
|
# (B, S, hidden_size) -> (B, num_kv_heads, S, head_size)
|
|||
|
|
key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
|
|||
|
|
value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
|
|||
|
|
|
|||
|
|
kv_seq_len = key_states.shape[-2]
|
|||
|
|
if past_key_value is not None:
|
|||
|
|
kv_seq_len += past_key_value[0].shape[-2]
|
|||
|
|
if self.position_embedding_type == 'rope':
|
|||
|
|
max_position = position_ids.max().item()+1 if position_ids is not None else kv_seq_len * get_sequence_parallel_size()
|
|||
|
|
cos, sin = self.rotary_emb(value_states, seq_len=max_position)
|
|||
|
|
query_states, key_states = apply_rotary_pos_emb(
|
|||
|
|
query_states, key_states, cos, sin,
|
|||
|
|
get_sequence_parallel_chunk(position_ids)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if past_key_value is not None:
|
|||
|
|
# reuse k, v, self_attention
|
|||
|
|
key_states = torch.cat([past_key_value[0], key_states], dim=2)
|
|||
|
|
value_states = torch.cat([past_key_value[1], value_states], dim=2)
|
|||
|
|
past_key_value = (key_states, value_states) if use_cache else None
|
|||
|
|
|
|||
|
|
# repeat k/v heads if n_kv_heads < n_heads
|
|||
|
|
key_states = self._repeat_kv(key_states, query_states.size(1))
|
|||
|
|
value_states = self._repeat_kv(value_states, query_states.size(1))
|
|||
|
|
|
|||
|
|
if seqlens is not None:
|
|||
|
|
seqlens = seqlens.to(dtype=torch.int32)
|
|||
|
|
max_seqlen = (seqlens[1:] - seqlens[:-1]).max().item()
|
|||
|
|
if self.position_embedding_type == 'alibi':
|
|||
|
|
alibi_slopes = torch.tensor(self.alibi_slopes, dtype=torch.float32).to(query_states.device)
|
|||
|
|
else:
|
|||
|
|
alibi_slopes = None
|
|||
|
|
attn_output = self.attention(
|
|||
|
|
query_states, key_states, value_states, seqlens, seqlens,
|
|||
|
|
max_seqlen, max_seqlen, causal=True, alibi_slopes=alibi_slopes, use_flash=True)
|
|||
|
|
else:
|
|||
|
|
attn_output = self.attention(
|
|||
|
|
query_states, key_states, value_states, attn_mask=attention_mask, use_flash=False)
|
|||
|
|
|
|||
|
|
attn_output = attn_output.reshape(bsz, q_len, -1)
|
|||
|
|
attn_output = self.o_proj(attn_output)
|
|||
|
|
|
|||
|
|
return attn_output, None, past_key_value
|
|||
|
|
|
|||
|
|
|
|||
|
|
class DecoderLayer(nn.Module):
|
|||
|
|
def __init__(self, config: OmniConfig, is_sparse=False):
|
|||
|
|
super().__init__()
|
|||
|
|
self.hidden_size = config.hidden_size
|
|||
|
|
self.self_attn = Attention(config=config, is_sparse=is_sparse)
|
|||
|
|
self.mlp = MLP(
|
|||
|
|
hidden_size=self.hidden_size,
|
|||
|
|
intermediate_size=config.intermediate_size,
|
|||
|
|
hidden_act=config.hidden_act,
|
|||
|
|
)
|
|||
|
|
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|||
|
|
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|||
|
|
|
|||
|
|
def forward(
|
|||
|
|
self,
|
|||
|
|
hidden_states: torch.Tensor,
|
|||
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|||
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|||
|
|
seqlens: Optional[torch.IntTensor] = None,
|
|||
|
|
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
|||
|
|
output_attentions: Optional[bool] = False,
|
|||
|
|
use_cache: Optional[bool] = False,
|
|||
|
|
group_index=None,
|
|||
|
|
) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
|
|||
|
|
|
|||
|
|
residual = hidden_states
|
|||
|
|
|
|||
|
|
hidden_states = self.input_layernorm(hidden_states)
|
|||
|
|
|
|||
|
|
# Self Attention
|
|||
|
|
hidden_states, self_attn_weights, present_key_value = self.self_attn(
|
|||
|
|
hidden_states=hidden_states,
|
|||
|
|
attention_mask=attention_mask,
|
|||
|
|
position_ids=position_ids,
|
|||
|
|
seqlens=seqlens,
|
|||
|
|
past_key_value=past_key_value,
|
|||
|
|
output_attentions=output_attentions,
|
|||
|
|
use_cache=use_cache,
|
|||
|
|
)
|
|||
|
|
hidden_states = residual + hidden_states
|
|||
|
|
|
|||
|
|
# Fully Connected
|
|||
|
|
residual = hidden_states
|
|||
|
|
hidden_states = self.post_attention_layernorm(hidden_states)
|
|||
|
|
hidden_states = self.mlp(hidden_states)
|
|||
|
|
hidden_states = residual + hidden_states
|
|||
|
|
|
|||
|
|
outputs = (hidden_states,)
|
|||
|
|
|
|||
|
|
if output_attentions:
|
|||
|
|
outputs += (self_attn_weights,)
|
|||
|
|
|
|||
|
|
if use_cache:
|
|||
|
|
outputs += (present_key_value,)
|
|||
|
|
|
|||
|
|
return outputs
|
|||
|
|
|
|||
|
|
|
|||
|
|
class OmniPreTrainedModel(PreTrainedModel):
|
|||
|
|
config_class = OmniConfig
|
|||
|
|
base_model_prefix = "model"
|
|||
|
|
supports_gradient_checkpointing = True
|
|||
|
|
_no_split_modules = ["DecoderLayer"]
|
|||
|
|
_keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
|
|||
|
|
|
|||
|
|
def _init_weights(self, module):
|
|||
|
|
std = self.config.initializer_range
|
|||
|
|
if isinstance(module, nn.Linear) or isinstance(module, nn.Conv1d) or isinstance(module, nn.ConvTranspose1d):
|
|||
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
|||
|
|
if module.bias is not None:
|
|||
|
|
module.bias.data.zero_()
|
|||
|
|
elif isinstance(module, nn.Embedding):
|
|||
|
|
module.weight.data.normal_(mean=0.0, std=std)
|
|||
|
|
if module.padding_idx is not None:
|
|||
|
|
module.weight.data[module.padding_idx].zero_()
|
|||
|
|
elif isinstance(module, nn.LayerNorm) or isinstance(module, nn.GroupNorm):
|
|||
|
|
module.weight.data.fill_(1.0)
|
|||
|
|
module.bias.data.zero_()
|
|||
|
|
elif isinstance(module, RMSNorm):
|
|||
|
|
module.weight.data.fill_(1.0)
|
|||
|
|
|
|||
|
|
def _set_gradient_checkpointing(self, module, value=False):
|
|||
|
|
if isinstance(module, OmniModel):
|
|||
|
|
module.gradient_checkpointing = value
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class OmniModelOutputWithPast(BaseModelOutputWithPast):
|
|||
|
|
audio_encoder_ret: Optional[Any] = None
|
|||
|
|
audio_decoder_ret: Optional[Any] = None
|
|||
|
|
|
|||
|
|
class OmniModel(OmniPreTrainedModel):
|
|||
|
|
def __init__(self, config: OmniConfig):
|
|||
|
|
super().__init__(config)
|
|||
|
|
self.padding_idx = config.pad_token_id
|
|||
|
|
self.vocab_size = config.vocab_size
|
|||
|
|
|
|||
|
|
if config.visual_config.enable:
|
|||
|
|
self.visual_model = OmniVisualEncoder(config.visual_config)
|
|||
|
|
self.visual_bridge_model = OmniVisualBridge(config.visual_config)
|
|||
|
|
if config.video_config.enable and not config.visual_config.enable: # in case 没有visual_config而只有video_config
|
|||
|
|
self.visual_model = OmniVisualEncoder(config.video_config)
|
|||
|
|
self.visual_bridge_model = OmniVisualBridge(config.video_config)
|
|||
|
|
|
|||
|
|
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
|||
|
|
self.layers = nn.ModuleList([
|
|||
|
|
DecoderLayer(config, is_sparse=layer_idx in config.sparse_attention_layers)
|
|||
|
|
for layer_idx in range(config.num_hidden_layers)
|
|||
|
|
])
|
|||
|
|
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
|||
|
|
|
|||
|
|
self.audio_embed_layers = nn.ModuleList([
|
|||
|
|
nn.Embedding(codedim + 1, config.hidden_size)
|
|||
|
|
for i, codedim in enumerate(config.audio_config.vq_config.codebook_sizes)
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
self.gradient_checkpointing = True
|
|||
|
|
# Initialize weights and apply final processing
|
|||
|
|
self.post_init()
|
|||
|
|
|
|||
|
|
def get_input_embeddings(self):
|
|||
|
|
return self.embed_tokens
|
|||
|
|
|
|||
|
|
def set_input_embeddings(self, value):
|
|||
|
|
self.embed_tokens = value
|
|||
|
|
|
|||
|
|
@torch.no_grad()
|
|||
|
|
def get_multimodal_mask(self, input_ids, pad_token_id, special_token_list):
|
|||
|
|
'''
|
|||
|
|
获取任意模态的特殊mask，包含以下
|
|||
|
|
1. pad mask 表示文本中图像/语音/视频模态提前留出的token位置
|
|||
|
|
2. special token mask 特殊token 例如对理解模型<start> <end> 不需要next token prediction
|
|||
|
|
3. embedding mask / lm_head mask 标记出特殊token在embedding中的mask
|
|||
|
|
'''
|
|||
|
|
pad_mask = torch.eq(input_ids, pad_token_id)
|
|||
|
|
sp_mask = torch.zeros_like(input_ids, dtype=torch.bool)
|
|||
|
|
lm_head_mask = torch.zeros([self.config.vocab_size, 1], dtype=torch.bool)
|
|||
|
|
for sp_id in special_token_list:
|
|||
|
|
sp_mask = torch.logical_or(sp_mask, torch.eq(input_ids, sp_id))
|
|||
|
|
lm_head_mask[sp_id, 0] = True
|
|||
|
|
return pad_mask, sp_mask, lm_head_mask
|
|||
|
|
|
|||
|
|
def get_multimodal_embed(
|
|||
|
|
self,
|
|||
|
|
input_ids,
|
|||
|
|
text_embedding, # 1. self.embed_tokens(input_ids) 2. 其他模态结果
|
|||
|
|
multimodal_embed,
|
|||
|
|
pad_token_id,
|
|||
|
|
fake_input,
|
|||
|
|
group_index=None, # 某种模态的编号
|
|||
|
|
):
|
|||
|
|
pad_mask, sp_mask, _ = self.get_multimodal_mask(input_ids, pad_token_id, self.config.multimodal_special_token_list)
|
|||
|
|
if not self.training: # 推理支持auto map 把多模态模块输出和input_ids 统一到一个device
|
|||
|
|
multimodal_embed = multimodal_embed.to(input_ids.device)
|
|||
|
|
if not fake_input: # 检查多模态token 和 pad mask数量一致 （不正确的截断会导致该问题）
|
|||
|
|
assert pad_mask.sum() == multimodal_embed.shape[0]
|
|||
|
|
else:
|
|||
|
|
assert pad_mask.sum() <= 0
|
|||
|
|
|
|||
|
|
# 合并 当前模态embeddings 和text embeddings
|
|||
|
|
input_ids = torch.where(pad_mask, torch.cumsum(pad_mask.view(-1).to(input_ids), dim=0).view(input_ids.shape)-1, input_ids)
|
|||
|
|
text_embedding = (1 - pad_mask.to(text_embedding)).unsqueeze(-1) * text_embedding # pad token位置填0
|
|||
|
|
multimodal_embedding = torch.embedding(multimodal_embed, input_ids * pad_mask) # 非 pad token 位置填idx=0位置结果
|
|||
|
|
multimodal_embedding = pad_mask.to(multimodal_embedding).unsqueeze(-1) * multimodal_embedding # 非pad token 位置填0
|
|||
|
|
final_embedding = multimodal_embedding.to(text_embedding) + text_embedding
|
|||
|
|
|
|||
|
|
if group_index is None:
|
|||
|
|
group_index = pad_mask.to(torch.int32)
|
|||
|
|
else:
|
|||
|
|
current_index = torch.max(group_index) + 1
|
|||
|
|
group_index += pad_mask.to(torch.int32) * current_index # 假设模态无重叠
|
|||
|
|
|
|||
|
|
return final_embedding, group_index
|
|||
|
|
|
|||
|
|
def get_visual_embed(
|
|||
|
|
self,
|
|||
|
|
input_ids,
|
|||
|
|
text_embedding, # 1. self.embed_tokens(input_ids) 2. 其他模态结果
|
|||
|
|
images = None,
|
|||
|
|
patch_nums = None,
|
|||
|
|
images_grid = None,
|
|||
|
|
videos = None,
|
|||
|
|
videos_patch_nums = None,
|
|||
|
|
videos_grid = None,
|
|||
|
|
group_index = None, # 某种模态的编号
|
|||
|
|
):
|
|||
|
|
if images is None or len(images) <= 0:
|
|||
|
|
images, images_grid, patch_nums = self.visual_model.fake_input(input_ids.device)
|
|||
|
|
image_fake_input = True
|
|||
|
|
else:
|
|||
|
|
image_fake_input = False
|
|||
|
|
|
|||
|
|
if videos is None or len(videos) <= 0 :
|
|||
|
|
videos, videos_grid, videos_patch_nums = self.visual_model.fake_input(input_ids.device)
|
|||
|
|
video_fake_input = True
|
|||
|
|
else:
|
|||
|
|
video_fake_input = False
|
|||
|
|
|
|||
|
|
visual_input = images + videos
|
|||
|
|
visual_grid = images_grid + videos_grid
|
|||
|
|
|
|||
|
|
visual_input = torch.cat(visual_input, dim=0)
|
|||
|
|
visual_grid = torch.tensor(np.array(visual_grid))
|
|||
|
|
|
|||
|
|
visual_embed = self.visual_model(visual_input, grid_thw=visual_grid)
|
|||
|
|
visual_embed = self.visual_bridge_model(visual_embed)
|
|||
|
|
|
|||
|
|
assert sum(patch_nums) + sum(videos_patch_nums) == visual_embed.shape[0]
|
|||
|
|
images_embed = visual_embed[:sum(patch_nums)]
|
|||
|
|
videos_embed = visual_embed[sum(patch_nums):]
|
|||
|
|
|
|||
|
|
final_embedding, group_index = self.get_multimodal_embed(input_ids, text_embedding, images_embed, self.config.visual_config.image_pad_token_id, image_fake_input, group_index=group_index)
|
|||
|
|
final_embedding, group_index = self.get_multimodal_embed(input_ids, final_embedding, videos_embed, self.config.video_config.video_place_token_id, video_fake_input, group_index=group_index)
|
|||
|
|
return final_embedding, group_index
|
|||
|
|
|
|||
|
|
|
|||
|
|
@torch.no_grad()
|
|||
|
|
def audio_fake_input(self, device):
|
|||
|
|
return torch.zeros(5, len(self.config.audio_config.vq_config.codebook_sizes), dtype=torch.int32, device=device)
|
|||
|
|
|
|||
|
|
def forward(
|
|||
|
|
self,
|
|||
|
|
input_ids: torch.LongTensor = None,
|
|||
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|||
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|||
|
|
seqlens: Optional[torch.IntTensor] = None,
|
|||
|
|
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
|||
|
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
|||
|
|
audios_tokens: Optional[List|torch.Tensor] = None, # 音频token bs*seqlen*vq_num
|
|||
|
|
images: Optional[List|torch.Tensor] = None,
|
|||
|
|
patch_nums: Optional[torch.Tensor] = None,
|
|||
|
|
images_grid: Optional[List|torch.Tensor] = None,
|
|||
|
|
videos: Optional[List|torch.Tensor] = None,
|
|||
|
|
videos_patch_nums: Optional[torch.Tensor] = None,
|
|||
|
|
videos_grid: Optional[List|torch.Tensor] = None,
|
|||
|
|
use_cache: Optional[bool] = None,
|
|||
|
|
output_attentions: Optional[bool] = None,
|
|||
|
|
output_hidden_states: Optional[bool] = None,
|
|||
|
|
return_dict: Optional[bool] = None,
|
|||
|
|
) -> Union[Tuple, OmniModelOutputWithPast]:
|
|||
|
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
|||
|
|
output_hidden_states = (
|
|||
|
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
|||
|
|
)
|
|||
|
|
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
|||
|
|
return_dict = True if (return_dict is not None or self.training) else self.config.use_return_dict
|
|||
|
|
|
|||
|
|
# retrieve input_ids and inputs_embeds
|
|||
|
|
if input_ids is not None and inputs_embeds is not None:
|
|||
|
|
raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
|
|||
|
|
elif input_ids is not None:
|
|||
|
|
batch_size, seq_length = input_ids.shape
|
|||
|
|
elif inputs_embeds is not None:
|
|||
|
|
batch_size, seq_length, _ = inputs_embeds.shape
|
|||
|
|
else:
|
|||
|
|
raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
|
|||
|
|
|
|||
|
|
seq_length_with_past = seq_length
|
|||
|
|
past_key_values_length = 0
|
|||
|
|
|
|||
|
|
if past_key_values is not None:
|
|||
|
|
past_key_values_length = past_key_values[0][0].shape[2]
|
|||
|
|
seq_length_with_past = seq_length_with_past + past_key_values_length
|
|||
|
|
|
|||
|
|
if position_ids is None:
|
|||
|
|
device = input_ids.device if input_ids is not None else inputs_embeds.device
|
|||
|
|
position_ids = torch.arange(
|
|||
|
|
past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
|
|||
|
|
)
|
|||
|
|
position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
|
|||
|
|
else:
|
|||
|
|
position_ids = position_ids.view(-1, seq_length).long()
|
|||
|
|
|
|||
|
|
group_index, audio_decoder_ret = None, None
|
|||
|
|
if inputs_embeds is None:
|
|||
|
|
sp_input_ids = get_sequence_parallel_chunk(input_ids)
|
|||
|
|
inputs_embeds = self.embed_tokens(sp_input_ids)
|
|||
|
|
if audios_tokens is None or len(audios_tokens) <= 0 :
|
|||
|
|
audios_tokens = torch.zeros(5, len(self.config.audio_config.vq_config.codebook_sizes), dtype=torch.int32, device=input_ids.device) # a fake input
|
|||
|
|
fake_input = True
|
|||
|
|
else:
|
|||
|
|
fake_input = False
|
|||
|
|
for i, audio_emb_layer in enumerate(self.audio_embed_layers):
|
|||
|
|
if i==0:
|
|||
|
|
audio_embs = audio_emb_layer(audios_tokens[..., i])
|
|||
|
|
else:
|
|||
|
|
audio_embs += audio_emb_layer(audios_tokens[..., i])
|
|||
|
|
inputs_embeds, group_index = self.get_multimodal_embed(sp_input_ids, inputs_embeds, audio_embs, self.config.audio_config.audio_pad_token_id, fake_input, group_index=group_index)
|
|||
|
|
|
|||
|
|
if self.config.visual_config.enable or self.config.video_config.enable:
|
|||
|
|
inputs_embeds, group_index = self.get_visual_embed(sp_input_ids, inputs_embeds, images, patch_nums, images_grid, videos, videos_patch_nums, videos_grid, group_index=group_index) # 注意更新group index
|
|||
|
|
|
|||
|
|
if seqlens is not None and seqlens.ndim == 2:
|
|||
|
|
cu_seqlens = []
|
|||
|
|
offset, seqlen = 0, seqlens.size(1)
|
|||
|
|
for lens in seqlens:
|
|||
|
|
cu_seqlens.append(offset)
|
|||
|
|
cu_seqlens.extend((lens[(lens > 0) & (lens < seqlen)] + offset).tolist())
|
|||
|
|
offset += seqlen
|
|||
|
|
cu_seqlens.append(offset)
|
|||
|
|
seqlens = torch.tensor(cu_seqlens, dtype=seqlens.dtype, device=seqlens.device)
|
|||
|
|
elif seqlens is None and self.training:
|
|||
|
|
seqlens = torch.arange(
|
|||
|
|
end=input_ids.size(0) + 1,
|
|||
|
|
dtype=torch.int32,
|
|||
|
|
device=input_ids.device
|
|||
|
|
) * input_ids.size(1)
|
|||
|
|
if seqlens is not None:
|
|||
|
|
attention_mask = None # unset attention_mask to save memory
|
|||
|
|
|
|||
|
|
if seqlens is None and attention_mask is None:
|
|||
|
|
attention_mask = torch.ones(
|
|||
|
|
(batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
|
|||
|
|
)
|
|||
|
|
if attention_mask is not None:
|
|||
|
|
attention_mask = _prepare_4d_causal_attention_mask(
|
|||
|
|
attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# embed positions
|
|||
|
|
hidden_states = inputs_embeds
|
|||
|
|
|
|||
|
|
if self.gradient_checkpointing and self.training:
|
|||
|
|
if use_cache:
|
|||
|
|
logger.warning_once(
|
|||
|
|
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
|
|||
|
|
)
|
|||
|
|
use_cache = False
|
|||
|
|
|
|||
|
|
# decoder layers
|
|||
|
|
all_hidden_states = () if output_hidden_states else None
|
|||
|
|
all_self_attns = () if output_attentions else None
|
|||
|
|
next_decoder_cache = () if use_cache else None
|
|||
|
|
|
|||
|
|
for idx, decoder_layer in enumerate(self.layers):
|
|||
|
|
if output_hidden_states:
|
|||
|
|
all_hidden_states += (hidden_states,)
|
|||
|
|
|
|||
|
|
past_key_value = past_key_values[idx] if past_key_values is not None else None
|
|||
|
|
|
|||
|
|
if self.gradient_checkpointing and self.training:
|
|||
|
|
|
|||
|
|
def create_custom_forward(module):
|
|||
|
|
def custom_forward(*inputs):
|
|||
|
|
# None for past_key_value
|
|||
|
|
return module(*inputs, output_attentions, False, group_index)
|
|||
|
|
|
|||
|
|
return custom_forward
|
|||
|
|
|
|||
|
|
layer_outputs = torch.utils.checkpoint.checkpoint(
|
|||
|
|
create_custom_forward(decoder_layer),
|
|||
|
|
hidden_states,
|
|||
|
|
attention_mask,
|
|||
|
|
position_ids,
|
|||
|
|
seqlens,
|
|||
|
|
None,
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
layer_outputs = decoder_layer(
|
|||
|
|
hidden_states,
|
|||
|
|
attention_mask=attention_mask,
|
|||
|
|
position_ids=position_ids,
|
|||
|
|
seqlens=seqlens,
|
|||
|
|
past_key_value=past_key_value,
|
|||
|
|
output_attentions=output_attentions,
|
|||
|
|
use_cache=use_cache,
|
|||
|
|
group_index=group_index,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
hidden_states = layer_outputs[0]
|
|||
|
|
|
|||
|
|
if use_cache:
|
|||
|
|
next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
|
|||
|
|
|
|||
|
|
if output_attentions:
|
|||
|
|
all_self_attns += (layer_outputs[1],)
|
|||
|
|
|
|||
|
|
hidden_states = self.norm(hidden_states)
|
|||
|
|
|
|||
|
|
# add hidden states from the last decoder layer
|
|||
|
|
if output_hidden_states:
|
|||
|
|
all_hidden_states += (hidden_states,)
|
|||
|
|
|
|||
|
|
next_cache = next_decoder_cache if use_cache else None
|
|||
|
|
if not return_dict:
|
|||
|
|
return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
|
|||
|
|
return BaseModelOutputWithPast(
|
|||
|
|
last_hidden_state=hidden_states,
|
|||
|
|
past_key_values=next_cache,
|
|||
|
|
hidden_states=all_hidden_states,
|
|||
|
|
attentions=all_self_attns,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class NormHead(nn.Module):
|
|||
|
|
def __init__(self, hidden_size, vocab_size, bias=False):
|
|||
|
|
super().__init__()
|
|||
|
|
self.hidden_size = hidden_size
|
|||
|
|
self.vocab_size = vocab_size
|
|||
|
|
self.weight = nn.Parameter(torch.empty((self.vocab_size, self.hidden_size)))
|
|||
|
|
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
|||
|
|
|
|||
|
|
def forward(self, hidden_states, mask=None):
|
|||
|
|
norm_weight = nn.functional.normalize(self.weight)
|
|||
|
|
if mask is not None:
|
|||
|
|
mask = mask.to(norm_weight)
|
|||
|
|
norm_weight = norm_weight * mask + (1 - mask) * norm_weight.detach()
|
|||
|
|
return nn.functional.linear(hidden_states, norm_weight)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def extra_repr(self) -> str:
|
|||
|
|
return f'in_features={self.hidden_size}, out_features={self.vocab_size}'
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class OmniMMCausalLMOutputWithPast(ModelOutput):
|
|||
|
|
loss: Optional[torch.FloatTensor] = None
|
|||
|
|
logits: Optional[torch.FloatTensor] = None
|
|||
|
|
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
|||
|
|
hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
|||
|
|
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
|||
|
|
audios_emb_for_infer: Optional[torch.FloatTensor] = None # 用于audio head 推理的 embeddings
|
|||
|
|
|
|||
|
|
|
|||
|
|
class CasualDepthTransformerLayer(nn.Module):
|
|||
|
|
def __init__(self, config, depth):
|
|||
|
|
super().__init__()
|
|||
|
|
self.config = config
|
|||
|
|
embed_size = config.hidden_size
|
|||
|
|
assert embed_size % 128 == 0
|
|||
|
|
num_heads = embed_size // 128
|
|||
|
|
self.self_attention = nn.MultiheadAttention(embed_dim=embed_size, num_heads=num_heads,batch_first=True)
|
|||
|
|
self.layernorm1 = RMSNorm(embed_size)
|
|||
|
|
self.layernorm2 = RMSNorm(embed_size)
|
|||
|
|
self.linear1 = nn.Linear(embed_size * depth, 2 * embed_size)
|
|||
|
|
self.linear2 = nn.Linear(2 * embed_size * depth, embed_size)
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
seq_len = x.size(1)
|
|||
|
|
res = x
|
|||
|
|
x = self.layernorm1(x)
|
|||
|
|
src_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool().to(x.device)
|
|||
|
|
_x, _ = self.self_attention(x, x, x, is_causal=True, attn_mask=src_mask)
|
|||
|
|
res = _x + res # (bs, sl, d)
|
|||
|
|
res = self.layernorm2(res)
|
|||
|
|
x = torch.einsum('bld,tld->blt', res, torch.reshape(self.linear1.weight, (2 * self.config.hidden_size, -1, self.config.hidden_size)))
|
|||
|
|
x = torch.nn.functional.gelu(x)
|
|||
|
|
x = torch.einsum('blt,dlt->bld', x, torch.reshape(self.linear2.weight, (self.config.hidden_size, -1, 2 * self.config.hidden_size)))
|
|||
|
|
return res + x
|
|||
|
|
|
|||
|
|
class OmniAudioHead(nn.Module):
|
|||
|
|
def __init__(self, config):
|
|||
|
|
super().__init__()
|
|||
|
|
self.config = config
|
|||
|
|
hidden_size = config.hidden_size
|
|||
|
|
self.transformer_layers = nn.ModuleList([
|
|||
|
|
CasualDepthTransformerLayer(config, len(config.audio_config.vq_config.codebook_sizes))
|
|||
|
|
for _ in range(config.audio_config.audio_head_transformer_layers)
|
|||
|
|
])
|
|||
|
|
self.headnorm = RMSNorm(hidden_size)
|
|||
|
|
self.heads = nn.ModuleList([
|
|||
|
|
nn.Linear(hidden_size, vq_size+1)
|
|||
|
|
for vq_size in config.audio_config.vq_config.codebook_sizes
|
|||
|
|
])
|
|||
|
|
self.gradient_checkpointing = True
|
|||
|
|
|
|||
|
|
def forward(self, x, audios_tokens, audio_emb_layers):
|
|||
|
|
cumsum_audio_embed = torch.stack([
|
|||
|
|
audio_emb_layers[i](audios_tokens[..., i])
|
|||
|
|
for i, vq_size in enumerate(self.config.audio_config.vq_config.codebook_sizes[:-1])
|
|||
|
|
], dim=1)
|
|||
|
|
cumsum_audio_embed = torch.cumsum(cumsum_audio_embed, dim=1) # (bs, depth-1, d)
|
|||
|
|
hidden_states = torch.concat([x.reshape(-1, 1, self.config.hidden_size), cumsum_audio_embed], dim=1) # (bs, depth, d)
|
|||
|
|
assert hidden_states.size(1) == len(self.config.audio_config.vq_config.codebook_sizes)
|
|||
|
|
for i, tlayer in enumerate(self.transformer_layers):
|
|||
|
|
hidden_states = tlayer(hidden_states,)
|
|||
|
|
hidden_states = self.headnorm(hidden_states)
|
|||
|
|
logits = [head(hidden_states[:,i]) for i, head in enumerate(self.heads)]
|
|||
|
|
return logits
|
|||
|
|
|
|||
|
|
|
|||
|
|
class OmniForCausalLM(OmniPreTrainedModel):
|
|||
|
|
def __init__(self, config):
|
|||
|
|
super().__init__(config)
|
|||
|
|
self.config = config
|
|||
|
|
self.model = OmniModel(config)
|
|||
|
|
self.audio_tokenizer = OmniAudioTokenizer(config)
|
|||
|
|
self.audio_head = OmniAudioHead(config)
|
|||
|
|
if config.use_norm_head:
|
|||
|
|
self.lm_head = NormHead(config.hidden_size, config.vocab_size, bias=False)
|
|||
|
|
else:
|
|||
|
|
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
|||
|
|
# Initialize weights and apply final processing
|
|||
|
|
self.post_init()
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def main_device(self):
|
|||
|
|
return self.lm_head.weight.device
|
|||
|
|
|
|||
|
|
def bind_processor(self, tokenizer, **kwargs):
|
|||
|
|
self.processor = OmniMMProcessor(
|
|||
|
|
tokenizer=tokenizer,
|
|||
|
|
config=self.config,
|
|||
|
|
**kwargs,
|
|||
|
|
)
|
|||
|
|
return self.processor
|
|||
|
|
|
|||
|
|
def get_input_embeddings(self):
|
|||
|
|
return self.model.embed_tokens
|
|||
|
|
|
|||
|
|
def set_input_embeddings(self, value):
|
|||
|
|
self.model.embed_tokens = value
|
|||
|
|
|
|||
|
|
def get_output_embeddings(self):
|
|||
|
|
return self.lm_head
|
|||
|
|
|
|||
|
|
def set_output_embeddings(self, new_embeddings):
|
|||
|
|
self.lm_head = new_embeddings
|
|||
|
|
|
|||
|
|
def set_decoder(self, decoder):
|
|||
|
|
self.model = decoder
|
|||
|
|
|
|||
|
|
def get_decoder(self):
|
|||
|
|
return self.model
|
|||
|
|
|
|||
|
|
def forward(
|
|||
|
|
self,
|
|||
|
|
input_ids: torch.LongTensor = None,
|
|||
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|||
|
|
position_ids: Optional[torch.LongTensor] = None,
|
|||
|
|
seqlens: Optional[torch.IntTensor] = None,
|
|||
|
|
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
|||
|
|
inputs_embeds: Optional[torch.FloatTensor] = None,
|
|||
|
|
labels: Optional[torch.LongTensor] = None,
|
|||
|
|
audios: Optional[List|torch.Tensor] = None,
|
|||
|
|
audios_tokens: Optional[List|torch.Tensor] = None,
|
|||
|
|
encoder_length: Optional[torch.Tensor] = None,
|
|||
|
|
bridge_length: Optional[torch.Tensor] = None,
|
|||
|
|
images: Optional[torch.Tensor] = None,
|
|||
|
|
patch_nums: Optional[torch.Tensor] = None,
|
|||
|
|
images_grid: Optional[torch.Tensor] = None,
|
|||
|
|
videos: Optional[torch.Tensor] = None,
|
|||
|
|
videos_patch_nums: Optional[torch.Tensor] = None,
|
|||
|
|
videos_grid: Optional[torch.Tensor] = None,
|
|||
|
|
use_cache: Optional[bool] = None,
|
|||
|
|
output_attentions: Optional[bool] = None,
|
|||
|
|
output_hidden_states: Optional[bool] = None,
|
|||
|
|
return_dict: Optional[bool] = None,
|
|||
|
|
) -> Union[Tuple, CausalLMOutputWithPast]:
|
|||
|
|
|
|||
|
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
|||
|
|
output_hidden_states = (
|
|||
|
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
|||
|
|
)
|
|||
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|||
|
|
|
|||
|
|
if audios_tokens is not None:
|
|||
|
|
assert isinstance(audios_tokens, torch.Tensor)
|
|||
|
|
else:
|
|||
|
|
if audios is None or len(audios) == 0:
|
|||
|
|
audios_tokens = None
|
|||
|
|
else:
|
|||
|
|
audios_tokens = self.audio_tokenizer(audios,encoder_length,bridge_length)
|
|||
|
|
|
|||
|
|
outputs = self.model(
|
|||
|
|
input_ids=input_ids,
|
|||
|
|
attention_mask=attention_mask,
|
|||
|
|
position_ids=position_ids,
|
|||
|
|
seqlens=seqlens,
|
|||
|
|
past_key_values=past_key_values,
|
|||
|
|
inputs_embeds=inputs_embeds,
|
|||
|
|
audios_tokens=audios_tokens,
|
|||
|
|
images=images,
|
|||
|
|
patch_nums=patch_nums,
|
|||
|
|
images_grid=images_grid,
|
|||
|
|
videos=videos,
|
|||
|
|
videos_patch_nums=videos_patch_nums,
|
|||
|
|
videos_grid=videos_grid,
|
|||
|
|
use_cache=use_cache,
|
|||
|
|
output_attentions=output_attentions,
|
|||
|
|
output_hidden_states=output_hidden_states,
|
|||
|
|
return_dict=return_dict,
|
|||
|
|
)
|
|||
|
|
hidden_states = outputs.last_hidden_state
|
|||
|
|
audios_emb_for_infer = hidden_states[:,-1,:]
|
|||
|
|
logits = self.lm_head(hidden_states)
|
|||
|
|
|
|||
|
|
return OmniMMCausalLMOutputWithPast(
|
|||
|
|
logits=logits,
|
|||
|
|
past_key_values=outputs.past_key_values,
|
|||
|
|
hidden_states=outputs.hidden_states,
|
|||
|
|
attentions=outputs.attentions,
|
|||
|
|
audios_emb_for_infer=audios_emb_for_infer
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def prepare_inputs_for_generation(
|
|||
|
|
self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
|
|||
|
|
):
|
|||
|
|
if past_key_values:
|
|||
|
|
input_ids = input_ids[:, past_key_values[0][0].shape[-2]:]
|
|||
|
|
|
|||
|
|
position_ids = kwargs.get("position_ids", None)
|
|||
|
|
if attention_mask is not None and position_ids is None:
|
|||
|
|
# create position_ids on the fly for batch generation
|
|||
|
|
position_ids = attention_mask.long().cumsum(-1)
|
|||
|
|
# position_ids.masked_fill_(attention_mask == 0, 1)
|
|||
|
|
if past_key_values:
|
|||
|
|
position_ids = position_ids[:, past_key_values[0][0].shape[-2]:]
|
|||
|
|
|
|||
|
|
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
|
|||
|
|
if inputs_embeds is not None and past_key_values is None:
|
|||
|
|
model_inputs = {"inputs_embeds": inputs_embeds}
|
|||
|
|
elif past_key_values is not None:
|
|||
|
|
model_inputs = {"input_ids": input_ids}
|
|||
|
|
else:
|
|||
|
|
model_inputs = {"input_ids": input_ids,
|
|||
|
|
"audios": kwargs.get("audios", None), "encoder_length": kwargs.get("encoder_length", None), "bridge_length": kwargs.get("bridge_length", None),
|
|||
|
|
"audios_tokens": kwargs.get("audios_tokens", None),
|
|||
|
|
"images": kwargs.get("images", None),
|
|||
|
|
"videos": kwargs.get("videos", None)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
model_inputs.update(
|
|||
|
|
{
|
|||
|
|
"position_ids": position_ids,
|
|||
|
|
"past_key_values": past_key_values,
|
|||
|
|
"use_cache": kwargs.get("use_cache"),
|
|||
|
|
"attention_mask": attention_mask,
|
|||
|
|
"images_grid": kwargs.get("images_grid"),
|
|||
|
|
"videos_grid": kwargs.get("videos_grid"),
|
|||
|
|
"patch_nums": kwargs.get("patch_nums"),
|
|||
|
|
"videos_patch_nums": kwargs.get("videos_patch_nums"),
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
return model_inputs
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _reorder_cache(past_key_values, beam_idx):
|
|||
|
|
reordered_past = ()
|
|||
|
|
for layer_past in past_key_values:
|
|||
|
|
reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
|
|||
|
|
return reordered_past
|
|||
|
|
|
|||
|
|
def chat(self, tokenizer, messages: List[dict], stream=False,
|
|||
|
|
generation_config: Optional[GenerationConfig]=None):
|
|||
|
|
generation_config = generation_config or self.generation_config
|
|||
|
|
input_ids = build_chat_input(self, tokenizer, messages, generation_config.max_new_tokens)
|
|||
|
|
if stream:
|
|||
|
|
streamer = TextIterStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
|||
|
|
Thread(target=self.generate, kwargs=dict(
|
|||
|
|
inputs=input_ids, streamer=streamer,
|
|||
|
|
generation_config=generation_config,
|
|||
|
|
)).start()
|
|||
|
|
return streamer
|
|||
|
|
else:
|
|||
|
|
outputs = self.generate(input_ids, generation_config=generation_config)
|
|||
|
|
response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)
|
|||
|
|
return response
|
|||
|
|
|
|||
|
|
|
|||
|
|
class OmniAudioTokenizer(OmniPreTrainedModel):
|
|||
|
|
"""
|
|||
|
|
Construct an audio tokenizer and decoder.
|
|||
|
|
"""
|
|||
|
|
def __init__(self, config: OmniConfig):
|
|||
|
|
super().__init__(config)
|
|||
|
|
self.padding_idx = None
|
|||
|
|
self.vocab_size = config.vocab_size
|
|||
|
|
self.training = False
|
|||
|
|
self.eval()
|
|||
|
|
self.audio_model = OmniAudioEncoder(config.audio_config)
|
|||
|
|
self.audio_bridge_model = OmniAudioVQBridgeTokenizer(config)
|
|||
|
|
if config.vocoder_config.enable:
|
|||
|
|
self.audio_decoder = OmniAudioDecoder(config)
|
|||
|
|
if config.flow_matching_config.enable:
|
|||
|
|
self.audio_flow_matching_decoder = OmniAudioFlowMatchingDecoder(config)
|
|||
|
|
|
|||
|
|
def encode(self, x, encoder_length: Optional[torch.Tensor] = None,
|
|||
|
|
bridge_length: Optional[torch.Tensor] = None):
|
|||
|
|
audio_emb = self.audio_model(x, encoder_length)
|
|||
|
|
audios_tokens = self.audio_bridge_model(audio_emb, bridge_length)
|
|||
|
|
return audios_tokens
|
|||
|
|
|
|||
|
|
def decode(self, audio_code_ids, bridge_length: Optional[torch.Tensor] = None):
|
|||
|
|
assert self.config.vocoder_config.enable, "Vocoder is not enabled in config."
|
|||
|
|
audio_emb = self.audio_bridge_model.decode(audio_code_ids)
|
|||
|
|
audio_dec = self.audio_decoder(
|
|||
|
|
audio_emb.to(next(self.audio_decoder.parameters())), bridge_length
|
|||
|
|
)
|
|||
|
|
if self.config.flow_matching_config.enable:
|
|||
|
|
if self.config.flow_matching_config.use_hidden_states_before_dconv2:
|
|||
|
|
hidden_states, hidden_states_length = (
|
|||
|
|
self.audio_flow_matching_decoder.unpack_hidden_states(
|
|||
|
|
audio_dec.hidden_states_before_dconv2,
|
|||
|
|
audio_dec.output_length_before_dconv2,
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
audio_flow_matching_decoder_ret = self.audio_flow_matching_decoder(
|
|||
|
|
hidden_states, hidden_states_length
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
audio_flow_matching_decoder_ret = self.audio_flow_matching_decoder(
|
|||
|
|
audio_dec.refined_mel, audio_dec.mel_length
|
|||
|
|
)
|
|||
|
|
return audio_flow_matching_decoder_ret
|
|||
|
|
else:
|
|||
|
|
return audio_dec
|
|||
|
|
|
|||
|
|
@torch.no_grad()
|
|||
|
|
def forward(self, audios, encoder_length: Optional[torch.Tensor] = None, bridge_length: Optional[torch.Tensor] = None):
|
|||
|
|
self.eval()
|
|||
|
|
audios_tokens = self.encode(audios, encoder_length, bridge_length)
|
|||
|
|
return audios_tokens
|