codebook assertion errors when generating long songs using quantised models

[sujitvasanth]

codebook assertion errors when generating long songs more than 2 minutes when using quantised models

my settings
max tokens 3000
segments 7
GPU rtx 3090
cache stage 1: 32768
cache stage 2: tried 16384, 32768
caches not quantised

using YuE-s1-7B-anneal-en-cot-exl2-4.0bpw

  File "/home/sujit/YuE-exllamav2/src/yue/infer_stage1.py", line 507, in <module>
    main()
  File "/home/sujit/YuE-exllamav2/src/yue/infer_stage1.py", line 500, in main
    pipeline.save(raw_output, args.output_dir, args.use_audio_prompt, args.use_dual_tracks_prompt)
  File "/home/sujit/YuE-exllamav2/src/yue/infer_stage1.py", line 175, in save
    vocals_ids = self.codec_tool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[0])
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sujit/YuE-exllamav2/src/yue/codecmanipulator.py", line 184, in ids2npy
    assert token_ids[0] >= codebook_0_range[0] \
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: token_ids[0]=54563, codebook_0_range=(45334, 46358)

was able to bypass stage 1 error with using
YuE-s1-7B-anneal-en-cot-exl2-4.25bpw
YuE-s2-1B-general-exl2-3.0bpw

Starting stage 2...
0it [00:00, ?it/s]
0it [00:00, ?it/s]
Traceback (most recent call last):
  File "/home/sujit/YuE-exllamav2/src/yue/infer_stage2.py", line 365, in <module>
    main()
  File "/home/sujit/YuE-exllamav2/src/yue/infer_stage2.py", line 356, in main
    outputs = pipeline.generate(output_dir=args.output_dir)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sujit/YuE-exllamav2/src/yue/infer_stage2.py", line 242, in generate
    prompt = self.get_codec_ids(prompt)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sujit/YuE-exllamav2/src/yue/infer_stage2.py", line 49, in get_codec_ids
    codec_ids = self.codec_tool.offset_tok_ids(
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sujit/YuE-exllamav2/src/yue/codecmanipulator.py", line 75, in offset_tok_ids
    assert x.max() < codebook_size, f"max(x)={x.max()}, codebook_size={codebook_size}"
           ^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: max(x)=11382, codebook_size=1024

just trying to remark out the assertions in codecmanipulator.py to see what happens...
removing the assertions does generate the output but I'm getting gibberish output after 2 mins 30... trying again with increased cache sizes

[sgsdxzy]

Can you verify if the problem is related to quantized models?

[brcisna]

Debian 13 Trixie
AMD Radeon Pro W6600 gfx1032 =/ gfx1030
Python3.10

I am getting the same error as @sujitvasanth ,, trying minmal parameters with the default models that are downloaded in the README.
Sometimes i will get a song to complete.. I know my GPU is only 8GB of VRAM,,not ideal,
I also noticed in the console each stage says " no hipblaslt found falling back to hipblas",,,although i have checked the ROCm hipblaslt config of this envionment and all checks pass.

Attempting to use hipBLASLt on an unsupported architecture! Overriding blas backend to hipblas (Triggered internally at /pytorch/aten/src/ATen/Context.cpp:310.) attn_weights = torch.matmul(q_states, k_states)
Cache Stage 1 = 5000 & Stage 2 = 8192

`YuE-extend/src/yue/codecmanipulator.py",`` line 75, in offset_tok_ids
assert x.max() < codebook_size, f"max(x)={x.max()}, codebook_size={codebook_size}"
AssertionError: max(x)=9001, codebook_size=1024

EDIT: I failed to mention I am actually using the -extended version of this code which provides a gradio interface so probably shouldnt have posted this error message. Apologizes.
`
TIA

[sujitvasanth]

@sgsdxzy I cant tell as I couldnt generate more than 2.5 minutes of audio with the non quantised models due to needing more VRAM (Ihave 24Gb)

[brcisna]

@sujitvasanth

Am curious do you actually get an OOM error,,,or is the assertion error,,,just another way of saying out of memory.
Sometimes,,i can rerun the same command with say two segments and the generation will complete.
For some reason,,i will almost always get about 20 seconds of music,before the text input actually starts.
I am very green at figuring this out,
Have tried many different cache settings, for stage one & two,, as well.

[sujitvasanth]

I get an assertion and no out of memory errors
when using the exl2 files (bpw 4,0 or 4.4 (this second one is more stable https://huggingface.co/Doctor-Shotgun/YuE-s1-7B-anneal-en-cot-exl2).

For me running larger number of segments without a quantised model leads to this issue #13

[DocShotgun]

This error essentially means that a token that shouldn't be sampled is being sampled, which could certainly be exacerbated by precision loss during quantization. I had this happen on occasion even with the unquantized model, and actually introduced d5c859e to address it by constraining the output during stage 1 to only tokens in the xcodec range.

However the first error message in the OP seems to suggest that only tokens in the first codebook of xcodec are valid? @turboderp did mention that he observed this to potentially be the case when examining the raw outputs. Perhaps the constraint needs to be made more aggressive.

Also 24gb of VRAM is definitely enough to generate with the unquantized model and fp16 cache, without limitation on song length, so not sure about your issue with that part.

[sujitvasanth]

@DocShotgun I tried unquantised modelsa and fp16 cache and I get this error for 3000 tokens, 6 segments when generating segment 6.

50%|█████     | 3/6 [04:11<04:10, 83.63s/it]
 87%|████████▋ | 2606/3000 [02:18<00:20, 18.78it/s]

 67%|██████▋   | 4/6 [06:30<03:30, 105.41s/it]
100%|██████████| 3000/3000 [03:17<00:00, 15.16it/s]

 83%|████████▎ | 5/6 [09:48<02:18, 138.78s/it]
 86%|████████▌ | 2579/3000 [03:29<00:34, 12.31it/s]

83%|████████▎ | 5/6 [13:17<02:39, 159.53s/it]
Traceback (most recent call last):
  File "/home/sujit/YuE-exllamav2/src/yue/infer_stage1.py", line 507, in <module>
    main()
  File "/home/sujit/YuE-exllamav2/src/yue/infer_stage1.py", line 484, in main
    raw_output = pipeline.generate(
                 ^^^^^^^^^^^^^^^^^^
  File "/home/sujit/YuE-exllamav2/src/yue/infer_stage1.py", line 418, in generate
    logits = logits.float().cpu()
             ^^^^^^^^^^^^^^^^^^^^
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

my settings can be seen here

Increasing the Stage 1 cache size to 64000 led to a successful generation of lentgtyh 2mins 30 seconds - just trying 7 segments,,,

[brcisna]

@sujitvasanth

I am curious. The Gradio interface you screenshot. Were did you download this gradio?
Reason asking. I setup the YuE-exllamav2-extended and it has a Gradio interface similar but does not have many of the settings options your Gradio GUI is showing.
There are getting to be so many forks of YuE is kind of confusing what is being added.

The items circled in screenshot i have below,,,i do not have in my Gradio GUI.

TIA, Barry

[DocShotgun]

@sujitvasanth There is no reason to increase stage1 cache size beyond 16384 since the model's max_position_embeddings is only 16384 anyways (without further modifications such as NTK-aware scaling).

I think if you edit this in infer_stage1.py

YuE-exllamav2/src/yue/infer_stage1.py

Line 357 in 8660f0f

gen_settings.allow_tokens(self.tokenizer, [32002] + list(range(45334, 56722)))

to

gen_settings.allow_tokens(self.tokenizer, [32002] + list(range(45334, 46358)))

Then the assertion error should go away, since we just constrain the sampling to tokens within that range. Again, I suspect that its a precision issue that's causing invalid tokens to have a high enough probability to be sampled. The longer the song is, the more tokens are generated, and the higher probability that an invalid token is sampled somewhere in the process.

Also for those of you using forks, please bring up issues with the fork authors - since there's no guarantee that the forks are on the same commit/don't have other commits that alter things.