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 your environment, install it manually:
```
$ pip install mlflow
$ # or:
$ uv pip install mlflow
```
MLflow tracking server (optional): For production use, set up a tracking server to centralize experiment data. For local development, MLflow will use a local file-based store 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, uses local file-based storage.
`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 data:

Local file storage (default): null or file:///path/to/mlruns
Remote tracking server: http://your-mlflow-server:5000
Database backend: postgresql://user:password@host:port/database

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 local storage
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: file:///path/to/mlruns/#/experiments/1/runs/abc123def456

View Results in MLflow UI

Launch the MLflow UI to visualize your experiments:

$ mlflow ui

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).

Install MLflow:

$ pip install mlflow
$ # or:
$ uv pip install mlflow

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