feat: MLFlow Integration for experiment tracking

README.md

@@ -16,6 +16,11 @@
   - [DPO](#dpo)
     - [DPO Single Node](#dpo-single-node)
     - [DPO Multi-node](#dpo-multi-node)
+  - [MLflow Integration](#mlflow-integration)
+    - [Installation](#installation)
+    - [Usage](#usage)
+    - [MLflow UI](#mlflow-ui)
+    - [Configuration](#configuration)
   - [Evaluation](#evaluation)
     - [Convert Model Format (Optional)](#convert-model-format-optional)
     - [Run Evaluation](#run-evaluation)
@@ -316,6 +321,58 @@ sbatch \
     ray.sub
 ```
 
+## MLflow Integration
+
+NeMo RL supports MLflow for experiment tracking and model management. MLflow provides a comprehensive platform for tracking experiments, managing models, and visualizing training progress.
+
+### Installation
+
+Install MLflow using uv:
+
+```sh
+uv pip install mlflow
+```
+
+### Usage
+
+To run training with MLflow logging, use the provided MLflow configuration:
+
+```sh
+# Run SFT training with MLflow logging
+uv run python examples/run_sft.py --config examples/configs/sft_mlflow.yaml
+```
+
+### MLflow UI
+
+After starting training, you can view the MLflow UI to monitor your experiments:
+
+```sh
+# Start MLflow UI (run in a separate terminal)
+mlflow ui --host 0.0.0.0 --port 5000
+```
+
+Then access the UI at `http://127.0.0.1:5000/` to view:
+- Training runs and experiments
+- Metrics (loss, validation metrics, etc.)
+- Hyperparameters
+- Model artifacts and checkpoints
+
+### Configuration
+
+The MLflow configuration is defined in `examples/configs/sft_mlflow.yaml`:
+
+```yaml
+logger:
+  mlflow_enabled: true
+  mlflow:
+    experiment_name: "nemo-rl-sft-experiment"
+    run_name: "sft-${data.dataset_name}-${policy.optimizer.kwargs.lr}"
+    tracking_uri: null  # Use local tracking
+    artifact_location: null  # Use default location
+```
+
+You can customize the MLflow settings by modifying these parameters or passing them as command-line overrides.
+
 ## Evaluation
 
 We provide evaluation tools to assess model capabilities.

examples/configs/sft_mlflow.yaml

@@ -0,0 +1,142 @@
+# SFT Algorithm Configuration with MLflow logging
+sft:
+  ## total number of steps to train will equal
+  ## min((max_num_epochs * len(train_dataloader)), max_num_steps)
+  max_num_epochs: 1
+  max_num_steps: 60
+
+  val_period: 10
+  val_batches: 8
+  val_global_batch_size: 32
+  val_micro_batch_size: 1
+  val_at_start: true
+  seed: 42
+
+checkpointing:
+  enabled: true
+  checkpoint_dir: "results/sft"
+  metric_name: "val_loss"
+  higher_is_better: false
+  keep_top_k: 3
+  save_period: 10
+
+policy:
+  model_name: "meta-llama/Llama-3.2-1B"
+  tokenizer:
+    name: ${policy.model_name} ## specify if you'd like to use a tokenizer different from the model's default
+    chat_template: "{% for message in messages %}{%- if message['role'] == 'system'  %}{{'Context: ' + message['content'].strip()}}{%- elif message['role'] == 'user'  %}{{' Question: ' + message['content'].strip() + ' Answer:'}}{%- elif message['role'] == 'assistant'  %}{{' ' + message['content'].strip()}}{%- endif %}{% endfor %}"
+  train_global_batch_size: 32
+  train_micro_batch_size: 1
+  max_total_sequence_length: 1024
+  precision: "bfloat16"
+  fsdp_offload_enabled: false
+  activation_checkpointing_enabled: false
+
+  dtensor_cfg:
+    enabled: true
+    cpu_offload: False
+    sequence_parallel: false
+    activation_checkpointing: false
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+    custom_parallel_plan: null
+
+  dynamic_batching:
+    enabled: false
+
+  # makes the training sequence length divisible by the tensor parallel size
+  # this is useful for sequence parallel training
+  make_sequence_length_divisible_by: ${policy.dtensor_cfg.tensor_parallel_size}
+  max_grad_norm: 1.0
+
+  optimizer:
+    name: "torch.optim.AdamW"
+    kwargs:
+      lr: 5.0e-6
+      weight_decay: 0.1
+      betas: [0.9, 0.98]
+      eps: 1e-5
+      # when using Dtensor, we need to set foreach
+      # and fused to False
+      foreach: False
+      fused: False
+
+  ## ignored since enabled=false, but needed for testing purposes
+  megatron_cfg:
+    enabled: false
+    empty_unused_memory_level: 1
+    activation_checkpointing: false
+    tensor_model_parallel_size: 2
+    pipeline_model_parallel_size: 2
+    context_parallel_size: 1
+    pipeline_dtype: ${policy.precision}
+    num_layers_in_first_pipeline_stage: null
+    num_layers_in_last_pipeline_stage: null
+    sequence_parallel: false
+
+    optimizer:
+      optimizer: "adam"
+      lr: 5.0e-6
+      min_lr: 4.9999e-6
+      weight_decay: 0.1
+      bf16: false
+      fp16: false
+      params_dtype: "float32"
+
+      #adam
+      adam_beta1: 0.9
+      adam_beta2: 0.98
+      adam_eps: 1e-5
+
+      #sgd
+      sgd_momentum: 0.9
+
+      #distributed optimizer
+      use_distributed_optimizer: true
+      use_precision_aware_optimizer: true
+
+      clip_grad: ${policy.max_grad_norm}
+
+    scheduler:
+      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
+      weight_decay_incr_style: "constant"
+      lr_decay_style: "constant"
+      lr_decay_iters: null
+      lr_warmup_iters: 50
+      lr_warmup_init: 4.9999e-6
+
+    distributed_data_parallel_config:
+      grad_reduce_in_fp32: false
+      overlap_grad_reduce: true
+      overlap_param_gather: false
+      average_in_collective: true
+      data_parallel_sharding_strategy: "optim_grads_params"
+
+data:
+  max_input_seq_length: ${policy.max_total_sequence_length}
+  dataset_name: "squad"
+  add_bos: true
+  add_eos: true
+  add_generation_prompt: false
+
+logger:
+  log_dir: "logs"  # Base directory for all logs
+  wandb_enabled: false  # Disable wandb for this example
+  tensorboard_enabled: true  # Keep tensorboard for comparison
+  mlflow_enabled: true  # Enable MLflow logging
+  monitor_gpus: true  
+  tensorboard:
+    log_dir: "tb_logs-sft-dev-${data.dataset_name}"
+  mlflow:
+    experiment_name: "nemo-rl-sft-experiment"
+    run_name: "sft-${data.dataset_name}-${policy.optimizer.kwargs.lr}"
+    tracking_uri: null  # Use local MLflow tracking server (default)
+    artifact_location: null  # Use default artifact location
+  gpu_monitoring:
+    collection_interval: 10  # How often to collect GPU usage metrics (in seconds)
+    flush_interval: 10  # How often to flush GPU usage metrics to the loggers (in seconds)
+
+cluster:
+  gpus_per_node: 1
+  num_nodes: 1