import torch
import torch.nn as nn
from typing import Optional
class HybridAttention(nn.Module):
"""Hybrid attention with 3:1 ratio of local to global attention."""
def __init__(
self,
hidden_dim: int,
num_heads: int,
window_size: int = 4096,
local_to_global_ratio: float = 3.0
):
super().__init__()
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.window_size = window_size
self.ratio = local_to_global_ratio
self.q_proj = nn.Linear(hidden_dim, hidden_dim)
self.k_proj = nn.Linear(hidden_dim, hidden_dim)
self.v_proj = nn.Linear(hidden_dim, hidden_dim)
self.local_heads = int(num_heads * (local_to_global_ratio / (1 + local_to_global_ratio)))
self.global_heads = num_heads - self.local_heads
self.out_proj = nn.Linear(hidden_dim, hidden_dim)
def local_attention(self, q, k, v, window_size: int):
"""
Apply sliding window attention to Q, K, V.
Args:
q, k, v: (batch, seq_len, hidden_dim)
window_size: Size of local window
Returns:
attended: (batch, seq_len, hidden_dim)
"""
batch_size, seq_len, _ = q.shape
q = q.view(batch_size, seq_len, self.local_heads, -1).transpose(1, 2)
k = k.view(batch_size, seq_len, self.local_heads, -1).transpose(1, 2)
v = v.view(batch_size, seq_len, self.local_heads, -1).transpose(1, 2)
scores = torch.matmul(q, k.transpose(-2, -1)) / (q.shape[-1] ** 0.5)
mask = torch.triu(
torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool),
diagonal=window_size + 1
)
mask |= torch.triu(
torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool),
diagonal=1
) & ~torch.triu(
torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool),
diagonal=window_size + 1 - window_size
)
scores = scores.masked_fill(mask.unsqueeze(0).unsqueeze(0), -1e9)
attn_weights = torch.softmax(scores, dim=-1)
attended = torch.matmul(attn_weights, v)
attended = attended.transpose(1, 2).contiguous()
attended = attended.view(batch_size, seq_len, -1)
return attended
def global_attention(self, q, k, v):
"""
Apply standard (full sequence) attention to remaining heads.
Args:
q, k, v: (batch, seq_len, hidden_dim)
Returns:
attended: (batch, seq_len, hidden_dim)
"""
batch_size, seq_len, _ = q.shape
q = q.view(batch_size, seq_len, self.global_heads, -1).transpose(1, 2)
k = k.view(batch_size, seq_len, self.global_heads, -1).transpose(1, 2)
v = v.view(batch_size, seq_len, self.global_heads, -1).transpose(1, 2)
scores = torch.matmul(q, k.transpose(-2, -1)) / (q.shape[-1] ** 0.5)
attn_weights = torch.softmax(scores, dim=-1)
attended = torch.matmul(attn_weights, v)
attended = attended.transpose(1, 2).contiguous()
attended = attended.view(batch_size, seq_len, -1)
return attended
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
"""
Apply hybrid attention combining local and global.
Args:
hidden_states: (batch, seq_len, hidden_dim)
Returns:
output: (batch, seq_len, hidden_dim)
"""
q = self.q_proj(hidden_states)
k = self.k_proj(hidden_states)
v = self.v_proj(hidden_states)
local_out = self.local_attention(q, k, v, self.window_size)
global_out = self.global_attention(q, k, v)
output = torch.cat([local_out, global_out], dim=-1)
output = self.out_proj(output)
return output
class QKReorderLN(nn.Module):
"""Custom QK-Reorder-LN normalization for improved attention stability."""
def __init__(self, hidden_dim: int, eps: float = 1e-6):
super().__init__()
self.hidden_dim = hidden_dim
self.eps = eps
self.weight = nn.Parameter(torch.ones(hidden_dim))
self.bias = nn.Parameter(torch.zeros(hidden_dim))
def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple:
"""
Reorder Q and K before LayerNorm for stability.
Args:
q: Query tensor
k: Key tensor
Returns:
Normalized (q, k)
"""
batch_size, seq_len, hidden = q.shape
head_dim = hidden // 8
q = q.view(batch_size, seq_len, 8, -1)
k = k.view(batch_size, seq_len, 8, -1)
q_mean = q.mean(dim=-1, keepdim=True)
q_std = q.std(dim=-1, keepdim=True) + self.eps
q_norm = (q - q_mean) / q_std
k_mean = k.mean(dim=-1, keepdim=True)
k_std = k.std(dim=-1, keepdim=True) + self.eps
k_norm = (k - k_mean) / k_std
q_norm = q_norm.view(batch_size, seq_len, hidden)
k_norm = k_norm.view(batch_size, seq_len, hidden)
q_out = q_norm * self.weight + self.bias
k_out = k_norm * self.weight + self.bias
return q_out, k_out
class ContextExtensionTrainer:
"""Progressive context length training with validation."""
def __init__(self, model, tokenizer, initial_context: int = 4096):
self.model = model
self.tokenizer = tokenizer
self.initial_context = initial_context
self.stages = [
{'context': 4096, 'name': 'base'},
{'context': 32768, 'name': 'intermediate'},
{'context': 131072, 'name': 'extended'}
]
def needle_in_haystack_test(
self,
needle: str,
haystack_length: int,
needle_position_ratio: float = 0.5
) -> bool:
"""
Test if model can retrieve needle from long haystack.
Validates context extension at each stage.
Args:
needle: Text to hide in context
haystack_length: Length of context to fill
needle_position_ratio: Where to place needle (0.5 = middle)
Returns:
True if model retrieves needle correctly
"""
filler = "The color of the sky is blue. " * (haystack_length // 30)
needle_pos = int(len(filler) * needle_position_ratio)
haystack = filler[:needle_pos] + needle + filler[needle_pos:]
haystack = haystack[:haystack_length]
prompt = f"{haystack}\n\nQuestion: Find the needle in this text and repeat it exactly.\n\nAnswer:"
input_ids = self.tokenizer.encode(prompt, return_tensors='pt')
if len(input_ids[0]) > haystack_length:
input_ids = input_ids[:, :haystack_length + 50]
with torch.no_grad():
output_ids = self.model.generate(
input_ids,
max_new_tokens=len(self.tokenizer.encode(needle)) + 10,
temperature=0.0
)
output = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
return needle in output
def train_stage(
self,
stage: dict,
training_data,
validation_data,
num_epochs: int = 3
) -> dict:
"""
Train model at specific context length.
Args:
stage: {'context': length, 'name': 'stage_name'}
training_data: Dataset with examples
validation_data: Validation dataset
num_epochs: Training epochs for this stage
Returns:
Stage metrics
"""
context_len = stage['context']
self.model.max_position_embeddings = context_len
from transformers import TextIterableDataset
print(f"\nTraining stage: {stage['name']} (context={context_len})")
for epoch in range(num_epochs):
total_loss = 0.0
for batch in training_data:
input_ids = batch['input_ids'][:, :context_len]
outputs = self.model(input_ids, labels=input_ids)
loss = outputs.loss
loss.backward()
self.model.optimizer.step()
self.model.optimizer.zero_grad()
total_loss += loss.item()
avg_loss = total_loss / len(training_data)
print(f" Epoch {epoch+1}: loss={avg_loss:.4f}")
accuracy = 0
for needle_pos in [0.25, 0.5, 0.75]:
if self.needle_in_haystack_test(
"NEEDLE",
int(context_len * 0.8),
needle_pos
):
accuracy += 1
needle_accuracy = accuracy / 3
print(f" Needle-In-Haystack accuracy: {needle_accuracy:.1%}")
return {
'stage': stage['name'],
'context': context_len,
'final_loss': avg_loss,
'needle_accuracy': needle_accuracy
}
def run_full_extension(self, training_data, validation_data):
"""Run all context extension stages."""
results = []
for stage in self.stages:
stage_results = self.train_stage(stage, training_data, validation_data)
results.append(stage_results)
return results
class AGAPO:
"""AGAPO: Improved GRPO with asymmetric sampling and group advantage."""
def __init__(self, model, learning_rate: float = 1e-6, beta_kl: float = 0.01):
self.model = model
self.reference_model = model
self.optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
self.beta_kl = beta_kl
def compute_group_advantage(
self,
rewards: torch.Tensor,
group_size: int = 8
) -> torch.Tensor:
"""
Compute group advantage: individual score minus group mean.
(Instead of clip-based GRPO)
Args:
rewards: (batch_size,) tensor of rewards
group_size: Size of comparison group
Returns:
advantages: (batch_size,) advantage per sample
"""
advantages = []
for i in range(0, len(rewards), group_size):
group_rewards = rewards[i:i+group_size]
group_mean = group_rewards.mean()
group_advantages = group_rewards - group_mean
advantages.extend(group_advantages.tolist())
return torch.tensor(advantages)
def asymmetric_sampling(
self,
model_outputs: list,
rewards: torch.Tensor,
correct_ratio: float = 0.5,
num_samples: int = 32
) -> tuple:
"""
Sample correct and incorrect responses asymmetrically.
Use more incorrect samples for contrastive learning.
Args:
model_outputs: All generated responses
rewards: Reward for each response
correct_ratio: How many correct samples to keep (e.g., 0.3)
num_samples: Total samples to select
Returns:
(correct_samples, incorrect_samples)
"""
correct_indices = torch.where(rewards > 0.5)[0]
incorrect_indices = torch.where(rewards <= 0.5)[0]
num_correct = max(1, int(num_samples * correct_ratio))
num_incorrect = num_samples - num_correct
correct_sample_idx = torch.randperm(len(correct_indices))[:num_correct]
incorrect_sample_idx = torch.randperm(len(incorrect_indices))[:num_incorrect]
correct_samples = [model_outputs[i] for i in correct_indices[correct_sample_idx]]
incorrect_samples = [model_outputs[i] for i in incorrect_indices[incorrect_sample_idx]]
return correct_samples, incorrect_samples
def training_step(
self,
prompts: list,
rewards: torch.Tensor,
model_outputs: list
) -> dict:
"""
Single AGAPO training step.
Args:
prompts: Input prompts
rewards: Scalar rewards per output
model_outputs: Generated responses
Returns:
Loss components
"""
self.model.train()
advantages = self.compute_group_advantage(rewards)
correct, incorrect = self.asymmetric_sampling(model_outputs, rewards)
total_loss = 0.0
for prompt, correct_response, advantage in zip(
prompts, correct, advantages
):
input_ids = self.model.tokenizer.encode(prompt + correct_response, return_tensors='pt')
outputs = self.model(input_ids)
logits = outputs.logits
log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
policy_loss = -advantage * log_probs.mean()
with torch.no_grad():
ref_logits = self.reference_model(input_ids).logits
ref_log_probs = torch.nn.functional.log_softmax(ref_logits, dim=-1)
kl_loss = torch.nn.functional.kl_div(log_probs, ref_log_probs.exp(), reduction='mean')
loss = policy_loss + self.beta_kl * kl_loss
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
total_loss += loss.item()
return {
'policy_loss': policy_loss.item(),
'kl_loss': kl_loss.item(),
'total_loss': total_loss / len(prompts)
}
class ModeSwitchingDecoder:
"""Generate in standard or reasoning mode."""
def __init__(self, model, tokenizer):
self.model = model
self.tokenizer = tokenizer
def generate(
self,
prompt: str,
mode: str = 'standard',
max_new_tokens: int = 100
) -> str:
"""
Generate with mode-specific hyperparameters.
Args:
prompt: Input text
mode: 'standard' (fast) or 'reasoning' (deep)
max_new_tokens: Generation length
Returns:
Generated text
"""
input_ids = self.tokenizer.encode(prompt, return_tensors='pt')
if mode == 'standard':
config = {
'max_new_tokens': min(max_new_tokens, 256),
'temperature': 0.0,
'top_p': 1.0,
'do_sample': False
}
elif mode == 'reasoning':
config = {
'max_new_tokens': min(max_new_tokens, 64000),
'temperature': 0.6,
'top_p': 0.95,
'do_sample': True,
'num_beams': 1
}
else:
raise ValueError(f"Unknown mode: {mode}")
with torch.no_grad():
output_ids = self.model.generate(input_ids, **config)
return self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
