File size: 3,317 Bytes
de40e7d c85068f de40e7d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | #!/usr/bin/env python3
"""
AETHER-Micro Decoder Layer
Transformer Decoder Layer: Attention + MoE
"""
import torch
import torch.nn as nn
from .configuration_aether_micro import AETHERMicroConfig
from .normalization import AETHERMicroRMSNorm
from .attention import AETHERMicroAttention
from .moe import AETHERMicroMoE
from .latent_thought import AETHERMicroLatentThought
from .self_evaluation import AETHERMicroSelfEvalHead
class AETHERMicroDecoderLayer(nn.Module):
"""
Transformer Decoder Layer
Structure:
1. Input LayerNorm
2. Self-Attention (RoPE + GQA)
3. Post-Attention LayerNorm
4. MoE FFN (Heterogeneous experts)
5. Residual connections
"""
def __init__(self, config: AETHERMicroConfig):
super().__init__()
self.hidden_size = config.hidden_size
# Self-Attention
self.input_layernorm = AETHERMicroRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.self_attn = AETHERMicroAttention(config)
# Latent Thought Loop (Block 1)
if config.enable_latent_thought:
self.latent_thought = AETHERMicroLatentThought(config)
# Self Evaluation (Block 4)
if config.enable_self_eval:
self.self_evaluation = AETHERMicroSelfEvalHead(config)
# MoE FFN
self.post_attention_layernorm = AETHERMicroRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.mlp = AETHERMicroMoE(config)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor = None,
position_ids: torch.LongTensor = None,
disable_ltl: bool = False,
) -> torch.Tensor:
"""
Args:
hidden_states: (batch_size, seq_length, hidden_size)
attention_mask: (batch_size, 1, seq_length, seq_length)
position_ids: (batch_size, seq_length)
Returns:
hidden_states: (batch_size, seq_length, hidden_size)
"""
residual = hidden_states
# Self-Attention
hidden_states = self.input_layernorm(hidden_states)
hidden_states = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
)
hidden_states = residual + hidden_states
# Latent Thought Loop (Block 1) - Checkpoint-safe
if hasattr(self, 'latent_thought') and not disable_ltl:
result = self.latent_thought(hidden_states)
# Tuple unpacking (gradient checkpointing 호환)
if isinstance(result, tuple):
hidden_states = result[0]
# metrics는 체크포인팅 모드에서는 무시됨
else:
hidden_states = result
# MoE FFN
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
# Self Evaluation (Block 4)
# NOTE: self_evaluation returns (quality, overall) tuple for metrics
# We don't modify hidden_states - just compute quality scores
if hasattr(self, 'self_evaluation'):
_ = self.self_evaluation(hidden_states) # Compute but don't assign
return hidden_states
|