Changing sample CLI design; Refactor devices; Trying out quantization (…

…#70)

* MidiDataset can initialize with an iterator and only expand when necessary.

* reduce some memory overhead (we are starting to have >100k MidiDict and may get more in the future)

* classmethod+property is better...

* remove functools import

* use separate workers to build dataset instead of process pool

* add jsonl.zst support; unit test; fix bug

* receive context length via commandline. It's more convenient than digging into the config file every time.

* fix a minor output format mismatch when grad_checkpoint is true

* remove hardcoded cuda() as well as autocast for cpu inferencing; int8 quantization works; fix a bug on gradient_checkpointing along with use_cache

* fixing device

* fixing device

* bitsandbytes unnecessary now.

* Add a warning to force CPU when quantization is used in aria.run sample

* formatting; Also add black formatter to Makefile

aria/model/model.py

@@ -311,7 +311,7 @@ def forward(self, src: torch.Tensor, use_cache=False, past_kv=None):
  # remove torch.compile from the train script as this is not currently
  # supported.
  # Implements gradient checkpoints on Encoder Layers.
- if self.model_config.grad_checkpoint is True:
+ if self.model_config.grad_checkpoint is True and not use_cache:
  for layer in self.encode_layers:
 
  def create_custom_forward(module):
@@ -326,7 +326,6 @@ def custom_forward(*args):
  preserve_rng_state=True,
  use_reentrant=True,
  )
-
  else:
  new_past_kv = []
  past_kv = (

aria/run.py

@@ -2,26 +2,93 @@
 
 import argparse
 import os
+import re
 import sys
+import pathlib
+import warnings
 
 
 def _parse_sample_args():
  argp = argparse.ArgumentParser(prog="aria sample")
- argp.add_argument("model", help="name of model config file")
- argp.add_argument("ckpt_path", help="path to model checkpoint")
- argp.add_argument("midi_path", help="path to midi file")
+ argp.add_argument("-m", help="name of model config file")
+ argp.add_argument("-c", help="path to model checkpoint")
+ argp.add_argument("-p", help="path to midi file")
  argp.add_argument(
- "-var", help="number of variations", type=int, required=True
+ "-var",
+ help="number of variations",
+ type=int,
+ default=1,
  )
  argp.add_argument(
- "-trunc", help="length to truncated prompt", type=int, required=True
+ "-trunc",
+ help="length to truncated prompt",
+ type=int,
+ default=200,
  )
  argp.add_argument("-e", action="store_true", help="enable force end")
- argp.add_argument("-l", type=int, help="generation length")
+ argp.add_argument("-l", type=int, help="generation length", default=1024)
+ argp.add_argument("-q", action="store_true", help="quantize the model")
 
  return argp.parse_args(sys.argv[2:])
 
 
+def _get_model_name(name: str | None, state: dict):
+ if name is not None:
+ return name
+
+ print("Model name is not provided. Trying to infer from checkpoint...")
+ _defaults = {
+ 16: "small",
+ 32: "medium",
+ 64: "large",
+ 96: "xlarge",
+ }
+ try:
+ pattern = re.compile(r"encode_layers\.(\d+)\.")
+ layer_keys = [pattern.search(k) for k in state.keys()]
+ layer_keys = set(p.group(1) for p in layer_keys if p is not None)
+ for i in range(len(layer_keys)):
+ assert str(i) in layer_keys
+
+ if len(layer_keys) in _defaults:
+ print(f"Selecting model name: {_defaults[len(layer_keys)]}")
+ return _defaults[len(layer_keys)]
+ assert False
+ except:
+ raise ValueError("Model name is not provided and cannot be inferred.")
+
+
+def _show_popup(prompt: str, files: list) -> str:
+ for i in range(len(files)):
+ print(f" [{i}] {files[i]}")
+
+ for tries in range(3): # 3 tries in case of fat fingers
+ try:
+ res = int(input(prompt + f" [0-{len(files) - 1}]: "))
+ assert 0 <= res < len(files)
+ return files[res]
+ except:
+ print("Invalid input. Try again...")
+
+ raise ValueError("Invalid input.")
+
+
+def _get_ckpt_path(ckpt_path: str | None) -> str:
+ if ckpt_path is None:
+ ckpts = list(pathlib.Path(".").glob("*.bin"))
+ ckpt_path = _show_popup("Choose a checkpoint", ckpts)
+ return ckpt_path
+
+
+def _get_midi_path(midi_path: str | None) -> str:
+ if midi_path is None:
+ midis = list(pathlib.Path(".").glob("*.mid")) + list(
+ pathlib.Path(".").glob("*.midi")
+ )
+ midi_path = _show_popup("Choose a midi-file", midis)
+ return midi_path
+
+
 def sample(args):
  """Entrypoint for sampling"""
 
@@ -34,20 +101,80 @@ def sample(args):
  from aria.data.midi import MidiDict
  from aria.utils import midi_to_audio
 
- assert cuda_is_available() is True, "CUDA device not available"
+ if not cuda_is_available():
+ print("CUDA device is not available. Using CPU instead.")
+ else:
+ greedy_sample = torch.autocast(device_type="cuda", dtype=torch.float16)(
+ greedy_sample
+ )
+ device = (
+ torch.device("cuda") if cuda_is_available() else torch.device("cpu")
+ )
+
+ ckpt_path = _get_ckpt_path(args.c) # let user input path if not provided
+ model_state = torch.load(ckpt_path, map_location=device)
+ model_name = _get_model_name(
+ args.m, model_state
+ ) # infer model name if not provided
 
- model_name = args.model
- ckpt_path = args.ckpt_path
- midi_path = args.midi_path
  num_variations = args.var
  truncate_len = args.trunc
  force_end = args.e
 
  tokenizer = TokenizerLazy(return_tensors=True)
  model_config = ModelConfig(**load_model_config(model_name))
  model_config.set_vocab_size(tokenizer.vocab_size)
- model = TransformerLM(model_config).cuda()
- model.load_state_dict(torch.load(ckpt_path))
+ model = TransformerLM(model_config).to(device)
+ model.load_state_dict(model_state)
+ if args.q:
+ if device.type != "cpu":
+ warnings.warn(
+ "Quantization is not supported on CUDA devices. Using CPU instead."
+ )
+ device = torch.device("cpu")
+
+ from torch.ao.quantization import get_default_qconfig_mapping
+ from torch.quantization.quantize_fx import prepare_fx, convert_fx
+
+ qconfig_mapping = get_default_qconfig_mapping()
+
+ def _quantize(module, key, input_shape):
+ inp = torch.randn(input_shape, dtype=torch.float, device=device)
+ m = prepare_fx(
+ getattr(module, key), qconfig_mapping, example_inputs=inp
+ )
+ m = convert_fx(m)
+ setattr(module, key, m)
+
+ for i in range(len(model.model.encode_layers)):
+ _quantize(
+ model.model.encode_layers[i],
+ "mixed_qkv",
+ input_shape=(1, 2048, model_config.n_heads),
+ )
+ _quantize(
+ model.model.encode_layers[i],
+ "att_proj_linear",
+ input_shape=(1, 2048, model_config.n_heads),
+ )
+ _quantize(
+ model.model.encode_layers[i],
+ "ff_linear_1",
+ input_shape=(1, 2048, model_config.n_heads),
+ )
+ _quantize(
+ model.model.encode_layers[i],
+ "ff_linear_2",
+ input_shape=(
+ 1,
+ 2048,
+ model_config.n_heads * model_config.ff_mult,
+ ),
+ )
+
+ midi_path = _get_midi_path(
+ args.p
+ ) # let user input midi path if not provided
 
  if args.l and 0 < args.l < model.max_seq_len:
  max_gen_len = args.l
@@ -70,6 +197,7 @@ def sample(args):
  model,
  tokenizer,
  prompts,
+ device=device,
  force_end=force_end,
  max_seq_len=model_config.max_seq_len,
  max_gen_len=max_gen_len,
@@ -124,7 +252,7 @@ def _parse_tokenized_dataset_args():
  argp.add_argument("load_path", help="path midi_dict dataset")
  argp.add_argument("save_path", help="path to save dataset")
  argp.add_argument("-s", help="also produce shuffled", action="store_true")
- argp.add_argument("-l", help="max sequence length", type=int)
+ argp.add_argument("-l", help="max sequence length", type=int, default=2048)
 
  return argp.parse_args(sys.argv[2:])
 

aria/sample.py

@@ -43,13 +43,13 @@ def _get_cfg_coeff(cfg_gamma, cfg_mode, cur_pos, start_pos, total_len):
 # temp=0.85, top_p=0.9, cfg_gamma=1.4
 
 
-@torch.autocast(device_type="cuda", dtype=torch.float16)
 def greedy_sample(
  model: TransformerLM,
  tokenizer: Tokenizer,
  prompts: List[list],
  max_seq_len: int,
  max_gen_len: int,
+ device: torch.device | None = None,
  cfg_gamma: float | None = 1.4,
  cfg_mode: str | None = None,
  neg_prompts: List[list] | None = None,
@@ -67,6 +67,7 @@ def greedy_sample(
  prompts (List[list]): A list of prompts to sample as a batch.
  max_seq_len (int): Maximum sequence length supported by the model.
  max_gen_len (int): Maximum desired sequence length of the samples.
+ device (torch.device, optional): Device to use. Defaults to None.
  cfg_gamma (float, optional): CFG gamma parameter. Defaults to 1.2.
  This parameter *determines* whether parameters related to CFG are used.
  None: No CFG or interpolation. `cfg_mode, neg_prompts, neg_prompt_len, alpha` are ignored.
@@ -88,6 +89,7 @@ def greedy_sample(
  List[list]: The list of samples, decoded by the tokenizer.
  """
  assert tokenizer.return_tensors is True, "tokenizer must return tensors."
+ device = device or torch.device("cuda")
  model.eval()
 
  pad_id = tokenizer.pad_id
@@ -121,24 +123,28 @@ def greedy_sample(
  [
  torch.concat(
  [
- torch.full((neg_max_len - len(neg_seq),), pad_id),
- tokenizer.encode(neg_seq),
+ torch.full(
+ (neg_max_len - len(neg_seq),), pad_id, device=device
+ ),
+ tokenizer.encode(neg_seq).to(device),
  ]
  )
  for neg_seq in neg_prompts
  ],
  axis=0,
- ).cuda()
+ )
  neg_len = (
  neg_min_len
  if neg_prompt_len is None
  else min(neg_min_len, neg_prompt_len)
  )
  neg_tokens = neg_prompt_tensors[:, :neg_len]
 
- tokens = torch.full((bsz, total_len), pad_id).cuda()
+ tokens = torch.full((bsz, total_len), pad_id, device=device)
  for idx, unencoded_seq in enumerate(prompts):
- tokens[idx, : len(unencoded_seq)] = tokenizer.encode(unencoded_seq)
+ tokens[idx, : len(unencoded_seq)] = tokenizer.encode(unencoded_seq).to(
+ device
+ )
 
  dim_tok_inserted = [False for _ in range(bsz)]
  input_text_mask = tokens != pad_id