Posted on Nov 19

Trusting AI Models With TEEs - Private, Decentralized, Verified Part 2

#ai #blockchain #verifiability #tee

In the first part, I referenced DeAI and how Oasis suggests a combination of RONL (runtime on-chain logic, Sapphire confidential EVM) and ROFL (runtime off-chain logic, framework) components, with SGX and TDX TEEs, pave the groundwork for trustworthy AI models.

Here, I will show fine-tuning of an LLM in Oasis TEE and publishing provenance information on-chain.

Fine-tuning an LLM in a TEE

In this experiment, we will try to integrate a GPU-enabled TEE with Oasis. We will use fine-tuning for a sample machine learning (ML) task where we will teach a new fact to an existing foundation model.

This ML task is set up using a server with AMD SEV-SNP and an NVIDIA H100 to run a confidential VM in accordance with NVIDIA’s deployment guide. The confidential VM shares an initialization script with the SecLM project, verifying its starting disk content and the GPU’s attestation.

The task at hand is to train a Low-Rank Adaptation (LoRA) adapter for Meta’s Llama 3 8B Instruct model inside the confidential VM by using Hugging Face’s Transformers library and Parameter-Efficient Fine-Tuning library.

An overview of the demo setup:

Key steps taken to execute the experiment:

The confidential VM launches with a specific configuration and a set of software to boot from: the firmware (OVMF.fd), the kernel (vmlinuz-...), and an initial RAM disk (initrd.img-...). This configuration and software are covered in the confidential VM’s “measurement". This will be part of the attestation report.
A script in the initial RAM disk (fs_setup_shasum_check) verifies the contents of two disk images against hashes in the kernel command line. The two disk images are a well-known general operating system disk image (ubuntu-....img) and a “seed” image (user-data.iso) containing application-specific configuration (user-data.cfg).
A script from the seed image runs an initialization script (init_cvm_gpu.sh) to set up NVIDIA’s nvtrust software and verify the local GPU.
A second script from the seed image receives the machine learning job’s code (job.tar.gz), verifies it against a hash, and executes it.
The ML task (job/job.sh) generates a key for the confidential VM and generates an attestation report bearing the fingerprint of this key.
The task then runs a Python program (job/train.py) to do the fine-tuning.
In the final step, the task creates a signed “manifest” file with information about the trained LoRA adapter, base model, and training data, and signs this file with the confidential VM’s private key.

Contents of the attestation report:

Pre-trained model's query-response (Querying stock Meta-Llama-3-8B-Instruct):

Training data (data.jsonl)

Querying the fine-tuned model

This ML task example using the confidential VM took, on average, 30 seconds across five trials, while an equivalently configured task running on a machine without confidentiality took 12 seconds on average. Point to note here is that the measurements are indicative for this particular example only, and will likely vary while performing other ML tasks.

Publishing Provenance Information On-chain

As mentioned in the first part of this series, verified provenance for AI models on-chain is seen as foreshadowing decentralized marketplaces for AI models and services. So, let's now build a sample ROFL application to validate and verify models built using our experimental setup.

The assumption here is that models are published in a repository. Now, the ROFL application continuously monitors for new entries. Detection of a new model triggers a series of verification steps.

It validates the Versioned Chip Endorsement Key (VCEK) by checking the certificate chain all the way up to the AMD root key.
It validates the guest trusted computing base by examining specific fields in the attestation report.
It verifies that the signature within the attestation report genuinely originates from the VCEK.
It checks that the confidential VM’s public key matches the one in the attestation report.
It confirms that the result manifest is signed by the confidential VM’s private key.

Only after successfully completing all these checks, the ROFL application will publish the hash values of the model and integrate the manifest report into a Sapphire smart contract. As the verification takes place in ROFL, the integrity of its execution is ensured. Users can be confident about the provenance and authenticity, thanks to the ability to access the details of how models were constructed on-chain.

Example model manifest is not in the contract before running ROFL: