my_decoder_only_model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
import torchviz
import yaml
import argparse
from transformers import PreTrainedModel, PretrainedConfig, AutoConfig, AutoModel
import logging
import hydra
import torch.utils.checkpoint as checkpoint
from typing import Optional, Dict, Any
from rotary_embedding_torch import RotaryEmbedding
from rich import traceback
from attention import ComplexMultiHeadAttentionV2
from accelerate import Accelerator, DeepSpeedPlugin
from accelerate.utils import DistributedType
import os
#from flash_attn import flash_attn_qkvpacked_func, flash_attn_func

#traceback.install(show_locals=True)

logger: logging.Logger = logging.getLogger(__name__)
logger.propagate = False
logger.addHandler(logging.FileHandler("logs/model.log"))


class CustomConfig(PretrainedConfig):
    model_type: str = "ComplexFormer"

    def __init__(self,
                 vocab_size: int = 30522,
                 hidden_dim: int = 512,
                 intermediate_size: int = 1024,
                 max_seq_len: int = 512,
                 n_layers: int = 8,
                 num_attention_heads: int = 8,
                 dropout: float = 0.0,
                 **kwargs: Any):
        super().__init__(**kwargs)
        self.vocab_size: int = vocab_size
        self.hidden_dim: int = hidden_dim
        self.intermediate_size: int = intermediate_size
        self.max_seq_len: int = max_seq_len
        self.n_layers: int = n_layers
        self.num_attention_heads: int = num_attention_heads
        self.dropout: float = dropout
        self.debug : bool = kwargs.get("debug", False)


class ComplexFormerModel(PreTrainedModel):
    config_class = CustomConfig
    #

    def __init__(self, config: CustomConfig):
        super().__init__(config)
        self.config: CustomConfig = config
        self.embedding: nn.Embedding = nn.Embedding(config.vocab_size, config.hidden_dim)
        self.pos_embedding = PositionalEncoding(config.hidden_dim, config.max_seq_len, device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
        self.head_dim = config.hidden_dim // config.num_attention_heads
        self.transformer_blocks: nn.ModuleList = nn.ModuleList(
            [TransformerBlock(config) for _ in range(config.n_layers)]  # type: ignore
        )
        self.linear: nn.Linear = nn.Linear(config.hidden_dim, config.vocab_size)
        self.softmax: nn.Softmax = nn.Softmax(dim=-1)

        #
        # 添加旋转位置编码
        self.rope = RotaryEmbedding(dim=self.head_dim)
        self.gradient_checkpointing = False
        self.debug = config.debug
        self._togger_loger()
        self.apply(self._init_weights)

    def can_generate(self)-> bool:
        return True



    def forward(self, input_ids: Tensor, attention_mask: Tensor, labels: Optional[Tensor] = None, use_checkpoint: bool = False,token_type_ids = None) -> Tensor:
        #
        if torch.isnan(input_ids).any() or torch.isinf(input_ids).any():
            raise ValueError("Tensor contains NaN or Inf values.")

        if (input_ids >= self.config.vocab_size).any() or (input_ids < 0).any():
            raise ValueError("input_ids contain values outside the valid range.")
        seq_len: int = input_ids.size(1)# 
        logger.info(f"Input IDs shape: {input_ids.shape}")#[batch_size, seq_len]
        device = input_ids.device
        #logger.info(f"Input IDs shape: {input_ids.device}")

        token_embeddings = self.embedding(input_ids).to(device)
        x: Tensor = token_embeddings 
        
        if self.config.complex_attention is False:
            position_embeddings = self.pos_embedding(input_ids).to(device)
            #logger.info(f"Position Embeddings shape: {position_embeddings.device}")

            # 将位置编码添加到输入嵌入
            ## 确保设备一致性
            position_embeddings = position_embeddings.to(device)
            #
            assert token_embeddings.device == position_embeddings.device, "input_ids and position_embedding must be on the same device."
            logger.info(f"Token Embeddings shape: {token_embeddings.device}")
            #logger.info(f"Token Embeddings shape: {token_embeddings.shape}")
            logger.info(f"Position Embeddings shape: {position_embeddings.device}")
            x: Tensor = token_embeddings + position_embeddings # [batch_size, seq_len, hidden_dim]
            logger.info(f"Input embeddings shape: {x.shape}")      
        #logger.info(f"Input shape: {x.shape}")
        #
        # 应用旋转位置编码
        # x = x.view(batch_size, seq_len, self.config.num_attention_heads, self.head_dim).transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        # x = self.rope(x, seq_len=seq_len)  # 应用 RoPE
        # x = x.transpose(1, 2).contiguous().view(batch_size, seq_len, hidden_dim)  # 恢复形状

        for block in self.transformer_blocks:
            if self.gradient_checkpointing:
                def create_custom_forward(module):
                    # 这个函数用来捕获外部变量
                    def custom_forward(*inputs): # inputs[0] 会是 x
                        # captured_key_padding_mask 是从外面捕获的 mask
                        return module(inputs[0], padding_mask=captured_key_padding_mask)
                    return custom_forward

                captured_key_padding_mask = attention_mask # 捕获当前的 mask
                layer_outputs = checkpoint.checkpoint(
                    create_custom_forward(block), # 包装后的函数
                    x,                          # 只有 x (需要梯度) 直接传入
                    use_reentrant=False,
                    preserve_rng_state=True # 保证 dropout 等随机性一致
                )
                x = layer_outputs
                logger.info(f"Transformer block output shape with ckpt: {x.shape}")
            else:
                x = block(x, padding_mask=attention_mask)
                logger.info(f"Transformer block output shape no ckpt: {x.shape}")
        logger.info(f"Transformer block output shape: {x.shape}")
        x = self.linear(x)
        logger.info(f"Linear layer output shape: {x.shape}")
        # x = self.softmax(x)
        # logger.info(f"Softmax output shape: {x.shape}") #gradient loss
       # assert x.shape == (self.config.batch_size, self.config.max_seq_len, self.config.vocab_size), f"Output shape mismatch{self.config.batch_size},{self.config.max_seq_len}, {self.config.vocab_size} is  {x.shape}"
        return x

    def gradient_checkpointing_enable(self,**kwargs):
        self.gradient_checkpointing = True
    
    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, nn.Linear):
            # Kaiming normal for layers with SiLU/ReLU
            nn.init.kaiming_normal_(module.weight, nonlinearity='leaky_relu') # SiLU 类似 leaky_relu
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.kaiming_normal_(module.weight, nonlinearity='leaky_relu') # 或者用 xavier_uniform_
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        # 可以为其他类型的层添加特定的初始化，例如 LayerNorm/RMSNorm
        elif isinstance(module, (RMSNorm, nn.LayerNorm)):
            if hasattr(module, 'weight') and module.weight is not None:
                nn.init.ones_(module.weight)
            if hasattr(module, 'bias') and module.bias is not None:
                nn.init.zeros_(module.bias)
    def _togger_loger(self):
        if self.debug == False:
            logger.setLevel(logging.WARNING)
            
    @torch.no_grad()
    def generate(self,
        input_ids: Tensor, # (batch_size, prompt_seq_len)
        attention_mask: Optional[Tensor] = None, # (batch_size, prompt_seq_len)
        max_length: int = 100,
        temperature: float = 1.0,
        top_k: int = 0, # 0 means no top-k filtering
        top_p: float = 1.0, # 1.0 means no top-p filtering
        repetition_penalty: float = 1.0,
        eos_token_id: Optional[int] = None,
        pad_token_id: Optional[int] = None,
        use_cache: bool = True,
        **kwargs):
        
        self.eval() # Set model to evaluation mode
        if pad_token_id is None and self.config.pad_token_id is not None:
            pad_token_id = self.config.pad_token_id
        if eos_token_id is None and self.config.eos_token_id is not None:
            eos_token_id = self.config.eos_token_id
        
        batch_size, cur_len = input_ids.shape
        
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids)

        # Keep track of whether we are using EOS
        model_kwargs = {"attention_mask": attention_mask, **kwargs}
        
        # `generated_ids` will store the complete sequence (prompt + generated)
        # Initialize with input_ids, ensuring it's on the correct device and not part of graph
        generated_ids = input_ids.clone().detach()
        
        # `past_key_values` for KV caching
        past_key_values = None

        for _ in range(max_length - cur_len): # Generate up to max_length - prompt_length new tokens
            # Prepare model inputs
            model_inputs = {"input_ids":generated_ids,"attention_mask":attention_mask}
            
            # Forward pass to get logits
            outputs = self(
                **model_inputs,
                # use_cache=use_cache # Already passed via model_inputs
            )
            logits = outputs
            next_token_logits = logits[:, -1, :] # Get logits for the last token: (batch_size, vocab_size)

            # Apply repetition penalty
            if repetition_penalty != 1.0:
                for i in range(batch_size):
                    for token_id in set(generated_ids[i].tolist()):
                        next_token_logits[i, token_id] /= repetition_penalty
            
            # Apply temperature
            if temperature != 1.0:
                next_token_logits = next_token_logits / temperature

            # Apply top-k filtering
            if top_k > 0:
                top_k_values, top_k_indices = torch.topk(next_token_logits, top_k, dim=-1)
                # Create a mask for elements not in top-k
                top_k_mask = torch.ones_like(next_token_logits, dtype=torch.bool).scatter_(-1, top_k_indices, False)
                next_token_logits.masked_fill_(top_k_mask, float('-inf'))

            # Apply top-p (nucleus) filtering
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True, dim=-1)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                
                # Remove tokens with cumulative probability above the threshold
                sorted_indices_to_remove = cumulative_probs > top_p
                # Shift the mask: keep the first token above threshold
                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                sorted_indices_to_remove[..., 0] = 0 # Never remove the most probable token
                
                # Scatter back to original order
                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                next_token_logits.masked_fill_(indices_to_remove, float('-inf'))

            # Sample next token
            probs = F.softmax(next_token_logits, dim=-1)
            next_token_id = torch.multinomial(probs, num_samples=1) # (batch_size, 1)

            # Append the new token
            generated_ids = torch.cat([generated_ids, next_token_id], dim=-1)
            
            # Update attention_mask if it's not None
            if attention_mask is not None:
                attention_mask = torch.cat(
                    [attention_mask, attention_mask.new_ones((batch_size, 1))], dim=-1
                )
                model_kwargs["attention_mask"] = attention_mask


            # # Update past_key_values
            # if use_cache:
            #     past_key_values = outputs.past_key_values
                
            # Check for EOS token
            if eos_token_id is not None and (next_token_id == eos_token_id).all():
                logger.info("EOS token generated for all sequences in batch.")
                break
            
            # Check if current length exceeds max_length (should be handled by loop, but good for sanity)
            if generated_ids.shape[1] >= max_length:
                break
                
        return generated_ids

        
class FFN(nn.Module):
    def __init__(self, config: CustomConfig):
        super().__init__()
        self.linear1: nn.Linear = nn.Linear(config.hidden_dim, config.intermediate_size)
        self.linear3: nn.Linear = nn.Linear(config.intermediate_size, config.hidden_dim)
        self.linear2: nn.Linear = nn.Linear(config.hidden_dim, config.intermediate_size)
        self.activation: nn.Module = nn.SiLU()  # or any other activation function

    def forward(self, x: Tensor) -> Tensor:
        gate = self.linear1(x)  # 输出通常是 float32 (如果 x 是 float32)
        up = self.linear2(x)    # 输出通常是 float32 (如果 x 是 float32)

        # 应用激活函数
        activated_up = self.activation(up) # 输出通常是 float32

        # 在乘法之前，可以将激活转换为 bfloat16
        # 注意：如果 gate 也是 float32，torch 会自动进行类型提升，结果可能是 float32
        # 为了确保乘法在 bfloat16 下进行（如果这是目标），两者都应转换
        # 或者至少乘法的一个操作数是 bfloat16，另一个会被提升或保持
        
        # 明确的类型转换 (推荐)
        if gate.dtype == torch.float32: # 仅当输入是 float32 时转换，避免不必要的转换
            gate_bf16 = gate.to(torch.bfloat16)
            activated_up_bf16 = activated_up.to(torch.bfloat16)
            intermediate = gate_bf16 * activated_up_bf16
            x = self.linear3(intermediate.to(torch.float32))
        else: # 如果输入已经是 bfloat16，则直接计算
            intermediate = gate * activated_up
            x = self.linear3(intermediate)
        
        return x

    



class PositionalEncoding(nn.Module):
    """
    compute sinusoid encoding.
    """

    def __init__(self, d_model, max_len, device):
        """
        constructor of sinusoid encoding class

        :param d_model: dimension of model
        :param max_len: max sequence length
        :param device: hardware device setting
        """
        super(PositionalEncoding, self).__init__()

        # same size with input matrix (for adding with input matrix)
        self.encoding = torch.zeros(max_len, d_model, device=device)
        self.encoding.requires_grad = False  # we don't need to compute gradient

        pos = torch.arange(0, max_len, device=device)
        pos = pos.float().unsqueeze(dim=1)
        # 1D => 2D unsqueeze to represent word's position

        _2i = torch.arange(0, d_model, step=2, device=device).float()
        # 'i' means index of d_model (e.g. embedding size = 50, 'i' = [0,50])
        # "step=2" means 'i' multiplied with two (same with 2 * i)

        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))
        # compute positional encoding to consider positional information of words

    def forward(self, x):
        # self.encoding
        # [max_len = 512, d_model = 512]

        batch_size, seq_len = x.size()
        # [batch_size = 128, seq_len = 30]

        return self.encoding[:seq_len, :]

class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int):
        super().__init__()
        self.weight: nn.Parameter = nn.Parameter(torch.ones(hidden_size))
        self.eps: float = 1e-5

    def forward(self, x: Tensor) -> Tensor:
        return x * (torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True)) * self.weight + self.eps)


class TransformerBlock(nn.Module):
    def __init__(self, config: CustomConfig):
        super().__init__()
        self.attention = ComplexMultiHeadAttentionV2(config.hidden_dim, config.num_attention_heads)
        self.ffn: FFN = FFN(config)
        self.rmsnorm: RMSNorm = RMSNorm(config.hidden_dim)
        self.config: CustomConfig = config

    def forward(self, x: Tensor, padding_mask: Tensor) -> Tensor:
        residual: Tensor = x
        x = self.rmsnorm(x)
        logger.info(f"TransformerBlock Input shape: {x.shape}")
        seq_len: int = x.shape[1]
        causal_mask: Tensor = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool().to(x.device)
        
        
        if self.config.complex_attention:
            logger.info('using complex attention')
            x = self.attention(x, x, x, mask=padding_mask)
        else:
            #TODO
            def atten_forward(x: Tensor) -> Tensor:
                return self.attention(x, x, x, attn_mask=causal_mask)[0]
            logger.info(f'before forward{x.dtype}')
            x = atten_forward(x.bfloat16())
            logger.info(f'after forward{x.dtype}')
        x = self.rmsnorm(x + residual)
        logger.info(f"TransformerBlock after rmsnorm shape: {x.shape}")

        if torch.isnan(x).any():
            logger.info(f"TransformerBlock after rmsnorm has nan")
            raise ValueError("TransformerBlock after rmsnorm has nan")
        residual = x
        x = self.ffn(x)
        return x






def load_config(config_path: str) -> CustomConfig:
    with open(config_path, 'r') as f:
        config_dict: Dict[str, Any] = yaml.safe_load(f)
    return CustomConfig(**config_dict)


@hydra.main(config_path='.', config_name="config.yaml")
def main(config: Dict[str, Any]) -> None:
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--config', type=str, default='./pretrain/config.yaml', help='Path to the config file')
    args = argparser.parse_args()

    AutoConfig.register(config.architecture, CustomConfig)
    AutoModel.register(CustomConfig, ComplexFormerModel)

    device: torch.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    #deepspeed config

    model: ComplexFormerModel = ComplexFormerModel(config=load_config(args.config)).to(device)
    

    #
    test_model(model,config)

    model.gradient_checkpointing_enable()
    model.save_pretrained(config.model.save_dir)
    logger.info(f"Model saved to {config.model.save_dir}")

     # --- Calculate Parameters ---
    num_params = sum(p.numel() for p in model.state_dict().values() if torch.is_tensor(p))
    print(f"Trainable model parameters: {num_params/1e9:,}B")
    print_model_parameters_llm(model)


    # 2. 打印模型参数的函数
def print_model_parameters_llm(model_to_inspect):
    if not isinstance(model_to_inspect, (AutoModel, torch.nn.Module)): # 检查是否是预期类型
        print("传入的不是一个有效的PyTorch模型。")
        return

    print(f"\n模型架构 (部分，如果太大):")
    # 对于非常大的LLM，直接打印模型对象可能会输出非常多内容
    # print(model_to_inspect)
    # 可以只打印模型类型
    print(type(model_to_inspect))


    print("\n参数详情:")
    print("---------------------------------------------------------------------------------------------------------------")
    # 调整列宽以适应更长的参数名
    print(f"{'Parameter Name':<70} | {'Shape':<25} | {'Numel':<12} | {'Requires Grad':<15} | {'Dtype':<10}")
    print("---------------------------------------------------------------------------------------------------------------")

    total_params = 0
    trainable_params = 0

    for name, param in model_to_inspect.named_parameters():
        numel = param.numel()
        total_params += numel
        if param.requires_grad:
            trainable_params += numel
        
        # 参数名可能非常长，如果需要可以截断显示
        display_name = name
        # if len(name) > 68:
        # display_name = name[:65] + "..."

        print(f"{display_name:<70} | {str(param.shape):<25} | {numel:<12,} | {str(param.requires_grad):<15} | {str(param.dtype).replace('torch.', ''):<10}")

    print("---------------------------------------------------------------------------------------------------------------")
    print(f"总参数量 (Total parameters): {total_params:,}")
    print(f"可训练参数量 (Trainable parameters): {trainable_params:,}")
    if total_params != trainable_params:
        print(f"不可训练/冻结参数量 (Non-trainable parameters): {total_params - trainable_params:,}")
    print("---------------------------------------------------------------------------------------------------------------")


@torch.no_grad()
def test_model(model: ComplexFormerModel,config) -> None:
    device: torch.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logger.info(f"Testing model...{device}")

    input_ids: Tensor = torch.randint(0, config.vocab_size, (config.batch_size, config.max_seq_len)).to(device)
    attention_mask: Tensor = torch.ones((config.batch_size, config.max_seq_len )).bool().to(device) #[seqlen,hidden_dim]
    labels: Tensor = torch.randint(0, config.vocab_size, (config.batch_size, config.max_seq_len)).to(device)
    output: Tensor = model(input_ids, attention_mask=attention_mask, labels=labels, use_checkpoint=True).to(device)
    logger.info(f"Model output shape: {output.shape}")
    assert output.shape == (config.batch_size, config.max_seq_len, config.vocab_size), "Output shape mismatch"

    
    # 使用损失函数
    criterion = nn.CrossEntropyLoss()
    loss = criterion(output.view(-1, config.vocab_size), labels.view(-1))  # 确保形状匹配
    logger.info(f"Loss: {loss.dim()}")

if __name__ == '__main__':
    main()

AutoConfig.register("ComplexFormer", CustomConfig)
AutoModel.register(CustomConfig, ComplexFormerModel)