|
import torch |
|
import torch.nn as nn |
|
import numpy as np |
|
import json |
|
import os |
|
from datetime import datetime |
|
|
|
class FrequencyHandler: |
|
"""Base class for parameter-specific frequency analysis functions""" |
|
|
|
def analyze(self, grad_sample, n_bands, eps=1e-8): |
|
"""Default frequency analysis implementation""" |
|
freq_repr = torch.fft.rfft(grad_sample.float()) |
|
freq_power = torch.abs(freq_repr) |
|
|
|
if freq_power.sum() > 0: |
|
freq_power = freq_power / (freq_power.sum() + eps) |
|
band_size = freq_power.shape[0] // n_bands |
|
if band_size <= 0: |
|
return [0.0] * n_bands |
|
|
|
band_powers = [] |
|
for i in range(n_bands): |
|
start_idx = i * band_size |
|
end_idx = min((i+1) * band_size, freq_power.shape[0]) |
|
if start_idx < end_idx: |
|
band_power = freq_power[start_idx:end_idx].sum().item() |
|
band_powers.append(band_power) |
|
else: |
|
band_powers.append(0.0) |
|
|
|
return band_powers |
|
|
|
def get_adaptive_momentum(self, band_values, base_alpha): |
|
"""Default adaptive momentum calculation""" |
|
n_bands = len(band_values) |
|
high_freq_activity = sum(band_values[n_bands//2:]) |
|
|
|
if high_freq_activity > 0.3: |
|
return min(0.95, base_alpha + 0.05) |
|
return base_alpha |
|
|
|
class ConvFrequencyHandler(FrequencyHandler): |
|
"""Specialized handler for convolutional layers""" |
|
|
|
def analyze(self, grad_sample, n_bands, eps=1e-8): |
|
freq_repr = torch.fft.rfft(grad_sample.float()) |
|
freq_power = torch.abs(freq_repr) |
|
|
|
if freq_power.sum() > 0: |
|
freq_power = freq_power / (freq_power.sum() + eps) |
|
band_powers = [] |
|
total_freqs = freq_power.shape[0] |
|
|
|
for i in range(n_bands): |
|
start_idx = int((total_freqs ** (i/n_bands)) - 1) |
|
end_idx = int((total_freqs ** ((i+1)/n_bands)) - 1) |
|
start_idx = max(0, start_idx) |
|
end_idx = min(end_idx, total_freqs) |
|
|
|
if start_idx < end_idx: |
|
band_power = freq_power[start_idx:end_idx].sum().item() |
|
band_powers.append(band_power) |
|
else: |
|
band_powers.append(0.0) |
|
|
|
return band_powers |
|
|
|
def get_adaptive_momentum(self, band_values, base_alpha): |
|
"""Convolutional layers benefit from more smoothing in mid-frequencies""" |
|
n_bands = len(band_values) |
|
mid_freq_activity = sum(band_values[n_bands//4:(3*n_bands)//4]) |
|
high_freq_activity = sum(band_values[(3*n_bands)//4:]) |
|
if mid_freq_activity > 0.4: |
|
return min(0.97, base_alpha + 0.07) |
|
elif high_freq_activity > 0.3: |
|
return min(0.95, base_alpha + 0.05) |
|
return base_alpha |
|
|
|
class AttentionFrequencyHandler(FrequencyHandler): |
|
"""Specialized handler for attention layers""" |
|
|
|
def analyze(self, grad_sample, n_bands, eps=1e-8): |
|
freq_repr = torch.fft.rfft(grad_sample.float()) |
|
freq_power = torch.abs(freq_repr) |
|
|
|
if freq_power.sum() > 0: |
|
freq_power = freq_power / (freq_power.sum() + eps) |
|
band_powers = [] |
|
half_bands = n_bands // 2 |
|
low_band_size = (freq_power.shape[0] // 2) // half_bands |
|
for i in range(half_bands): |
|
start_idx = i * low_band_size |
|
end_idx = min((i+1) * low_band_size, freq_power.shape[0] // 2) |
|
if start_idx < end_idx: |
|
band_power = freq_power[start_idx:end_idx].sum().item() |
|
band_powers.append(band_power) |
|
else: |
|
band_powers.append(0.0) |
|
high_band_size = (freq_power.shape[0] - (freq_power.shape[0] // 2)) // (n_bands - half_bands) |
|
for i in range(half_bands, n_bands): |
|
start_idx = (freq_power.shape[0] // 2) + (i - half_bands) * high_band_size |
|
end_idx = min((freq_power.shape[0] // 2) + (i - half_bands + 1) * high_band_size, freq_power.shape[0]) |
|
if start_idx < end_idx: |
|
band_power = freq_power[start_idx:end_idx].sum().item() |
|
band_powers.append(band_power) |
|
else: |
|
band_powers.append(0.0) |
|
|
|
return band_powers |
|
|
|
def get_adaptive_momentum(self, band_values, base_alpha): |
|
"""Custom adaptive momentum for attention layers""" |
|
n_bands = len(band_values) |
|
max_band_idx = np.argmax(band_values) |
|
if max_band_idx < n_bands // 4: |
|
return max(0.85, base_alpha - 0.05) |
|
elif max_band_idx > 3*n_bands // 4: |
|
return min(0.98, base_alpha + 0.08) |
|
return base_alpha |
|
|
|
class EmbeddingFrequencyHandler(FrequencyHandler): |
|
"""Specialized handler for embedding layers""" |
|
|
|
def get_adaptive_momentum(self, band_values, base_alpha): |
|
"""Embeddings often benefit from very stable updates""" |
|
n_bands = len(band_values) |
|
high_freq_activity = sum(band_values[(3*n_bands)//4:]) |
|
if high_freq_activity > 0.2: |
|
return min(0.98, base_alpha + 0.08) |
|
return base_alpha |
|
|
|
class FAMOptimizer(torch.optim.Optimizer): |
|
""" |
|
Frequency-Adaptive Momentum optimizer with parameter-specific handlers. |
|
|
|
Args: |
|
... (existing parameters) |
|
debug (bool, optional): Whether to collect debug information (default: False) |
|
debug_dir (str, optional): Directory to save debug info (default: './fam_debug') |
|
debug_interval (int, optional): Steps between debug dumps (default: 1000) |
|
""" |
|
def __init__(self, params, lr=1e-3, alpha=0.9, beta=0.99, eps=1e-8, |
|
weight_decay=0.0, n_bands=8, fam_start_step=100, |
|
layer_boost=True, min_size=256, debug=False, |
|
debug_dir='./fam_debug', debug_interval=1000): |
|
defaults = dict(lr=lr, alpha=alpha, beta=beta, eps=eps, |
|
weight_decay=weight_decay, n_bands=n_bands, |
|
fam_start_step=fam_start_step, |
|
layer_boost=layer_boost, min_size=min_size) |
|
self.debug = debug |
|
self.debug_info = {} if debug else None |
|
self.debug_dir = debug_dir |
|
self.debug_interval = debug_interval |
|
self.last_dump_step = 0 |
|
|
|
if debug and debug_dir: |
|
os.makedirs(debug_dir, exist_ok=True) |
|
self.debug_file = os.path.join( |
|
debug_dir, |
|
f"fam_debug_{datetime.now().strftime('%m%d_%H%M%S')}.json" |
|
) |
|
with open(self.debug_file, 'w') as f: |
|
json.dump({ |
|
"optimizer": "FAMOptimizer", |
|
"settings": { |
|
"lr": lr, |
|
"alpha": alpha, |
|
"beta": beta, |
|
"n_bands": n_bands, |
|
"fam_start_step": fam_start_step, |
|
}, |
|
"parameters": {}, |
|
"steps_recorded": [] |
|
}, f, indent=2) |
|
self.handlers = { |
|
"default": FrequencyHandler(), |
|
"conv": ConvFrequencyHandler(), |
|
"attention": AttentionFrequencyHandler(), |
|
"embedding": EmbeddingFrequencyHandler() |
|
} |
|
param_groups = self._add_handlers_to_groups(params) |
|
super(FAMOptimizer, self).__init__(params=param_groups, defaults=defaults) |
|
def _add_handlers_to_groups(self, params): |
|
"""Add appropriate handlers to parameter groups based on type""" |
|
if isinstance(params, list) and all(isinstance(pg, dict) for pg in params): |
|
for pg in params: |
|
if 'handler' not in pg: |
|
if any('conv' in name.lower() for name in pg.get('names', [])): |
|
pg['handler'] = 'conv' |
|
elif any(name in name.lower() for name in pg.get('names', []) |
|
for name in ['attention', 'mha', 'self_attn']): |
|
pg['handler'] = 'attention' |
|
elif any(name in name.lower() for name in pg.get('names', []) |
|
for name in ['embed', 'token']): |
|
pg['handler'] = 'embedding' |
|
else: |
|
pg['handler'] = 'default' |
|
return params |
|
else: |
|
return [{'params': params, 'handler': 'default'}] |
|
|
|
def get_handler(self, group): |
|
"""Get the appropriate frequency handler for the parameter group""" |
|
handler_name = group.get('handler', 'default') |
|
return self.handlers[handler_name] |
|
|
|
def dump_debug_info(self, force=False): |
|
"""Save the current debug information to file""" |
|
if not self.debug or not hasattr(self, 'debug_file'): |
|
return |
|
current_step = max([self.state[p]['step'] for p in self.state], default=0) |
|
if force or (current_step - self.last_dump_step >= self.debug_interval): |
|
try: |
|
with open(self.debug_file, 'r') as f: |
|
debug_data = json.load(f) |
|
debug_data["steps_recorded"].append(current_step) |
|
|
|
for param_name, param_info in self.debug_info.items(): |
|
if param_name not in debug_data["parameters"]: |
|
debug_data["parameters"][param_name] = { |
|
"handler": param_info.get('handler', 'default'), |
|
"steps": [], |
|
"bands": [], |
|
"alpha": [] |
|
} |
|
last_recorded = len(debug_data["parameters"][param_name]["steps"]) |
|
if last_recorded < len(param_info['steps']): |
|
debug_data["parameters"][param_name]["steps"].extend(param_info['steps'][last_recorded:]) |
|
debug_data["parameters"][param_name]["bands"].extend(param_info['bands'][last_recorded:]) |
|
debug_data["parameters"][param_name]["alpha"].extend(param_info['alpha'][last_recorded:]) |
|
with open(self.debug_file, 'w') as f: |
|
json.dump(debug_data, f) |
|
|
|
self.last_dump_step = current_step |
|
for param_info in self.debug_info.values(): |
|
param_info['steps'] = param_info['steps'][-10:] |
|
param_info['bands'] = param_info['bands'][-10:] |
|
param_info['alpha'] = param_info['alpha'][-10:] |
|
|
|
except Exception as e: |
|
print(f"Error dumping FAM debug info: {e}") |
|
|
|
@torch.no_grad() |
|
def step(self, closure=None): |
|
"""Perform a single optimization step.""" |
|
loss = None |
|
if closure is not None: |
|
with torch.enable_grad(): |
|
loss = closure() |
|
|
|
for group in self.param_groups: |
|
for p_idx, p in enumerate(group['params']): |
|
if p.grad is None: |
|
continue |
|
|
|
grad = p.grad |
|
if grad.is_sparse: |
|
raise RuntimeError('FAMOptimizer does not support sparse gradients') |
|
|
|
state = self.state[p] |
|
|
|
if len(state) == 0: |
|
state['step'] = 0 |
|
state['exp_avg'] = torch.zeros_like(p) |
|
state['freq_history'] = {} |
|
state['param_name'] = f"param_{p_idx}" |
|
|
|
state['step'] += 1 |
|
|
|
if group['weight_decay'] != 0: |
|
grad = grad.add(p, alpha=group['weight_decay']) |
|
|
|
exp_avg = state['exp_avg'] |
|
alpha = group['alpha'] |
|
beta = group['beta'] |
|
lr = group['lr'] |
|
n_bands = group['n_bands'] |
|
handler = self.get_handler(group) |
|
|
|
should_apply_fam = ( |
|
state['step'] > group['fam_start_step'] and |
|
p.numel() > group['min_size'] |
|
) |
|
|
|
if should_apply_fam: |
|
try: |
|
if p.numel() > 10000: |
|
if p.dim() > 1: |
|
row_indices = torch.randperm(p.size(0))[:min(p.size(0), 64)] |
|
col_indices = torch.randperm(p.size(1))[:min(p.size(1), 64)] |
|
grad_sample = grad[row_indices][:, col_indices].flatten() |
|
else: |
|
sample_idx = torch.randperm(p.numel())[:1000] |
|
grad_sample = grad.flatten()[sample_idx] |
|
else: |
|
grad_sample = grad.flatten() |
|
band_powers = handler.analyze(grad_sample, n_bands, group['eps']) |
|
if state['step'] <= 10 and p_idx == 0: |
|
print(f"Step {state['step']}: Found {len(band_powers)} frequency bands") |
|
print(f"Band powers: {[f'{v:.4f}' for v in band_powers]}") |
|
for i, power in enumerate(band_powers): |
|
band_key = f'band_{i}' |
|
if band_key not in state['freq_history']: |
|
state['freq_history'][band_key] = power |
|
else: |
|
state['freq_history'][band_key] = ( |
|
beta * state['freq_history'][band_key] + |
|
(1-beta) * power |
|
) |
|
band_values = [state['freq_history'].get(f'band_{i}', 0) |
|
for i in range(n_bands)] |
|
effective_alpha = handler.get_adaptive_momentum(band_values, alpha) |
|
|
|
if self.debug: |
|
param_name = state['param_name'] |
|
if param_name not in self.debug_info: |
|
self.debug_info[param_name] = { |
|
'steps': [], |
|
'bands': [], |
|
'handler': group.get('handler', 'default'), |
|
'alpha': [] |
|
} |
|
|
|
if state['step'] % 10 == 0: |
|
self.debug_info[param_name]['steps'].append(state['step']) |
|
self.debug_info[param_name]['bands'].append(band_values) |
|
self.debug_info[param_name]['alpha'].append(effective_alpha) |
|
exp_avg.mul_(effective_alpha).add_(grad, alpha=1-effective_alpha) |
|
except Exception as e: |
|
import traceback |
|
print(f"Error in FAM processing for parameter {p_idx}:") |
|
print(f"Error type: {type(e).__name__}") |
|
print(f"Error message: {e}") |
|
print(f"Parameter shape: {p.shape}, numel: {p.numel()}") |
|
print(traceback.format_exc()) |
|
exp_avg.mul_(alpha).add_(grad, alpha=1-alpha) |
|
else: |
|
exp_avg.mul_(alpha).add_(grad, alpha=1-alpha) |
|
p.add_(exp_avg, alpha=-lr) |
|
|
|
if self.debug: |
|
self.dump_debug_info() |
|
|
|
return loss |
|
|
|
def __del__(self): |
|
"""Clean up and final debug dump when optimizer is destroyed""" |
|
if self.debug: |
|
self.dump_debug_info(force=True) |
|
|
|
def get_parameter_groups(model, lr=1e-3, weight_decay=0.0): |
|
""" |
|
Create parameter groups for FAMOptimizer with appropriate handlers based on layer type |
|
""" |
|
param_groups = [] |
|
conv_params = [] |
|
conv_names = [] |
|
|
|
attn_params = [] |
|
attn_names = [] |
|
|
|
embed_params = [] |
|
embed_names = [] |
|
|
|
norm_params = [] |
|
norm_names = [] |
|
|
|
other_params = [] |
|
other_names = [] |
|
for name, param in model.named_parameters(): |
|
if not param.requires_grad: |
|
continue |
|
|
|
if any(x in name.lower() for x in ['conv', 'cnn']): |
|
conv_params.append(param) |
|
conv_names.append(name) |
|
elif any(x in name.lower() for x in ['attention', 'mha', 'self_attn']): |
|
attn_params.append(param) |
|
attn_names.append(name) |
|
elif any(x in name.lower() for x in ['embed', 'token']): |
|
embed_params.append(param) |
|
embed_names.append(name) |
|
elif any(x in name.lower() for x in ['norm', 'batch', 'layer']): |
|
norm_params.append(param) |
|
norm_names.append(name) |
|
else: |
|
other_params.append(param) |
|
other_names.append(name) |
|
if conv_params: |
|
param_groups.append({ |
|
'params': conv_params, |
|
'names': conv_names, |
|
'lr': lr, |
|
'weight_decay': weight_decay, |
|
'alpha': 0.9, |
|
'handler': 'conv', |
|
'n_bands': 10 |
|
}) |
|
|
|
if attn_params: |
|
param_groups.append({ |
|
'params': attn_params, |
|
'names': attn_names, |
|
'lr': lr, |
|
'weight_decay': weight_decay, |
|
'alpha': 0.92, |
|
'handler': 'attention', |
|
'n_bands': 12 |
|
}) |
|
|
|
if embed_params: |
|
param_groups.append({ |
|
'params': embed_params, |
|
'names': embed_names, |
|
'lr': lr * 0.8, |
|
'weight_decay': weight_decay * 1.5, |
|
'alpha': 0.95, |
|
'handler': 'embedding', |
|
'n_bands': 8 |
|
}) |
|
|
|
if norm_params: |
|
param_groups.append({ |
|
'params': norm_params, |
|
'names': norm_names, |
|
'lr': lr, |
|
'weight_decay': 0.0, |
|
'alpha': 0.9, |
|
'handler': 'default', |
|
'n_bands': 4 |
|
}) |
|
|
|
if other_params: |
|
param_groups.append({ |
|
'params': other_params, |
|
'names': other_names, |
|
'lr': lr, |
|
'weight_decay': weight_decay, |
|
'alpha': 0.9, |
|
'handler': 'default', |
|
'n_bands': 8 |
|
}) |
|
|
|
return param_groups |
|
|
|
import torch |
|
from torch.optim.lr_scheduler import _LRScheduler |
|
import math |
|
|
|
class FAMSchedulerb(_LRScheduler): |
|
""" |
|
Scheduler with linear warmup followed by cosine annealing. |
|
|
|
Args: |
|
optimizer: Wrapped optimizer |
|
warmup_epochs: Number of epochs for the linear warmup |
|
max_epochs: Total number of epochs |
|
warmup_start_lr: Initial learning rate for warmup |
|
eta_min: Minimum learning rate after cosine annealing |
|
""" |
|
def __init__(self, optimizer, warmup_epochs, max_epochs, warmup_start_lr=1e-8, eta_min=1e-8, last_epoch=-1): |
|
self.warmup_epochs = warmup_epochs |
|
self.max_epochs = max_epochs |
|
self.warmup_start_lr = warmup_start_lr |
|
self.eta_min = eta_min |
|
super(FAMScheduler, self).__init__(optimizer, last_epoch) |
|
|
|
def get_lr(self): |
|
if self.last_epoch < self.warmup_epochs: |
|
alpha = self.last_epoch / self.warmup_epochs |
|
return [self.warmup_start_lr + (base_lr - self.warmup_start_lr) * alpha for base_lr in self.base_lrs] |
|
else: |
|
return [self.eta_min + (base_lr - self.eta_min) * |
|
(1 + math.cos(math.pi * (self.last_epoch - self.warmup_epochs) / |
|
(self.max_epochs - self.warmup_epochs))) / 2 |
|
for base_lr in self.base_lrs] |
|
import torch |
|
import math |
|
|
|
class SimpleFAM(torch.optim.Optimizer): |
|
""" |
|
Simplified Frequency-Adaptive Momentum optimizer |
|
|
|
A lightweight implementation that focuses on the core concepts |
|
without complex debugging or parameter-specific handlers. |
|
""" |
|
def __init__(self, params, lr=0.001, alpha=0.9, beta=0.99): |
|
defaults = dict(lr=lr, alpha=alpha, beta=beta) |
|
super(SimpleFAM, self).__init__(params, defaults) |
|
print(f"SimpleFAM initialized with lr={lr}, alpha={alpha}") |
|
|
|
@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 |
|
|
|
state = self.state[p] |
|
if len(state) == 0: |
|
state['step'] = 0 |
|
state['exp_avg'] = torch.zeros_like(p) |
|
|
|
state['step'] += 1 |
|
exp_avg = state['exp_avg'] |
|
alpha = group['alpha'] |
|
if p.numel() > 1000 and state['step'] > 100: |
|
grad_sample = p.grad.flatten()[:min(1000, p.numel())] |
|
freq = torch.fft.rfft(grad_sample.float()) |
|
power = torch.abs(freq) |
|
half = power.shape[0] // 2 |
|
high_ratio = power[half:].sum() / (power.sum() + 1e-8) |
|
effective_alpha = min(0.98, alpha + 0.05 * high_ratio) |
|
exp_avg.mul_(effective_alpha).add_(p.grad, alpha=1-effective_alpha) |
|
else: |
|
exp_avg.mul_(alpha).add_(p.grad, alpha=1-alpha) |
|
p.add_(exp_avg, alpha=-group['lr']) |
|
|
|
return loss |
|
|
|
class FAMScheduler(torch.optim.lr_scheduler._LRScheduler): |
|
""" |
|
Step-based learning rate scheduler for FAM optimizer |
|
with warmup and cosine annealing. |
|
""" |
|
def __init__(self, optimizer, warmup_steps=1000, total_steps=100000, |
|
decay_start_step=None, warmup_start_lr=1e-6, eta_min=1e-6, |
|
last_epoch=-1): |
|
self.warmup_steps = warmup_steps |
|
self.total_steps = total_steps |
|
self.decay_start_step = decay_start_step if decay_start_step is not None else warmup_steps |
|
self.warmup_start_lr = warmup_start_lr |
|
self.eta_min = eta_min |
|
super(FAMScheduler, self).__init__(optimizer, last_epoch) |
|
|
|
def get_lr(self): |
|
if self.last_epoch < self.warmup_steps: |
|
alpha = self.last_epoch / self.warmup_steps |
|
return [self.warmup_start_lr + (base_lr - self.warmup_start_lr) * alpha |
|
for base_lr in self.base_lrs] |
|
|
|
elif self.last_epoch < self.decay_start_step: |
|
return self.base_lrs |
|
|
|
else: |
|
return [self.eta_min + (base_lr - self.eta_min) * |
|
(1 + math.cos(math.pi * (self.last_epoch - self.decay_start_step) / |
|
(self.total_steps - self.decay_start_step))) / 2 + 1e-8 |
|
for base_lr in self.base_lrs] |
