wehub-resource-sync
351 lines
16 KiB
Python
351 lines
16 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
import torch.nn as nn
|
|
from typing import Optional, Tuple
|
|
|
|
from .transformers import Attention, RMSNorm, Rotary, Embedding
|
|
from utils.custom_op import FakeLinear
|
|
from utils.spinner import spinner_run
|
|
from .torch_utils import onnx_export
|
|
from transformers.activations import ACT2FN
|
|
|
|
|
|
class DFlashAttention(torch.nn.Module):
|
|
"""DFlash non-causal attention: Q from noise, K/V from cat(context, noise)"""
|
|
def __init__(self, config, layer_idx):
|
|
super().__init__()
|
|
self.hidden_size = config.hidden_size
|
|
self.head_dim = config.head_dim
|
|
self.num_attention_heads = config.num_attention_heads
|
|
self.num_key_value_heads = config.num_key_value_heads
|
|
self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
|
|
self.scaling = self.head_dim ** -0.5
|
|
|
|
self.q_proj = nn.Linear(self.hidden_size, self.num_attention_heads * self.head_dim, bias=False)
|
|
self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
|
self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
|
self.o_proj = nn.Linear(self.num_attention_heads * self.head_dim, self.hidden_size, bias=False)
|
|
self.q_norm = RMSNorm(self.head_dim)
|
|
self.k_norm = RMSNorm(self.head_dim)
|
|
|
|
def forward(self, hidden_states, context_hidden, q_cos, q_sin, k_cos, k_sin, attention_mask):
|
|
"""
|
|
hidden_states: [1, block_size, hidden_size] (noise)
|
|
context_hidden: [1, context_len, hidden_size]
|
|
q_cos/q_sin: [1, 1, block_size, head_dim] - RoPE for Q
|
|
k_cos/k_sin: [1, 1, context_len + block_size, head_dim] - RoPE for K
|
|
attention_mask: [1, 1, block_size, context_len + block_size]
|
|
"""
|
|
bsz = 1
|
|
q_len = hidden_states.shape[1]
|
|
ctx_len = context_hidden.shape[1]
|
|
total_len = ctx_len + q_len
|
|
|
|
# Q from noise only
|
|
q = self.q_proj(hidden_states)
|
|
q = q.view(bsz, q_len, self.num_attention_heads, self.head_dim)
|
|
q = self.q_norm(q).transpose(1, 2) # [1, num_heads, q_len, head_dim]
|
|
|
|
# K/V from cat(context, noise)
|
|
kv_input = torch.cat([context_hidden, hidden_states], dim=1) # [1, total_len, hidden_size]
|
|
k = self.k_proj(kv_input)
|
|
v = self.v_proj(kv_input)
|
|
k = k.view(bsz, total_len, self.num_key_value_heads, self.head_dim)
|
|
k = self.k_norm(k).transpose(1, 2) # [1, num_kv_heads, total_len, head_dim]
|
|
v = v.view(bsz, total_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
|
|
|
|
# Apply RoPE (pre-computed, no dynamic slicing needed)
|
|
q = self._apply_rope(q, q_cos, q_sin)
|
|
k = self._apply_rope(k, k_cos, k_sin)
|
|
|
|
# GQA repeat
|
|
if self.num_key_value_groups > 1:
|
|
k = k.repeat_interleave(self.num_key_value_groups, dim=1)
|
|
v = v.repeat_interleave(self.num_key_value_groups, dim=1)
|
|
|
|
# Attention
|
|
attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scaling
|
|
attn_weights = attn_weights + attention_mask
|
|
attn_weights = torch.softmax(attn_weights, dim=-1)
|
|
attn_output = torch.matmul(attn_weights, v)
|
|
|
|
attn_output = attn_output.transpose(1, 2).reshape(bsz, q_len, -1)
|
|
return self.o_proj(attn_output)
|
|
|
|
@staticmethod
|
|
def _apply_rope(x, cos, sin):
|
|
"""Apply rotary position embedding."""
|
|
x1 = x[..., : x.shape[-1] // 2]
|
|
x2 = x[..., x.shape[-1] // 2 :]
|
|
rotated = torch.cat((-x2, x1), dim=-1)
|
|
return x * cos + rotated * sin
|
|
|
|
|
|
class DFlashDecoderLayer(torch.nn.Module):
|
|
def __init__(self, config, layer_idx):
|
|
super().__init__()
|
|
self.self_attn = DFlashAttention(config, layer_idx)
|
|
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)
|
|
self.mlp = nn.Module()
|
|
self.mlp.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
|
self.mlp.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
|
self.mlp.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
|
self.mlp.act_fn = ACT2FN[config.hidden_act]
|
|
|
|
def forward(self, hidden_states, context_hidden, q_cos, q_sin, k_cos, k_sin, attention_mask):
|
|
residual = hidden_states
|
|
hidden_states = self.input_layernorm(hidden_states)
|
|
hidden_states = self.self_attn(hidden_states, context_hidden, q_cos, q_sin, k_cos, k_sin, attention_mask)
|
|
hidden_states = residual + hidden_states
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.post_attention_layernorm(hidden_states)
|
|
hidden_states = self.mlp.down_proj(
|
|
self.mlp.act_fn(self.mlp.gate_proj(hidden_states)) * self.mlp.up_proj(hidden_states)
|
|
)
|
|
hidden_states = residual + hidden_states
|
|
return hidden_states
|
|
|
|
|
|
class DFlashFc(torch.nn.Module):
|
|
"""Feature compression module: fc + hidden_norm"""
|
|
def __init__(self, fc, hidden_norm):
|
|
super().__init__()
|
|
self.fc = fc
|
|
self.hidden_norm = hidden_norm
|
|
|
|
def forward(self, target_hidden):
|
|
return self.hidden_norm(self.fc(target_hidden))
|
|
|
|
|
|
class DFlash(torch.nn.Module):
|
|
"""DFlash Draft Model for export."""
|
|
def __init__(self, dflash_path, base):
|
|
super().__init__()
|
|
from transformers.configuration_utils import PretrainedConfig
|
|
|
|
# Load DFlash config
|
|
config_path = os.path.join(dflash_path, "config.json")
|
|
with open(config_path, 'r') as f:
|
|
config_dict = json.load(f)
|
|
|
|
self.dflash_config = config_dict
|
|
self.model_type = base.config.model_type
|
|
|
|
# Base model config
|
|
self.hidden_size = base.config.hidden_size
|
|
self.head_dim = base.config.head_dim
|
|
self.num_attention_heads = base.config.num_attention_heads
|
|
self.num_key_value_heads = base.config.num_key_value_heads
|
|
self.rms_norm_eps = getattr(base.config, 'rms_norm_eps', 1e-6)
|
|
|
|
# DFlash-specific config
|
|
dflash_cfg = config_dict.get('dflash_config', {})
|
|
self.block_size = config_dict.get('block_size', 16)
|
|
self.mask_token_id = dflash_cfg.get('mask_token_id', 0)
|
|
|
|
num_hidden_layers = config_dict.get('num_hidden_layers', 1)
|
|
num_target_layers = config_dict.get('num_target_layers', 3)
|
|
# Use origin_config (the original HF config) for attributes not in LlmConfig
|
|
origin_cfg = getattr(base.config, 'origin_config', base.config)
|
|
intermediate_size = config_dict.get('intermediate_size', getattr(origin_cfg, 'intermediate_size', 9728))
|
|
hidden_act = config_dict.get('hidden_act', 'silu')
|
|
|
|
# Build target layer ids
|
|
target_layer_ids = dflash_cfg.get('target_layer_ids', None)
|
|
if target_layer_ids is None:
|
|
# Use build_target_layer_ids logic
|
|
target_num_layers = getattr(base.config, 'num_hidden_layers', 32)
|
|
if num_hidden_layers == 1:
|
|
target_layer_ids = [target_num_layers // 2]
|
|
else:
|
|
start = 1
|
|
end = target_num_layers - 3
|
|
span = end - start
|
|
target_layer_ids = [
|
|
int(round(start + (i * span) / (num_target_layers - 1)))
|
|
for i in range(num_target_layers)
|
|
]
|
|
self.target_layer_ids = target_layer_ids
|
|
|
|
# Build a simple config namespace for sub-modules
|
|
class SimpleConfig:
|
|
pass
|
|
cfg = SimpleConfig()
|
|
cfg.hidden_size = self.hidden_size
|
|
cfg.head_dim = self.head_dim
|
|
cfg.num_attention_heads = self.num_attention_heads
|
|
cfg.num_key_value_heads = self.num_key_value_heads
|
|
cfg.intermediate_size = intermediate_size
|
|
cfg.hidden_act = hidden_act
|
|
cfg.rms_norm_eps = self.rms_norm_eps
|
|
|
|
# FC: Linear(num_target_layers * hidden_size, hidden_size)
|
|
self.fc = nn.Linear(len(self.target_layer_ids) * self.hidden_size, self.hidden_size, bias=False)
|
|
self.hidden_norm = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
|
|
|
|
# Decoder layers
|
|
self.layers = nn.ModuleList([
|
|
DFlashDecoderLayer(cfg, i) for i in range(num_hidden_layers)
|
|
])
|
|
|
|
# Final norm
|
|
self.norm = RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
|
|
|
|
# Shared lm_head from base model (for inclusion in dflash.onnx output)
|
|
self.lm_head = base.lm.lm
|
|
|
|
# Shared embed_tokens from base model (for embedding block tokens)
|
|
self.embed_tokens = base.embed.embed
|
|
|
|
# Rotary embedding
|
|
# Compatibility: transformers>=5.x moved rope_theta into rope_parameters dict
|
|
self.rope_theta = getattr(base.config, 'rope_theta', None)
|
|
if self.rope_theta is None or self.rope_theta == 10000.0:
|
|
origin_cfg = getattr(base.config, 'origin_config', base.config)
|
|
rp = getattr(origin_cfg, 'rope_parameters', None) or getattr(origin_cfg, 'rope_scaling', None)
|
|
if isinstance(rp, dict) and 'rope_theta' in rp:
|
|
self.rope_theta = rp['rope_theta']
|
|
if self.rope_theta is None:
|
|
self.rope_theta = 10000.0
|
|
self.max_position_embeddings = getattr(base.config, 'max_position_embeddings', 32768)
|
|
|
|
# Load weights
|
|
self._load_weights(dflash_path)
|
|
|
|
self.unloaded_ops = {}
|
|
|
|
def _load_weights(self, dflash_path):
|
|
"""Load DFlash model weights from safetensors or bin file."""
|
|
safetensors_path = os.path.join(dflash_path, "model.safetensors")
|
|
bin_path = os.path.join(dflash_path, "pytorch_model.bin")
|
|
state_dict = None
|
|
if os.path.exists(safetensors_path):
|
|
from safetensors.torch import load_file
|
|
state_dict = load_file(safetensors_path, device="cpu")
|
|
elif os.path.exists(bin_path):
|
|
state_dict = torch.load(bin_path, map_location="cpu")
|
|
else:
|
|
raise FileNotFoundError(
|
|
f"DFlash path '{dflash_path}' has no 'model.safetensors' or 'pytorch_model.bin'."
|
|
)
|
|
|
|
# Map weights to our structure
|
|
new_state_dict = {}
|
|
for key, value in state_dict.items():
|
|
new_key = key
|
|
new_state_dict[new_key] = value
|
|
|
|
# Filter to only load our parameters (exclude lm_head, embed_tokens, rotary)
|
|
own_keys = set(k for k, _ in self.named_parameters())
|
|
filtered = {}
|
|
for key, value in new_state_dict.items():
|
|
if key in own_keys:
|
|
filtered[key] = value
|
|
missing, unexpected = self.load_state_dict(filtered, strict=False)
|
|
# lm_head and embed_tokens are shared from base, so they'll be in missing - that's fine
|
|
|
|
def unload_param(self):
|
|
"""Replace linear layers with FakeLinear for memory-efficient export."""
|
|
def build_faker(real, name):
|
|
faker = FakeLinear(real.in_features, real.out_features, real.bias is not None, name)
|
|
self.unloaded_ops[name] = real
|
|
return faker
|
|
|
|
with torch.no_grad():
|
|
for i in range(len(self.layers)):
|
|
for name, child in self.layers[i].self_attn.named_children():
|
|
if isinstance(child, torch.nn.Linear):
|
|
setattr(self.layers[i].self_attn, name, build_faker(child, f'/dflash_layers.{i}/self_attn/{name}/Linear'))
|
|
for name, child in self.layers[i].mlp.named_children():
|
|
if isinstance(child, torch.nn.Linear):
|
|
setattr(self.layers[i].mlp, name, build_faker(child, f'/dflash_layers.{i}/mlp/{name}/Linear'))
|
|
self.fc = build_faker(self.fc, '/dflash/fc/Linear')
|
|
self.lm_head = build_faker(self.lm_head, '/lm/lm_head/Linear')
|
|
|
|
def forward(self, noise_embedding, context_hidden, attention_mask, q_position_ids, k_position_ids):
|
|
"""
|
|
DFlash main forward pass.
|
|
Args:
|
|
noise_embedding: [1, block_size, hidden_size] - embedded block tokens
|
|
context_hidden: [1, context_len, hidden_size] - output from fc module
|
|
attention_mask: [1, 1, block_size, context_len + block_size] - all zeros (non-causal)
|
|
q_position_ids: [1, block_size] - position ids for Q (block positions only)
|
|
k_position_ids: [1, context_len + block_size] - position ids for K/V (all positions)
|
|
Returns:
|
|
logits: [1, block_size, vocab_size]
|
|
"""
|
|
hidden_states = noise_embedding
|
|
|
|
# Compute rotary embeddings separately for Q and K
|
|
q_cos, q_sin = self._compute_rope(q_position_ids) # [1, 1, block_size, head_dim]
|
|
k_cos, k_sin = self._compute_rope(k_position_ids) # [1, 1, total_len, head_dim]
|
|
|
|
for layer in self.layers:
|
|
hidden_states = layer(hidden_states, context_hidden, q_cos, q_sin, k_cos, k_sin, attention_mask)
|
|
|
|
hidden_states = self.norm(hidden_states)
|
|
# Apply lm_head to get logits
|
|
logits = self.lm_head(hidden_states)
|
|
return logits
|
|
|
|
def _compute_rope(self, position_ids):
|
|
"""Compute rotary position embeddings (cos, sin) for given positions."""
|
|
# position_ids: [1, seq_len]
|
|
inv_freq = 1.0 / (self.rope_theta ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=position_ids.device) / self.head_dim))
|
|
# [seq_len] x [head_dim/2] -> [seq_len, head_dim/2]
|
|
freqs = position_ids.float().squeeze(0).unsqueeze(-1) * inv_freq.unsqueeze(0)
|
|
# [seq_len, head_dim]
|
|
emb = torch.cat([freqs, freqs], dim=-1)
|
|
cos = emb.cos().unsqueeze(0).unsqueeze(1) # [1, 1, seq_len, head_dim]
|
|
sin = emb.sin().unsqueeze(0).unsqueeze(1)
|
|
return cos, sin
|
|
|
|
@spinner_run(f'export onnx model to ')
|
|
def export(self, onnx_path):
|
|
dflash_model = f'{onnx_path}/dflash.onnx'
|
|
dflash_fc_model = f'{onnx_path}/dflash_fc.onnx'
|
|
|
|
block_size = self.block_size
|
|
context_len = 3 # dummy context length for export
|
|
|
|
# Export dflash_fc.onnx
|
|
fc_module = DFlashFc(self.fc, self.hidden_norm)
|
|
fc_hidden = torch.ones([1, context_len, len(self.target_layer_ids) * self.hidden_size], dtype=torch.float)
|
|
with torch.no_grad():
|
|
onnx_export(
|
|
fc_module, (fc_hidden,),
|
|
dflash_fc_model,
|
|
input_names=['target_hidden'],
|
|
output_names=['context_hidden'],
|
|
dynamic_axes={"target_hidden": {1: "seq_len"}}
|
|
)
|
|
|
|
# Unload params for main model export
|
|
self.unload_param()
|
|
|
|
# Export dflash.onnx (main model)
|
|
noise_embedding = torch.ones([1, block_size, self.hidden_size], dtype=torch.float)
|
|
context_hidden = torch.ones([1, context_len, self.hidden_size], dtype=torch.float)
|
|
attention_mask = torch.zeros([1, 1, block_size, context_len + block_size], dtype=torch.float)
|
|
q_position_ids = torch.arange(context_len, context_len + block_size, dtype=torch.int).unsqueeze(0)
|
|
k_position_ids = torch.arange(context_len + block_size, dtype=torch.int).unsqueeze(0)
|
|
|
|
with torch.no_grad():
|
|
onnx_export(
|
|
self, (noise_embedding, context_hidden, attention_mask, q_position_ids, k_position_ids),
|
|
dflash_model,
|
|
input_names=['noise_embedding', 'context_hidden', 'attention_mask', 'q_position_ids', 'k_position_ids'],
|
|
output_names=['logits'],
|
|
dynamic_axes={
|
|
"noise_embedding": {1: "block_size"},
|
|
"context_hidden": {1: "context_len"},
|
|
"attention_mask": {2: "block_size", 3: "total_len"},
|
|
"q_position_ids": {1: "block_size"},
|
|
"k_position_ids": {1: "total_len"},
|
|
}
|
|
)
|
|
|
|
return dflash_model, dflash_fc_model |