Ensuring Data Governance in LLM Fine-Tuning with Amazon SageMaker AI and Databricks Unity Catalog

Overview of the Integration Challenge

Solution Overview

Prerequisites for Implementation

Step-by-Step Walkthrough of the Fine-Tuning Process

Step 1: Setting Up AWS for Fine-Tuning

Step 2: Configuring Databricks Unity Catalog

Step 3: Setting Up EMR Serverless Applications

Step 4: Preprocessing Data with EMR Serverless

Step 5: Fine-Tuning the Model Using SageMaker AI

Step 6: Registering Model Artifacts in Unity Catalog

Step 7: Creating Data Lineage in Unity Catalog

Cleanup of Resources After Testing

Conclusion: Achieving Governed LLM Workflows

About the Authors

Fine-Tuning Large Language Models with Amazon SageMaker AI and Databricks Unity Catalog

When working with large language models (LLMs) like Ministral-3B-Instruct in a governed data ecosystem, challenges often arise, particularly concerning data governance and compliance. In this post, we’ll explore how to fine-tune LLMs using Amazon SageMaker AI while leveraging Databricks Unity Catalog, addressing those unique challenges while maintaining robust governance protocols.

Context and Challenges

Integrating Amazon SageMaker AI with Databricks Unity Catalog can raise concerns about data governance, especially when the underlying data is stored in Amazon Simple Storage Service (Amazon S3). Unity Catalog aids in managing metadata and permissions, ensuring that sensitive information is handled appropriately. However, if SageMaker AI Training jobs bypass Unity Catalog’s fine-grained authorization model, critical compliance risks can arise—especially in regulated industries. This can lead to:

• Inconsistent policy enforcement
• Audit gaps
• Compliance exposure due to the lack of visibility into the training data

To navigate these challenges, maintaining a structured integration pattern is essential. This allows organizations to ensure compliance without losing capabilities or flexibility.

A Secure Workflow for Fine-Tuning

This post outlines a secure and compliant workflow for fine-tuning LLMs. The integration of Unity Catalog with Amazon SageMaker AI, coupled with Amazon EMR Serverless for preprocessing, allows for secure data access and maintains data lineage across services.

Solution Overview

The proposed workflow accomplishes the following:

1. Reads training data from a Unity Catalog-managed table.
2. Preprocesses data using EMR Serverless with Apache Spark.
3. Fine-tunes the Ministral-3B-Instruct model using SageMaker AI.
4. Tracks data lineage in Unity Catalog from source data to the trained model.

Architecture Diagram:
(Insert diagram illustrating data flow between SageMaker AI Studio, EMR Serverless, and Databricks Unity Catalog)

Key Components and Their Roles

Walkthrough of Business Logic

To guide you through implementing this workflow, follow these steps:

Prerequisites

Before initiating the process, ensure you have the following set up in your AWS environment:

• An Amazon S3 bucket for data storage
• AWS Secrets Manager for credential management
• Required IAM roles for SageMaker and EMR

Step 1: AWS Setup

1. Create S3 Buckets
   Set up an S3 bucket with the appropriate structure (e.g., raw, curated, and ML).

2. Store Databricks Credentials
   Use AWS Secrets Manager to securely store OAuth credentials for Databricks service principals.

3. Create IAM Roles
   Implement policies that allow SageMaker and EMR access to the Unity Catalog-managed resources.

Step 2: Databricks Unity Catalog Setup

1. Configure Unity Catalog
   Create a Unity Catalog structure and grant the necessary permissions, ensuring proper governance over data access.

2. Test the Connection
   Use the Databricks SDK to confirm successful access to Unity Catalog tables.

Step 3: EMR Serverless Application Setup

1. Create an EMR Serverless Application
   Use a VPC with internet access to facilitate external resource downloading required for Delta Lake support.

Step 4: Data Pre-processing

1. Submit an EMR Serverless Job
   Create a preprocessing script that cleans and formats the risk factors from SEC EDGAR data into an instruction-style prompt.

Step 5: Fine-tuning with SageMaker AI

1. Fine-Tune the LLM
   Implement a SageMaker AI training job to fine-tune the Ministral model, utilizing memory-efficient techniques.

Step 6: Register Artifacts in Unity Catalog

1. Model Registration
   After completing training, register the model in Unity Catalog for effective management and lifecycle tracking.

Step 7: Create Data Lineage

1. Tracking Lineage
   Utilize Unity Catalog’s External Metadata and Lineage APIs to create a complete audit trail of the data and models, enhancing compliance capabilities.

Conclusion

Integrating Databricks Unity Catalog with Amazon SageMaker AI provides a robust architecture for fine-tuning large language models while maintaining governance and compliance. The structured workflow described here allows organizations to leverage the strengths of multiple services, ensuring secure data access, effective lineage tracking, and adherence to governance policies.

Get Started Today

Ready to implement this pattern? Download the notebook, deploy the reference architecture in your AWS environment, and test the workflow with a Unity Catalog-managed dataset. This approach serves as a strong foundation for developing governed, production-ready ML and generative AI workloads.

For any questions or feedback, feel free to share your thoughts in the comments!

About the Authors

• Genta Watanabe: Senior Technical Account Manager at AWS focusing on Machine Learning architectures.
• Mayank Gupta: Senior AI/ML Specialist with expertise in model development and deployment.
• Ram Vittal: Principal GenAI/ML Specialist SA at AWS with extensive experience in cloud applications.
• Venkatavaradhan Viswanathan: Global Partner Solutions Architect at AWS specializing in data and ML technologies.

This post aims to equip you with the knowledge and tools needed to navigate the challenges of fine-tuning LLMs securely while complying with strict data governance standards. Happy coding!