Moshi-Finetune

To monitor training metrics (loss, learning rate, etc.) in real-time, use Weights & Biases.

1. Install the library: pip install wandb.
2. Configure the wandb: section in your training YAML file with your key and project name.

After training, you can run the Moshi server for interactive inference. First, ensure moshi is installed:

pip install git+https://github.com/kyutai-labs/moshi.git#egg=moshi&subdirectory=moshi

Running LoRA adapters

If you trained using LoRA and set save_adapters: true, use the --lora-weight flag:

python -m moshi.server \
  --lora-weight=$CHECKPOINT_DIR/consolidated/lora.safetensors \
  --config-path=$CHECKPOINT_DIR/consolidated/config.json

Running full fine-tuned models

If you performed full fine-tuning or did not save only adapters, use the --moshi-weight flag:

python -m moshi.server \
--moshi-weight=$CHECKPOINT_DIR/consolidated/consolidated.safetensors \
--config-path=$CHECKPOINT_DIR/consolidated/consolidated/config.json 

To install Moshi-Finetune, clone the repository and install dependencies.

Using uv (Recommended)

If you have uv installed, you do not need to run an explicit install command. Simply prefix your commands with uv run (e.g., uv run torchrun ...) to automatically manage dependencies based on pyproject.toml.

Using pip

If using pip, ensure you have Python 3.10+ and a virtual environment. Install the package in editable mode:

cd moshi-finetune
pip install -e .

git clone git@github.com:kyutai-labs/moshi-finetune.git
cd moshi-finetune
pip install -e .

Training is executed using torchrun. You can run on a single GPU or multiple GPUs.

Single GPU

torchrun --nproc-per-node 1 -m train example/moshi_7B.yaml

Multiple GPUs (e.g., 8 GPUs)

torchrun --nproc-per-node 8 --master_port $RANDOM -m train example/moshi_7B.yaml

Troubleshooting Out-of-Memory (OOM)

If you encounter OOM errors:

1. Reduce batch_size in your config.
2. If the issue persists, lower duration_sec (note: this may impact inference quality/silence behavior).

The training setup is defined via a YAML configuration file. To use an official Moshi model from the Hugging Face Hub, include the moshi_paths section in your config file with the hf_repo_id key.

moshi_paths:
   hf_repo_id: "kyutai/moshiko-pytorch-bf16"

The get_fsdp_policy function determines how the model is partitioned across GPUs using PyTorch's auto_wrap_policy. The strategy changes depending on whether you are doing LoRA or full fine-tuning:

1. Standard Transformer Wrapping: In both modes, each StreamingTransformerLayer is treated as an individual FSDP group to ensure efficient sharding of transformer blocks.
2. LoRA-specific Wrapping: When is_lora=True, an additional lambda_auto_wrap_policy is applied. This is required because FSDP requires that parameters with different requires_grad statuses (trainable LoRA weights vs. frozen base weights) reside in different FSDP groups.

The resulting policy is an _or_policy that combines the transformer block wrapping with the LoRA parameter grouping.

To prevent disk space exhaustion, you can limit the number of saved checkpoints by passing num_ckpt_keep to the Checkpointer constructor.

When save_checkpoint is called, the class automatically identifies the oldest checkpoint directories (sorted by creation time) and deletes them if the total count exceeds num_ckpt_keep.

# Keep only the 5 most recent checkpoints
checkpointer = Checkpointer(
    model=model,
    state=state,
    run_dir="./runs/exp",
    config=config,
    num_ckpt_keep=5
)

# This call will trigger deletion of old checkpoints if necessary
checkpointer.save_checkpoint(save_only_lora=False)

Datasets should be downloaded and organized in a local directory. The training process expects a .jsonl file containing the training data. In the provided example, the kyutai/DailyTalkContiguous dataset is used.

from pathlib import Path
from huggingface_hub import snapshot_download

# Create directory and download dataset
Path("/content/data/daily-talk-contiguous").mkdir(parents=True, exist_ok=True)
local_dir = snapshot_download(
    "kyutai/DailyTalkContiguous",
    repo_type="dataset",
    local_dir="/content/data/daily-talk-contiguous",
)

When performing LoRA fine-tuning, you can save only the adapter weights instead of the full model. This is done by calling save_checkpoint with save_only_lora=True.

Note: You cannot save only LoRA weights if full_finetuning=True was set during Checkpointer initialization. The resulting file will be named lora.safetensors inside the checkpoint directory.

# To save only the LoRA adapters
checkpointer.save_checkpoint(save_only_lora=True)