Support loading large model weights

tfjs-backend-webgpu/src/softmax_webgpu.ts

@@ -59,16 +59,17 @@ export class SoftmaxProgram implements WebGPUProgram {
       }
       workgroupBarrier();
 
-      let reduceSize = min(cols, blockSize);
-      for (var currSize = reduceSize >> 1;  currSize > 0; currSize = currSize >> 1) {
+      var reduceSize = min(cols, blockSize);
+      for (var currSize = reduceSize >> 1;  currSize > 0; currSize = reduceSize >> 1) {
+        reduceSize = currSize + (reduceSize & 1);
         if (tid < currSize) {
-          buf[tid] = max(buf[tid], buf[tid + currSize]);
+          buf[tid] = max(buf[tid], buf[tid + reduceSize]);
         }
         workgroupBarrier();
       }
 
       if (tid == 0) {
-        rowMaxShared = max(buf[0], buf[reduceSize - 1]);
+        rowMaxShared = buf[0];
       }
       workgroupBarrier();
 

tfjs-converter/src/executor/graph_model_test.ts

@@ -586,7 +586,8 @@ describe('Model', () => {
         expect(handler.savedArtifacts.modelTopology).toEqual(CUSTOM_OP_MODEL);
         expect(handler.savedArtifacts.weightSpecs).toEqual(weightsManifest);
         tfc.test_util.expectArraysClose(
-            new Int32Array(handler.savedArtifacts.weightData), bias.dataSync());
+            new Int32Array(io.concatenateArrayBuffers(
+                handler.savedArtifacts.weightData)), bias.dataSync());
       });
     });
   });
@@ -616,7 +617,8 @@ describe('Model', () => {
         });
         expect(handler.savedArtifacts.weightSpecs).toEqual(weightsManifest);
         tfc.test_util.expectArraysClose(
-            new Int32Array(handler.savedArtifacts.weightData), bias.dataSync());
+          new Int32Array(io.concatenateArrayBuffers(
+            handler.savedArtifacts.weightData)), bias.dataSync());
       });
     });
 
@@ -904,7 +906,8 @@ describe('Model', () => {
         });
         expect(handler.savedArtifacts.weightSpecs).toEqual(weightsManifest);
         tfc.test_util.expectArraysClose(
-            new Int32Array(handler.savedArtifacts.weightData), bias.dataSync());
+          new Int32Array(io.concatenateArrayBuffers(handler.savedArtifacts
+            .weightData)), bias.dataSync());
       });
     });
 

tfjs-core/src/io/browser_files.ts

@@ -25,7 +25,7 @@ import {env} from '../environment';
 
 import {basename, concatenateArrayBuffers, getModelArtifactsForJSON, getModelArtifactsInfoForJSON, getModelJSONForModelArtifacts} from './io_utils';
 import {IORouter, IORouterRegistry} from './router_registry';
-import {IOHandler, ModelArtifacts, ModelJSON, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
+import {IOHandler, ModelArtifacts, ModelJSON, SaveResult, WeightData, WeightsManifestConfig, WeightsManifestEntry} from './types';
 
 const DEFAULT_FILE_NAME_PREFIX = 'model';
 const DEFAULT_JSON_EXTENSION_NAME = '.json';
@@ -70,8 +70,13 @@ export class BrowserDownloads implements IOHandler {
           'Browser downloads are not supported in ' +
           'this environment since `document` is not present');
     }
+
+    // TODO(mattsoulanille): Support saving models over 2GB that exceed
+    // Chrome's ArrayBuffer size limit.
+    const weightBuffer = concatenateArrayBuffers(modelArtifacts.weightData);
+
     const weightsURL = window.URL.createObjectURL(new Blob(
-        [modelArtifacts.weightData], {type: 'application/octet-stream'}));
+        [weightBuffer], {type: 'application/octet-stream'}));
 
     if (modelArtifacts.modelTopology instanceof ArrayBuffer) {
       throw new Error(
@@ -169,7 +174,7 @@ class BrowserFiles implements IOHandler {
   }
 
   private loadWeights(weightsManifest: WeightsManifestConfig): Promise<[
-    /* weightSpecs */ WeightsManifestEntry[], /* weightData */ ArrayBuffer
+    /* weightSpecs */ WeightsManifestEntry[], WeightData,
   ]> {
     const weightSpecs: WeightsManifestEntry[] = [];
     const paths: string[] = [];
@@ -185,7 +190,7 @@ class BrowserFiles implements IOHandler {
         paths.map(path => this.loadWeightsFile(path, pathToFile[path]));
 
     return Promise.all(promises).then(
-        buffers => [weightSpecs, concatenateArrayBuffers(buffers)]);
+        buffers => [weightSpecs, buffers]);
   }
 
   private loadWeightsFile(path: string, file: File): Promise<ArrayBuffer> {

tfjs-core/src/io/browser_files_test.ts

@@ -23,6 +23,7 @@ import * as tf from '../index';
 import {BROWSER_ENVS, describeWithFlags} from '../jasmine_util';
 import {browserDownloads, BrowserDownloads, browserDownloadsRouter} from './browser_files';
 import {WeightsManifestConfig, WeightsManifestEntry} from './types';
+import {concatenateArrayBuffers} from './io_utils';
 
 const modelTopology1: {} = {
   'class_name': 'Sequential',
@@ -310,7 +311,7 @@ describeWithFlags('browserFiles', BROWSER_ENVS, () => {
     expect(modelArtifacts.modelInitializer).toEqual({});
     expect(modelArtifacts.trainingConfig).toEqual(trainingConfig1);
 
-    expect(new Uint8Array(modelArtifacts.weightData))
+    expect(new Uint8Array(concatenateArrayBuffers(modelArtifacts.weightData)))
         .toEqual(new Uint8Array(weightData1));
   });
 
@@ -351,9 +352,10 @@ describeWithFlags('browserFiles', BROWSER_ENVS, () => {
     const modelArtifacts = await filesHandler.load();
     expect(modelArtifacts.modelTopology).toEqual(modelTopology1);
     expect(modelArtifacts.weightSpecs).toEqual(weightSpecs);
-    expect(new Uint8Array(modelArtifacts.weightData)).toEqual(new Uint8Array([
-      1, 2, 3, 4, 10, 20, 30, 40
-    ]));
+    expect(new Uint8Array(concatenateArrayBuffers(modelArtifacts.weightData)))
+        .toEqual(new Uint8Array([
+            1, 2, 3, 4, 10, 20, 30, 40
+        ]));
   });
 
   it(`Two groups, four paths, reverseOrder=false`, async () => {
@@ -418,9 +420,10 @@ describeWithFlags('browserFiles', BROWSER_ENVS, () => {
     expect(modelArtifacts.modelTopology).toEqual(modelTopology1);
     expect(modelArtifacts.weightSpecs)
         .toEqual(weightSpecs1.concat(weightSpecs2));
-    expect(new Uint8Array(modelArtifacts.weightData)).toEqual(new Uint8Array([
-      1, 3, 5, 7, 10, 30, 50, 70, 2, 4, 6, 8, 20, 40, 60, 80
-    ]));
+    expect(new Uint8Array(concatenateArrayBuffers(modelArtifacts.weightData)))
+        .toEqual(new Uint8Array([
+            1, 3, 5, 7, 10, 30, 50, 70, 2, 4, 6, 8, 20, 40, 60, 80
+        ]));
   });
 
   it(`Two groups, four paths, reverseOrder=true`, async () => {
@@ -485,9 +488,10 @@ describeWithFlags('browserFiles', BROWSER_ENVS, () => {
     expect(modelArtifacts.modelTopology).toEqual(modelTopology1);
     expect(modelArtifacts.weightSpecs)
         .toEqual(weightSpecs1.concat(weightSpecs2));
-    expect(new Uint8Array(modelArtifacts.weightData)).toEqual(new Uint8Array([
-      1, 3, 5, 7, 10, 30, 50, 70, 2, 4, 6, 8, 20, 40, 60, 80
-    ]));
+    expect(new Uint8Array(concatenateArrayBuffers(modelArtifacts.weightData)))
+        .toEqual(new Uint8Array([
+            1, 3, 5, 7, 10, 30, 50, 70, 2, 4, 6, 8, 20, 40, 60, 80
+        ]));
   });
 
   it('Upload model topology only', async () => {

tfjs-core/src/io/composite_array_buffer.ts

@@ -0,0 +1,206 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+import {TypedArray} from '../types';
+import * as util from '../util';
+
+type BufferShard = {
+  start: number,
+  end: number,
+  buffer: ArrayBuffer,
+};
+
+/**
+ * Wraps a list of ArrayBuffers into a `slice()`-able object without allocating
+ * a large ArrayBuffer.
+ *
+ * Allocating large ArrayBuffers (~2GB) can be unstable on Chrome. TFJS loads
+ * its weights as a list of (usually) 4MB ArrayBuffers and then slices the
+ * weight tensors out of them. For small models, it's safe to concatenate all
+ * the weight buffers into a single ArrayBuffer and then slice the weight
+ * tensors out of it, but for large models, a different approach is needed.
+ */
+
+export class CompositeArrayBuffer {
+  private shards: BufferShard[] = [];
+  private previousShardIndex = 0;
+  private bufferUniformSize?: number;
+  public readonly byteLength: number;
+
+  constructor(buffers: ArrayBuffer | ArrayBuffer[] | TypedArray |
+    TypedArray[]) {
+    // Normalize the `buffers` input to be `ArrayBuffer[]`.
+    if (!(buffers instanceof Array)) {
+      buffers = [buffers];
+    }
+    buffers = buffers.map((bufferOrTypedArray) => {
+      if (util.isTypedArray(bufferOrTypedArray)) {
+        return bufferOrTypedArray.buffer;
+      }
+      return bufferOrTypedArray;
+    });
+
+    // Skip setting up shards if there are no buffers.
+    if (buffers.length === 0) {
+      return;
+    }
+
+    this.bufferUniformSize = buffers[0].byteLength;
+    let start = 0;
+
+    for (let i = 0; i < buffers.length; i++) {
+      const buffer = buffers[i];
+      // Check that all buffers except the last one have the same length.
+      if (i !== buffers.length - 1 &&
+        buffer.byteLength !== this.bufferUniformSize) {
+        // Unset the buffer uniform size, since the buffer sizes are not
+        // uniform.
+        this.bufferUniformSize = undefined;
+      }
+
+      // Create the shards, including their start and end points.
+      const end = start + buffer.byteLength;
+      this.shards.push({ buffer, start, end });
+      start = end;
+    }
+
+    // Set the byteLenghth
+    if (this.shards.length === 0) {
+      this.byteLength = 0;
+    }
+    this.byteLength = this.shards[this.shards.length - 1].end;
+  }
+
+  slice(start = 0, end = this.byteLength): ArrayBuffer {
+    // NaN is treated as zero for slicing. This matches ArrayBuffer's behavior.
+    start = isNaN(Number(start)) ? 0 : start;
+    end = isNaN(Number(end)) ? 0 : end;
+
+    // Fix the bounds to within the array.
+    start = Math.max(0, start);
+    end = Math.min(this.byteLength, end);
+    if (end <= start) {
+      return new ArrayBuffer(0);
+    }
+
+    const startShardIndex = this.findShardForByte(start);
+    if (startShardIndex === -1) {
+      // This should not happen since the start and end indices are always
+      // within 0 and the composite array's length.
+      throw new Error(`Could not find start shard for byte ${start}`);
+    }
+
+    const size = end - start;
+    const outputBuffer = new ArrayBuffer(size);
+    const outputArray = new Uint8Array(outputBuffer);
+    let sliced = 0;
+    for (let i = startShardIndex; i < this.shards.length; i++) {
+      const shard = this.shards[i];
+
+      const globalStart = start + sliced;
+      const localStart = globalStart - shard.start;
+      const outputStart = sliced;
+
+      const globalEnd = Math.min(end, shard.end);
+      const localEnd = globalEnd - shard.start;
+
+      const outputSlice = new Uint8Array(shard.buffer.slice(localStart,
+        localEnd));
+      outputArray.set(outputSlice, outputStart);
+      sliced += outputSlice.length;
+
+      if (end < shard.end) {
+        break;
+      }
+    }
+    return outputBuffer;
+  }
+
+  /**
+   * Get the index of the shard that contains the byte at `byteIndex`.
+   */
+  private findShardForByte(byteIndex: number): number {
+    if (this.shards.length === 0 || byteIndex < 0 ||
+      byteIndex >= this.byteLength) {
+      return -1;
+    }
+
+    // If the buffers have a uniform size, compute the shard directly.
+    if (this.bufferUniformSize != null) {
+      this.previousShardIndex = Math.floor(byteIndex / this.bufferUniformSize);
+      return this.previousShardIndex;
+    }
+
+    // If the buffers don't have a uniform size, we need to search for the
+    // shard. That means we need a function to check where the byteIndex lies
+    // relative to a given shard.
+    function check(shard: BufferShard) {
+      if (byteIndex < shard.start) {
+        return -1;
+      }
+      if (byteIndex >= shard.end) {
+        return 1;
+      }
+      return 0;
+    }
+
+    // For efficiency, try the previous shard first.
+    if (check(this.shards[this.previousShardIndex]) === 0) {
+      return this.previousShardIndex;
+    }
+
+    // Otherwise, use a generic search function.
+    // This should almost never end up being used in practice since the weight
+    // entries should always be in order.
+    const index = search(this.shards, check);
+    if (index === -1) {
+      return -1;
+    }
+
+    this.previousShardIndex = index;
+    return this.previousShardIndex;
+  }
+}
+
+/**
+ * Search for an element of a sorted array.
+ *
+ * @param sortedArray The sorted array to search
+ * @param compare A function to compare the current value against the searched
+ *     value. Return 0 on a match, negative if the searched value is less than
+ *     the value passed to the function, and positive if the searched value is
+ *     greater than the value passed to the function.
+ * @returns The index of the element, or -1 if it's not in the array.
+ */
+export function search<T>(sortedArray: T[], compare: (t: T) => number): number {
+  // Binary search
+  let min = 0;
+  let max = sortedArray.length;
+
+  while (min <= max) {
+    const middle = Math.floor((max - min) / 2) + min;
+    const side = compare(sortedArray[middle]);
+
+    if (side === 0) {
+      return middle;
+    } else if (side < 0) {
+      max = middle;
+    } else {
+      min = middle + 1;
+    }
+  }
+  return -1;
+}