Draft: GPT2 training on MCQ med data

.knip.json

@@ -13,7 +13,10 @@
       "entry": [
         "src/benchmark_gpt.ts",
         "src/hellaswag_gpt.ts",
-        "src/train_gpt.ts"
+        "src/train_gpt.ts",
+        "src/evaluate_finetuned_gpt2_full_answer.ts",
+        "src/evaluate_finetuned_gpt2.ts",
+        "src/measure_memorization_gpt2.ts"
       ]
     },
     "onnx-converter": {

cli/package.json

@@ -9,6 +9,10 @@
     "benchmark_gpt": "pnpm run build && node dist/benchmark_gpt.js",
     "train_gpt": "pnpm run build && node dist/train_gpt.js",
     "hellaswag_gpt": "pnpm run build && node dist/hellaswag_gpt.js",
+    "eval_finetuned_gpt2": "pnpm run build && node dist/evaluate_finetuned_gpt2.js",
+    "eval_finetuned_gpt2_full_answer": "pnpm run build && node dist/evaluate_finetuned_gpt2_full_answer.js",
+    "measure_memorization_gpt2": "pnpm run build && node dist/measure_memorization_gpt2.js",
+    "finetune_gpt": "pnpm run build && node dist/finetune_gpt.js",
     "build": "tsc --build",
     "test": ": nothing"
   },
@@ -17,7 +21,9 @@
   "dependencies": {
     "@epfml/discojs": "workspace:",
     "@epfml/discojs-node": "workspace:",
+    "@tensorflow/tfjs": "catalog:",
     "@tensorflow/tfjs-node": "catalog:",
+    "debug": "catalog:",
     "immutable": "catalog:",
     "server": "workspace:"
   },

cli/src/args.ts

@@ -2,7 +2,7 @@ import { parse } from "ts-command-line-args";
 import { Map, Set } from "immutable";
 
 import type { DataType, Network, TaskProvider } from "@epfml/discojs";
-import { defaultTasks } from "@epfml/discojs";
+import { defaultTasks, models } from "@epfml/discojs";
 
 type AggregationStrategy = "mean" | "byzantine" | "secure";
 
@@ -17,8 +17,18 @@ export interface BenchmarkArguments {
   numberOfUsers: number;
   epochs: number;
   roundDuration: number;
+  roundIterations?: number;
   batchSize: number;
   validationSplit: number;
+  validationFrequency?: number;
+  datasetPath?: string;
+  validationDatasetPath?: string;
+  outputPath?: string;
+  goldfishLoss: boolean;
+  goldfishK: number;
+  goldfishH: number;
+  goldfishPadTokenId?: number;
+  learningRate?: number;
 
   // DP
   epsilon?: number;
@@ -33,24 +43,24 @@ export interface BenchmarkArguments {
   // Secure aggregator
   maxShareValue?: number;
 
-  save: boolean;
+  saveLogs: boolean;
+  saveModel: boolean;
+  saveCheckpoints: boolean;
   host: URL;
 }
 
 type BenchmarkUnsafeArguments = Omit<BenchmarkArguments, "provider"> & {
   task: string;
+  datasetPath?: string;
+  validationDatasetPath?: string;
   help?: boolean;
 };
 
 const argExample = "e.g. pnpm start -u 2 -e 3 # runs 2 users for 3 epochs";
 
 const unsafeArgs = parse<BenchmarkUnsafeArguments>(
   {
-    testID: {
-      type: String,
-      alias: "i",
-      description: "ID of the testcase",
-    },
+    testID: { type: String, alias: "i", description: "ID of the testcase" },
     task: {
       type: String,
       alias: "t",
@@ -76,6 +86,12 @@ const unsafeArgs = parse<BenchmarkUnsafeArguments>(
       description: "Round duration (in epochs)",
       defaultValue: 2,
     },
+    roundIterations: {
+      type: Number,
+      description:
+        "For GPT text tasks, aggregate every N training batches without rewinding the dataset",
+      optional: true,
+    },
     batchSize: {
       type: Number,
       alias: "b",
@@ -88,12 +104,75 @@ const unsafeArgs = parse<BenchmarkUnsafeArguments>(
       description: "Validation dataset ratio",
       defaultValue: 0.2,
     },
-    save: {
+    validationFrequency: {
+      type: Number,
+      description:
+        "Run validation every N aggregation rounds. Defaults to every round; use 0 to disable validation metrics.",
+      optional: true,
+    },
+    datasetPath: {
+      type: String,
+      alias: "d",
+      description: "Path to the dataset",
+      optional: true,
+    },
+    validationDatasetPath: {
+      type: String,
+      alias: "V",
+      description: "Path to the validation dataset",
+      optional: true,
+    },
+    outputPath: {
+      type: String,
+      alias: "o",
+      description: "Path to save logs and models. Defaults to ./<testID>",
+      optional: true,
+    },
+    goldfishLoss: {
+      type: Boolean,
+      description: "Use Goldfish loss for GPT text tasks",
+      defaultValue: false,
+    },
+    goldfishK: {
+      type: Number,
+      description:
+        "Goldfish loss drop modulus k. Drops target if hash(context) mod k == 0",
+      defaultValue: 4,
+    },
+    goldfishH: {
+      type: Number,
+      description: "Goldfish loss localized hash context length",
+      defaultValue: 13,
+    },
+    goldfishPadTokenId: {
+      type: Number,
+      description:
+        "Optional padding token id to exclude from Goldfish loss denominator",
+      optional: true,
+    },
+    learningRate: {
+      type: Number,
+      description: "Override learning rate for GPT text tasks",
+      optional: true,
+    },
+    saveLogs: {
       type: Boolean,
       alias: "s",
       description: "Save logs of benchmark",
       defaultValue: false,
     },
+    saveModel: {
+      type: Boolean,
+      alias: "m",
+      description: "Save trained model to disk",
+      defaultValue: false,
+    },
+    saveCheckpoints: {
+      type: Boolean,
+      description:
+        "Save each client model after every completed round/aggregation",
+      defaultValue: false,
+    },
     host: {
       type: (raw: string) => new URL(raw),
       typeLabel: "URL",
@@ -173,18 +252,20 @@ const unsafeArgs = parse<BenchmarkUnsafeArguments>(
 
 const supportedTasks = Map(
   await Promise.all(
-    Set.of<TaskProvider<"image" | "tabular", Network>>(
+    Set.of<TaskProvider<"image" | "tabular" | "text", Network>>(
       defaultTasks.cifar10,
       defaultTasks.lusCovid,
       defaultTasks.simpleFace,
       defaultTasks.titanic,
       defaultTasks.tinderDog,
       defaultTasks.mnist,
+      defaultTasks.privacyrun,
+      defaultTasks.centralizedGPT2FineTune,
     ).map(
       async (t) =>
         [(await t.getTask()).id, t] as [
           string,
-          TaskProvider<"image" | "tabular", Network>,
+          TaskProvider<"image" | "tabular" | "text", Network>,
         ],
     ),
   ),
@@ -206,6 +287,45 @@ export const args: BenchmarkArguments = {
       task.trainingInformation.roundDuration = unsafeArgs.roundDuration;
       task.trainingInformation.epochs = unsafeArgs.epochs;
       task.trainingInformation.validationSplit = unsafeArgs.validationSplit;
+      task.trainingInformation.roundIterations = unsafeArgs.roundIterations;
+      task.trainingInformation.validationFrequency =
+        unsafeArgs.validationFrequency;
+
+      if (unsafeArgs.goldfishLoss) {
+        if (
+          task.dataType !== "text" ||
+          task.trainingInformation.tensorBackend !== "gpt"
+        )
+          throw new Error("Goldfish loss is only supported for GPT text tasks");
+        if (!Number.isInteger(unsafeArgs.goldfishK) || unsafeArgs.goldfishK < 1)
+          throw new Error("goldfishK must be a positive integer");
+        if (!Number.isInteger(unsafeArgs.goldfishH) || unsafeArgs.goldfishH < 1)
+          throw new Error("goldfishH must be a positive integer");
+
+        task.trainingInformation.goldfishLoss = {
+          enabled: true,
+          k: unsafeArgs.goldfishK,
+          h: unsafeArgs.goldfishH,
+          padTokenId: unsafeArgs.goldfishPadTokenId,
+        };
+      }
+
+      if (unsafeArgs.learningRate !== undefined) {
+        if (
+          task.dataType !== "text" ||
+          task.trainingInformation.tensorBackend !== "gpt"
+        )
+          throw new Error(
+            "learningRate override is only supported for GPT text tasks",
+          );
+        if (
+          !Number.isFinite(unsafeArgs.learningRate) ||
+          unsafeArgs.learningRate <= 0
+        )
+          throw new Error("learningRate must be a positive finite number");
+
+        task.trainingInformation.learningRate = unsafeArgs.learningRate;
+      }
 
       const { aggregator, clippingRadius, maxIterations, beta, maxShareValue } =
         unsafeArgs;
@@ -279,6 +399,27 @@ export const args: BenchmarkArguments = {
 
       return task;
     },
-    getModel: () => provider.getModel(),
+    async getModel() {
+      const model = await provider.getModel();
+
+      if (unsafeArgs.learningRate !== undefined) {
+        if (!(model instanceof models.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}`,
+        );
+      }
+
+      return model;
+    },
   },
 };