# How to Handle Out-of-Memory Errors During Training in Fish-Speech

> Learn to handle out-of-memory errors in Fish-Speech training. Use mixed-precision, adjust batch size and max length, and enable gradient checkpointing to optimize memory usage.

- Repository: [Fish Audio/fish-speech](https://github.com/fishaudio/fish-speech)
- Tags: how-to-guide
- Published: 2026-03-12

---

**Enable mixed-precision training with `bf16-mixed`, reduce `batch_size` and `max_length` in your Hydra config, and ensure `use_gradient_checkpointing` remains enabled to trade compute for memory.**

Training Fish-Speech models on large audio datasets can quickly exhaust GPU memory when processing long sequences or using high batch sizes. To handle out-of-memory errors during training, the repository provides several built-in mechanisms to manage memory consumption without sacrificing model quality. By adjusting configuration parameters in the Hydra YAML files and leveraging PyTorch Lightning's optimization features, you can systematically prevent CUDA OOM crashes.

## Use Mixed-Precision Training

Fish-Speech supports **bfloat16 mixed-precision training** through PyTorch Lightning's precision settings, which halves the memory footprint of activations and weights while maintaining numerical stability.

In [`fish_speech/configs/base.yaml`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/configs/base.yaml), the precision is controlled via the trainer configuration at line 22:

```yaml
trainer:
  precision: bf16-mixed  # or bf16-true for full bfloat16

```

Setting `precision: bf16-mixed` keeps the model weights in full precision while casting operations to bfloat16, whereas `bf16-true` casts everything to bfloat16 for maximum memory savings.

## Enable Gradient Checkpointing

The model implements **gradient checkpointing** to reduce activation memory at the cost of additional computation during the backward pass. This feature is enabled by default in the model definition.

In [`fish_speech/models/text2semantic/llama.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py), the `BaseModelArgs` dataclass exposes the `use_gradient_checkpointing` parameter at lines 54-56:

```python
@dataclass
class BaseModelArgs:
    use_gradient_checkpointing: bool = True
    # ... other args

```

Keep this set to `true` unless you have abundant GPU memory, as it significantly lowers peak memory usage during training by recomputing intermediate activations instead of storing them.

## Reduce Batch Size and Sequence Length

The most direct way to handle OOM errors is to reduce the **batch size** and **maximum sequence length**, which directly lowers activation memory.

These parameters are exposed in the Hydra configuration files. In [`fish_speech/configs/text2semantic_finetune.yaml`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/configs/text2semantic_finetune.yaml) at lines 52-56, you can adjust:

```yaml
data:
  batch_size: 2  # Reduce from 4 to 2 or 1

  max_length: 2048  # Reduce from 4096 to 2048 or lower

```

Lowering `max_length` truncates audio sequences earlier in the pipeline, while reducing `batch_size` decreases the number of samples processed in parallel. Both changes immediately reduce the memory allocated per training step.

## Monitor Memory and Handle Crashes Safely

Fish-Speech includes utilities to monitor GPU usage and gracefully handle OOM exceptions.

The training script logs maximum reserved GPU memory after generation steps in [`fish_speech/models/text2semantic/inference.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/inference.py) at line 730, helping you identify memory leaks or gradual growth:

```python
logger.info(f"GPU Memory used: {torch.cuda.max_memory_reserved() / 1e9:.2f} GB")

```

Additionally, the `task_wrapper` decorator in [`fish_speech/utils/utils.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/utils/utils.py) at line 77 catches unexpected exceptions including OOM errors, records them, and prevents silent failures:

```python
@task_wrapper
def train(cfg: DictConfig) -> None:
    # Training logic that catches OOM and other exceptions

```

While the inference engine in [`fish_speech/inference_engine/__init__.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/inference_engine/__init__.py) at line 121 demonstrates explicit cache clearing with `torch.cuda.empty_cache()` after heavy inference steps, showing the repository's built-in awareness of memory pressure management.

## Practical Configuration Examples

### Minimal Low-Memory Config

Create a custom Hydra configuration file at [`fish_speech/configs/low_mem.yaml`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/configs/low_mem.yaml):

```yaml
defaults:
  - base
  - _self_

project: low_mem_train
max_length: 2048

trainer:
  precision: bf16-mixed
  max_steps: 5000

data:
  batch_size: 1
  num_workers: 2
  max_length: ${max_length}

model:
  _target_: fish_speech.models.text2semantic.lit_module.TextToSemantic
  model:
    _target_: fish_speech.models.text2semantic.llama.BaseTransformer.from_pretrained
    path: checkpoints/openaudio-s1-mini
    load_weights: true
    max_length: ${max_length}
    # Gradient checkpointing remains enabled by default

```

Run with:

```bash
python -m fish_speech.train hydra.run.dir=./runs low_mem.yaml

```

### Programmatic Adjustment

For Jupyter notebooks or Python scripts, modify configurations programmatically:

```python
from omegaconf import OmegaConf
import fish_speech.train
from fish_speech.utils import utils as fish_utils

# Load and modify config

cfg = OmegaConf.load("fish_speech/configs/text2semantic_finetune.yaml")
cfg.trainer.precision = "bf16-mixed"
cfg.data.batch_size = 1
cfg.max_length = 2048

# Run with safety wrapper

fish_utils.task_wrapper(fish_speech.train)(cfg)

```

### Custom Memory Monitoring Callback

Add a Lightning callback to track GPU usage during training:

```python
from lightning import Callback
import torch

class MemoryLogger(Callback):
    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        if torch.cuda.is_available():
            used = torch.cuda.max_memory_reserved() / 1e9
            pl_module.log("gpu_mem_gb", used, prog_bar=True)
            print(f"GPU Memory: {used:.2f} GB")

```

Register this in your Hydra config under the `callbacks` section to monitor memory in real-time.

## Summary

- **Mixed-precision training**: Set `precision: bf16-mixed` or `bf16-true` in [`fish_speech/configs/base.yaml`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/configs/base.yaml) to reduce tensor memory usage.
- **Gradient checkpointing**: Keep `use_gradient_checkpointing: true` in [`fish_speech/models/text2semantic/llama.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py) to trade compute for activation memory.
- **Batch and sequence reduction**: Lower `data.batch_size` and `max_length` in your finetuning config to directly cut memory allocation.
- **Safety mechanisms**: The `task_wrapper` in [`fish_speech/utils/utils.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/utils/utils.py) catches OOM errors gracefully, while built-in logging tracks GPU usage.

## Frequently Asked Questions

### What is the first setting I should change when I encounter an OOM error?

**Reduce the batch size first.** In [`fish_speech/configs/text2semantic_finetune.yaml`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/configs/text2semantic_finetune.yaml), change `data.batch_size` from the default (often 4) down to 2 or 1. If OOM persists, reduce `max_length` from 4096 to 2048. These changes have immediate impact without requiring code modifications.

### Does enabling bf16-mixed affect model quality?

**No, bfloat16 mixed-precision maintains full model quality** for most training scenarios. The `bf16-mixed` setting in the trainer config keeps master weights in full precision while using bfloat16 for computations, preventing the gradient underflow issues sometimes seen with fp16. Only use `bf16-true` if you understand the trade-offs with numerical stability.

### Where is the gradient checkpointing setting located?

**In [`fish_speech/models/text2semantic/llama.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py) within the `BaseModelArgs` dataclass (lines 54-56).** The parameter `use_gradient_checkpointing` defaults to `True`, which is optimal for memory-constrained training. Set it to `False` only if you have excess GPU memory and want to speed up training by avoiding recomputation overhead.

### How can I monitor GPU memory usage during training?

**Check the logs for "GPU Memory used" messages** printed by [`fish_speech/models/text2semantic/inference.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/inference.py), or implement a custom Lightning callback that logs `torch.cuda.max_memory_reserved()`. The repository also provides the `task_wrapper` decorator in [`fish_speech/utils/utils.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/utils/utils.py) that captures OOM exceptions and records them before exiting, preventing silent failures in automated training pipelines.