role: user

type: text

text:
<user_query>
Examine the VibeVoice fine-tuning code at /home/ubuntu/VibeVoice-finetuning/src/ to determine if it supports full fine-tuning (not just LoRA).

Specifically check:
1. /home/ubuntu/VibeVoice-finetuning/src/finetune_vibevoice_lora.py - look for any flags/options that allow full fine-tuning instead of LoRA
2. Check if there's a separate full fine-tuning script (finetune_vibevoice.py or similar)
3. Look at the training setup code to understand:
   - How the model is loaded
   - How LoRA is applied (can it be skipped?)
   - What parameters control LoRA vs full fine-tuning
   - How the diffusion head and connectors are trained (they might already be fully fine-tuned)
4. Check /home/ubuntu/vibevoice-community/vibevoice/finetune/ for any additional training scripts
5. Check how voice_prompt_drop_rate is implemented

Return:
- Whether full fine-tuning is supported out of the box
- If not, what code changes would be needed
- The recommended approach for full fine-tuning with this codebase
- Memory/GPU requirements for full fine-tuning of 1.5B model
</user_query>


role: assistant

text: Let me systematically examine the VibeVoice fine-tuning codebase to answer all the user's questions. I'll start by exploring the directory structure and then read the relevant files.


text: Now let me look at the model architecture to understand parameter counts.


text: Now let me check the diffusion head architecture to understand its size.


text:
Here is a thorough analysis of the VibeVoice fine-tuning codebase regarding full fine-tuning support.

---

## 1. Does the script support full fine-tuning out of the box?

**No, the LLM backbone does NOT support full fine-tuning out of the box.** However, the diffusion head and connectors already support full fine-tuning via existing flags.

The critical blocker is in the training setup logic. Here's the sequence in both `finetune_vibevoice_lora.py` scripts (both copies are nearly identical):

**Step 1** -- LoRA is applied to the LLM (line 381 in the finetuning repo):

```376:383:/home/ubuntu/VibeVoice-finetuning/src/finetune_vibevoice_lora.py
    lora_cfg = build_lora_config(model_args)
    tm_lower = [s.strip().lower() for s in model_args.lora_target_modules.split(",") if s.strip()]
    skip_lm_lora = (len(tm_lower) == 0) or all(t in ("none", "off", "disable", "disabled") for t in tm_lower)
    if not skip_lm_lora:
        model.model.language_model = get_peft_model(model.model.language_model, lora_cfg)
    else:
        logger.info("Skipping LLM LoRA wrapping (lora_target_modules indicates none).")
```

**Step 2** -- ALL parameters are frozen, then only LoRA params are unfrozen (lines 391-399):

```390:399:/home/ubuntu/VibeVoice-finetuning/src/finetune_vibevoice_lora.py
    for _, p in model.named_parameters():
        p.requires_grad = False

    try:
        for n, p in model.model.language_model.named_parameters():
            if "lora_A" in n or "lora_B" in n:
                p.requires_grad = True
    except Exception:
        logger.warning("Could not re-enable LoRA params on language_model.")
```

**Step 3** -- Embeddings and LM head are always frozen (lines 459-468):

```459:468:/home/ubuntu/VibeVoice-finetuning/src/finetune_vibevoice_lora.py
    try:
        emb = model.get_input_embeddings()
        if hasattr(emb, "weight"):
            emb.weight.requires_grad_(False)
        head = model.get_output_embeddings()
        if head is not None and hasattr(head, "weight"):
            head.weight.requires_grad_(False)
    except Exception:
        pass
```

---

## 2. Existing flags that partially enable full training

The script already has these `ModelArguments` flags:

| Flag | Default | Effect |
|------|---------|--------|
| `--lora_target_modules` | `"q_proj,k_proj,..."` | Setting to `"none"` skips PEFT wrapping of the LLM |
| `--train_diffusion_head` | `False` | Unfreezes **all** diffusion head parameters (full fine-tune) |
| `--train_connectors` | `False` | Unfreezes acoustic + semantic connectors (full fine-tune) |
| `--lora_wrap_diffusion_head` | `False` | Alternative: LoRA-wrap the diffusion head instead |
| `--layers_to_freeze` | `None` | Selectively re-freeze diffusion head layer indices |
| `--freeze_acoustic_tokenizer` | `True` | Controls frozen state of acoustic tokenizer |
| `--freeze_semantic_tokenizer` | `True` | Controls frozen state of semantic tokenizer |

So the **diffusion head** and **connectors** already support full fine-tuning. But the **LLM backbone** does not -- even if you set `--lora_target_modules none` to skip LoRA wrapping, the "freeze everything then only unfreeze LoRA params" logic at lines 391-399 means the LLM stays entirely frozen since there are no `lora_A`/`lora_B` parameters to unfreeze.

---

## 3. What code changes are needed for full LLM fine-tuning

You need to modify the `requires_grad` logic in `finetune_vibevoice_lora.py`. The key change is: **when LoRA is skipped, unfreeze all LLM parameters instead of looking for `lora_A`/`lora_B` params.**

Specifically, replace the block at lines 390-399 with something like:

```python
for _, p in model.named_parameters():
    p.requires_grad = False

if skip_lm_lora:
    # Full fine-tuning: unfreeze all LLM parameters
    for n, p in model.model.language_model.named_parameters():
        p.requires_grad = True
else:
    # LoRA: only unfreeze adapter params
    for n, p in model.model.language_model.named_parameters():
        if "lora_A" in n or "lora_B" in n:
            p.requires_grad = True
```

You'd also need to decide whether to unfreeze the embeddings and LM head (lines 459-468). For full fine-tuning, you likely want them trainable. You could add a flag like `--train_embeddings` to control this.

Additionally, the `_save` method and post-training save logic are LoRA-oriented (they call `save_pretrained` for PEFT adapters). For full fine-tuning, you'd want to either:
- Call `model.save_pretrained(output_dir)` to save the entire model, or
- Save the full LLM state dict similarly to how connectors are saved

---

## 4. How `voice_prompt_drop_rate` is implemented

It's a simple random dropout applied in the `VibeVoiceCollator.__call__` method. In `data_vibevoice.py` (line 232-233):

```232:233:/home/ubuntu/VibeVoice-finetuning/src/data_vibevoice.py
            proc = self.processor(
                text=[text],
                voice_samples=[voice_prompts] if voice_prompts is not None and random.random() >= _drop_rate else None,
```

For each training example, a `random.random()` draw is compared against `_drop_rate`. If the random value is below the drop rate, `voice_samples` is passed as `None` to the processor (effectively removing the voice conditioning prompt). This is classifier-free guidance training -- at drop rate 0.0, prompts are always kept; at 1.0, they're always dropped.

The rate is clamped to [0.0, 1.0] at lines 226-230:

```226:230:/home/ubuntu/VibeVoice-finetuning/src/data_vibevoice.py
            _drop_rate = self.voice_prompt_drop_rate
            if _drop_rate < 0.0:
                _drop_rate = 0.0
            elif _drop_rate > 1.0:
                _drop_rate = 1.0
```

---

## 5. Model architecture and parameter count summary

From the 1.5B config (`qwen2.5_1.5b_64k.json`):

| Component | Architecture | Estimated Parameters |
|-----------|-------------|---------------------|
| **LLM (Qwen2 1.5B)** | 28 layers, hidden=1536, intermediate=8960, 12 heads, 2 KV heads | ~1.5B |
| **Diffusion head** | 4 HeadLayers, hidden=1536, ffn_ratio=3.0 (ffn_dim=4608), latent=64 | ~85-100M |
| **Acoustic connector** | Linear(64, 1536) + RMSNorm + Linear(1536, 1536) | ~4.8M |
| **Semantic connector** | Linear(128, 1536) + RMSNorm + Linear(1536, 1536) | ~4.9M |
| **Acoustic tokenizer** | CNN encoder/decoder (frozen) | ~20-30M |
| **Semantic tokenizer** | CNN encoder (frozen) | ~15-20M |
| **LM head** | Linear(1536, 151936) -- tied to embeddings | ~233M (shared) |

---

## 6. Memory/GPU requirements for full fine-tuning

For full fine-tuning of the 1.5B model (LLM + diffusion head + connectors, ~1.6B trainable params):

**With AdamW optimizer (standard):**
- Model weights (bf16): ~3 GB
- Optimizer states (2x fp32 for Adam): ~12 GB
- Gradients (bf16): ~3 GB
- Activations (varies with batch size and sequence length): ~4-12 GB
- Acoustic/semantic tokenizer (frozen, inference only): ~0.5 GB
- **Total estimate: ~24-35 GB per GPU** at batch_size=1 with gradient checkpointing

**Practical recommendations:**
- **Single GPU**: 1x A100 80GB should work comfortably with `--gradient_checkpointing True`, `--bf16 True`, `--per_device_train_batch_size 1`, and `--gradient_accumulation_steps 4-8`
- **Single GPU**: 1x A6000 48GB is possible but tight -- requires gradient checkpointing and batch size 1
- **Multi-GPU**: 2x A100 40GB with FSDP or DeepSpeed ZeRO Stage 2/3

**If using DeepSpeed ZeRO-3**, memory per GPU drops significantly (sharding optimizer states + gradients + parameters across GPUs), making it feasible on 2-4x RTX 4090 (24GB each).

---

## 7. Recommended approach

The cleanest path for full fine-tuning:

1. **Minimal code change** -- Modify the `requires_grad` logic as described above (5-10 lines changed) so that `--lora_target_modules none` enables full LLM training.

2. **Add a `--train_lm_full` flag** to `ModelArguments` for explicit control, keeping the LoRA path as default.

3. **Update the save logic** to handle full model checkpoints when not using LoRA.

4. **Keep the tokenizers frozen** (`--freeze_acoustic_tokenizer True --freeze_semantic_tokenizer True`) -- these are codec models and should not be fine-tuned.

5. **Run with:**
   ```bash
   --lora_target_modules none \
   --train_diffusion_head True \
   --train_connectors True \
   --bf16 True \
   --gradient_checkpointing True \

Now I have a comprehensive understanding of the codebase. Let me compile my findings.