added option to use SGD optimizer through the YAML, added option to pass in additional optimizer parameters through the YAML, added experimental unified AR+NAR model (does not seem fruitful in testing)

vall_e/config.py

@@ -254,6 +254,10 @@ class Models:
 	def ar(self):
 		return self.get("ar")
 
+	@property
+	def ar_nar(self):
+		return self.get("ar+nar")
+
 	@property
 	def nar(self):
 		return self.get("nar")
@@ -283,6 +287,7 @@ class Hyperparameters:
 	gradient_clipping: int = 100
 
 	optimizer: str = "Adamw"
+	optimizer_params: dict = field(default_factory=lambda: {})
 	learning_rate: float = 3.25e-4
 
 	scheduler_type: str = ""

vall_e/models/__init__.py

@@ -1,11 +1,14 @@
 from .ar import AR
 from .nar import NAR
+from .ar_nar import AR_NAR
 
 def get_model(cfg):
 	if cfg.name == "ar":
 		Model = AR
 	elif cfg.name == "nar":
 		Model = NAR
+	elif cfg.name == "ar+nar":
+		Model = AR_NAR
 	else:
 		raise f"invalid model name: {cfg.name}"
 	name = cfg.name

vall_e/models/ar.py

@@ -13,10 +13,6 @@ class AR(Base):
 	def causal(self):
 		return True
 
-	@property
-	def use_stop_token(self):
-		return True
-
 	@property
 	def norm_type(self):
 		return "ln"
@@ -45,10 +41,6 @@ class AR(Base):
 	def n_tasks(self) -> int:
 		return cfg.models.tasks
 
-	@property
-	def resp_loss_only(self) -> bool:
-		return False
-
 	@property
 	def recurrent_chunk_size(self) -> int:
 		if cfg.mode == "training":
@@ -103,8 +95,6 @@ class AR(Base):
 				resps_list=self._unsqueeze_list(resps_list),
 				targ_list=resps_list,
 				quant_levels=None,
-				shift_targ_list=True,
-				return_all_resp=False,
 			)
 
 		device = text_list[0].device
@@ -122,9 +112,10 @@ class AR(Base):
 			# get next in sequence
 
 			r = super().forward(
-				text_list,
+				text_list=text_list,
-				proms_list,
+				proms_list=proms_list,
-				self._unsqueeze_list(resps_list),
+				resps_list=self._unsqueeze_list(resps_list),
+				quant_levels=None,
 				sampling_temperature=sampling_temperature,
 				state=state
 			)
@@ -188,12 +179,14 @@ def example_usage():
 		'n_heads': 16,
 		'n_layers': 24,
 	}
+	
 	try:
 		kwargs['config'] = cfg.models.ar
 	except Exception as e:
 		pass
+
 	model = AR(**kwargs).to(device)
-	engine = Engine(model=model, optimizer=torch.optim.AdamW(model.parameters(), lr=1e-4))
+	engine = Engine(model=model, optimizer=torch.optim.SGD(model.parameters(), lr=0.1))
 	
 	def sample( name, steps=400 ):
 		engine.eval()

vall_e/models/ar_nar.py

@@ -0,0 +1,258 @@
+from ..config import cfg
+from .base import Base, list_to_tensor, Categorical
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+import random
+from einops import rearrange
+from torch import Tensor
+from tqdm import trange
+
+class AR_NAR(Base):
+	@property
+	def causal(self):
+		return True
+
+	@property
+	def norm_type(self):
+		return "ln"
+
+	@property
+	def arch_type(self) -> str:
+		if hasattr(self, "config") and self.config:
+			return self.config.arch_type
+		return cfg.models.ar_nar.arch_type
+
+	@property
+	def n_prom_levels(self) -> int:
+		return cfg.models.prom_levels
+
+	@property
+	def n_resp_levels(self) -> int:
+		if hasattr(self, "config") and self.config:
+			return self.config.resp_levels
+		return cfg.models.ar_nar.resp_levels
+
+	@property
+	def n_max_levels(self) -> int:
+		return cfg.models.max_levels
+
+	@property
+	def n_tasks(self) -> int:
+		return cfg.models.tasks
+
+	@property
+	def recurrent_chunk_size(self) -> int:
+		if cfg.mode == "training":
+			return 0
+		return cfg.inference.recurrent_chunk_size
+
+	@property
+	def interleave(self) -> bool:
+		if hasattr(self, "config") and self.config:
+			return self.config.interleave
+		return False
+
+	def _prune(self, l: Tensor):
+		indices = (l == self.stop_token).nonzero()
+		if len(indices) == 0:
+			return l
+		return l[: indices.min().item()]
+
+	def _interleave( self, codes ):
+		if not self.interleave:
+			return codes
+
+		return codes.flatten()
+
+	def _deinterleave( self, codes, length = 0 ):
+		if not self.interleave:
+			return codes
+
+		return torch.unflatten( codes[:codes.shape[0] // self.n_prom_levels * self.n_prom_levels], 0, ( codes.shape[0] // self.n_prom_levels, self.n_prom_levels ) )
+
+	@staticmethod
+	def _unsqueeze_list(x_list, axis=-1):
+		return [x.unsqueeze(dim=axis) for x in x_list]
+
+	def forward(
+		self,
+		text_list: list[Tensor],
+		proms_list: list[Tensor],
+		resps_list: list[Tensor] | None = None,
+		max_steps: int = 1000,
+		sampling_temperature: float = 0.0,
+	):
+		device = text_list[0].device
+		batch_size = len(text_list)
+
+		# is training or NAR
+		if resps_list is not None:
+			n_levels_set = {r.shape[-1] for r in resps_list}
+			n_levels = next(iter(n_levels_set))
+
+			# is training
+			if n_levels == self.n_resp_levels:
+				if random.random() < 0.5:
+					quant_levels = None
+
+					targ_list = [r[..., 0] for r in resps_list] # guarantees we only have the first levels
+					resps_list = self._unsqueeze_list(targ_list)
+				else:
+					quant_levels = torch.randint(1, self.n_resp_levels, (batch_size,))
+
+					targ_list = [o[...,   l] for o, l in zip(resps_list, quant_levels)]
+					resps_list = [o[..., : l] for o, l in zip(resps_list, quant_levels)]
+
+				if quant_levels is not None:
+					quant_levels.to(device=device)
+
+				return super().forward(
+					text_list=text_list,
+					proms_list=proms_list,
+					resps_list=resps_list,
+					targ_list=targ_list,
+					quant_levels=quant_levels,
+				)
+			# is NAR
+			prev_list = resps_list
+
+			while True:
+				level = prev_list[0].shape[-1] - 1
+
+				if level >= self.n_resp_levels:
+					break
+
+				quant_levels = torch.full((len(text_list),), level, device=device)
+
+				resps_list = super().forward(
+					text_list,
+					proms_list,
+					prev_list,
+					quant_levels=quant_levels,
+					sampling_temperature=sampling_temperature,
+				)
+
+				prev_list = [
+					torch.cat([rs, r.unsqueeze(-1)], dim=-1)
+					for rs, r in zip(prev_list, resps_list)
+				]
+
+			return prev_list
+
+		# is AR
+		resps_list = [ torch.zeros(0, device=device).to(torch.int16) for _ in text_list ]
+		stopped = torch.zeros(batch_size, device=device).bool()
+
+		state = {} if cfg.inference.recurrent_forward else None
+
+		if self.interleave:
+			max_steps *= self.n_prom_levels
+
+		for n in trange(max_steps // max(1, self.recurrent_chunk_size)):
+			# get next in sequence
+
+			r = super().forward(
+				text_list,
+				proms_list,
+				self._unsqueeze_list(resps_list),
+				sampling_temperature=sampling_temperature,
+				state=state
+			)
+
+			# append tokens
+			for i, ri in enumerate(r):
+				if self.stop_token in ri:
+					stopped[i] = True
+				resps_list[i] = torch.cat([resps_list[i], ri])
+
+			# stop token found
+			stopped |= r == self.stop_token
+			if stopped.all().item():
+				break
+
+		return [self._prune(r) for r in resps_list]
+
+
+def example_usage():
+	cfg.trainer.backend = "local"
+	from functools import partial
+
+	from einops import repeat
+
+	from ..emb.qnt import decode_to_file
+	from ..engines import Engine
+	from tqdm import tqdm
+
+	device = "cuda"
+	x8 = partial(repeat, pattern="t -> t l", l=cfg.models.prom_levels) 
+	symmap = {'<s>': 1, '</s>': 2, ' ': 3, '.': 4, ',': 5, '!': 6, '?': 7, 'p': 7, 'iː': 8, 'ɚ': 9, 'ˌ': 10, 'dˌ': 11, 'mˌ': 12, 'd': 13, 'ɹ': 14, 'tˈ': 15, 'pˌ': 16, 'uː': 17, 'l': 18, 'æ': 19, 'ɛ': 20, 'ɪ': 21, 'j': 22, 'ʊ': 23, 't': 24, 'n': 25, 'v': 26, 'a': 27, 'o': 28, 'ŋ': 29, 'w': 30, 'ʌ': 31, 'hˈ': 32, 'ɡˈ': 33, 'ə': 34, 'θˈ': 35, 'dˈ': 36, 'wˌ': 37, 'h': 38, 'z': 39, 'k': 40, 'ð': 41, 'ɡˌ': 42, 'ˈ': 43, 'fˈ': 44, 'i': 45, 's': 46, 'ʃ': 47, 'wˈ': 48, 'ðˈ': 49, 'ɹˈ': 50, 'lˈ': 51, 'ɡ': 52, 'oː': 53, 'mˈ': 54, 'e': 55, 'ɑː': 56, 'nˈ': 57, 'm': 58, 'θˌ': 59, 'sˈ': 60, 'f': 61, 'ɔː': 62, 'hˌ': 63, 'b': 64, 'jˈ': 65, 'ɐ': 66, 'ʒˈ': 67, 'θ': 68, 'bˈ': 69, 'ɾ': 70, 'ɜː': 71, 'ʌˈ': 72, 'ʃˌ': 73, 'bˌ': 74, 'kˈ': 75, 'ɔ': 76, 'zˈ': 77, 'ᵻ': 78, 'kˌ': 79, 'vˈ': 80, 'fˌ': 81, 'ʒ': 82, 'ʃˈ': 83, 'ɹˌ': 84, 'tˌ': 85, 'pˈ': 86, 'ðˌ': 87, 'sˌ': 88, 'nˌ': 89, 'lˌ': 90, '̩': 91, 'ʔ': 92, 'vˌ': 93, 'ɪˈ': 94, '"': 95, 'ɪˌ': 96, 'ʒˌ': 97, 'uːˌ': 98, 'ʊˈ': 99, 'jˌ': 100, 'uːˈ': 101, 'iːˈ': 102, 'zˌ': 103, '.ˈ': 104, '…': 105, 'ŋˌ': 106, 'ɐˌ': 107, '—ˈ': 108, 'iˌ': 109, 'iːˌ': 110, 'ɛː': 111, ')': 112, ')ˈ': 113, '(': 114, 'u': 115, '-': 116, 'ɖˈ': 117, 'iˈ': 118, 'ʰˈ': 119, 'ɟˈ': 120, '̃': 121, 'eː': 122, 'ɾˈ': 123, 'r': 124, 'ʰ': 125, '-ˌ': 126, 'ɫ': 127, 'q': 128, '—': 129, 'ʊˌ': 130, 'aː': 131, 'cˈ': 132, '…ˈ': 133, 'c': 134, 'ɳ': 135, 'ɐˈ': 136, 'x': 137, 'ʔˌ': 138, '.ˌ': 139, 'ɑ': 140, '?ˈ': 141, '̩ˈ': 142, '"ˈ': 143, ',ˈ': 144, 'ŋˈ': 145, 'əˌ': 146, '!ˈ': 147, '"ˌ': 148, '?ˌ': 149, ',ˌ': 150, '—ˌ': 151, '̩ˌ': 152, 'əˈ': 153, '!ˌ': 154, 'ɬ': 155, 'ʲ': 156, '¡': 157, 'ɯ': 158, 'qˌ': 159, 'ʑ': 160, 'ʑˈ': 161, '¿': 162, 'ɑːˈ': 163, 'iːː': 164, 'ɛˈ': 165, '¡ˈ': 166, 'æˈ': 167, 'ç': 168, 'ɾˌ': 169, 'ᵻˈ': 170, 'xˈ': 171, 'ɔːˈ': 172, ';': 173, 'ɬˌ': 174, ':': 175, 'ʔˈ': 176, 'ɑːˌ': 177, 'ɬˈ': 178}
+	def tokenize(content, lang_marker="en"):
+		split = content.split(" ")
+		phones = [f"<s>"] + [ " " if not p else p for p in split ] + [f"</s>"]
+		return torch.tensor([*map(symmap.get, phones)]).to()
+
+	qnt = torch.load("data/qnt.pt")[0].t()[:, :cfg.models.prom_levels].to(device)
+
+	text_list = [
+		#torch.tensor([1, 2, 3], device=device),
+		tokenize("ˈ a ɪ   w ɪ l   nˌ ɑː t  ˈ æ s k   ɐ   sˈ ɛ k ə n d   tˈ a ɪ m").to(device),
+	]
+	proms_list = [
+		#x8(torch.tensor([1, 2, 3], device=device)),
+		qnt.to(device),
+	]
+	resps_list = [
+		qnt.to(device),
+	]
+
+	text_list = text_list[:1]
+	proms_list = proms_list[:1]
+	resps_list = resps_list[:1]
+
+	kwargs = {
+		'n_tokens': 1024,
+		'd_model': 1024,
+		'n_heads': 16,
+		'n_layers': 24,
+	}
+	
+	"""
+	try:
+		kwargs['config'] = cfg.models.ar_nar
+	except Exception as e:
+		pass
+	"""
+
+	model = AR_NAR(**kwargs).to(device)
+	engine = Engine(model=model, optimizer=torch.optim.AdamW(model.parameters(), lr=0.001))
+	
+	def sample( name, steps=600 ):
+		engine.eval()
+		resps_list = engine(text_list, proms_list, max_steps=steps, sampling_temperature=0.95 )
+		
+		for i, o in enumerate(resps_list):
+			_ = decode_to_file(o, f"data/ar.{i}.{name}.wav", device=device)
+
+		resps_list = [r.unsqueeze(-1) for r in resps_list]
+		resps_list = engine( text_list, proms_list, resps_list=resps_list, sampling_temperature=0.2 )
+
+		for i, o in enumerate(resps_list):
+			_ = decode_to_file(o, f"data/ar+nar.{i}.{name}.wav", device=device)
+
+	def train():
+		engine.train()
+		t = trange(5000)
+		for i in t:
+			stats = {"step": i}
+			stats |= engine.traverse(text_list=text_list, proms_list=proms_list, resps_list=resps_list)
+
+			tqdm.write(f"{stats}")
+
+	sample("init", 75)
+	train()
+	sample("final")
+
+if __name__ == "__main__":
+	example_usage()

vall_e/models/base.py

@@ -94,14 +94,6 @@ class Base(nn.Module):
 	def causal(self) -> bool:
 		raise NotImplementedError
 
-	@property
-	def n_resp_levels(self) -> int:
-		raise NotImplementedError
-
-	@property
-	def use_stop_token(self) -> bool:
-		raise NotImplementedError
-
 	@property
 	def arch_type(self) -> str:
 		raise NotImplementedError
@@ -114,6 +106,10 @@ class Base(nn.Module):
 	def n_prom_levels(self) -> int:
 		raise NotImplementedError
 
+	@property
+	def n_resp_levels(self) -> int:
+		raise NotImplementedError
+
 	@property
 	def n_max_levels(self) -> int:
 		raise NotImplementedError
@@ -122,10 +118,6 @@ class Base(nn.Module):
 	def n_tasks(self) -> int:
 		raise NotImplementedError
 
-	@property
-	def resp_loss_only(self):
-		raise NotImplementedError
-
 	@property
 	def recurrent_chunk_size(self) -> int:
 		raise NotImplementedError
@@ -134,6 +126,24 @@ class Base(nn.Module):
 	def interleave(self) -> bool:
 		return False
 
+	@property
+	def stop_token(self):
+		if not self.causal:
+			raise ValueError("Not using stop token!")
+		return self.n_tokens
+
+	@property
+	def ignore_index(self):
+		return -100
+
+	@staticmethod
+	def _samplewise_merge_tensors(*l, sep: Tensor | None):
+		if sep is None:
+			cat = torch.cat
+		else:
+			cat = partial(_join, sep=sep)
+		return [*map(cat, zip(*l))]
+
 	def __init__(
 		self,
 		n_tokens: int = 1024,
@@ -155,7 +165,7 @@ class Base(nn.Module):
 
 		# +1 to include the stop token
 		n_prom_tokens = n_tokens + (self.n_tasks - 1) # - 1 because tts is an inherent task
-		n_resp_tokens = n_tokens + (1 if self.use_stop_token else 0) # AR requires a stop token to... know when to stop
+		n_resp_tokens = n_tokens + (1 if self.causal else 0) # AR requires a stop token to... know when to stop
 
 		self.text_emb = Embedding(n_tokens, d_model)
 		self.proms_emb = MultiEmbedding(self.n_prom_levels, n_prom_tokens, d_model)
@@ -208,24 +218,6 @@ class Base(nn.Module):
 			ignore_index=self.ignore_index,
 		)
 
-	@property
-	def stop_token(self):
-		if not self.use_stop_token:
-			raise ValueError("Not using stop token!")
-		return self.n_tokens
-
-	@property
-	def ignore_index(self):
-		return -100
-
-	@staticmethod
-	def _samplewise_merge_tensors(*l, sep: Tensor | None):
-		if sep is None:
-			cat = torch.cat
-		else:
-			cat = partial(_join, sep=sep)
-		return [*map(cat, zip(*l))]
-
 	@overload
 	def forward(
 		self,
@@ -234,9 +226,6 @@ class Base(nn.Module):
 		resps_list: list[Tensor],
 		targ_list: list[Tensor] | None = None,
 		quant_levels: Tensor | None = None,
-		shift_targ_list: bool = False,
-		return_all: Literal[False] = False,
-		return_all_resp: Literal[False] = False,
 		sampling_temperature: float = 1.0,
 	) -> Tensor:
 		...
@@ -249,9 +238,6 @@ class Base(nn.Module):
 		resps_list: list[Tensor],
 		targ_list: list[Tensor] | None = None,
 		quant_levels: Tensor | None = None,
-		shift_targ_list: bool = False,
-		return_all: Literal[True] = True,
-		return_all_resp: Literal[True] = True,
 		sampling_temperature: float = 1.0,
 	) -> list[Tensor]:
 		...
@@ -262,28 +248,12 @@ class Base(nn.Module):
 		proms_list: list[Tensor],
 		resps_list: list[Tensor],
 		targ_list: list[Tensor] | None = None,
+
 		quant_levels: Tensor | None = None,
-		shift_targ_list: bool = False,
-		return_all: bool = False,
-		return_all_resp: bool = False,
 		sampling_temperature: float = 1.0,
 
 		state: dict | None = None,
 	):
-		"""
-		Args:
-			text_list: [t] * b
-			proms_list: [t' l] * b, l quantization levels.
-			resps_list: [t'' l] * b, l quantization levels.
-			targ_list: [t''] * b, one quantization level only; when given, loss will be computed
-			quant_levels: specify which quant_levels to feed forward, used in NAR mode.
-			shift_targ_list: whether to shift target list when computing loss. True if AR.
-			return_all_resp: True if NAR.
-			sampling_temperature: a lower temperature makes the result more robust but less diverse.
-		Returns:
-			y: sampled tokens
-		"""
-
 		x_list = self._samplewise_merge_tensors(
 			self.text_emb(text_list),
 			self.proms_emb(proms_list),
@@ -334,17 +304,16 @@ class Base(nn.Module):
 
 			# process each batch
 			for i in range(len(text_prom_list)):
-				# for the NAR, ignore completely computing the loss against the text prompt
-				if self.resp_loss_only:
-					text_prom_list[i][:] = self.ignore_index
-
 				# for the AR, shift the text/input prompt into the future by 1, and ignore the rolled back text token
-				else:
+				if quant_levels is None:
 					text_prom_list[i] = text_prom_list[i].roll(-1, dims=0)
 					text_prom_list[i][-1] = self.ignore_index
+				# for the NAR, ignore completely computing the loss against the text prompt
+				else:
+					text_prom_list[i][:] = self.ignore_index
 
 			# adjust the target sequence if needed for the AR
-			if shift_targ_list:
+			if quant_levels is None:
 				# creates a copy because this is aliased against input response sequence
 				targ_list = [*targ_list] 
 				# shift the target response into the future by 1, and mark the rolled back token / last token as a stop token
@@ -370,10 +339,11 @@ class Base(nn.Module):
 			)
 
 		# return the entire generated token string
+		return_all = False
 		if return_all:
 			logits = [hi[:] for hi, li in zip(h_list, map(len, resps_list))]
 		# return the entire generated response
-		elif return_all_resp:
+		elif quant_levels is not None:
 			logits = [hi[-li:] for hi, li in zip(h_list, map(len, resps_list))]
 		# return the last chunkwise piece
 		elif self.causal and self.recurrent_chunk_size > 0:

vall_e/models/nar.py

@@ -11,10 +11,6 @@ class NAR(Base):
 	def causal(self):
 		return False
 
-	@property
-	def use_stop_token(self):
-		return False
-
 	@property
 	def arch_type(self) -> str:
 		if hasattr(self, "config") and self.config:
@@ -43,10 +39,6 @@ class NAR(Base):
 	def n_tasks(self) -> int:
 		return cfg.models.tasks
 
-	@property
-	def resp_loss_only(self) -> bool:
-		return True
-
 	@property
 	def recurrent_chunk_size(self) -> int:
 		return 0

vall_e/utils/trainer.py

@@ -62,16 +62,27 @@ def load_engines(invert=False):
 		optimizer = None
 		lr_scheduler = None
 
-		# yuck, should instead check be optimier == "adamw" and backend != "deepspeed"
+		# cfg.deepspeed.torch_adam
-		# and then have ds_cfg pass in the config flag to use pytorch adamw
-		# I genuinely cannot validate if this ever actually gets used in DeepSpeed
 		if (cfg.trainer.backend == "local" and cfg.hyperparameters.optimizer.lower() == "adamw") or (cfg.trainer.backend == "deepspeed" and cfg.hyperparameters.optimizer.lower() == "adamw-torch"):
+			params = {
+				"lr": cfg.hyperparameters.learning_rate,
+				"betas": (0.9, 0.96),
+				"eps": 1e-07,
+				"weight_decay": 0.01,
+			}
+			params.update(cfg.hyperparameters.optimizer_params)
 			optimizer = ml.AdamW(
 				model.parameters(),
-				lr=cfg.hyperparameters.learning_rate,
+				**params,
-				betas=(0.9, 0.96),
+			)
-				eps=1e-07,
+		elif (cfg.trainer.backend == "local" and cfg.hyperparameters.optimizer.lower() == "sgd") or (cfg.trainer.backend == "deepspeed" and cfg.hyperparameters.optimizer.lower() == "sgd-torch"):
-				weight_decay=0.01,
+			params = {
+				"lr": cfg.hyperparameters.learning_rate,
+			}
+			params.update(cfg.hyperparameters.optimizer_params)
+			optimizer = ml.SGD(
+				model.parameters(),
+				**params,
 			)
 
 		if not model._cfg.training:

vall_e/utils/wrapper.py

@@ -25,14 +25,17 @@ if cfg.bitsandbytes.enabled:
 			self.sparse,
 		)).to(self.weight.dtype) )
 
-Adam = torch.optim.Adam
-AdamW = torch.optim.AdamW
 
 if cfg.bitsandbytes.enabled:
 	import bitsandbytes as bnb
 
 	Adam = bnb.optim.Adam
 	AdamW = bnb.optim.AdamW
+	SGD = bnb.optim.SGD
+else:
+	Adam = torch.optim.Adam
+	AdamW = torch.optim.AdamW
+	SGD = torch.optim.SGD
 
 # handles generically converting to a specific tensor type and converting back (implemented solely for bfloat16)
 @contextmanager
@@ -73,3 +76,4 @@ if cfg.bitsandbytes.injects and cfg.bitsandbytes.enabled:
 
 	torch.optim.Adam = Adam
 	torch.optim.AdamW = AdamW
+	torch.optim.SGD = SGD