575 lines
17 KiB
Python
575 lines
17 KiB
Python
"""
|
||
This is an experiment to:
|
||
* entertain a thought to try and abide by HF's transformers API (to benefit from caching better)
|
||
* conform to a single embedding (instead of a bunch of them) by folding/unfolding inputs
|
||
* stop trying to make a mixed AR+NAR model work since it seems lobotomized if I keep trying to enforce both recurrent and parallel inferencing (despite a penalty cost)
|
||
+ I will not cave and go with codebook patterns, not yet.
|
||
"""
|
||
|
||
from ..config import cfg
|
||
|
||
from ..data import fold_inputs, unfold_outputs
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from torch.nn.utils.rnn import pad_sequence
|
||
from torch import Tensor
|
||
from torch.nn import CrossEntropyLoss
|
||
from torch.utils.checkpoint import checkpoint
|
||
from torchmetrics.classification import BinaryAccuracy, MulticlassAccuracy, MulticlassPrecision
|
||
|
||
import random
|
||
import math
|
||
import logging
|
||
|
||
_logger = logging.getLogger(__name__)
|
||
|
||
from einops import rearrange
|
||
from tqdm import trange
|
||
|
||
from .arch import *
|
||
|
||
if cfg.model.arch_type not in AVAILABLE_ARCHES:
|
||
raise ValueError(f"Requesting arch `{cfg.model.arch_type}` but not available")
|
||
|
||
if cfg.model.arch_type in ["mamba","mamba2"]:
|
||
LlmArchClass = MambaLMHeadModel
|
||
elif cfg.model.arch_type == "llama":
|
||
LlmArchClass = LlamaForCausalLM
|
||
elif cfg.model.arch_type == "retnet":
|
||
LlmArchClass = RetNetForCausalLM
|
||
else:
|
||
raise ValueError(f"Requesting arch `{cfg.model.arch_type}` but not available")
|
||
|
||
class Model(LlmArchClass):
|
||
def __init__(
|
||
self,
|
||
|
||
n_text_tokens = 256,
|
||
n_audio_tokens = 1024,
|
||
|
||
d_model=1024,
|
||
n_layers=12,
|
||
n_heads=16,
|
||
p_dropout=0.1,
|
||
|
||
config = cfg.model,
|
||
):
|
||
self.hyper_config = config
|
||
|
||
hf_attention = config.attention if config is not None else None
|
||
gradient_checkpointing = config.gradient_checkpointing if config is not None else True
|
||
# text_tokens + rvq levels + [audio tokens * codebooks] (prom) + [audio tokens * codebooks] (resp) + stop
|
||
# vocab_size = n_text_tokens + cfg.model.max_levels + (n_audio_tokens * cfg.model.max_levels) + (n_audio_tokens * cfg.model.max_levels) + 1
|
||
|
||
if hf_attention == "auto":
|
||
if AVAILABLE_ATTENTIONS:
|
||
hf_attention = AVAILABLE_ATTENTIONS[0]
|
||
else:
|
||
hf_attention = "eager"
|
||
|
||
if hf_attention == "xformers":
|
||
hf_attention = "mem_efficient"
|
||
|
||
text_start = 0
|
||
text_end = text_start + config.text_tokens
|
||
|
||
lang_start = text_end
|
||
lang_end = lang_start + config.langs
|
||
|
||
rvq_start = lang_end
|
||
rvq_end = rvq_start + config.resp_levels
|
||
|
||
prom_start = rvq_end
|
||
prom_end = prom_start + config.audio_tokens * config.resp_levels
|
||
|
||
task_start = prom_end
|
||
task_end = task_start + config.tasks
|
||
|
||
tone_start = task_end
|
||
tone_end = tone_start + config.tones
|
||
|
||
resp_start = tone_end
|
||
resp_end = resp_start + config.audio_tokens * config.resp_levels
|
||
|
||
vocab_size = resp_end
|
||
|
||
if config.arch_type == "llama":
|
||
super().__init__(config=LlamaConfig(
|
||
vocab_size=vocab_size,
|
||
hidden_size=d_model,
|
||
max_position_embeddings=cfg.dataset.frames_per_second * config.max_levels * 60, # max-length of 60 seconds
|
||
intermediate_size=d_model*4,
|
||
num_hidden_layers=n_layers,
|
||
num_attention_heads=n_heads,
|
||
attention_dropout=p_dropout,
|
||
num_key_value_heads=n_heads,
|
||
sliding_window=cfg.dataset.frames_per_second * config.max_levels * 12,
|
||
hidden_act="gelu",
|
||
is_encoder_decoder=False,
|
||
is_decoder=True,
|
||
attn_implementation=hf_attention,
|
||
))
|
||
|
||
if gradient_checkpointing:
|
||
self.gradient_checkpointing_enable(gradient_checkpointing_kwargs=dict(
|
||
use_reentrant=False
|
||
))
|
||
elif config.arch_type == "retnet":
|
||
super().__init__(config=RetNetConfig(
|
||
vocab_size=vocab_size,
|
||
decoder_embed_dim=d_model,
|
||
decoder_value_embed_dim =d_model * 2,
|
||
decoder_retention_heads=n_heads,
|
||
decoder_ffn_embed_dim=d_model * 4,
|
||
decoder_layers=n_layers,
|
||
dropout=p_dropout,
|
||
checkpoint_activations=gradient_checkpointing,
|
||
activation_fn="gelu",
|
||
use_layernorm=False,
|
||
use_biases=False,
|
||
use_glu=True,
|
||
|
||
#chunkwise_recurrent=self.causal and self.recurrent_chunk_size > 0,
|
||
#recurrent_chunkwise_size=self.recurrent_chunk_size if self.causal else 0,
|
||
#no_output_layer=True,
|
||
#rotary_embedding_base=self.rotary_embedding_base, # 10000
|
||
|
||
decoder_normalize_before=True,
|
||
))
|
||
elif config.arch_type in ["mamba","mamba2"]:
|
||
super().__init__(config=MambaConfig(
|
||
vocab_size=vocab_size,
|
||
d_model=d_model,
|
||
n_layer=n_layers*2,
|
||
d_intermediate=0, # d_model*4,
|
||
ssm_cfg={"layer": "Mamba2", "use_mem_eff_path": True} if config.arch_type == "mamba2" else {},
|
||
rms_norm=True,
|
||
fused_add_norm=True,
|
||
residual_in_fp32=False,
|
||
))
|
||
|
||
self.backbone.gradient_checkpointing = gradient_checkpointing
|
||
|
||
self.accuracy_metric = None if True else MulticlassAccuracy(
|
||
vocab_size,
|
||
top_k=10,
|
||
average="micro",
|
||
multidim_average="global",
|
||
ignore_index=-100,
|
||
)
|
||
|
||
def generate(
|
||
self,
|
||
*args,
|
||
**kwargs
|
||
):
|
||
if cfg.model.arch_type in ["mamba","mamba2"]:
|
||
kwargs["cg"] = True
|
||
|
||
if "attention_mask" in kwargs:
|
||
kwargs.pop("attention_mask")
|
||
|
||
if "do_sample" in kwargs:
|
||
kwargs.pop("do_sample")
|
||
|
||
if "min_length" in kwargs:
|
||
kwargs.pop("min_length")
|
||
|
||
"""
|
||
if "position_ids" in kwargs:
|
||
kwargs.pop("position_ids")
|
||
|
||
if "max_new_tokens" in kwargs:
|
||
kwargs.pop("max_new_tokens")
|
||
|
||
if "max_length" not in kwargs:
|
||
kwargs["max_length"] = 500 * (self.hyper_config.resp_levels if self.hyper_config.experimental.interleave else 1)
|
||
|
||
if "num_last_tokens" not in kwargs:
|
||
kwargs["num_last_tokens"] = self.hyper_config.experimental.causal_size
|
||
"""
|
||
|
||
input_ids = kwargs.pop("input_ids")
|
||
attention_mask = kwargs.pop("attention_mask", None)
|
||
position_ids = kwargs.pop("position_ids", None)
|
||
|
||
stop_token = kwargs.pop("eos_token_id", 3)
|
||
max_steps = kwargs.pop("max_new_tokens", 500)
|
||
|
||
device = input_ids.device
|
||
batch_size = input_ids.shape[0]
|
||
|
||
sequence_list = [ inputs for inputs in input_ids ]
|
||
position_list = [ positions for positions in position_ids ]
|
||
|
||
start_positions = [ inputs.shape[0] for inputs in input_ids ]
|
||
|
||
stopped = torch.zeros(batch_size, device=device).bool()
|
||
|
||
config = self.hyper_config
|
||
state = None
|
||
disable_tqdm = False
|
||
causal_size = config.experimental.causal_size
|
||
|
||
# get next in sequence
|
||
for n in trange(max_steps // max(1, causal_size), desc="AR", disable=disable_tqdm):
|
||
output = super().forward(
|
||
input_ids=torch.stack(sequence_list),
|
||
#attention_mask=attention_mask,
|
||
#past_key_values=state,
|
||
#position_ids=torch.stack(position_list),
|
||
#use_cache=False,
|
||
#return_dict=False
|
||
)
|
||
|
||
logits = output[0]
|
||
# state = output[1]
|
||
|
||
r = [ logit[-causal_size:].argmax(dim=1) for logit in logits ]
|
||
|
||
# append tokens
|
||
for i, ri in enumerate(r):
|
||
if stop_token in ri:
|
||
stopped[i] = True
|
||
|
||
last_position_id = position_list[i][-1].item() + 1
|
||
sequence_list[i] = torch.cat([ sequence_list[i], ri.to(device) ], dim=0)
|
||
#position_list[i] = torch.cat([ position_list[i], torch.tensor([ last_position_id + _ for _ in range( ri.shape[0] ) ], device=device, dtype=torch.int32) ])
|
||
|
||
# stop token found
|
||
stopped |= r == stop_token
|
||
if stopped.all().item():
|
||
break
|
||
|
||
def _prune(l: Tensor, stop = stop_token):
|
||
indices = (l == stop).nonzero()
|
||
|
||
if len(indices) == 0:
|
||
return l
|
||
|
||
return l[: indices.min().item()]
|
||
|
||
sequence_list = [ _prune(seq[start_positions[i]:], stop_token) for i, seq in enumerate(sequence_list) ]
|
||
return torch.stack(sequence_list)
|
||
|
||
"""
|
||
return super().generate(*args, **kwargs)
|
||
"""
|
||
|
||
def forward(
|
||
self,
|
||
*args,
|
||
**kwargs,
|
||
):
|
||
config = self.hyper_config
|
||
|
||
if "text_list" in kwargs:
|
||
text_list = kwargs.pop("text_list", None)
|
||
proms_list = kwargs.pop("proms_list", None)
|
||
resps_list = kwargs.pop("resps_list", None)
|
||
lang_list = kwargs.pop("lang_list", None)
|
||
tone_list = kwargs.pop("tone_list", None)
|
||
|
||
training = kwargs.pop("training", False)
|
||
steps = kwargs.pop("steps", 500)
|
||
|
||
batch_size = len(text_list)
|
||
|
||
if training:
|
||
quant_levels = None if config.experimental.interleave else [ random.randint( 0 if "ar" in config.capabilities else 1, config.max_levels - 1) for _ in range(batch_size) ]
|
||
|
||
input_ids, attention_mask, position_ids = fold_inputs(
|
||
text_list=text_list,
|
||
prom_list=proms_list,
|
||
resp_list=resps_list,
|
||
targ_list=resps_list,
|
||
quant_levels=quant_levels,
|
||
)
|
||
target_ids, target_attention_mask, target_position_ids = fold_inputs(
|
||
text_list=text_list,
|
||
prom_list=proms_list,
|
||
resp_list=resps_list,
|
||
targ_list=resps_list,
|
||
quant_levels=quant_levels,
|
||
ignore_index=-100
|
||
)
|
||
return self.forward(
|
||
input_ids=input_ids,
|
||
labels=target_ids,
|
||
position_ids=position_ids,
|
||
|
||
quant_levels=quant_levels,
|
||
)
|
||
|
||
if config.experimental.interleave:
|
||
input_ids, attention_mask, position_ids = fold_inputs( text_list=text_list, prom_list=proms_list )
|
||
output = self.generate(
|
||
input_ids=input_ids,
|
||
position_ids=position_ids,
|
||
attention_mask=attention_mask,
|
||
eos_token_id=3,
|
||
do_sample=True,
|
||
max_new_tokens=steps*config.max_levels,
|
||
)
|
||
return unfold_outputs( output )["resp_list"]
|
||
|
||
resps_list = [ [] for _ in range(batch_size) ]
|
||
for l in range(config.max_levels):
|
||
quant_levels = [ l for _ in range(batch_size) ]
|
||
|
||
input_ids, attention_mask, position_ids = fold_inputs(text_list=text_list, prom_list=proms_list, resp_list=resps_list, quant_levels=quant_levels)
|
||
min_length = 1
|
||
for batch in input_ids:
|
||
min_length = max( min_length, batch.shape[0] + 1 )
|
||
|
||
# to-do: figure out a way to do one forward pass but sample N tokens to replicate the NAR sample pass
|
||
output = self.generate(
|
||
input_ids=input_ids,
|
||
attention_mask=attention_mask,
|
||
position_ids=position_ids,
|
||
eos_token_id=3,
|
||
do_sample=True,
|
||
max_new_tokens=steps,
|
||
)
|
||
|
||
unfolded = unfold_outputs( output, quant_levels=quant_levels )
|
||
|
||
if l == 0:
|
||
steps = 0
|
||
|
||
for batch, resp in enumerate(unfolded["resp_list"]):
|
||
length = resp.shape[-1]
|
||
|
||
# store length
|
||
if l == 0:
|
||
steps = max( steps, length )
|
||
# pad
|
||
else:
|
||
resp = resp[:steps]
|
||
if length < steps:
|
||
resp = torch.cat([ resp, torch.Tensor([ 0 for _ in range(steps-length) ]).to(resp) ])
|
||
resps_list[batch].append( resp )
|
||
|
||
for i, resp in enumerate( resps_list ):
|
||
resps_list[i] = torch.stack( resp ).t()
|
||
|
||
return resps_list
|
||
|
||
if config.arch_type in ["mamba","mamba2"]:
|
||
kwargs.pop("attention_mask", None)
|
||
|
||
labels = kwargs.pop("labels", None)
|
||
quant_levels = kwargs.pop("quant_levels", None)
|
||
|
||
output = super().forward(*args, **kwargs)
|
||
logits = output.logits
|
||
|
||
# i HATE the correct way
|
||
if labels is not None:
|
||
if quant_levels is None:
|
||
quant_levels = [0 for _ in range(labels.shape[0])]
|
||
|
||
# predict the next token for AR, else predict in place
|
||
loss = sum([ F.cross_entropy(
|
||
logit[:-config.experimental.causal_size, :] if quant_level == 0 or "nar" not in config.capabilities else logit,
|
||
label[config.experimental.causal_size:] if quant_level == 0 or "nar" not in config.capabilities else label,
|
||
ignore_index=-100
|
||
) for logit, label, quant_level in zip( logits, labels, quant_levels ) ])
|
||
|
||
self.loss = dict(
|
||
nll = loss,
|
||
)
|
||
|
||
if self.accuracy_metric is not None:
|
||
self.stats = dict(
|
||
acc = (sum([ self.accuracy_metric( logit, target ) for logit, target in zip( logits, labels ) ] ) / len( logits )).item()
|
||
)
|
||
|
||
"""
|
||
if config.loss_factors:
|
||
sep = 3
|
||
# determine specific sections to focus on
|
||
indices = [ [ idx for idx, token in enumerate( batch ) if token == sep ] for i, batch in enumerate( labels ) ]
|
||
|
||
text_index = 0
|
||
resp_index = 1 # 1 includes everything non text, -3 includes pre_resp + resp (ignores prom, probably better to include prom here)
|
||
|
||
labels_text = [ batch[:indices[i][text_index] + 1 ] for i, batch in enumerate( labels ) ]
|
||
labels_resp = [ batch[indices[i][resp_index] + 1:] for i, batch in enumerate( labels ) ]
|
||
|
||
logits_text = [ batch[:indices[i][text_index] + 1 ] for i, batch in enumerate( logits ) ]
|
||
logits_resp = [ batch[indices[i][resp_index] + 1:] for i, batch in enumerate( logits ) ]
|
||
|
||
loss_text = sum([ F.cross_entropy( logit[:-1, :], label[1:], ignore_index=-100 ) for logit, label in zip( logits_text, labels_text ) ]) / len(logits_text) * self.hyper_config.loss_factor("text")
|
||
loss_resp = sum([ F.cross_entropy( logit[:-1, :], label[1:], ignore_index=-100 ) for logit, label in zip( logits_resp, labels_resp ) ]) / len(logits_resp) * self.hyper_config.loss_factor("resp")
|
||
|
||
self.loss = dict(
|
||
text = loss_text,
|
||
resp = loss_resp,
|
||
)
|
||
|
||
if self.accuracy_metric is not None:
|
||
self.stats = dict(
|
||
acc = dict(
|
||
text = (sum([ self.accuracy_metric( logit, target ) for logit, target in zip( logits_text, labels_text ) ] ) / len( logits_text )).item(),
|
||
resp = (sum([ self.accuracy_metric( logit, target ) for logit, target in zip( logits_resp, labels_resp ) ] ) / len( logits_resp )).item(),
|
||
)
|
||
)
|
||
"""
|
||
|
||
return output
|
||
|
||
def example_usage():
|
||
cfg.trainer.backend = "local"
|
||
cfg.hyperparameters.gradient_accumulation_steps = 1
|
||
if cfg.audio_backend == "dac":
|
||
cfg.sample_rate = 44_100
|
||
|
||
from functools import partial
|
||
from einops import repeat
|
||
from tqdm import tqdm
|
||
|
||
from ..emb.qnt import decode_to_file, unload_model
|
||
from ..engines import Engine
|
||
from ..utils import wrapper as ml
|
||
|
||
import numpy as np
|
||
import re
|
||
|
||
device = "cuda"
|
||
|
||
def tokenize(content):
|
||
return torch.tensor( cfg.tokenizer.encode(content) )
|
||
|
||
def _load_quants(path) -> Tensor:
|
||
qnt = np.load(path, allow_pickle=True)[()]
|
||
return torch.from_numpy(qnt["codes"].astype(np.int16))[0, :cfg.model.max_levels, :].t().to(torch.int16)
|
||
|
||
qnt = _load_quants(f"./data/qnt.{'dac' if cfg.audio_backend == 'dac' else 'enc'}")
|
||
|
||
|
||
text_list = [
|
||
tokenize("ˈaɪ wɪl nˌɑːt ˈæsk ɐ sˈɛkənd tˈaɪm").to(device),
|
||
#tokenize("ˈaɪ wɪl nˌɑːt ˈæsk ɐ sˈɛkənd tˈaɪm").to(device),
|
||
]
|
||
prom_list = [
|
||
qnt[:cfg.dataset.frames_per_second, :].to(device),
|
||
#qnt[:cfg.dataset.frames_per_second, :].to(device),
|
||
]
|
||
resp_list = [
|
||
qnt[:, :].to(device),
|
||
#qnt[cfg.dataset.frames_per_second:, :].to(device),
|
||
#qnt[:cfg.dataset.frames_per_second, :].to(device),
|
||
]
|
||
|
||
text_list = text_list[:1]
|
||
prom_list = prom_list[:1]
|
||
resp_list = resp_list[:1]
|
||
|
||
kwargs = {}
|
||
model = Model(**kwargs).to(device)
|
||
steps = 50 # 100 if cfg.model.experimental.interleave else 300
|
||
|
||
optimizer = cfg.hyperparameters.optimizer.lower() if cfg.yaml_path is not None else "prodigy"
|
||
scheduler = cfg.hyperparameters.scheduler.lower() if cfg.yaml_path is not None else ""
|
||
learning_rate = cfg.hyperparameters.learning_rate if cfg.yaml_path is not None else None
|
||
|
||
if cfg.optimizations.dadaptation:
|
||
# do not combine the two
|
||
if scheduler == "schedulefree":
|
||
scheduler = ""
|
||
|
||
learning_rate = 1.0
|
||
|
||
if optimizer == "prodigy":
|
||
if learning_rate is None:
|
||
learning_rate = 1.0
|
||
|
||
optimizer = ml.Prodigy
|
||
elif optimizer == "adagrad":
|
||
if learning_rate is None:
|
||
learning_rate = 1.0e-2
|
||
|
||
optimizer = ml.Adagrad
|
||
elif optimizer == "adamw":
|
||
if learning_rate is None:
|
||
learning_rate = 1.0e-4
|
||
|
||
optimizer = ml.AdamW
|
||
elif optimizer == "sdg":
|
||
if learning_rate is None:
|
||
learning_rate = 1.0e-4
|
||
|
||
optimizer = ml.SGD
|
||
else:
|
||
raise ValueError(f"Unrecognized optimizer: {optimizer}")
|
||
|
||
_logger.info(f"Optimizer: {optimizer}\tLearning rate: {learning_rate}")
|
||
|
||
optimizer = optimizer(model.parameters(), lr=learning_rate)
|
||
|
||
if scheduler == "schedulefree":
|
||
if isinstance(optimizer, ml.AdamW):
|
||
scheduler = ml.schedulefree.AdamWScheduleFree
|
||
elif isinstance(optimizer, ml.SGD):
|
||
scheduler = ml.schedulefree.SGDScheduleFree
|
||
else:
|
||
scheduler = None
|
||
|
||
if scheduler is not None:
|
||
_logger.info(f"Scheduler: {scheduler}")
|
||
optimizer = scheduler( model.parameters(), lr = learning_rate )
|
||
|
||
if cfg.optimizations.replace and cfg.optimizations.linear:
|
||
model = ml.replace_linear( model )
|
||
|
||
if cfg.optimizations.replace and cfg.optimizations.embedding:
|
||
model = ml.replace_embedding( model )
|
||
|
||
engine = Engine(model=model, optimizer=optimizer)
|
||
|
||
"""
|
||
torch.save( {
|
||
'module': model.state_dict()
|
||
}, f"./data/{cfg.model.arch_type}.pth" )
|
||
"""
|
||
|
||
_logger.info(f"{LlmArchClass} parameter count: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
|
||
|
||
@torch.inference_mode()
|
||
def sample( name, steps=cfg.model.max_levels*cfg.dataset.frames_per_second*6 ):
|
||
engine.eval()
|
||
|
||
resp_list = model( text_list=text_list, proms_list=prom_list )
|
||
|
||
for i, batch in enumerate(resp_list):
|
||
_ = decode_to_file(batch.to(device=device), f"data/{cfg.model.arch_type}.{cfg.audio_backend}.{i}.{name}.wav", device=device)
|
||
|
||
unload_model()
|
||
|
||
def train():
|
||
engine.train()
|
||
t = trange(steps)
|
||
for i in t:
|
||
stats = {"step": i}
|
||
|
||
stats |= engine.traverse(text_list=text_list, proms_list=prom_list, resps_list=resp_list, training=True)
|
||
stats |= engine.gather_attribute("stats")
|
||
stats |= {"grad_norm": engine.get_global_grad_norm()}
|
||
|
||
tqdm.write(f"{stats}")
|
||
|
||
"""
|
||
torch.save( {
|
||
'module': model.state_dict()
|
||
}, f"./data/{cfg.model.arch_type}.pth" )
|
||
"""
|
||
|
||
#sample("init", 5)
|
||
train()
|
||
sample("final")
|
||
|
||
if __name__ == "__main__":
|
||
example_usage()
|