from typing import Any, Optional, Callable, List, Tuple
import os
import time
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from accelerate import init_empty_weights
from transformers.activations import ACT2FN
from transformers.generation import GenerationConfig
from transformers.models.opt.modeling_opt import (
OPTAttention,
OPTDecoder,
OPTDecoderLayer,
OPTForCausalLM,
OPTModel,
)
from transformers.models.opt.configuration_opt import OPTConfig
from huggingface_hub import snapshot_download
from configuration_tricksy import TricksyConfig
from util import batch_copy, compute_index_diffs, load_mlp_sparsity_predictor, mmap_to_tensor, topk_and_threshold
TRICKSY_WEIGHTS_PATH = 'tricksy-weights/'
class SparseMLPCache:
def __init__(
self,
indexed_fc1_weight: Optional[torch.Tensor] = None,
indexed_fc1_bias: Optional[torch.Tensor] = None,
indexed_fc2_weight: Optional[torch.Tensor] = None,
gpu_cached_mlp_indices: Optional[torch.Tensor] = None,
):
# [ffn_embed_dim * min_mlp_sparsity, hidden_size]
self.indexed_fc1_weight = indexed_fc1_weight
# [ffn_embed_dim * min_mlp_sparsity]
self.indexed_fc1_bias = indexed_fc1_bias
# [ffn_embed_dim * min_mlp_sparsity, hidden_size] (stored in transpose for efficient indexing)
self.indexed_fc2_weight = indexed_fc2_weight
# Indices that are already on GPU (this tensor is stored on the CPU)
# [ffn_embed_dim * min_mlp_sparsity]
self.gpu_cached_mlp_indices = gpu_cached_mlp_indices
class SparseIndices:
def __init__(self, tricksy_config: TricksyConfig, opt_config: OPTConfig):
self.mlp_indices_buffer_gpu = torch.empty(
(int(opt_config.ffn_dim * tricksy_config.min_mlp_sparsity_gpu),),
dtype=torch.int32,
device='cuda'
)
self.mlp_indices_buffer_cpu = torch.empty(
(int(opt_config.ffn_dim * tricksy_config.min_mlp_sparsity_gpu),),
dtype=torch.int32,
device='cpu',
pin_memory=True,
)
# Default stream blocks until indices are copied to CPU
self.index_copy_stream = torch.cuda.default_stream()
def copy_mlp_indices_to_cpu(self):
self.mlp_indices_buffer_cpu = batch_copy([self.mlp_indices_buffer_gpu], self.index_copy_stream, device='cpu')[0]
class OPTDiskWeights:
def __init__(self, model_name: str):
self.model_name = model_name
self.model_suffix = model_name.split('/')[-1]
self.config = OPTConfig.from_pretrained(model_name)
try:
print(f'downloading from austinsilveria/tricksy-{self.model_suffix}')
self.weight_path = snapshot_download(repo_id=f'austinsilveria/tricksy-{self.model_suffix}') + '/'
except:
print(f'failed to download from austinsilveria/tricksy-{self.model_suffix}')
self.weight_path = f'{TRICKSY_WEIGHTS_PATH}{self.model_suffix}/'
with init_empty_weights():
model = OPTModel(self.config)
self.state_dict = model.state_dict()
if not os.path.exists(f'{self.weight_path}decoder.embed_tokens.weight'):
# Download original weights and write memmap files
print(f'downloading and preprocessing original weights')
self.cache_weights()
head_dim = self.config.hidden_size // self.config.num_attention_heads
for i in range(self.config.num_hidden_layers):
layer_prefix = f'decoder.layers.{i}.'
self.delete_weights([
f'{layer_prefix}self_attn.q_proj.weight',
f'{layer_prefix}self_attn.k_proj.weight',
f'{layer_prefix}self_attn.v_proj.weight',
f'{layer_prefix}self_attn.out_proj.weight',
f'{layer_prefix}self_attn.q_proj.bias',
f'{layer_prefix}self_attn.k_proj.bias',
f'{layer_prefix}self_attn.v_proj.bias'
])
self.add_weights([
(f'{layer_prefix}fc2.weight', (self.config.ffn_dim, self.config.hidden_size)),
(f'{layer_prefix}self_attn.catted_head_weights', (self.config.num_attention_heads, head_dim * 4, self.config.hidden_size)),
(f'{layer_prefix}self_attn.catted_head_biases', (self.config.num_attention_heads, 3, head_dim)),
])
self.memmap_weights = { key: self.load_memmap_weight(key) for key in self.state_dict.keys() }
def load_memmap_weight(self, key: str):
return torch.from_numpy(np.memmap(f'{self.weight_path}{key}', dtype='float16', mode='r', shape=(self.state_dict[key].shape)))
def add_weights(self, weights: List[Tuple[str, torch.Size]]):
for key, shape in weights:
self.state_dict[key] = torch.empty(shape, dtype=torch.float16, device='meta')
def delete_weights(self, keys: List[str]):
for key in keys:
if key in self.state_dict:
del self.state_dict[key]
path = f'{self.weight_path}{key}'
if os.path.exists(path):
os.remove(path)
def cache_weights(self):
os.makedirs(self.weight_path, exist_ok=True)
weights_location = snapshot_download(repo_id=self.model_name, ignore_patterns=['flax*', 'tf*'])
shards = [file for file in os.listdir(weights_location) if file.startswith("pytorch_model") and file.endswith(".bin")]
for shard in shards:
print(f'caching {shard}')
shard_path = os.path.join(weights_location, shard)
shard_state_dict = torch.load(shard_path)
for key in shard_state_dict.keys():
path = f'{self.weight_path}{key.replace("model.", "")}'
memmap = np.memmap(path, dtype='float16', mode='w+', shape=(shard_state_dict[key].shape))
memmap[:] = shard_state_dict[key].cpu().numpy()
# Store weights in shape for efficient indexing
for i in range(self.config.num_hidden_layers):
layer_prefix = f'decoder.layers.{i}.'
# FC2 in transpose
fc2t = torch.from_numpy(np.array(self.load_memmap_weight(f'{layer_prefix}fc2.weight')[:])).t().contiguous().clone()
np.memmap(f'{self.weight_path}decoder.layers.{i}.fc2.weight', dtype='float16', mode='w+', shape=fc2t.shape)[:] = fc2t.numpy()
# Attention weights by head
head_dim = self.config.hidden_size // self.config.num_attention_heads
qw = mmap_to_tensor(self.load_memmap_weight(f'{layer_prefix}self_attn.q_proj.weight')[:])
kw = mmap_to_tensor(self.load_memmap_weight(f'{layer_prefix}self_attn.k_proj.weight')[:])
vw = mmap_to_tensor(self.load_memmap_weight(f'{layer_prefix}self_attn.v_proj.weight')[:])
ow = mmap_to_tensor(self.load_memmap_weight(f'{layer_prefix}self_attn.out_proj.weight')[:])
pre_cat_shape = (self.config.num_attention_heads, head_dim, self.config.hidden_size)
# [head, head_dim * 4, hidden_size]
catted_head_weights = torch.cat(
[qw.view(pre_cat_shape).clone(), kw.view(pre_cat_shape).clone(), vw.view(pre_cat_shape).clone(), ow.T.view(pre_cat_shape).clone(),],
dim=1,
).contiguous().clone()
np.memmap(f'{self.weight_path}{layer_prefix}self_attn.catted_head_weights', dtype='float16', mode='w+', shape=catted_head_weights.shape)[:] =\
catted_head_weights.numpy()
# Attention biases by head
qb = mmap_to_tensor(self.load_memmap_weight(f'{layer_prefix}self_attn.q_proj.bias')[:])
kb = mmap_to_tensor(self.load_memmap_weight(f'{layer_prefix}self_attn.k_proj.bias')[:])
vb = mmap_to_tensor(self.load_memmap_weight(f'{layer_prefix}self_attn.v_proj.bias')[:])
pre_cat_shape = (self.config.num_attention_heads, 1, head_dim)
# [head, 3, head_dim]
catted_head_biases = torch.cat(
# Don't index out bias since we need all dims after projecting back up to hidden size
[qb.view(pre_cat_shape).clone(), kb.view(pre_cat_shape).clone(), vb.view(pre_cat_shape).clone()],
dim=1,
).contiguous().clone()
np.memmap(f'{self.weight_path}{layer_prefix}self_attn.catted_head_biases', dtype='float16', mode='w+', shape=catted_head_biases.shape)[:] =\
catted_head_biases.numpy()
self.delete_weights([
f'{layer_prefix}self_attn.q_proj.weight',
f'{layer_prefix}self_attn.k_proj.weight',
f'{layer_prefix}self_attn.v_proj.weight',
f'{layer_prefix}self_attn.out_proj.weight',
f'{layer_prefix}self_attn.q_proj.bias',
f'{layer_prefix}self_attn.k_proj.bias',
f'{layer_prefix}self_attn.v_proj.bias'
])
self.add_weights([
(f'{layer_prefix}self_attn.catted_head_weights', catted_head_weights.shape),
(f'{layer_prefix}self_attn.catted_head_biases', catted_head_biases.shape),
])
class TricksyContext:
def __init__(self, tricksy_config: TricksyConfig, opt_config: OPTConfig):
self.indices = SparseIndices(tricksy_config, opt_config)
self.load_weight_stream = torch.cuda.Stream()
self.layer = 0
self.is_prompt_phase = True
self.forward_times = []
class TricksyLayer:
def __call__(self, *args: Any, **kwds: Any) -> Any:
return self.forward(*args, **kwds)
def load_weights(self, tricksy_context: TricksyContext):
pass
class TricksyLayerInputs:
def __init__(
self,
disk_weights: OPTDiskWeights,
layer_key_prefix: str = None,
next_layer: TricksyLayer = None,
sparsity_predictors: List[Callable[[torch.Tensor], torch.Tensor]] = None,
) -> None:
self.disk_weights = disk_weights
# self.get_weight = lambda key: self.disk_weights.load_memmap_weight(f'{layer_key_prefix}{key}')
self.get_weight = lambda key: self.disk_weights.memmap_weights[(f'{layer_key_prefix}{key}')]
self.layer_key_prefix = layer_key_prefix
self.next_layer = next_layer
self.sparsity_predictors = sparsity_predictors
class TricksyOPTLearnedPositionalEmbedding(TricksyLayer):
"""
This module learns positional embeddings up to a fixed maximum size.
"""
def __init__(self, tricksy_context):
# OPT is set up so that if padding_idx is specified then offset the embedding ids by 2
# and adjust num_embeddings appropriately. Other models don't have this hack
self.offset = 2
self.tricksy_context = tricksy_context
self.weight = None
def __call__(self, *args: Any, **kwds: Any) -> Any:
return self.forward(*args, **kwds)
def forward(self, attention_mask: torch.LongTensor, past_key_values_length: int = 0):
"""`input_ids_shape` is expected to be [bsz x seqlen]."""
attention_mask = attention_mask.long()
# create positions depending on attention_mask
positions = (torch.cumsum(attention_mask, dim=1).type_as(attention_mask) * attention_mask).long() - 1
# cut positions if `past_key_values_length` is > 0
positions = positions[:, past_key_values_length:]
out = F.embedding(positions + self.offset, self.weight)
return out
class TricksyOPTAttention(OPTAttention, TricksyLayer):
def __init__(self, tricksy_config: TricksyConfig, inputs: TricksyLayerInputs, tricksy_context: TricksyContext, is_decoder: bool = False, **kwargs):
nn.Module.__init__(self)
self.tricksy_config = tricksy_config
self.config = tricksy_config.opt_config
def _handle_deprecated_argument(config_arg_name, config, fn_arg_name, kwargs):
"""
If a the deprecated argument `fn_arg_name` is passed, raise a deprecation
warning and return that value, otherwise take the equivalent config.config_arg_name
"""
val = None
if fn_arg_name in kwargs:
print(
"Passing in {} to {self.__class__.__name__} is deprecated and won't be supported from v4.38."
" Please set it in the config instead"
)
val = kwargs.pop(fn_arg_name)
else:
val = getattr(config, config_arg_name)
return val
self.embed_dim = _handle_deprecated_argument("hidden_size", self.config, "embed_dim", kwargs)
self.num_heads = _handle_deprecated_argument("num_attention_heads", self.config, "num_heads", kwargs)
self.dropout = _handle_deprecated_argument("attention_dropout", self.config, "dropout", kwargs)
self.enable_bias = _handle_deprecated_argument("enable_bias", self.config, "bias", kwargs)
self.head_dim = self.embed_dim // self.num_heads
self.is_causal = True
if (self.head_dim * self.num_heads) != self.embed_dim:
raise ValueError(
f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
f" and `num_heads`: {self.num_heads})."
)
self.scaling = self.head_dim**-0.5
self.is_decoder = is_decoder
# [Tricksy]
self.tricksy_context = tricksy_context
self.inputs = inputs
self.head_dim = self.config.hidden_size // self.config.num_attention_heads
self.qw = self.kw = self.vw = self.ow = self.qb = self.kb = self.vb = self.out_proj_bias = self.layer_norm_weight = self.layer_norm_bias = None
self.q_proj = lambda x: F.linear(x, self.qw, self.qb)
self.k_proj = lambda x: F.linear(x, self.kw, self.kb)
self.v_proj = lambda x: F.linear(x, self.vw, self.vb)
self.out_proj = lambda x: F.linear(x, self.ow, self.out_proj_bias)
self.layer_norm = lambda x: F.layer_norm(x, (self.config.hidden_size,), self.layer_norm_weight, self.layer_norm_bias)
def clear(self):
self.qw = self.kw = self.vw = self.ow = self.qb = self.kb = self.vb = self.out_proj_bias = self.layer_norm_weight = self.layer_norm_bias = None
def load_weights(self, tricksy_context: TricksyContext):
if self.tricksy_context.is_prompt_phase:
# Full weights for prompt phase
self.catted_weights, self.catted_biases, self.out_proj_bias, self.layer_norm_weight, self.layer_norm_bias = batch_copy(
[
mmap_to_tensor(self.inputs.get_weight('self_attn.catted_head_weights')[:], pin_memory=True),
mmap_to_tensor(self.inputs.get_weight('self_attn.catted_head_biases')[:], pin_memory=True),
mmap_to_tensor(self.inputs.get_weight('self_attn.out_proj.bias')[:], pin_memory=True),
mmap_to_tensor(self.inputs.get_weight('self_attn_layer_norm.weight')[:], pin_memory=True),
mmap_to_tensor(self.inputs.get_weight('self_attn_layer_norm.bias')[:], pin_memory=True),
],
tricksy_context.load_weight_stream,
)
torch.cuda.synchronize()
# Weights stored in shape for efficient indexing to support offloading attention heads (not currently being done)
self.qw = self.catted_weights[:, :self.head_dim, :].reshape(self.config.hidden_size, self.config.hidden_size).contiguous()
self.kw = self.catted_weights[:, self.head_dim:self.head_dim * 2, :].reshape(self.config.hidden_size, self.config.hidden_size).contiguous()
self.vw = self.catted_weights[:, self.head_dim * 2:self.head_dim * 3, :].reshape(self.config.hidden_size, self.config.hidden_size).contiguous()
self.ow = self.catted_weights[:, self.head_dim * 3:, :].reshape(self.config.hidden_size, self.config.hidden_size).t().contiguous()
self.catted_weights = None
self.qb = self.catted_biases[:, 0, :].reshape(self.config.hidden_size).contiguous()
self.kb = self.catted_biases[:, 1, :].reshape(self.config.hidden_size).contiguous()
self.vb = self.catted_biases[:, 2, :].reshape(self.config.hidden_size).contiguous()
self.catted_biases = None
def forward(self, hidden_states, **kwargs):
# Wait for attention weights to get to GPU
torch.cuda.synchronize()
# Predict MLP sparsity based on attention input
self.tricksy_context.indices.mlp_indices_buffer_gpu = topk_and_threshold(
self.inputs.sparsity_predictors[0](hidden_states)[0, -1, :],
int(self.config.ffn_dim * self.tricksy_config.min_mlp_sparsity_gpu),
)
self.tricksy_context.indices.copy_mlp_indices_to_cpu()
torch.cuda.synchronize()
# Load MLP weights while computing attention
self.inputs.next_layer.load_weights(self.tricksy_context)
out = super().forward(self.layer_norm(hidden_states), **kwargs)
# Wait for MLP weights to get to GPU
torch.cuda.synchronize()
return out
class TricksyOPTDecoderLayer(OPTDecoderLayer):
def __init__(self, tricksy_config: TricksyConfig, inputs: TricksyLayerInputs, tricksy_context: TricksyContext):
nn.Module.__init__(self)
self.tricksy_config = tricksy_config
self.config = tricksy_config.opt_config
self.embed_dim = self.config.hidden_size
self.tricksy_context = tricksy_context
self.self_attn_layer_inputs = TricksyLayerInputs(
disk_weights=inputs.disk_weights,
layer_key_prefix=inputs.layer_key_prefix,
# While computing attention, load MLP
next_layer=self,
sparsity_predictors=inputs.sparsity_predictors,
)
self.self_attn = TricksyOPTAttention(tricksy_config, self.self_attn_layer_inputs, tricksy_context, is_decoder=True)
self.do_layer_norm_before = self.config.do_layer_norm_before
self.dropout = self.config.dropout
self.activation_fn = ACT2FN[self.config.activation_function]
self.inputs = inputs
random_mlp_indices_gpu =\
torch.randperm(self.config.ffn_dim, device='cpu', dtype=torch.int32)[:int(self.config.ffn_dim * self.tricksy_config.min_mlp_sparsity_gpu)]
self.index_cache = SparseMLPCache(gpu_cached_mlp_indices=random_mlp_indices_gpu)
# identity since we move this to attention layer
# extreme tricksy
self.self_attn_layer_norm = lambda x: x
self.fc1_weight = self.fc2_weight = self.final_layer_norm_weight = self.fc1_bias = self.fc2_bias = self.final_layer_norm_bias = None
self.ring_idx = 0
self.fc1_weight_diff = self.fc2_weight_diff = self.fc1_bias_diff = None
self.fc1 = lambda x: F.linear(x, torch.cat([self.fc1_weight, self.fc1_weight_diff]), torch.cat([self.fc1_bias, self.fc1_bias_diff]))
self.fc2 = lambda x: F.linear(x, torch.cat([self.fc2_weight, self.fc2_weight_diff]).T, self.fc2_bias)
self.final_layer_norm = lambda x: F.layer_norm(x, (self.embed_dim,), self.final_layer_norm_weight, self.final_layer_norm_bias)
def clear(self):
self.fc1_weight = self.fc2_weight = self.final_layer_norm_weight = self.fc1_bias = self.fc2_bias = self.final_layer_norm_bias = None
self.fc1_weight_diff = self.fc2_weight_diff = self.fc1_bias_diff = None
def load_weights(self, tricksy_context: TricksyContext):
if self.tricksy_context.is_prompt_phase:
# Full weights for prompt phase
fc1w = mmap_to_tensor(self.inputs.get_weight('fc1.weight')[:], pin_memory=True)
fc1b = mmap_to_tensor(self.inputs.get_weight('fc1.bias')[:], pin_memory=True)
fc2w = mmap_to_tensor(self.inputs.get_weight('fc2.weight')[:], pin_memory=True)
fc2b = mmap_to_tensor(self.inputs.get_weight('fc2.bias')[:], pin_memory=True)
lnw = mmap_to_tensor(self.inputs.get_weight('final_layer_norm.weight')[:], pin_memory=True)
lnb = mmap_to_tensor(self.inputs.get_weight('final_layer_norm.bias')[:], pin_memory=True)
self.fc1_weight, self.fc1_bias, self.fc2_weight, self.fc2_bias, self.final_layer_norm_weight, self.final_layer_norm_bias =\
batch_copy([fc1w, fc1b, fc2w, fc2b, lnw, lnb], tricksy_context.load_weight_stream)
self.fc1_weight_diff = torch.tensor([], dtype=self.tricksy_config.dtype, device='cuda')
self.fc1_bias_diff = torch.tensor([], dtype=self.tricksy_config.dtype, device='cuda')
self.fc2_weight_diff = torch.tensor([], dtype=self.tricksy_config.dtype, device='cuda')
index_diffs = compute_index_diffs(tricksy_context.indices.mlp_indices_buffer_cpu, [self.index_cache.gpu_cached_mlp_indices])
if len(index_diffs) > 0:
gpu_index_diff = index_diffs[0]
self.index_cache.gpu_cached_mlp_indices[gpu_index_diff.off_positions] = gpu_index_diff.off_elements
self.index_cache.indexed_fc1_weight = fc1w.contiguous().pin_memory()
self.index_cache.indexed_fc1_bias = fc1b.contiguous().pin_memory()
self.index_cache.indexed_fc2_weight = fc2w.contiguous().pin_memory()
return
elif self.fc1_weight is None:
# Full weights if full offload
self.fc1_weight, self.fc1_bias, self.fc2_weight = batch_copy(
[self.index_cache.indexed_fc1_weight, self.index_cache.indexed_fc1_bias, self.index_cache.indexed_fc2_weight],
tricksy_context.load_weight_stream
)
self.fc1_weight_diff = torch.tensor([], dtype=self.tricksy_config.dtype, device='cuda')
self.fc1_bias_diff = torch.tensor([], dtype=self.tricksy_config.dtype, device='cuda')
self.fc2_weight_diff = torch.tensor([], dtype=self.tricksy_config.dtype, device='cuda')
off_elements = torch.tensor(
list(set(tricksy_context.indices.mlp_indices_buffer_cpu.tolist()).difference(set(self.index_cache.gpu_cached_mlp_indices.tolist()))),
device='cpu',
dtype=torch.int32,
pin_memory=True
)
if off_elements.size(0) == 0:
self.fc1_weight_diff = torch.tensor([], dtype=self.tricksy_config.dtype, device='cuda')
self.fc1_bias_diff = torch.tensor([], dtype=self.tricksy_config.dtype, device='cuda')
self.fc2_weight_diff = torch.tensor([], dtype=self.tricksy_config.dtype, device='cuda')
return
new_ring_idx = (self.ring_idx + off_elements.size(0)) % self.index_cache.gpu_cached_mlp_indices.size(0)
if new_ring_idx > self.ring_idx:
# single contiguous update
self.index_cache.gpu_cached_mlp_indices[self.ring_idx:new_ring_idx] = off_elements
elif off_elements.size(0) > 0:
split = self.index_cache.gpu_cached_mlp_indices.size(0) - self.ring_idx
# end of ring
self.index_cache.gpu_cached_mlp_indices[self.ring_idx:] = off_elements[:split]
# beginning of ring
self.index_cache.gpu_cached_mlp_indices[:new_ring_idx] = off_elements[split:]
# Allocate
self.fc1_weight_diff = torch.empty((off_elements.size(0), self.config.hidden_size), dtype=self.tricksy_config.dtype, device='cuda')
self.fc1_bias_diff = torch.empty((off_elements.size(0)), dtype=self.tricksy_config.dtype, device='cuda')
self.fc2_weight_diff = torch.empty((off_elements.size(0), self.config.hidden_size), dtype=self.tricksy_config.dtype, device='cuda')
# Index
fc1wd = self.index_cache.indexed_fc1_weight[off_elements].pin_memory()
fc1bd = self.index_cache.indexed_fc1_bias[off_elements].pin_memory()
fc2wd = self.index_cache.indexed_fc2_weight[off_elements].pin_memory()
# Copy
self.fc1_weight_diff, self.fc1_bias_diff, self.fc2_weight_diff = batch_copy([fc1wd, fc1bd, fc2wd], tricksy_context.load_weight_stream)
def forward(self, *args, **kwargs):
# Wait for attention weights to get to GPU
torch.cuda.synchronize()
# Load next layer's attention weights
self.inputs.next_layer.load_weights(self.tricksy_context)
out = super().forward(*args, **kwargs)
if self.tricksy_config.full_offload:
self.fc1_weight = self.fc1_bias = self.fc2_weight = None
elif self.tricksy_context.is_prompt_phase:
# Only keep sparse MLP weights on GPU after prompt phase
self.fc1_weight = self.fc1_weight[self.index_cache.gpu_cached_mlp_indices.to('cuda')]
self.fc1_bias = self.fc1_bias[self.index_cache.gpu_cached_mlp_indices.to('cuda')]
self.fc2_weight = self.fc2_weight[self.index_cache.gpu_cached_mlp_indices.to('cuda')]
# Update ring buffers
if not self.tricksy_config.full_offload:
prev_ring_idx = self.ring_idx
self.ring_idx = (self.ring_idx + self.fc1_weight_diff.size(0)) % self.fc1_weight.size(0)
if self.ring_idx > prev_ring_idx:
# does not wrap around ring
self.fc1_weight[prev_ring_idx:self.ring_idx] = self.fc1_weight_diff
self.fc1_bias[prev_ring_idx:self.ring_idx] = self.fc1_bias_diff
self.fc2_weight[prev_ring_idx:self.ring_idx] = self.fc2_weight_diff
elif self.fc1_weight_diff.size(0) > 0:
# wraps around ring
split = self.fc1_weight_diff.size(0) - self.ring_idx
self.fc1_weight[prev_ring_idx:] = self.fc1_weight_diff[:split]
self.fc1_weight[:self.ring_idx] = self.fc1_weight_diff[split:]
self.fc1_bias[prev_ring_idx:] = self.fc1_bias_diff[:split]
self.fc1_bias[:self.ring_idx] = self.fc1_bias_diff[split:]
self.fc2_weight[prev_ring_idx:] = self.fc2_weight_diff[:split]
self.fc2_weight[:self.ring_idx] = self.fc2_weight_diff[split:]
self.fc1_weight_diff = self.fc2_weight_diff = self.fc1_bias_diff = None
self.tricksy_context.layer += 1
return out
class TricksyOPTDecoder(OPTDecoder, TricksyLayer):
def __init__(self, tricksy_config: TricksyConfig, disk_weights: OPTDiskWeights, tricksy_opt_for_causal_lm, tricksy_context: TricksyContext):
nn.Module.__init__(self)
self.config = tricksy_config.opt_config
self.dropout = self.config.dropout
self.layerdrop = self.config.layerdrop
self.padding_idx = self.config.pad_token_id
self.max_target_positions = self.config.max_position_embeddings
self.vocab_size = self.config.vocab_size
self._use_flash_attention_2 = False
self.gradient_checkpointing = False
self.project_out = None
self.project_in = None
self.embed_tokens_weight = None
self.embed_positions = TricksyOPTLearnedPositionalEmbedding(tricksy_context)
self.tricksy_context = tricksy_context
self.layers: List[TricksyOPTDecoderLayer] = []
for i in range(self.config.num_hidden_layers):
pretrained_layer_num = self.config.num_hidden_layers - i - 1
sparsity_predictors = [load_mlp_sparsity_predictor(disk_weights.weight_path, pretrained_layer_num, tricksy_config.dtype)]
if sparsity_predictors[0] is None:
sparsity_predictors[0] = lambda x: F.linear(x, torch.rand((self.config.ffn_dim, self.config.hidden_size), device='cuda', dtype=tricksy_config.dtype))
self.layers.append(TricksyOPTDecoderLayer(
tricksy_config,
TricksyLayerInputs(
disk_weights=disk_weights,
layer_key_prefix=f'decoder.layers.{pretrained_layer_num}.',
# While computing MLP, load next attention
# While computing last MLP, load output embeddings (stored in TricksyOPTForCausalLM)
next_layer=self.layers[i - 1].self_attn if i > 0 else tricksy_opt_for_causal_lm,
sparsity_predictors=sparsity_predictors,
),
tricksy_context,
))
self.layers.reverse()
self.final_layer_norm = lambda x: x
self.inputs = TricksyLayerInputs(disk_weights=disk_weights, layer_key_prefix='decoder.')
def clear(self):
self.embed_tokens_weight = self.embed_positions.weight = None
for layer in self.layers:
layer.clear()
def embed_tokens(self, x):
return F.embedding(x, self.embed_tokens_weight, self.padding_idx)
def load_weights(self, tricksy_context: TricksyContext):
if self.embed_tokens_weight is None:
self.embed_tokens_weight, self.embed_positions.weight = batch_copy(
[
mmap_to_tensor(self.inputs.get_weight('embed_tokens.weight')[:], pin_memory=True),
mmap_to_tensor(self.inputs.get_weight('embed_positions.weight')[:], pin_memory=True),
],
tricksy_context.load_weight_stream,
)
def forward(self, *args, **kwargs):
# Wait for input embedding weights to get to GPU
torch.cuda.synchronize()
# While computing input embeddings, load first attention
self.layers[0].self_attn.load_weights(self.tricksy_context)
out = super().forward(*args, **kwargs)
# Wait for output embedding weights to get to GPU
torch.cuda.synchronize()
# No longer prompt phase after first full pass
self.tricksy_context.is_prompt_phase = False
# Load input embeddings while computing output
self.load_weights(self.tricksy_context)
return out
class TricksyOPTModel(OPTModel):
def __init__(self, tricksy_config: TricksyConfig, disk_weights: OPTDiskWeights, tricksy_opt_for_causal_lm, tricksy_context: TricksyContext):
nn.Module.__init__(self)
self.config = tricksy_config.opt_config
self.tricksy_context = tricksy_context
self.decoder = TricksyOPTDecoder(tricksy_config, disk_weights, tricksy_opt_for_causal_lm, tricksy_context)
def clear(self):
self.decoder.clear()
def forward(self, *args, **kwargs):
out = super().forward(*args, **kwargs)
return out
# who's got the weights?
# [InputEmbedding, Attention.0, MLP.0, Attention.1, MLP.1, ..., OutputEmbedding]
# [TricksyOPTDecoder, TricksyOPTAttention.0, TricksyOPTDecoderLayer.0, TricksyOPTAttention.1, TricksyDecoderLayer.1, ..., TricksyOPTForCausalLM]
#
# 1. Prompt pass: Before computing layer, send full dense weights to GPU. After computing layer, only keep sparse weights on GPU.
# 2. Generation passes: Before computing layer, compute and send sparse weight diff to GPU.
class TricksyOPTForCausalLM(OPTForCausalLM, TricksyLayer):
def __init__(self, tricksy_config: TricksyConfig, disk_weights: OPTDiskWeights):
nn.Module.__init__(self)
self.config = disk_weights.config
self.generation_config = GenerationConfig.from_model_config(self.config) if self.can_generate() else None
self.tricksy_context = TricksyContext(tricksy_config, self.config)
self.model = TricksyOPTModel(tricksy_config, disk_weights, self, self.tricksy_context)
self.final_layer_norm_weight = self.lm_head_weight = self.final_layer_norm_bias = None
# double stacking tricksy!
self.final_layer_norm = lambda x: F.layer_norm(x, (self.config.hidden_size,), self.final_layer_norm_weight, self.final_layer_norm_bias)
self.lm_head = lambda x: F.linear(self.final_layer_norm(x), self.lm_head_weight)
self.inputs = TricksyLayerInputs(disk_weights=disk_weights, layer_key_prefix='decoder.', next_layer=self.model.decoder)
def clear(self):
self.model.clear()
def load_weights(self, tricksy_context: TricksyContext):
if self.final_layer_norm_weight is None:
self.final_layer_norm_weight, self.lm_head_weight, self.final_layer_norm_bias = batch_copy(
[
mmap_to_tensor(self.inputs.get_weight('final_layer_norm.weight')[:], pin_memory=True),
mmap_to_tensor(self.inputs.get_weight('embed_tokens.weight')[:], pin_memory=True),
mmap_to_tensor(self.inputs.get_weight('final_layer_norm.bias')[:], pin_memory=True),
],
tricksy_context.load_weight_stream,
)
def forward(self, *args, **kwargs):
torch.cuda.synchronize()
start = time.time()
out = super().forward(*args, **kwargs)
torch.cuda.synchronize()
self.tricksy_context.forward_times.append(time.time() - start)
self.tricksy_context.layer = 0
return out
def generate(self, *args, **kwargs):
# Load input embeddings for first token
self.model.decoder.load_weights(self.tricksy_context)
torch.cuda.synchronize()
out = super().generate(*args, **kwargs)
return out