fegger
a763ab0774
- Detect AMD GPUs via ROCm device name in config - Replace single-timestep LSTM in GNN model with leaner post-GNN MLP - Pass edge_attr through AMDGATConv propagate and use ones for self-loops - Fix live trading to sell existing positions on negative signals instead of skipping them entirely - Use per-file try/except in data pipeline and batch SQLite inserts - Import torch directly in backtester instead of dynamic __import__ - Update AMP autocast import for PyTorch 2.0+ compatibility
344 lines
12 KiB
Python
344 lines
12 KiB
Python
import logging
|
|
import time
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch_geometric.nn import MessagePassing
|
|
from torch_geometric.utils import softmax
|
|
|
|
from config import config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AMDOptimizer:
|
|
"""
|
|
AMD-specific optimizations for PyTorch models
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.device = torch.device(config.DEVICE)
|
|
self._configure_rocm()
|
|
|
|
def _configure_rocm(self):
|
|
"""Configure ROCm for optimal performance"""
|
|
if config.DEVICE == "cuda" and config.AMD_GPU:
|
|
try:
|
|
# Set ROCm optimization level
|
|
torch.backends.hip.set_optimization_level(config.ROCM_OPT_LEVEL)
|
|
|
|
# Enable memory efficient attention if available
|
|
try:
|
|
from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
|
|
|
|
torch.backends.cuda.enable_flash_sdp(True)
|
|
logger.info("Enabled Flash Attention for AMD GPU")
|
|
except ImportError:
|
|
logger.warning("xformers not available, using standard attention")
|
|
|
|
# Configure memory limits
|
|
total_memory = torch.cuda.get_device_properties(0).total_memory
|
|
memory_limit = int(total_memory * config.GPU_MEMORY_LIMIT)
|
|
torch.cuda.set_per_process_memory_fraction(config.GPU_MEMORY_LIMIT, 0)
|
|
|
|
logger.info(
|
|
f"Configured ROCm with optimization level {config.ROCM_OPT_LEVEL}"
|
|
)
|
|
logger.info(
|
|
f"GPU Memory: {total_memory / 1024**3:.2f}GB, Limit: {memory_limit / 1024**3:.2f}GB"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error configuring ROCm: {str(e)}")
|
|
|
|
def optimize_model(self, model: nn.Module):
|
|
"""Apply AMD-specific optimizations to a model"""
|
|
model = model.to(self.device)
|
|
|
|
if config.DEVICE != "cuda" or not config.AMD_GPU:
|
|
return model
|
|
|
|
try:
|
|
# Apply memory optimizations
|
|
model = self._apply_memory_optimizations(model)
|
|
|
|
logger.info("Applied AMD optimizations to model")
|
|
return model
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error optimizing model: {str(e)}")
|
|
return model
|
|
|
|
def _apply_mixed_precision(self, model: nn.Module):
|
|
"""Mixed precision is handled via torch.autocast in the training loop.
|
|
Permanently casting parameters causes numerical instability; this is a no-op."""
|
|
logger.info(
|
|
f"Mixed precision ({config.PRECISION}) handled via torch.autocast in training loop"
|
|
)
|
|
return model
|
|
|
|
def _apply_memory_optimizations(self, model: nn.Module):
|
|
"""Apply memory optimizations to the model"""
|
|
# Enable gradient checkpointing for memory efficiency
|
|
if (
|
|
hasattr(model, "supports_gradient_checkpointing")
|
|
and model.supports_gradient_checkpointing
|
|
):
|
|
model.gradient_checkpointing_enable()
|
|
logger.info("Enabled gradient checkpointing")
|
|
|
|
# Apply activation checkpointing to specific modules
|
|
for name, module in model.named_modules():
|
|
if isinstance(module, (nn.LSTM, nn.GRU)):
|
|
module.activation_checkpointing = True
|
|
|
|
return model
|
|
|
|
def get_precision_dtype(self):
|
|
"""Get the precision dtype for mixed precision training"""
|
|
if config.PRECISION == "fp16":
|
|
return torch.float16
|
|
elif config.PRECISION == "bf16":
|
|
return torch.bfloat16
|
|
else:
|
|
return torch.float32
|
|
|
|
def benchmark_model(self, model: nn.Module, input_data, num_runs: int = 100):
|
|
"""Benchmark model performance on AMD GPU"""
|
|
if config.DEVICE != "cuda":
|
|
logger.warning("Benchmarking only supported on GPU")
|
|
return {}
|
|
|
|
try:
|
|
# Warm up
|
|
for _ in range(10):
|
|
_ = model(input_data)
|
|
|
|
# Benchmark inference
|
|
start_time = time.time()
|
|
for _ in range(num_runs):
|
|
with torch.no_grad():
|
|
_ = model(input_data)
|
|
inference_time = (time.time() - start_time) / num_runs
|
|
|
|
# Benchmark training
|
|
model.train()
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=config.LEARNING_RATE)
|
|
criterion = nn.MSELoss()
|
|
|
|
start_time = time.time()
|
|
for _ in range(num_runs):
|
|
optimizer.zero_grad()
|
|
out = model(input_data)
|
|
loss = criterion(out, torch.randn_like(out))
|
|
loss.backward()
|
|
optimizer.step()
|
|
training_time = (time.time() - start_time) / num_runs
|
|
|
|
# Memory usage
|
|
memory_allocated = torch.cuda.memory_allocated(0)
|
|
max_memory = torch.cuda.max_memory_allocated(0)
|
|
|
|
return {
|
|
"inference_time": inference_time,
|
|
"training_time": training_time,
|
|
"throughput_inference": 1 / inference_time,
|
|
"throughput_training": 1 / training_time,
|
|
"memory_allocated": memory_allocated,
|
|
"max_memory": max_memory,
|
|
"memory_usage_percent": (
|
|
memory_allocated / torch.cuda.get_device_properties(0).total_memory
|
|
)
|
|
* 100,
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error benchmarking model: {str(e)}")
|
|
return {}
|
|
|
|
|
|
class AMDSparseAttention(nn.Module):
|
|
"""
|
|
Sparse attention implementation optimized for AMD GPUs
|
|
"""
|
|
|
|
def __init__(self, embed_dim, num_heads, dropout=0.1):
|
|
super().__init__()
|
|
self.embed_dim = embed_dim
|
|
self.num_heads = num_heads
|
|
self.head_dim = embed_dim // num_heads
|
|
self.scaling = self.head_dim**-0.5
|
|
|
|
self.qkv_proj = nn.Linear(embed_dim, embed_dim * 3)
|
|
self.out_proj = nn.Linear(embed_dim, embed_dim)
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
# Initialize weights
|
|
self._init_weights()
|
|
|
|
def _init_weights(self):
|
|
nn.init.xavier_uniform_(self.qkv_proj.weight)
|
|
nn.init.xavier_uniform_(self.out_proj.weight)
|
|
nn.init.zeros_(self.qkv_proj.bias)
|
|
nn.init.zeros_(self.out_proj.bias)
|
|
|
|
def forward(self, x, mask=None):
|
|
batch_size, seq_len, embed_dim = x.size()
|
|
|
|
# Project queries, keys, values
|
|
qkv = self.qkv_proj(x)
|
|
q, k, v = qkv.chunk(3, dim=-1)
|
|
|
|
# Reshape for multi-head attention
|
|
q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
|
|
k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
|
|
v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
|
|
|
|
# Compute attention scores
|
|
attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scaling
|
|
|
|
# Apply mask if provided
|
|
if mask is not None:
|
|
attn_scores = attn_scores.masked_fill(mask == 0, float("-inf"))
|
|
|
|
# Compute attention weights
|
|
attn_weights = F.softmax(attn_scores, dim=-1)
|
|
attn_weights = self.dropout(attn_weights)
|
|
|
|
# Apply attention to values
|
|
output = torch.matmul(attn_weights, v)
|
|
|
|
# Concatenate heads
|
|
output = (
|
|
output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
|
|
)
|
|
|
|
# Final projection
|
|
output = self.out_proj(output)
|
|
|
|
return output
|
|
|
|
|
|
class AMDGATConv(MessagePassing):
|
|
"""
|
|
GATConv implementation optimized for AMD GPUs
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels,
|
|
out_channels,
|
|
heads=1,
|
|
concat=True,
|
|
dropout=0.6,
|
|
add_self_loops=True,
|
|
):
|
|
super().__init__(aggr="add", node_dim=0)
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.heads = heads
|
|
self.concat = concat
|
|
self.dropout = dropout
|
|
self.add_self_loops = add_self_loops
|
|
|
|
# Linear transformations for each head
|
|
self.lin_src = nn.Parameter(torch.Tensor(in_channels, heads * out_channels))
|
|
self.lin_dst = nn.Parameter(torch.Tensor(in_channels, heads * out_channels))
|
|
|
|
# Attention parameters
|
|
self.att_src = nn.Parameter(torch.Tensor(1, heads, out_channels))
|
|
self.att_dst = nn.Parameter(torch.Tensor(1, heads, out_channels))
|
|
|
|
# Bias
|
|
self.bias = nn.Parameter(torch.Tensor(heads * out_channels))
|
|
|
|
# Initialize weights
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self):
|
|
nn.init.xavier_uniform_(self.lin_src)
|
|
nn.init.xavier_uniform_(self.lin_dst)
|
|
nn.init.xavier_uniform_(self.att_src)
|
|
nn.init.xavier_uniform_(self.att_dst)
|
|
nn.init.zeros_(self.bias)
|
|
|
|
def forward(self, x, edge_index, edge_attr=None, size=None):
|
|
# Linear transformation
|
|
if size is None and torch.is_tensor(x):
|
|
x_src = x_dst = torch.matmul(x, self.lin_src).view(
|
|
-1, self.heads, self.out_channels
|
|
)
|
|
else:
|
|
x_src, x_dst = x[0], x[1]
|
|
x_src = torch.matmul(x_src, self.lin_src).view(
|
|
-1, self.heads, self.out_channels
|
|
)
|
|
x_dst = torch.matmul(x_dst, self.lin_dst).view(
|
|
-1, self.heads, self.out_channels
|
|
)
|
|
|
|
# Add self loops if needed
|
|
if self.add_self_loops:
|
|
num_nodes = x_src.size(0)
|
|
edge_index, edge_attr = self._add_self_loops(
|
|
edge_index, edge_attr, num_nodes
|
|
)
|
|
|
|
# Compute attention coefficients
|
|
alpha_src = (x_src * self.att_src).sum(dim=-1)
|
|
alpha_dst = (x_dst * self.att_dst).sum(dim=-1)
|
|
alpha = (alpha_src, alpha_dst)
|
|
|
|
# Propagate — pass edge_attr so correlation weights reach message()
|
|
out = self.propagate(edge_index, x=(x_src, x_dst), alpha=alpha, edge_attr=edge_attr, size=size)
|
|
|
|
# Concatenate or average heads
|
|
if self.concat:
|
|
out = out.view(-1, self.heads * self.out_channels)
|
|
else:
|
|
out = out.mean(dim=1)
|
|
|
|
# Add bias
|
|
out = out + self.bias
|
|
|
|
return out
|
|
|
|
def _add_self_loops(self, edge_index, edge_attr, num_nodes):
|
|
# Add self loops to edge_index
|
|
loop_index = torch.arange(
|
|
0, num_nodes, dtype=torch.long, device=edge_index.device
|
|
)
|
|
loop_index = loop_index.unsqueeze(0).repeat(2, 1)
|
|
|
|
if edge_attr is not None:
|
|
# Self-loop weight = 1.0 so self-messages are not zeroed out
|
|
loop_attr = edge_attr.new_ones((num_nodes,) + edge_attr.size()[1:])
|
|
edge_attr = torch.cat([edge_attr, loop_attr], dim=0)
|
|
|
|
edge_index = torch.cat([edge_index, loop_index], dim=1)
|
|
return edge_index, edge_attr
|
|
|
|
def propagate(self, edge_index, size=None, **kwargs):
|
|
return super().propagate(edge_index, size=size, **kwargs)
|
|
|
|
def message(self, x_j, alpha_j, alpha_i, edge_attr, index, ptr, size_i):
|
|
# Compute attention weights
|
|
alpha = alpha_j + alpha_i
|
|
alpha = F.leaky_relu(alpha, negative_slope=0.2)
|
|
alpha = self._softmax(alpha, index, ptr, size_i)
|
|
alpha = F.dropout(alpha, p=self.dropout, training=self.training)
|
|
|
|
# (E, heads, out_channels) scaled by attention
|
|
msg = x_j * alpha.unsqueeze(-1)
|
|
|
|
# Scale by edge correlation weight: (E, 1) → broadcasts to (E, heads, out_channels)
|
|
if edge_attr is not None:
|
|
msg = msg * edge_attr.unsqueeze(1)
|
|
|
|
return msg
|
|
|
|
def _softmax(self, src, index, ptr, num_nodes):
|
|
# Memory-efficient softmax
|
|
return softmax(src, index, ptr, num_nodes)
|