Lưu các mô hình AI vào bộ nhớ đệm trong trình duyệt

Hầu hết các mô hình AI đều có một điểm chung: chúng khá lớn đối với một tài nguyên được chuyển qua Internet. Mô hình phát hiện đối tượng MediaPipe nhỏ nhất (SSD MobileNetV2 float16) có kích thước 5,6 MB và lớn nhất là khoảng 25 MB.

LLM nguồn mở gemma-2b-it-gpu-int4.bin có dung lượng 1,35 GB và đây được coi là dung lượng rất nhỏ đối với một LLM. Mô hình AI tạo sinh có thể rất lớn. Đây là lý do tại sao nhiều hoạt động sử dụng AI hiện nay diễn ra trong đám mây. Ngày càng có nhiều ứng dụng chạy các mô hình được tối ưu hoá cao ngay trên thiết bị. Mặc dù đã có các bản minh hoạ về LLM chạy trong trình duyệt, nhưng sau đây là một số ví dụ cấp độ sản xuất về các mô hình khác chạy trong trình duyệt:

Adobe Photoshop trên web với công cụ chọn đối tượng dựa trên AI đang mở, trong đó có 3 đối tượng được chọn: 2 con hươu cao cổ và 1 mặt trăng.

Để các lần khởi chạy ứng dụng trong tương lai diễn ra nhanh hơn, bạn nên lưu dữ liệu mô hình vào bộ nhớ đệm trên thiết bị một cách rõ ràng, thay vì dựa vào bộ nhớ đệm trình duyệt HTTP ngầm ẩn.

Mặc dù hướng dẫn này sử dụng gemma-2b-it-gpu-int4.bin model để tạo chatbot, nhưng phương pháp này có thể được khái quát hoá để phù hợp với các mô hình và trường hợp sử dụng khác trên thiết bị. Cách phổ biến nhất để kết nối một ứng dụng với một mô hình là phân phát mô hình cùng với các tài nguyên ứng dụng còn lại. Điều quan trọng là bạn phải tối ưu hoá việc phân phối.

Định cấu hình đúng tiêu đề bộ nhớ đệm

Nếu bạn phân phát mô hình AI từ máy chủ, điều quan trọng là phải định cấu hình đúng tiêu đề Cache-Control. Ví dụ sau đây cho thấy một chế độ cài đặt mặc định chắc chắn mà bạn có thể xây dựng dựa trên nhu cầu của ứng dụng.

Cache-Control: public, max-age=31536000, immutable

Mỗi phiên bản phát hành của mô hình AI là một tài nguyên tĩnh. Nội dung không bao giờ thay đổi phải được cấp một max-age dài kết hợp với tính năng đập vỡ bộ nhớ đệm trong URL yêu cầu. Nếu cần cập nhật mô hình, bạn phải cung cấp URL mới cho mô hình đó.

Khi người dùng tải lại trang, ứng dụng sẽ gửi yêu cầu xác thực lại, mặc dù máy chủ biết rằng nội dung đã ổn định. Chỉ thị immutable cho biết rõ rằng không cần xác thực lại vì nội dung sẽ không thay đổi. Lệnh immutable không được hỗ trợ rộng rãi bởi các trình duyệt và bộ nhớ đệm trung gian hoặc máy chủ proxy, nhưng bằng cách kết hợp lệnh này với lệnh max-age được hiểu rộng rãi, bạn có thể đảm bảo khả năng tương thích tối đa. Lệnh phản hồi public cho biết rằng phản hồi có thể được lưu trữ trong bộ nhớ đệm dùng chung.

Công cụ cho nhà phát triển Chrome hiển thị các tiêu đề Cache-Control sản xuất do Hugging Face gửi khi yêu cầu một mô hình AI. (Nguồn)

Lưu mô hình AI vào bộ nhớ đệm phía máy khách

Khi phân phát mô hình AI, điều quan trọng là bạn phải lưu mô hình vào bộ nhớ đệm trong trình duyệt một cách rõ ràng. Điều này đảm bảo dữ liệu mô hình luôn sẵn có sau khi người dùng tải lại ứng dụng.

Bạn có thể sử dụng một số kỹ thuật để đạt được điều này. Đối với các mẫu mã sau đây, giả sử mỗi tệp mô hình được lưu trữ trong một đối tượng Blob có tên là blob trong bộ nhớ.

Để hiểu rõ hiệu suất, mỗi mã mẫu được chú thích bằng phương thức performance.mark()performance.measure(). Các biện pháp này phụ thuộc vào thiết bị và không thể áp dụng chung.

Trong Công cụ của Chrome cho nhà phát triển Ứng dụng > Bộ nhớ, hãy xem lại sơ đồ sử dụng có các phân đoạn cho IndexedDB, Bộ nhớ đệm và Hệ thống tệp. Mỗi phân đoạn được cho là tiêu thụ 1354 megabyte dữ liệu, tổng cộng là 4063 megabyte.

Bạn có thể chọn sử dụng một trong các API sau để lưu các mô hình AI vào bộ nhớ đệm trong trình duyệt: API bộ nhớ đệm, API hệ thống tệp riêng tư gốcAPI IndexedDB. Nhìn chung, bạn nên sử dụng Cache API, nhưng hướng dẫn này sẽ thảo luận về ưu và khuyết điểm của tất cả các tuỳ chọn.

API bộ nhớ đệm

Cache API (API bộ nhớ đệm) cung cấp bộ nhớ ổn định cho các cặp đối tượng RequestResponse được lưu vào bộ nhớ đệm trong bộ nhớ dài hạn. Mặc dù API này được xác định trong thông số kỹ thuật của Worker, nhưng bạn có thể sử dụng API này từ luồng chính hoặc worker thông thường. Để sử dụng đối tượng này bên ngoài ngữ cảnh của worker dịch vụ, hãy gọi phương thức Cache.put() với đối tượng Response tổng hợp, được ghép nối với một URL tổng hợp thay vì đối tượng Request.

Hướng dẫn này giả định một blob trong bộ nhớ. Sử dụng URL giả làm khoá bộ nhớ đệm và Response tổng hợp dựa trên blob. Nếu tải trực tiếp mô hình xuống, bạn sẽ sử dụng Response mà bạn nhận được khi tạo yêu cầu fetch().

Ví dụ: dưới đây là cách lưu trữ và khôi phục tệp mô hình bằng API bộ nhớ đệm.

const storeFileInSWCache = async (blob) => {
  try {
    performance.mark('start-sw-cache-cache');
    const modelCache = await caches.open('models');
    await modelCache.put('model.bin', new Response(blob));
    performance.mark('end-sw-cache-cache');

    const mark = performance.measure(
      'sw-cache-cache',
      'start-sw-cache-cache',
      'end-sw-cache-cache'
    );
    console.log('Model file cached in sw-cache.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromSWCache = async () => {
  try {
    performance.mark('start-sw-cache-restore');
    const modelCache = await caches.open('models');
    const response = await modelCache.match('model.bin');
    if (!response) {
      throw new Error(`File model.bin not found in sw-cache.`);
    }
    const file = await response.blob();
    performance.mark('end-sw-cache-restore');
    const mark = performance.measure(
      'sw-cache-restore',
      'start-sw-cache-restore',
      'end-sw-cache-restore'
    );
    console.log(mark.name, mark.duration.toFixed(2));
    console.log('Cached model file found in sw-cache.');
    return file;
  } catch (err) {    
    throw err;
  }
};

API Hệ thống tệp riêng tư của Origin

Hệ thống tệp riêng của nguồn gốc (OPFS) là một tiêu chuẩn tương đối mới cho điểm cuối bộ nhớ. Đây là tệp riêng tư đối với nguồn gốc của trang và do đó người dùng sẽ không nhìn thấy tệp này, không giống như hệ thống tệp thông thường. Thư viện này cung cấp quyền truy cập vào một tệp đặc biệt được tối ưu hoá cao cho hiệu suất và cung cấp quyền ghi vào nội dung của tệp đó.

Ví dụ: sau đây là cách lưu trữ và khôi phục tệp mô hình trong OPFS.

const storeFileInOPFS = async (blob) => {
  try {
    performance.mark('start-opfs-cache');
    const root = await navigator.storage.getDirectory();
    const handle = await root.getFileHandle('model.bin', { create: true });
    const writable = await handle.createWritable();
    await blob.stream().pipeTo(writable);
    performance.mark('end-opfs-cache');
    const mark = performance.measure(
      'opfs-cache',
      'start-opfs-cache',
      'end-opfs-cache'
    );
    console.log('Model file cached in OPFS.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromOPFS = async () => {
  try {
    performance.mark('start-opfs-restore');
    const root = await navigator.storage.getDirectory();
    const handle = await root.getFileHandle('model.bin');
    const file = await handle.getFile();
    performance.mark('end-opfs-restore');
    const mark = performance.measure(
      'opfs-restore',
      'start-opfs-restore',
      'end-opfs-restore'
    );
    console.log('Cached model file found in OPFS.', mark.name, mark.duration.toFixed(2));
    return file;
  } catch (err) {    
    throw err;
  }
};

IndexedDB API

IndexedDB là một tiêu chuẩn đã được thiết lập để lưu trữ dữ liệu tuỳ ý một cách ổn định trong trình duyệt. IndexedDB nổi tiếng với API hơi phức tạp, nhưng bằng cách sử dụng thư viện trình bao bọc như idb-keyval, bạn có thể coi IndexedDB như một kho khoá-giá trị cổ điển.

Ví dụ:

import { get, set } from 'https://cdn.jsdelivr.net/npm/idb-keyval@latest/+esm';

const storeFileInIDB = async (blob) => {
  try {
    performance.mark('start-idb-cache');
    await set('model.bin', blob);
    performance.mark('end-idb-cache');
    const mark = performance.measure(
      'idb-cache',
      'start-idb-cache',
      'end-idb-cache'
    );
    console.log('Model file cached in IDB.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromIDB = async () => {
  try {
    performance.mark('start-idb-restore');
    const file = await get('model.bin');
    if (!file) {
      throw new Error('File model.bin not found in IDB.');
    }
    performance.mark('end-idb-restore');
    const mark = performance.measure(
      'idb-restore',
      'start-idb-restore',
      'end-idb-restore'
    );
    console.log('Cached model file found in IDB.', mark.name, mark.duration.toFixed(2));
    return file;
  } catch (err) {    
    throw err;
  }
};

Đánh dấu bộ nhớ là bộ nhớ cố định

Gọi navigator.storage.persist() ở cuối bất kỳ phương thức lưu vào bộ nhớ đệm nào trong số này để yêu cầu quyền sử dụng bộ nhớ cố định. Phương thức này trả về một lời hứa phân giải thành true nếu quyền được cấp và false nếu không. Trình duyệt có thể tuân thủ hoặc không tuân thủ yêu cầu, tuỳ thuộc vào các quy tắc dành riêng cho trình duyệt.

if ('storage' in navigator && 'persist' in navigator.storage) {
  try {
    const persistent = await navigator.storage.persist();
    if (persistent) {
      console.log("Storage will not be cleared except by explicit user action.");
      return;
    }
    console.log("Storage may be cleared under storage pressure.");  
  } catch (err) {
    console.error(err.name, err.message);
  }
}

Trường hợp đặc biệt: Sử dụng mô hình trên ổ đĩa cứng

Bạn có thể tham chiếu các mô hình AI ngay từ ổ đĩa cứng của người dùng thay vì bộ nhớ của trình duyệt. Kỹ thuật này có thể giúp các ứng dụng tập trung vào nghiên cứu giới thiệu khả thi của việc chạy các mô hình nhất định trong trình duyệt hoặc cho phép nghệ sĩ sử dụng các mô hình tự huấn luyện trong các ứng dụng sáng tạo chuyên nghiệp.

API Truy cập hệ thống tệp

Với API Truy cập hệ thống tệp, bạn có thể mở các tệp từ ổ đĩa cứng và lấy FileSystemFileHandle mà bạn có thể lưu vào IndexedDB.

Với mẫu này, người dùng chỉ cần cấp quyền truy cập vào tệp mô hình một lần. Nhờ quyền ổn định, người dùng có thể chọn cấp quyền truy cập vĩnh viễn vào tệp. Sau khi tải lại ứng dụng và một cử chỉ bắt buộc của người dùng, chẳng hạn như nhấp chuột, bạn có thể khôi phục FileSystemFileHandle từ IndexedDB bằng quyền truy cập vào tệp trên ổ đĩa cứng.

Các quyền truy cập vào tệp sẽ được truy vấn và yêu cầu nếu cần, giúp quá trình này diễn ra liền mạch trong các lần tải lại sau này. Ví dụ sau đây cho biết cách lấy handle cho một tệp từ ổ đĩa cứng, sau đó lưu trữ và khôi phục handle đó.

import { fileOpen } from 'https://cdn.jsdelivr.net/npm/browser-fs-access@latest/dist/index.modern.js';
import { get, set } from 'https://cdn.jsdelivr.net/npm/idb-keyval@latest/+esm';

button.addEventListener('click', async () => {
  try {
    const file = await fileOpen({
      extensions: ['.bin'],
      mimeTypes: ['application/octet-stream'],
      description: 'AI model files',
    });
    if (file.handle) {
      // It's an asynchronous method, but no need to await it.
      storeFileHandleInIDB(file.handle);
    }
    return file;
  } catch (err) {
    if (err.name !== 'AbortError') {
      console.error(err.name, err.message);
    }
  }
});

const storeFileHandleInIDB = async (handle) => {
  try {
    performance.mark('start-file-handle-cache');
    await set('model.bin.handle', handle);
    performance.mark('end-file-handle-cache');
    const mark = performance.measure(
      'file-handle-cache',
      'start-file-handle-cache',
      'end-file-handle-cache'
    );
    console.log('Model file handle cached in IDB.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromFileHandle = async () => {
  try {
    performance.mark('start-file-handle-restore');
    const handle = await get('model.bin.handle');
    if (!handle) {
      throw new Error('File handle model.bin.handle not found in IDB.');
    }
    if ((await handle.queryPermission()) !== 'granted') {
      const decision = await handle.requestPermission();
      if (decision === 'denied' || decision === 'prompt') {
        throw new Error(Access to file model.bin.handle not granted.');
      }
    }
    const file = await handle.getFile();
    performance.mark('end-file-handle-restore');
    const mark = performance.measure(
      'file-handle-restore',
      'start-file-handle-restore',
      'end-file-handle-restore'
    );
    console.log('Cached model file handle found in IDB.', mark.name, mark.duration.toFixed(2));
    return file;
  } catch (err) {    
    throw err;
  }
};

Các phương thức này không loại trừ lẫn nhau. Có thể xảy ra trường hợp bạn vừa lưu một mô hình vào bộ nhớ đệm trong trình duyệt một cách rõ ràng vừa sử dụng một mô hình từ ổ đĩa cứng của người dùng.

Bản minh hoạ

Bạn có thể thấy cả ba phương thức lưu trữ trường hợp thông thường và phương thức ổ đĩa cứng được triển khai trong bản minh hoạ MediaPipe LLM.

Phần thưởng: Tải tệp lớn xuống theo từng phần

Nếu bạn cần tải một mô hình AI lớn xuống từ Internet, hãy tải xuống song song thành các phần riêng biệt, sau đó ghép lại với nhau trên ứng dụng.

Sau đây là một hàm trợ giúp mà bạn có thể sử dụng trong mã của mình. Bạn chỉ cần truyền url vào đó. chunkSize (mặc định: 5 MB), maxParallelRequests (mặc định: 6), hàm progressCallback (báo cáo về downloadedBytes và tổng fileSize) và signal cho tín hiệu AbortSignal đều là không bắt buộc.

Bạn có thể sao chép hàm sau trong dự án hoặc cài đặt gói fetch-in-chunks từ gói npm.

async function fetchInChunks(
  url,
  chunkSize = 5 * 1024 * 1024,
  maxParallelRequests = 6,
  progressCallback = null,
  signal = null
) {
  // Helper function to get the size of the remote file using a HEAD request
  async function getFileSize(url, signal) {
    const response = await fetch(url, { method: 'HEAD', signal });
    if (!response.ok) {
      throw new Error('Failed to fetch the file size');
    }
    const contentLength = response.headers.get('content-length');
    if (!contentLength) {
      throw new Error('Content-Length header is missing');
    }
    return parseInt(contentLength, 10);
  }

  // Helper function to fetch a chunk of the file
  async function fetchChunk(url, start, end, signal) {
    const response = await fetch(url, {
      headers: { Range: `bytes=${start}-${end}` },
      signal,
    });
    if (!response.ok && response.status !== 206) {
      throw new Error('Failed to fetch chunk');
    }
    return await response.arrayBuffer();
  }

  // Helper function to download chunks with parallelism
  async function downloadChunks(
    url,
    fileSize,
    chunkSize,
    maxParallelRequests,
    progressCallback,
    signal
  ) {
    let chunks = [];
    let queue = [];
    let start = 0;
    let downloadedBytes = 0;

    // Function to process the queue
    async function processQueue() {
      while (start < fileSize) {
        if (queue.length < maxParallelRequests) {
          let end = Math.min(start + chunkSize - 1, fileSize - 1);
          let promise = fetchChunk(url, start, end, signal)
            .then((chunk) => {
              chunks.push({ start, chunk });
              downloadedBytes += chunk.byteLength;

              // Update progress if callback is provided
              if (progressCallback) {
                progressCallback(downloadedBytes, fileSize);
              }

              // Remove this promise from the queue when it resolves
              queue = queue.filter((p) => p !== promise);
            })
            .catch((err) => {              
              throw err;              
            });
          queue.push(promise);
          start += chunkSize;
        }
        // Wait for at least one promise to resolve before continuing
        if (queue.length >= maxParallelRequests) {
          await Promise.race(queue);
        }
      }

      // Wait for all remaining promises to resolve
      await Promise.all(queue);
    }

    await processQueue();

    return chunks.sort((a, b) => a.start - b.start).map((chunk) => chunk.chunk);
  }

  // Get the file size
  const fileSize = await getFileSize(url, signal);

  // Download the file in chunks
  const chunks = await downloadChunks(
    url,
    fileSize,
    chunkSize,
    maxParallelRequests,
    progressCallback,
    signal
  );

  // Stitch the chunks together
  const blob = new Blob(chunks);

  return blob;
}

export default fetchInChunks;

Chọn phương thức phù hợp với bạn

Hướng dẫn này đã khám phá nhiều phương pháp để lưu các mô hình AI vào bộ nhớ đệm một cách hiệu quả trong trình duyệt. Đây là một nhiệm vụ quan trọng để nâng cao trải nghiệm người dùng và hiệu suất của ứng dụng. Nhóm bộ nhớ của Chrome đề xuất Cache API để đạt được hiệu suất tối ưu, đảm bảo truy cập nhanh vào các mô hình AI, giảm thời gian tải và cải thiện khả năng phản hồi.

OPFS và IndexedDB là các tuỳ chọn ít hữu dụng hơn. OPFS và API IndexedDB cần chuyển đổi tuần tự dữ liệu trước khi có thể lưu trữ. IndexedDB cũng cần phải chuyển đổi tuần tự dữ liệu khi truy xuất, khiến đây trở thành nơi tệ nhất để lưu trữ các mô hình lớn.

Đối với các ứng dụng chuyên biệt, API Truy cập hệ thống tệp cung cấp quyền truy cập trực tiếp vào các tệp trên thiết bị của người dùng, phù hợp với những người dùng quản lý mô hình AI của riêng họ.

Nếu bạn cần bảo mật mô hình AI, hãy giữ mô hình đó trên máy chủ. Sau khi lưu trữ trên máy khách, bạn có thể dễ dàng trích xuất dữ liệu từ cả Bộ nhớ đệm và IndexedDB bằng DevTools hoặc tiện ích DevTools OFPS. Các API bộ nhớ này vốn có tính bảo mật như nhau. Bạn có thể muốn lưu trữ phiên bản đã mã hoá của mô hình, nhưng sau đó, bạn cần phải chuyển khoá giải mã đến ứng dụng. Khoá này có thể bị chặn. Điều này có nghĩa là kẻ xấu sẽ khó đánh cắp mô hình của bạn hơn một chút, nhưng không phải là không thể.

Bạn nên chọn một chiến lược lưu vào bộ nhớ đệm phù hợp với các yêu cầu của ứng dụng, hành vi của đối tượng mục tiêu và đặc điểm của các mô hình AI được sử dụng. Điều này đảm bảo ứng dụng của bạn phản hồi nhanh và mạnh mẽ trong nhiều điều kiện mạng và các quy tắc ràng buộc của hệ thống.


Lời cảm ơn

Bài viết này đã được Joshua Bell, Reilly Grant, Evan Stade, Nathan Memmott, Austin Sullivan, Etienne Noël, André Bandarra, Alexandra Klepper, François Beaufort, Paul Kinlan và Rachel Andrew xem xét.