Gejala Masalah: Worker Terminated saat Step Training Audio

Pada pipeline pemrosesan model audio berbasis difusi (seperti DiffWave atau AudioLDM), worker queue sering berhenti mendadak akibat eksepsi torch.cuda.OutOfMemoryError. Pada kartu grafis kelas entry-level atau consumer dengan alokasi VRAM 6GB (misalnya NVIDIA RTX 2060/3060 6GB), proses komputasi audio sangat rentan terhadap lonjakan kebutuhan memori secara mendadak.

Gejala umum yang muncul pada log worker:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 MiB (GPU 0; 5.79 GiB total capacity; 4.82 GiB already allocated; 112.12 MiB free; 5.10 GiB reserved in total by PyTorch)

Masalah ini umumnya memuncak saat fase backpropagation ketika batch spectrogram atau raw waveform dengan durasi panjang diproses secara bersamaan.

Analisis Root Cause: Mengapa VRAM 6GB Habis

Terdapat tiga faktor utama yang memicu OOM pada skenario worker audio:

  1. Akumulasi Computation Graph (Autograd Leak): Penyimpanan variabel loss atau output langsung ke dalam array log (misalnya history.append(loss) alih-alih history.append(loss.item())) menahan seluruh directed acyclic graph (DAG) autograd di memori GPU sepanjang lifecycle worker.
  2. Overhead Memory Activations pada Tensor Audio: Representasi audio berdimensi tinggi menghasilkan aktivasi per-layer yang sangat besar. Menyimpan intermediate activation untuk setiap layer konvolusional dan attention hingga step loss.backward() melampaui batas fisik 6GB.
  3. Fragmentasi PyTorch Caching Allocator: PyTorch tidak langsung mengembalikan memori GPU ke OS saat tensor dihapus, melainkan menyimpannya dalam memory pool. Karena data audio sering bervariasi dalam panjang sekuens (variable-length audio), alokator mengalami fragmentasi internal: kapasitas total bebas mencukupi, tetapi tidak tersedia blok memori kontinu yang cukup besar.

Langkah Remediasi dan Proteksi Memori

1. Konfigurasi Allocator via Environment Variable

Gunakan fitur alokator PyTorch untuk mencegah fragmentasi blok memori virtual tanpa perlu membatasi throughput worker:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Parameter expandable_segments:True memungkinkan PyTorch menggunakan virtual memory mapping CUDA untuk memperluas segmen memori yang sudah dialokasikan alih-alih gagal mencari blok kontinu baru.

2. Mixed-Precision Training (torch.amp.autocast)

Konversikan aktivasi dan weight forward pass ke FP16/BF16 menggunakan Automatic Mixed Precision (AMP). Pendekatan ini memangkas konsumsi VRAM untuk aktivasi layer hingga 50% sekaligus mempercepat eksekusi Tensor Core.

3. Gradient Checkpointing

Terapkan torch.utils.checkpoint pada blok U-Net atau Transformer backbone. Gradient checkpointing menukar waktu komputasi ekstra (menghitung ulang aktivasi saat backward pass) dengan penghematan VRAM signifikan, karena intermediate activation tidak lagi disimpan terus-menerus selama forward pass.

4. Manajemen State Worker Eksplisit

Pastikan metrik logging selalu di-detach dari computation graph, dan jalankan pembersihan cache pada interval batch tertentu atau ketika exception tertangkap.

Implementasi Kode: Worker Loop Audio Terproteksi

Berikut implementasi minimal loop pemrosesan batch audio pada worker PyTorch:

import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

import gc
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

class AudioDiffusionBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.conv = nn.Conv1d(dim, dim, kernel_size=3, padding=1)
        self.act = nn.SiLU()

    def forward(self, x):
        # ponytail: implementasi langsung checkpointing pada block model
        return checkpoint(self._inner_forward, x, use_reentrant=False)

    def _inner_forward(self, x):
        return self.act(self.conv(x))


def process_batch(model, optimizer, scaler, batch, device):
    model.train()
    optimizer.zero_grad(set_to_none=True)  # Hemat alokasi memori gradien

    audio_inputs, targets = batch[0].to(device, non_blocking=True), batch[1].to(device, non_blocking=True)

    # Gunakan autocast untuk mixed-precision FP16
    with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(audio_inputs)
        loss = nn.functional.mse_loss(outputs, targets)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

    # Hindari memory leak: selalu detach dan ekstrak float murni
    loss_metric = loss.item()
    del audio_inputs, targets, outputs, loss

    return loss_metric


def cleanup_worker_memory():
    gc.collect()
    torch.cuda.empty_cache()

Verifikasi dan Monitoring Metrik VRAM

Pantau metrik penggunaan memori GPU secara programatik pada loop worker untuk membedakan antara memori yang dialokasikan (sedang dipakai tensor) dan memori yang direservasi oleh caching allocator:

def log_cuda_metrics(step: int):
    allocated_mb = torch.cuda.memory_allocated() / (1024 ** 2)
    reserved_mb = torch.cuda.memory_reserved() / (1024 ** 2)
    max_allocated_mb = torch.cuda.max_memory_allocated() / (1024 ** 2)

    print(
        f"[Step {step}] Allocated: {allocated_mb:.2f} MiB | "
        f"Reserved: {reserved_mb:.2f} MiB | "
        f"Peak: {max_allocated_mb:.2f} MiB"
    )

Aturan verifikasi:

  • Jika Allocated stabil namun Reserved terus mendekati 6000 MiB, fragmentasi sedang terjadi (aktifkan expandable_segments).
  • Jika Allocated meningkat secara linear setiap step, terjadi leak autograd graph (periksa apakah ada tensor yang disimpan tanpa .item() atau .detach()).
  • Jika Peak melonjak tajam hanya pada input berdurasi panjang, batasi panjang audio maksimum (audio slicing / dynamic chunking) sebelum masuk ke model.