adu2021 RAMP: Reinforcement Adaptive Mixed Precision Quantization
Quantizing large language models for on-device inference requires assigning different bit-widths to different layers—a combinatorial optimization problem that standard uniform quantization handles poorly. Layers vary dramatically in their sensitivity to quantization: some can survive 2-bit precision while others need 8-bits to maintain quality.
RAMP solves this through an off-policy reinforcement learning approach that learns which layers should receive which bit-widths. The key innovation is that the learned policy generalizes across models—a policy trained on Llama 2 7B zero-shot transfers to Llama 2 13B and Mistral 7B without retraining. This is possible because quantization sensitivity is fundamentally architectural rather than model-specific.
Core Concept
RAMP reframes mixed-precision quantization as a Markov Decision Process (MDP):
State: Layer index + activation statistics (11-dim embedding capturing weight properties, activation distributions, structural info)
Action: Bit-width assignment (typically 2, 4, 6, or 8 bits)
Reward: Quality-prioritized with asymmetric penalties: missing a quality target is costlier than exceeding a computational budget
Policy: Learned via off-policy Soft Actor-Critic (SAC) RL; outputs optimal bit-width for each layer given state
The state representation is model-agnostic, enabling zero-shot transfer. The reward function is shaped to handle the tradeoff between model quality (perplexity) and model size (bit budget).
Architecture Overview
- State Encoder: 11-dimensional embedding of activation/weight statistics (model-agnostic)
- Policy Network: Maps state -> bit-width distribution via SAC
- Value Networks: Q-functions for off-policy learning
- Reward Shaping: Multi-objective: minimize perplexity loss, respect bit budget
- Scale Folding: Novel preconditioning that stabilizes sub-4-bit quantization
- Zero-Shot Transfer: Learned policy applies to different model architectures
Implementation Steps
Step 1: Compute Model-Agnostic State Embedding
Extract 11-dimensional state features from model layers.
import torch
import torch.nn as nn
import numpy as np
def compute_layer_state_embedding(layer, activation_sample):
"""
Compute 11-dimensional state embedding for RL policy.
Captures layer characteristics in model-agnostic form.
layer: torch.nn.Module (e.g., Linear layer)
activation_sample: (batch_size, hidden_dim) tensor of activations through this layer
"""
state_features = []
# 1. Weight magnitude statistics (mean, std, max)
weight = layer.weight.data
state_features.append(weight.abs().mean().item())
state_features.append(weight.abs().std().item())
state_features.append(weight.abs().max().item())
# 2. Activation magnitude statistics
state_features.append(activation_sample.abs().mean().item())
state_features.append(activation_sample.abs().std().item())
state_features.append(activation_sample.abs().max().item())
# 3. Layer structure (param count, input/output dims, layer type)
param_count = sum(p.numel() for p in layer.parameters())
state_features.append(np.log10(max(param_count, 1))) # Log scale
# 4. Input-output dimension ratio
if hasattr(layer, 'in_features') and hasattr(layer, 'out_features'):
dim_ratio = layer.out_features / max(layer.in_features, 1)
state_features.append(dim_ratio)
else:
state_features.append(1.0)
# 5. Sparsity of activations
sparsity = (activation_sample.abs() < 1e-6).float().mean().item()
state_features.append(sparsity)
# 6. Dynamic range (max / min non-zero)
nonzero = activation_sample[activation_sample.abs() > 1e-8]
if len(nonzero) > 0:
dynamic_range = nonzero.abs().max() / (nonzero.abs().min() + 1e-8)
state_features.append(np.log10(dynamic_range))
else:
state_features.append(0.0)
# 7. Outlier ratio (percent beyond 2 std dev)
outlier_ratio = (activation_sample.abs() > 2 * activation_sample.std()).float().mean().item()
state_features.append(outlier_ratio)
return torch.tensor(state_features, dtype=torch.float32)
Step 2: Implement Scale Folding for Stability
Pre-condition the model to enable sub-4-bit quantization.
class ScaleFoldingPreconditioning:
"""
Scale folding: migrates activation outliers into weights via per-channel scaling.
Stabilizes sub-4-bit quantization by reducing dynamic range of activations.
"""
@staticmethod
def compute_per_channel_scales(activations):
"""
Compute per-channel scaling factors from activation statistics.
activations: (batch_size, hidden_dim)
Returns: (hidden_dim,) scaling factors
"""
# Per-channel scaling: scale = std / target_std
channel_stds = activations.std(dim=0, keepdim=True)
target_std = channel_stds.median() # Normalize to median
scales = target_std / (channel_stds + 1e-8)
return scales.squeeze()
@staticmethod
def apply_scale_folding(layer, scales):
"""
Apply scale folding by adjusting weights and scales.
layer: torch.nn.Module with weight and bias
scales: (out_features,) scaling factors
"""
# Absorb scales into weight matrix
# For Linear: weight_new = weight / scales
if hasattr(layer, 'weight'):
layer.weight.data = layer.weight.data / (scales.unsqueeze(1) + 1e-8)
# If layer has bias, adjust accordingly
if hasattr(layer, 'bias') and layer.bias is not None:
layer.bias.data = layer.bias.data / (scales + 1e-8)
return scales
@staticmethod
def fold_activations(x, scales):
"""
Fold scales into activations (multiply by scales).
x: (batch_size, hidden_dim) activations
scales: (hidden_dim,) pre-computed scales
Returns: scaled activations
"""
return x * scales.unsqueeze(0)
Step 3: Define RL Environment for Quantization
Set up the MDP for bit-width optimization.
class QuantizationEnvironment:
"""
RL environment for mixed-precision quantization.
State: layer embedding | Action: bit-width | Reward: perplexity - budget_penalty
"""
def __init__(self, model, calibration_data, target_bit_budget=8.0 * 1024):
self.model = model
self.calibration_data = calibration_data
self.target_bit_budget = target_bit_budget
self.num_layers = len(list(model.parameters()))
self.current_layer_idx = 0
self.bit_assignments = {}
self.layer_states = {}
# Pre-compute layer embeddings
self._compute_all_layer_states()
def _compute_all_layer_states(self):
"""Pre-compute state embeddings for all layers."""
for idx, (name, layer) in enumerate(self.model.named_modules()):
if isinstance(layer, torch.nn.Linear):
# Get a sample of activations
with torch.no_grad():
activations = self._get_layer_activations(layer)
state_embedding = compute_layer_state_embedding(layer, activations)
self.layer_states[idx] = state_embedding
def _get_layer_activations(self, target_layer):
"""Extract activations for a specific layer."""
activations = []
def hook_fn(module, input, output):
activations.append(output.detach())
handle = target_layer.register_forward_hook(hook_fn)
with torch.no_grad():
for batch in self.calibration_data[:10]: # Use first 10 batches
self.model(batch)
handle.remove()
return torch.cat(activations, dim=0)
def reset(self):
"""Start quantization from first layer."""
self.current_layer_idx = 0
self.bit_assignments = {}
return self.layer_states[0]
def step(self, bit_width_action):
"""
Assign bit-width to current layer, move to next.
bit_width_action: int (2, 4, 6, or 8)
Returns: (next_state, reward, done)
"""
self.bit_assignments[self.current_layer_idx] = bit_width_action
self.current_layer_idx += 1
# Get next state
if self.current_layer_idx >= self.num_layers:
done = True
next_state = None
else:
done = False
next_state = self.layer_states.get(self.current_layer_idx)
# Compute reward
reward = self._compute_reward(done)
return next_state, reward, done
def _compute_reward(self, done):
"""
Compute reward balancing quality and budget.
Quality (perplexity) is strongly prioritized over budget.
"""
if not done:
return 0.0 # No intermediate rewards
# Compute total bit budget used
total_bits = sum(
bits * self._get_layer_param_count(layer_idx)
for layer_idx, bits in self.bit_assignments.items()
)
# Compute quality loss (perplexity on calibration set)
perplexity = self._compute_perplexity_quantized()
# Reward: maximize quality (minimize perplexity), respect budget
quality_reward = -perplexity # Negative because we minimize
budget_penalty = max(0, (total_bits - self.target_bit_budget) / self.target_bit_budget) * 100
# Quality strongly prioritized (10x weight on quality vs budget)
reward = quality_reward - 0.1 * budget_penalty
return reward
def _get_layer_param_count(self, layer_idx):
"""Get number of parameters in a layer."""
layer = list(self.model.modules())[layer_idx]
return sum(p.numel() for p in layer.parameters()) / 1e6 # In millions
def _compute_perplexity_quantized(self):
"""Evaluate perplexity of quantized model on calibration set."""
total_loss = 0.0
total_tokens = 0
with torch.no_grad():
for batch_input, batch_target in self.calibration_data:
logits = self.model(batch_input)
loss = nn.CrossEntropyLoss()(logits.view(-1, logits.size(-1)), batch_target.view(-1))
total_loss += loss.item() * batch_target.numel()
total_tokens += batch_target.numel()
perplexity = torch.exp(torch.tensor(total_loss / total_tokens))
return perplexity.item()
Step 4: Train RL Policy with Soft Actor-Critic
Use off-policy RL to learn optimal bit-width assignments.
from torch.distributions import Categorical
class SAC_QuantizationPolicy:
"""
Soft Actor-Critic for learning bit-width assignments.
Off-policy allows zero-shot transfer across models.
"""
def __init__(self, state_dim=11, num_actions=4, learning_rate=1e-4):
self.state_dim = state_dim
self.num_actions = num_actions # 4 actions: 2, 4, 6, 8 bits
# Actor network: state -> action probabilities
self.actor = nn.Sequential(
nn.Linear(state_dim, 128),
nn.ReLU(),
nn.Linear(128, num_actions)
)
# Q-networks for value estimation
self.q1 = nn.Sequential(
nn.Linear(state_dim + num_actions, 128),
nn.ReLU(),
nn.Linear(128, 1)
)
self.q2 = nn.Sequential(
nn.Linear(state_dim + num_actions, 128),
nn.ReLU(),
nn.Linear(128, 1)
)
self.optimizer_actor = torch.optim.Adam(self.actor.parameters(), lr=learning_rate)
self.optimizer_q = torch.optim.Adam(
list(self.q1.parameters()) + list(self.q2.parameters()),
lr=learning_rate
)
def select_action(self, state, temperature=1.0):
"""Select action with exploration."""
logits = self.actor(state)
probs = torch.softmax(logits / temperature, dim=-1)
dist = Categorical(probs)
action = dist.sample()
return action.item(), probs
def train_step(self, batch_states, batch_actions, batch_rewards, batch_next_states, batch_done):
"""Single training step with SAC loss."""
# Q-function loss
state_action = torch.cat([batch_states, batch_actions], dim=-1)
q1_pred = self.q1(state_action)
q2_pred = self.q2(state_action)
target_q = batch_rewards + 0.99 * (1 - batch_done.float()) * torch.min(
self.q1(torch.cat([batch_next_states, self._greedy_action(batch_next_states)], dim=-1)),
self.q2(torch.cat([batch_next_states, self._greedy_action(batch_next_states)], dim=-1))
)
q_loss = nn.MSELoss()(q1_pred, target_q) + nn.MSELoss()(q2_pred, target_q)
self.optimizer_q.zero_grad()
q_loss.backward()
self.optimizer_q.step()
# Actor loss (entropy regularized)
logits = self.actor(batch_states)
entropy = -(torch.softmax(logits, dim=-1) * torch.log_softmax(logits, dim=-1)).sum(dim=-1).mean()
actor_loss = -entropy # Maximize entropy
self.optimizer_actor.zero_grad()
actor_loss.backward()
self.optimizer_actor.step()
return q_loss.item(), actor_loss.item()
def _greedy_action(self, states):
"""Select greedy action (max probability)."""
logits = self.actor(states)
actions = torch.argmax(logits, dim=-1, keepdim=True)
return torch.nn.functional.one_hot(actions.squeeze(-1), num_classes=self.num_actions).float()
def get_bit_assignments(self, layer_states):
"""Get optimal bit-widths for all layers (greedy policy evaluation)."""
bit_map = {0: 2, 1: 4, 2: 6, 3: 8}
assignments = {}
with torch.no_grad():
for layer_idx, state in layer_states.items():
action_idx = torch.argmax(self.actor(state.unsqueeze(0))).item()
assignments[layer_idx] = bit_map[action_idx]
return assignments
Step 5: Zero-Shot Transfer Evaluation
Validate policy generalizes to new models without retraining.
def evaluate_policy_zero_shot(policy, source_model, target_models):
"""
Evaluate if policy trained on source_model transfers to target_models.
"""
results = {}
for target_name, target_model in target_models.items():
# Compute layer states for target model
target_states = {}
for idx, layer in enumerate(target_model.modules()):
if isinstance(layer, torch.nn.Linear):
activations = get_target_activations(target_model)
state = compute_layer_state_embedding(layer, activations)
target_states[idx] = state
# Use policy without retraining
bit_assignments = policy.get_bit_assignments(target_states)
# Evaluate perplexity with assigned bit-widths
perplexity = evaluate_quantized_model(target_model, bit_assignments)
results[target_name] = {
'bit_assignments': bit_assignments,
'perplexity': perplexity
}
return results
Practical Guidance
Hyperparameters:
- State dimension: 11 (fixed, model-agnostic)
- Bit-width options: [2, 4, 6, 8] (common for on-device models)
- Target bit budget: scale with model size (e.g., 8 bits * model params in billions)
- SAC temperature: 1.0 (exploration) during training, 0.0 (greedy) during evaluation
- Scale folding scaling factor: 1.0-2.0 (controls how aggressively to fold)
When to Use:
- On-device LLM deployment (mobile, edge devices)
- Need for mixed-precision quantization with dynamic per-layer decisions
- Have calibration data and want to optimize bit-width allocation
- Planning to deploy to multiple model architectures (zero-shot transfer)
When NOT to Use:
- Server-side inference where 16-bit or FP32 is standard
- Models with highly irregular architectures (state embedding may not generalize)
- Very small models where quantization overhead dominates
- Real-time quantization (training the policy takes time upfront)
Pitfalls:
- Scale folding can be unstable: monitor activation ranges after folding
- Calibration data must be representative; poor calibration data leads to poor bit assignments
- Policy overfitting: if training on single model size, may not transfer; train on diverse sizes
- Reward shaping critical: asymmetric penalties prevent policy from overshooting budget; tune carefully
Reference
Paper: arxiv.org/abs/2603.17891