MLflow Logging in NeMo AutoModel

Introduction

MLflow is an open-source platform for managing the machine learning lifecycle, including experiment tracking, model versioning, and deployment. NeMo AutoModel integrates with MLflow to log training metrics, parameters, and artifacts during model training.

With MLflow integration, you can:

Track and compare experiments across multiple runs
Log hyperparameters and training configurations
Monitor training and validation metrics in real-time
Store model checkpoints and artifacts
Visualize experiment results through the MLflow UI
Share results with team members

Prerequisites

Before using MLflow logging in NeMo AutoModel, ensure you have:

MLflow installed: MLflow is installed with nemo-automodel by default. If you see an import error in a released-package environment, recreate that environment and install NeMo AutoModel there:
```
$ uv venv
$ source .venv/bin/activate
$ uv pip install "nemo-automodel"
```
MLflow tracking server (optional): For production use, set up a tracking server to centralize experiment data. For local development, current MLflow releases use a local SQLite database at sqlite:///mlflow.db by default.

Configuration

Enable MLflow logging by adding an mlflow section to your recipe YAML configuration:

1 mlflow:
2   experiment_name: "automodel-llm-llama3_2_1b_squad-finetune"
3   run_name: ""
4   tracking_uri: null
5   artifact_location: null
6   tags:
7     task: "squad-finetune"
8     model_family: "llama3.2"
9     model_size: "1b"
10     dataset: "squad"
11     framework: "automodel"

Configuration Parameters

Parameter	Type	Default	Description
`experiment_name`	str	”automodel-experiment”	Name of the MLflow experiment. All runs are grouped under this experiment.
`run_name`	str	""	Optional name for the current run. If empty, MLflow generates a unique name.
`tracking_uri`	str	null	URI of the MLflow tracking server. If null, delegates to MLflow, which currently defaults to `sqlite:///mlflow.db` for a new local store.
`artifact_location`	str	null	Location to store artifacts. If null, uses default MLflow location.
`tags`	dict		Dictionary of tags to attach to the run for organization and filtering.

Tracking URI Options

The tracking_uri parameter determines where MLflow stores experiment metadata such as runs, parameters, metrics, tags, and statuses:

Local SQLite database (default): null or sqlite:///mlflow.db
Remote tracking server: http://your-mlflow-server:5000
Shared database backend: postgresql://user:password@host:port/database

Artifacts remain separate from tracking metadata. For example, you can use SQLite for tracking while storing configuration files and models in a local directory or object store through artifact_location.

The legacy filesystem tracking backend (file:///path/to/mlruns) is in maintenance mode in MLflow and requires an explicit opt-out in current releases. Migrate existing tracking data to SQLite rather than using the filesystem escape hatch for new workflows.

For team collaboration, we recommend setting up a remote tracking server.

What Gets Logged

NeMo AutoModel automatically logs the following information to MLflow:

Metrics

Training loss at each step
Validation loss and metrics
Learning rate schedule
Gradient norms (if gradient clipping is enabled)

Parameters

Model configuration (architecture, size, pretrained checkpoint)
Training hyperparameters (learning rate, batch size, optimizer settings)
Dataset information
Parallelism configuration (DP, TP, CP settings)

Artifacts

Model checkpoints (if configured)
Training configuration files

Only rank 0 in distributed training logs to MLflow to avoid duplicate entries and reduce overhead.

Usage Example

Here’s a complete example of training with MLflow logging enabled:

Configure Your Recipe

Add the MLflow configuration to your YAML file (e.g., llama3_2_1b_squad.yaml):

1 step_scheduler:
2   global_batch_size: 64
3   local_batch_size: 8
4   ckpt_every_steps: 1000
5   val_every_steps: 10
6   num_epochs: 1
7 
8 model:
9   _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_pretrained
10   pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
11 
12 mlflow:
13   experiment_name: "llama3-squad-finetune"
14   run_name: "baseline-run-1"
15   tracking_uri: null  # Uses MLflow's default local SQLite database
16   tags:
17     task: "question-answering"
18     dataset: "squad"
19     model: "llama-3.2-1b"

Run Training

$ automodel --nproc-per-node=8 examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml

During training, you’ll see MLflow logging messages:

MLflow run started: abc123def456
View run at: sqlite:///mlflow.db/#/experiments/1/runs/abc123def456

View Results in MLflow UI

Launch the MLflow UI to visualize your experiments:

$ mlflow ui --backend-store-uri sqlite:///mlflow.db

By default, the UI runs at http://localhost:5000. Open this URL in your browser to:

Compare metrics across runs
View parameter configurations
Download artifacts
Filter and search experiments by tags

Integration with Other Loggers

MLflow can be used alongside other logging tools like Weights & Biases (WandB). Simply enable both in your configuration:

1 # Enable both MLflow and WandB
2 mlflow:
3   experiment_name: "my-experiment"
4   tags:
5     framework: "automodel"
6 
7 wandb:
8   project: "my-project"
9   entity: "my-team"
10   name: "my-run"

Both loggers will track the same metrics independently, allowing you to leverage the strengths of each platform.

Best Practices

Experiment Organization

Use descriptive experiment names: Group related runs under meaningful experiment names.
```
1 experiment_name: "llama3-squad-ablation-study"
```

Tag your runs: Add tags for easy filtering and comparison.

1 tags:
2   model_size: "1b"
3   learning_rate: "1e-5"
4   optimizer: "adam"

Use run names for variants: Differentiate runs within an experiment.
```
1 run_name: "lr-1e5-bs64"
```

Remote Tracking Server

For team collaboration, set up a shared MLflow tracking server:

1 mlflow:
2   tracking_uri: "http://mlflow-server.example.com:5000"
3   experiment_name: "team-llm-experiments"

Artifact Storage

For large-scale experiments, configure a dedicated artifact location:

1 mlflow:
2   artifact_location: "s3://my-bucket/mlflow-artifacts"

Supported storage backends include S3, Azure Blob Storage, Google Cloud Storage, and network file systems.

Performance Considerations

MLflow logging adds minimal overhead since only rank 0 logs.
Metrics are logged asynchronously to avoid blocking training.
For very frequent logging (every step), consider increasing val_every_steps to reduce I/O.

Troubleshooting

MLflow Not Installed

If you see an import error:

ImportError: MLflow is not installed. Please install it (e.g. pip install mlflow).

Recreate the released-package environment so its declared MLflow dependency is installed:

$ uv venv
$ source .venv/bin/activate
$ uv pip install "nemo-automodel"

Existing Filesystem Tracking Store

Current MLflow releases no longer enable the legacy filesystem tracking backend by default. Migrate an existing local mlruns store losslessly to SQLite:

$ mlflow migrate-filestore --source ./mlruns --target sqlite:///mlflow.db

Then point AutoModel and the MLflow UI at the migrated database:

1 mlflow:
2   tracking_uri: "sqlite:///mlflow.db"

$ mlflow ui --backend-store-uri sqlite:///mlflow.db

If migration is not possible, set MLFLOW_ALLOW_FILE_STORE=true before running AutoModel to explicitly retain the maintenance-only filesystem backend.

Connection Issues

If you can’t connect to a remote tracking server:

Verify the tracking_uri is correct
Check network connectivity and firewall rules
Ensure the tracking server is running

Missing Metrics

If metrics aren’t appearing in MLflow:

Verify you’re running on rank 0 or check rank 0 logs
Ensure the MLflow run started successfully (check for “MLflow run started” message)
Check that metrics are being computed during training