Spaces:

Sattigeri07
/

Music-project

Running

App Files Files Community

Music-project / ACE-Step-1.5 /docs /en /Large_Scale_SFT_Training_Guide.md

Shubham Sattigeri

Initial commit

116a8d1 about 1 month ago

preview code

Raw

History Blame Contribute Delete

23.7 kB

A newer version of the Gradio SDK is available: 6.22.0

Upgrade

Large-Scale SFT Training Guide with Muse Dataset

This guide covers two training scenarios using the Muse dataset (116K licensed synthetic songs):

Large-Scale LoRA Fine-Tuning — Adapter-only training on the full Muse corpus
Full-Parameter SFT — Training all decoder parameters (requires code modifications)

Prerequisites
Part 1: Dataset Preparation
Part 2: Large-Scale LoRA Training
Part 3: Full-Parameter SFT
Part 4: Multi-GPU Training
Appendix: Architecture Reference

Prerequisites

Hardware Requirements

Scenario	Minimum	Recommended
LoRA (rank 128, multi-GPU)	4x 24 GB (e.g. 4x RTX 4090)	8x 80 GB (e.g. 8x H100/A100)
Full SFT	8x 40 GB (e.g. 8x A100-40G)	8x 80 GB (e.g. 8x H100)
Preprocessing only	1x 12 GB	1x 24 GB

Software Requirements

# Clone ACE-Step 1.5
git clone https://github.com/ace-step/ACE-Step-1.5.git
cd ACE-Step-1.5

# Install dependencies
pip install -r requirements.txt

# Ensure Hugging Face CLI is installed
pip install huggingface_hub[cli]

Storage Requirements

Component	Size
Muse audio (116K songs, MP3)	~500 GB
Preprocessed .pt tensors	~800 GB - 1.2 TB
Training checkpoints	~50 GB (LoRA) / ~200 GB (Full SFT)
Total	~1.5 - 2 TB

Part 1: Dataset Preparation

1.1 Download the Muse Dataset

# Create working directories
mkdir -p muse_dataset/{audio,metadata}
cd muse_dataset

# Download metadata files
huggingface-cli download bolshyC/Muse \
    train_cn.jsonl train_en.jsonl \
    validation_cn.jsonl validation_en.jsonl \
    --local-dir ./metadata --repo-type dataset

# Download Chinese audio archives (25 parts)
for i in $(seq -w 1 25); do
    huggingface-cli download bolshyC/Muse \
        "cn_part${i}_of_025.tar" \
        --local-dir ./audio --repo-type dataset
done

# Download English audio archives (35 parts)
for i in $(seq -w 1 35); do
    huggingface-cli download bolshyC/Muse \
        "en_part${i}_of_035.tar" \
        --local-dir ./audio --repo-type dataset
done

1.2 Extract Audio Files

# IMPORTANT: maintain the directory structure that metadata expects
mkdir -p suno_cn_songs suno_en_songs

# Extract Chinese audio
for f in audio/cn_part*.tar; do
    tar -xf "$f" -C suno_cn_songs/
done

# Extract English audio
for f in audio/en_part*.tar; do
    tar -xf "$f" -C suno_en_songs/
done

1.3 Convert Muse JSONL to ACE-Step Dataset Format

ACE-Step expects a specific directory layout. Create a conversion script:

#!/usr/bin/env python3
"""convert_muse_to_acestep.py -- Convert Muse JSONL to ACE-Step dataset format."""

import json
import os
import shutil
from pathlib import Path
from typing import Optional


def convert_muse_jsonl(
    jsonl_path: str,
    audio_root: str,
    output_dir: str,
    max_samples: Optional[int] = None,
):
    """Convert a Muse JSONL file into ACE-Step's per-song directory layout.

    ACE-Step expects:
        output_dir/
        ├── song_0001.mp3
        ├── song_0001.lyrics.txt
        ├── song_0001.json
        ├── song_0002.mp3
        ...

    Args:
        jsonl_path: Path to train_cn.jsonl or train_en.jsonl
        audio_root: Root directory containing the extracted audio files
        output_dir: Where to write the ACE-Step compatible dataset
        max_samples: Limit number of samples (None = all)
    """
    os.makedirs(output_dir, exist_ok=True)
    count = 0

    with open(jsonl_path, "r", encoding="utf-8") as f:
        for line in f:
            if max_samples and count >= max_samples:
                break

            record = json.loads(line.strip())
            song_id = record["song_id"]
            audio_path = os.path.join(audio_root, record["audio_path"])

            # Skip if audio file doesn't exist
            if not os.path.exists(audio_path):
                print(f"[SKIP] Audio not found: {audio_path}")
                continue

            # Build caption from style field
            style = record.get("style", "")
            caption = style  # e.g. "Pop, Rock, Electronic, Male Vocal, Energetic"

            # Extract lyrics from sections
            lyrics_parts = []
            for section in record.get("sections", []):
                section_type = section.get("section", "")
                text = section.get("text", "").strip()
                if text:
                    lyrics_parts.append(f"[{section_type}]")
                    lyrics_parts.append(text)
                elif section_type == "Intro":
                    lyrics_parts.append("[Intro]")
                elif section_type == "Outro":
                    lyrics_parts.append("[Outro]")
                elif section_type in ("Interlude", "Break"):
                    lyrics_parts.append(f"[{section_type}]")

            lyrics = "\n".join(lyrics_parts)
            if not lyrics.strip():
                lyrics = "[Instrumental]"

            # Detect if instrumental (no lyrics text in any section)
            is_instrumental = all(
                not s.get("text", "").strip()
                for s in record.get("sections", [])
            )

            # Determine BPM from style (not directly in Muse metadata,
            # so we leave it for auto-detection or set a reasonable default)
            # You can use a BPM detection library like librosa if needed.

            # Copy audio file
            ext = Path(audio_path).suffix
            base_name = f"{song_id}_{record.get('track_index', 0)}"
            dest_audio = os.path.join(output_dir, f"{base_name}{ext}")
            if not os.path.exists(dest_audio):
                shutil.copy2(audio_path, dest_audio)

            # Write lyrics file
            lyrics_file = os.path.join(output_dir, f"{base_name}.lyrics.txt")
            with open(lyrics_file, "w", encoding="utf-8") as lf:
                lf.write(lyrics)

            # Write metadata JSON
            meta = {
                "caption": caption,
                "is_instrumental": is_instrumental,
            }

            # Add section-level descriptions as custom_tag for extra context
            descs = [
                s.get("desc", "")
                for s in record.get("sections", [])
                if s.get("desc", "").strip()
            ]
            if descs:
                # Use the first section description as additional context
                meta["custom_tag"] = descs[0][:200]

            meta_file = os.path.join(output_dir, f"{base_name}.json")
            with open(meta_file, "w", encoding="utf-8") as mf:
                json.dump(meta, mf, ensure_ascii=False, indent=2)

            count += 1
            if count % 1000 == 0:
                print(f"[INFO] Converted {count} samples...")

    print(f"[DONE] Converted {count} samples to {output_dir}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--jsonl", required=True, help="Path to Muse JSONL file")
    parser.add_argument("--audio-root", required=True, help="Root dir with extracted audio")
    parser.add_argument("--output", required=True, help="Output directory")
    parser.add_argument("--max-samples", type=int, default=None)
    args = parser.parse_args()

    convert_muse_jsonl(args.jsonl, args.audio_root, args.output, args.max_samples)

Run the conversion:

# Convert Chinese training data
python convert_muse_to_acestep.py \
    --jsonl metadata/train_cn.jsonl \
    --audio-root . \
    --output acestep_dataset_cn

# Convert English training data
python convert_muse_to_acestep.py \
    --jsonl metadata/train_en.jsonl \
    --audio-root . \
    --output acestep_dataset_en

# Merge into a single dataset directory (optional)
mkdir -p acestep_dataset_all
cp -r acestep_dataset_cn/* acestep_dataset_all/
cp -r acestep_dataset_en/* acestep_dataset_all/

1.4 Preprocess Audio to .pt Tensors

ACE-Step's preprocessing converts raw audio into pre-computed VAE latents and text embeddings. This is a two-pass process:

Pass 1 (~3 GB VRAM): VAE encoding + text tokenization → .tmp.pt
Pass 2 (~6 GB VRAM): DIT encoder → final .pt

# Preprocess using the SideStep CLI
# Process in batches to manage disk I/O
python -m acestep.training_v2.cli.train_fixed \
    --preprocess \
    --audio-dir ./acestep_dataset_all \
    --tensor-output ./preprocessed_tensors \
    --checkpoint-dir ./checkpoints \
    --model-variant turbo \
    --max-duration 240 \
    --device cuda:0

# For very large datasets, you can split and run in parallel:
# GPU 0: Chinese subset
python -m acestep.training_v2.cli.train_fixed \
    --preprocess \
    --audio-dir ./acestep_dataset_cn \
    --tensor-output ./preprocessed_tensors_cn \
    --checkpoint-dir ./checkpoints \
    --model-variant turbo \
    --device cuda:0 &

# GPU 1: English subset
python -m acestep.training_v2.cli.train_fixed \
    --preprocess \
    --audio-dir ./acestep_dataset_en \
    --tensor-output ./preprocessed_tensors_en \
    --checkpoint-dir ./checkpoints \
    --model-variant turbo \
    --device cuda:1 &

wait
# Then merge the .pt files
mkdir -p preprocessed_tensors_all
cp preprocessed_tensors_cn/*.pt preprocessed_tensors_all/
cp preprocessed_tensors_en/*.pt preprocessed_tensors_all/

Preprocessing Time Estimate: ~2-5 minutes per song on a single GPU. For 116K songs, expect 4-10 days on 1 GPU, or ~12-24 hours on 8 GPUs in parallel.

Part 2: Large-Scale LoRA Training

2.1 Single-GPU Quick Test

Start with a small subset to verify everything works:

python -m acestep.training_v2.cli.train_fixed \
    --dataset-dir ./preprocessed_tensors_all \
    --output-dir ./output/lora_test \
    --checkpoint-dir ./checkpoints \
    --model-variant turbo \
    --adapter-type lora \
    --r 64 \
    --alpha 128 \
    --batch-size 4 \
    --gradient-accumulation-steps 4 \
    --max-epochs 5 \
    --learning-rate 1e-4 \
    --warmup-steps 100 \
    --device cuda:0 \
    --yes

2.2 Multi-GPU LoRA Training (Recommended)

For 116K songs, multi-GPU training is essential. With PR #941 merged:

# 8x GPU DDP training
python -m acestep.training_v2.cli.train_fixed \
    --dataset-dir ./preprocessed_tensors_all \
    --output-dir ./output/lora_muse_full \
    --checkpoint-dir ./checkpoints \
    --model-variant turbo \
    --adapter-type lora \
    --r 128 \
    --alpha 256 \
    --dropout 0.05 \
    --batch-size 16 \
    --gradient-accumulation-steps 2 \
    --max-epochs 50 \
    --learning-rate 5e-5 \
    --warmup-steps 500 \
    --weight-decay 0.01 \
    --optimizer-type adamw \
    --scheduler-type cosine \
    --gradient-checkpointing \
    --num-devices 8 \
    --strategy ddp \
    --save-every-n-epochs 5 \
    --log-every 10 \
    --yes

2.3 Recommended Hyperparameters for Large-Scale LoRA

Parameter	Small Dataset (<1K)	Medium (1K-10K)	Large (10K-116K)
LoRA rank (`--r`)	8-32	64	128
LoRA alpha (`--alpha`)	2x rank	2x rank	2x rank
Dropout	0.1	0.05	0.05
Batch size (per GPU)	1-4	8-16	16-32
Gradient accumulation	4-8	2-4	1-2
Effective batch size	8-16	32-64	128-256
Learning rate	1e-4	5e-5	2e-5 ~ 5e-5
Warmup steps	50-100	200-500	500-1000
Max epochs	100-500	30-100	10-50
Optimizer	adamw	adamw	adamw / adamw8bit

2.4 SLURM Job Script Example

For HPC clusters:

#!/bin/bash
#SBATCH --job-name=acestep-lora-muse
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=64
#SBATCH --mem=512G
#SBATCH --time=72:00:00
#SBATCH --partition=gpu
#SBATCH --output=logs/%j.out

module load cuda/12.1

cd $SLURM_SUBMIT_DIR

python -m acestep.training_v2.cli.train_fixed \
    --dataset-dir /data/preprocessed_tensors_all \
    --output-dir /data/output/lora_muse_$SLURM_JOBID \
    --checkpoint-dir /data/checkpoints \
    --model-variant turbo \
    --adapter-type lora \
    --r 128 \
    --alpha 256 \
    --batch-size 16 \
    --gradient-accumulation-steps 2 \
    --max-epochs 50 \
    --learning-rate 5e-5 \
    --warmup-steps 500 \
    --num-devices 8 \
    --strategy ddp \
    --save-every-n-epochs 5 \
    --yes

2.5 Monitor Training

# TensorBoard
tensorboard --logdir ./output/lora_muse_full/runs --port 6006

# Key metrics to watch:
# - train/loss: should decrease steadily, target ~0.45-0.50
# - train/lr: should follow cosine schedule
# - Gradient norms: should be stable, no spikes

Part 3: Full-Parameter SFT

Warning: Full-parameter SFT is NOT natively supported by ACE-Step 1.5. This section describes the code modifications required.

3.1 Why Full-Parameter SFT?

Aspect	LoRA	Full SFT
Trainable params	~50M (rank 128)	~3.5B (full decoder)
VRAM per GPU	20-30 GB	60-80 GB
Quality ceiling	Good for style transfer	Maximum quality
Risk of catastrophic forgetting	Low	High
Training time	Hours-days	Days-weeks

Full SFT makes sense when:

You have a very large, diverse dataset (like Muse's 116K songs)
You want to fundamentally change the model's capabilities
You're creating a new base model for downstream LoRA fine-tuning

3.2 Code Modifications Required

Step 1: Modify `model_loader.py` — Add full SFT mode

# In acestep/training_v2/model_loader.py
# After the existing load_decoder_for_training() function, add:

def load_decoder_for_full_sft(
    checkpoint_dir: str | Path,
    variant: str = "turbo",
    device: str = "cuda",
    precision: str = "bf16",
) -> Any:
    """Load the decoder with ALL parameters unfrozen for full SFT.

    Unlike load_decoder_for_training() which freezes everything,
    this keeps all decoder parameters trainable.
    """
    model_dir, dtype = _resolve_model_dir(checkpoint_dir, variant, precision)

    model = AutoModel.from_pretrained(
        str(model_dir),
        trust_remote_code=True,
        attn_implementation="sdpa",
        dtype=dtype,
    )

    # Do NOT freeze parameters — keep everything trainable
    # Only freeze non-decoder components (encoder, VAE) if desired
    for name, param in model.named_parameters():
        if "decoder" in name:
            param.requires_grad = True
        else:
            param.requires_grad = False

    model = model.to(device=device, dtype=dtype)
    model.train()

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"[OK] Full SFT mode: {trainable:,} / {total:,} params trainable")

    return model

Step 2: Add `--full-sft` CLI flag

# In acestep/training_v2/cli/args.py
# In _add_common_training_args(), add:

g.add_argument(
    "--full-sft",
    action="store_true",
    default=False,
    help="Train all decoder parameters instead of just LoRA/LoKR adapters",
)

Step 3: Modify `configs.py`

# In acestep/training_v2/configs.py, in TrainingConfigV2:

full_sft: bool = False
"""Train all decoder parameters (no adapter injection)."""

Step 4: Modify `trainer_fixed.py` — Skip LoRA injection in full SFT mode

The key change is in the train() method. When full_sft=True:

Skip inject_lora_into_dit() / inject_lokr_into_dit()
Keep all decoder parameters trainable
Save full model checkpoints instead of adapter-only checkpoints

# In FixedLoRATrainer.train(), after model loading:

if cfg.full_sft:
    # Full SFT: unfreeze all decoder parameters
    for param in self.module.model.decoder.parameters():
        param.requires_grad = True
    self.module.model.decoder.train()
else:
    # Adapter mode: inject LoRA/LoKR as before
    inject_lora_into_dit(self.module.model.decoder, adapter_cfg)

And modify _save_final() / _save_checkpoint():

def _save_final(self, path: str):
    if self.training_config.full_sft:
        # Save full decoder state dict
        os.makedirs(path, exist_ok=True)
        torch.save(
            self.module.model.decoder.state_dict(),
            os.path.join(path, "decoder_state_dict.pt"),
        )
        # Also save in HuggingFace format for easy loading
        self.module.model.save_pretrained(path)
    else:
        # Existing adapter save logic
        save_lora_weights(...)

Step 5: Use DeepSpeed ZeRO for memory efficiency (recommended)

For full SFT with 3.5B decoder parameters, single-GPU training is impractical. Replace Lightning Fabric with DeepSpeed ZeRO-2 or ZeRO-3:

# In trainer_fixed.py, modify _train_fabric():

if cfg.full_sft and num_devices > 1:
    # Use DeepSpeed ZeRO-2 for full SFT
    from lightning.fabric.strategies import DeepSpeedStrategy

    ds_config = {
        "zero_optimization": {
            "stage": 2,
            "offload_optimizer": {
                "device": "cpu",
                "pin_memory": True,
            },
            "allgather_partitions": True,
            "allgather_bucket_size": 2e8,
            "reduce_scatter": True,
            "reduce_bucket_size": 2e8,
            "overlap_comm": True,
        },
        "bf16": {"enabled": True},
        "gradient_clipping": cfg.max_grad_norm,
    }

    self.fabric = Fabric(
        accelerator="cuda",
        devices=num_devices,
        strategy=DeepSpeedStrategy(config=ds_config),
        precision="bf16-true",
    )

3.3 Full SFT Training Command

After applying the code modifications:

# 8x H100 full SFT
python -m acestep.training_v2.cli.train_fixed \
    --dataset-dir ./preprocessed_tensors_all \
    --output-dir ./output/full_sft_muse \
    --checkpoint-dir ./checkpoints \
    --model-variant turbo \
    --full-sft \
    --batch-size 8 \
    --gradient-accumulation-steps 4 \
    --max-epochs 20 \
    --learning-rate 1e-5 \
    --warmup-steps 2000 \
    --weight-decay 0.01 \
    --optimizer-type adamw \
    --scheduler-type cosine \
    --gradient-checkpointing \
    --num-devices 8 \
    --strategy ddp \
    --save-every-n-epochs 2 \
    --yes

3.4 Recommended Hyperparameters for Full SFT

Parameter	Conservative	Balanced	Aggressive
Learning rate	5e-6	1e-5	2e-5
Warmup steps	3000	2000	1000
Weight decay	0.01	0.01	0.05
Batch size (per GPU)	4	8	16
Gradient accumulation	8	4	2
Effective batch size	256	256	256
Max epochs	10-20	20-30	30-50
Max grad norm	1.0	1.0	0.5

Important: Use a lower learning rate for full SFT (1e-5) compared to LoRA (5e-5) to avoid catastrophic forgetting. The warmup should be longer.

Part 4: Multi-GPU Training

4.1 DDP (Data Distributed Parallel)

DDP replicates the model on each GPU and splits the dataset. Each GPU processes different batches, and gradients are synchronized via all-reduce.

# Simple DDP (single node, multiple GPUs)
python -m acestep.training_v2.cli.train_fixed \
    --num-devices 8 \
    --strategy ddp \
    ... # other args

4.2 FSDP / DeepSpeed ZeRO (for Full SFT)

When the model doesn't fit on a single GPU, use model parallelism:

# DeepSpeed ZeRO-2 (optimizer states sharded across GPUs)
# Requires the code modifications from Part 3

# DeepSpeed ZeRO-3 (model parameters + optimizer states sharded)
# Maximum memory savings, but slower communication

4.3 Multi-Node Training

For training across multiple machines:

# Node 0 (master)
export MASTER_ADDR=node0-ip
export MASTER_PORT=29500
export WORLD_SIZE=16
export NODE_RANK=0
python -m acestep.training_v2.cli.train_fixed \
    --num-devices 8 \
    --strategy ddp \
    ... # other args

# Node 1
export MASTER_ADDR=node0-ip
export MASTER_PORT=29500
export WORLD_SIZE=16
export NODE_RANK=1
python -m acestep.training_v2.cli.train_fixed \
    --num-devices 8 \
    --strategy ddp \
    ... # other args

Appendix: Architecture Reference

ACE-Step Training Pipeline

Raw Audio (MP3/WAV)
    │
    ▼
[Preprocessing - Pass 1: ~3 GB VRAM]
    ├─ Audio → 48kHz stereo → VAE encoder → target_latents [T, 64]
    ├─ Caption → Text tokenizer + encoder → text_hidden_states [L, 768]
    └─ Lyrics → Lyric tokenizer + encoder → lyric_hidden_states [L, 768]
    │
    ▼
[Preprocessing - Pass 2: ~6 GB VRAM]
    ├─ text_hs + lyric_hs → DIT encoder → encoder_hidden_states [L, 768]
    └─ silence_latent + chunk_mask → context_latents [T, 65]
    │
    ▼
[.pt Tensor File]
    ├─ target_latents: [T, 64]         # ground truth audio latents
    ├─ attention_mask: [T]              # valid audio positions
    ├─ encoder_hidden_states: [L, 768]  # packed text + lyrics embeddings
    ├─ encoder_attention_mask: [L]      # valid condition positions
    ├─ context_latents: [T, 65]         # context audio + chunk mask
    └─ metadata: {...}                  # song info
    │
    ▼
[Training Loop - Flow Matching]
    1. Sample noise x1 ~ N(0, 1) with shape [B, T, 64]
    2. Sample timestep t ~ logit-normal(μ=-0.4, σ=1.0)
    3. Interpolate: xt = t * x1 + (1-t) * x0    (x0 = target_latents)
    4. CFG dropout: replace 15% of conditions with null embedding
    5. Decoder forward: predict_flow = decoder(xt, t, conditions)
    6. Loss = MSE(predict_flow, x1 - x0)         (flow matching loss)
    7. Backprop through LoRA params (or all params in full SFT)

Key Model Dimensions

Component	Dimension
VAE latent	64
Context latent (with mask)	65
Decoder input (concat)	129
Hidden size	768 (after patch projection)
Decoder layers	32
Attention heads	32
Patch size	4 (reduces sequence length by 4x)
Audio sample rate	48,000 Hz
Latent frame rate	~160 Hz (before patching)

Muse Dataset → ACE-Step Mapping

Muse Field	ACE-Step Usage
`style`	`caption` (text conditioning)
`sections[].text`	Lyrics file (`.lyrics.txt`)
`sections[].section`	Structural tags (`[Verse]`, `[Chorus]`, etc.)
`sections[].desc`	`custom_tag` in metadata JSON
`audio_path`	Source audio for VAE encoding
`style_sim`	Can be used to filter low-quality samples

Tips for Best Results

Filter by style_sim: The Muse dataset includes a style_sim score. Consider filtering out samples with style_sim < 0.4 for higher-quality training data.
Balance CN/EN: If you want bilingual capability, ensure roughly equal representation of Chinese and English samples.
Progressive training: Start with a smaller subset (e.g., 10K songs) to validate the pipeline, then scale up.
Monitor for mode collapse: If loss plateaus too early or generated samples lack diversity, reduce learning rate or increase dropout.
Checkpoint frequently: Save checkpoints every 2-5 epochs. Large-scale training can fail due to hardware issues.
Validate periodically: Use the validation split (validation_cn.jsonl, validation_en.jsonl) to monitor generalization. Generate sample audio every N epochs to subjectively evaluate quality.