DEV Community

Muhammad Umair Ashraf
Muhammad Umair Ashraf

Posted on

Build a Fine-Tuned Model Training and Serving Pipeline on AWS

From raw training data to production-ready machine learning using Amazon SageMaker.

Building a machine learning model is only the beginning of an AI project.

In a real-world environment, we also need to prepare datasets, fine-tune models, evaluate performance, manage model versions, deploy inference endpoints, and monitor the model after deployment.

Managing these steps manually can become difficult as projects grow.

This is where Amazon SageMaker AI and AWS managed services can help.

In this article, we'll build an end-to-end machine learning workflow on AWS, using a customer support classification use case to understand each stage.

By the end, you'll understand how to move from raw training data to a deployed model endpoint using AWS services.

1. What Is Model Fine-Tuning?

Fine-tuning means taking an existing pre-trained model and adapting it to perform a specific task using your own dataset.

Instead of training a model from scratch, we start with a model that has already learned useful patterns from a large dataset.

We then train it further using task-specific examples.

For example, imagine a company receives thousands of customer support messages every day.

We want to automatically classify each message into one of four categories:

  • Billing
  • Technical Support
  • Account Access
  • General Inquiry

A pre-trained language model can be fine-tuned using labeled customer messages to recognize these categories.

Common fine-tuning approaches include:

  • Full fine-tuning: Updates most or all model parameters.
  • PEFT: Updates a smaller number of parameters.
  • LoRA: Uses low-rank adapters to reduce trainable parameters.
  • QLoRA: Combines quantization and LoRA to reduce memory requirements.
  • Continued pre-training: Continues training on domain-specific data before task-specific fine-tuning.

The right approach depends on model size, dataset, available compute, and performance requirements.

2. AWS Architecture Overview

We'll use AWS managed services to create a repeatable ML workflow.

The high-level architecture looks like this:

Amazon S3
   |
   v
SageMaker Processing
   |
   v
SageMaker Training
   |
   v
Model Evaluation
   |
   v
Quality Check
   |
   v
SageMaker Model Registry
   |
   v
Model Approval
   |
   v
SageMaker Endpoint
   |
   v
Application Predictions
   |
   v
Amazon CloudWatch
Enter fullscreen mode Exit fullscreen mode

Here's what each service does:

AWS Service Purpose
Amazon S3 Stores raw data, processed datasets, and model artifacts
SageMaker Processing Cleans and prepares training data
SageMaker Training Fine-tunes the pre-trained model
SageMaker Pipelines Automates the ML workflow
SageMaker Model Registry Manages model versions and approval
SageMaker Endpoint Serves model predictions
Amazon CloudWatch Collects logs and operational metrics
AWS IAM Controls access to AWS resources
AWS KMS Supports encryption and key management

The goal is to make the complete workflow reproducible and easier to manage.

3. Step One: Prepare the Training Dataset

Let's start with a simple customer support dataset.

Create a CSV file named customer-support.csv:

text,label
"My payment was deducted twice",Billing
"I cannot access my account",Account Access
"The application keeps crashing",Technical Support
"I want to know about your services",General Inquiry
"My invoice contains an incorrect amount",Billing
"Password reset is not working",Account Access
Enter fullscreen mode Exit fullscreen mode

A real production dataset should contain enough representative examples for every category.

Before training, check for:

  • Missing values
  • Duplicate records
  • Incorrect labels
  • Class imbalance
  • Poor-quality text
  • Personally identifiable information

Store the dataset in Amazon S3:

s3://my-ml-project/
│
├── raw-data/
│   └── customer-support.csv
│
├── processed-data/
│   ├── train/
│   └── validation/
│
├── model-artifacts/
│
└── evaluation/
Enter fullscreen mode Exit fullscreen mode

Amazon S3 provides durable storage for datasets, processed files, and trained model artifacts.

4. Step Two: Process the Data with SageMaker Processing

Before fine-tuning, we need to clean and split the dataset.

Amazon SageMaker Processing allows us to execute data preparation scripts through managed processing jobs.

Here's a simple Python example using pandas and scikit-learn:

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.read_csv("customer-support.csv")

# Remove missing and duplicate records
df = df.dropna(subset=["text", "label"])
df = df.drop_duplicates()

# Clean text
df["text"] = df["text"].astype(str).str.strip()

# Split dataset
train_df, validation_df = train_test_split(
    df,
    test_size=0.2,
    random_state=42,
    stratify=df["label"]
)

train_df.to_csv("train.csv", index=False)
validation_df.to_csv("validation.csv", index=False)
Enter fullscreen mode Exit fullscreen mode

The stratify parameter helps preserve class proportions across the training and validation datasets.

For larger projects, preprocessing may also include tokenization, text normalization, data validation, and removal of sensitive information.

The processing job can read the raw data from S3 and write the prepared datasets back to S3.

5. Step Three: Fine-Tune the Model

Now let's fine-tune a pre-trained model.

For this example, we'll use Hugging Face Transformers with distilbert-base-uncased.

DistilBERT is a smaller transformer-based model that can be useful for text classification tasks.

The following is an illustrative training script:

import pandas as pd

from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments
)

MODEL_NAME = "distilbert-base-uncased"

labels = [
    "Billing",
    "Technical Support",
    "Account Access",
    "General Inquiry"
]

label_to_id = {
    label: index
    for index, label in enumerate(labels)
}

train_df = pd.read_csv("train.csv")
validation_df = pd.read_csv("validation.csv")

train_df["label"] = train_df["label"].map(label_to_id)
validation_df["label"] = validation_df["label"].map(label_to_id)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

def tokenize(batch):
    return tokenizer(
        batch["text"],
        truncation=True,
        padding="max_length",
        max_length=128
    )

train_dataset = Dataset.from_pandas(train_df)
validation_dataset = Dataset.from_pandas(validation_df)

train_dataset = train_dataset.map(
    tokenize,
    batched=True
)

validation_dataset = validation_dataset.map(
    tokenize,
    batched=True
)

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=len(labels)
)

training_args = TrainingArguments(
    output_dir="./model-output",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    eval_strategy="epoch",
    save_strategy="epoch",
    logging_steps=50,
    report_to="none"
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=validation_dataset,
    processing_class=tokenizer
)

trainer.train()

trainer.save_model("./model-output")
tokenizer.save_pretrained("./model-output")
Enter fullscreen mode Exit fullscreen mode

This example demonstrates the core fine-tuning process.

For a SageMaker Training Job, adapt the script to the selected training container, dependency versions, input-channel paths, and expected model output directory.

Pin compatible versions of PyTorch, Transformers, and other dependencies in your training environment.

After training, SageMaker can store the model artifacts in Amazon S3.

6. Step Four: Evaluate Model Performance

A completed training job does not automatically mean the model is ready for production.

We need to evaluate its performance.

For classification, useful metrics include:

  • Accuracy
  • Precision
  • Recall
  • F1-score
  • Confusion matrix

Example evaluation code:

from sklearn.metrics import (
    accuracy_score,
    precision_recall_fscore_support
)

predictions = trainer.predict(validation_dataset)

predicted_labels = predictions.predictions.argmax(axis=1)
actual_labels = predictions.label_ids

accuracy = accuracy_score(
    actual_labels,
    predicted_labels
)

precision, recall, f1, _ = precision_recall_fscore_support(
    actual_labels,
    predicted_labels,
    average="weighted",
    zero_division=0
)

print("Accuracy:", accuracy)
print("Precision:", precision)
print("Recall:", recall)
print("F1 Score:", f1)
Enter fullscreen mode Exit fullscreen mode

We can save the evaluation results as a JSON file in S3.

Example:

{
  "accuracy": 0.94,
  "precision": 0.93,
  "recall": 0.94,
  "f1_score": 0.93
}
Enter fullscreen mode Exit fullscreen mode

These values are illustrative, not actual benchmark results.

We can define a minimum quality threshold before allowing the model to proceed.

For example:

Minimum F1 Score: 0.90

If F1 >= 0.90:
    Continue to model registration
Else:
    Stop the pipeline
Enter fullscreen mode Exit fullscreen mode

This prevents an underperforming model from automatically reaching production.

7. Step Five: Automate Everything with SageMaker Pipelines

Running every step manually is not a scalable approach.

Amazon SageMaker Pipelines helps orchestrate machine learning workflows.

A typical pipeline includes:

  1. Data processing
  2. Model training
  3. Model evaluation
  4. Quality validation
  5. Model registration

A simplified pipeline structure looks like this:

from sagemaker.workflow.pipeline import Pipeline
from sagemaker.workflow.steps import (
    ProcessingStep,
    TrainingStep
)
from sagemaker.workflow.conditions import (
    ConditionGreaterThanOrEqualTo
)
from sagemaker.workflow.condition_step import ConditionStep
from sagemaker.workflow.functions import JsonGet

# Configure these objects for your environment:
# processor, processing_args, estimator, training_args,
# evaluation_args, evaluation_report, register_step, session

processing_step = ProcessingStep(
    name="PrepareData",
    step_args=processing_args
)

training_step = TrainingStep(
    name="FineTuneModel",
    step_args=training_args
)

evaluation_step = ProcessingStep(
    name="EvaluateModel",
    step_args=evaluation_args
)

quality_check = ConditionStep(
    name="CheckModelQuality",
    conditions=[
        ConditionGreaterThanOrEqualTo(
            left=JsonGet(
                step_name=evaluation_step.name,
                property_file=evaluation_report,
                json_path="f1_score"
            ),
            right=0.90
        )
    ],
    if_steps=[register_step],
    else_steps=[]
)

pipeline = Pipeline(
    name="FineTunedModelPipeline",
    steps=[
        processing_step,
        training_step,
        evaluation_step,
        quality_check
    ],
    sagemaker_session=session
)

pipeline.upsert(
    role_arn="YOUR_SAGEMAKER_EXECUTION_ROLE_ARN"
)

pipeline.start()
Enter fullscreen mode Exit fullscreen mode

Note: This is a structural example, not a complete executable pipeline. You must configure the processing, training, evaluation, model registration, and evaluation PropertyFile objects. The registration step must be defined before it is referenced in the quality gate.

The key idea is that SageMaker Pipelines can connect all the stages into one repeatable workflow.

8. Step Six: Register and Manage Model Versions

After successful evaluation, we can register the model using Amazon SageMaker Model Registry.

Model Registry helps us manage model versions and track their lifecycle.

For each model version, we can maintain:

  • Model artifacts
  • Training information
  • Evaluation metrics
  • Model package details
  • Approval status
  • Deployment history

A typical lifecycle looks like this:

Model Training
      |
      v
Model Evaluation
      |
      v
Quality Threshold
      |
      v
Model Registry
      |
      v
Manual Approval
      |
      v
Production Deployment
Enter fullscreen mode Exit fullscreen mode

For production systems, manual approval is useful because an ML engineer or authorized reviewer can inspect the model before deployment.

9. Step Seven: Deploy the Model for Predictions

Once approved, we can deploy the model to an inference endpoint.

Amazon SageMaker supports several inference options:

Inference Type Best Use Case
Real-time inference Low-latency online predictions
Serverless inference Intermittent or variable workloads
Asynchronous inference Large requests or longer processing
Batch Transform Large offline prediction workloads

For our customer support classification application, real-time inference may be appropriate when an immediate response is required.

Here's a simplified deployment example:

from sagemaker.huggingface import HuggingFaceModel

huggingface_model = HuggingFaceModel(
    model_data="s3://YOUR_BUCKET/model/model.tar.gz",
    role="YOUR_SAGEMAKER_EXECUTION_ROLE_ARN",
    transformers_version="YOUR_SUPPORTED_VERSION",
    pytorch_version="YOUR_SUPPORTED_VERSION",
    py_version="YOUR_SUPPORTED_PYTHON_VERSION"
)

predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.m5.large"
)
Enter fullscreen mode Exit fullscreen mode

The framework versions, Python version, instance type, and inference handler must match the selected SageMaker container and model artifacts.

Example input:

{
  "text": "My payment was deducted twice"
}
Enter fullscreen mode Exit fullscreen mode

Illustrative prediction:

{
  "label": "Billing",
  "confidence": 0.97
}
Enter fullscreen mode Exit fullscreen mode

The actual response format depends on your inference handler and model implementation.

10. Step Eight: Monitor the Deployed Model

Deployment is not the end of the machine learning lifecycle.

Real-world data can change over time. Customer language, user behavior, and business requirements may evolve.

A model that performs well today may become less accurate in the future.

Amazon CloudWatch can help monitor:

  • Endpoint invocations
  • Latency
  • Errors
  • CPU and memory utilization
  • Endpoint availability
  • Application-level prediction metrics

You can also capture inference inputs and outputs, subject to privacy and data-retention requirements, for later analysis.

Monitor data quality and prediction distributions to identify possible data drift.

Important: Check the current availability of SageMaker Model Monitor in your AWS account and Region. AWS documentation notes that Model Monitor is no longer available to new customers. If it is unavailable, you can build alternative monitoring workflows using captured data, custom processing jobs, CloudWatch metrics, and your own evaluation logic.

A good monitoring strategy should cover both infrastructure health and model quality.

11. Security Best Practices

Machine learning pipelines may process sensitive customer and business information.

Keep these practices in mind:

IAM least privilege: Use dedicated IAM roles for processing, training, pipeline execution, and inference. Grant only the permissions each component requires.

Data encryption: Encrypt S3 objects and other supported storage resources. Use AWS KMS where appropriate.

Protect sensitive data: Remove or mask personally identifiable information before training whenever possible.

Private networking: Consider VPC configuration and controlled network access for workloads handling sensitive data.

Model quality gates: Require measurable evaluation criteria before registration or deployment.

Model approval: Use Model Registry approval workflows to control production releases.

Cost monitoring: Review training, processing, endpoint, and storage costs regularly.

12. Cost Optimization

The total cost of a fine-tuning pipeline depends on:

  • Training instance type and duration
  • Processing job duration
  • Model size
  • Storage consumption
  • Inference traffic
  • Endpoint configuration
  • Data transfer and monitoring requirements

Some practical optimization ideas:

  1. Start with a smaller pre-trained model.
  2. Use PEFT or LoRA where appropriate.
  3. Use smaller datasets for initial experiments.
  4. Consider managed spot training where supported and suitable.
  5. Choose serverless or asynchronous inference for workloads that do not require continuously provisioned real-time capacity.
  6. Shut down unused real-time endpoints.
  7. Review S3 storage and CloudWatch usage.

Always balance cost against model quality, latency, and availability requirements.

13. Final Architecture

Let's bring everything together:

             Amazon S3
          Raw Training Data
                 |
                 v
       SageMaker Processing
          Data Preparation
                 |
                 v
        SageMaker Training
          Model Fine-Tuning
                 |
                 v
         Model Evaluation
        Accuracy / F1 Score
                 |
                 v
          Quality Check
                 |
          Pass   |   Fail
            |    |    |
            v    |    v
       Model Registry  Stop
            |
            v
      Model Approval
            |
            v
     SageMaker Endpoint
            |
            v
      Application / API
            |
            v
       CloudWatch Monitoring
Enter fullscreen mode Exit fullscreen mode

This architecture gives us a structured and repeatable approach to managing the machine learning lifecycle.

Conclusion

Fine-tuning a model is only one part of delivering a successful AI solution.

A production-ready system also needs automated data preparation, repeatable training, evaluation, model versioning, controlled deployment, and continuous monitoring.

By combining Amazon S3, SageMaker Processing, SageMaker Training, SageMaker Pipelines, Model Registry, and SageMaker Endpoints, we can build a scalable workflow for these requirements.

The same architecture can be adapted for:

  • Customer support automation
  • Document classification
  • Sentiment analysis
  • Fraud detection
  • Product categorization
  • Domain-specific text classification
  • Intelligent enterprise applications

Don't just train a model. Build a complete pipeline that makes it reproducible, measurable, manageable, and ready for production.

Useful AWS Resources


What does your current ML workflow look like? Are you using SageMaker Pipelines, or are you still managing training and deployment manually? Share your experience in the comments!

AWS #MachineLearning #MLOps #AmazonSageMaker

Top comments (0)