AcademiaSD/YuE2-3B-NF4-for-LoRA-Training

YuE2-3B NF4 for LoRA Training

YuE2-3B (the Mixture-of-Transformers song model by m-a-p) quantized to 4-bit NF4 for the YuE2-3B trainer of AcademiaSD LoRAlab Trainer Studio: song LoRA training on consumer NVIDIA GPUs. The trainer downloads it automatically.

You do not need to download this repository by hand: the trainer downloads it on first use (or copies an existing YuE2-3B_NF4/ folder next to its settings).

Contents

Folder	Content	Size
model/	YuE2 backbone cache: attention and MLP layers of both branches (AR and NAR) in NF4 (bitsandbytes, pre-quantized), latent bridges (llm2vae, vae2llm, time_embedder), lm_head, embeddings, norms and RoPE in BF16	3.0 GB
vae/	YuE2-VAE, byte-identical to the original	0.5 GB
modeling_yue2.py, config, tokenizer	Copied from the original model for offline loading without trust_remote_code	0.01 GB
Total		3.4 GB (original: 7.3 GB)

The quantization selects layers by the same name rules the trainer uses: 392 Linear4bit layers, 5 BF16 Linear layers (the latent bridges, the time embedder MLP and the output head) and 227 BF16 tensors in others.safetensors. The latent bridges and the output head stay in BF16 because the VAE latent scale and the flow matching are the precision-sensitive parts.

Format

model/ is not a standard transformers checkpoint. It is a pre-quantized cache, identical in layout to the other Trainer Studio NF4 caches: one .safetensors file per layer in weights/ (NF4 weight plus its packed bitsandbytes QuantState, or BF16 for the excluded layers), others.safetensors with the norms, embeddings and position buffers, and index.json listing every layer. The trainer builds an empty YuE2ForCausalLM from the config and fills it from this cache, so the BF16 backbone is never downloaded or loaded. A round-trip check rebuilds the model with QuantState.from_dict and compares AR logits against BF16: the relative error is ~7.6 %, bounded and expected of NF4; training losses match the on-the-fly quantization.

• The converter that produced it is tools/yue2/5_conversor_YuE2_NF4.py in the Trainer Studio repository.
• Loading from this cache takes ~5 s and 3.1 GB of VRAM (the raw model quantizes on the fly in ~23 s from the BF16 weights).

NF4 is meant for training. For full-quality inference, use the original model.

LoRAs trained with it

The trainer exports two kinds of files, one per stage — YuE2 is a Mixture-of-Transformers and every block carries an AR route (semantic codec tokens) and a NAR route (flow matching on the 64-dim audio latent); each stage LoRA only touches its own branch:

• YuE2_FINAL_LoRA.safetensors (NAR, timbre) and YuE2_FINAL_AR_LoRA.safetensors (AR, lyrics phrasing), with native keys (model.layers.N.nar_self_attn.…, …lora_A/lora_B.weight plus per-layer alpha tensor) ready for ComfyUI-style loaders.
• PEFT consumers add a base_model.model. prefix to the keys when injecting with set_peft_model_state_dict.
• Intermediate checkpoints (resume_checkpoint/, resume_checkpoint_ar/) are standard PEFT save_pretrained folders, so training resumes exactly step-to-step.

Measured on an RTX 5080 with 12 s clips, BF16 backbone: ~7.3 GB of VRAM and ~0.9 s per step at rank 16 (NAR); with this NF4 cache the backbone itself only takes 3.1 GB.

License

CC BY-NC 4.0, the same license as the original m-a-p/YuE2-3B. It also covers the LoRAs trained with these weights, as model derivatives. All credit for the model goes to the m-a-p team; this repository only changes the storage precision.