Implementation:Facebookresearch Audiocraft Export lm

Overview

Export_lm extracts the best model state from a full MusicGen or AudioGen training checkpoint and saves it as a lightweight export file suitable for inference and distribution.

Implements

Principle:Facebookresearch_Audiocraft_Language_Model_Export

API Signature

def export_lm(
    checkpoint_path: Union[Path, str],
    out_file: Union[Path, str]
) -> Union[Path, str]

Source: audiocraft/utils/export.py, lines 61-79

Import:

from audiocraft.utils.export import export_lm

Parameters

Parameter	Type	Default	Description
`checkpoint_path`	`Union[Path, str]`	required	Path to the full training checkpoint file (e.g., `checkpoint_best.th`) produced by MusicGenSolver or AudioGenSolver
`out_file`	`Union[Path, str]`	required	Destination path for the exported lightweight checkpoint file

Return Value

Returns the out_file path as provided, after successfully writing the exported checkpoint.

Dependencies

torch -- for loading and saving checkpoint files
omegaconf.OmegaConf -- for serializing the experiment config to YAML
audiocraft.__version__ -- for embedding the version tag in the export

How It Works

from audiocraft.utils.export import export_lm

# Export a training checkpoint to deployment format
export_lm(
    checkpoint_path='/path/to/xps/abc123/checkpoint_best.th',
    out_file='/path/to/exports/my_model/state_dict.bin'
)

The implementation:

def export_lm(checkpoint_path, out_file):
    pkg = torch.load(checkpoint_path, 'cpu')
    if pkg['fsdp_best_state']:
        best_state = pkg['fsdp_best_state']['model']
    else:
        assert pkg['best_state']
        best_state = pkg['best_state']['model']
    new_pkg = {
        'best_state': best_state,
        'xp.cfg': OmegaConf.to_yaml(pkg['xp.cfg']),
        'version': __version__,
        'exported': True,
    }
    Path(out_file).parent.mkdir(exist_ok=True, parents=True)
    torch.save(new_pkg, out_file)
    return out_file

Output Format

The exported checkpoint is a torch.save-serialized dictionary with the following keys:

Key	Type	Description
`best_state`	`dict`	Flattened model state dictionary (without the outer `model` nesting)
`xp.cfg`	`str`	YAML-serialized OmegaConf configuration containing all model architecture and dataset parameters
`version`	`str`	AudioCraft version string (e.g., `"1.3.0"`)
`exported`	`bool`	Always `True`, used by loaders to identify exported vs. training checkpoints

FSDP Priority Logic

The function prioritizes fsdp_best_state over best_state. This is because in FSDP training, the regular best_state contains sharded parameters that are incomplete on any single rank, while fsdp_best_state is gathered to contain the full model parameters. The truthiness check (if pkg['fsdp_best_state']) evaluates to False when the key exists but contains an empty dict or None, correctly falling back to the standard best_state for non-FSDP training runs.

Related Implementations

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment