Taming Momentum can be used to reduce the memory footprint of LLMs used in biological research.

Applying the "Taming Momentum" optimizer modification (which clips or rescales the momentum buffer in adaptive gradient methods such as Adam/AdamW) to large language models used in biological research workflows reduces peak GPU memory consumption by ≥15% relative to standard Adam/AdamW, without degrading task performance (perplexity, downstream biological NLP accuracy) by more than 2% absolute, across models with ≥1B parameters operating on biological text corpora (genomics literature, clinical notes, protein sequence annotations).

MINIMUM VIABLE TEST (MVT): Fine-tune a 7B-parameter biological LLM (BioMedLM-7B or LLaMA-3-8B fine-tuned on PubMed) on a standardized biological NLP task (BioASQ Question Answering) using: Condition A: Standard AdamW (β1=0.9, β2=0.999, ε=1e-8) Condition B: Taming Momentum AdamW (identical hyperparameters + momentum norm clipping at τ=5.0) Condition C: AdamW + gradient checkpointing (positive control for memory reduction) Condition D: SGD with momentum (negative control — different memory profile)

Measure: (1) peak GPU memory per training step, (2) wall-clock time per step, (3) final perplexity on PubMed held-out set, (4) BioASQ accuracy at convergence.

Replicate across 3 model sizes: 1.3B (BioGPT-Large), 7B (LLaMA-3-8B-BioMed), 13B (BioMedLM-13B or equivalent). Run 5 seeds per condition per model size = 60 total training runs. Each run: 10,000 gradient steps (sufficient for convergence signal on fine-tuning).

FULL VALIDATION adds:

• Ablation over τ ∈ {1.0, 2.0, 5.0, 10.0, 50.0}
• Multi-GPU (DDP) scaling test at 4×A100 and 8×A100
• Comparison against competing memory-reduction methods: 8-bit Adam (bitsandbytes), Adafactor, CAME optimizer
• Biological task battery: BioASQ, MedMCQA, BLURB (6 subtasks), PubMedQA

# ============================================================
# Taming Momentum Optimizer — Core Implementation
# ============================================================

import torch
from torch.optim import AdamW
from typing import Optional

class TamingMomentumAdamW(AdamW):
    """
    AdamW with Taming Momentum: clips momentum buffer norm
    to prevent unbounded growth, reducing effective optimizer
    state memory pressure and improving numerical stability.
    
    Args:
        params: model parameters
        lr: learning rate (default 1e-4 for biological LLM fine-tuning)
        betas: (beta1, beta2) momentum coefficients
        eps: numerical stability epsilon
        weight_decay: L2 regularization coefficient
        tau: momentum clipping threshold (key hyperparameter)
               tau=inf → standard AdamW
               tau=5.0 → recommended starting point
               tau=1.0 → aggressive clipping (may hurt convergence)
    """
    def __init__(self, params, lr=1e-4, betas=(0.9, 0.999),
                 eps=1e-8, weight_decay=0.01, tau=5.0):
        super().__init__(params, lr=lr, betas=betas, eps=eps,
                        weight_decay=weight_decay)
        self.tau = tau
    
    @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
                state = self.state[p]
                
                # Initialize state
                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p)      # m_t
                    state['exp_avg_sq'] = torch.zeros_like(p)   # v_t
                
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                beta1, beta2 = group['betas']
                state['step'] += 1
                
                # Standard Adam momentum update
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                
                # === TAMING MOMENTUM: clip momentum buffer norm ===
                if self.tau < float('inf'):
                    d = exp_avg.numel()  # parameter dimension
                    momentum_norm = exp_avg.norm(2).item()
                    clip_threshold = self.tau * (d ** 0.5)
                    
                    if momentum_norm > clip_threshold:
                        # Rescale momentum buffer in-place (no new allocation)
                        scale_factor = clip_threshold / (momentum_norm + 1e-8)
                        exp_avg.mul_(scale_factor)
                        # NOTE: in-place operation avoids additional memory allocation
                        # This is the key memory-relevant operation
                # ===================================================
                
                # Bias correction
                bias_correction1 = 1 - beta1 ** state['step']
                bias_correction2 = 1 - beta2 ** state['step']
                
                # Compute step size
                step_size = group['lr'] / bias_correction1
                
                # AdamW weight decay (decoupled)
                p.mul_(1 - group['lr'] * group['weight_decay'])
                
                # Parameter update
                denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(group['eps'])
                p.addcdiv_(exp_avg, denom, value=-step_size)
        
        return loss


# ============================================================
# Memory Profiling Harness
# ============================================================

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.utils.data import DataLoader
import wandb

def profile_optimizer_memory(
    model_name: str,          # e.g., "microsoft/biogpt-large"
    optimizer_class,          # TamingMomentumAdamW or AdamW
    optimizer_kwargs: dict,   # e.g., {"tau": 5.0, "lr": 1e-4}
    dataset_path: str,        # path to tokenized PubMed dataset
    n_steps: int = 100,
    seed: int = 42
) -> dict:
    """
    Profile peak GPU memory for a given model + optimizer combination.
    Returns dict with memory breakdown and performance metrics.
    """
    torch.manual_seed(seed)
    torch.cuda.reset_peak_memory_stats()
    
    # Load model in BF16 (standard for LLM fine-tuning)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16,
        device_map="cuda:0"
    )
    
    # Record model parameter memory
    param_