Fix tensor leakage in discojs

cli/src/args.ts

@@ -444,21 +444,26 @@ export const args: BenchmarkArguments = {
       async getModel() {
         const model = await provider.modelCard.getModel();
 
-        if (unsafeArgs.learningRate !== undefined) {
-          if (!(model instanceof GPT))
-            throw new Error(
-              "learningRate override is only supported for GPT models",
+        try {
+          if (unsafeArgs.learningRate !== undefined) {
+            if (!(model instanceof GPT))
+              throw new Error(
+                "learningRate override is only supported for GPT models",
+              );
+            if (
+              !Number.isFinite(unsafeArgs.learningRate) ||
+              unsafeArgs.learningRate <= 0
+            )
+              throw new Error("learningRate must be a positive finite number");
+
+            model.setLearningRate(unsafeArgs.learningRate);
+            console.log(
+              `Overriding GPT learning rate to ${unsafeArgs.learningRate}`,
             );
-          if (
-            !Number.isFinite(unsafeArgs.learningRate) ||
-            unsafeArgs.learningRate <= 0
-          )
-            throw new Error("learningRate must be a positive finite number");
-
-          model.setLearningRate(unsafeArgs.learningRate);
-          console.log(
-            `Overriding GPT learning rate to ${unsafeArgs.learningRate}`,
-          );
+          }
+        } catch (e) {
+          model.dispose();
+          throw e;
         }
 
         return model;

cli/src/benchmark_gpt.ts

@@ -130,7 +130,7 @@ async function main(args: Required<CLIArguments>): Promise<void> {
       .batch(batchSize);
 
     // Init and train the model
-    const model = new GPT(config);
+    using model = new GPT(config);
     console.log(
       `\tmodel type ${modelType} \n\tbatch size ${batchSize} \n\tcontext length ${contextLength}`,
     );
@@ -151,7 +151,7 @@ async function main(args: Required<CLIArguments>): Promise<void> {
      * Inference benchmark
      */
   } else {
-    const model = await loadModelFromDisk(modelPath);
+    using model = await loadModelFromDisk(modelPath);
     if (!(model instanceof GPT)) {
       throw new Error("Loaded model isn't a GPT model");
     }

cli/src/evaluate_finetuned_gpt2.ts

@@ -93,6 +93,7 @@ function predictTokenLogits(
 ): tf.Tensor3D {
   const logits = tfModel.predict(inputTensor);
   if (Array.isArray(logits)) {
+    tf.dispose(logits);
     throw new Error("Expected GPT model to return a single logits tensor");
   }
   if (logits.rank !== 3) {
@@ -244,12 +245,6 @@ async function scoreContinuations(
     ...Array(maxInputLength - truncatedInputTokens.length).fill(0),
   ]);
 
-  const inputTensor = tf.tensor2d(
-    paddedInputs,
-    [paddedInputs.length, maxInputLength],
-    "int32",
-  );
-
   const targetIndexes: number[][] = [];
   const targetTokenIds: number[] = [];
   const targetOwners: number[] = [];
@@ -275,7 +270,6 @@ async function scoreContinuations(
   });
 
   if (targetIndexes.length === 0) {
-    inputTensor.dispose();
     return scoredInputs.map((scoredInput) => ({
       score: Number.NEGATIVE_INFINITY,
       promptTokens: promptTokens.length,
@@ -284,8 +278,13 @@ async function scoreContinuations(
     }));
   }
 
-  const logits = predictTokenLogits(tfModel, inputTensor);
   const targetLogProbs = tf.tidy(() => {
+    const inputTensor = tf.tensor2d(
+      paddedInputs,
+      [paddedInputs.length, maxInputLength],
+      "int32",
+    );
+    const logits = predictTokenLogits(tfModel, inputTensor);
     const targetIndexTensor = tf.tensor2d(
       targetIndexes,
       [targetIndexes.length, 2],
@@ -301,7 +300,12 @@ async function scoreContinuations(
     return tf.gatherND(logProbs, targetTokenIndexTensor);
   });
 
-  const targetScores = (await targetLogProbs.array()) as number[];
+  let targetScores: number[];
+  try {
+    targetScores = (await targetLogProbs.array()) as number[];
+  } finally {
+    targetLogProbs.dispose();
+  }
   const scoreSums = Array(scoredInputs.length).fill(0) as number[];
   const scoreCounts = Array(scoredInputs.length).fill(0) as number[];
 
@@ -321,10 +325,6 @@ async function scoreContinuations(
     usedInputTokens: scoredInput.truncatedInputTokens.length,
   }));
 
-  inputTensor.dispose();
-  logits.dispose();
-  targetLogProbs.dispose();
-
   return results;
 }
 
@@ -493,7 +493,7 @@ async function main() {
   const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2");
 
   console.log("Loading model...");
-  const model = await loadModelFromDisk(args.modelPath);
+  using model = await loadModelFromDisk(args.modelPath);
 
   if (!(model instanceof GPT)) {
     throw new Error("Model must be GPT");

cli/src/hellaswag_gpt.ts

@@ -115,7 +115,11 @@ async function main(): Promise<void> {
       model = (await modelDecode(encodedModel)) as GPT;
       break;
   }
-  await evaluateModel(model, args.numDataPoints);
+  try {
+    await evaluateModel(model, args.numDataPoints);
+  } finally {
+    model.dispose();
+  }
 
   console.log("Benchmark completed!");
 }

cli/src/measure_memorization_gpt2.ts

@@ -187,17 +187,11 @@ async function sampleGenerateGPT2(
 
   for (let i = 0; i < maxNewTokens; i++) {
     const modelInput = generated.slice(-maxContextLength);
-    const input = tf.tensor2d([modelInput], [1, modelInput.length], "int32");
-
-    const logits = tf.tidy(() => {
+    const nextTokenTensor = tf.tidy(() => {
+      const input = tf.tensor2d([modelInput], [1, modelInput.length], "int32");
       const output = tfModel.predict(input);
-      if (Array.isArray(output)) {
-        return output[0];
-      }
-      return output;
-    });
+      const logits = Array.isArray(output) ? output[0] : output;
 
-    const nextTokenTensor = tf.tidy(() => {
       const last = logits.slice([0, modelInput.length - 1, 0], [1, 1, -1]);
       const scaled = last.squeeze<tf.Tensor1D>().div(temperature);
       const { values: topKLogits, indices: topKTokens } = tf.topk(scaled, topK);
@@ -213,12 +207,12 @@ async function sampleGenerateGPT2(
       return topKTokens.gather(sampledIndex).squeeze<tf.Scalar>();
     });
 
-    const nextTokenData = await nextTokenTensor.data();
-    const nextToken = nextTokenData[0];
-
-    input.dispose();
-    logits.dispose();
-    nextTokenTensor.dispose();
+    let nextToken: number;
+    try {
+      nextToken = (await nextTokenTensor.data())[0];
+    } finally {
+      nextTokenTensor.dispose();
+    }
 
     generated.push(nextToken);
   }
@@ -349,7 +343,7 @@ async function main() {
   const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2");
 
   console.log("Loading model...");
-  const loadedModel = await loadModelFromDisk(args.modelPath);
+  using loadedModel = await loadModelFromDisk(args.modelPath);
   if (!(loadedModel instanceof GPT)) {
     throw new Error("modelPath must point to a Disco GPT model");
   }

cli/src/train_gpt.ts

@@ -27,7 +27,7 @@ async function main(): Promise<void> {
     .repeat()
     .batch(8);
 
-  const model = new GPT(config);
+  using model = new GPT(config);
   for await (const logs of model.train(tokenDataset, undefined)) {
     console.log(logs);
   }

discojs/src/aggregator.spec.ts

@@ -4,13 +4,20 @@ import {
   type Aggregator,
   MeanAggregator,
   SecureAggregator,
+  ByzantineRobustAggregator,
+  PercentileClippingAggregator,
 } from "#aggregator/index";
 import type { NodeID } from "#client/index";
 import { WeightsContainer } from "#weights/index";
 
 const AGGREGATORS: Set<[name: string, new () => Aggregator]> = Set.of<
   new () => Aggregator
->(MeanAggregator, SecureAggregator).map((Aggregator) => [
+>(
+  MeanAggregator,
+  SecureAggregator,
+  ByzantineRobustAggregator,
+  PercentileClippingAggregator,
+).map((Aggregator) => [
   // MeanAggregator waits for 100% of the node's contributions by default
   Aggregator.name,
   Aggregator,
@@ -122,11 +129,15 @@ export async function communicate<A extends Aggregator>(
       if (contribution === undefined)
         throw new Error(`no contribution for ${id}`);
 
-      for (const [to, payload] of agg.makePayloads(contribution))
-        network.get(to)?.add(id, payload.clone(), aggregationRound, r);
+      for (const [to, payload] of agg.makePayloads(contribution)) {
+        network.get(to)?.add(id, payload, aggregationRound, r);
+        payload.dispose();
+      }
     }
 
-    contributions = Map(await Promise.all(nextContributions));
+    const aggregated = Map(await Promise.all(nextContributions));
+    if (r > 0) contributions.forEach((c) => c.dispose());
+    contributions = aggregated;
   }
 
   return contributions;

discojs/src/aggregator/aggregator.ts

@@ -264,6 +264,7 @@ export abstract class Aggregator extends EventEmitter<{
 
   /**
    * Constructs the payloads sent to other nodes as contribution.
+   * The payloads are owned by the caller, who has to dispose them.
    * @param base Object from which the payload is computed
    */
   abstract makePayloads(base: WeightsContainer): Map<NodeID, WeightsContainer>;

discojs/src/aggregator/byzantine.spec.ts

@@ -1,6 +1,7 @@
 import { Set } from "immutable";
 import { describe, expect, it } from "vitest";
 import fc from "fast-check";
+import * as tf from "@tensorflow/tfjs";
 
 import { WeightsContainer } from "#weights/index";
 import { ByzantineRobustAggregator } from "#aggregator/byzantine";
@@ -342,4 +343,50 @@ describe("ByzantineRobustAggregator", () => {
 
     expect(await run()).to.be.closeTo(await run(), 1e-6);
   });
+
+  it("keeps working after the caller disposes the previous aggregate", async () => {
+    const agg = new ByzantineRobustAggregator(0, 2, "absolute", 1.0, 2, 0.5);
+    agg.setNodes(Set(["a", "b"]));
+
+    const p1 = agg.getPromiseForAggregation();
+    agg.add("a", WeightsContainer.of([1]), 0);
+    agg.add("b", WeightsContainer.of([1]), 0);
+    // the caller owns the aggregate, disposing it must not break the next round
+    (await p1).dispose();
+
+    const p2 = agg.getPromiseForAggregation();
+    agg.add("a", WeightsContainer.of([1]), 1);
+    agg.add("b", WeightsContainer.of([1]), 1);
+    const out = await p2;
+    // m = 0.5 * 1 + 0.5 * 0.5
+    expect((await WSIntoArrays(out))[0][0]).to.be.closeTo(0.75, 1e-6);
+    out.dispose();
+    agg.dispose();
+  });
+
+  it("aggregation leaves no dangling tensors", async () => {
+    const baseline = tf.memory().numTensors;
+
+    const agg = new ByzantineRobustAggregator(0, 3, "absolute", 1.0, 3, 0.5);
+    const ids = ["a", "b", "c"];
+    agg.setNodes(Set(ids));
+
+    for (let round = 0; round < 3; round++) {
+      const contributions = ids.map((_, i) =>
+        WeightsContainer.of([i, round], [100 * i]),
+      );
+      const p = agg.getPromiseForAggregation();
+      // update a contribution within the same round
+      const replaced = WeightsContainer.of([0, 0], [0]);
+      agg.add("a", replaced, round);
+      ids.forEach((id, i) => agg.add(id, contributions[i], round));
+      (await p).dispose();
+      replaced.dispose();
+      contributions.forEach((c) => c.dispose());
+    }
+
+    // the aggregator only keeps the momentums and the previous aggregate
+    agg.dispose();
+    expect(tf.memory().numTensors).to.equal(baseline);
+  });
 });