Draft: GPT2 training on MCQ med data

.knip.json

@@ -12,7 +12,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": {

Dockerfile

@@ -1,7 +1,7 @@
 # TODO freeze to 22 until tfjs#8425 is merged
 FROM node:22 AS builder
 
-RUN npm i -g pnpm
+RUN npm i -g pnpm@11.11.0
 
 WORKDIR /disco
 

cli/README.md

@@ -36,14 +36,34 @@ Non-mandatory fields will automatically use values from the task specification.
 - `testID`: (mandatory) arbitrary test ID defined by the user for the test run
 - `task`: (mandatory) pre-defined task (adding a new task is described in the next section)
 - `numberOfUsers`: number of users participating in the learning round
-- `save`: whether to save the logs of the test run
+- `host`: URL of the server to connect to, defaults to `http://localhost:8080`
+- `outputPath`: path to save logs and models, defaults to `./<testID>`
+- `saveLogs`: whether to save the logs of the test run
+- `saveModel`: whether to save the trained model to disk
+- `saveCheckpoints`: whether to save each client model after every completed round/aggregation
+
+### Dataset arguments
+
+- `datasetPath`: path to the training dataset
+- `validationDatasetPath`: path to a separate validation dataset shared by all clients, takes precedence over `validationSplit`
 
 ### Learning hyperparameters
 
 - `epochs`: total number of training epochs
-- `roundDuration`: number of epochs per round
+- `roundDuration`: number of epochs per round, ignored if `roundIterations` is set.
+- `roundIterations`: number of iterations per round, takes precedence over `roundDuration`
 - `batchSize`: batch size
-- `validationSplit`: ratio of the validation set used for evaluation
+- `validationSplit`: fraction of each client's training data used for validation, ignored when `validationDatasetPath` is set; 0 disables split-based validation
+- `validationFrequency`: how often to validate. Validate the first aggregation round and every N rounds after it; defaults to every round, 0 disables validation metrics
+- `validationMode`: when to run the validation: `before` model aggregation (default), `after`, or `both`
+- `learningRate`: override the learning rate (GPT text tasks only)
+
+### Goldfish loss parameters (GPT text tasks only)
+
+- `goldfishLoss`: train with the [goldfish loss](https://arxiv.org/abs/2406.10209), which drops a subset of target tokens from the loss to mitigate memorization
+- `goldfishK`: drop modulus k, a target token is dropped if hash(context) mod k == 0
+- `goldfishH`: localized hash context length
+- `goldfishPadTokenId`: (optional) padding token id to exclude from the goldfish loss denominator
 
 ### Aggregator parameters
 
@@ -57,7 +77,7 @@ Non-mandatory fields will automatically use values from the task specification.
 
 ## Adding new tasks
 
-The CLI can be used on several pre-defined tasks: titanic, lus_covid and CIFAR10. In order
+The CLI can be used on several pre-defined tasks: `cifar10`, `lus_covid`, `mnist`, `simple_face`, `tinder_dog`, `titanic` and `goldfish` (GPT-2 fine-tuning). In order
 to understand how to add a new task have a look at [TASK.md](../docs/TASK.md).
 
 Once a new task has been defined in `discojs`, it can be loaded in [data.ts](./src/data.ts) as it is already implemented for current tasks. There are currently [multiple classes](../discojs-node/src/loaders) you can use to load data using Node.js and preprocess data: loadImagesInDir, loadCSV and loadText.

cli/package.json

@@ -9,6 +9,8 @@
     "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",
+    "measure_memorization_gpt2": "pnpm run build && node dist/measure_memorization_gpt2.js",
     "build": "tsc --build",
     "test": ": nothing"
   },
@@ -17,7 +19,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, GPT } from "@epfml/discojs";
 
 type AggregationStrategy = "mean" | "byzantine" | "secure";
 
@@ -11,14 +11,56 @@ function parseAggregator(raw: string): AggregationStrategy {
   else throw new Error(`Aggregator ${raw} is not supported.`);
 }
 
+type ValidationMode = "before" | "after" | "both";
+
+function parseValidationMode(raw: string): ValidationMode {
+  if (raw === "before" || raw === "after" || raw === "both") return raw;
+  else
+    throw new Error(
+      `Validation mode ${raw} is not supported, expected "before", "after" or "both".`,
+    );
+}
+
 export interface BenchmarkArguments {
   provider: TaskProvider<DataType, Network>;
   testID: string;
   numberOfUsers: number;
   epochs: number;
   roundDuration: number;
+  roundIterations?: number;
   batchSize: number;
+  /**
+   * Fraction of each client's training dataset reserved for validation.
+   * Ignored when `validationDatasetPath` is set. A value of 0 leaves the
+   * client without validation data unless `validationDatasetPath` is set.
+   */
   validationSplit: number;
+  /**
+   * Validate the first aggregation round and every N rounds after it. If
+   * omitted, validation runs every round; 0 disables validation metrics. This
+   * only controls when validation runs, not whether its data comes from
+   * `validationSplit` or `validationDatasetPath`.
+   */
+  validationFrequency?: number;
+  /**
+   * When to run validation relative to weight aggregation: "before" (on the
+   * local model), "after" (on the freshly aggregated global model), or "both".
+   * Defaults to "before".
+   */
+  validationMode?: ValidationMode;
+  datasetPath?: string;
+  /**
+   * Path to a separate validation dataset. When set, this dataset is shared
+   * by all clients and takes precedence over `validationSplit`, including
+   * when `validationSplit` is non-zero.
+   */
+  validationDatasetPath?: string;
+  outputPath?: string;
+  goldfishLoss: boolean;
+  goldfishK: number;
+  goldfishH: number;
+  goldfishPadTokenId?: number;
+  learningRate?: number;
 
   // DP
   epsilon?: number;
@@ -33,24 +75,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",
@@ -75,6 +117,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",
@@ -84,15 +132,87 @@ const unsafeArgs = parse<BenchmarkUnsafeArguments>(
     validationSplit: {
       type: Number,
       alias: "v",
-      description: "Validation dataset ratio",
+      description:
+        "Fraction of each client's training data used for validation. Ignored when --validationDatasetPath is set; 0 disables split-based validation.",
       defaultValue: 0.2,
     },
-    save: {
+    validationFrequency: {
+      type: Number,
+      description:
+        "Validate the first aggregation round and every N rounds after it. Defaults to every round; use 0 to disable validation metrics.",
+      optional: true,
+    },
+    validationMode: {
+      type: parseValidationMode,
+      typeLabel: "before|after|both",
+      description:
+        "When to run validation relative to weight aggregation: before (local model), after (aggregated global model), or both. Defaults to before.",
+      optional: true,
+    },
+    datasetPath: {
+      type: String,
+      alias: "d",
+      description: "Path to the dataset",
+      optional: true,
+    },
+    validationDatasetPath: {
+      type: String,
+      alias: "V",
+      description:
+        "Path to a separate validation dataset shared by all clients. Takes precedence over --validationSplit.",
+      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",
@@ -172,17 +292,18 @@ 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.titanic,
       defaultTasks.tinderDog,
       defaultTasks.mnist,
+      defaultTasks.goldfish,
     ).map(
       async (t) =>
         [(await t.getTask()).id, t] as [
           string,
-          TaskProvider<"image" | "tabular", Network>,
+          TaskProvider<"image" | "tabular" | "text", Network>,
         ],
     ),
   ),
@@ -204,6 +325,46 @@ 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;
+      task.trainingInformation.validationMode = unsafeArgs.validationMode;
+
+      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;
@@ -277,6 +438,30 @@ export const args: BenchmarkArguments = {
 
       return task;
     },
-    modelCard: provider.modelCard,
+    modelCard: {
+      card: provider.modelCard.card,
+      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",
+            );
+          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;
+      },
+    },
   },
 };