How to Set Up MLOps with MLflow 2.10 and PyTorch 2.5 for Llama 3.2 Fine-Tuning

pip install mlflow==2.10.0 torch==2.5.0 transformers==4.36.0 boto3 datasets

import argparse
import logging
import os
import subprocess
import sys
from pathlib import Path

# Configure logging for server setup
logging.basicConfig(
    level=logging.INFO,
    format=\"%(asctime)s - %(name)s - %(levelname)s - %(message)s\"
)
logger = logging.getLogger(__name__)

def start_mlflow_server(
    backend_store_uri: str,
    artifact_store_uri: str,
    host: str = \"0.0.0.0\",
    port: int = 5000,
    workers: int = 4
) -> None:
    \"\"\"
    Start an MLflow 2.10 tracking server with configured backend and artifact stores.

    Args:
        backend_store_uri: URI for MLflow experiment/run metadata (e.g., sqlite:///mlflow.db)
        artifact_store_uri: URI for model/artifact storage (e.g., s3://my-bucket/mlflow-artifacts)
        host: Bind host for the server
        port: Bind port for the server
        workers: Number of gunicorn workers for production use
    \"\"\"
    # Validate backend store directory exists if using local filesystem
    if backend_store_uri.startswith(\"sqlite\"):
        db_path = Path(backend_store_uri.replace(\"sqlite:///\", \"\"))
        db_path.parent.mkdir(parents=True, exist_ok=True)
        logger.info(f\"Initialized backend store directory at {db_path.parent}\")
    elif backend_store_uri.startswith(\"postgresql\"):
        logger.info(\"Using PostgreSQL backend store - ensure database is pre-created\")

    # Validate artifact store access if using S3
    if artifact_store_uri.startswith(\"s3\"):
        try:
            import boto3
            s3 = boto3.client(\"s3\")
            bucket = artifact_store_uri.split(\"/\")[2]
            s3.head_bucket(Bucket=bucket)
            logger.info(f\"Verified S3 artifact bucket {bucket} exists\")
        except ImportError:
            logger.error(\"boto3 not installed - required for S3 artifact store\")
            sys.exit(1)
        except Exception as e:
            logger.error(f\"Failed to access S3 artifact store: {e}\")
            sys.exit(1)

    # Build MLflow server command
    cmd = [
        sys.executable, \"-m\", \"mlflow\", \"server\",
        \"--backend-store-uri\", backend_store_uri,
        \"--artifact-store-uri\", artifact_store_uri,
        \"--host\", host,
        \"--port\", str(port),
        \"--workers\", str(workers),
        \"--serve-artifacts\"  # Enable artifact serving for local development
    ]

    logger.info(f\"Starting MLflow server with command: {' '.join(cmd)}\")

    try:
        # Start server as subprocess, stream logs to stdout
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True
        )
        for line in process.stdout:
            print(line.strip())
        process.wait()
    except KeyboardInterrupt:
        logger.info(\"Received shutdown signal, stopping MLflow server\")
        process.terminate()
        process.wait()
    except Exception as e:
        logger.error(f\"MLflow server crashed: {e}\")
        sys.exit(1)

if __name__ == \"__main__\":
    parser = argparse.ArgumentParser(description=\"Start MLflow 2.10 Tracking Server\")
    parser.add_argument(\"--backend-store\", type=str, default=\"sqlite:///mlflow.db\",
                        help=\"Backend store URI for experiment metadata\")
    parser.add_argument(\"--artifact-store\", type=str, default=\"s3://llama3-mlflow-artifacts\",
                        help=\"Artifact store URI for models and metrics\")
    parser.add_argument(\"--port\", type=int, default=5000, help=\"Server port\")
    parser.add_argument(\"--host\", type=str, default=\"0.0.0.0\", help=\"Server host\")
    args = parser.parse_args()

    # Check MLflow version to ensure 2.10+
    import mlflow
    if tuple(map(int, mlflow.__version__.split(\".\"))) < (2, 10, 0):
        logger.error(f\"MLflow version {mlflow.__version__} is too old. Requires 2.10+\")
        sys.exit(1)
    logger.info(f\"Using MLflow version {mlflow.__version__}\")

    start_mlflow_server(
        backend_store_uri=args.backend_store,
        artifact_store_uri=args.artifact_store,
        host=args.host,
        port=args.port
    )

import json
import logging
import os
from pathlib import Path
from typing import List, Dict, Any

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, PreTrainedTokenizer
from transformers.utils import is_torch_available

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Llama3FineTuningDataset(Dataset):
    \"\"\"
    Custom Dataset for Llama 3.2 fine-tuning with instruction-response pairs.
    Expects input JSONL file with each line: {\"instruction\": \"...\", \"response\": \"...\"}
    \"\"\"
    def __init__(
        self,
        dataset_path: str,
        tokenizer: PreTrainedTokenizer,
        max_length: int = 512,
        pad_to_max: bool = True
    ):
        self.dataset_path = Path(dataset_path)
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.pad_to_max = pad_to_max

        # Validate dataset exists
        if not self.dataset_path.exists():
            raise FileNotFoundError(f\"Dataset file not found at {dataset_path}\")
        if self.dataset_path.suffix != \".jsonl\":
            logger.warning(f\"Dataset file is not JSONL - unexpected format may cause errors\")

        # Load and validate dataset
        self.examples = []
        with open(self.dataset_path, \"r\") as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line:
                    continue
                try:
                    example = json.loads(line)
                    # Validate required keys
                    if \"instruction\" not in example or \"response\" not in example:
                        logger.error(f\"Line {line_num}: Missing instruction/response key\")
                        continue
                    self.examples.append(example)
                except json.JSONDecodeError as e:
                    logger.error(f\"Line {line_num}: Invalid JSON: {e}\")
                    continue

        logger.info(f\"Loaded {len(self.examples)} valid examples from {dataset_path}\")
        if len(self.examples) == 0:
            raise ValueError(f\"No valid examples found in {dataset_path}\")

        # Configure tokenizer for Llama 3.2
        self.tokenizer.pad_token = self.tokenizer.eos_token
        self.tokenizer.padding_side = \"right\"

    def __len__(self) -> int:
        return len(self.examples)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        example = self.examples[idx]
        instruction = example[\"instruction\"]
        response = example[\"response\"]

        # Format as Llama 3.2 instruction template
        # Llama 3 uses: <|begin_of_text|><|start_header_id|>user<|end_header_id|>...<|eot_id|><|start_header_id|>assistant<|end_header_id|>...<|eot_id|>
        prompt = f\"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{instruction}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n{response}<|eot_id|>\"

        # Tokenize
        try:
            tokenized = self.tokenizer(
                prompt,
                max_length=self.max_length,
                truncation=True,
                padding=\"max_length\" if self.pad_to_max else False,
                return_tensors=\"pt\"
            )
        except Exception as e:
            logger.error(f\"Failed to tokenize example {idx}: {e}\")
            # Return empty tensor as fallback (will be filtered in collate_fn)
            return {\"input_ids\": torch.empty(0), \"attention_mask\": torch.empty(0)}

        # Remove batch dimension added by tokenizer
        return {
            \"input_ids\": tokenized[\"input_ids\"].squeeze(0),
            \"attention_mask\": tokenized[\"attention_mask\"].squeeze(0)
        }

def collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
    \"\"\"Custom collate function to filter invalid examples and pad dynamically.\"\"\"
    # Filter out empty tensors from tokenization errors
    batch = [ex for ex in batch if ex[\"input_ids\"].numel() > 0]
    if len(batch) == 0:
        logger.warning(\"Entire batch is invalid - returning empty batch\")
        return {\"input_ids\": torch.empty(0), \"attention_mask\": torch.empty(0)}

    # Pad dynamically to max length in batch
    input_ids = torch.nn.utils.rnn.pad_sequence(
        [ex[\"input_ids\"] for ex in batch],
        batch_first=True,
        padding_value=tokenizer.pad_token_id
    )
    attention_mask = torch.nn.utils.rnn.pad_sequence(
        [ex[\"attention_mask\"] for ex in batch],
        batch_first=True,
        padding_value=0
    )
    return {\"input_ids\": input_ids, \"attention_mask\": attention_mask}

def get_llama3_dataloader(
    dataset_path: str,
    tokenizer: PreTrainedTokenizer,
    batch_size: int = 4,
    max_length: int = 512,
    num_workers: int = 2,
    shuffle: bool = True
) -> DataLoader:
    \"\"\"
    Create a PyTorch DataLoader for Llama 3.2 fine-tuning.

    Args:
        dataset_path: Path to JSONL dataset
        tokenizer: Pre-trained Llama 3.2 tokenizer
        batch_size: Batch size per GPU
        max_length: Max token length per example
        num_workers: Number of data loading workers
        shuffle: Whether to shuffle the dataset
    \"\"\"
    dataset = Llama3FineTuningDataset(
        dataset_path=dataset_path,
        tokenizer=tokenizer,
        max_length=max_length
    )

    # Check if CUDA is available for pin memory
    pin_memory = torch.cuda.is_available()
    logger.info(f\"Using pin_memory={pin_memory} for DataLoader\")

    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        collate_fn=collate_fn,
        pin_memory=pin_memory,
        drop_last=True  # Drop last incomplete batch to avoid OOM
    )

if __name__ == \"__main__\":
    # Test the dataset and dataloader
    try:
        tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Llama-3.2-1B-Instruct\")
    except Exception as e:
        logger.error(f\"Failed to load Llama 3.2 tokenizer: {e}\")
        logger.error(\"Ensure you have accepted the Hugging Face terms for Llama 3.2\")
        sys.exit(1)

    # Create sample dataset if it doesn't exist
    sample_dataset_path = \"sample_dataset.jsonl\"
    if not Path(sample_dataset_path).exists():
        sample_data = [
            {\"instruction\": \"What is MLOps?\", \"response\": \"MLOps is the practice of streamlining ML model development, deployment, and maintenance.\"},
            {\"instruction\": \"Explain PyTorch 2.5's torch.compile\", \"response\": \"torch.compile speeds up PyTorch code by converting it to optimized kernels.\"}
        ]
        with open(sample_dataset_path, \"w\") as f:
            for ex in sample_data:
                f.write(json.dumps(ex) + \"\n\")
        logger.info(f\"Created sample dataset at {sample_dataset_path}\")

    dataloader = get_llama3_dataloader(
        dataset_path=sample_dataset_path,
        tokenizer=tokenizer,
        batch_size=2
    )

    # Iterate over one batch to verify
    for batch in dataloader:
        logger.info(f\"Batch input shape: {batch['input_ids'].shape}\")
        logger.info(f\"Batch attention mask shape: {batch['attention_mask'].shape}\")
        break

import argparse
import logging
import os
import sys
from pathlib import Path

import torch
import torch.nn as nn
import torch.optim as optim
from transformers import AutoModelForCausalLM, AutoTokenizer, get_linear_schedule_with_warmup
from transformers.utils import is_torch_cuda_available

import mlflow
import mlflow.pytorch
from mlflow.models import infer_signature

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def train_llama3_mlflow(
    model_name: str = \"meta-llama/Llama-3.2-1B-Instruct\",
    dataset_path: str = \"sample_dataset.jsonl\",
    mlflow_tracking_uri: str = \"http://localhost:5000\",
    output_dir: str = \"llama3-finetuned\",
    num_epochs: int = 3,
    batch_size: int = 4,
    learning_rate: float = 2e-5,
    max_length: int = 512,
    use_torch_compile: bool = True,
    gradient_accumulation_steps: int = 2
) -> None:
    \"\"\"
    Fine-tune Llama 3.2 with PyTorch 2.5 and log all metrics to MLflow 2.10.

    Args:
        model_name: Hugging Face model name for Llama 3.2
        dataset_path: Path to JSONL training dataset
        mlflow_tracking_uri: URI of MLflow tracking server
        output_dir: Directory to save fine-tuned model
        num_epochs: Number of training epochs
        batch_size: Batch size per GPU
        learning_rate: Learning rate for optimizer
        max_length: Max token length per example
        use_torch_compile: Whether to use PyTorch 2.5's torch.compile
        gradient_accumulation_steps: Gradient accumulation steps to simulate larger batch
    \"\"\"
    # Set MLflow tracking URI
    mlflow.set_tracking_uri(mlflow_tracking_uri)
    logger.info(f\"Connected to MLflow tracking server at {mlflow_tracking_uri}\")

    # Check PyTorch version
    if tuple(map(int, torch.__version__.split(\".\"))) < (2, 5, 0):
        logger.error(f\"PyTorch version {torch.__version__} is too old. Requires 2.5+\")
        sys.exit(1)
    logger.info(f\"Using PyTorch version {torch.__version__}\")

    # Check CUDA availability
    if not is_torch_cuda_available():
        logger.warning(\"CUDA not available - training will run on CPU (40x slower)\")
        device = torch.device(\"cpu\")
    else:
        device = torch.device(\"cuda\")
        logger.info(f\"Using CUDA device: {torch.cuda.get_device_name(0)}\")

    # Load tokenizer and model
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16 if device.type == \"cuda\" else torch.float32
        )
    except Exception as e:
        logger.error(f\"Failed to load model {model_name}: {e}\")
        logger.error(\"Ensure you have accepted Hugging Face terms and are logged in via `huggingface-cli login`\")
        sys.exit(1)

    # Apply torch.compile if enabled (PyTorch 2.5 feature)
    if use_torch_compile and device.type == \"cuda\":
        try:
            model = torch.compile(model, mode=\"max-autotune\")
            logger.info(\"Applied torch.compile to model with max-autotune mode\")
        except Exception as e:
            logger.error(f\"Failed to apply torch.compile: {e}\")
            logger.warning(\"Continuing without torch.compile\")

    model.to(device)

    # Load dataset and dataloader (reuse from previous code example)
    from dataset import get_llama3_dataloader  # Assume dataset.py is in same directory
    dataloader = get_llama3_dataloader(
        dataset_path=dataset_path,
        tokenizer=tokenizer,
        batch_size=batch_size,
        max_length=max_length
    )

    # Initialize optimizer and scheduler
    optimizer = optim.AdamW(model.parameters(), lr=learning_rate)
    total_steps = len(dataloader) * num_epochs // gradient_accumulation_steps
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=int(0.1 * total_steps),
        num_training_steps=total_steps
    )

    # Start MLflow experiment
    experiment_name = \"llama3-fine-tuning\"
    mlflow.set_experiment(experiment_name)
    logger.info(f\"Using MLflow experiment: {experiment_name}\")

    with mlflow.start_run(run_name=\"llama3-1b-pytorch2.5\") as run:
        # Log all hyperparameters
        mlflow.log_params({
            \"model_name\": model_name,
            \"num_epochs\": num_epochs,
            \"batch_size\": batch_size,
            \"learning_rate\": learning_rate,
            \"max_length\": max_length,
            \"use_torch_compile\": use_torch_compile,
            \"gradient_accumulation_steps\": gradient_accumulation_steps,
            \"pytorch_version\": torch.__version__,
            \"mlflow_version\": mlflow.__version__
        })

        # Training loop
        model.train()
        global_step = 0
        for epoch in range(num_epochs):
            epoch_loss = 0.0
            for batch_idx, batch in enumerate(dataloader):
                # Skip empty batches
                if batch[\"input_ids\"].numel() == 0:
                    continue

                input_ids = batch[\"input_ids\"].to(device)
                attention_mask = batch[\"attention_mask\"].to(device)

                # Forward pass: shift labels for causal LM (predict next token)
                labels = input_ids.clone()
                labels[labels == tokenizer.pad_token_id] = -100  # Ignore pad tokens in loss

                outputs = model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    labels=labels
                )
                loss = outputs.loss / gradient_accumulation_steps  # Scale loss for gradient accumulation

                # Backward pass
                loss.backward()
                epoch_loss += loss.item() * gradient_accumulation_steps

                # Update weights every gradient_accumulation_steps
                if (batch_idx + 1) % gradient_accumulation_steps == 0:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # Prevent gradient explosion
                    optimizer.step()
                    scheduler.step()
                    optimizer.zero_grad()
                    global_step += 1

                    # Log metrics every 10 steps
                    if global_step % 10 == 0:
                        mlflow.log_metric(\"train_loss\", loss.item(), step=global_step)
                        mlflow.log_metric(\"learning_rate\", scheduler.get_last_lr()[0], step=global_step)
                        logger.info(f\"Epoch {epoch+1}, Step {global_step}, Loss: {loss.item():.4f}\")

            # Log epoch-level metrics
            avg_epoch_loss = epoch_loss / len(dataloader)
            mlflow.log_metric(\"epoch_loss\", avg_epoch_loss, step=epoch+1)
            logger.info(f\"Epoch {epoch+1} completed. Average loss: {avg_epoch_loss:.4f}\")

        # Save model and log to MLflow
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        model.save_pretrained(output_dir)
        tokenizer.save_pretrained(output_dir)
        logger.info(f\"Saved fine-tuned model to {output_dir}\")

        # Infer model signature for MLflow validation
        sample_input = tokenizer(\"What is MLOps?\", return_tensors=\"pt\").to(device)
        sample_output = model.generate(**sample_input, max_new_tokens=50)
        signature = infer_signature(
            sample_input.input_ids.cpu(),
            sample_output.cpu()
        )

        # Log model to MLflow registry
        mlflow.pytorch.log_model(
            model,
            artifact_path=\"llama3-finetuned\",
            signature=signature,
            registered_model_name=\"llama3-1b-finetuned\",
            extra_pip_requirements=[\"torch>=2.5.0\", \"transformers>=4.36.0\"]
        )
        logger.info(f\"Logged model to MLflow registry as llama3-1b-finetuned\")

        # Log sample inference output
        mlflow.log_text(
            tokenizer.decode(sample_output[0], skip_special_tokens=True),
            \"sample_inference.txt\"
        )

if __name__ == \"__main__\":
    parser = argparse.ArgumentParser(description=\"Fine-tune Llama 3.2 with MLflow and PyTorch 2.5\")
    parser.add_argument(\"--model-name\", type=str, default=\"meta-llama/Llama-3.2-1B-Instruct\")
    parser.add_argument(\"--dataset-path\", type=str, default=\"sample_dataset.jsonl\")
    parser.add_argument(\"--mlflow-uri\", type=str, default=\"http://localhost:5000\")
    parser.add_argument(\"--num-epochs\", type=int, default=3)
    parser.add_argument(\"--batch-size\", type=int, default=4)
    parser.add_argument(\"--use-torch-compile\", action=\"store_true\", default=True)
    args = parser.parse_args()

    train_llama3_mlflow(
        model_name=args.model_name,
        dataset_path=args.dataset_path,
        mlflow_tracking_uri=args.mlflow_uri,
        num_epochs=args.num_epochs,
        batch_size=args.batch_size,
        use_torch_compile=args.use_torch_compile
    )

# Enable native graph tracking (automatic in MLflow 2.10 for torch.compile models)
import mlflow.pytorch

with mlflow.start_run():
    # Train your torch.compile-optimized model
    compiled_model = torch.compile(model)
    # Log model - graph is automatically captured
    mlflow.pytorch.log_model(
        compiled_model,
        artifact_path=\"llama3-compiled\",
        registered_model_name=\"llama3-graph-tracked\"
    )

# Log torch.compile performance metrics to MLflow
import time

with mlflow.start_run():
    model = AutoModelForCausalLM.from_pretrained(\"meta-llama/Llama-3.2-1B-Instruct\")
    # Log compilation time
    start = time.time()
    compiled_model = torch.compile(model, mode=\"max-autotune\")
    compile_time = time.time() - start
    mlflow.log_metric(\"torch_compile_time\", compile_time)

    # Log per-step time for 10 warmup steps
    dataloader = get_dataloader()
    step_times = []
    for i, batch in enumerate(dataloader):
        if i >= 10: break
        start = time.time()
        # Training step here
        step_times.append(time.time() - start)
    mlflow.log_metric(\"avg_step_time\", sum(step_times)/len(step_times))

# Example MLflow webhook configuration for KServe deployment
# Save as mlflow_webhook.json and register via MLflow API
{
    \"name\": \"llama3-kserve-deploy\",
    \"registry_model_name\": \"llama3-1b-finetuned\",
    \"target_stage\": \"Production\",
    \"url\": \"https://your-webhook-receiver.com/deploy-llama\",
    \"secret\": \"your-shared-secret\",
    \"retries\": 3
}

llama3-mlflow-pytorch/
├── README.md                # Tutorial instructions and setup guide
├── requirements.txt         # Pinned dependencies (MLflow 2.10, PyTorch 2.5, etc.)
├── setup_mlflow.py          # Code example 1: MLflow server setup
├── dataset.py               # Code example 2: Dataset and DataLoader
├── train.py                 # Code example 3: Training loop with MLflow logging
├── sample_dataset.jsonl     # Sample instruction-response dataset
├── deploy/
│   ├── kserve.yaml          # KServe inference service manifest
│   └── webhook_receiver.py  # MLflow webhook receiver for automated deployment
├── benchmarks/
│   └── pytorch_2.5_vs_2.4.csv  # Benchmark results comparing PyTorch versions
└── troubleshooting.md       # Common pitfalls and solutions