Applying low-rank approximation to optimizer states in LLMs will reduce memory overhead in multi-agent financial trading simulations.

Applying low-rank approximation (specifically rank-r decomposition where r << d_model) to optimizer states (Adam/AdamW first and second moment estimates) during fine-tuning or continual learning of large language models (≥7B parameters) deployed in multi-agent financial trading simulations will reduce peak optimizer-state memory footprint by ≥40% relative to full-rank Adam baselines, without degrading agent decision quality (Sharpe ratio degradation <5%, trade execution accuracy degradation <3%) or increasing wall-clock training time by more than 15%.

Minimum Viable Test (MVT) — 3-phase design:

Phase A (Gradient Rank Audit, 2 days): Profile gradient matrices of a 7B-parameter LLM (Llama-3-8B or Mistral-7B) fine-tuned on financial instruction data for 500 steps. Compute SVD of gradient tensors for all attention and MLP weight matrices. Measure cumulative singular value energy to determine empirical rank r* at 80%, 90%, 95% thresholds. This is a go/no-go gate: if r* > 0.25 × min(d) for >50% of layers, abort.

Phase B (Memory Benchmark, 3 days): Implement GaLore-style or FLORA-style low-rank optimizer state projection. Compare peak GPU memory (nvidia-smi dmon at 100ms intervals) for full-rank Adam vs. low-rank Adam at r ∈ {4, 8, 16, 32, 64} across 1,000 fine-tuning steps on FinPile or Bloomberg financial corpus subset. Single agent, single GPU baseline.

Phase C (Multi-Agent Trading Simulation, 5 days): Deploy 4, 8, and 16 concurrent agents in a FinRL or custom OpenAI Gym trading environment (S&P 500 constituents, 2020–2024 minute bars). Each agent runs continual learning updates every 1,000 steps. Compare full-rank vs. low-rank (best r from Phase B) on: peak aggregate GPU memory, Sharpe ratio, max drawdown, trade accuracy, and wall-clock time per update cycle.

# ============================================================
# LOW-RANK OPTIMIZER STATE COMPRESSION FOR MULTI-AGENT LLM TRADING
# Architecture Outline
# ============================================================

import torch
import torch.nn as nn
from torch.optim import AdamW
from typing import Dict, Tuple, Optional
import numpy as np

# --- CORE: Low-Rank Adam Optimizer ---
class LowRankAdam(torch.optim.Optimizer):
    """
    Adam optimizer with low-rank projection of moment states.
    Based on GaLore (Zhao et al., 2024) extended to multi-agent setting.
    
    Memory: O(r * (m + n)) vs O(m * n) for full-rank Adam moments.
    Reduction ratio: 1 - 2r/(m+n) per layer.
    """
    def __init__(
        self,
        params,
        lr: float = 2e-5,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 0.01,
        rank: int = 16,                    # r << min(m, n)
        update_proj_gap: int = 200,        # T: steps between SVD updates
        scale: float = 1.0,
    ):
        defaults = dict(
            lr=lr, betas=betas, eps=eps,
            weight_decay=weight_decay,
            rank=rank,
            update_proj_gap=update_proj_gap,
            scale=scale,
        )
        super().__init__(params, defaults)
    
    def _get_low_rank_projection(
        self,
        grad: torch.Tensor,
        rank: int,
        state: Dict,
        step: int,
        update_gap: int,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Compute or retrieve low-rank projection matrices P, Q.
        G ≈ P @ G_lr @ Q^T  where G_lr ∈ R^{r×r}
        
        SVD cost: O(m*n*min(m,n)) amortized over update_gap steps.
        """
        if step % update_gap == 0 or 'proj_P' not in state:
            # Compute SVD of current gradient
            # grad shape: (m, n)
            U, S, Vh = torch.linalg.svd(
                grad.float(),
                full_matrices=False
            )
            # Keep top-r singular vectors
            state['proj_P'] = U[:, :rank].to(grad.dtype)   # (m, r)
            state['proj_Q'] = Vh[:rank, :].T.to(grad.dtype) # (n, r)
            
            # Log singular value energy for monitoring
            energy_ratio = (S[:rank]**2).sum() / (S**2).sum()
            state['sv_energy_ratio'] = energy_ratio.item()
        
        P = state['proj_P']  # (m, r)
        Q = state['proj_Q']  # (n, r)
        
        # Project gradient to low-rank subspace
        # G_lr = P^T @ G @ Q ∈ R^{r×r}
        G_lr = P.T @ grad @ Q
        
        return G_lr, P, Q
    
    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        
        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue
                
                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError("LowRankAdam: sparse gradients not supported")
                
                state = self.state[p]
                rank = group['rank']
                
                # Initialize state
                if len(state) == 0:
                    state['step'] = 0
                    # Only store moments for low-rank projection
                    # Memory: 2 * r * r vs 2 * m * n
                    state['exp_avg'] = torch.zeros(rank, rank, device=p.device, dtype=p.dtype)
                    state['exp_avg_sq'] = torch.zeros(rank, rank, device=p.device, dtype=p.dtype)
                
                state['step'] += 1
                step = state['step']
                beta1, beta2 = group['betas']
                
                # Handle 2D weight matrices (attention, MLP layers)
                if grad.dim() == 2:
                    G_lr, P, Q = self._get_low_rank_projection(
                        grad, rank, state,
                        step, group['update_proj_gap']
                    )
                    
                    # Adam update in low-rank subspace
                    exp_avg = state['exp_avg']
                    exp_avg_sq = state['exp_avg_sq']
                    
                    exp_avg.mul_(beta1).add_(G_lr, alpha=1 - beta1)
                    exp_avg_sq.mul_(beta2).addcmul_(G_lr, G_lr, value=1 - beta2)
                    
                    bias_correction1 = 1 - beta1 ** step
                    bias_correction2 = 1 - beta2 ** step
                    
                    step_size = group['lr'] / bias_correction1
                    denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(group['eps'])
                    
                    # Compute update in low-rank space
                    update_lr = exp_avg / denom  # (r, r)
                    
                    # Project back to full parameter space
                    # update = P @ update_lr @ Q^T ∈ R^{m×n}
                    update_full = P @ update_lr @ Q.T
                    
                    # Apply weight decay + parameter update
                    p.mul_(1 - group['lr'] * group['weight_decay'])
                    p.add_(update_full, alpha=-step_size * group['scale'])
                
                else:
                    # Fallback: standard Adam for 1D params (biases, LayerNorm)
                    # These are small — no memory savings needed
                    if 'exp_avg_full' not in state:
                        state['exp_avg_full'] = torch.zeros_like(p)
                        state['exp_avg_sq_full'] = torch.zeros_like(p)
                    
                    exp_avg = state['exp_avg_full']
                    exp_avg_sq = state['exp_avg_sq_full']
                    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                    
                    bias_correction1 = 1 - beta1 ** step
                    bias_correction2 = 1 - beta2 ** step
                    step_size = group['lr'] / bias_correction1
                    denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(group['eps'])
                    
                    p.mul_(1 - group['lr'] * group['weight_decay'])
                    p.addcdiv_(exp_avg, denom, value=-step_size)
        
        return loss


# --- MEMORY PROFILER ---
class OptimizerMemoryProfiler:
    """Track optimizer state memory across training steps."""
    
    def __init__(self, optimizer: torch.optim.Optimizer, model: nn.Module):
        self.optimizer = optimizer
        self.model = model
        self.memory_log = []
    
    def compute_optimizer_state_bytes(self) -> int:
        """Compute exact bytes used by optimizer states."""
        total_bytes = 0
        for group in self.optimizer.param_groups:
            for p in group['params']:
                state = self.optimizer.state[p]
                for key, val in state.items():
                    if isinstance(val, torch.Tensor):
                        total_bytes += val.numel() * val.element_size()
        return total_bytes
    
    def log_step(self, step: int):
        opt_bytes = self.compute_optimizer_state_bytes()
        gpu_bytes = torch.cuda.max_memory_allocated()
        self.memory_log.append({
            'step': step,
            'optimizer_state_mb': opt_bytes / 1e6,
            'peak_gpu_mb': gpu_bytes / 1e6,
        })
    
    def report(self) -> Dict:
        if not self.memory_log:
            return {}
        peak_opt = max(d['optimizer_state_mb'] for d in self.memory_log)
        peak_gpu = max(d['peak_gpu_mb'] for d in self.memory_log)
        return {
            'peak_optimizer_state_mb': peak_opt,
            'peak_gpu_mb': peak_gpu,
            'n_steps_logged': len(self.memory_log),
        }


# --- MULTI-AGENT TRADING ENVIRONMENT WRAPPER ---
class MultiAgentLLMTradingEnv:
    """
    Orchestrates N concurrent LLM-based trading agents.
    Each agent: frozen LLM trunk + trainable LoRA adapters.
    Continual learning via LowRankAdam on adapter parameters.
    
    Memory bottleneck: N × optimizer_state_size (per-agent Adam moments)
    Low-rank compression reduces this by ~40% per agent.
    """
    
    def __init__(
        self,
        n_agents: int,
        model_name: str = "meta-llama/Meta-Llama-3-8B",
        rank: int = 16,
        use_low_rank_optimizer: bool = True,
        device_map: str = "auto",
    ):
        self.n_agents = n_agents
        self.rank = rank
        self.use_low_rank = use_low_rank_optimizer
        self.agents = []
        self.optimizers = []
        self.profilers = []
        
        # Initialize agents (shared frozen trunk, separate adapters)
        self._init_agents(model_name, device_map)
    
    def _init_agents(self, model_name: str, device_map: str):
        """
        Load shared LLM backbone once, create per-agent adapter layers.
        Optimizer states only allocated for adapter parameters.
        
        Adapter param count: ~4M for rank-16 LoRA on 8B model
        Full-rank Adam states: 4M × 2 × 4 bytes = 32 MB per agent
        Low-rank Adam states: ~0.4M × 2 × 4 bytes = 3.2 MB per agent
        N=16 agents: 512 MB → 51 MB (90% reduction for adapter states)
        """
        from transformers import AutoModelForCausalLM
        from peft import get_peft_model, LoraConfig
        
        # Load shared backbone (frozen)
        backbone = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16,
            device_map=device_map,
        )
        backbone.requires_grad_(False)
        
        for agent_id in range(self.n_agents):
            # Per-agent LoRA adapter (trainable)
            lora_config = LoraConfig(
                r=self.rank,
                lora_alpha=32,
                target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
                lora_dropout=0.05,
                bias="none",
            )
            agent_model = get_peft_model(backbone, lora_config)
            trainable_params = [p for p in agent_model.parameters() if p.requires_grad]
            
            # Select optimizer
            if self.use_low_rank:
                optimizer = LowRankAdam(
                    trainable_params,
                    lr=2e-5,
                    rank=min(self.rank, 8),  # r for optimizer states
                    update_proj_gap=200,
                )
            else:
                optimizer = AdamW(trainable_params, lr=2e-5)
            
            profiler = OptimizerMemoryProfiler(optimizer, agent_model)
            
            self.agents.append(agent_model)
            self.optimizers.append(optimizer)
            self.profilers.append(profiler)
    
    def continual_update(
        self,
        agent_id: int,
        market_observations: torch.Tensor,
        rewards: torch.Tensor,
        step: int,
    ):
        """
        Perform one continual learning update for agent agent_id.
        Called every 1,000 environment steps.
        """
        agent = self.agents[agent_id]
        optimizer = self.optimizers[agent_id]
        profiler = self.profilers[agent_id]
        
        optimizer.zero_grad()
        
        # Forward pass: LLM generates trading action
        outputs = agent(
            input_ids=market_observations,
            labels=market_observations,  # simplified; use RL loss in practice
        )
        loss = outputs.loss
        
        # Backward pass
        loss.backward()
        
        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(
            [p for p in agent.parameters() if p.requires_grad],
            max_norm=1.0,
        )
        
        # Optimizer step (low-rank or full-rank)
        optimizer.step()
        
        # Log memory
        profiler.log_step(step)
        
        return loss.item()
    
    def aggregate_memory_report(self) -> Dict:
        """Aggregate memory stats across all agents."""
        reports = [p.report() for p in self.profilers]
        total_opt_mb = sum(r.get('peak_optimizer_state_mb', 0) for r in reports)
        return {
            'n_agents': self.n_agents,
            'total_optimizer_state_mb': total_opt_mb,
            'per_agent_optimizer_state_mb': total_opt_mb / self.n_agents,
            'use_low_rank': self.use_low_rank,
        }


# --- TRADING PERFORMANCE EVALUATOR ---
class TradingPerformanceEvaluator:
    """
    Evaluate agent trading performance on held-out 2024 data.
    Metrics: Sharpe, Sortino, max drawdown, win rate, total return.
    """
    
    def __init__(self, initial_capital: float = 1_000_000.0):
        self.initial_capital = initial_capital
        self.returns_history = []
    
    def compute_sharpe(self, daily_returns: np.ndarray, risk_free_rate: float = 0.05) -> float:
        """Annualized Sharpe ratio."""
        excess = daily_returns - risk_free_rate / 252
        if excess.std() == 0:
            return 0.0
        return np.sqrt(252) * excess.mean() / excess.std()
    
    def compute_max_drawdown(self, cumulative_returns: np.ndarray) -> float:
        """Maximum drawdown as fraction of peak."""
        peak = np.maximum.accumulate(cumulative_returns)
        drawdown = (cumulative_returns - peak) / peak
        return drawdown.min()
    
    def evaluate(self, daily_returns: np.ndarray) -> Dict:
        sharpe = self.compute_sharpe(daily_returns)
        cum_returns = (1 + daily_returns).cumprod()
        max_dd = self.compute_max_drawdown(cum_returns)
        win_rate = (daily_returns > 0).mean()
        total_return = cum_returns[-1] - 1
        
        return {
            'sharpe_ratio': sharpe,
            'max_drawdown': max_dd,
            'win_rate': win_rate,
            'total_return': total_return,
            'annualized_return': (1 + total_return) ** (252 / len(daily_returns)) - 1,
        }


# --- MAIN EXPERIMENT RUNNER ---
def run_experiment(
    n_agents: int = 4,
    rank: int = 16,
    n_seeds: int = 10,
    n_trading_days: int = 252,
):
    """
    Full experiment: compare full-rank vs. low-rank Adam
    in multi-agent LLM trading simulation.
    
    Expected runtime: ~32 GPU-hours per configuration.
    """
    results = {}
    
    for use_low_rank in [False, True]:
        config_name = f"low_rank_r{rank}" if use_low_rank else "full_rank"
        results[config_name] = {
            'memory_reports': [],
            'trading_metrics': [],
        }
        
        for seed in range(n_seeds):
            torch.manual_seed(seed)
            np.random.seed(seed)
            
            # Initialize environment
            env = MultiAgentLLMTradingEnv(
                n_agents=n_agents,
                rank=rank,
                use_low_rank_optimizer=use_low_rank,
            )
            
            evaluator = TradingPerformanceEvaluator()
            
            # Simulate trading + continual learning
            daily_returns = []
            for day in range(n_trading_days):
                # [Simplified: replace with actual FinRL environment step]
                market_obs = torch.randint(0, 32000, (1, 512))  # tokenized market data
                rewards = torch.randn(n_agents)
                
                # Continual learning update every 10 days
                if day % 10 == 0:
                    for agent_id in range(n_agents):
                        env.continual_update(agent_id, market_obs, rewards, step=day)
                
                # [Simplified: replace with actual portfolio return computation]
                daily_return = np.random.normal(0.001, 0.015)  # placeholder
                daily_returns.append(daily_return)
            
            # Collect results
            mem_report = env.aggregate_memory_report()
            trade_metrics = evaluator.evaluate(np.array(daily_returns))
            
            results[config_name]['memory_reports'].append(mem_report)
            results[config_name]['trading_metrics'].append(trade_metrics)
            
            print(f"[{config_name}] Seed {seed}: "
                  f"Optimizer memory = {mem_report['total_optimizer_state_mb']:.1f} MB, "
                  f"Sharpe = {trade_metrics['sharpe_ratio']:.3f}")
    
    # Compute summary statistics
    for config_name, data in results.items():
        mem_values = [r['total_optimizer_state_mb'] for r in data['memory_reports']]
        sharpe_values = [m['sharpe_ratio'] for m in data['trading_metrics']]
        
        print(f"\n=== {config_name} SUMMARY ===")
        print(f"Optimizer Memory: {np.mean(mem_values):.1f} ± {np.std(mem_values):.1f} MB")
        print(f"Sharpe Ratio:     {np.mean(sharpe_values):.3f} ± {np.std(sharpe_values):.3f}")
    
    return results


# --- GRADIENT RANK AUDIT ---
def audit_gradient_rank(
    model: nn.Module,
    dataloader,
    n_steps: int = 500,
    energy_threshold: float = 0.90,
) -> Dict:
    """
    Phase A: Measure effective rank of gradient matrices.
    GO/NO-GO gate: median r* ≤ 0.15 × min(d_row, d_col)
    """
    rank_ratios = []
    
    for step, batch in enumerate(dataloader):
        if step >= n_steps:
            break
        
        outputs = model(**batch)
        outputs.loss.backward()
        
        if step % 100 == 0:
            for name, param in model.named_parameters():
                if param.grad is not None and param.grad.dim() == 2:
                    grad = param.grad.float()
                    m, n = grad.shape
                    
                    # SVD
                    _, S, _ = torch.linalg.svd(grad, full_matrices=False)
                    
                    # Find r* at energy_threshold
                    cumulative_energy = (S**2).cumsum(0) / (S**2).sum()
                    r_star = (cumulative_energy < energy_threshold).sum().item() + 1
                    rank_ratio = r_star / min(m, n)
                    rank_ratios.append(rank_ratio)
        
        model.zero_grad()
    
    median_ratio = np.median(rank_ratios)
    go_decision = median_ratio <= 0.15
    
    return {
        'median_rank_ratio': median_ratio,
        'mean_rank_ratio': np.mean(rank_ratios),
        'p90_rank_ratio': np.percentile(rank_ratios, 90),
        'go_no_go': 'GO' if go_decision else 'NO-GO',
        'n_layers_audited': len(rank_ratios),
        'energy_threshold': energy_threshold,
    }


if __name__ == "__main__":
    # Phase A: Gradient rank audit (abort if NO-GO)
    # audit_result = audit_gradient_rank(model, financial_dataloader)
    # assert audit_result['go_no_go'] == 'GO', f"ABORT: {audit_result}"
    
    # Phase B+C: Full experiment
    results = run_experiment(
        n_agents=4,
        rank=16,
        n_seeds=10,
        n_trading_days=252,
    )