# How the TextToSemantic Lightning Module Handles Training in Fish Speech

> Learn how the TextToSemantic Lightning module trains models. Discover its approach to optimizer configuration, loss computation, and LoRA checkpointing for efficient speech model development.

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

---

**The TextToSemantic Lightning module orchestrates training by wrapping a `NaiveTransformer` in a PyTorch Lightning interface, handling optimizer configuration with parameter grouping, computing combined token and codebook losses, and managing LoRA-specific checkpointing.**

The `TextToSemantic` module in the `fishaudio/fish-speech` repository provides the training infrastructure for converting text into semantic tokens. Located in [`fish_speech/models/text2semantic/lit_module.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/lit_module.py), this PyTorch Lightning module coordinates the complex interaction between text embeddings, transformer forward passes, and multi-codebook semantic outputs while abstracting away boilerplate training logic.

## Architecture Overview

The `TextToSemantic` class acts as a lightweight wrapper around the core `NaiveTransformer` model defined in [`fish_speech/models/text2semantic/llama.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/llama.py). Instead of implementing raw training loops, the module leverages PyTorch Lightning's structured hooks for initialization, optimization, and checkpointing. This separation allows the transformer to focus on architecture while the Lightning module handles training semantics, loss computation, and hardware abstraction.

## Core Training Components

### Module Initialization and Forward Pass

The constructor receives three key components: the `NaiveTransformer` instance, a callable optimizer builder, and a learning-rate scheduler builder. These are stored for later use in `configure_optimizers`. The `forward` method remains minimal, delegating directly to the underlying transformer:

```python
def forward(self, x):
    return self.model(x)

```

This delegation pattern ensures that the Lightning module maintains the same interface as the base model while adding training-specific orchestration.

### Optimizer Configuration with Weight Decay Grouping

The `configure_optimizers` method in [`fish_speech/models/text2semantic/lit_module.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/lit_module.py) implements parameter-sensitive optimization. Biases, LayerNorm weights, and embeddings are identified and excluded from weight decay, while transformer weights receive regularization. This creates two distinct parameter groups:

- **Group 1**: Weights with decay (standard transformer parameters)
- **Group 2**: Weights without decay (biases, norms, embeddings)

The method then instantiates the optimizer using the provided builder function and attaches the learning-rate scheduler. This configuration is critical for stable training of large language models, as improper weight decay on bias terms can lead to underfitting.

### LoRA Checkpoint Management

For efficient fine-tuning, the module implements `on_save_checkpoint` to support LoRA (Low-Rank Adaptation) workflows. When saving checkpoints, the method filters the state dictionary to retain only LoRA-specific parameters, stripping out the full base model weights. This results in dramatically smaller checkpoint files suitable for storage and transfer, while the base model can be reloaded separately during inference.

## The Training Step Pipeline

### Batch Processing and Loss Computation

The core training logic resides in the `_step` method, which handles both training and validation phases. This method orchestrates a multi-part loss calculation that combines traditional language modeling with semantic codebook prediction:

1. **Mode Management**: Sets the model to training mode explicitly during the "train" stage to ensure LoRA updates apply correctly.
2. **Forward Pass**: Calls the transformer with `inputs`, `attention_masks`, and `labels` to obtain dual outputs: token logits and codebook logits.
3. **Base Loss Calculation**: Computes cross-entropy loss on the primary token prediction (the first element of the output).

### Semantic Loss Masking and Codebook Handling

The semantic component requires careful label masking to isolate valid semantic tokens from the text sequence:

- **Region Identification**: Identifies semantic tokens using `semantic_begin_id` and `semantic_end_id` boundaries.
- **Label Extraction**: Selects matching codebook labels and permutes them from `[batch, seq, codebook]` to `[batch, codebook, seq]` to align with logits.
- **Cross-Entropy Application**: Computes semantic loss across all codebooks using cross-entropy, treating each codebook as a separate classification task.

The final loss is the sum of the base token loss and the semantic codebook loss, providing a unified optimization target for both text understanding and semantic representation learning.

### Accuracy Metrics and Logging

The `get_accuracy` method computes top-5 accuracy on codebook predictions, masking out padding tokens (marked as `-100` or `CODEBOOK_PAD_TOKEN_ID`). The `_step` method logs these metrics along with the component losses (base, semantic, and total) to the Lightning logger, enabling real-time monitoring of both classification accuracy and semantic reconstruction quality.

The `training_step` and `validation_step` hooks simply delegate to `_step` with the appropriate stage identifier, allowing Lightning to handle gradient backpropagation, metric aggregation, and validation loop management automatically.

## Integration with the Training Script

The generic training entry point in [`fish_speech/train.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/train.py) orchestrates the complete pipeline. It builds the data module (typically using `SemanticDataset` from [`fish_speech/datasets/semantic.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/datasets/semantic.py)), instantiates the `TextToSemantic` model via Hydra configuration, and attaches callbacks, loggers, and the Lightning `Trainer`.

When `trainer.fit()` is invoked, it triggers the entire flow: optimization configuration, checkpoint saving with LoRA filtering, and the iterative training loop with dual-loss computation. The Hydra configuration in [`configs/llama_pretrain.yaml`](https://github.com/fishaudio/fish-speech/blob/main/configs/llama_pretrain.yaml) wires the `TextToSemantic` module to the data, optimizer, and scheduler without requiring code changes.

## Code Examples

### Minimal Instantiation and Training with Lightning

```python
import lightning as L
import torch
from fish_speech.models.text2semantic.lit_module import TextToSemantic
from fish_speech.models.text2semantic.llama import NaiveTransformer

# Build the underlying transformer

transformer = NaiveTransformer(
    vocab_size=51200,
    num_codebooks=8,
    d_model=768,
    n_head=12,
    n_layer=12,
    max_seq_len=1024,
)

# Optimizer and scheduler factories

def make_optimizer(params):
    return torch.optim.AdamW(params, lr=1e-4, weight_decay=0.01)

def make_scheduler(optimizer):
    return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=1000)

# Wrap with Lightning

model = TextToSemantic(
    model=transformer,
    optimizer=make_optimizer,
    lr_scheduler=make_scheduler,
)

# Dummy data loader for demonstration

def dummy_loader():
    while True:
        batch = {
            "inputs": torch.randint(0, 51200, (4, 1024)),
            "attention_masks": torch.ones(4, 1024, dtype=torch.bool),
            "labels": torch.randint(0, 51200, (4, 1 + transformer.config.num_codebooks)),
        }
        yield batch

loader = torch.utils.data.DataLoader(list(dummy_loader()), batch_size=4, num_workers=0)

# Trainer

trainer = L.Trainer(
    max_epochs=5,
    accelerator="cpu",
    logger=False,
    enable_checkpointing=False,
)
trainer.fit(model, train_dataloaders=loader)

```

### Using the Full Training Script via Hydra

```bash

# Install dependencies

pip install -r requirements.txt

# Run training with default configuration

python -m fish_speech.train \
    data=semantic \
    model=text2semantic \
    trainer.max_epochs=10 \
    optimizer=adamw \
    lr_scheduler=cosine

```

This command leverages the Hydra configuration in [`configs/llama_pretrain.yaml`](https://github.com/fishaudio/fish-speech/blob/main/configs/llama_pretrain.yaml) to instantiate the `TextToSemantic` module with the `SemanticDataset`, AdamW optimizer, and cosine learning rate schedule.

## Summary

- The `TextToSemantic` Lightning module in [`fish_speech/models/text2semantic/lit_module.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/text2semantic/lit_module.py) wraps the `NaiveTransformer` to handle text-to-semantic training.
- **Dual-loss computation** combines cross-entropy for text tokens with masked cross-entropy for semantic codebook tokens.
- **Smart optimizer configuration** separates parameters into weight-decay and no-decay groups (biases, LayerNorm, embeddings) for stable training.
- **LoRA support** via `on_save_checkpoint` strips non-LoRA weights to produce lightweight checkpoint files.
- **Hydra integration** allows full configuration through YAML files without code modification, orchestrated by [`fish_speech/train.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/train.py).

## Frequently Asked Questions

### How does the TextToSemantic module handle the dual loss for text and semantic tokens?

The `_step` method computes a **base loss** using standard cross-entropy on the primary text token predictions, then extracts semantic tokens using `semantic_begin_id` and `semantic_end_id` boundaries to compute a **semantic loss** across all codebooks. These losses are summed into a single optimization target that updates both the linguistic and semantic representations simultaneously.

### What is the purpose of the weight decay grouping in `configure_optimizers`?

The `configure_optimizers` method separates parameters into two groups: one with weight decay for standard transformer weights, and one without decay for biases, LayerNorm parameters, and embeddings. This follows best practices for training transformer models, preventing the regularization of bias terms that should freely adjust to data shifts, thereby improving convergence stability.

### How does the module support LoRA fine-tuning checkpoints?

The `on_save_checkpoint` hook filters the model state dictionary to retain only LoRA-specific parameters (identified by name patterns or adapters), discarding the full base model weights. This produces checkpoint files that are orders of magnitude smaller, suitable for distribution and storage, while the base model can be reloaded separately during inference to reconstruct the full state.

### Where does the training loop actually start in the fish-speech repository?

The entry point is [`fish_speech/train.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/train.py), which uses Hydra to instantiate the `TextToSemantic` module, data loaders (typically from [`fish_speech/datasets/semantic.py`](https://github.com/fishaudio/fish-speech/blob/main/fish_speech/datasets/semantic.py)), and the Lightning `Trainer`. When `trainer.fit()` is called, it triggers the `configure_optimizers` setup and iterates through `training_step` calls, which delegate to the `_step` method for loss computation and logging.